From a28a0466b54e62d23755307c0343ae28c810a662 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:59:59 +0800 Subject: [PATCH 001/126] feat(promql): define native histogram semantics (#8758) * feat(promql): define native histogram semantics Signed-off-by: shuiyisong * chore: change author of the rfc Signed-off-by: shuiyisong * chore: update rfc Signed-off-by: shuiyisong * chore: update rfc Signed-off-by: shuiyisong * chore: update comments and tests Signed-off-by: shuiyisong * chore: fix issues Signed-off-by: shuiyisong * chore: add test Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- docs/rfcs/2026-08-04-native-histograms.md | 231 +++++++++++++ src/common/query/src/native_histogram.rs | 304 +++++++++++++----- src/promql/src/extension_plan.rs | 24 ++ .../src/extension_plan/instant_manipulate.rs | 116 +++++-- src/promql/src/extension_plan/normalize.rs | 143 +++++++- src/promql/src/extension_plan/test_util.rs | 21 ++ src/promql/src/functions/native_histogram.rs | 262 +++++++++++++-- src/query/src/promql/planner.rs | 102 ++++-- 8 files changed, 1023 insertions(+), 180 deletions(-) create mode 100644 docs/rfcs/2026-08-04-native-histograms.md diff --git a/docs/rfcs/2026-08-04-native-histograms.md b/docs/rfcs/2026-08-04-native-histograms.md new file mode 100644 index 00000000000..1e7692bf516 --- /dev/null +++ b/docs/rfcs/2026-08-04-native-histograms.md @@ -0,0 +1,231 @@ +--- +Feature Name: Native Histogram Support and Compatibility Decisions +Tracking Issue: TBD +Date: 2026-08-04 +Author: codex +--- + +# Summary + +GreptimeDB stores a Prometheus native histogram in one Struct-valued field and +evaluates it as a first-class PromQL sample. This document records the supported +protocols, the storage invariant, and the compatibility decisions needed to make +that path predictable. + +Native histograms are experimental. Prometheus Remote Write 2.0 is the only +supported native-histogram ingestion protocol. Remote Write 1.0 histogram +payloads are rejected instead of being acknowledged and dropped. +Native-histogram Remote Read is deferred; the existing Remote Read path +continues to return scalar samples only. + +# Goals + +1. Prevent silent native-histogram data loss at protocol boundaries. +2. Keep one stable persisted representation for integer, float, exponential, + and custom-bucket native histograms. +3. Match Prometheus query behavior where it is observable and practical. +4. State intentional limitations explicitly so incomplete behavior is not + mistaken for support. + +# Non-Goals + +- Supporting native histograms in Remote Write 1.0. +- Returning native histograms through Prometheus Remote Read. +- Persisting Remote Write metric metadata. +- Ingesting OTLP exponential histograms. +- Persisting exemplars. +- Removing mixed float/histogram handling from PromQL expressions. +- Propagating PromQL annotations produced on datanodes back to the frontend. + +# Data Model and Invariants + +Each native histogram is stored in the configured native-histogram field +(`greptime_native_histogram` by default), a Struct whose children preserve the +wire-level schema, zero threshold, sum, reset hint, start timestamp, custom +bounds, spans, and either the integer or float count family. +Integer bucket deltas are converted to absolute integer counts for storage; +float bucket counts are already absolute. No separate sample-kind discriminator +is stored because the populated count family identifies it. + +Within one resolved catalog, schema, and physical-table routing context, a +metric name has exactly one persisted sample kind: + +- a float metric uses `greptime_value`; +- a native-histogram metric uses the configured native-histogram field. + +Labels do not change that choice. Remote Write 2.0 rejects a request when two +label sets use different sample kinds for the same metric, and the existing +table schema rejects kind changes across requests. Prometheus metadata type is +not used for this decision: a classic histogram family is described as a +histogram but is represented by float-valued `_bucket`, `_sum`, and `_count` +series. + +The storage invariant does not remove the need for mixed-sample PromQL support. +An expression can combine different metrics or branches, for example with +`or`, and therefore produce a vector containing float samples and histogram +samples from different series. + +# Protocol Decisions + +## Remote Write 1.0 + +Remote Write 1.0 scalar samples remain supported. Any `TimeSeries.histograms` +field causes the complete request to fail with an invalid-arguments response. +Rejecting at protobuf decode time prevents both histogram-only data loss and +partial ingestion of a request that mixes scalar and histogram series. Clients +that send native histograms must use Remote Write 2.0. + +## Remote Write 2.0 + +Remote Write 2.0 accepts integer and float native histograms while +`http.experimental_enable_prometheus_native_histogram` is enabled. Supported +exponential schemas are `-4` through `8`; schema `-53` represents native +histograms with custom buckets. + +Exponential spans may end at the schema's overflow bucket but must not continue +beyond it. The overflow index is `(1024 << schema) + 1` for non-negative schemas +and `(1024 >> -schema) + 1` for negative schemas. Inputs beyond that index are +rejected rather than clamped or merged into infinity. + +Exemplars and declared metric metadata are accepted on the wire but are not +persisted by this feature. + +## Remote Read + +Native-histogram Remote Read is deferred. GreptimeDB's existing Remote Read path +continues to negotiate `SAMPLES` responses and serialize scalar samples only. + +A follow-up can map the stored Struct into `TimeSeries.histograms`. It must +preserve integer counts without passing through `f64`, retain the stored +histogram fields, and decide sampled versus streamed response coverage. The +sampled protobuf has no start-timestamp field, so that omission must remain +explicit if sampled responses are implemented. + +## OTLP + +OTLP exponential histograms are not converted in this version. The ingestion +branch intentionally remains deferred until it can map temporality, scale, +reset behavior, attributes, and rejected-point reporting into the canonical +native-histogram path. The code carries an explicit TODO rather than a partial +encoder. + +# PromQL Compatibility Decisions + +## Staleness + +A native histogram is stale only when its `sum` has Prometheus's stale-NaN bit +pattern. Ordinary NaN sums remain samples. Range selectors remove stale markers +before range functions run. Instant selectors treat the marker as the end of +the series for lookback purposes and do not resurrect an older sample. + +## Start timestamps + +Reset detection already uses consecutive start timestamps in addition to +bucket and count monotonicity. Counter `rate` and `increase` also use the first +histogram's start timestamp as a synthetic zero when it is nonzero and strictly +inside the query window before the first sample. This permits the same +single-sample calculation as Prometheus and prevents left extrapolation past the +known start. When no usable start timestamp exists, a positive histogram-count +increase lets counter extrapolation infer a zero point from the first count to +cap left extrapolation; a synthetic zero uses its known timestamp instead of +that heuristic. + +This behavior is native-histogram-only. GreptimeDB does not persist start +timestamps for float samples, so float `rate` and `increase` still require two +samples. Warnings for overlapping start timestamps are deferred. + +## Arithmetic and averages + +Histogram addition and subtraction always reconcile layouts, including empty +histograms. Mixing exponential and custom layouts fails; custom layouts are +reconciled according to their shared bounds. The only layout bypass is local to +counter rate: when the first-to-second pair is a reset, the first sample's +layout is irrelevant and the second sample starts the accumulated segment. + +`avg` and `avg_over_time` use a mergeable weighted running mean. This prevents a +finite mean from becoming infinity solely because the intermediate sum +overflowed. Native-histogram aggregates are not currently split into datanode +partial aggregation and frontend state merging, so this state does not cross +nodes or versions. If distributed stepping is added later, its state must be +versioned before rolling upgrades can safely mix implementations. + +The rejected alternative is Prometheus's direct-sum/Kahan-compensation design +with an overflow-triggered mode change. It offers closer rounding parity but +requires additional state and substantially more merge logic. + +## Equality, changes, and resets + +Equality follows Prometheus's represented-layout semantics. Histograms compare +equal only when their represented bucket-index sequences and bucket value bit +patterns match. An explicitly represented zero bucket therefore differs from an +omitted bucket, and `changes()` observes that layout change. Redundant +zero-length span encodings are normalized, malformed layouts compare unequal, +and reset hints and start timestamps remain excluded. + +Custom bucket bounds use ordinary floating-point equality rather than bitwise +equality, matching Prometheus's `CustomBucketBoundsMatch`. Histogram payload +values continue to use bitwise equality so NaN payloads behave deterministically. + +`resets()` counts ordinary counter resets and, unlike Prometheus, also counts +each transition between gauge and non-gauge native histograms as one reset. + +## Annotations + +Warnings and infos produced by frontend-executed PromQL functions are returned +through the corresponding Prometheus JSON response fields. Annotations produced +while a pushed-down plan executes on a datanode are not yet transported back to +the frontend. The sample-dropping behavior is unchanged, but those remote +annotations remain silent until a query-result transport is defined for them. + +## Standard deviation and bucket bounds + +Custom-bucket standard deviation and variance use the arithmetic midpoint +`(lower + upper) / 2` for every bucket. Underflow and overflow buckets therefore +naturally produce infinite or NaN estimates. Exponential buckets retain the +signed geometric midpoint, with zero used for a bucket spanning zero. + +Boundary calculation returns the last finite boundary as `f64::MAX`, the +overflow boundary as positive infinity, and rejects indices beyond the +overflow bucket. Widened integer arithmetic prevents large negative-schema +indices from wrapping into an unrelated finite boundary. + +# Metric Metadata + +The Prometheus `/metadata` endpoint currently derives metric names from logical +table names and reads type and unit from semantic table options stamped at table +creation. Native-histogram tables without a declared type fall back to +`histogram`. Help is always empty, and unit is absent for tables without a +semantic unit option. + +Persisting declared Remote Write metadata is separate work. Reusing table +options alone is insufficient because metadata-only writes, later updates, and +existing tables do not pass through auto-create. An in-memory registry would +lose data on restart. The recommended follow-up is a persistent registry keyed +by catalog, schema, and metric family, populated by Remote Write 1.0 metadata-only +requests and Remote Write 2.0 per-series metadata. This follows the update +boundary already recorded in +[Table Semantic Layer](2026-05-28-table-semantic-layer.md). + +# Testing and Compatibility + +Compatibility coverage includes protocol rejection, kind exclusivity, +stale-marker selector semantics, synthetic-zero rates, incompatible empty +layouts, overflow-safe averages, layout-sensitive equality, infinite custom +midpoints, and exponential overflow indices for schemas `-4`, `0`, and `8`. +Behavioral coverage also verifies frontend PromQL warning and info responses. + +This work does not change GreptimeDB's persisted Struct, protobuf dependencies, +or public configuration. Existing native-histogram data remains readable. + +# Future Work + +- Native-histogram Remote Read, including exact integer round-trips and streamed + chunks with start timestamps. +- OTLP exponential-histogram conversion with partial-success reporting. +- Persistent Remote Write metadata and accurate help/unit updates. +- Native-histogram exemplars and exemplar query APIs. +- Start-timestamp overlap annotations. +- Datanode-to-frontend PromQL annotation propagation. +- Versioned two-phase native-histogram aggregation state. +- Exact Prometheus summation compensation if measured precision differences + justify the extra aggregate state. diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index 3fdec9b309a..931fb452ede 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -306,17 +306,6 @@ impl NativeHistogram { /// Returns `None` when the layouts are invalid or cannot be reconciled. pub fn add(&self, other: &Self) -> Option { let reset_hint = add_reset_hint(self.reset_hint, other.reset_hint); - if self.is_empty_payload() { - let mut result = other.clone(); - result.reset_hint = reset_hint; - return result.compact(); - } - if other.is_empty_payload() { - let mut result = self.clone(); - result.reset_hint = reset_hint; - return result.compact(); - } - let (left, right) = self.reconcile(other)?; left.combine_exact(&right, reset_hint, |left, right| left + right)? .compact() @@ -326,13 +315,6 @@ impl NativeHistogram { /// /// Returns `None` when the layouts are invalid or cannot be reconciled. pub fn sub(&self, other: &Self) -> Option { - if self.is_empty_payload() { - return other.clone().negated().compact(); - } - if other.is_empty_payload() { - return self.clone().into_gauge().compact(); - } - let (left, right) = self.reconcile(other)?; left.combine_exact(&right, CounterResetHint::Gauge, |left, right| left - right)? .compact() @@ -372,7 +354,7 @@ impl NativeHistogram { /// Compares PromQL-visible payload values by their exact bit patterns. /// - /// Reset hints, start timestamps, and sparse zero placement are ignored. + /// Reset hints and start timestamps are ignored. pub fn promql_eq(&self, other: &Self) -> bool { self.schema == other.schema && self.zero_threshold == other.zero_threshold @@ -380,13 +362,13 @@ impl NativeHistogram { && self.count.to_bits() == other.count.to_bits() && self.zero_count.to_bits() == other.zero_count.to_bits() && self.sum.to_bits() == other.sum.to_bits() - && side_counts_equal( + && side_layout_equal( &self.positive_spans, &self.positive_buckets, &other.positive_spans, &other.positive_buckets, ) - && side_counts_equal( + && side_layout_equal( &self.negative_spans, &self.negative_buckets, &other.negative_spans, @@ -792,18 +774,12 @@ impl NativeHistogram { } fn bucket_midpoint(&self, bucket: &Bucket) -> f64 { - if bucket.lower == f64::NEG_INFINITY && bucket.upper == f64::INFINITY { - return 0.0; - } - if bucket.lower == f64::NEG_INFINITY { - return bucket.upper; - } - if bucket.upper == f64::INFINITY { - return bucket.lower; - } - if self.uses_custom_buckets() || (bucket.lower <= 0.0 && bucket.upper >= 0.0) { + if self.uses_custom_buckets() { return (bucket.lower + bucket.upper) / 2.0; } + if bucket.lower <= 0.0 && bucket.upper >= 0.0 { + return 0.0; + } if bucket.upper < 0.0 { -((bucket.lower.abs() * bucket.upper.abs()).sqrt()) } else { @@ -846,14 +822,6 @@ impl NativeHistogram { _ => None, } } - - fn is_empty_payload(&self) -> bool { - self.count == 0.0 - && self.zero_count == 0.0 - && self.sum == 0.0 - && self.positive_buckets.iter().all(|count| *count == 0.0) - && self.negative_buckets.iter().all(|count| *count == 0.0) - } } fn reconcile_exponential( @@ -1011,7 +979,10 @@ fn side_bucket_indices(spans: &[Span]) -> Option> { let mut indices = Vec::new(); let mut current_index = 0i32; let mut first = true; - for span in spans { + for (span_index, span) in spans.iter().enumerate() { + if span_index > 0 && span.offset < 0 { + return None; + } if first { current_index = span.offset; first = false; @@ -1026,6 +997,12 @@ fn side_bucket_indices(spans: &[Span]) -> Option> { Some(indices) } +fn span_bucket_len(spans: &[Span]) -> Option { + spans + .iter() + .try_fold(0usize, |sum, span| sum.checked_add(span.length as usize)) +} + fn spans_from_indices_counts(values: Vec<(i32, f64)>) -> Option<(Vec, Vec)> { let mut spans = Vec::::new(); let mut buckets = Vec::new(); @@ -1069,27 +1046,28 @@ fn side_counts(spans: &[Span], buckets: &[f64]) -> Option> { Some(values) } -fn side_counts_equal( +fn side_layout_equal( left_spans: &[Span], left_buckets: &[f64], right_spans: &[Span], right_buckets: &[f64], ) -> bool { - match ( - side_counts(left_spans, left_buckets), - side_counts(right_spans, right_buckets), - ) { - (Some(left), Some(right)) => { - left.len() == right.len() - && left.iter().all(|(idx, count)| { - right - .get(idx) - .is_some_and(|other| count.to_bits() == other.to_bits()) - }) - } - (None, None) => true, - _ => false, + if span_bucket_len(left_spans) != Some(left_buckets.len()) + || span_bucket_len(right_spans) != Some(right_buckets.len()) + { + return false; } + let Some(left_indices) = side_bucket_indices(left_spans) else { + return false; + }; + let Some(right_indices) = side_bucket_indices(right_spans) else { + return false; + }; + left_indices == right_indices + && left_buckets + .iter() + .zip(right_buckets) + .all(|(left, right)| left.to_bits() == right.to_bits()) } fn merge_side( @@ -1194,6 +1172,19 @@ fn ceil_div(value: i32, divisor: i32) -> i32 { value.div_euclid(divisor) + i32::from(value.rem_euclid(divisor) != 0) } +/// Returns the index of the exponential bucket containing positive infinity. +pub fn exponential_overflow_bucket_index(schema: i32) -> Option { + if !(MIN_EXPONENTIAL_SCHEMA..=MAX_EXPONENTIAL_SCHEMA).contains(&schema) { + return None; + } + let last_finite = if schema >= 0 { + 1024_i64.checked_shl(schema as u32)? + } else { + 1024_i64.checked_shr((-schema) as u32)? + }; + i32::try_from(last_finite.checked_add(1)?).ok() +} + fn get_bound(idx: i32, schema: i32, custom_values: &[f64]) -> Option { if schema == CUSTOM_BUCKETS_SCHEMA { return match idx { @@ -1206,23 +1197,25 @@ fn get_bound(idx: i32, schema: i32, custom_values: &[f64]) -> Option { }; } - if !(MIN_EXPONENTIAL_SCHEMA..=MAX_EXPONENTIAL_SCHEMA).contains(&schema) { + let overflow_index = exponential_overflow_bucket_index(schema)?; + if idx > overflow_index { return None; } + if idx == overflow_index { + return Some(f64::INFINITY); + } + if idx == overflow_index - 1 { + return Some(f64::MAX); + } if schema < 0 { - let exponent = idx.checked_shl((-schema) as u32)?; - return Some(if exponent == f64::MAX_EXP { - f64::MAX - } else { - 2.0_f64.powi(exponent) - }); + let exponent = i64::from(idx).checked_shl((-schema) as u32)?; + let Ok(exponent) = i32::try_from(exponent) else { + return Some(0.0); + }; + return Some(2.0_f64.powi(exponent)); } let exponent = idx as f64 / (1u32 << schema) as f64; - Some(if exponent == f64::MAX_EXP as f64 { - f64::MAX - } else { - 2.0_f64.powf(exponent) - }) + Some(2.0_f64.powf(exponent)) } /// Returns the Arrow type for a complete native histogram value. @@ -1336,12 +1329,9 @@ fn read_spans(offsets: Vec, lengths: Vec, name: &str) -> DfResult DfResult<()> { - let span_len = spans - .iter() - .try_fold(0usize, |sum, span| sum.checked_add(span.length as usize)) - .ok_or_else(|| { - DataFusionError::Execution(format!("native histogram {name} spans overflow")) - })?; + let span_len = span_bucket_len(spans).ok_or_else(|| { + DataFusionError::Execution(format!("native histogram {name} spans overflow")) + })?; if span_len != buckets { return Err(DataFusionError::Execution(format!( "native histogram {name} spans describe {span_len} buckets, found {buckets}" @@ -1719,7 +1709,7 @@ mod tests { } #[test] - fn add_combines_reset_hints_for_empty_payloads() { + fn add_combines_reset_hints_for_compatible_empty_payloads() { let mut empty = histogram(Vec::new(), Vec::new()); empty.reset_hint = CounterResetHint::Gauge; let mut populated = histogram( @@ -1741,6 +1731,45 @@ mod tests { ); } + #[test] + fn empty_arithmetic_rejects_incompatible_schemas() { + let exponential = histogram(Vec::new(), Vec::new()); + let mut custom = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![2.0], + ); + custom.schema = CUSTOM_BUCKETS_SCHEMA; + custom.custom_values = vec![1.0]; + + assert!(exponential.add(&custom).is_none()); + assert!(custom.add(&exponential).is_none()); + assert!(exponential.sub(&custom).is_none()); + assert!(custom.sub(&exponential).is_none()); + } + + #[test] + fn empty_custom_histogram_still_reconciles_bounds() { + let mut empty = histogram(Vec::new(), Vec::new()); + empty.schema = CUSTOM_BUCKETS_SCHEMA; + empty.custom_values = vec![1.0]; + let mut populated = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![2.0], + ); + populated.schema = CUSTOM_BUCKETS_SCHEMA; + populated.custom_values = vec![2.0]; + + let result = empty.add(&populated).unwrap(); + assert!(result.custom_values.is_empty()); + assert_eq!(result.positive_buckets, vec![2.0]); + } + #[test] fn zero_threshold_expands_to_populated_bucket_boundary() { let left = NativeHistogram { @@ -1805,7 +1834,7 @@ mod tests { } #[test] - fn promql_eq_ignores_reset_hint_start_timestamp_and_sparse_zeros() { + fn promql_eq_ignores_metadata_but_compares_sparse_layout() { let mut left = histogram( vec![Span { offset: 0, @@ -1826,6 +1855,27 @@ mod tests { right.reset_hint = CounterResetHint::NotCounterReset; right.start_timestamp = Some(2000); + assert!(!left.promql_eq(&right)); + + right.positive_spans = vec![ + Span { + offset: 0, + length: 0, + }, + Span { + offset: 0, + length: 1, + }, + Span { + offset: 42, + length: 0, + }, + ]; + left.positive_spans = vec![Span { + offset: 0, + length: 1, + }]; + left.positive_buckets = vec![1.0]; assert!(left.promql_eq(&right)); } @@ -1875,7 +1925,33 @@ mod tests { }], vec![1.0], ); - assert!(undecodable.promql_eq(&undecodable.clone())); + assert!(!undecodable.promql_eq(&undecodable.clone())); + + let mut mismatched = histogram( + vec![Span { + offset: 0, + length: 2, + }], + vec![1.0], + ); + mismatched.count = 1.0; + mismatched.sum = 1.0; + assert!(!mismatched.promql_eq(&mismatched.clone())); + + let invalid_offset = histogram( + vec![ + Span { + offset: 0, + length: 1, + }, + Span { + offset: -1, + length: 1, + }, + ], + vec![1.0, 1.0], + ); + assert!(!invalid_offset.promql_eq(&invalid_offset.clone())); } #[test] @@ -1944,11 +2020,51 @@ mod tests { } #[test] - fn terminal_exponential_bucket_uses_max_finite_bound() { - for (schema, index) in [(-1, 512), (0, 1024), (2, 4096)] { - assert_eq!(get_bound(index, schema, &[]), Some(f64::MAX)); - assert_eq!(get_bound(index + 1, schema, &[]), Some(f64::INFINITY)); + fn exponential_bucket_bounds_stop_after_overflow_bucket() { + for schema in [-4, 0, 8] { + let overflow = exponential_overflow_bucket_index(schema).unwrap(); + assert_eq!(get_bound(overflow - 1, schema, &[]), Some(f64::MAX)); + assert_eq!(get_bound(overflow, schema, &[]), Some(f64::INFINITY)); + assert_eq!(get_bound(overflow + 1, schema, &[]), None); } + assert_eq!(exponential_overflow_bucket_index(-5), None); + assert_eq!(exponential_overflow_bucket_index(9), None); + assert_eq!(get_bound(i32::MIN, -4, &[]), Some(0.0)); + } + + #[test] + fn custom_bucket_midpoints_preserve_infinite_bounds() { + let mut histogram = histogram(Vec::new(), Vec::new()); + histogram.schema = CUSTOM_BUCKETS_SCHEMA; + + assert_eq!( + histogram.bucket_midpoint(&Bucket { + lower: f64::NEG_INFINITY, + upper: -1.0, + count: 1.0, + boundary_rule: BoundaryRule::OpenLeft, + }), + f64::NEG_INFINITY + ); + assert_eq!( + histogram.bucket_midpoint(&Bucket { + lower: 1.0, + upper: f64::INFINITY, + count: 1.0, + boundary_rule: BoundaryRule::OpenLeft, + }), + f64::INFINITY + ); + assert!( + histogram + .bucket_midpoint(&Bucket { + lower: f64::NEG_INFINITY, + upper: f64::INFINITY, + count: 1.0, + boundary_rule: BoundaryRule::OpenLeft, + }) + .is_nan() + ); } #[test] @@ -1989,7 +2105,25 @@ mod tests { } #[test] - fn subtraction_treats_empty_left_as_zero() { + fn subtraction_treats_compatible_empty_left_as_zero() { + let left = histogram(vec![], vec![]); + let right = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![2.0], + ); + + let result = left.sub(&right).unwrap(); + assert_eq!(result.reset_hint, CounterResetHint::Gauge); + assert_eq!(result.count, -2.0); + assert_eq!(result.sum, -2.0); + assert_eq!(result.positive_buckets, vec![-2.0]); + } + + #[test] + fn subtraction_rejects_incompatible_empty_left() { let left = histogram(vec![], vec![]); let mut right = histogram( vec![Span { @@ -2001,13 +2135,7 @@ mod tests { right.schema = CUSTOM_BUCKETS_SCHEMA; right.custom_values = vec![1.0]; - let result = left.sub(&right).unwrap(); - - assert_eq!(result.schema, CUSTOM_BUCKETS_SCHEMA); - assert_eq!(result.reset_hint, CounterResetHint::Gauge); - assert_eq!(result.count, -2.0); - assert_eq!(result.sum, -2.0); - assert_eq!(result.positive_buckets, vec![-2.0]); + assert!(left.sub(&right).is_none()); } } diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index a72ddbeaaf3..973758032c2 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -26,9 +26,13 @@ mod test_util; mod union_distinct_on; pub use absent::{Absent, AbsentExec, AbsentStream}; +use common_query::native_histogram::{SUM_FIELD, native_histogram_value_type}; +use common_query::prometheus::is_prometheus_stale_nan; +use datafusion::arrow::array::{Array, Float64Array, StructArray}; use datafusion::arrow::datatypes::{ArrowPrimitiveType, TimestampMillisecondType}; use datafusion::common::DFSchemaRef; use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datatypes::data_type::DataType as _; pub use empty_metric::{EmptyMetric, EmptyMetricExec, EmptyMetricStream, build_special_time_expr}; pub use histogram_fold::{HistogramFold, HistogramFoldExec, HistogramFoldStream}; pub use instant_manipulate::{InstantManipulate, InstantManipulateExec, InstantManipulateStream}; @@ -43,6 +47,26 @@ pub type Millisecond = ::Native; const METRIC_NUM_SERIES: &str = "num_series"; +fn prometheus_stale_sample_column(column: &dyn Array) -> Option<(&dyn Array, &Float64Array)> { + let values = if let Some(values) = column.as_any().downcast_ref::() { + values + } else { + let histograms = column.as_any().downcast_ref::()?; + if histograms.data_type() != &native_histogram_value_type().as_arrow_type() { + return None; + } + histograms + .column_by_name(SUM_FIELD)? + .as_any() + .downcast_ref::()? + }; + Some((column, values)) +} + +fn is_prometheus_stale_sample((column, values): (&dyn Array, &Float64Array), row: usize) -> bool { + column.is_valid(row) && values.is_valid(row) && is_prometheus_stale_nan(values.value(row)) +} + /// Utilities for handling unfix logic in extension plans /// Convert column name to index for serialization pub fn serialize_column_index(schema: &DFSchemaRef, column_name: &str) -> u64 { diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 1897b453f38..f92a4e4490f 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -18,8 +18,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use common_query::prometheus::is_prometheus_stale_nan; -use datafusion::arrow::array::{Array, Float64Array, TimestampMillisecondArray, UInt64Array}; +use datafusion::arrow::array::{Array, TimestampMillisecondArray, UInt64Array}; use datafusion::arrow::datatypes::{DataType, SchemaRef}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::stats::Precision; @@ -46,7 +45,8 @@ use snafu::ResultExt; use crate::error::{DeserializeSnafu, Result}; use crate::extension_plan::series_divide::SeriesDivide; use crate::extension_plan::{ - METRIC_NUM_SERIES, Millisecond, resolve_column_name, serialize_column_index, + METRIC_NUM_SERIES, Millisecond, is_prometheus_stale_sample, prometheus_stale_sample_column, + resolve_column_name, serialize_column_index, }; use crate::metrics::PROMQL_SERIES_COUNT; @@ -510,11 +510,12 @@ impl Stream for InstantManipulateStream { } impl InstantManipulateStream { - // Refer to Prometheus `vectorSelectorSingle` / lookback semantics. - // - // Prometheus `v3.9.1` uses a start-exclusive lookback window: - // (eval_ts - lookback_delta, eval_ts] - // i.e. a sample at exactly `eval_ts - lookback_delta` is considered too old. + /// Manipulates one complete series sorted by timestamp. The planner enforces + /// this input contract with a sort followed by [`SeriesDivide`]. + /// + /// Prometheus `v3.9.1`'s `vectorSelectorSingle` uses a start-exclusive + /// lookback window `(eval_ts - lookback_delta, eval_ts]`; a sample at exactly + /// `eval_ts - lookback_delta` is too old. pub fn manipulate(&self, input: RecordBatch) -> DataFusionResult { let ts_column = input .column(self.time_index) @@ -531,10 +532,11 @@ impl InstantManipulateStream { return Ok(input); } - // field column for staleness check - let field_column = self + // Field column for staleness checks, classified once per batch. + let stale_sample_column = self .field_index - .and_then(|index| input.column(index).as_any().downcast_ref::()); + .map(|index| input.column(index).as_ref()) + .and_then(prometheus_stale_sample_column); // Optimize iteration range based on actual data bounds let first_ts = ts_column.value(0); @@ -579,9 +581,8 @@ impl InstantManipulateStream { let curr = ts_column.value(cursor); match curr.cmp(&expected_ts) { Ordering::Equal => { - if let Some(field_column) = &field_column - && field_column.is_valid(cursor) - && is_prometheus_stale_nan(field_column.value(cursor)) + if stale_sample_column + .is_some_and(|column| is_prometheus_stale_sample(column, cursor)) { // Ignore the stale marker. } else { @@ -614,9 +615,8 @@ impl InstantManipulateStream { let prev_ts = ts_column.value(prev_cursor); if prev_ts + self.lookback_delta > expected_ts { // only use the point in the time range - if let Some(field_column) = &field_column - && field_column.is_valid(prev_cursor) - && is_prometheus_stale_nan(field_column.value(prev_cursor)) + if stale_sample_column + .is_some_and(|column| is_prometheus_stale_sample(column, prev_cursor)) { // Do not use a stale marker as the newest value. continue; @@ -626,9 +626,8 @@ impl InstantManipulateStream { aligned_ts.push(expected_ts); } } - } else if let Some(field_column) = &field_column - && field_column.is_valid(cursor) - && is_prometheus_stale_nan(field_column.value(cursor)) + } else if stale_sample_column + .is_some_and(|column| is_prometheus_stale_sample(column, cursor)) { // Do not use a stale marker as the newest value. } else { @@ -692,6 +691,9 @@ fn reuse_constant_column(array: &Arc, len: usize) -> DataFusionResult #[cfg(test)] mod test { + use common_query::native_histogram::build_histogram_array; + use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; + use datafusion::arrow::array::Float64Array; use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::ToDFSchema; @@ -702,7 +704,7 @@ mod test { use super::*; use crate::extension_plan::test_util::{ - TIME_INDEX_COLUMN, prepare_test_data, prepare_test_data_with_stale_marker, + TIME_INDEX_COLUMN, native_histogram, prepare_test_data, prepare_test_data_with_stale_marker, }; async fn do_normalize_test( @@ -1386,7 +1388,7 @@ mod test { } #[tokio::test] - async fn prometheus_stale_nan_suppresses_exact_and_lookback() { + async fn prometheus_stale_nan_selects_before_and_suppresses_after_marker() { let schema = Arc::new(Schema::new(vec![ Field::new( TIME_INDEX_COLUMN, @@ -1409,6 +1411,69 @@ mod test { let input = Arc::new(DataSourceExec::new(Arc::new( MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); + let exec = Arc::new(InstantManipulateExec { + start: 750, + end: 1_500, + lookback_delta: 1_001, + interval: 250, + time_index_column: TIME_INDEX_COLUMN.to_string(), + field_column: Some("value".to_string()), + reuse_tsid_column: false, + input, + metric: ExecutionPlanMetricsSet::new(), + }); + + let context = SessionContext::default(); + let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) + .await + .unwrap(); + + let row_count = batches.iter().map(RecordBatch::num_rows).sum::(); + let batch = batches.iter().find(|batch| batch.num_rows() > 0).unwrap(); + let timestamp = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + let value = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!( + (row_count, timestamp, value), + (1, 750, 42.0), + "only the evaluation before the stale marker should select 42.0" + ); + } + + #[tokio::test] + async fn native_histogram_stale_nan_suppresses_exact_and_lookback() { + let histograms = build_histogram_array(&[ + Some(native_histogram(42.0)), + Some(native_histogram(f64::from_bits(PROMETHEUS_STALE_NAN_BITS))), + ]); + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("value", histograms.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![500, 1_000])), + histograms, + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); let exec = Arc::new(InstantManipulateExec { start: 1_000, end: 1_500, @@ -1426,12 +1491,7 @@ mod test { .await .unwrap(); - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - 0, - "the stale marker must suppress both the exact and lookback selections rather than \ - falling back to 42.0" - ); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); } #[tokio::test] diff --git a/src/promql/src/extension_plan/normalize.rs b/src/promql/src/extension_plan/normalize.rs index a84d5fc5c01..573fa5fb0ed 100644 --- a/src/promql/src/extension_plan/normalize.rs +++ b/src/promql/src/extension_plan/normalize.rs @@ -17,8 +17,8 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use common_query::prometheus::is_prometheus_stale_nan; -use datafusion::arrow::array::{Array, BooleanArray, Float64Array}; +use common_query::native_histogram::{START_TIMESTAMP_FIELD, native_histogram_arrow_type}; +use datafusion::arrow::array::{Array, BooleanArray, StructArray}; use datafusion::arrow::compute; use datafusion::common::{DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics}; use datafusion::error::DataFusionError; @@ -34,7 +34,7 @@ use datafusion::physical_plan::{ }; use datafusion_expr::col; use datatypes::arrow::array::TimestampMillisecondArray; -use datatypes::arrow::datatypes::SchemaRef; +use datatypes::arrow::datatypes::{SchemaRef, TimestampMillisecondType}; use datatypes::arrow::record_batch::RecordBatch; use futures::{Stream, StreamExt, ready}; use greptime_proto::substrait_extension as pb; @@ -43,7 +43,8 @@ use snafu::ResultExt; use crate::error::{DeserializeSnafu, Result}; use crate::extension_plan::{ - METRIC_NUM_SERIES, Millisecond, resolve_column_name, serialize_column_index, + METRIC_NUM_SERIES, Millisecond, is_prometheus_stale_sample, prometheus_stale_sample_column, + resolve_column_name, serialize_column_index, }; use crate::metrics::PROMQL_SERIES_COUNT; @@ -51,7 +52,7 @@ use crate::metrics::PROMQL_SERIES_COUNT; /// the input batch only contains sample points from one time series. /// /// Roughly speaking, this method does these things: -/// - bias sample's timestamp by offset +/// - bias sample and native histogram start timestamps by offset /// - sort the record batch based on timestamp column /// - remove Prometheus stale markers (optional) #[derive(Debug, PartialEq, Eq, Hash, PartialOrd)] @@ -397,17 +398,62 @@ impl SeriesNormalizeStream { ) })?; + let bias_timestamp = |timestamp: i64| { + timestamp.checked_add(self.offset).ok_or_else(|| { + DataFusionError::Execution("SeriesNormalize: timestamp offset overflow".into()) + }) + }; + // bias the timestamp column by offset let ts_column_biased = if self.offset == 0 { Arc::new(ts_column.clone()) as _ } else { - Arc::new(TimestampMillisecondArray::from_iter( - ts_column.iter().map(|ts| ts.map(|ts| ts + self.offset)), - )) + Arc::new(ts_column.try_unary::<_, TimestampMillisecondType, _>(&bias_timestamp)?) }; let mut columns = input.columns().to_vec(); columns[self.time_index] = ts_column_biased; + // Offset selectors move samples into the evaluation timeline. Keep native histogram + // start timestamps on the same timeline for rate and reset calculations. + if self.offset != 0 { + let native_histogram_type = native_histogram_arrow_type(); + for column in &mut columns { + let Some((histograms, start_timestamp_index, start_timestamps)) = column + .as_any() + .downcast_ref::() + .filter(|histograms| histograms.data_type() == &native_histogram_type) + .and_then(|histograms| { + let (index, _) = histograms.fields().find(START_TIMESTAMP_FIELD)?; + let timestamps = histograms + .column(index) + .as_any() + .downcast_ref::()?; + Some((histograms, index, timestamps)) + }) + else { + continue; + }; + let start_timestamps = start_timestamps + .try_unary::<_, TimestampMillisecondType, _>(|timestamp| { + // Prometheus uses zero to mean that the start timestamp is unknown. + if timestamp == 0 { + Ok(0) + } else { + bias_timestamp(timestamp) + } + })?; + // Replace only the start timestamp child to preserve the histogram payload and + // null bitmap. + let mut children = histograms.columns().to_vec(); + children[start_timestamp_index] = Arc::new(start_timestamps); + *column = Arc::new(StructArray::new( + histograms.fields().clone(), + children, + histograms.nulls().cloned(), + )); + } + } + let result_batch = RecordBatch::try_new(input.schema(), columns)?; if !self.filter_stale_markers { return Ok(result_batch); @@ -416,11 +462,12 @@ impl SeriesNormalizeStream { // Filter out Prometheus stale markers. let mut stale_marker_filter = vec![true; input.num_rows()]; for column in result_batch.columns() { - if let Some(float_column) = column.as_any().downcast_ref::() { - for (i, flag) in stale_marker_filter.iter_mut().enumerate() { - if float_column.is_valid(i) && is_prometheus_stale_nan(float_column.value(i)) { - *flag = false; - } + let Some(stale_sample_column) = prometheus_stale_sample_column(column.as_ref()) else { + continue; + }; + for (i, flag) in stale_marker_filter.iter_mut().enumerate() { + if is_prometheus_stale_sample(stale_sample_column, i) { + *flag = false; } } } @@ -462,6 +509,8 @@ impl Stream for SeriesNormalizeStream { #[cfg(test)] mod test { + use common_query::native_histogram::{build_histogram_array, read_histogram}; + use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use datafusion::arrow::array::Float64Array; use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{ @@ -476,6 +525,7 @@ mod test { use datatypes::arrow_array::StringArray; use super::*; + use crate::extension_plan::test_util::native_histogram; const TIME_INDEX_COLUMN: &str = "timestamp"; @@ -658,4 +708,71 @@ mod test { assert_eq!(auxiliary.value(0).to_bits(), 0x7ff8_0000_0000_0000); assert!(!value.is_valid(1)); } + + #[tokio::test] + async fn offsets_native_histogram_timestamps_and_filters_stale_markers() { + let mut regular = native_histogram(42.0); + regular.start_timestamp = Some(500); + let mut ordinary_nan = native_histogram(f64::NAN); + ordinary_nan.start_timestamp = Some(0); + let histograms = build_histogram_array(&[ + Some(regular), + Some(native_histogram(f64::from_bits(PROMETHEUS_STALE_NAN_BITS))), + Some(ordinary_nan), + None, + ]); + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + TimestampMillisecondType::DATA_TYPE, + false, + ), + Field::new("value", histograms.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 1_000, 2_000, 3_000, 4_000, + ])), + histograms, + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let exec = Arc::new(SeriesNormalizeExec { + offset: 1_000, + time_index_column_name: TIME_INDEX_COLUMN.to_string(), + filter_stale_markers: true, + tag_columns: Vec::new(), + input, + metric: ExecutionPlanMetricsSet::new(), + }); + + let context = SessionContext::default(); + let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) + .await + .unwrap(); + let batch = batches.iter().find(|batch| batch.num_rows() == 3).unwrap(); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + + let timestamps = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(timestamps.values(), &[2_000, 4_000, 5_000]); + let regular = read_histogram(values, 0).unwrap().unwrap(); + assert_eq!((regular.sum, regular.start_timestamp), (42.0, Some(1_500))); + let ordinary_nan = read_histogram(values, 1).unwrap().unwrap(); + assert!(ordinary_nan.sum.is_nan()); + assert_eq!(ordinary_nan.start_timestamp, Some(0)); + assert!(read_histogram(values, 2).unwrap().is_none()); + } } diff --git a/src/promql/src/extension_plan/test_util.rs b/src/promql/src/extension_plan/test_util.rs index 43d0d686038..90519d7a98a 100644 --- a/src/promql/src/extension_plan/test_util.rs +++ b/src/promql/src/extension_plan/test_util.rs @@ -16,6 +16,7 @@ use std::sync::Arc; +use common_query::native_histogram::{NativeHistogram, Span, UNKNOWN_COUNTER_RESET_HINT}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use common_recordbatch::DfRecordBatch as RecordBatch; use datafusion::arrow::array::Float64Array; @@ -29,6 +30,26 @@ use datatypes::arrow_array::StringArray; pub(crate) const TIME_INDEX_COLUMN: &str = "timestamp"; +pub(crate) fn native_histogram(sum: f64) -> NativeHistogram { + NativeHistogram { + schema: 0, + zero_threshold: 0.0, + sum, + reset_hint: UNKNOWN_COUNTER_RESET_HINT, + start_timestamp: None, + custom_values: Vec::new(), + positive_spans: vec![Span { + offset: 0, + length: 1, + }], + negative_spans: Vec::new(), + count: 1.0, + zero_count: 0.0, + positive_buckets: vec![1.0], + negative_buckets: Vec::new(), + } +} + pub(crate) fn prepare_test_data() -> DataSourceExec { let schema = Arc::new(Schema::new(vec![ Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true), diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index 5208db8c6b1..4a4e3adad5c 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -1462,24 +1462,25 @@ fn histogram_delta( .map(NativeHistogram::into_gauge); } - let mut result = samples.first()?.zero_like(); - let mut segment_start = samples.first()?; - for (sample_pair, ts_pair) in samples.windows(2).zip(timestamps.windows(2)) { - let previous_ts = ts_pair[0]; - let current_ts = ts_pair[1]; - let previous = &sample_pair[0]; - let current = &sample_pair[1]; - if current.detect_counter_reset(previous, previous_ts, current_ts) { - result = result.add(&previous.sub(segment_start)?)?; - result = result.add(current)?; - segment_start = current; + let first_reset = samples[1].detect_counter_reset(&samples[0], timestamps[0], timestamps[1]); + let (initial, reset_scan_start) = if first_reset { + // The first sample is irrelevant after an immediate reset. Adopt the + // second sample's layout so an incompatible pre-reset layout is ignored. + (samples[1].zero_like(), 2) + } else { + (samples[0].clone(), 1) + }; + let mut result = samples.last()?.sub(&initial)?; + for index in reset_scan_start..samples.len() { + if samples[index].detect_counter_reset( + &samples[index - 1], + timestamps[index - 1], + timestamps[index], + ) { + result = result.add(&samples[index - 1])?; } } - Some( - result - .add(&samples.last()?.sub(segment_start)?)? - .into_gauge(), - ) + Some(result.into_gauge()) } fn idelta_value( @@ -1569,7 +1570,7 @@ fn native_extrapolated_rate( let (raw_offset, raw_length) = unpack(keys[index]); let offset = raw_offset as usize; let length = raw_length as usize; - if length < 2 { + if length == 0 { result.push(None); continue; } @@ -1588,6 +1589,20 @@ fn native_extrapolated_rate( continue; } + let first_ts = all_timestamps[offset]; + let last_ts = all_timestamps[offset + length - 1]; + let range_end = eval_ts.value(index); + let range_start = range_end - range_length; + let synthetic_zero_timestamp = IS_COUNTER + .then_some(samples[0].start_timestamp) + .flatten() + .filter(|start| *start != 0 && range_start < *start && *start < first_ts); + let synthetic_zero_start = synthetic_zero_timestamp.is_some(); + if length < 2 && !synthetic_zero_start { + result.push(None); + continue; + } + let wrong_flavor = if IS_COUNTER { samples .iter() @@ -1608,7 +1623,11 @@ fn native_extrapolated_rate( for pair in samples.windows(2) { record_custom_reconciliation(&collector, func_name, &pair[0], &pair[1]); } - let Some(mut histogram) = histogram_delta(&samples, timestamps, IS_COUNTER) else { + let mut histogram = if length == 1 { + samples[0].clone() + } else if let Some(histogram) = histogram_delta(&samples, timestamps, IS_COUNTER) { + histogram + } else { record_warning( &collector, format!("{func_name}: dropped native histogram range with incompatible schemas"), @@ -1616,21 +1635,41 @@ fn native_extrapolated_rate( result.push(None); continue; }; + if synthetic_zero_start && length > 1 { + let Some(with_synthetic_zero) = histogram.add(&samples[0]) else { + record_warning( + &collector, + format!( + "{func_name}: dropped native histogram range with incompatible schemas" + ), + ); + result.push(None); + continue; + }; + histogram = with_synthetic_zero; + } - let first_ts = all_timestamps[offset]; - let last_ts = all_timestamps[offset + length - 1]; - let range_end = eval_ts.value(index); - let range_start = range_end - range_length; - let sampled_interval_ms = (last_ts - first_ts) as f64; + let real_sampled_interval_ms = (last_ts - first_ts) as f64; + let sampled_interval_ms = synthetic_zero_timestamp + .map(|start| (last_ts - start) as f64) + .unwrap_or(real_sampled_interval_ms); if sampled_interval_ms <= 0.0 { result.push(None); continue; } - let average_interval_ms = sampled_interval_ms / (length - 1) as f64; - let mut duration_to_start_ms = (first_ts - range_start) as f64; + let average_interval_ms = if length > 1 { + real_sampled_interval_ms / (length - 1) as f64 + } else { + 0.0 + }; + let mut duration_to_start_ms = if synthetic_zero_start { + 0.0 + } else { + (first_ts - range_start) as f64 + }; let duration_to_end_ms = (range_end - last_ts) as f64; - if IS_COUNTER && histogram.count > 0.0 && samples[0].count >= 0.0 { + if IS_COUNTER && !synthetic_zero_start && histogram.count > 0.0 && samples[0].count >= 0.0 { let duration_to_zero = sampled_interval_ms * (samples[0].count / histogram.count); if duration_to_zero < duration_to_start_ms { duration_to_start_ms = duration_to_zero; @@ -2010,14 +2049,47 @@ mod tests { histograms: Vec, ) -> NativeHistogram { let range_length = histograms.len() as i64 * 1000; - let mut input = histogram_range_input(histograms.into_iter().map(Some).collect::>()); + let timestamps = (0..histograms.len()) + .map(|idx| (idx as i64 + 1) * 1000) + .collect(); + extrapolated_histogram_result(udf, timestamps, histograms, range_length, range_length) + .unwrap() + } + + fn extrapolated_histogram_result( + udf: ScalarUDF, + timestamps: Vec, + histograms: Vec, + range_end: i64, + range_length: i64, + ) -> Option { + assert_eq!(timestamps.len(), histograms.len()); + let range = [(0, u32::try_from(histograms.len()).unwrap())]; + let timestamps = Arc::new(TimestampMillisecondArray::from(timestamps)); + let histograms = + build_histogram_array(&histograms.into_iter().map(Some).collect::>()); + let mut input = vec![ + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(timestamps, range) + .unwrap() + .into_dict(), + )), + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(histograms, range) + .unwrap() + .into_dict(), + )), + ]; input.push(ColumnarValue::Array(Arc::new( - TimestampMillisecondArray::from(vec![range_length]), + TimestampMillisecondArray::from(vec![range_end]), ))); input.push(ColumnarValue::Array(Arc::new(Int64Array::from(vec![ range_length, ])))); - run_histogram_udf(udf, input) + let result = run_udf(udf, input, native_histogram_arrow_type()); + let result = extract_array(&result).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + read_histogram(result, 0).unwrap() } fn collected_warnings(collector: &PromqlAnnotationCollector) -> Vec { @@ -2051,7 +2123,7 @@ mod tests { } #[test] - fn comparison_uses_promql_histogram_equality() { + fn comparison_observes_explicit_sparse_zero_buckets() { let mut left = sample_histogram(1.0, 1.0, vec![1.0, 0.0]); left.reset_hint = COUNTER_RESET_HINT; left.start_timestamp = Some(1000); @@ -2069,7 +2141,7 @@ mod tests { ); let values = extract_array(&result).unwrap(); let values = values.as_any().downcast_ref::().unwrap(); - assert!(values.value(0)); + assert!(!values.value(0)); } #[test] @@ -2397,4 +2469,132 @@ mod tests { assert_eq!(idelta.count, 7.0); assert_eq!(idelta.sum, 12.0); } + + #[test] + fn extrapolated_rate_uses_start_timestamp_synthetic_zero() { + let mut single = sample_histogram(1.0, 1.0, vec![1.0]); + single.start_timestamp = Some(1_000); + + let rate = extrapolated_histogram_result( + NativeHistogramRate::scalar_udf(), + vec![2_000], + vec![single.clone()], + 3_000, + 3_000, + ) + .unwrap(); + assert_eq!(rate.count, 1.0 / 3.0); + assert_eq!(rate.sum, 1.0 / 3.0); + + let increase = extrapolated_histogram_result( + NativeHistogramIncrease::scalar_udf(), + vec![2_000], + vec![single], + 3_000, + 3_000, + ) + .unwrap(); + assert_eq!(increase.count, 1.0); + assert_eq!(increase.sum, 1.0); + + let mut first = sample_histogram(2.0, 2.0, vec![2.0]); + first.start_timestamp = Some(1_000); + let last = sample_histogram(4.0, 4.0, vec![4.0]); + let increase = extrapolated_histogram_result( + NativeHistogramIncrease::scalar_udf(), + vec![2_000, 3_000], + vec![first, last], + 3_000, + 3_000, + ) + .unwrap(); + assert_eq!(increase.count, 4.0); + assert_eq!(increase.sum, 4.0); + } + + #[test] + fn extrapolated_rate_requires_strictly_in_range_start_timestamp() { + for (start_timestamp, range_end, range_length) in [ + (0, 3_000, 3_000), + (500, 3_000, 2_000), + (1_000, 3_000, 2_000), + (2_000, 3_000, 3_000), + (2_500, 3_000, 3_000), + ] { + let mut sample = sample_histogram(1.0, 1.0, vec![1.0]); + sample.start_timestamp = Some(start_timestamp); + assert!( + extrapolated_histogram_result( + NativeHistogramRate::scalar_udf(), + vec![2_000], + vec![sample], + range_end, + range_length, + ) + .is_none(), + "start_timestamp={start_timestamp}" + ); + } + + let mut gauge = sample_histogram(1.0, 1.0, vec![1.0]); + gauge.start_timestamp = Some(1_000); + gauge.reset_hint = GAUGE_RESET_HINT; + assert!( + extrapolated_histogram_result( + NativeHistogramDelta::scalar_udf(), + vec![2_000], + vec![gauge], + 3_000, + 3_000, + ) + .is_none() + ); + } + + #[test] + fn first_reset_ignores_incompatible_pre_reset_layout() { + let first = sample_histogram(10.0, 10.0, vec![10.0]); + let second = NativeHistogram { + schema: CUSTOM_BUCKETS_SCHEMA, + zero_threshold: 0.0, + sum: 2.0, + reset_hint: COUNTER_RESET_HINT, + start_timestamp: None, + custom_values: vec![1.0], + positive_spans: vec![Span { + offset: 0, + length: 1, + }], + negative_spans: Vec::new(), + count: 2.0, + zero_count: 0.0, + positive_buckets: vec![2.0], + negative_buckets: Vec::new(), + }; + + let delta = histogram_delta(&[first, second], &[1_000, 2_000], true).unwrap(); + assert_eq!(delta.schema, CUSTOM_BUCKETS_SCHEMA); + assert_eq!(delta.count, 2.0); + assert_eq!(delta.sum, 2.0); + assert_eq!(delta.positive_buckets, vec![2.0]); + } + + #[test] + fn counter_delta_handles_reset_segment_boundaries() { + let first = sample_histogram(5.0, 5.0, vec![5.0]); + let second = sample_histogram(7.0, 7.0, vec![7.0]); + let mut reset = sample_histogram(2.0, 2.0, vec![2.0]); + reset.reset_hint = COUNTER_RESET_HINT; + let last = sample_histogram(4.0, 4.0, vec![4.0]); + + let delta = histogram_delta( + &[first, second, reset, last], + &[1_000, 2_000, 3_000, 4_000], + true, + ) + .unwrap(); + assert_eq!(delta.count, 6.0); + assert_eq!(delta.sum, 6.0); + assert_eq!(delta.positive_buckets, vec![6.0]); + } } diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 3c3729e9af0..a463f176c71 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -140,6 +140,7 @@ const DB_COLUMN_MATCHER: &str = "__database__"; const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v"; const OR_FLOAT_FIELD_PREFIX: &str = "__promql_or_float_"; const OR_HISTOGRAM_FIELD_PREFIX: &str = "__promql_or_histogram_"; +const TIMESTAMP_VALUE_PREFIX: &str = "__promql_timestamp_value_"; /// Threshold for scatter scan mode const MAX_SCATTER_POINTS: i64 = 400; @@ -1516,42 +1517,77 @@ impl PromPlanner { let normalize = self .selector_to_series_normalize_plan(offset, matchers, false) .await?; + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; - let normalize = if timestamp_fn { - // If evaluating the PromQL `timestamp()` function, project the time index column as the value column - // before wrapping with [`InstantManipulate`], so the output matches PromQL's `timestamp()` semantics. - self.create_timestamp_func_plan(normalize)? + let (normalize, timestamp_value_column) = if timestamp_fn { + // Keep the original sample for stale-marker detection while carrying + // its timestamp through InstantManipulate in a private value column. + let occupied = normalize + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(); + let mut timestamp_value_column = TIMESTAMP_VALUE_PREFIX.to_string(); + while occupied.contains(timestamp_value_column.as_str()) { + timestamp_value_column.push('_'); + } + let mut project_exprs = normalize + .schema() + .iter() + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + project_exprs + .push(build_special_time_expr(&time_index_column).alias(×tamp_value_column)); + let normalize = LogicalPlanBuilder::from(normalize) + .project(project_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + (normalize, Some(timestamp_value_column)) } else { - normalize + (normalize, None) }; + let field_column = self.ctx.field_columns.first().cloned(); let manipulate = InstantManipulate::new( self.ctx.start, self.ctx.end, self.ctx.lookback_delta, self.ctx.interval, - self.ctx - .time_index_column - .clone() - .expect("time index should be set in `setup_context`"), + time_index_column, if self.ctx.use_tsid { vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()] } else { self.ctx.tag_columns.clone() }, - self.ctx.field_columns.first().cloned(), + field_column, normalize, ); - Ok(LogicalPlan::Extension(Extension { + let manipulate = LogicalPlan::Extension(Extension { node: Arc::new(manipulate), - })) + }); + if let Some(timestamp_value_column) = timestamp_value_column { + self.create_timestamp_func_plan(manipulate, ×tamp_value_column) + } else { + Ok(manipulate) + } } /// Builds a projection plan for the PromQL `timestamp()` function. /// Projects the time index column as the value column for each row. /// /// # Arguments - /// * `normalize` - Input [`LogicalPlan`] for the normalized series. + /// * `input` - Input [`LogicalPlan`] after instant-vector selection. + /// * `timestamp_value_column` - Private column containing each selected sample's timestamp. /// /// # Returns /// Returns a [`Result`] where the resulting logical plan projects the timestamp @@ -1568,16 +1604,19 @@ impl PromPlanner { /// # Side Effects /// Updates the planner context's field columns to the timestamp column name. /// - fn create_timestamp_func_plan(&mut self, normalize: LogicalPlan) -> Result { - let time_expr = build_special_time_expr(self.ctx.time_index_column.as_ref().unwrap()) - .alias(DEFAULT_FIELD_COLUMN); + fn create_timestamp_func_plan( + &mut self, + input: LogicalPlan, + timestamp_value_column: &str, + ) -> Result { + let time_expr = col(timestamp_value_column).alias(DEFAULT_FIELD_COLUMN); self.ctx.field_columns = vec![time_expr.schema_name().to_string()]; let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2); project_exprs.push(self.create_time_index_column_expr()?); project_exprs.push(time_expr); project_exprs.extend(self.create_tag_column_exprs()?); - LogicalPlanBuilder::from(normalize) + LogicalPlanBuilder::from(input) .project(project_exprs) .context(DataFusionPlanningSnafu)? .build() @@ -6740,9 +6779,32 @@ mod test { } #[tokio::test] - #[should_panic] - async fn single_timestamp() { - do_single_instant_function_call("timestamp", "").await; + async fn single_timestamp_plan_preserves_source_value() { + let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#); + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = String::from( + "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(some_metric.timestamp AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + assert_eq!(plan.display_indent_schema().to_string(), expected); } #[tokio::test] From cabc2f6cc667291d7bc3dcbb8f9cc82e27bd93a0 Mon Sep 17 00:00:00 2001 From: Palak Jha Date: Thu, 6 Aug 2026 12:54:59 +0530 Subject: [PATCH 002/126] feat(flow): add information_schema.flow_statistics and SHOW FLOW STATUS (#7987) (#8392) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(flow): add information_schema.flow_statistics and SHOW FLOW STATUS (fixes #7987) Signed-off-by: onepizzateam * feat(flow): add information_schema.flow_statistics and SHOW FLOW STATUS (fixes #7987) Signed-off-by: onepizzateam * test(flow): add sqlness golden result for flow_status Signed-off-by: Palak Jha * fix(catalog): remove unused OptionExt import in flow_statistics Signed-off-by: onepizzateam * docs(flow): fix stale 'recent errors' comment on QueryFlowExecStats Signed-off-by: onepizzateam * test: regenerate golden results for flow_statistics table Signed-off-by: onepizzateam * style: apply rustfmt to flow_statistics changes Signed-off-by: onepizzateam * refactor(catalog): hoist current_time_millis out of flow loop and clamp uptime Signed-off-by: onepizzateam * chore: remove accidentally committed fmt_check.log Signed-off-by: onepizzateam * Update flow_status.result del eof trailing blank line as per review Signed-off-by: onepizzateam * test(flow): restore runner-generated trailing blank line for sqlness Signed-off-by: onepizzateam * postgres: include SHOW FLOW STATUS in extended-query describe (return flow_statistics fields) Signed-off-by: onepizzateam * fix: address reviewer feedback on flow_statistics PR Signed-off-by: onepizzateam * fix(flow): resolve merge conflicts with main Signed-off-by: polar Signed-off-by: onepizzateam * sqlness check post gen (information_schema.result) Signed-off-by: onepizzateam * fix(flow): record start_time after req/snapshot_seqs built, before dispatch Signed-off-by: onepizzateam * fix(sql): handle ShowFlowStatus in match statement at util.rs Signed-off-by: onepizzateam * feat: review patch implementation Signed-off-by: onepizzateam * chore: remove accidentally committed local tool output files and fix fmt Signed-off-by: onepizzateam * fix worker.rs return type formatting Signed-off-by: onepizzateam * fix(flow): apply rustfmt to get_full_flow_stat return type Signed-off-by: onepizzateam * fix worker.rs return type formatting Signed-off-by: onepizzateam * fix(flow): re-apply rustfmt to get_full_flow_stat return type after merge Signed-off-by: onepizzateam * fix: merge conflicts Signed-off-by: onepizzateam * fix(auth): warn when credential load disables Postgres SCRAM or drops a line (#8652) * fix(auth): warn when credential load disables Postgres SCRAM or drops a line Static and watch user providers degraded silently in two ways: - A single non-SCRAM verifier (mysql_native_password, or a legacy pbkdf2_sha256 hash that predates SCRAM) disables Postgres SCRAM for every user and falls back to cleartext, with no signal to the operator. - A malformed credential line (commonly a plaintext password containing '=', which splits into more than two parts) was dropped without a trace. Emit a warning at each credential load for both cases so operators don't unknowingly serve cleartext passwords over Postgres or lose a user. This is logging only; authentication behavior is unchanged. The SCRAM check never logs secrets, and the malformed-line warning logs the line number and file, never the line content. Signed-off-by: Dennis Zhuang * fix(auth): warn on credential file read error before truncating A read error from lines() (I/O failure or invalid UTF-8) ends the iterator via map_while, silently dropping every remaining credential. Warn with the line number and file before truncating, matching the malformed-line handling, so the drop is observable. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang Signed-off-by: onepizzateam * test(object-store): fix racy SecureFs abort test (#8720) test_writer_abort_is_unsupported_without_atomic_write asserted the file content immediately after abort() returned Unsupported. SecureFsWriter writes through tokio::fs::File, whose write_all() only enqueues a blocking write task (tokio's poll_write returns Ready before the write completes), so the data may not be visible yet when the test reads the file. Drop the race-prone content assertion and only verify the Unsupported contract. Signed-off-by: Lei, HUANG Signed-off-by: onepizzateam * feat(query): plan native histogram functions (#8705) * feat(query): plan native histogram functions Signed-off-by: shuiyisong * fix: cr issue & add tests Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong Signed-off-by: onepizzateam * perf(promql): avoid repeated scans in sliding range evaluation (#8646) * perf(promql): use two pointers for sliding range boundaries Replace the stale cursor heuristic in RangeManipulateStream::calculate_range with monotonic left/right cursors. The old path rescanned each evaluation window (O(E x samples-per-window)) and could lose valid samples after sparse gaps or trailing empty windows. The two pointers keep strict monotonic progress, reducing boundary generation to O(N + E) while preserving (curr-range, curr] semantics, start/end shortening, and empty-window output. Controlled release benchmarks (fixed CPU, ABBA): - Public RangeManipulate wall time: ~28% faster at 1m/15s, ~66% at 5m/15s, ~96% at 1h/15s. - Warmed distributed TQL ANALYZE 1h queries: ~17-21% faster end to end; shorter windows stayed within run-order noise. Signed-off-by: discord9 * perf(promql): specialize changes/resets with adaptive edge counting The generic range_fn macro slices, downcasts, and rescans every overlapping window for changes() and resets(). Replace the macro path for these two functions with hand-written UDF wrappers backed by a shared private edge-count kernel: direct raw-offset scans when requested edges are few, otherwise one global u64 edge prefix so each window is answered by a prefix difference. Behavior is preserved bit-for-bit, including raw null-buffer values, NaN semantics, signed zero, infinities, empty/singleton windows, independent timestamp/value offsets, arbitrary window layouts, and exact DataFusion error messages. The shared proc macro, planner, serializer, and other range functions are untouched. Controlled release benchmarks (fixed CPU, ABBA): - Dense sliding windows (k=4/20/240): 91.7-95.6% less public UDF wall time. - Low-coverage fallback (N=4096, 8 windows): 73.9-74.4% faster. - Warmed distributed TQL ANALYZE 5m/1h changes/resets: 12.1-19.7% client and 12.0-20.9% server latency improvement; controls stayed within drift. Signed-off-by: discord9 * ci(query-regression): include PromQL range boundary case in defaults An audit of historical query-regression runs found zero range-query coverage: all 208 PromQL ANALYZE samples were bare selectors, so range evaluation could regress without CI noticing. Wire the promql_range_boundary case (introduced in #8646) into DEFAULT_CASES so label-triggered runs measure the range path. The case is cheap: a ~0.3s synthetic fixture and about a minute of query execution per base/candidate pass. Signed-off-by: discord9 * chore(promql): address sliding range review nits Move test-only imports into their test modules and remove the unused pre-specialization changes and resets helpers. Signed-off-by: discord9 * style(promql): apply pinned rustfmt Signed-off-by: discord9 * test(promql): cover sparse range results Share the changes and resets test scaffolding while keeping their behavior oracles independent. Add an end-to-end sqlness regression for sparse samples, empty intermediate windows, and a valid trailing sample. Signed-off-by: discord9 --------- Signed-off-by: discord9 Signed-off-by: onepizzateam * ci: optimize fuzz and split workflows (#8710) * ci: batch fuzz targets in GitHub Actions Signed-off-by: WenyXu * ci: improve fuzz test observability Signed-off-by: WenyXu * fix(ci): preserve fuzz setup failure artifacts Signed-off-by: WenyXu * test(ci): keep fuzz mock output in logs Signed-off-by: WenyXu * ci: optimize fuzz worker cache Signed-off-by: WenyXu * ci: warm fuzz target binaries Signed-off-by: WenyXu * ci: isolate fuzz workflow Signed-off-by: WenyXu * ci: centralize fuzz target preparation Signed-off-by: WenyXu * ci: split general workflows Signed-off-by: WenyXu * ci: streamline docs required checks Signed-off-by: WenyXu * fix: transfer fuzz targets as artifacts Signed-off-by: WenyXu * fix: preserve fuzz binary permissions Signed-off-by: WenyXu * ci: streamline fuzz workers Signed-off-by: WenyXu * ci: cache PR build dependencies Signed-off-by: WenyXu * ci: retain main build cache policy Signed-off-by: WenyXu * ci: address fuzz review feedback Signed-off-by: WenyXu --------- Signed-off-by: WenyXu Signed-off-by: onepizzateam * fix: add public constructor for compactor (#8724) Signed-off-by: onepizzateam * feat(logging): add enable_file_logging option to disable file logging (#8721) Signed-off-by: xhwhis Signed-off-by: onepizzateam * avoid cloning final Prometheus remote write row (#8733) perf: avoid cloning final Prometheus remote write row Signed-off-by: lyang24 Signed-off-by: onepizzateam * feat(function): add json_object_keys scalar function (#8722) Expose JSON object key listing for outermost objects, with sqlness coverage. Signed-off-by: onepizzateam * refactor(mito2): revise compaction trigger behavior (#8706) * refactor(mito2): revise compaction trigger behavior Distinguish automatic and manual triggers, coalesce explicit automatic follow-ups, and reject concurrent manual compactions. Remove implicit post-execution continuation and transient idle statuses so scheduler entries always represent an active lifecycle. Signed-off-by: Lei, HUANG * fix(mito2): track automatic compaction follow-ups Signed-off-by: Lei, HUANG * docs(mito2): fix compaction transition rustdoc Signed-off-by: Lei, HUANG * fix(mito2): mark manual compaction conflict retryable Signed-off-by: Lei, HUANG * refactor(mito2): drop unused RequestCancelResult::NotRunning variant request_cancel is only called in tests where the region is guaranteed to be running, so the NotRunning case was dead code. Simplify to unwrap() and remove the variant. Signed-off-by: Lei, HUANG * fix(mito2): gate test-only cancellation import Signed-off-by: Lei, HUANG * fix(mito2): prioritize DDL after compaction planning Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG Signed-off-by: Lei, HUANG Signed-off-by: onepizzateam * feat: update dashboard to v0.13.11 (#8737) Signed-off-by: onepizzateam * fix(object-store): skip removed-entry lister test on Windows (#8735) DirEntry on Windows is a snapshot from FindFirstFileW: file_type() and metadata() keep returning cached data after the file is removed, so read_list_entry() cannot observe the deletion. The test asserts the Unix behavior (lstat returns ENOENT) and fails deterministically on Windows nightly CI (4/4 tries). Gate it with #[cfg(not(windows))]. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: onepizzateam * fix(query): preserve remote dynamic filter target (#8615) * fix(query): preserve remote dynamic filter target Signed-off-by: discord9 * fix(query): check RDF subscriber registration Signed-off-by: discord9 * fix(query): refresh initial dyn filter snapshot before dispatch and handle RDF unregister The remote dynamic filter dispatch ordering regression: freeze the target, pre-register subscribers, refresh the initial snapshot, then dispatch. Also implement handle_remote_dyn_filter_unregister to keep unregister targets consistent with do_get/update. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): update test-only RegionQueryHandler impl to new trait signatures Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: onepizzateam * docs: rework README release badges, drop star history, fix grpc flag (#8743) * docs: show stable, latest and nightly version badges in README The single release badge rendered whatever GitHub considered newest, so a pre-release such as v1.2.0-beta.1 looked like the recommended version. Split it into three self-updating badges using the shields.io `filter` parameter, keyed off the existing tag naming: - stable: `!*-*` matches tags without a hyphen (v1.1.4) - latest: `!*-*-*` excludes nightly and dev builds (v1.2.0-beta.1) - nightly: `*-nightly-*` matches the weekly build (v1.2.0-nightly-20260706) No workflow changes are needed; the badges track new releases on their own. A one-line caption below them says which channel to pick. The release-date badge is dropped as the three version badges already carry that signal. Signed-off-by: Dennis Zhuang * docs: remove star history chart from README The chart carried a sealed_token in three URLs and added a large third-party image to the Project Status section without saying anything the badges and case studies do not already cover. Signed-off-by: Dennis Zhuang * docs: use --grpc-bind-addr in README quickstart --rpc-bind-addr is now only a hidden alias of --grpc-bind-addr and no longer shows up in --help. Signed-off-by: Dennis Zhuang * Update README.md Co-authored-by: Ning Sun --------- Signed-off-by: Dennis Zhuang Co-authored-by: Ning Sun Signed-off-by: onepizzateam * chore: check enterprise-gated files are listed in both license configs (#8750) * chore: check enterprise-gated files are listed in both license configs A file reachable only through `#[cfg(feature = "enterprise")] mod ...;` is governed by the GreptimeDB Enterprise License, so it must appear in the `includes` of licenserc-enterprise.toml and the `excludes` of licenserc.toml. hawkeye stays silent when it does not: the file keeps its Apache-2.0 header and passes the default check precisely because it was never excluded from it. scripts/check-enterprise-license.py walks enterprise-gated `mod` declarations, resolves them to files (submodules included) and diffs that set against both configs, also reporting stale entries. It runs in the license job in CI and as `make check-enterprise-license`. Documents the split it cannot decide for you — whole enterprise features get their own file, a gated match arm stays inline — in .agents/architecture-invariants.md. Signed-off-by: Dennis Zhuang * fix: tighten enterprise license checks Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang Signed-off-by: onepizzateam * fix(operator): invalidate local cache after dropping view (#8748) Signed-off-by: WenyXu Signed-off-by: onepizzateam * chore!: gate soft-drop table behind the enterprise feature (#8747) * chore: gate soft-drop table behind the enterprise feature Soft-drop table becomes an enterprise-only feature: - metasrv rejects gc.experimental_soft_drop.enable=true at startup in non-enterprise builds, and ddl_soft_drop_enabled is hard-disabled without the enterprise feature as a second line of defense - the UNDROP TABLE parser/AST/statement variant, ADMIN purge_table() registration, and information_schema.recycle_bin registration are compiled out unless the enterprise feature is enabled - common-meta procedures, tombstone keys, and DdlTask serde stay unconditional for persisted-procedure recovery and wire compatibility - the [gc.experimental_soft_drop] section is removed from the OSS example config and generated docs (moving to the enterprise repo) - the soft-drop sqlness cases and their CI job are removed from OSS (moving to the enterprise repo); affected information_schema .result files are regenerated Signed-off-by: Lei, HUANG * refactor: limit unused_variables allow to non-enterprise builds Addresses review comment: apply the allow via cfg_attr so enterprise builds still catch accidental unused variables in register_admin_only. Signed-off-by: Lei, HUANG * refactor: include the config key in the soft-drop enterprise gate error Addresses review comment: name gc.experimental_soft_drop.enable in the startup validation error so users can locate the setting quickly when it is set via env vars or layered config. Signed-off-by: Lei, HUANG * test: limit unused_mut allow to non-enterprise builds Addresses review comment: apply the allow via cfg_attr so enterprise builds still catch unused mut in the table_ddl_event test setup. Signed-off-by: Lei, HUANG * feat: reject soft-drop DDL submissions in non-enterprise builds Addresses review comment: clients could bypass the SQL-level gates by submitting DdlTask::UndropTable or DdlTask::PurgeDroppedTable directly to the procedure service. Reject fresh submissions at the DdlManager boundary in non-enterprise builds while keeping the procedure loaders registered for crash recovery and wire compatibility. Signed-off-by: Lei, HUANG * test: stop --enable-gc from enabling soft drop in the sqlness template Addresses review comment: the metasrv test template rendered [gc.experimental_soft_drop] enable = true under the generic --enable-gc flag, which non-enterprise metasrv now rejects at startup, making the documented --enable-gc mode unusable in OSS. Keep the flag scoped to plain GC; enterprise soft-drop coverage moves to the enterprise repo. Signed-off-by: Lei, HUANG * fix: gate fresh soft-drop procedures Signed-off-by: Lei, HUANG * test: gate soft-drop fallback coverage Signed-off-by: Lei, HUANG * fix: gate soft-drop procedure implementation Signed-off-by: Lei, HUANG * refactor: gate drop table soft-drop behavior Signed-off-by: Lei, HUANG * refactor: gate expired soft-drop gc behavior Signed-off-by: Lei, HUANG * ci: test enterprise table ddl lifecycle Signed-off-by: Lei, HUANG * chore: mark purge_table as enterprise licensed The purge_table module is compiled only with the enterprise feature, so apply the Enterprise License header and register it with both license header configurations. Signed-off-by: Lei, HUANG * chore: mark recycle_bin as enterprise licensed The recycle_bin module is compiled only with the enterprise feature, so apply the Enterprise License header and register it with both license header configurations. Signed-off-by: Lei, HUANG * chore: mark soft-drop procedure sources as enterprise licensed The purge and undrop procedure implementations plus the recycle-bin test module compile only with the enterprise feature. Apply the Enterprise License header and register them with both license configurations. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG Signed-off-by: onepizzateam * feat: make frontend heartbeat extensible and lifecycle-safe (#8726) * feat: make frontend heartbeat extensible and lifecycle-safe Signed-off-by: jeremyhi * fix: isolate heartbeat extension response handlers Signed-off-by: jeremyhi * fix: cancel in-flight heartbeat response handling Signed-off-by: jeremyhi * test: cover heartbeat wire compatibility Signed-off-by: jeremyhi * fix: clean up failed heartbeat startup Signed-off-by: jeremyhi * fix: address frontend heartbeat review feedback Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi Signed-off-by: onepizzateam * fix(meta): release region guards after drop rollback (#8751) Signed-off-by: WenyXu Signed-off-by: onepizzateam * refactor!: move native histogram config and `prom_validation_mode` to prom_store (#8744) * chore: adjust the position of experimental_enable_prometheus_native_histogram Signed-off-by: shuiyisong * chore: move prom_validation_mode as well Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong Signed-off-by: onepizzateam * feat: add health-aware gRPC client routing (#8684) * feat: add gRPC client health routing Signed-off-by: WenyXu * fix: harden gRPC client health routing Signed-off-by: WenyXu * fix: defer gRPC client health checks until first use Signed-off-by: WenyXu --------- Signed-off-by: WenyXu Signed-off-by: onepizzateam * feat(mito2): discard unflushed region data safely (#8600) * feat: support discarding unflushed region data Signed-off-by: evenyag * fix(mito2): wake stalled writers after discard Signed-off-by: evenyag * refactor(mito2): drop redundant manifest check for discarding unflushed data Signed-off-by: evenyag --------- Signed-off-by: evenyag Signed-off-by: onepizzateam * docs: align wal.sync_period documented default with actual fallback (5s) (#8753) The example TOMLs and generated config.md documented the default of wal.sync_period as "10s", but since #5677 moved the WAL sync task to a background RepeatedTask, an unset sync_period falls back to 5s in RaftEngineLogStore. The two paths therefore had different fsync periods: deployments based on the example configs used 10s while bare configs used 5s. Align the documentation with the actual code behavior (5s) instead of changing the code fallback to 10s, so that no existing deployment silently gets a larger data-loss window on host power loss. - config/datanode.example.toml, config/standalone.example.toml: 10s -> 5s - config/config.md: regenerated via make config-docs - src/cmd/tests/load_config_test.rs: update assertions accordingly Signed-off-by: jeremyhi Signed-off-by: onepizzateam * fix: support Utf8View labels in Prometheus response (#8754) Signed-off-by: evenyag Signed-off-by: onepizzateam * fix(query): validate merge scan remote schema (#8579) * fix(query): validate merge scan remote schema Signed-off-by: discord9 * fix(query): treat JSON columns as schema-compatible across wire/decode forms CI (Sqlness json2_limit standalone + distributed) failed on the new remote-schema validation: a JSON column is Binary + extension metadata (ARROW:extension:name=greptime.json, greptime:type=Json) on the wire but decodes to Struct(...) with the extension metadata — validate_remote_schema compared raw arrow data_type and rejected it as a mismatch. Adds json_fields_compatible(): JSON fields are equal when name and nullability match, greptime:type matches, and the JSON2 settings (ARROW:extension:metadata type hints) match, ignoring the physical arrow type. Only JSON fields may bypass the raw-type comparison; non-JSON validation stays strict. Adds 4 regression tests mirroring the CI failure (wire-binary vs decoded-struct accepted both directions; different JSON2 settings rejected; JSON vs plain Binary rejected). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: onepizzateam * feat(grafana): add events dashboard (#8725) * feat(grafana): add events dashboard Signed-off-by: WenyXu * fix(grafana): tolerate evolving event schemas Signed-off-by: WenyXu * fix(grafana): address events dashboard review Signed-off-by: WenyXu * fix(grafana): restore events dashboard panels Signed-off-by: WenyXu * fix(grafana): bound events dashboard queries Signed-off-by: WenyXu * fix(grafana): include historical event catalogs Signed-off-by: WenyXu * fix(grafana): preserve events drill-down context Signed-off-by: WenyXu * fix(grafana): correct events lifecycle outcomes Signed-off-by: WenyXu * fix(grafana): handle empty event type ranges Signed-off-by: WenyXu * fix(grafana): scope event catalogs to submissions Signed-off-by: WenyXu * fix(grafana): handle empty events dashboard Signed-off-by: WenyXu * fix(grafana): refresh event schema variables Signed-off-by: WenyXu --------- Signed-off-by: WenyXu Signed-off-by: onepizzateam * refactor: separate a json2 extension type (#8745) Signed-off-by: luofucong Signed-off-by: onepizzateam * feat: add admin function registrar (#8762) * feat: add admin function registrar Signed-off-by: jeremyhi * fix: reject admin function name collisions Signed-off-by: jeremyhi * chore: fix typo in admin function test Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi Signed-off-by: onepizzateam * fix: information_schema.rs table initialization issue Signed-off-by: onepizzateam * rustfmt fix Signed-off-by: onepizzateam --------- Signed-off-by: onepizzateam Signed-off-by: Palak Jha Signed-off-by: polar Signed-off-by: Dennis Zhuang Signed-off-by: Lei, HUANG Signed-off-by: shuiyisong Signed-off-by: discord9 Signed-off-by: WenyXu Signed-off-by: xhwhis Signed-off-by: lyang24 Signed-off-by: Lei, HUANG Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: jeremyhi Signed-off-by: evenyag Signed-off-by: luofucong Co-authored-by: dennis zhuang Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com> Co-authored-by: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Co-authored-by: discord9 Co-authored-by: Weny Xu Co-authored-by: Ning Sun Co-authored-by: Whis Liao Co-authored-by: Lanqing Yang Co-authored-by: sun Co-authored-by: Ning Sun Co-authored-by: jeremyhi Co-authored-by: Yingwen Co-authored-by: LFC <990479+MichaelScofield@users.noreply.github.com> --- .../src/system_schema/information_schema.rs | 11 + .../information_schema/flow_statistics.rs | 285 ++++++++++++++++++ .../information_schema/table_names.rs | 1 + src/cli/src/metadata/control/put/key.rs | 4 +- src/common/catalog/src/consts.rs | 2 + src/common/meta/src/key/flow/flow_state.rs | 67 +++- src/flow/src/adapter/flownode_impl.rs | 4 + src/flow/src/adapter/stat.rs | 26 +- src/flow/src/adapter/worker.rs | 76 +++-- src/flow/src/batching_mode/engine.rs | 22 +- src/flow/src/batching_mode/state.rs | 22 +- src/flow/src/batching_mode/task.rs | 8 + src/flow/src/compute/state.rs | 12 + src/frontend/src/instance.rs | 5 + .../src/handler/flow_state_handler.rs | 6 +- src/operator/src/statement.rs | 1 + src/operator/src/statement/show.rs | 16 +- src/query/src/sql.rs | 48 ++- src/servers/src/postgres/handler.rs | 47 +++ src/sql/src/parsers/show_parser.rs | 32 +- src/sql/src/statements/show.rs | 15 + src/sql/src/statements/statement.rs | 8 +- src/sql/src/util.rs | 1 + .../standalone/common/flow/flow_status.result | 76 +++++ .../standalone/common/flow/flow_status.sql | 42 +++ .../common/show/show_databases_tables.result | 3 + .../common/system/information_schema.result | 7 + .../standalone/common/view/create.result | 1 + 28 files changed, 767 insertions(+), 81 deletions(-) create mode 100644 src/catalog/src/system_schema/information_schema/flow_statistics.rs create mode 100644 tests/cases/standalone/common/flow/flow_status.result create mode 100644 tests/cases/standalone/common/flow/flow_status.sql diff --git a/src/catalog/src/system_schema/information_schema.rs b/src/catalog/src/system_schema/information_schema.rs index a51810e048c..d0539a062c5 100644 --- a/src/catalog/src/system_schema/information_schema.rs +++ b/src/catalog/src/system_schema/information_schema.rs @@ -14,6 +14,7 @@ mod cluster_info; pub mod columns; +pub mod flow_statistics; pub mod flows; mod information_memory_table; pub mod key_column_usage; @@ -73,6 +74,7 @@ use crate::CatalogManager; use crate::error::{Error, Result}; use crate::process_manager::ProcessManagerRef; use crate::system_schema::information_schema::cluster_info::InformationSchemaClusterInfo; +use crate::system_schema::information_schema::flow_statistics::InformationSchemaFlowStatistics; use crate::system_schema::information_schema::flows::InformationSchemaFlows; use crate::system_schema::information_schema::information_memory_table::get_schema_columns; use crate::system_schema::information_schema::key_column_usage::InformationSchemaKeyColumnUsage; @@ -271,6 +273,11 @@ impl SystemSchemaProviderInner for InformationSchemaProvider { self.catalog_manager.clone(), self.flow_metadata_manager.clone(), )) as _), + FLOW_STATISTICS => Some(Arc::new(InformationSchemaFlowStatistics::new( + self.catalog_name.clone(), + self.catalog_manager.clone(), + self.flow_metadata_manager.clone(), + )) as _), PROCEDURE_INFO => Some( Arc::new(procedure_info::InformationSchemaProcedureInfo::new( self.catalog_manager.clone(), @@ -406,6 +413,10 @@ impl InformationSchemaProvider { self.build_table(STATISTICS).unwrap(), ); tables.insert(FLOWS.to_string(), self.build_table(FLOWS).unwrap()); + tables.insert( + FLOW_STATISTICS.to_string(), + self.build_table(FLOW_STATISTICS).unwrap(), + ); #[cfg(feature = "enterprise")] tables.insert( RECYCLE_BIN.to_string(), diff --git a/src/catalog/src/system_schema/information_schema/flow_statistics.rs b/src/catalog/src/system_schema/information_schema/flow_statistics.rs new file mode 100644 index 00000000000..6163461eda7 --- /dev/null +++ b/src/catalog/src/system_schema/information_schema/flow_statistics.rs @@ -0,0 +1,285 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::{Arc, Weak}; + +use common_catalog::consts::INFORMATION_SCHEMA_FLOW_STATISTICS_TABLE_ID; +use common_error::ext::BoxedError; +use common_meta::key::FlowId; +use common_meta::key::flow::FlowMetadataManager; +use common_meta::key::flow::flow_state::FlowStat; +use common_recordbatch::adapter::RecordBatchStreamAdapter; +use common_recordbatch::{DfSendableRecordBatchStream, RecordBatch, SendableRecordBatchStream}; +use common_time::util::current_time_millis; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter; +use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream; +use datatypes::prelude::ConcreteDataType as CDT; +use datatypes::scalars::ScalarVectorBuilder; +use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; +use datatypes::timestamp::TimestampMillisecond; +use datatypes::value::Value; +use datatypes::vectors::{ + Int64VectorBuilder, StringVectorBuilder, TimestampMillisecondVectorBuilder, + UInt32VectorBuilder, UInt64VectorBuilder, VectorRef, +}; +use futures::TryStreamExt; +use snafu::ResultExt; +use store_api::storage::{ScanRequest, TableId}; + +use crate::CatalogManager; +use crate::error::{CreateRecordBatchSnafu, InternalSnafu, ListFlowsSnafu, Result}; +use crate::information_schema::{FLOW_STATISTICS, Predicates}; +use crate::system_schema::information_schema::InformationTable; +use crate::system_schema::utils; + +const INIT_CAPACITY: usize = 42; + +// rows of information_schema.flow_statistics +pub const FLOW_ID: &str = "flow_id"; +pub const FLOW_NAME: &str = "flow_name"; +pub const START_TIME: &str = "start_time"; +pub const LAST_EXECUTION_TIME: &str = "last_execution_time"; +pub const UPTIME_SECONDS: &str = "uptime_seconds"; +pub const STATE_SIZE: &str = "state_size"; + +/// The `information_schema.flow_statistics` provides runtime statistics about flows. +#[derive(Debug)] +pub(super) struct InformationSchemaFlowStatistics { + schema: SchemaRef, + catalog_name: String, + catalog_manager: Weak, + flow_metadata_manager: Arc, +} + +impl InformationSchemaFlowStatistics { + pub(super) fn new( + catalog_name: String, + catalog_manager: Weak, + flow_metadata_manager: Arc, + ) -> Self { + Self { + schema: Self::schema(), + catalog_name, + catalog_manager, + flow_metadata_manager, + } + } + + pub(crate) fn schema() -> SchemaRef { + Arc::new(Schema::new( + vec![ + (FLOW_ID, CDT::uint32_datatype(), false), + (FLOW_NAME, CDT::string_datatype(), false), + (START_TIME, CDT::timestamp_millisecond_datatype(), true), + ( + LAST_EXECUTION_TIME, + CDT::timestamp_millisecond_datatype(), + true, + ), + (UPTIME_SECONDS, CDT::int64_datatype(), true), + (STATE_SIZE, CDT::uint64_datatype(), true), + ] + .into_iter() + .map(|(name, ty, nullable)| ColumnSchema::new(name, ty, nullable)) + .collect(), + )) + } + + fn builder(&self) -> InformationSchemaFlowStatisticsBuilder { + InformationSchemaFlowStatisticsBuilder::new( + self.schema.clone(), + self.catalog_name.clone(), + self.catalog_manager.clone(), + &self.flow_metadata_manager, + ) + } +} + +impl InformationTable for InformationSchemaFlowStatistics { + fn table_id(&self) -> TableId { + INFORMATION_SCHEMA_FLOW_STATISTICS_TABLE_ID + } + + fn table_name(&self) -> &'static str { + FLOW_STATISTICS + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn to_stream(&self, request: ScanRequest) -> Result { + let schema = self.schema.arrow_schema().clone(); + let mut builder = self.builder(); + let stream = Box::pin(DfRecordBatchStreamAdapter::new( + schema, + futures::stream::once(async move { + builder + .make_flow_statistics(Some(request)) + .await + .map(|x| x.into_df_record_batch()) + .map_err(|err| datafusion::error::DataFusionError::External(Box::new(err))) + }), + )); + Ok(Box::pin( + RecordBatchStreamAdapter::try_new(stream) + .map_err(BoxedError::new) + .context(InternalSnafu)?, + )) + } +} + +/// Builds the `information_schema.flow_statistics` table row by row. +struct InformationSchemaFlowStatisticsBuilder { + schema: SchemaRef, + catalog_name: String, + catalog_manager: Weak, + flow_metadata_manager: Arc, + + flow_ids: UInt32VectorBuilder, + flow_names: StringVectorBuilder, + start_times: TimestampMillisecondVectorBuilder, + last_execution_times: TimestampMillisecondVectorBuilder, + uptime_seconds: Int64VectorBuilder, + state_sizes: UInt64VectorBuilder, +} + +impl InformationSchemaFlowStatisticsBuilder { + fn new( + schema: SchemaRef, + catalog_name: String, + catalog_manager: Weak, + flow_metadata_manager: &Arc, + ) -> Self { + Self { + schema, + catalog_name, + catalog_manager, + flow_metadata_manager: flow_metadata_manager.clone(), + + flow_ids: UInt32VectorBuilder::with_capacity(INIT_CAPACITY), + flow_names: StringVectorBuilder::with_capacity(INIT_CAPACITY), + start_times: TimestampMillisecondVectorBuilder::with_capacity(INIT_CAPACITY), + last_execution_times: TimestampMillisecondVectorBuilder::with_capacity(INIT_CAPACITY), + uptime_seconds: Int64VectorBuilder::with_capacity(INIT_CAPACITY), + state_sizes: UInt64VectorBuilder::with_capacity(INIT_CAPACITY), + } + } + + /// Construct the `information_schema.flow_statistics` virtual table. + async fn make_flow_statistics(&mut self, request: Option) -> Result { + let catalog_name = self.catalog_name.clone(); + let predicates = Predicates::from_scan_request(&request); + + let flow_info_manager = self.flow_metadata_manager.clone(); + + let mut stream = flow_info_manager + .flow_name_manager() + .flow_names(&catalog_name) + .await; + + let flow_stat = { + let information_extension = utils::information_extension(&self.catalog_manager)?; + information_extension.flow_stats().await? + }; + + let now = current_time_millis(); + + while let Some((flow_name, flow_id)) = stream + .try_next() + .await + .map_err(BoxedError::new) + .context(ListFlowsSnafu { + catalog: &catalog_name, + })? + { + self.add_flow_statistic(&predicates, flow_id.flow_id(), &flow_name, &flow_stat, now); + } + + self.finish() + } + + fn add_flow_statistic( + &mut self, + predicates: &Predicates, + flow_id: FlowId, + flow_name: &str, + flow_stat: &Option, + now: i64, + ) { + let row = [ + (FLOW_ID, &Value::from(flow_id)), + (FLOW_NAME, &Value::from(flow_name.to_string())), + ]; + if !predicates.eval(&row) { + return; + } + + let start_time = flow_stat + .as_ref() + .and_then(|stat| stat.start_time_map.get(&flow_id).copied()); + + self.flow_ids.push(Some(flow_id)); + self.flow_names.push(Some(flow_name)); + self.start_times + .push(start_time.map(TimestampMillisecond::new)); + self.last_execution_times + .push(flow_stat.as_ref().and_then(|stat| { + stat.last_exec_time_map + .get(&flow_id) + .map(|v| TimestampMillisecond::new(*v)) + })); + self.uptime_seconds + .push(start_time.map(|start| ((now - start) / 1000).max(0))); + self.state_sizes.push( + flow_stat + .as_ref() + .and_then(|stat| stat.state_size.get(&flow_id).map(|v| *v as u64)), + ); + } + + fn finish(&mut self) -> Result { + let columns: Vec = vec![ + Arc::new(self.flow_ids.finish()), + Arc::new(self.flow_names.finish()), + Arc::new(self.start_times.finish()), + Arc::new(self.last_execution_times.finish()), + Arc::new(self.uptime_seconds.finish()), + Arc::new(self.state_sizes.finish()), + ]; + RecordBatch::new(self.schema.clone(), columns).context(CreateRecordBatchSnafu) + } +} + +impl DfPartitionStream for InformationSchemaFlowStatistics { + fn schema(&self) -> &arrow_schema::SchemaRef { + self.schema.arrow_schema() + } + + fn execute(&self, _: Arc) -> DfSendableRecordBatchStream { + let schema: Arc = self.schema.arrow_schema().clone(); + let mut builder = self.builder(); + Box::pin(DfRecordBatchStreamAdapter::new( + schema, + futures::stream::once(async move { + builder + .make_flow_statistics(None) + .await + .map(|x| x.into_df_record_batch()) + .map_err(Into::into) + }), + )) + } +} diff --git a/src/catalog/src/system_schema/information_schema/table_names.rs b/src/catalog/src/system_schema/information_schema/table_names.rs index 5fa1aa517cc..f161fecb54b 100644 --- a/src/catalog/src/system_schema/information_schema/table_names.rs +++ b/src/catalog/src/system_schema/information_schema/table_names.rs @@ -44,6 +44,7 @@ pub const TABLE_CONSTRAINTS: &str = "table_constraints"; pub const CLUSTER_INFO: &str = "cluster_info"; pub const VIEWS: &str = "views"; pub const FLOWS: &str = "flows"; +pub const FLOW_STATISTICS: &str = "flow_statistics"; pub const PROCEDURE_INFO: &str = "procedure_info"; pub const REGION_INFO: &str = "region_info"; pub const REGION_STATISTICS: &str = "region_statistics"; diff --git a/src/cli/src/metadata/control/put/key.rs b/src/cli/src/metadata/control/put/key.rs index 7becfd72dce..49d4b61e2f4 100644 --- a/src/cli/src/metadata/control/put/key.rs +++ b/src/cli/src/metadata/control/put/key.rs @@ -345,9 +345,7 @@ mod tests { #[test] fn test_validate_exact_flow_state_key() { - let value = FlowStateValue::new(BTreeMap::new(), BTreeMap::new()) - .try_as_raw_value() - .unwrap(); + let value = FlowStateValue::default().try_as_raw_value().unwrap(); validate_metadata_value(&flow_state_full_key(), &value).unwrap(); } diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index adbd3a8d473..3149f6ccd4e 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -120,6 +120,8 @@ pub const INFORMATION_SCHEMA_TABLE_SEMANTICS_TABLE_ID: u32 = 42; pub const INFORMATION_SCHEMA_STATISTICS_TABLE_ID: u32 = 43; /// id for information_schema.recycle_bin pub const INFORMATION_SCHEMA_RECYCLE_BIN_TABLE_ID: u32 = 44; +/// id for information_schema.flow_statistics +pub const INFORMATION_SCHEMA_FLOW_STATISTICS_TABLE_ID: u32 = 45; // ----- End of information_schema tables ----- diff --git a/src/common/meta/src/key/flow/flow_state.rs b/src/common/meta/src/key/flow/flow_state.rs index 1b161929fec..77ccd6206ab 100644 --- a/src/common/meta/src/key/flow/flow_state.rs +++ b/src/common/meta/src/key/flow/flow_state.rs @@ -93,22 +93,28 @@ impl<'a> MetadataKey<'a, FlowStateKey> for FlowStateKey { } /// The value of flow state size -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] pub struct FlowStateValue { /// For each key, the bytes of the state in memory pub state_size: BTreeMap, /// For each key, the last execution time of flow in unix timestamp milliseconds. pub last_exec_time_map: BTreeMap, + /// For each flow, the time the flow first executed, in unix timestamp milliseconds. + /// TODO(#7987-followup): not yet propagated via the heartbeat wire format in distributed mode. + #[serde(default)] + pub start_time_map: BTreeMap, } impl FlowStateValue { pub fn new( state_size: BTreeMap, last_exec_time_map: BTreeMap, + start_time_map: BTreeMap, ) -> Self { Self { state_size, last_exec_time_map, + start_time_map, } } } @@ -147,12 +153,15 @@ impl FlowStateManager { } /// Flow's state report, send regularly through heartbeat message -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub struct FlowStat { /// For each key, the bytes of the state in memory pub state_size: BTreeMap, /// For each key, the last execution time of flow in unix timestamp milliseconds. pub last_exec_time_map: BTreeMap, + /// For each flow, the time the flow first executed, in unix timestamp milliseconds. + /// TODO(#7987-followup): not yet propagated via the heartbeat wire format in distributed mode. + pub start_time_map: BTreeMap, } impl From for FlowStat { @@ -160,6 +169,7 @@ impl From for FlowStat { Self { state_size: value.state_size, last_exec_time_map: value.last_exec_time_map, + start_time_map: value.start_time_map, } } } @@ -169,6 +179,59 @@ impl From for FlowStateValue { Self { state_size: value.state_size, last_exec_time_map: value.last_exec_time_map, + start_time_map: value.start_time_map, } } } + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use crate::key::FlowId; + use crate::key::flow::flow_state::FlowStateValue; + + #[test] + fn test_deserialize_legacy_flow_state_value() { + // Legacy format: only state_size and last_exec_time_map are present, + // without the start_time_map field added in PR #8392. + let legacy_json = + r#"{"state_size":{"1":1024,"2":2048},"last_exec_time_map":{"1":1700000000000}}"#; + let value: FlowStateValue = serde_json::from_str(legacy_json).unwrap(); + + let mut expected_state_size = BTreeMap::new(); + expected_state_size.insert(FlowId::from(1u32), 1024usize); + expected_state_size.insert(FlowId::from(2u32), 2048usize); + assert_eq!(value.state_size, expected_state_size); + + let mut expected_last_exec_time_map = BTreeMap::new(); + expected_last_exec_time_map.insert(FlowId::from(1u32), 1700000000000i64); + assert_eq!(value.last_exec_time_map, expected_last_exec_time_map); + + // serde(default) kicks in: old persisted data must not break, + // and the new field defaults to empty. + assert!(value.start_time_map.is_empty()); + } + + #[test] + fn test_flow_state_value_roundtrip_includes_start_time_map() { + let mut state_size = BTreeMap::new(); + state_size.insert(FlowId::from(1u32), 1024usize); + let mut last_exec_time_map = BTreeMap::new(); + last_exec_time_map.insert(FlowId::from(1u32), 1700000000000i64); + let mut start_time_map = BTreeMap::new(); + start_time_map.insert(FlowId::from(1u32), 1700000000000i64); + + let value = FlowStateValue { + state_size, + last_exec_time_map, + start_time_map, + }; + + let json = serde_json::to_string(&value).unwrap(); + assert!(json.contains("start_time_map")); + + let decoded: FlowStateValue = serde_json::from_str(&json).unwrap(); + assert_eq!(decoded, value); + } +} diff --git a/src/flow/src/adapter/flownode_impl.rs b/src/flow/src/adapter/flownode_impl.rs index 9e3399c1822..7e83d5131d8 100644 --- a/src/flow/src/adapter/flownode_impl.rs +++ b/src/flow/src/adapter/flownode_impl.rs @@ -174,9 +174,13 @@ impl FlowDualEngine { let mut last_exec_time_map = streaming.last_exec_time_map; last_exec_time_map.extend(batching.last_exec_time_map); + let mut start_time_map = streaming.start_time_map; + start_time_map.extend(batching.start_time_map); + FlowStat { state_size, last_exec_time_map, + start_time_map, } } diff --git a/src/flow/src/adapter/stat.rs b/src/flow/src/adapter/stat.rs index 68f8160c870..9521a24c5c9 100644 --- a/src/flow/src/adapter/stat.rs +++ b/src/flow/src/adapter/stat.rs @@ -21,33 +21,27 @@ use crate::engine::FlowStatProvider; impl FlowStatProvider for StreamingEngine { async fn flow_stat(&self) -> FlowStat { - let mut full_report = BTreeMap::new(); + let mut state_size_map = BTreeMap::new(); let mut last_exec_time_map = BTreeMap::new(); + let mut start_time_map = BTreeMap::new(); for worker in self.worker_handles.iter() { - match worker.get_state_size().await { - Ok(state_size) => { - full_report.extend(state_size.into_iter().map(|(k, v)| (k as u32, v))); + match worker.get_full_flow_stat().await { + Ok((sizes, exec_times, start_times)) => { + state_size_map.extend(sizes.into_iter().map(|(k, v)| (k as u32, v))); + last_exec_time_map.extend(exec_times.into_iter().map(|(k, v)| (k as u32, v))); + start_time_map.extend(start_times.into_iter().map(|(k, v)| (k as u32, v))); } Err(err) => { - common_telemetry::error!(err; "Get flow stat size error"); - } - } - - match worker.get_last_exec_time_map().await { - Ok(last_exec_time) => { - last_exec_time_map - .extend(last_exec_time.into_iter().map(|(k, v)| (k as u32, v))); - } - Err(err) => { - common_telemetry::error!(err; "Get last exec time error"); + common_telemetry::error!(err; "Get full flow stat error"); } } } FlowStat { - state_size: full_report, + state_size: state_size_map, last_exec_time_map, + start_time_map, } } } diff --git a/src/flow/src/adapter/worker.rs b/src/flow/src/adapter/worker.rs index 32cc4eb7f51..71fbacb18ff 100644 --- a/src/flow/src/adapter/worker.rs +++ b/src/flow/src/adapter/worker.rs @@ -202,30 +202,24 @@ impl WorkerHandle { } } - pub async fn get_state_size(&self) -> Result, Error> { + pub async fn get_full_flow_stat( + &self, + ) -> Result< + ( + BTreeMap, + BTreeMap, + BTreeMap, + ), + Error, + > { let ret = self .itc_client - .call_with_resp(Request::QueryStateSize) + .call_with_resp(Request::QueryFullFlowStat) .await?; - ret.into_query_state_size().map_err(|ret| { + ret.into_query_full_flow_stat().map_err(|ret| { InternalSnafu { reason: format!( - "Flow Node/Worker itc failed, expect Response::QueryStateSize, found {ret:?}" - ), - } - .build() - }) - } - - pub async fn get_last_exec_time_map(&self) -> Result, Error> { - let ret = self - .itc_client - .call_with_resp(Request::QueryLastExecTimeMap) - .await?; - ret.into_query_last_exec_time_map().map_err(|ret| { - InternalSnafu { - reason: format!( - "Flow Node/Worker get_last_exec_time_map failed, expect Response::QueryLastExecTimeMap, found {ret:?}" + "Flow Node/Worker get_full_flow_stat failed, expected Response::QueryFullFlowStat, found {ret:?}" ), } .build() @@ -408,21 +402,24 @@ impl<'s> Worker<'s> { Some(Response::ContainTask { result: ret }) } Request::Shutdown => return Err(()), - Request::QueryStateSize => { - let mut ret = BTreeMap::new(); + Request::QueryFullFlowStat => { + let mut state_size = BTreeMap::new(); + let mut last_exec_time_map = BTreeMap::new(); + let mut start_time_map = BTreeMap::new(); for (flow_id, task_state) in self.task_states.iter() { - ret.insert(*flow_id, task_state.state.get_state_size()); - } - Some(Response::QueryStateSize { result: ret }) - } - Request::QueryLastExecTimeMap => { - let mut ret = BTreeMap::new(); - for (flow_id, task_state) in self.task_states.iter() { - if let Some(last_exec_time) = task_state.state.last_exec_time() { - ret.insert(*flow_id, last_exec_time); + state_size.insert(*flow_id, task_state.state.get_state_size()); + if let Some(t) = task_state.state.last_exec_time() { + last_exec_time_map.insert(*flow_id, t); + } + if let Some(t) = task_state.state.start_time() { + start_time_map.insert(*flow_id, t); } } - Some(Response::QueryLastExecTimeMap { result: ret }) + Some(Response::QueryFullFlowStat { + state_size, + last_exec_time_map, + start_time_map, + }) } }; Ok(ret) @@ -455,8 +452,7 @@ pub enum Request { flow_id: FlowId, }, Shutdown, - QueryStateSize, - QueryLastExecTimeMap, + QueryFullFlowStat, } #[derive(Debug, EnumAsInner)] @@ -472,13 +468,10 @@ enum Response { result: bool, }, RunAvail, - QueryStateSize { - /// each flow tasks' state size - result: BTreeMap, - }, - QueryLastExecTimeMap { - /// each flow tasks' last execution time - result: BTreeMap, + QueryFullFlowStat { + state_size: BTreeMap, + last_exec_time_map: BTreeMap, + start_time_map: BTreeMap, }, } @@ -604,7 +597,8 @@ mod test { ); tx.send(Batch::empty()).unwrap(); handle.run_available(0, true).await.unwrap(); - assert_eq!(handle.get_state_size().await.unwrap().len(), 1); + let (state_size, _, _) = handle.get_full_flow_stat().await.unwrap(); + assert_eq!(state_size.len(), 1); assert_eq!(sink_rx.recv().await.unwrap(), Batch::empty()); drop(handle); worker_thread_handle.join().unwrap(); diff --git a/src/flow/src/batching_mode/engine.rs b/src/flow/src/batching_mode/engine.rs index 3f65761b558..d05eec92518 100644 --- a/src/flow/src/batching_mode/engine.rs +++ b/src/flow/src/batching_mode/engine.rs @@ -414,14 +414,24 @@ impl BatchingEngine { impl FlowStatProvider for BatchingEngine { async fn flow_stat(&self) -> FlowStat { + let runtime = self.runtime.read().await; + let mut last_exec_time_map = BTreeMap::new(); + let mut start_time_map = BTreeMap::new(); + + for (flow_id, task) in runtime.tasks.iter() { + let id = *flow_id as u32; + if let Some(ts) = task.last_execution_time_millis() { + last_exec_time_map.insert(id, ts); + } + if let Some(ts) = task.start_time_millis() { + start_time_map.insert(id, ts); + } + } + FlowStat { state_size: BTreeMap::new(), - last_exec_time_map: self - .get_last_exec_time_map() - .await - .into_iter() - .map(|(flow_id, timestamp)| (flow_id as u32, timestamp)) - .collect(), + last_exec_time_map, + start_time_map, } } } diff --git a/src/flow/src/batching_mode/state.rs b/src/flow/src/batching_mode/state.rs index 225d76a8647..59c4f0fe10f 100644 --- a/src/flow/src/batching_mode/state.rs +++ b/src/flow/src/batching_mode/state.rs @@ -46,6 +46,8 @@ pub struct TaskState { last_query_duration: Duration, /// Last successful execution time in unix timestamp milliseconds. last_exec_time_millis: Option, + /// First execution time in unix timestamp milliseconds, set once. + start_time_millis: Option, /// Dirty Time windows need to be updated /// mapping of `start -> end` and non-overlapping pub(crate) dirty_time_windows: DirtyTimeWindows, @@ -79,6 +81,7 @@ impl TaskState { last_update_time: Instant::now(), last_query_duration: Duration::from_secs(0), last_exec_time_millis: None, + start_time_millis: None, dirty_time_windows, checkpoint_mode: CheckpointMode::FullSnapshot, pending_fenced_repair: None, @@ -90,8 +93,18 @@ impl TaskState { } } - /// called after last query is done - /// `is_succ` indicate whether the last query is successful + /// Record the first-execution start time. Call this once, just before + /// the first frontend query is dispatched, not after it completes. + pub fn record_start_time_if_first(&mut self) { + if self.start_time_millis.is_none() { + // start_time is recorded just before the first frontend query is dispatched + // (pre-execution), so it may be marginally earlier than the streaming engine's + // start_time which is set post-execution. Both are valid approximations of + // "when this flow first ran". + self.start_time_millis = Some(common_time::util::current_time_millis()); + } + } + pub fn after_query_exec(&mut self, elapsed: Duration, is_succ: bool) { self.exec_state = ExecState::Idle; self.last_query_duration = elapsed; @@ -105,6 +118,11 @@ impl TaskState { self.last_exec_time_millis } + /// First execution time in unix timestamp milliseconds, set once. + pub fn start_time_millis(&self) -> Option { + self.start_time_millis + } + pub fn checkpoint_mode(&self) -> CheckpointMode { self.checkpoint_mode } diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index 8428198de3b..196c9ea42d3 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -290,6 +290,10 @@ impl BatchingTask { self.state.read().unwrap().last_execution_time_millis() } + pub fn start_time_millis(&self) -> Option { + self.state.read().unwrap().start_time_millis() + } + /// Collect flow-related extensions from the task's query context that should be /// forwarded to the frontend (e.g. scheduled time). fn frontend_extensions(&self) -> HashMap { @@ -717,6 +721,10 @@ impl BatchingTask { }; let snapshot_seqs = coverage.snapshot_seqs(); + { + let mut state = self.state.write().unwrap(); + state.record_start_time_if_first(); + } frontend_client .query_with_terminal_metrics( catalog, diff --git a/src/flow/src/compute/state.rs b/src/flow/src/compute/state.rs index 2ccd366194f..b71633f517b 100644 --- a/src/flow/src/compute/state.rs +++ b/src/flow/src/compute/state.rs @@ -47,6 +47,8 @@ pub struct DataflowState { expire_after: Option, /// the last time each subgraph executed last_exec_time: Option, + /// the time the flow first executed, in unix timestamp milliseconds + start_time: Option, } impl DataflowState { @@ -120,11 +122,21 @@ impl DataflowState { pub fn set_last_exec_time(&mut self, time: Timestamp) { self.last_exec_time = Some(time); + if self.start_time.is_none() { + // start_time is recorded at the completion of the first execution + // (post-execution), consistent with how last_exec_time is recorded. + self.start_time = Some(time); + } } pub fn last_exec_time(&self) -> Option { self.last_exec_time } + + /// Returns the time the flow first executed, in unix timestamp milliseconds. + pub fn start_time(&self) -> Option { + self.start_time + } } #[derive(Debug, Clone)] diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 7a4b2f9a412..706e651462b 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1559,6 +1559,11 @@ pub fn check_permission( Statement::ShowFlows(stmt) => { validate_db_permission!(stmt, query_ctx); } + Statement::ShowFlowStatus(_stmt) => { + // Flow statistics are organized based on the catalog dimension and + // filtered by the current catalog, so there is no need to check the + // permission of the database(schema). + } #[cfg(feature = "enterprise")] Statement::ShowTriggers(_stmt) => { // The trigger is organized based on the catalog dimension, so there diff --git a/src/meta-srv/src/handler/flow_state_handler.rs b/src/meta-srv/src/handler/flow_state_handler.rs index 9683fd760a1..a13c18a9e03 100644 --- a/src/meta-srv/src/handler/flow_state_handler.rs +++ b/src/meta-srv/src/handler/flow_state_handler.rs @@ -55,7 +55,11 @@ impl HeartbeatHandler for FlowStateHandler { .iter() .map(|(k, v)| (*k, *v)) .collect(); - let value: FlowStateValue = FlowStateValue::new(state_size, last_exec_time_map); + // TODO(#7987-followup): start_time_map is not yet propagated through the heartbeat + // wire format (`api::v1::meta::FlowStat`); it will always be empty in distributed + // mode until a follow-up PR adds heartbeat propagation. + let value: FlowStateValue = + FlowStateValue::new(state_size, last_exec_time_map, Default::default()); self.flow_state_manager .put(value) .await diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 92e0c04a92b..320e7c2e642 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -259,6 +259,7 @@ impl StatementExecutor { Statement::ShowViews(stmt) => self.show_views(stmt, query_ctx).await, Statement::ShowFlows(stmt) => self.show_flows(stmt, query_ctx).await, + Statement::ShowFlowStatus(stmt) => self.show_flow_status(stmt, query_ctx).await, #[cfg(feature = "enterprise")] Statement::ShowTriggers(stmt) => self.show_triggers(stmt, query_ctx).await, diff --git a/src/operator/src/statement/show.rs b/src/operator/src/statement/show.rs index 33cdb0f3677..4a3fa5ae9b5 100644 --- a/src/operator/src/statement/show.rs +++ b/src/operator/src/statement/show.rs @@ -26,8 +26,9 @@ use sql::ast::ObjectNamePartExt; use sql::statements::OptionMap; use sql::statements::create::Partitions; use sql::statements::show::{ - ShowColumns, ShowCreateFlow, ShowCreateView, ShowDatabases, ShowFlows, ShowIndex, ShowKind, - ShowProcessList, ShowRegion, ShowTableStatus, ShowTables, ShowVariables, ShowViews, + ShowColumns, ShowCreateFlow, ShowCreateView, ShowDatabases, ShowFlowStatus, ShowFlows, + ShowIndex, ShowKind, ShowProcessList, ShowRegion, ShowTableStatus, ShowTables, ShowVariables, + ShowViews, }; use table::TableRef; use table::metadata::{TableInfo, TableType}; @@ -250,6 +251,17 @@ impl StatementExecutor { .context(ExecuteStatementSnafu) } + #[tracing::instrument(skip_all)] + pub(super) async fn show_flow_status( + &self, + stmt: ShowFlowStatus, + query_ctx: QueryContextRef, + ) -> Result { + query::sql::show_flow_status(stmt, &self.query_engine, &self.catalog_manager, query_ctx) + .await + .context(ExecuteStatementSnafu) + } + #[cfg(feature = "enterprise")] #[tracing::instrument(skip_all)] pub(super) async fn show_triggers( diff --git a/src/query/src/sql.rs b/src/query/src/sql.rs index 90c3b4ccd47..fc755deb8a7 100644 --- a/src/query/src/sql.rs +++ b/src/query/src/sql.rs @@ -20,8 +20,9 @@ use std::sync::Arc; use catalog::CatalogManagerRef; use catalog::information_schema::{ - CHARACTER_SETS, COLLATIONS, COLUMNS, FLOWS, REGION_PEERS, SCHEMATA, STATISTICS, TABLES, VIEWS, - columns, flows, process_list, region_peers, schemata, statistics, tables, + CHARACTER_SETS, COLLATIONS, COLUMNS, FLOW_STATISTICS, FLOWS, REGION_PEERS, SCHEMATA, + STATISTICS, TABLES, VIEWS, columns, flow_statistics, flows, process_list, region_peers, + schemata, statistics, tables, }; use common_catalog::consts::{ INFORMATION_SCHEMA_NAME, SEMANTIC_TYPE_FIELD, SEMANTIC_TYPE_PRIMARY_KEY, @@ -57,8 +58,8 @@ use sql::parser::ParserContext; use sql::statements::OptionMap; use sql::statements::create::{CreateDatabase, CreateFlow, CreateView, Partitions, SqlOrTql}; use sql::statements::show::{ - ShowColumns, ShowDatabases, ShowFlows, ShowIndex, ShowKind, ShowProcessList, ShowRegion, - ShowTableStatus, ShowTables, ShowVariables, ShowViews, + ShowColumns, ShowDatabases, ShowFlowStatus, ShowFlows, ShowIndex, ShowKind, ShowProcessList, + ShowRegion, ShowTableStatus, ShowTables, ShowVariables, ShowViews, }; use sql::statements::statement::Statement; use sqlparser::ast::ObjectName; @@ -971,6 +972,45 @@ pub async fn show_flows( .await } +/// Execute [`ShowFlowStatus`] statement and return the [`Output`] if success. +pub async fn show_flow_status( + stmt: ShowFlowStatus, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let projects = vec![ + (flow_statistics::FLOW_ID, flow_statistics::FLOW_ID), + (flow_statistics::FLOW_NAME, flow_statistics::FLOW_NAME), + (flow_statistics::START_TIME, flow_statistics::START_TIME), + ( + flow_statistics::LAST_EXECUTION_TIME, + flow_statistics::LAST_EXECUTION_TIME, + ), + ( + flow_statistics::UPTIME_SECONDS, + flow_statistics::UPTIME_SECONDS, + ), + (flow_statistics::STATE_SIZE, flow_statistics::STATE_SIZE), + ]; + let like_field = Some(flow_statistics::FLOW_NAME); + let sort = vec![col(flow_statistics::FLOW_NAME).sort(true, true)]; + + query_from_information_schema_table( + query_engine, + catalog_manager, + query_ctx, + FLOW_STATISTICS, + vec![], + projects, + vec![], + like_field, + sort, + stmt.kind, + ) + .await +} + #[cfg(feature = "enterprise")] pub async fn show_triggers( stmt: sql::statements::show::trigger::ShowTriggers, diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index e372a542375..484bb6a1f11 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -600,6 +600,53 @@ fn describe_fields( format.format_for(1), ), ]), + // SHOW FLOW STATUS returns six columns; return their descriptions so + // prepared/extended-protocol clients receive the correct row description. + SqlPlan::Statement(Statement::ShowFlowStatus(_), _) => Ok(vec![ + FieldInfo::new( + "flow_id".to_string(), + None, + None, + Type::INT8, // matches type_gt_to_pg(UInt32) — do not use INT4 + format.format_for(0), + ), + FieldInfo::new( + "flow_name".to_string(), + None, + None, + Type::TEXT, + format.format_for(1), + ), + FieldInfo::new( + "start_time".to_string(), + None, + None, + Type::TIMESTAMP, + format.format_for(2), + ), + FieldInfo::new( + "last_execution_time".to_string(), + None, + None, + Type::TIMESTAMP, + format.format_for(3), + ), + FieldInfo::new( + "uptime_seconds".to_string(), + None, + None, + Type::INT8, + format.format_for(4), + ), + FieldInfo::new( + "state_size".to_string(), + None, + None, + Type::NUMERIC, + format.format_for(5), + ), + ]), + // single column show statements SqlPlan::Statement( Statement::ShowTables(_) | Statement::ShowFlows(_) | Statement::ShowViews(_), diff --git a/src/sql/src/parsers/show_parser.rs b/src/sql/src/parsers/show_parser.rs index d6fc35c6753..075c17dabe8 100644 --- a/src/sql/src/parsers/show_parser.rs +++ b/src/sql/src/parsers/show_parser.rs @@ -26,8 +26,8 @@ use crate::error::{ use crate::parser::ParserContext; use crate::statements::show::{ ShowColumns, ShowCreateDatabase, ShowCreateFlow, ShowCreateTable, ShowCreateTableVariant, - ShowCreateView, ShowDatabases, ShowFlows, ShowIndex, ShowKind, ShowProcessList, ShowRegion, - ShowSearchPath, ShowStatus, ShowTableStatus, ShowTables, ShowVariables, ShowViews, + ShowCreateView, ShowDatabases, ShowFlowStatus, ShowFlows, ShowIndex, ShowKind, ShowProcessList, + ShowRegion, ShowSearchPath, ShowStatus, ShowTableStatus, ShowTables, ShowVariables, ShowViews, }; use crate::statements::statement::Statement; @@ -57,6 +57,12 @@ impl ParserContext<'_> { self.parse_show_views() } else if self.consume_token("FLOWS") { self.parse_show_flows() + } else if self.consume_token("FLOW") { + if self.consume_token("STATUS") { + self.parse_show_flow_status() + } else { + self.unsupported(self.peek_token_as_string()) + } } else if self.matches_keyword(Keyword::CHARSET) { self.parser.next_token(); Ok(Statement::ShowCharset(self.parse_show_kind()?)) @@ -587,6 +593,12 @@ impl ParserContext<'_> { Ok(Statement::ShowFlows(ShowFlows { kind, database })) } + fn parse_show_flow_status(&mut self) -> Result { + let kind = self.parse_show_kind()?; + + Ok(Statement::ShowFlowStatus(ShowFlowStatus { kind })) + } + fn parse_show_processlist(&mut self, full: bool) -> Result { match self.parser.next_token().token { Token::EOF | Token::SemiColon => { @@ -1250,6 +1262,22 @@ mod tests { assert_eq!(sql, stmts[0].to_string()); } + #[test] + pub fn test_show_flow_status() { + let sql = "SHOW FLOW STATUS"; + let result = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()); + let stmts = result.unwrap(); + assert_eq!(1, stmts.len()); + assert_eq!( + stmts[0], + Statement::ShowFlowStatus(ShowFlowStatus { + kind: ShowKind::All, + }) + ); + assert_eq!(sql, stmts[0].to_string()); + } + #[test] pub fn test_show_processlist() { let sql = "SHOW PROCESSLIST"; diff --git a/src/sql/src/statements/show.rs b/src/sql/src/statements/show.rs index 77880e4a500..67624192b82 100644 --- a/src/sql/src/statements/show.rs +++ b/src/sql/src/statements/show.rs @@ -256,6 +256,21 @@ impl Display for ShowFlows { } } +/// SQL structure for `SHOW FLOW STATUS`. +#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] +pub struct ShowFlowStatus { + pub kind: ShowKind, +} + +impl Display for ShowFlowStatus { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "SHOW FLOW STATUS")?; + format_kind!(self, f); + + Ok(()) + } +} + /// SQL structure for `SHOW CREATE VIEW`. #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] pub struct ShowCreateView { diff --git a/src/sql/src/statements/statement.rs b/src/sql/src/statements/statement.rs index 0f652f0179e..096acb654d1 100644 --- a/src/sql/src/statements/statement.rs +++ b/src/sql/src/statements/statement.rs @@ -40,8 +40,8 @@ use crate::statements::query::Query; use crate::statements::set_variables::SetVariables; use crate::statements::show::{ ShowColumns, ShowCreateDatabase, ShowCreateFlow, ShowCreateTable, ShowCreateView, - ShowDatabases, ShowFlows, ShowIndex, ShowKind, ShowProcessList, ShowRegion, ShowSearchPath, - ShowStatus, ShowTableStatus, ShowTables, ShowVariables, ShowViews, + ShowDatabases, ShowFlowStatus, ShowFlows, ShowIndex, ShowKind, ShowProcessList, ShowRegion, + ShowSearchPath, ShowStatus, ShowTableStatus, ShowTables, ShowVariables, ShowViews, }; use crate::statements::tql::Tql; use crate::statements::truncate::TruncateTable; @@ -118,6 +118,8 @@ pub enum Statement { ShowCreateTrigger(crate::statements::show::trigger::ShowCreateTrigger), /// SHOW FLOWS ShowFlows(ShowFlows), + /// SHOW FLOW STATUS + ShowFlowStatus(ShowFlowStatus), // SHOW TRIGGERS #[cfg(feature = "enterprise")] ShowTriggers(crate::statements::show::trigger::ShowTriggers), @@ -178,6 +180,7 @@ impl Statement { | Statement::ShowCreateTable(_) | Statement::ShowCreateFlow(_) | Statement::ShowFlows(_) + | Statement::ShowFlowStatus(_) | Statement::ShowCreateView(_) | Statement::ShowStatus(_) | Statement::ShowSearchPath(_) @@ -268,6 +271,7 @@ impl Display for Statement { #[cfg(feature = "enterprise")] Statement::ShowCreateTrigger(s) => s.fmt(f), Statement::ShowFlows(s) => s.fmt(f), + Statement::ShowFlowStatus(s) => s.fmt(f), #[cfg(feature = "enterprise")] Statement::ShowTriggers(s) => s.fmt(f), Statement::ShowCreateDatabase(s) => s.fmt(f), diff --git a/src/sql/src/util.rs b/src/sql/src/util.rs index 03071cfaa2c..c0c237cb0c0 100644 --- a/src/sql/src/util.rs +++ b/src/sql/src/util.rs @@ -335,6 +335,7 @@ fn extract_tables_from_statement(stmt: &Statement, names: &mut HashSet Date: Thu, 6 Aug 2026 16:29:00 +0800 Subject: [PATCH 003/126] ci: automate compatibility version window (#8608) * ci: automate compatibility version window Signed-off-by: discord9 * ci: address compat version window review feedback Address all four review comments on the compat version window automation: - Keep the PR window to the sliding window only (latest patch of the two newest stable minor lines). Exact =vX.Y.Z anchors from case.toml are no longer unioned into from_versions; they are validated by the new --check-anchors mode and exercised by nightly runs via --nightly-window. - Add --published-only: the window is computed over stable git tags that have a published, non-draft GitHub release carrying the sqlness compat artifacts (greptime-linux-amd64 tar.gz and sha256sum), so failed releases cannot land in the window. - Run the updater Python tests plus the window/anchor consistency check in PR and merge-group CI (new compat-updater-check job in integration.yml). - Regenerate tests/compatibility/ci.toml to the current sliding window [v1.0.2, v1.1.4]. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .github/scripts/run-compat.py | 31 +- .../tests/test_update_compat_versions.py | 332 +++++++++++++ .github/scripts/update-compat-versions.py | 468 ++++++++++++++++++ .github/workflows/integration.yml | 32 +- .github/workflows/release.yml | 9 + .github/workflows/update-compat-versions.yml | 136 +++++ tests/compatibility/AGENTS.md | 9 +- tests/compatibility/ci.toml | 11 +- 8 files changed, 1019 insertions(+), 9 deletions(-) create mode 100644 .github/scripts/tests/test_update_compat_versions.py create mode 100644 .github/scripts/update-compat-versions.py create mode 100644 .github/workflows/update-compat-versions.yml diff --git a/.github/scripts/run-compat.py b/.github/scripts/run-compat.py index 1b83798b356..ee2f39a754e 100644 --- a/.github/scripts/run-compat.py +++ b/.github/scripts/run-compat.py @@ -61,6 +61,10 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Pass --preserve-state to the real compat run for artifact upload.", ) + parser.add_argument( + "--from-versions-extra", + help="Comma-separated extra from-versions appended to the ci.toml window (for nightly runs).", + ) return parser.parse_args() @@ -107,6 +111,28 @@ def load_versions( return validated +def load_extra_versions(raw: str | None, window: list[str]) -> list[str]: + """Parse --from-versions-extra, skipping versions already in the window. + + Each entry must match VERSION_RE; like load_versions, no version may appear + twice in the effective list, so window entries and duplicates are skipped. + Order is preserved. + """ + if not raw: + return [] + extra: list[str] = [] + seen: set[str] = set() + for item in raw.split(","): + version = item.strip() + if VERSION_RE.fullmatch(version) is None: + raise SystemExit(f"Invalid compat from-versions-extra version: {version!r}") + if version in window or version in seen: + continue + seen.add(version) + extra.append(version) + return extra + + def check_inputs(runner: Path, to_bins_dir: Path) -> None: if not runner.is_file(): raise SystemExit(f"sqlness-runner binary not found: {runner}") @@ -215,6 +241,7 @@ def main() -> int: check_inputs(runner, to_bins_dir) from_versions = load_versions(config_path, "from_versions") + extra_versions = load_extra_versions(args.from_versions_extra, from_versions) downgrade_to_versions = load_versions( config_path, "downgrade_to_versions", required=False ) @@ -222,8 +249,10 @@ def main() -> int: print("Compatibility from-version window:", flush=True) for version in from_versions: print(f" - {version}", flush=True) + for version in extra_versions: + print(f" - {version} (nightly extra)", flush=True) - for from_version in from_versions: + for from_version in [*from_versions, *extra_versions]: run_for_version( runner=runner, to_bins_dir=to_bins_dir, diff --git a/.github/scripts/tests/test_update_compat_versions.py b/.github/scripts/tests/test_update_compat_versions.py new file mode 100644 index 00000000000..7bfc3d1414b --- /dev/null +++ b/.github/scripts/tests/test_update_compat_versions.py @@ -0,0 +1,332 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import contextlib +import importlib.util +import io +import json +import tempfile +import unittest +import urllib.error +from pathlib import Path +from unittest import mock + + +SCRIPT = Path(__file__).resolve().parents[1] / "update-compat-versions.py" +SPEC = importlib.util.spec_from_file_location("update_compat_versions", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +UPDATER = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(UPDATER) + + +class UpdateCompatVersionsTest(unittest.TestCase): + def write_case(self, root: Path, name: str, from_range: str) -> None: + case_dir = root / "tests" / "compatibility" / "cases" / name + case_dir.mkdir(parents=True) + (case_dir / "case.toml").write_text( + f'name = "{name}"\nfrom_range = {from_range}\n', encoding="utf-8" + ) + + def write_config(self, root: Path, versions: str, trailer: str = "") -> Path: + config = root / "tests" / "compatibility" / "ci.toml" + config.parent.mkdir(parents=True, exist_ok=True) + config.write_text( + "# Keep these comments when the window is regenerated.\n" + f"from_versions = {versions}\n{trailer}", + encoding="utf-8", + ) + return config + + def test_stable_tags_filter_prereleases_and_nightlies(self) -> None: + self.assertEqual( + ["v1.0.0", "v1.0.2", "v1.1.0"], + UPDATER.stable_tags( + ["v1.1.0-nightly-20260101", "v1.0.2", "v1.0.0", "v1.1.0", "v1.0.0-rc.1"] + ), + ) + + def test_sliding_versions_roll_over_to_two_newest_minor_lines(self) -> None: + self.assertEqual( + ["v1.9.4", "v2.0.1"], + UPDATER.sliding_versions(["v1.8.9", "v1.9.0", "v1.9.4", "v2.0.0", "v2.0.1"]), + ) + + def test_sliding_versions_supports_one_minor_line(self) -> None: + self.assertEqual(["v1.1.3"], UPDATER.sliding_versions(["v1.1.0", "v1.1.3"])) + + def test_sliding_versions_rejects_no_stable_tags(self) -> None: + with self.assertRaisesRegex(UPDATER.CompatVersionError, "No stable release tags"): + UPDATER.sliding_versions(["v1.1.0-rc.1", "v1.1.0-nightly-20260101"]) + + def test_effective_versions_returns_only_the_sliding_window(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_case(root, "historical", '["=v1.0.2", ">=v1.1.0"]') + versions = UPDATER.effective_versions(["v1.0.2", "v1.1.0", "v1.1.1", "v1.1.3"]) + self.assertEqual(["v1.0.2", "v1.1.3"], versions) + + def test_nightly_versions_includes_anchors_and_sliding(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_case(root, "historical", '["=v1.0.2", ">=v1.1.0"]') + self.write_case(root, "more_history", '["=v1.1.0", "=v1.1.1"]') + versions = UPDATER.nightly_versions( + ["v1.0.2", "v1.1.0", "v1.1.1", "v1.1.3"], + root / "tests" / "compatibility" / "cases", + ) + self.assertEqual(["v1.0.2", "v1.1.0", "v1.1.1", "v1.1.3"], versions) + + def test_exact_anchor_syntaxes_are_normalized(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_case( + root, + "historical", + '[" = v1.0.2 ", " == v1.1.0 ", " v1.1.1 ", "1.1.2", "v01.01.003-rc.1"]', + ) + anchors = UPDATER.extract_pinned_anchors(root / "tests" / "compatibility" / "cases") + self.assertEqual( + {"v1.0.2", "v1.1.0", "v1.1.1", "v1.1.2", "v1.1.3"}, anchors + ) + + def test_ordered_ranges_and_wildcard_are_not_anchors(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_case( + root, + "non_exact", + '[">= v1.0.0", "<=v1.1.0", ">v1.0.2", "< v2.0.0", "*"]', + ) + anchors = UPDATER.extract_pinned_anchors(root / "tests" / "compatibility" / "cases") + self.assertEqual(set(), anchors) + + def test_update_is_sorted_and_idempotent(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + config = self.write_config( + root, + '[\n "v1.1.0",\n "v1.0.2",\n]', + "# This trailing comment must remain.\nother_setting = true\n", + ) + self.write_case(root, "historical", '["=v1.0.2"]') + tags = ["v1.0.2", "v1.1.0", "v1.1.1", "v1.1.3"] + with mock.patch.object(UPDATER, "git_tags", return_value=tags): + self.assertEqual(0, UPDATER.main(["--repo-root", str(root), "--update"])) + first_update = config.read_text(encoding="utf-8") + self.assertEqual(0, UPDATER.main(["--repo-root", str(root), "--update"])) + self.assertEqual(first_update, config.read_text(encoding="utf-8")) + self.assertIn('from_versions = ["v1.0.2", "v1.1.3"]', first_update) + self.assertIn("# Keep these comments when the window is regenerated.", first_update) + self.assertIn("# This trailing comment must remain.\nother_setting = true\n", first_update) + + def test_check_reports_a_unified_diff_for_stale_config(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_config(root, '["v1.0.2"]') + self.write_case(root, "historical", '["=v1.0.2"]') + stderr = io.StringIO() + with mock.patch.object(UPDATER, "git_tags", return_value=["v1.0.2", "v1.1.3"]): + with contextlib.redirect_stderr(stderr): + result = UPDATER.main(["--repo-root", str(root), "--check"]) + self.assertEqual(1, result) + self.assertIn("--- a/tests/compatibility/ci.toml", stderr.getvalue()) + self.assertIn("+++ b/tests/compatibility/ci.toml", stderr.getvalue()) + + def test_rejects_duplicate_config_versions(self) -> None: + with tempfile.TemporaryDirectory() as directory: + config = self.write_config(Path(directory), '["v1.0.2", "v1.0.2"]') + with self.assertRaisesRegex(UPDATER.CompatVersionError, "Duplicate from_versions"): + UPDATER.load_from_versions(config) + + def test_rejects_unavailable_pinned_anchor(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_case(root, "historical", '["=v1.0.2"]') + with self.assertRaisesRegex(UPDATER.CompatVersionError, "unavailable"): + UPDATER.nightly_versions( + ["v1.1.0"], root / "tests" / "compatibility" / "cases" + ) + + def test_rejects_malformed_exact_constraint(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_case(root, "historical", '["=v1.0"]') + with self.assertRaisesRegex(UPDATER.CompatVersionError, "Malformed version"): + UPDATER.extract_pinned_anchors(root / "tests" / "compatibility" / "cases") + + def test_rejects_malformed_bare_exact_constraint(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_case(root, "historical", '["v1.0"]') + with self.assertRaisesRegex(UPDATER.CompatVersionError, "Malformed version"): + UPDATER.extract_pinned_anchors(root / "tests" / "compatibility" / "cases") + + def test_rejects_malformed_ordered_constraint(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_case(root, "historical", '[">=v1.0"]') + with self.assertRaisesRegex(UPDATER.CompatVersionError, "Malformed version"): + UPDATER.extract_pinned_anchors(root / "tests" / "compatibility" / "cases") + + def test_published_only_excludes_failed_release_tags(self) -> None: + self.assertEqual( + ["v1.0.2", "v1.1.3"], + UPDATER.effective_versions( + ["v1.0.2", "v1.1.3", "v1.1.4"], published={"v1.0.2", "v1.1.3"} + ), + ) + + def test_published_only_empty_intersection_raises(self) -> None: + with self.assertRaisesRegex(UPDATER.CompatVersionError, "published release assets"): + UPDATER.effective_versions(["v1.0.2"], published={"v1.1.4"}) + + def test_required_asset_names(self) -> None: + self.assertEqual( + { + "greptime-linux-amd64-v1.1.4.tar.gz", + "greptime-linux-amd64-v1.1.4.sha256sum", + }, + UPDATER.required_asset_names("v1.1.4"), + ) + + class _FakeResponse: + def __init__(self, payload) -> None: + self._payload = payload + + def __enter__(self): + return self + + def __exit__(self, *exc) -> None: + return None + + def read(self) -> bytes: + return json.dumps(self._payload).encode("utf-8") + + def _release(self, tag: str, *, draft=False, prerelease=False, assets=()) -> dict: + return { + "tag_name": tag, + "draft": draft, + "prerelease": prerelease, + "assets": [{"name": name} for name in assets], + } + + def test_published_release_tags_parses_and_filters(self) -> None: + required = { + "greptime-linux-amd64-v1.1.4.tar.gz", + "greptime-linux-amd64-v1.1.4.sha256sum", + } + page_one = [ + self._release(f"preview-{i}", assets=required) for i in range(96) + ] + [ + self._release("v1.1.4", assets=required), + self._release("v1.1.3", draft=True, assets=required), + self._release("v1.1.2", prerelease=True, assets=UPDATER.required_asset_names("v1.1.2")), + self._release("v1.1.1", assets={"greptime-linux-amd64-v1.1.1.tar.gz"}), + ] + self.assertEqual(100, len(page_one)) + page_two = [ + self._release("v1.0.2", assets=UPDATER.required_asset_names("v1.0.2")), + self._release("v1.0.1-beta", assets=required), + ] + calls: list[str] = [] + + def fake_urlopen(request, *args, **kwargs): + calls.append(request.full_url) + payload = page_two if "page=2" in request.full_url else page_one + return self._FakeResponse(payload) + + with mock.patch.object(UPDATER.urllib.request, "urlopen", side_effect=fake_urlopen): + published = UPDATER.published_release_tags("GreptimeTeam", "greptimedb", None) + # Draft releases are excluded; prerelease releases with the required + # assets are included; releases missing assets or with non-stable tags + # are excluded. + self.assertEqual({"v1.1.4", "v1.1.2", "v1.0.2"}, published) + self.assertEqual(2, len(calls)) + self.assertTrue(all("per_page=100" in url for url in calls)) + + def test_published_release_tags_http_error_raises(self) -> None: + with mock.patch.object( + UPDATER.urllib.request, + "urlopen", + side_effect=urllib.error.HTTPError( + "https://api.github.com/repos/GreptimeTeam/greptimedb/releases", + 403, + "Forbidden", + None, + None, + ), + ): + with self.assertRaisesRegex( + UPDATER.CompatVersionError, "Unable to fetch published releases" + ): + UPDATER.published_release_tags("GreptimeTeam", "greptimedb", None) + + def test_check_anchors_rejects_unpublished_anchor(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_case(root, "historical", '["=v1.0.2", "=v1.1.0"]') + with mock.patch.object(UPDATER, "git_tags", return_value=["v1.0.2", "v1.1.0"]): + with mock.patch.object( + UPDATER, "published_release_tags", return_value={"v1.0.2"} + ): + result = UPDATER.main(["--check-anchors", "--repo-root", str(root)]) + self.assertEqual(2, result) + + def test_check_anchors_passes(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_case(root, "historical", '["=v1.0.2", "=v1.1.0"]') + with mock.patch.object(UPDATER, "git_tags", return_value=["v1.0.2", "v1.1.0"]): + with mock.patch.object( + UPDATER, "published_release_tags", return_value={"v1.0.2", "v1.1.0"} + ): + result = UPDATER.main(["--check-anchors", "--repo-root", str(root)]) + self.assertEqual(0, result) + + def test_nightly_window_prints_csv(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.write_case(root, "historical", '["=v1.0.2"]') + self.write_case(root, "more_history", '["=v1.1.0", "=v1.1.1"]') + stdout = io.StringIO() + with mock.patch.object( + UPDATER, "git_tags", return_value=["v1.0.2", "v1.1.0", "v1.1.1", "v1.1.3"] + ): + with contextlib.redirect_stdout(stdout): + result = UPDATER.main(["--nightly-window", "--repo-root", str(root)]) + self.assertEqual(0, result) + self.assertEqual("v1.0.2,v1.1.0,v1.1.1,v1.1.3\n", stdout.getvalue()) + + def test_update_with_published_only(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + config = self.write_config(root, '["v1.0.2"]') + self.write_case(root, "historical", '["=v1.0.2"]') + with mock.patch.object( + UPDATER, "git_tags", return_value=["v1.0.2", "v1.1.3", "v1.1.4"] + ): + with mock.patch.object( + UPDATER, "published_release_tags", return_value={"v1.0.2", "v1.1.3"} + ): + result = UPDATER.main( + ["--update", "--published-only", "--repo-root", str(root)] + ) + updated = config.read_text(encoding="utf-8") + self.assertEqual(0, result) + self.assertIn('from_versions = ["v1.0.2", "v1.1.3"]', updated) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/update-compat-versions.py b/.github/scripts/update-compat-versions.py new file mode 100644 index 00000000000..2e8fec928fe --- /dev/null +++ b/.github/scripts/update-compat-versions.py @@ -0,0 +1,468 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Update the compatibility CI version window from local release tags and cases. + +The PR window keeps only the sliding window over stable release tags whose +GitHub releases carry the sqlness compat artifacts (greptime-linux-amd64 +tar.gz and sha256sum). Exact =vX.Y.Z `from_range` anchors in case.toml are not +retained in the PR window: they are validated with --check-anchors and +exercised by nightly runs via --nightly-window. +""" + +from __future__ import annotations + +import argparse +import ast +import difflib +import json +import os +import re +import subprocess +import sys +import urllib.error +import urllib.request +from pathlib import Path +from typing import Iterable + + +VERSION_RE = re.compile(r"^v(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)$") +FROM_VERSIONS_RE = re.compile( + r"(?ms)^(?P[ \t]*)from_versions[ \t]*=[ \t]*(?P\[(?:[^\]]|\n)*?\])" +) +FROM_VERSIONS_ASSIGNMENT_RE = re.compile(r"(?m)^[ \t]*from_versions[ \t]*=") +FROM_RANGE_RE = re.compile( + r"(?ms)^[ \t]*from_range[ \t]*=[ \t]*(?P\[(?:[^\]]|\n)*?\])" +) +FROM_RANGE_ASSIGNMENT_RE = re.compile(r"(?m)^[ \t]*from_range[ \t]*=") + + +class CompatVersionError(ValueError): + """A compatibility version source does not satisfy the updater policy.""" + + +def version_key(version: str) -> tuple[int, int, int]: + """Return a sortable semantic-version key for a stable version tag.""" + match = VERSION_RE.fullmatch(version) + if match is None: + raise CompatVersionError(f"Not a stable release tag: {version!r}") + major, minor, patch = (int(part) for part in match.groups()) + return major, minor, patch + + +def stable_tags(tags: Iterable[str]) -> list[str]: + """Keep only exact stable release tags, excluding prereleases and nightlies.""" + return sorted({tag for tag in tags if VERSION_RE.fullmatch(tag)}, key=version_key) + + +def sliding_versions(tags: Iterable[str]) -> list[str]: + """Return the latest patch in each of the two newest stable minor lines.""" + stable = stable_tags(tags) + if not stable: + raise CompatVersionError("No stable release tags found (expected exact v.. tags)") + + latest_by_minor: dict[tuple[int, int], str] = {} + for tag in stable: + major, minor, _ = version_key(tag) + minor_line = (major, minor) + if minor_line not in latest_by_minor or version_key(tag) > version_key( + latest_by_minor[minor_line] + ): + latest_by_minor[minor_line] = tag + + newest_minor_lines = sorted(latest_by_minor, reverse=True)[:2] + return sorted((latest_by_minor[line] for line in newest_minor_lines), key=version_key) + + +def _parse_string_array(array: str, context: str) -> list[str]: + without_comments = "\n".join(line.split("#", 1)[0] for line in array.splitlines()) + try: + values = ast.literal_eval(without_comments) + except (SyntaxError, ValueError) as err: + raise CompatVersionError(f"Malformed string array in {context}: {err}") from err + + if not isinstance(values, list) or any(not isinstance(value, str) for value in values): + raise CompatVersionError(f"{context} must be a TOML string array") + return values + + +def _from_versions_match(content: str, config_path: Path) -> re.Match[str]: + matches = list(FROM_VERSIONS_RE.finditer(content)) + assignments = list(FROM_VERSIONS_ASSIGNMENT_RE.finditer(content)) + if len(assignments) != 1 or len(matches) != 1: + raise CompatVersionError(f"{config_path} must define exactly one from_versions array") + match = matches[0] + remainder = content[match.end() :].split("\n", 1)[0].strip() + if remainder and not remainder.startswith("#"): + raise CompatVersionError(f"Malformed from_versions assignment in {config_path}") + return match + + +def load_from_versions(config_path: Path) -> list[str]: + """Validate and load the current checked-in compatibility window.""" + if not config_path.is_file(): + raise CompatVersionError(f"Compatibility CI config not found: {config_path}") + + match = _from_versions_match(config_path.read_text(encoding="utf-8"), config_path) + versions = _parse_string_array(match.group("array"), str(config_path)) + if not versions: + raise CompatVersionError(f"{config_path} must define a non-empty from_versions list") + + duplicates = {version for version in versions if versions.count(version) > 1} + if duplicates: + raise CompatVersionError( + f"Duplicate from_versions in {config_path}: {', '.join(sorted(duplicates))}" + ) + for version in versions: + version_key(version) + return versions + + +def extract_pinned_anchors(cases_dir: Path) -> set[str]: + """Extract exact from_range constraints that must remain in the CI window.""" + if not cases_dir.is_dir(): + raise CompatVersionError(f"Compatibility cases directory not found: {cases_dir}") + anchors: set[str] = set() + for case_path in sorted(cases_dir.glob("**/case.toml")): + content = case_path.read_text(encoding="utf-8") + assignments = list(FROM_RANGE_ASSIGNMENT_RE.finditer(content)) + matches = list(FROM_RANGE_RE.finditer(content)) + if len(assignments) != 1 or len(matches) != 1: + raise CompatVersionError(f"{case_path} must define exactly one from_range array") + + for constraint in _parse_string_array(matches[0].group("array"), str(case_path)): + anchor = exact_anchor(constraint, case_path) + if anchor is not None: + anchors.add(anchor) + return anchors + + +def _runner_version(raw: str, case_path: Path, constraint: str) -> str: + """Parse and normalize a version exactly like compat_case.rs Version::parse.""" + stripped = raw.strip() + if stripped.startswith("v"): + stripped = stripped[1:] + core = stripped.split("-", 1)[0].split("+", 1)[0] + parts = core.split(".") + if len(parts) != 3: + raise CompatVersionError( + f"Malformed version in from_range constraint in {case_path}: {constraint!r}" + ) + + numbers: list[int] = [] + for part in parts: + if not part.isascii() or not part.isdigit(): + raise CompatVersionError( + f"Malformed version in from_range constraint in {case_path}: {constraint!r}" + ) + value = int(part) + if value > 2**64 - 1: + raise CompatVersionError( + f"Version component overflows u64 in {case_path}: {constraint!r}" + ) + numbers.append(value) + return f"v{numbers[0]}.{numbers[1]}.{numbers[2]}" + + +def exact_anchor(constraint: str, case_path: Path) -> str | None: + """Validate a runner constraint and return its normalized exact anchor.""" + normalized = constraint.strip() + if not normalized: + raise CompatVersionError(f"Empty from_range constraint in {case_path}") + if normalized == "*": + return None + + for operator in (">=", "<=", "==", "=", ">", "<"): + if normalized.startswith(operator): + version = _runner_version(normalized[len(operator) :], case_path, constraint) + return version if operator in ("==", "=") else None + + return _runner_version(normalized, case_path, constraint) + + +def required_asset_names(tag: str) -> set[str]: + """Return the release assets the sqlness compat runner downloads for a tag.""" + return { + f"greptime-linux-amd64-{tag}.tar.gz", + f"greptime-linux-amd64-{tag}.sha256sum", + } + + +def published_release_tags(owner: str, repo: str, token: str | None) -> set[str]: + """Return stable tags whose non-draft GitHub releases carry the required assets. + + Only draft releases are skipped: prerelease-flagged releases with the + required assets are included, and VERSION_RE already excludes prerelease + version tags (e.g. v1.2.0-beta.1). + """ + published: set[str] = set() + page = 1 + while True: + url = ( + f"https://api.github.com/repos/{owner}/{repo}/releases" + f"?per_page=100&page={page}" + ) + headers = { + "Accept": "application/vnd.github+json", + "User-Agent": "update-compat-versions", + } + if token: + headers["Authorization"] = f"Bearer {token}" + request = urllib.request.Request(url, headers=headers) + try: + with urllib.request.urlopen(request) as response: + payload = json.loads(response.read().decode("utf-8")) + except (urllib.error.HTTPError, urllib.error.URLError, OSError, ValueError) as err: + raise CompatVersionError( + f"Unable to fetch published releases from {url}: {err}" + ) from err + if not isinstance(payload, list): + raise CompatVersionError( + f"Unexpected response from {url}: expected a JSON release list" + ) + + for release in payload: + if not isinstance(release, dict): + continue + if release.get("draft"): + continue + tag = release.get("tag_name") + if not isinstance(tag, str) or VERSION_RE.fullmatch(tag) is None: + continue + assets = release.get("assets") + asset_names: set[str] = set() + if isinstance(assets, list): + asset_names = { + asset.get("name") + for asset in assets + if isinstance(asset, dict) and isinstance(asset.get("name"), str) + } + if required_asset_names(tag) <= asset_names: + published.add(tag) + + if len(payload) < 100: + break + page += 1 + return published + + +def effective_versions(tags: Iterable[str], published: set[str] | None = None) -> list[str]: + """Return the PR window: the sliding window over stable release tags. + + Exact =vX.Y.Z case anchors are intentionally excluded; they are validated + with --check-anchors and exercised by nightly runs (--nightly-window). + When ``published`` is given, only stable tags with published release + assets are considered. + """ + stable = stable_tags(tags) + if published is not None: + stable = sorted(set(stable) & set(published), key=version_key) + if not stable: + raise CompatVersionError( + "No stable tags have published release assets " + "(greptime-linux-amd64 tar.gz and sha256sum); cannot compute the sliding window" + ) + return sliding_versions(stable) + + +def nightly_versions( + tags: Iterable[str], cases_dir: Path, published: set[str] | None = None +) -> list[str]: + """Return the nightly window: exact anchors plus the sliding release window.""" + anchors = extract_pinned_anchors(cases_dir) + window = effective_versions(tags, published) + available = set(stable_tags(tags)) + if published is not None: + available &= set(published) + unavailable = sorted(anchors - available, key=version_key) + if unavailable: + raise CompatVersionError( + "Pinned compatibility anchor(s) are unavailable as stable git tags: " + + ", ".join(unavailable) + ) + return sorted(anchors | set(window), key=version_key) + + +def updated_config(content: str, config_path: Path, versions: list[str]) -> str: + """Return config content with only the from_versions assignment replaced.""" + match = _from_versions_match(content, config_path) + # Validate the existing setting before replacing it, so malformed or duplicate + # checked-in configs never get silently repaired by automation. + load_from_versions(config_path) + rendered = ( + f'{match.group("indent")}from_versions = [' + + ", ".join(f'"{version}"' for version in versions) + + "]" + ) + return content[: match.start()] + rendered + content[match.end() :] + + +def git_tags(repo_root: Path) -> list[str]: + try: + result = subprocess.run( + ["git", "tag", "--list"], + cwd=repo_root, + check=True, + text=True, + capture_output=True, + ) + except (OSError, subprocess.CalledProcessError) as err: + raise CompatVersionError(f"Unable to read local git tags in {repo_root}: {err}") from err + return result.stdout.splitlines() + + +def _owner_repo() -> tuple[str, str]: + """Return the (owner, repo) pair from GITHUB_REPOSITORY or the default repo.""" + repository = os.environ.get("GITHUB_REPOSITORY", "GreptimeTeam/greptimedb") + owner, _, repo = repository.partition("/") + if not owner or not repo: + return "GreptimeTeam", "greptimedb" + return owner, repo + + +def _github_token() -> str | None: + """Return the GitHub token from the environment, if any.""" + return os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN") + + +def check_anchors(repo_root: Path, cases_dir: Path) -> int: + """Validate every exact anchor is a stable git tag with a published release. + + Returns 0 when all anchors are runnable; prints problems to stderr and + returns 2 when any anchor fails. + """ + tags = set(git_tags(repo_root)) + anchors = extract_pinned_anchors(cases_dir) + if not anchors: + return 0 + published = published_release_tags(*_owner_repo(), _github_token()) + problems: list[str] = [] + for anchor in sorted(anchors, key=version_key): + if anchor not in tags: + problems.append(f"{anchor}: not a stable git tag") + elif anchor not in published: + problems.append(f"{anchor}: missing published release with required assets") + if problems: + print( + "error: exact from_range anchor(s) are not runnable in compatibility CI:", + file=sys.stderr, + ) + for problem in problems: + print(f" - {problem}", file=sys.stderr) + return 2 + return 0 + + +def _path_from_repo(repo_root: Path, value: str | None, default: str) -> Path: + path = Path(value) if value is not None else Path(default) + return path if path.is_absolute() else repo_root / path + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + modes = parser.add_mutually_exclusive_group() + modes.add_argument("--check", action="store_true", help="Fail if ci.toml is stale.") + modes.add_argument("--update", action="store_true", help="Rewrite only ci.toml from_versions.") + modes.add_argument( + "--nightly-window", + action="store_true", + help="Print the comma-separated nightly window (anchors plus sliding) to stdout.", + ) + parser.add_argument( + "--check-anchors", + action="store_true", + help="Validate exact case anchors are stable git tags with published release assets.", + ) + parser.add_argument( + "--published-only", + action="store_true", + help=( + "With --check/--update, restrict the window to stable tags with " + "published release assets (requires network)." + ), + ) + parser.add_argument( + "--repo-root", + default=str(Path(__file__).resolve().parents[2]), + help="Repository root used for default paths and git tags.", + ) + parser.add_argument("--config", help="Compatibility CI TOML path, relative to --repo-root.") + parser.add_argument("--cases-dir", help="Compatibility cases path, relative to --repo-root.") + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + repo_root = Path(args.repo_root).resolve() + config_path = _path_from_repo(repo_root, args.config, "tests/compatibility/ci.toml") + cases_dir = _path_from_repo(repo_root, args.cases_dir, "tests/compatibility/cases") + + if not (args.check or args.update or args.nightly_window or args.check_anchors): + print( + "error: one of --check, --update, --nightly-window, or --check-anchors is required", + file=sys.stderr, + ) + return 2 + + try: + if args.nightly_window: + versions = nightly_versions(git_tags(repo_root), cases_dir) + print(",".join(versions)) + return 0 + + if args.check or args.update: + content = config_path.read_text(encoding="utf-8") + published = ( + published_release_tags(*_owner_repo(), _github_token()) + if args.published_only + else None + ) + versions = effective_versions(git_tags(repo_root), published) + updated = updated_config(content, config_path, versions) + else: + content = updated = "" + + if args.check: + if content == updated: + if args.check_anchors: + return check_anchors(repo_root, cases_dir) + return 0 + try: + display_path = config_path.relative_to(repo_root) + except ValueError: + display_path = config_path + sys.stderr.writelines( + difflib.unified_diff( + content.splitlines(keepends=True), + updated.splitlines(keepends=True), + fromfile=f"a/{display_path}", + tofile=f"b/{display_path}", + ) + ) + return 1 + + if content != updated: + config_path.write_text(updated, encoding="utf-8") + + if args.check_anchors: + return check_anchors(repo_root, cases_dir) + except (CompatVersionError, OSError) as err: + print(f"error: {err}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 9095c4c19c4..b8868123a4a 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -237,6 +237,23 @@ jobs: with: artifact-name: bins + compat-updater-check: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'merge_group' || github.event_name == 'pull_request') }} + name: Check compat version updater + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + fetch-depth: 0 + - name: Run compat updater unit tests + run: python3 -m unittest discover -s .github/scripts/tests -p 'test_*.py' + - name: Check compat version window and anchors + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: python3 .github/scripts/update-compat-versions.py --check --published-only --check-anchors + compat: if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'merge_group' || github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || (github.event_name == 'pull_request' && (github.event.action == 'ready_for_review' || (github.event.action == 'opened' && github.event.pull_request.draft == false)))) }} name: Compatibility Test (recent releases) @@ -247,6 +264,7 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false + fetch-depth: 0 - name: Download pre-built binaries uses: actions/download-artifact@v4 with: @@ -254,8 +272,20 @@ jobs: path: . - name: Unzip binaries run: tar -xvf ./bins.tar.gz + - name: Compute nightly compat window + id: nightly + if: ${{ github.event_name == 'schedule' }} + shell: bash + run: echo "extra=$(python3 .github/scripts/update-compat-versions.py --nightly-window)" >> "$GITHUB_OUTPUT" - name: Run compatibility test - run: python3 .github/scripts/run-compat.py --preserve-state + env: + NIGHTLY_EXTRA: ${{ steps.nightly.outputs.extra }} + run: | + EXTRA="" + if [[ "${GITHUB_EVENT_NAME}" == "schedule" && -n "${NIGHTLY_EXTRA}" ]]; then + EXTRA="--from-versions-extra=${NIGHTLY_EXTRA}" + fi + python3 .github/scripts/run-compat.py --preserve-state ${EXTRA} - name: Upload compatibility logs if: failure() uses: actions/upload-artifact@v4 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index ee1754dd6b8..dbb0f3c3146 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -599,6 +599,15 @@ jobs: with: version: ${{ needs.allocate-runners.outputs.version }} + update-compat-versions: + name: Update compatibility versions + needs: [allocate-runners, publish-github-release] + if: ${{ needs.publish-github-release.result == 'success' }} + uses: ./.github/workflows/update-compat-versions.yml + with: + release_tag: ${{ needs.allocate-runners.outputs.version }} + secrets: inherit + ### Stop runners ### # It's very necessary to split the job of releasing runners into 'stop-linux-amd64-runner' and 'stop-linux-arm64-runner'. # Because we can terminate the specified EC2 instance immediately after the job is finished without unnecessary waiting. diff --git a/.github/workflows/update-compat-versions.yml b/.github/workflows/update-compat-versions.yml new file mode 100644 index 00000000000..460a2c30815 --- /dev/null +++ b/.github/workflows/update-compat-versions.yml @@ -0,0 +1,136 @@ +name: Update compatibility versions + +on: + workflow_call: + inputs: + release_tag: + description: Stable release tag that triggered the caller. + required: true + type: string + release: + types: [published] + schedule: + # Release events created with GITHUB_TOKEN do not trigger workflows. + - cron: "17 3 * * *" + workflow_dispatch: + +concurrency: + group: update-compat-versions + cancel-in-progress: false + +jobs: + validate-release: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} + runs-on: ubuntu-latest + outputs: + should-run: ${{ steps.gate.outputs.should-run }} + steps: + - id: gate + shell: bash + env: + EVENT_NAME: ${{ github.event_name }} + RELEASE_DRAFT: ${{ github.event.release.draft }} + RELEASE_PRERELEASE: ${{ github.event.release.prerelease }} + RELEASE_TAG: ${{ github.event.release.tag_name }} + WORKFLOW_RELEASE_TAG: ${{ inputs.release_tag }} + run: | + set -euo pipefail + should_run=true + stable_tag_re='^v(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)$' + if [[ -n "${WORKFLOW_RELEASE_TAG}" ]]; then + if [[ ! "${WORKFLOW_RELEASE_TAG}" =~ ${stable_tag_re} ]]; then + should_run=false + fi + elif [[ "${EVENT_NAME}" == "release" ]]; then + if [[ "${RELEASE_DRAFT}" == "true" || "${RELEASE_PRERELEASE}" == "true" || ! "${RELEASE_TAG}" =~ ${stable_tag_re} ]]; then + should_run=false + fi + elif [[ "${EVENT_NAME}" == "workflow_call" ]]; then + should_run=false + fi + echo "should-run=${should_run}" >> "${GITHUB_OUTPUT}" + + update: + needs: validate-release + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && needs.validate-release.outputs.should-run == 'true' }} + runs-on: ubuntu-latest + permissions: + contents: write + pull-requests: write + env: + BRANCH: ci/update-compat-versions + steps: + # Always use main: release events otherwise provide the release-tag revision. + - name: Checkout main + uses: actions/checkout@v4 + with: + ref: main + fetch-depth: 0 + token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + + - name: Validate updater on main + shell: bash + env: + GH_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + run: | + set -euo pipefail + python3 -m unittest discover -s .github/scripts/tests -p 'test_*.py' + python3 -m py_compile .github/scripts/update-compat-versions.py + # Log-only: a stale window is expected here (it is exactly what the + # updater fixes below), so a stale --check must not fail the job. + python3 .github/scripts/update-compat-versions.py --check --published-only || true + python3 .github/scripts/update-compat-versions.py --check-anchors + + - name: Prepare updater branch + id: prepare + shell: bash + env: + GH_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + run: | + set -euo pipefail + git fetch origin main + git config user.name "greptimedb-ci" + git config user.email "greptimedb-ci@users.noreply.github.com" + remote_oid="$(git ls-remote --heads origin "refs/heads/${BRANCH}" | cut -f1)" + if [[ -n "${remote_oid}" ]]; then + git fetch origin "refs/heads/${BRANCH}:refs/remotes/origin/${BRANCH}" + fi + git checkout -B "${BRANCH}" origin/main + git reset --hard origin/main + python3 .github/scripts/update-compat-versions.py --update --published-only + if git diff --quiet -- tests/compatibility/ci.toml; then + echo "Compatibility version window is already current." + exit 0 + fi + desired_blob="$(git hash-object tests/compatibility/ci.toml)" + if [[ -n "${remote_oid}" ]]; then + remote_blob="$(git rev-parse "${remote_oid}:tests/compatibility/ci.toml" 2>/dev/null || true)" + if [[ "${desired_blob}" == "${remote_blob}" ]]; then + echo "Remote updater branch already has the desired compatibility window." + echo "branch-ready=true" >> "${GITHUB_OUTPUT}" + exit 0 + fi + fi + git add tests/compatibility/ci.toml + git commit -s -m "chore(ci): update compatibility versions" + git push --force-with-lease="refs/heads/${BRANCH}:${remote_oid}" origin "HEAD:refs/heads/${BRANCH}" + echo "branch-ready=true" >> "${GITHUB_OUTPUT}" + + - name: Create compatibility version PR + if: ${{ steps.prepare.outputs.branch-ready == 'true' }} + shell: bash + env: + GH_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + run: | + set -euo pipefail + open_prs="$(gh pr list --head "${BRANCH}" --base main --state open --json number --jq 'length')" + if [[ "${open_prs}" != "0" ]]; then + echo "An open compatibility updater PR already exists." + exit 0 + fi + gh pr create \ + --base main \ + --head "${BRANCH}" \ + --title "chore(ci): update compatibility versions" \ + --reviewer "GreptimeTeam/db-approver" \ + --body $'I hereby agree to the terms of the [GreptimeDB CLA](https://github.com/GreptimeTeam/.github/blob/main/CLA.md).\n\n## Refer to a related PR or issue link (optional)\n\nAutomated after a stable release or the daily fallback.\n\n## What\047s changed and what\047s your intention?\n\nUpdates `tests/compatibility/ci.toml` from checked-out `main` tags that have published release assets. The PR window keeps the latest patch of the two most recent stable minor lines; exact `=vX.Y.Z` case anchors are validated nightly instead of being added to the PR window.\n\n## PR Checklist\n\n- [ ] I have written the necessary rustdoc comments.\n- [ ] I have added the necessary unit tests and integration tests.\n- [ ] This PR requires documentation updates.\n- [x] API changes are backward compatible.\n- [x] Schema or data changes are backward compatible.' diff --git a/tests/compatibility/AGENTS.md b/tests/compatibility/AGENTS.md index 41880813eb2..8ef9b6c0117 100644 --- a/tests/compatibility/AGENTS.md +++ b/tests/compatibility/AGENTS.md @@ -54,9 +54,12 @@ would be selected before a real run. ## CI Version Window -- `tests/compatibility/ci.toml` controls the small sliding window of recent - released versions used by the CI job. Do not hard-code old versions - directly in workflow YAML. +- `tests/compatibility/ci.toml` controls the small PR window: the latest patch + of the two most recent stable minor lines whose tags have published release + assets. Do not hard-code old versions directly in workflow YAML. +- Exact `=vX.Y.Z` `from_range` anchors in case.toml are not retained in the PR + window; `--check-anchors` validates that they are released tags and + `--nightly-window` exercises them on nightly schedule runs. - Keep the PR/merge-queue window short; wider compatibility coverage belongs in nightly or release-validation workflows. - Case `from_range` / `to_range` metadata still controls whether each case runs diff --git a/tests/compatibility/ci.toml b/tests/compatibility/ci.toml index 2990fe8be5f..67c0602a509 100644 --- a/tests/compatibility/ci.toml +++ b/tests/compatibility/ci.toml @@ -1,9 +1,12 @@ # Compatibility CI version window. # -# The CI job tests this small sliding window of recent released `from` -# versions against the PR-built `to` binary. Broader historical windows should -# run in nightly or release-validation workflows. -from_versions = ["v1.0.0", "v1.1.0"] +# The PR window keeps the latest patch in each of the two most recent stable +# minor lines whose tags have published release assets. Exact =vX.Y.Z +# `from_range` anchors in case.toml are validated by --check-anchors and +# exercised by nightly runs (--nightly-window); they are not retained in the PR +# window. Run .github/scripts/update-compat-versions.py --update --published-only +# after release tags land. +from_versions = ["v1.0.2", "v1.1.4"] # Recent releases that must be able to reopen tables written by the PR build. downgrade_to_versions = ["v1.1.4"] From c7c1a61d56cac9249c428fe96c7568b41ceb85c9 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 6 Aug 2026 16:29:46 +0800 Subject: [PATCH 004/126] feat: support enabling skip_wal with ALTER TABLE (#8730) * feat: support enabling skip_wal with alter table Signed-off-by: jeremyhi * test: cover skip wal on metric physical region Signed-off-by: jeremyhi * fix: retry skip wal alter on route changes Signed-off-by: jeremyhi * fix: address skip wal review feedback Signed-off-by: jeremyhi * fix: preserve skip wal create options Signed-off-by: jeremyhi * fix: prefer typed skip wal option Signed-off-by: jeremyhi * fix: address skip wal review comments Signed-off-by: jeremyhi * fix: canonicalize legacy skip wal option Signed-off-by: jeremyhi * fix: keep typed skip wal canonical Signed-off-by: jeremyhi * fix: simplify skip wal option tracking Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/common/meta/src/ddl/alter_table.rs | 195 ++++++++++- .../meta/src/ddl/alter_table/executor.rs | 48 ++- src/common/meta/src/ddl/tests/alter_table.rs | 322 +++++++++++++++++- src/common/meta/src/ddl_manager.rs | 62 +++- src/metric-engine/src/engine/alter.rs | 14 +- src/mito2/src/compaction/window.rs | 1 + src/mito2/src/engine/skip_wal_test.rs | 183 +++++++++- src/mito2/src/region.rs | 5 + src/mito2/src/region/options.rs | 18 +- src/mito2/src/region_write_ctx.rs | 5 + src/mito2/src/worker/handle_alter.rs | 42 ++- src/mito2/src/worker/handle_close.rs | 5 +- src/mito2/src/worker/handle_write.rs | 6 +- src/operator/src/expr_helper.rs | 20 +- src/query/src/sql/show_create_table.rs | 29 +- src/store-api/src/region_request.rs | 33 +- src/table/src/metadata.rs | 43 ++- src/table/src/requests.rs | 44 ++- tests/cases/standalone/common/skip_wal.result | 91 +++++ tests/cases/standalone/common/skip_wal.sql | 32 +- 20 files changed, 1133 insertions(+), 65 deletions(-) diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index 1516772e518..59fd2af2985 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -16,12 +16,14 @@ mod executor; mod metadata; mod region_request; +use std::collections::HashSet; use std::vec; use api::region::RegionResponse; use api::v1::alter_table_expr::Kind; use api::v1::{RenameTable, SetTableOptions, UnsetTableOptions}; use async_trait::async_trait; +use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE}; use common_error::ext::BoxedError; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ @@ -33,9 +35,10 @@ use serde::{Deserialize, Serialize}; use snafu::{ResultExt, ensure}; use store_api::metadata::ColumnMetadata; use store_api::metric_engine_consts::TABLE_COLUMN_METADATA_EXTENSION_KEY; +use store_api::storage::RegionId; use strum::AsRefStr; use table::metadata::{TableId, TableInfo}; -use table::requests::REPARTITION_COLUMN_HINT_KEY; +use table::requests::{REPARTITION_COLUMN_HINT_KEY, SKIP_WAL_KEY}; use table::table_reference::TableReference; use crate::ddl::DdlContext; @@ -47,10 +50,12 @@ use crate::ddl::utils::{ MultipleResults, extract_column_metadatas, handle_multiple_results, map_to_procedure_error, sync_follower_regions, }; -use crate::error::{AbortProcedureSnafu, NoLeaderSnafu, PutPoisonSnafu, Result, RetryLaterSnafu}; +use crate::error::{ + AbortProcedureSnafu, NoLeaderSnafu, PutPoisonSnafu, Result, RetryLaterSnafu, UnsupportedSnafu, +}; use crate::key::table_info::TableInfoValue; use crate::key::{DeserializedValueWithBytes, RegionDistribution}; -use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; +use crate::lock_key::{CatalogLock, RegionLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; use crate::poison_key::table_poison_key; use crate::rpc::ddl::AlterTableTask; @@ -70,6 +75,9 @@ pub struct AlterTableProcedure { executor: AlterTableExecutor, } +#[derive(Debug)] +pub(crate) struct RegionRouteChanged; + /// Builds the executor from the [`AlterTableData`]. /// /// # Panics @@ -90,8 +98,17 @@ impl AlterTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::AlterTable"; pub fn new(table_id: TableId, task: AlterTableTask, context: DdlContext) -> Result { + Self::new_with_region_locks(table_id, task, vec![], context) + } + + pub(crate) fn new_with_region_locks( + table_id: TableId, + task: AlterTableTask, + region_locks: Vec, + context: DdlContext, + ) -> Result { task.validate()?; - let data = AlterTableData::new(task, table_id); + let data = AlterTableData::new(task, table_id, region_locks); let executor = build_executor_from_alter_expr(&data); Ok(Self { context, @@ -130,11 +147,55 @@ impl AlterTableProcedure { // Safety: Checked in `AlterTableProcedure::new`. let alter_kind = self.data.task.alter_table.kind.as_ref().unwrap(); - if is_metadata_only_alter(alter_kind) { - self.data.state = AlterTableState::UpdateMetadata; - } else { - self.data.state = AlterTableState::SubmitAlterRegionRequests; - }; + if enables_skip_wal(alter_kind) { + ensure!( + only_enables_skip_wal(alter_kind), + UnsupportedSnafu { + operation: "combining skip_wal = 'true' with other table options".to_string() + } + ); + let engine = table_info_value.table_info.meta.engine.as_str(); + ensure!( + engine == MITO_ENGINE || engine == METRIC_ENGINE, + UnsupportedSnafu { + operation: format!("setting skip_wal on {engine} engine tables") + } + ); + // Persist the irreversible intent before any region stops writing WAL. + let (physical_table_id, physical_table_route) = self + .context + .table_metadata_manager + .table_route_manager() + .get_physical_table_route(self.data.table_id()) + .await?; + ensure!( + physical_table_id == self.data.table_id(), + UnsupportedSnafu { + operation: "setting skip_wal on logical tables".to_string() + } + ); + let current_region_ids = physical_table_route + .region_routes + .iter() + .map(|route| route.region.id) + .collect::>(); + let locked_region_ids = self + .data + .region_locks + .iter() + .copied() + .collect::>(); + if physical_table_route.region_routes.len() != self.data.region_locks.len() + || current_region_ids != locked_region_ids + { + // Procedure lock keys are fixed at submission. Finish this attempt so the + // DDL manager can submit another procedure with locks for the current route. + return Ok(Status::done_with_output(RegionRouteChanged)); + } + self.data.region_distribution = + Some(region_distribution(&physical_table_route.region_routes)); + } + self.data.state = self.data.flow().after_prepare(); Ok(Status::executing(true)) } @@ -182,6 +243,40 @@ impl AlterTableProcedure { ensure!(!leaders.is_empty(), NoLeaderSnafu { table_id }); // Puts the poison before submitting alter region requests to datanodes. self.put_poison(ctx_provider, procedure_id).await?; + let flow = self.data.flow(); + if flow == AlterTableFlow::MetadataFirst { + let results = self + .executor + .on_alter_skip_wal_regions( + &self.context.node_manager, + &physical_table_route.region_routes, + alter_kind, + ) + .await; + + return match handle_multiple_results(results) { + MultipleResults::PartialRetryable(error) => Err(error), + MultipleResults::PartialNonRetryable(error) + | MultipleResults::AllNonRetryable(error) => { + // The metadata already enables skip-WAL. Keep retrying until every + // replica applies the runtime option instead of leaving the table in + // a permanently inconsistent state. + Err(BoxedError::new(error)).context(RetryLaterSnafu { + clean_poisons: true, + }) + } + MultipleResults::AllRetryable(error) => { + Err(BoxedError::new(error)).context(RetryLaterSnafu { + clean_poisons: true, + }) + } + MultipleResults::Ok(results) => { + self.handle_alter_region_response(results)?; + Ok(Status::executing_with_clean_poisons(true)) + } + }; + } + let results = self .executor .on_alter_regions( @@ -240,7 +335,7 @@ impl AlterTableProcedure { "altering table result doesn't contains extension key `{TABLE_COLUMN_METADATA_EXTENSION_KEY}`,leaving the table's column metadata unchanged" ); } - self.data.state = AlterTableState::UpdateMetadata; + self.data.state = self.data.flow().after_regions(); Ok(()) } @@ -272,7 +367,8 @@ impl AlterTableProcedure { let table_info_value = self.data.table_info_value.as_ref().unwrap(); // Safety: Checked in `AlterTableProcedure::new`. let alter_kind = self.data.task.alter_table.kind.as_ref().unwrap(); - let metadata_only_alter = is_metadata_only_alter(alter_kind); + let flow = self.data.flow(); + let metadata_only_alter = flow == AlterTableFlow::MetadataOnly; // Gets the table info from the cache or builds it. let new_info = match &self.new_table_info { @@ -302,7 +398,7 @@ impl AlterTableProcedure { info!( "Updated table metadata for table {table_ref}, table_id: {table_id}, kind: {alter_kind:?}" ); - self.data.state = AlterTableState::InvalidateTableCache; + self.data.state = flow.after_metadata(); Ok(Status::executing(true)) } @@ -322,6 +418,10 @@ impl AlterTableProcedure { lock_key.push(SchemaLock::read(table_ref.catalog, table_ref.schema).into()); lock_key.push(TableLock::Write(table_id).into()); + for region_id in &self.data.region_locks { + lock_key.push(RegionLock::Write(*region_id).into()); + } + // Safety: Checked in `AlterTableProcedure::new`. let alter_kind = self.data.task.alter_table.kind.as_ref().unwrap(); if let Kind::RenameTable(RenameTable { new_table_name }) = alter_kind { @@ -344,6 +444,26 @@ impl AlterTableProcedure { } } +fn enables_skip_wal(alter_kind: &Kind) -> bool { + let Kind::SetTableOptions(SetTableOptions { table_options }) = alter_kind else { + return false; + }; + + table_options + .iter() + .any(|option| option.key == SKIP_WAL_KEY && option.value == "true") +} + +pub(crate) fn only_enables_skip_wal(alter_kind: &Kind) -> bool { + let Kind::SetTableOptions(SetTableOptions { table_options }) = alter_kind else { + return false; + }; + + table_options.len() == 1 + && table_options[0].key == SKIP_WAL_KEY + && table_options[0].value == "true" +} + fn is_metadata_only_alter(alter_kind: &Kind) -> bool { match alter_kind { Kind::RenameTable { .. } => true, @@ -357,6 +477,46 @@ fn is_metadata_only_alter(alter_kind: &Kind) -> bool { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AlterTableFlow { + RegionFirst, + MetadataFirst, + MetadataOnly, +} + +impl AlterTableFlow { + fn from_kind(kind: &Kind) -> Self { + if only_enables_skip_wal(kind) { + Self::MetadataFirst + } else if is_metadata_only_alter(kind) { + Self::MetadataOnly + } else { + Self::RegionFirst + } + } + + fn after_prepare(self) -> AlterTableState { + match self { + Self::RegionFirst => AlterTableState::SubmitAlterRegionRequests, + Self::MetadataFirst | Self::MetadataOnly => AlterTableState::UpdateMetadata, + } + } + + fn after_regions(self) -> AlterTableState { + match self { + Self::RegionFirst => AlterTableState::UpdateMetadata, + Self::MetadataFirst | Self::MetadataOnly => AlterTableState::InvalidateTableCache, + } + } + + fn after_metadata(self) -> AlterTableState { + match self { + Self::MetadataFirst => AlterTableState::SubmitAlterRegionRequests, + Self::RegionFirst | Self::MetadataOnly => AlterTableState::InvalidateTableCache, + } + } +} + #[async_trait] impl Procedure for AlterTableProcedure { fn type_name(&self) -> &str { @@ -451,10 +611,13 @@ pub struct AlterTableData { table_info_value: Option>, /// Region distribution for table in case we need to update region options. region_distribution: Option, + /// Region locks held by irreversible region-option alters. + #[serde(default)] + region_locks: Vec, } impl AlterTableData { - pub fn new(task: AlterTableTask, table_id: TableId) -> Self { + pub fn new(task: AlterTableTask, table_id: TableId, region_locks: Vec) -> Self { Self { state: AlterTableState::Prepare, task, @@ -462,6 +625,7 @@ impl AlterTableData { column_metadatas: vec![], table_info_value: None, region_distribution: None, + region_locks, } } @@ -479,6 +643,11 @@ impl AlterTableData { .map(|value| &value.table_info) } + fn flow(&self) -> AlterTableFlow { + // Safety: Checked in `AlterTableProcedure::new`. + AlterTableFlow::from_kind(self.task.alter_table.kind.as_ref().unwrap()) + } + #[cfg(test)] pub(crate) fn column_metadatas(&self) -> &[ColumnMetadata] { &self.column_metadatas diff --git a/src/common/meta/src/ddl/alter_table/executor.rs b/src/common/meta/src/ddl/alter_table/executor.rs index 673d894a23b..fbacdd1b2f5 100644 --- a/src/common/meta/src/ddl/alter_table/executor.rs +++ b/src/common/meta/src/ddl/alter_table/executor.rs @@ -38,7 +38,7 @@ use crate::key::table_info::TableInfoValue; use crate::key::table_name::TableNameKey; use crate::key::{DeserializedValueWithBytes, RegionDistribution, TableMetadataManagerRef}; use crate::node_manager::NodeManagerRef; -use crate::rpc::router::{RegionRoute, find_leaders, region_distribution}; +use crate::rpc::router::{RegionRoute, find_followers, find_leaders, region_distribution}; /// [AlterTableExecutor] performs: /// - Alters the metadata of the table. @@ -169,26 +169,62 @@ impl AlterTableExecutor { node_manager: &NodeManagerRef, region_routes: &[RegionRoute], kind: Option, + ) -> Vec> { + self.dispatch_alter_region_requests(node_manager, region_routes, kind, false) + .await + } + + /// Alters all replicas for the irreversible skip-WAL flow. + pub(crate) async fn on_alter_skip_wal_regions( + &self, + node_manager: &NodeManagerRef, + region_routes: &[RegionRoute], + kind: Option, + ) -> Vec> { + self.dispatch_alter_region_requests(node_manager, region_routes, kind, true) + .await + } + + async fn dispatch_alter_region_requests( + &self, + node_manager: &NodeManagerRef, + region_routes: &[RegionRoute], + kind: Option, + include_followers: bool, ) -> Vec> { let region_distribution = region_distribution(region_routes); - let leaders = find_leaders(region_routes) + let mut peers = find_leaders(region_routes) .into_iter() .map(|p| (p.id, p)) .collect::>(); + if include_followers { + peers.extend(find_followers(region_routes).into_iter().map(|p| (p.id, p))); + } let total_num_region = region_distribution .values() - .map(|r| r.leader_regions.len()) + .map(|r| { + r.leader_regions.len() + + if include_followers { + r.follower_regions.len() + } else { + 0 + } + }) .sum::(); let mut alter_region_tasks = Vec::with_capacity(total_num_region); for (datanode_id, region_role_set) in region_distribution { - if region_role_set.leader_regions.is_empty() { + let mut region_ids = region_role_set.leader_regions; + if include_followers { + region_ids.extend(region_role_set.follower_regions); + } + if region_ids.is_empty() { continue; } // Safety: must exists. - let peer = leaders.get(&datanode_id).unwrap(); + let peer = peers.get(&datanode_id).unwrap(); let requester = node_manager.datanode(peer).await; - for region_id in region_role_set.leader_regions { + for region_id in region_ids { let region_id = RegionId::new(self.table_id, region_id); let request = make_alter_region_request(region_id, kind.clone()); diff --git a/src/common/meta/src/ddl/tests/alter_table.rs b/src/common/meta/src/ddl/tests/alter_table.rs index 173ca870b89..38fc14e760a 100644 --- a/src/common/meta/src/ddl/tests/alter_table.rs +++ b/src/common/meta/src/ddl/tests/alter_table.rs @@ -24,11 +24,11 @@ use api::v1::{ AddColumn, AddColumns, AlterTableExpr, ColumnDataType, ColumnDef as PbColumnDef, DropColumn, DropColumns, SemanticType, SetTableOptions, }; -use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, FILE_ENGINE}; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_procedure::store::poison_store::PoisonStore; -use common_procedure::{ProcedureId, Status}; +use common_procedure::{Procedure, ProcedureId, Status}; use common_procedure_test::MockContextProvider; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; @@ -38,21 +38,24 @@ use store_api::metric_engine_consts::{ }; use store_api::region_engine::RegionManifestInfo; use store_api::storage::RegionId; -use table::requests::TTL_KEY; +use table::requests::{SKIP_WAL_KEY, TTL_KEY}; use tokio::sync::mpsc::{self}; -use crate::ddl::alter_table::AlterTableProcedure; +use crate::ddl::alter_table::{AlterTableProcedure, RegionRouteChanged}; use crate::ddl::test_util::alter_table::TestAlterTableExprBuilder; use crate::ddl::test_util::create_table::test_create_table_task; use crate::ddl::test_util::datanode_handler::{ AllFailureDatanodeHandler, DatanodeWatcher, PartialSuccessDatanodeHandler, RequestOutdatedErrorDatanodeHandler, }; -use crate::ddl::test_util::{assert_column_name, assert_column_name_and_id}; +use crate::ddl::test_util::{ + assert_column_name, assert_column_name_and_id, create_logical_table, create_physical_table, +}; use crate::error::{Error, Result}; use crate::key::datanode_table::DatanodeTableKey; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; +use crate::lock_key::RegionLock; use crate::node_manager::NodeManagerRef; use crate::peer::Peer; use crate::poison_key::table_poison_key; @@ -594,6 +597,315 @@ async fn test_on_update_table_options() { ); } +#[tokio::test] +async fn test_skip_wal_rejects_mixed_table_options() { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let table_name = "foo"; + let table_id = 1024; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + prepare_table_route(table_id), + HashMap::new(), + ) + .await + .unwrap(); + + let mut procedure = AlterTableProcedure::new( + table_id, + AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![ + api::v1::Option { + key: SKIP_WAL_KEY.to_string(), + value: "true".to_string(), + }, + api::v1::Option { + key: TTL_KEY.to_string(), + value: "1d".to_string(), + }, + ], + })), + }, + }, + ddl_context, + ) + .unwrap(); + + let error = procedure.on_prepare().await.unwrap_err(); + assert_matches!(error, Error::Unsupported { .. }); +} + +#[tokio::test] +async fn test_skip_wal_rejects_logical_table() { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let physical_table_id = create_physical_table(&ddl_context, "physical").await; + let logical_table_id = + create_logical_table(ddl_context.clone(), physical_table_id, "logical").await; + let task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "logical".to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![api::v1::Option { + key: SKIP_WAL_KEY.to_string(), + value: "true".to_string(), + }], + })), + }, + }; + + let mut procedure = + AlterTableProcedure::new(logical_table_id, task, ddl_context.clone()).unwrap(); + let error = procedure.on_prepare().await.unwrap_err(); + assert_matches!(error, Error::Unsupported { .. }); + + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(logical_table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert!(!table_info.meta.options.skip_wal); +} + +#[tokio::test] +async fn test_skip_wal_rejects_file_engine_table() { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let table_name = "file_table"; + let table_id = 1024; + let mut create_task = test_create_table_task(table_name, table_id); + create_task.table_info.meta.engine = FILE_ENGINE.to_string(); + ddl_context + .table_metadata_manager + .create_table_metadata( + create_task.table_info, + prepare_table_route(table_id), + HashMap::new(), + ) + .await + .unwrap(); + let task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![api::v1::Option { + key: SKIP_WAL_KEY.to_string(), + value: "true".to_string(), + }], + })), + }, + }; + + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); + let error = procedure.on_prepare().await.unwrap_err(); + assert_matches!(error, Error::Unsupported { .. }); + + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert!(!table_info.meta.options.skip_wal); +} + +#[test] +fn test_skip_wal_holds_region_locks() { + let table_id = 1024; + let region_ids = vec![RegionId::new(table_id, 2), RegionId::new(table_id, 1)]; + let task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "foo".to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![api::v1::Option { + key: SKIP_WAL_KEY.to_string(), + value: "true".to_string(), + }], + })), + }, + }; + let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let procedure = AlterTableProcedure::new_with_region_locks( + table_id, + task, + region_ids.clone(), + context.clone(), + ) + .unwrap(); + + let lock_key = procedure.lock_key(); + for region_id in ®ion_ids { + assert!( + lock_key + .keys_to_lock() + .any(|key| key == &RegionLock::Write(*region_id).into()) + ); + } + + let recovered = AlterTableProcedure::from_json(&procedure.dump().unwrap(), context).unwrap(); + assert_eq!(lock_key, recovered.lock_key()); +} + +#[tokio::test] +async fn test_skip_wal_detects_region_route_change() { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let table_name = "foo"; + let table_id = 1024; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + prepare_table_route(table_id), + HashMap::new(), + ) + .await + .unwrap(); + + let alter_task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![api::v1::Option { + key: SKIP_WAL_KEY.to_string(), + value: "true".to_string(), + }], + })), + }, + }; + let stale_region_locks = vec![RegionId::new(table_id, 4)]; + let mut procedure = AlterTableProcedure::new_with_region_locks( + table_id, + alter_task, + stale_region_locks, + ddl_context.clone(), + ) + .unwrap(); + + let status = procedure.on_prepare().await.unwrap(); + assert!(status.downcast_output_ref::().is_some()); + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert!(!table_info.meta.options.skip_wal); +} + +#[tokio::test] +async fn test_skip_wal_updates_metadata_before_all_replicas() { + let (tx, mut rx) = mpsc::channel(8); + let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); + let ddl_context = new_ddl_context(node_manager); + let table_name = "foo"; + let table_id = 1024; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + prepare_table_route(table_id), + HashMap::new(), + ) + .await + .unwrap(); + + let alter_task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![api::v1::Option { + key: SKIP_WAL_KEY.to_string(), + value: "true".to_string(), + }], + })), + }, + }; + + let region_locks = prepare_table_route(table_id) + .region_routes() + .unwrap() + .iter() + .map(|route| route.region.id) + .collect(); + let mut procedure = AlterTableProcedure::new_with_region_locks( + table_id, + alter_task, + region_locks, + ddl_context.clone(), + ) + .unwrap(); + procedure.on_prepare().await.unwrap(); + let persisted: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + assert_eq!("UpdateMetadata", persisted["state"]); + assert!(persisted.get("alter_regions_after_metadata").is_none()); + + procedure.on_update_metadata().await.unwrap(); + // A crash can happen after the metadata transaction commits but before + // the procedure state is persisted. Replaying this step must be safe. + procedure.on_update_metadata().await.unwrap(); + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert!(table_info.meta.options.skip_wal); + let persisted: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + assert_eq!("SubmitAlterRegionRequests", persisted["state"]); + + let provider = Arc::new(MockContextProvider::default()); + procedure + .submit_alter_region_requests(ProcedureId::random(), provider.as_ref()) + .await + .unwrap(); + let mut requests = Vec::new(); + while let Ok(request) = rx.try_recv() { + requests.push(request); + } + requests.sort_unstable_by_key(|(peer, _)| peer.id); + assert_eq!(5, requests.len()); + let expected = [ + (1, RegionId::new(table_id, 1)), + (2, RegionId::new(table_id, 2)), + (3, RegionId::new(table_id, 3)), + (4, RegionId::new(table_id, 2)), + (5, RegionId::new(table_id, 1)), + ]; + for ((peer, request), (peer_id, region_id)) in requests.into_iter().zip(expected) { + assert_alter_request(peer, request, peer_id, region_id); + } +} + async fn prepare_alter_table_procedure( node_manager: NodeManagerRef, ) -> (AlterTableProcedure, ProcedureId) { diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 3c01d6acef6..e333b3ccb7f 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -27,12 +27,12 @@ use common_telemetry::tracing_context::{FutureExt, TracingContext}; use common_telemetry::{debug, info, tracing}; use derive_builder::Builder; use snafu::{OptionExt, ResultExt, ensure}; -use store_api::storage::TableId; +use store_api::storage::{RegionId, TableId}; use table::table_name::TableName; use crate::ddl::alter_database::AlterDatabaseProcedure; use crate::ddl::alter_logical_tables::AlterLogicalTablesProcedure; -use crate::ddl::alter_table::AlterTableProcedure; +use crate::ddl::alter_table::{AlterTableProcedure, RegionRouteChanged, only_enables_skip_wal}; use crate::ddl::comment_on::CommentOnProcedure; use crate::ddl::create_database::{CreateDatabaseMetadataCommitterRef, CreateDatabaseProcedure}; use crate::ddl::create_flow::CreateFlowProcedure; @@ -80,6 +80,8 @@ use crate::rpc::ddl::{ TruncateTableTask, UndropTableTask, }; +const MAX_REGION_ROUTE_CHANGE_RETRIES: usize = 3; + /// A configurator that customizes or enhances a [`DdlManager`]. #[async_trait::async_trait] pub trait DdlManagerConfigurator: Send + Sync { @@ -417,12 +419,60 @@ impl DdlManager { .await; } - let context = self.create_context(); - let procedure = AlterTableProcedure::new(table_id, alter_table_task, context)?; + let lock_regions = alter_table_task + .alter_table + .kind + .as_ref() + .is_some_and(only_enables_skip_wal); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let mut route_change_retries = 0; + loop { + // Hold the same physical region locks as migration while validating that + // this route snapshot is still the one the procedure will mutate. + let region_ids_to_lock = if lock_regions { + let (_, route) = self + .table_metadata_manager() + .table_route_manager() + .get_physical_table_route(table_id) + .await?; + route + .region_routes + .iter() + .map(|route| route.region.id) + .collect::>() + } else { + vec![] + }; - self.execute_procedure_and_wait(procedure_with_id).await + let context = self.create_context(); + let procedure = AlterTableProcedure::new_with_region_locks( + table_id, + alter_table_task.clone(), + region_ids_to_lock, + context, + )?; + + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let result = self.execute_procedure_and_wait(procedure_with_id).await?; + if result + .1 + .as_ref() + .is_some_and(|output| output.is::()) + { + if route_change_retries == MAX_REGION_ROUTE_CHANGE_RETRIES { + let source = error::UnexpectedSnafu { + err_msg: format!( + "Region route kept changing while altering table {table_id}" + ), + } + .build(); + return Err(error::Error::retry_later(source)); + } + route_change_retries += 1; + continue; + } + return Ok(result); + } } /// Submits and executes a create table task. diff --git a/src/metric-engine/src/engine/alter.rs b/src/metric-engine/src/engine/alter.rs index 517d66ae865..8bc9ae00b7b 100644 --- a/src/metric-engine/src/engine/alter.rs +++ b/src/metric-engine/src/engine/alter.rs @@ -270,8 +270,6 @@ impl MetricEngineInner { #[cfg(test)] mod test { - use std::time::Duration; - use api::v1::SemanticType; use common_meta::ddl::test_util::assert_column_name_and_id; use common_meta::ddl::utils::{parse_column_metadatas, parse_manifest_infos_from_extensions}; @@ -306,16 +304,16 @@ mod test { "Alter request to physical region is forbidden".to_string() ); - // alter physical region's option should work + // skip WAL on the physical region should be forwarded to the data region let alter_region_option_request = RegionAlterRequest { kind: AlterKind::SetRegionOptions { - options: vec![SetRegionOption::Ttl(Some(Duration::from_secs(500).into()))], + options: vec![SetRegionOption::SkipWal], }, }; - let result = engine_inner - .alter_physical_region(physical_region_id, alter_region_option_request.clone()) - .await; - assert!(result.is_ok()); + engine_inner + .alter_physical_region(physical_region_id, alter_region_option_request) + .await + .unwrap(); // alter logical region let metadata_region = env.metadata_region(); diff --git a/src/mito2/src/compaction/window.rs b/src/mito2/src/compaction/window.rs index b5fb803dbf2..422d600a577 100644 --- a/src/mito2/src/compaction/window.rs +++ b/src/mito2/src/compaction/window.rs @@ -327,6 +327,7 @@ mod tests { compaction_override: false, storage: None, append_mode: false, + skip_wal: false, wal_options: Default::default(), index_options: Default::default(), memtable: None, diff --git a/src/mito2/src/engine/skip_wal_test.rs b/src/mito2/src/engine/skip_wal_test.rs index 2309d91569b..a564411473c 100644 --- a/src/mito2/src/engine/skip_wal_test.rs +++ b/src/mito2/src/engine/skip_wal_test.rs @@ -17,9 +17,12 @@ use std::time::Duration; use api::v1::Rows; use common_wal::options::{WAL_OPTIONS_KEY, WalOptions}; +use store_api::logstore::provider::Provider; +use store_api::mito_engine_options::SKIP_WAL_KEY; use store_api::region_engine::{RegionEngine, RegionRole}; use store_api::region_request::{ - RegionCloseRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, RegionTruncateRequest, + AlterKind, RegionAlterRequest, RegionCloseRequest, RegionOpenRequest, RegionPutRequest, + RegionRequest, RegionTruncateRequest, SetRegionOption, }; use store_api::storage::{RegionId, ScanRequest}; @@ -39,6 +42,184 @@ async fn test_close_region_skip_wal_without_pending_data() { test_close_region_skip_wal(false).await; } +#[tokio::test] +async fn test_alter_skip_wal_stops_wal_and_flushes_on_close() { + let mut env = TestEnv::with_prefix("alter-skip-wal").await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let mut request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows(0, 3), + }, + ) + .await; + + let before_alter = engine + .get_region(region_id) + .unwrap() + .version_control + .current(); + assert!(!matches!( + &engine.get_region(region_id).unwrap().provider, + Provider::Noop + )); + assert!(!before_alter.version.memtables.is_empty()); + + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::SetRegionOptions { + options: vec![SetRegionOption::SkipWal], + }, + }), + ) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::SetRegionOptions { + options: vec![SetRegionOption::SkipWal], + }, + }), + ) + .await + .unwrap(); + + let after_alter = engine + .get_region(region_id) + .unwrap() + .version_control + .current(); + assert!(after_alter.version.options.skip_wal); + assert_eq!(before_alter.last_entry_id, after_alter.last_entry_id); + assert_eq!( + before_alter.version.flushed_sequence, + after_alter.version.flushed_sequence + ); + + put_rows( + &engine, + region_id, + Rows { + schema, + rows: build_rows(3, 6), + }, + ) + .await; + assert_eq!( + before_alter.last_entry_id, + engine + .get_region(region_id) + .unwrap() + .version_control + .current() + .last_entry_id + ); + + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + + request + .options + .insert(SKIP_WAL_KEY.to_string(), "true".to_string()); + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir: request.table_dir, + path_type: store_api::region_request::PathType::Bare, + options: request.options, + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + + let batches = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap(); + assert_eq!( + 6, + batches.iter().map(|batch| batch.num_rows()).sum::() + ); +} + +#[tokio::test] +async fn test_alter_skip_wal_on_follower_survives_promotion() { + let mut env = TestEnv::with_prefix("alter-skip-wal-follower").await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + engine + .set_region_role(region_id, RegionRole::Follower) + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::SetRegionOptions { + options: vec![SetRegionOption::SkipWal], + }, + }), + ) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert!(region.is_follower()); + assert!(region.version().options.skip_wal); + + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + let last_entry_id = region.version_control.current().last_entry_id; + put_rows( + &engine, + region_id, + Rows { + schema, + rows: build_rows(0, 1), + }, + ) + .await; + assert_eq!( + last_entry_id, + region.version_control.current().last_entry_id + ); +} + async fn test_close_region_skip_wal(insert: bool) { common_telemetry::init_default_ut_logging(); let mut env = TestEnv::with_prefix(&format!("close-skip-wal-{}", insert)).await; diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 6ce366799cc..c5a9021bde5 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -278,6 +278,11 @@ impl MitoRegion { version_data.version } + /// Returns whether writes to this region should skip WAL. + pub(crate) fn skip_wal(&self) -> bool { + self.provider == Provider::Noop || self.version().options.skip_wal + } + /// Returns last flush timestamp in millis. pub(crate) fn last_flush_millis(&self) -> i64 { self.last_flush_millis.load(Ordering::Relaxed) diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 1089dcd3ac2..9023a877113 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -94,6 +94,8 @@ pub struct RegionOptions { pub storage: Option, /// If append mode is enabled, the region keeps duplicate rows. pub append_mode: bool, + /// Whether to skip writing new WAL entries. + pub skip_wal: bool, /// Wal options. pub wal_options: WalOptions, /// Index options. @@ -275,6 +277,7 @@ impl RegionOptions { compaction_override, storage: options.storage, append_mode: options.append_mode, + skip_wal: options.skip_wal, wal_options, index_options, memtable, @@ -389,6 +392,8 @@ struct RegionOptionsWithoutEnum { storage: Option, #[serde_as(as = "DisplayFromStr")] append_mode: bool, + #[serde_as(as = "DisplayFromStr")] + skip_wal: bool, #[serde_as(as = "NoneAsEmptyString")] merge_mode: Option, #[serde_as(as = "NoneAsEmptyString")] @@ -406,6 +411,7 @@ impl Default for RegionOptionsWithoutEnum { auto_flush_interval: options.auto_flush_interval, storage: options.storage, append_mode: options.append_mode, + skip_wal: options.skip_wal, merge_mode: options.merge_mode, sst_format: options.sst_format, max_row_group_row_count: options.max_row_group_row_count, @@ -552,7 +558,7 @@ mod tests { use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_wal::options::KafkaWalOptions; - use store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY; + use store_api::mito_engine_options::{SKIP_WAL_KEY, WRITE_BUFFER_SIZE_KEY}; use super::*; @@ -581,6 +587,13 @@ mod tests { assert_eq!(expect, options); } + #[test] + fn test_with_skip_wal() { + let map = make_map(&[(SKIP_WAL_KEY, "true")]); + let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap(); + assert!(options.skip_wal); + } + #[test] fn test_with_auto_flush_interval() { let map = make_map(&[("auto_flush_interval", "5m")]); @@ -937,6 +950,7 @@ mod tests { compaction_override: true, storage: Some("S3".to_string()), append_mode: false, + skip_wal: false, wal_options, index_options: IndexOptions { inverted_index: InvertedIndexOptions { @@ -974,6 +988,7 @@ mod tests { compaction_override: false, storage: Some("S3".to_string()), append_mode: false, + skip_wal: false, wal_options: WalOptions::Kafka(KafkaWalOptions::new("test_topic".to_string())), index_options: IndexOptions { inverted_index: InvertedIndexOptions { @@ -1040,6 +1055,7 @@ mod tests { compaction_override: false, storage: Some("S3".to_string()), append_mode: false, + skip_wal: false, wal_options: WalOptions::Kafka(KafkaWalOptions::new("test_topic".to_string())), index_options: IndexOptions { inverted_index: InvertedIndexOptions { diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index d7fdb1d7097..c6abdb0872c 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -199,6 +199,11 @@ impl RegionWriteCtx { &self.version } + /// Returns whether writes in this context should skip WAL. + pub(crate) fn skip_wal(&self) -> bool { + self.provider == Provider::Noop || self.version.options.skip_wal + } + /// Sets error and marks all write operations are failed. pub(crate) fn set_error(&mut self, err: Arc) { // Set error for all notifiers. diff --git a/src/mito2/src/worker/handle_alter.rs b/src/mito2/src/worker/handle_alter.rs index cf9fdb8ccda..69065436da2 100644 --- a/src/mito2/src/worker/handle_alter.rs +++ b/src/mito2/src/worker/handle_alter.rs @@ -50,8 +50,16 @@ impl RegionWorkerLoop { request: RegionAlterRequest, sender: OptionOutputTx, ) { - let region = match self.regions.writable_non_staging_region(region_id) { - Ok(region) => region, + let skip_wal_only = only_enables_skip_wal(&request.kind); + let (region, is_follower) = match self.regions.writable_non_staging_region(region_id) { + Ok(region) => (region, false), + Err(_) if skip_wal_only => match self.regions.follower_region(region_id) { + Ok(region) => (region, true), + Err(e) => { + sender.send(Err(e)); + return; + } + }, Err(e) => { sender.send(Err(e)); return; @@ -60,6 +68,19 @@ impl RegionWorkerLoop { info!("Try to alter region: {}, request: {:?}", region_id, request); + // Followers only accept skip-WAL, which is an in-memory option change and must + // never enter the leader path that may flush memtables. + if is_follower { + let mut options = region.version().options.clone(); + if !options.skip_wal { + info!("Stop writing WAL for follower region: {}", region_id); + options.skip_wal = true; + region.version_control.alter_options(options); + } + sender.send(Ok(0)); + return; + } + // Gets the version before alter. let version = region.version(); @@ -279,6 +300,12 @@ impl RegionWorkerLoop { all_options_altered = false; } } + SetRegionOption::SkipWal => { + if !current_options.skip_wal { + info!("Stop writing WAL for region: {}", region.region_id); + current_options.skip_wal = true; + } + } } } if all_options_altered { @@ -294,6 +321,14 @@ impl RegionWorkerLoop { } } +fn only_enables_skip_wal(kind: &AlterKind) -> bool { + matches!( + kind, + AlterKind::SetRegionOptions { options } + if matches!(options.as_slice(), [SetRegionOption::SkipWal]) + ) +} + /// Returns the new region options if there are updates to the options. fn new_region_options_on_empty_memtable( current_options: &RegionOptions, @@ -315,7 +350,8 @@ fn new_region_options_on_empty_memtable( SetRegionOption::WriteBufferSize(_) | SetRegionOption::Ttl(_) | SetRegionOption::Twsc(_, _) - | SetRegionOption::AutoFlushInterval(_) => (), + | SetRegionOption::AutoFlushInterval(_) + | SetRegionOption::SkipWal => (), SetRegionOption::Format(format_str) => { // Safety: handle_alter_region_options_fast() has validated this. let new_format = format_str.parse::().unwrap(); diff --git a/src/mito2/src/worker/handle_close.rs b/src/mito2/src/worker/handle_close.rs index ef5702c8367..5a0406c655f 100644 --- a/src/mito2/src/worker/handle_close.rs +++ b/src/mito2/src/worker/handle_close.rs @@ -16,7 +16,6 @@ use common_telemetry::info; use store_api::logstore::LogStore; -use store_api::logstore::provider::Provider; use store_api::region_request::{RegionCloseRequest, RegionFlushReason, RegionFlushRequest}; use store_api::storage::RegionId; @@ -37,10 +36,10 @@ impl RegionWorkerLoop { info!("Try to close region {}, worker: {}", region_id, self.id); - // If the close request asks for a flush, or the region is using Noop WAL, + // If the close request asks for a flush, or the region skips WAL, // and has data in memtable and region is flushable (like, not in follower state), // we should flush it before closing to ensure durability. - if (request.flush_on_close || region.provider == Provider::Noop) + if (request.flush_on_close || region.skip_wal()) && !region .version_control .current() diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 0eb4c7bc39e..c52572df114 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -19,7 +19,6 @@ use std::sync::Arc; use api::v1::OpType; use common_telemetry::{debug, error}; -use common_wal::options::WalOptions; use snafu::ensure; use store_api::codec::PrimaryKeyEncoding; use store_api::logstore::LogStore; @@ -113,8 +112,7 @@ impl RegionWorkerLoop { .start_timer(); let mut wal_writer = self.wal.writer(); for region_ctx in region_ctxs.values_mut() { - if let WalOptions::Noop = ®ion_ctx.version().options.wal_options { - // Skip wal write for noop region. + if region_ctx.skip_wal() { continue; } if let Err(e) = region_ctx.add_wal_entry(&mut wal_writer).map_err(Arc::new) { @@ -124,7 +122,7 @@ impl RegionWorkerLoop { match wal_writer.write_to_wal().await.map_err(Arc::new) { Ok(response) => { for (region_id, region_ctx) in region_ctxs.iter_mut() { - if let WalOptions::Noop = ®ion_ctx.version().options.wal_options { + if region_ctx.skip_wal() { continue; } diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index f6ed0d8092e..45e3e340a33 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -215,12 +215,11 @@ pub fn create_to_expr( .context(ExternalSnafu)?; let time_index = find_time_index(&create.constraints)?; - let table_options = HashMap::from( + let mut table_options = HashMap::from( &TableOptions::try_from_iter(create.options.to_str_map()) .context(UnrecognizedTableOptionSnafu)?, ); - let mut table_options = table_options; if table_options.contains_key(COMPACTION_TYPE) { table_options.insert(COMPACTION_OVERRIDE.to_string(), "true".to_string()); } @@ -1495,6 +1494,23 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; "1.0MiB", expr.table_options.get("write_buffer_size").unwrap() ); + + let sql = "CREATE TABLE monitor (ts TIMESTAMP TIME INDEX) WITH(skip_wal='false');"; + let stmt = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .pop() + .unwrap(); + let Statement::CreateTable(create_table) = stmt else { + unreachable!() + }; + let expr = create_to_expr(&create_table, &QueryContext::arc()).unwrap(); + assert_eq!( + Some("false"), + expr.table_options + .get(store_api::mito_engine_options::SKIP_WAL_KEY) + .map(String::as_str) + ); } #[test] diff --git a/src/query/src/sql/show_create_table.rs b/src/query/src/sql/show_create_table.rs index 0c66ae87387..eb0392f15a1 100644 --- a/src/query/src/sql/show_create_table.rs +++ b/src/query/src/sql/show_create_table.rs @@ -39,7 +39,8 @@ use sql::statements::{self, OptionMap, concrete_data_type_to_sql_data_type}; use store_api::metric_engine_consts::{is_metric_engine, is_metric_engine_internal_column}; use table::metadata::{TableInfoRef, TableMeta}; use table::requests::{ - COMMENT_KEY as TABLE_COMMENT_KEY, FILE_TABLE_META_KEY, TTL_KEY, WRITE_BUFFER_SIZE_KEY, + COMMENT_KEY as TABLE_COMMENT_KEY, FILE_TABLE_META_KEY, SKIP_WAL_KEY, TTL_KEY, + WRITE_BUFFER_SIZE_KEY, }; use crate::error::{ @@ -66,7 +67,6 @@ fn create_sql_options(table_meta: &TableMeta, schema_options: Option for ModifyColumnType { } } +/// Region option changes used by ALTER requests. +/// +/// This type currently derives serde for request persistence. Keep future changes +/// backward compatible with previously serialized variants. #[derive(Debug, Eq, PartialEq, Clone, Serialize, Deserialize)] pub enum SetRegionOption { WriteBufferSize(Option), @@ -1450,6 +1454,8 @@ pub enum SetRegionOption { AutoFlushInterval(Option), // Modifying the max number of rows in a parquet row group. MaxRowGroupRowCount(Option), + // Stops writing new WAL entries. This operation is irreversible. + SkipWal, } impl TryFrom<&PbOption> for SetRegionOption { @@ -1507,6 +1513,7 @@ impl TryFrom<&PbOption> for SetRegionOption { .ok_or_else(|| InvalidSetRegionOptionRequestSnafu { key, value }.build())?; Ok(Self::MaxRowGroupRowCount(Some(row_count))) } + SKIP_WAL_KEY if value == "true" => Ok(Self::SkipWal), _ => InvalidSetRegionOptionRequestSnafu { key, value }.fail(), } } @@ -1879,6 +1886,28 @@ mod tests { assert!(SetRegionOption::try_from(&pb).is_err()); } + #[test] + fn test_set_region_option_skip_wal_try_from() { + let pb = PbOption { + key: SKIP_WAL_KEY.to_string(), + value: "true".to_string(), + }; + assert_eq!( + SetRegionOption::SkipWal, + SetRegionOption::try_from(&pb).unwrap() + ); + + for value in ["false", "", "invalid"] { + let pb = PbOption { + key: SKIP_WAL_KEY.to_string(), + value: value.to_string(), + }; + assert!(SetRegionOption::try_from(&pb).is_err()); + } + + assert!(UnsetRegionOption::try_from(SKIP_WAL_KEY).is_err()); + } + #[test] fn test_set_region_option_max_row_group_row_count_try_from() { let pb = PbOption { diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index 43a88611266..7b8958e22ce 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -30,7 +30,7 @@ use snafu::{OptionExt, ResultExt, ensure}; use store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY; use store_api::mito_engine_options::{ APPEND_MODE_KEY, AUTO_FLUSH_INTERVAL_KEY, COMPACTION_TYPE, COMPACTION_TYPE_TWCS, - MAX_ROW_GROUP_ROW_COUNT, MERGE_MODE_KEY, SST_FORMAT_KEY, + MAX_ROW_GROUP_ROW_COUNT, MERGE_MODE_KEY, SKIP_WAL_KEY, SST_FORMAT_KEY, }; use store_api::region_request::{SetRegionOption, UnsetRegionOption}; use store_api::storage::{ColumnDescriptor, ColumnDescriptorBuilder, ColumnId}; @@ -401,6 +401,14 @@ impl TableMeta { new_options.extra_options.remove(MAX_ROW_GROUP_ROW_COUNT); } } + SetRegionOption::SkipWal => { + new_options.skip_wal = true; + // Keep the explicit table option so it remains distinguishable + // from a value inherited from the schema. + new_options + .extra_options + .insert(SKIP_WAL_KEY.to_string(), true.to_string()); + } } } let mut builder = self.new_meta_builder(); @@ -1687,6 +1695,39 @@ mod tests { ); } + #[test] + fn test_set_skip_wal_updates_typed_and_extra_options() { + let mut meta = TableMetaBuilder::empty() + .schema(Arc::new(new_test_schema())) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + meta.options + .extra_options + .insert(SKIP_WAL_KEY.to_string(), false.to_string()); + + let alter_kind = AlterKind::SetTableOptions { + options: vec![SetRegionOption::SkipWal], + }; + let new_meta = meta + .builder_with_alter_kind("my_table", &alter_kind) + .unwrap() + .build() + .unwrap(); + + assert!(new_meta.options.skip_wal); + assert_eq!( + Some("true"), + new_meta + .options + .extra_options + .get(SKIP_WAL_KEY) + .map(String::as_str) + ); + } + #[test] fn test_set_repartition_column_hint() { let meta = TableMetaBuilder::empty() diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index afaa5da7b5d..8c736e26749 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -227,7 +227,7 @@ impl fmt::Display for TableOptions { key_vals.push(format!("{}={}", TTL_KEY, ttl)); } - if self.skip_wal { + if self.skip_wal && !self.extra_options.contains_key(SKIP_WAL_KEY) { key_vals.push(format!("{}={}", SKIP_WAL_KEY, self.skip_wal)); } @@ -241,7 +241,7 @@ impl fmt::Display for TableOptions { impl From<&TableOptions> for HashMap { fn from(opts: &TableOptions) -> Self { - let mut res = HashMap::with_capacity(2 + opts.extra_options.len()); + let mut res = HashMap::with_capacity(3 + opts.extra_options.len()); if let Some(write_buffer_size) = opts.write_buffer_size { let _ = res.insert( WRITE_BUFFER_SIZE_KEY.to_string(), @@ -251,11 +251,10 @@ impl From<&TableOptions> for HashMap { if let Some(ttl_str) = opts.ttl.map(|ttl| ttl.to_string()) { let _ = res.insert(TTL_KEY.to_string(), ttl_str); } - res.extend( - opts.extra_options - .iter() - .map(|(k, v)| (k.clone(), v.clone())), - ); + if opts.skip_wal { + let _ = res.insert(SKIP_WAL_KEY.to_string(), true.to_string()); + } + res.extend(opts.extra_options.clone()); res } } @@ -555,6 +554,20 @@ mod tests { let serialized = TableOptions::try_from_iter(&serialized_map).unwrap(); assert_eq!(options, serialized); + let options = TableOptions { + write_buffer_size: None, + ttl: None, + extra_options: HashMap::from([(SKIP_WAL_KEY.to_string(), true.to_string())]), + skip_wal: true, + }; + let serialized_map = HashMap::from(&options); + assert_eq!( + Some("true"), + serialized_map.get(SKIP_WAL_KEY).map(String::as_str) + ); + let serialized = TableOptions::try_from_iter(&serialized_map).unwrap(); + assert_eq!(options, serialized); + let options = TableOptions { write_buffer_size: None, ttl: Default::default(), @@ -574,6 +587,15 @@ mod tests { let serialized_map = HashMap::from(&options); let serialized = TableOptions::try_from_iter(&serialized_map).unwrap(); assert_eq!(options, serialized); + + let options = TableOptions { + extra_options: HashMap::from([(SKIP_WAL_KEY.to_string(), false.to_string())]), + skip_wal: false, + ..Default::default() + }; + let serialized_map = HashMap::from(&options); + let serialized = TableOptions::try_from_iter(&serialized_map).unwrap(); + assert_eq!(options, serialized); } #[test] @@ -612,5 +634,13 @@ mod tests { "write_buffer_size=128.0MiB ttl=16m 40s skip_wal=true", options.to_string() ); + + let options = TableOptions { + write_buffer_size: None, + ttl: None, + extra_options: HashMap::from([(SKIP_WAL_KEY.to_string(), "false".to_string())]), + skip_wal: false, + }; + assert_eq!("skip_wal=false", options.to_string()); } } diff --git a/tests/cases/standalone/common/skip_wal.result b/tests/cases/standalone/common/skip_wal.result index dc7866f7fec..5f3ae29f52b 100644 --- a/tests/cases/standalone/common/skip_wal.result +++ b/tests/cases/standalone/common/skip_wal.result @@ -56,3 +56,94 @@ DROP TABLE system_metrics; Affected Rows: 0 +CREATE TABLE alter_skip_wal ( + host STRING, + val DOUBLE, + ts TIMESTAMP, + PRIMARY KEY(host), + TIME INDEX(ts) +); + +Affected Rows: 0 + +INSERT INTO alter_skip_wal VALUES ('host1', 1, 1000); + +Affected Rows: 1 + +ALTER TABLE alter_skip_wal SET 'skip_wal' = 'true'; + +Affected Rows: 0 + +SHOW CREATE TABLE alter_skip_wal; + ++----------------+-----------------------------------------------+ +| Table | Create Table | ++----------------+-----------------------------------------------+ +| alter_skip_wal | CREATE TABLE IF NOT EXISTS "alter_skip_wal" ( | +| | "host" STRING NULL, | +| | "val" DOUBLE NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | skip_wal = 'true' | +| | ) | ++----------------+-----------------------------------------------+ + +ALTER TABLE alter_skip_wal SET 'skip_wal' = 'false'; + +Error: 1004(InvalidArguments), Invalid set table option request: Invalid set region option request, key: skip_wal, value: false + +ALTER TABLE alter_skip_wal UNSET 'skip_wal'; + +Error: 1004(InvalidArguments), Invalid unset table option request: Invalid set region option request, key: skip_wal + +INSERT INTO alter_skip_wal VALUES ('host2', 2, 2000); + +Affected Rows: 1 + +SELECT * FROM alter_skip_wal ORDER BY ts; + ++-------+-----+---------------------+ +| host | val | ts | ++-------+-----+---------------------+ +| host1 | 1.0 | 1970-01-01T00:00:01 | +| host2 | 2.0 | 1970-01-01T00:00:02 | ++-------+-----+---------------------+ + +-- The post-ALTER row can be lost because restart does not flush skip-WAL memtables. +-- SQLNESS ARG restart=true +SHOW CREATE TABLE alter_skip_wal; + ++----------------+-----------------------------------------------+ +| Table | Create Table | ++----------------+-----------------------------------------------+ +| alter_skip_wal | CREATE TABLE IF NOT EXISTS "alter_skip_wal" ( | +| | "host" STRING NULL, | +| | "val" DOUBLE NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | skip_wal = 'true' | +| | ) | ++----------------+-----------------------------------------------+ + +SELECT * FROM alter_skip_wal ORDER BY ts; + ++-------+-----+---------------------+ +| host | val | ts | ++-------+-----+---------------------+ +| host1 | 1.0 | 1970-01-01T00:00:01 | ++-------+-----+---------------------+ + +DROP TABLE alter_skip_wal; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/skip_wal.sql b/tests/cases/standalone/common/skip_wal.sql index 2b55cb7a12b..cdf0d551ab6 100644 --- a/tests/cases/standalone/common/skip_wal.sql +++ b/tests/cases/standalone/common/skip_wal.sql @@ -29,4 +29,34 @@ ADMIN flush_table('system_metrics'); -- SQLNESS ARG restart=true SELECT * FROM system_metrics; -DROP TABLE system_metrics; \ No newline at end of file +DROP TABLE system_metrics; + +CREATE TABLE alter_skip_wal ( + host STRING, + val DOUBLE, + ts TIMESTAMP, + PRIMARY KEY(host), + TIME INDEX(ts) +); + +INSERT INTO alter_skip_wal VALUES ('host1', 1, 1000); + +ALTER TABLE alter_skip_wal SET 'skip_wal' = 'true'; + +SHOW CREATE TABLE alter_skip_wal; + +ALTER TABLE alter_skip_wal SET 'skip_wal' = 'false'; + +ALTER TABLE alter_skip_wal UNSET 'skip_wal'; + +INSERT INTO alter_skip_wal VALUES ('host2', 2, 2000); + +SELECT * FROM alter_skip_wal ORDER BY ts; + +-- The post-ALTER row can be lost because restart does not flush skip-WAL memtables. +-- SQLNESS ARG restart=true +SHOW CREATE TABLE alter_skip_wal; + +SELECT * FROM alter_skip_wal ORDER BY ts; + +DROP TABLE alter_skip_wal; From 62e70027a4a14969ab90842a66f51655ace257f5 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 6 Aug 2026 16:36:56 +0800 Subject: [PATCH 005/126] chore: bump version to 1.3.0 on default branch (#8749) chore: dump version to 1.3.0 on default branch --- Cargo.lock | 152 ++++++++++++++++++++++++++--------------------------- Cargo.toml | 2 +- 2 files changed, 77 insertions(+), 77 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ccfb23c51af..77474ea66fd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -219,7 +219,7 @@ checksum = "d301b3b94cb4b2f23d7917810addbbaff90738e0ca2be692bd027e70d7e0330c" [[package]] name = "api" -version = "1.2.0" +version = "1.3.0" dependencies = [ "arrow-schema 58.3.0", "common-base", @@ -942,7 +942,7 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "auth" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "async-trait", @@ -1592,7 +1592,7 @@ dependencies = [ [[package]] name = "cache" -version = "1.2.0" +version = "1.3.0" dependencies = [ "catalog", "common-error", @@ -1658,7 +1658,7 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "catalog" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "arrow 58.3.0", @@ -1993,7 +1993,7 @@ checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675" [[package]] name = "cli" -version = "1.2.0" +version = "1.3.0" dependencies = [ "async-stream", "async-trait", @@ -2052,7 +2052,7 @@ dependencies = [ [[package]] name = "client" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "arc-swap", @@ -2085,7 +2085,7 @@ dependencies = [ "serde_json", "snafu 0.8.6", "store-api", - "substrait 1.2.0", + "substrait 1.3.0", "tokio", "tokio-stream", "tonic 0.14.2", @@ -2134,7 +2134,7 @@ dependencies = [ [[package]] name = "cmd" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "async-trait", @@ -2286,7 +2286,7 @@ dependencies = [ [[package]] name = "common-base" -version = "1.2.0" +version = "1.3.0" dependencies = [ "ahash 0.8.12", "anymap2", @@ -2306,14 +2306,14 @@ dependencies = [ [[package]] name = "common-catalog" -version = "1.2.0" +version = "1.3.0" dependencies = [ "const_format", ] [[package]] name = "common-config" -version = "1.2.0" +version = "1.3.0" dependencies = [ "common-base", "common-error", @@ -2337,7 +2337,7 @@ dependencies = [ [[package]] name = "common-datasource" -version = "1.2.0" +version = "1.3.0" dependencies = [ "arrow 58.3.0", "arrow-schema 58.3.0", @@ -2373,7 +2373,7 @@ dependencies = [ [[package]] name = "common-decimal" -version = "1.2.0" +version = "1.3.0" dependencies = [ "bigdecimal 0.4.8", "common-error", @@ -2386,7 +2386,7 @@ dependencies = [ [[package]] name = "common-error" -version = "1.2.0" +version = "1.3.0" dependencies = [ "common-macro", "http 1.3.1", @@ -2398,7 +2398,7 @@ dependencies = [ [[package]] name = "common-event-recorder" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "async-trait", @@ -2422,7 +2422,7 @@ dependencies = [ [[package]] name = "common-frontend" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "async-trait", @@ -2444,7 +2444,7 @@ dependencies = [ [[package]] name = "common-function" -version = "1.2.0" +version = "1.3.0" dependencies = [ "ahash 0.8.12", "api", @@ -2507,7 +2507,7 @@ dependencies = [ [[package]] name = "common-greptimedb-telemetry" -version = "1.2.0" +version = "1.3.0" dependencies = [ "async-trait", "common-runtime", @@ -2524,7 +2524,7 @@ dependencies = [ [[package]] name = "common-grpc" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "arrow-flight", @@ -2559,7 +2559,7 @@ dependencies = [ [[package]] name = "common-grpc-expr" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "common-base", @@ -2579,7 +2579,7 @@ dependencies = [ [[package]] name = "common-macro" -version = "1.2.0" +version = "1.3.0" dependencies = [ "greptime-proto", "once_cell", @@ -2590,7 +2590,7 @@ dependencies = [ [[package]] name = "common-mem-prof" -version = "1.2.0" +version = "1.3.0" dependencies = [ "anyhow", "common-error", @@ -2606,7 +2606,7 @@ dependencies = [ [[package]] name = "common-memory-manager" -version = "1.2.0" +version = "1.3.0" dependencies = [ "common-error", "common-macro", @@ -2618,7 +2618,7 @@ dependencies = [ [[package]] name = "common-meta" -version = "1.2.0" +version = "1.3.0" dependencies = [ "anymap2", "api", @@ -2691,7 +2691,7 @@ dependencies = [ [[package]] name = "common-options" -version = "1.2.0" +version = "1.3.0" dependencies = [ "common-base", "common-grpc", @@ -2702,11 +2702,11 @@ dependencies = [ [[package]] name = "common-plugins" -version = "1.2.0" +version = "1.3.0" [[package]] name = "common-pprof" -version = "1.2.0" +version = "1.3.0" dependencies = [ "common-error", "common-macro", @@ -2717,7 +2717,7 @@ dependencies = [ [[package]] name = "common-procedure" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "async-stream", @@ -2746,7 +2746,7 @@ dependencies = [ [[package]] name = "common-procedure-test" -version = "1.2.0" +version = "1.3.0" dependencies = [ "async-trait", "common-procedure", @@ -2756,7 +2756,7 @@ dependencies = [ [[package]] name = "common-query" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "async-trait", @@ -2786,7 +2786,7 @@ dependencies = [ [[package]] name = "common-recordbatch" -version = "1.2.0" +version = "1.3.0" dependencies = [ "arc-swap", "common-base", @@ -2811,7 +2811,7 @@ dependencies = [ [[package]] name = "common-runtime" -version = "1.2.0" +version = "1.3.0" dependencies = [ "async-trait", "clap", @@ -2840,7 +2840,7 @@ dependencies = [ [[package]] name = "common-session" -version = "1.2.0" +version = "1.3.0" dependencies = [ "serde", "strum 0.27.1", @@ -2848,7 +2848,7 @@ dependencies = [ [[package]] name = "common-sql" -version = "1.2.0" +version = "1.3.0" dependencies = [ "arrow-schema 58.3.0", "common-base", @@ -2868,7 +2868,7 @@ dependencies = [ [[package]] name = "common-stat" -version = "1.2.0" +version = "1.3.0" dependencies = [ "common-base", "common-runtime", @@ -2883,7 +2883,7 @@ dependencies = [ [[package]] name = "common-telemetry" -version = "1.2.0" +version = "1.3.0" dependencies = [ "backtrace", "common-base", @@ -2912,7 +2912,7 @@ dependencies = [ [[package]] name = "common-test-util" -version = "1.2.0" +version = "1.3.0" dependencies = [ "client", "common-grpc", @@ -2925,7 +2925,7 @@ dependencies = [ [[package]] name = "common-time" -version = "1.2.0" +version = "1.3.0" dependencies = [ "arrow 58.3.0", "chrono", @@ -2943,7 +2943,7 @@ dependencies = [ [[package]] name = "common-version" -version = "1.2.0" +version = "1.3.0" dependencies = [ "cargo-manifest", "const_format", @@ -2953,7 +2953,7 @@ dependencies = [ [[package]] name = "common-wal" -version = "1.2.0" +version = "1.3.0" dependencies = [ "common-base", "common-error", @@ -2976,7 +2976,7 @@ dependencies = [ [[package]] name = "common-workload" -version = "1.2.0" +version = "1.3.0" dependencies = [ "common-telemetry", "serde", @@ -4451,7 +4451,7 @@ dependencies = [ [[package]] name = "datanode" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "arrow-flight", @@ -4526,7 +4526,7 @@ checksum = "c286de4e81ea2590afc24d754e0f83810c566f50a1388fa75ebd57928c0d9745" [[package]] name = "datatypes" -version = "1.2.0" +version = "1.3.0" dependencies = [ "arrow 58.3.0", "arrow-array 58.3.0", @@ -5263,7 +5263,7 @@ checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" [[package]] name = "file-engine" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "async-trait", @@ -5394,7 +5394,7 @@ checksum = "8bf7cc16383c4b8d58b9905a8509f02926ce3058053c056376248d958c9df1e8" [[package]] name = "flow" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "arrow 58.3.0", @@ -5464,7 +5464,7 @@ dependencies = [ "sql", "store-api", "strum 0.27.1", - "substrait 1.2.0", + "substrait 1.3.0", "table", "tokio", "tokio-stream", @@ -5547,7 +5547,7 @@ checksum = "28dd6caf6059519a65843af8fe2a3ae298b14b80179855aeb4adc2c1934ee619" [[package]] name = "frontend" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "arc-swap", @@ -6850,7 +6850,7 @@ dependencies = [ [[package]] name = "index" -version = "1.2.0" +version = "1.3.0" dependencies = [ "async-trait", "asynchronous-codec", @@ -7971,7 +7971,7 @@ checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" [[package]] name = "log-query" -version = "1.2.0" +version = "1.3.0" dependencies = [ "chrono", "common-error", @@ -7983,7 +7983,7 @@ dependencies = [ [[package]] name = "log-store" -version = "1.2.0" +version = "1.3.0" dependencies = [ "async-stream", "async-trait", @@ -8308,7 +8308,7 @@ dependencies = [ [[package]] name = "meta-client" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "async-trait", @@ -8340,7 +8340,7 @@ dependencies = [ [[package]] name = "meta-srv" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "async-trait", @@ -8440,7 +8440,7 @@ dependencies = [ [[package]] name = "metric-engine" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "aquamarine", @@ -8541,7 +8541,7 @@ dependencies = [ [[package]] name = "mito-codec" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "bytes", @@ -8566,7 +8566,7 @@ dependencies = [ [[package]] name = "mito2" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "aquamarine", @@ -9356,7 +9356,7 @@ dependencies = [ [[package]] name = "object-store" -version = "1.2.0" +version = "1.3.0" dependencies = [ "anyhow", "async-trait", @@ -9927,7 +9927,7 @@ dependencies = [ [[package]] name = "operator" -version = "1.2.0" +version = "1.3.0" dependencies = [ "ahash 0.8.12", "api", @@ -9986,7 +9986,7 @@ dependencies = [ "sql", "sqlparser", "store-api", - "substrait 1.2.0", + "substrait 1.3.0", "table", "tokio", "tokio-util", @@ -10271,7 +10271,7 @@ checksum = "e3c406c9e2aa74554e662d2c2ee11cd3e73756988800be7e6f5eddb16fed4699" [[package]] name = "partition" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "async-trait", @@ -10665,7 +10665,7 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "pipeline" -version = "1.2.0" +version = "1.3.0" dependencies = [ "ahash 0.8.12", "api", @@ -10822,7 +10822,7 @@ dependencies = [ [[package]] name = "plugins" -version = "1.2.0" +version = "1.3.0" dependencies = [ "auth", "catalog", @@ -11162,7 +11162,7 @@ dependencies = [ [[package]] name = "promql" -version = "1.2.0" +version = "1.3.0" dependencies = [ "ahash 0.8.12", "async-trait", @@ -11515,7 +11515,7 @@ dependencies = [ [[package]] name = "puffin" -version = "1.2.0" +version = "1.3.0" dependencies = [ "async-compression", "async-trait", @@ -11577,7 +11577,7 @@ dependencies = [ [[package]] name = "query" -version = "1.2.0" +version = "1.3.0" dependencies = [ "ahash 0.8.12", "api", @@ -11648,7 +11648,7 @@ dependencies = [ "sql", "sqlparser", "store-api", - "substrait 1.2.0", + "substrait 1.3.0", "table", "tokio", "tokio-stream", @@ -13268,7 +13268,7 @@ dependencies = [ [[package]] name = "servers" -version = "1.2.0" +version = "1.3.0" dependencies = [ "ahash 0.8.12", "api", @@ -13408,7 +13408,7 @@ dependencies = [ [[package]] name = "session" -version = "1.2.0" +version = "1.3.0" dependencies = [ "ahash 0.8.12", "api", @@ -13761,7 +13761,7 @@ dependencies = [ [[package]] name = "sql" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "arrow-buffer 58.3.0", @@ -13822,7 +13822,7 @@ dependencies = [ [[package]] name = "sqlness-runner" -version = "1.2.0" +version = "1.3.0" dependencies = [ "async-trait", "clap", @@ -14104,7 +14104,7 @@ dependencies = [ [[package]] name = "standalone" -version = "1.2.0" +version = "1.3.0" dependencies = [ "async-trait", "catalog", @@ -14150,7 +14150,7 @@ checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" [[package]] name = "store-api" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "aquamarine", @@ -14343,7 +14343,7 @@ dependencies = [ [[package]] name = "substrait" -version = "1.2.0" +version = "1.3.0" dependencies = [ "async-trait", "bytes", @@ -14465,7 +14465,7 @@ dependencies = [ [[package]] name = "table" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "arc-swap", @@ -14748,7 +14748,7 @@ checksum = "8f50febec83f5ee1df3015341d8bd429f2d1cc62bcba7ea2076759d315084683" [[package]] name = "tests-fuzz" -version = "1.2.0" +version = "1.3.0" dependencies = [ "arbitrary", "async-trait", @@ -14796,7 +14796,7 @@ dependencies = [ [[package]] name = "tests-integration" -version = "1.2.0" +version = "1.3.0" dependencies = [ "api", "arrow-flight", @@ -14878,7 +14878,7 @@ dependencies = [ "sqlx", "standalone", "store-api", - "substrait 1.2.0", + "substrait 1.3.0", "table", "tempfile", "time", diff --git a/Cargo.toml b/Cargo.toml index a9d509031cd..9582b12bd85 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -75,7 +75,7 @@ members = [ resolver = "2" [workspace.package] -version = "1.2.0" +version = "1.3.0" edition = "2024" license = "Apache-2.0" From d3b10d13a5712868d72969b4a453ae786f0d60e3 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 6 Aug 2026 16:45:19 +0800 Subject: [PATCH 006/126] fix(mito2): fail open when Bloom IN predicate has non-literal or unencodable members (#8709) * fix(mito2): fail open when Bloom IN predicate has non-literal or unencodable members collect_in_list previously filtered out non-literal and encoding-failed IN members and could build a partial hard-pruning predicate, so scanning with the Bloom filter enabled could prune rows that actually match the query (false-negative results). Now any non-literal member or any encoding failure disables Bloom pruning for the whole IN expression (fail open), while independent AND subpredicates such as `col = 42` are still extracted as before. Encoding errors are logged and ignored, so queries never fail. Adds builder unit tests covering pure-literal, pure-nonliteral, mixed literal+null, all-null, mixed-nonliteral-with-AND, and encoding-failure cases, plus an engine-level witness test asserting Bloom-on and Bloom-off scans return identical rows for a mixed IN filter. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(mito2): rename bloom filter tests to semantic names without bug-id prefixes Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/engine/index_build_test.rs | 215 +++++++++++++++++- .../sst/index/bloom_filter/applier/builder.rs | 147 ++++++++++-- 2 files changed, 339 insertions(+), 23 deletions(-) diff --git a/src/mito2/src/engine/index_build_test.rs b/src/mito2/src/engine/index_build_test.rs index 90815711649..ce950f8a246 100644 --- a/src/mito2/src/engine/index_build_test.rs +++ b/src/mito2/src/engine/index_build_test.rs @@ -24,6 +24,8 @@ use api::v1::{Rows, SemanticType}; use async_trait::async_trait; use common_base::readable_size::ReadableSize; use common_recordbatch::RecordBatches; +use datafusion_expr::expr::InList; +use datafusion_expr::{Expr, col, lit}; use datatypes::arrow::array::AsArray; use datatypes::arrow::datatypes::TimestampMillisecondType; use datatypes::prelude::ConcreteDataType; @@ -38,14 +40,16 @@ use store_api::storage::{RegionId, ScanRequest}; use tokio::sync::{Notify, Semaphore}; use crate::cache::file_cache::{FileType, IndexKey}; +use crate::cache::{CacheManager, CacheStrategy}; use crate::config::{IndexBuildMode, MitoConfig, Mode}; use crate::engine::MitoEngine; use crate::engine::compaction_test::put_and_flush; use crate::engine::listener::{EventListener, GateIndexBuildListener, IndexBuildListener}; use crate::manifest::action::RegionEdit; -use crate::read::scan_region::Scanner; +use crate::read::scan_region::{ScanRegion, Scanner}; use crate::sst::file::{FileMeta, RegionFileId, RegionIndexId}; use crate::sst::location; +use crate::test_util::batch_util::sort_batches_and_print; use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, reopen_region, rows_schema, }; @@ -98,6 +102,215 @@ fn assert_listener_counts( assert_eq!(listener.finish_count(), expected_success_count); } +fn bloom_filter_test_cache() -> Arc { + const CACHE_SIZE: u64 = 64 * 1024; + + Arc::new( + CacheManager::builder() + .index_metadata_size(CACHE_SIZE) + .index_content_size(CACHE_SIZE) + .index_content_page_size(CACHE_SIZE) + .puffin_metadata_size(CACHE_SIZE) + .build(), + ) +} + +async fn scan_with_bloom_filter( + engine: &MitoEngine, + region_id: RegionId, + request: ScanRequest, + cache_manager: Arc, + ignore_bloom_filter: bool, +) -> RecordBatches { + let region = engine.get_region(region_id).unwrap(); + let scanner = ScanRegion::new( + region.version(), + region.access_layer.clone(), + request, + CacheStrategy::EnableAll(cache_manager), + ) + .with_ignore_bloom_filter(ignore_bloom_filter) + .scanner() + .await + .unwrap(); + + RecordBatches::try_collect(scanner.scan().await.unwrap()) + .await + .unwrap() +} + +#[tokio::test] +async fn bloom_mixed_in_index_on_off_equivalent() { + let mut env = TestEnv::with_prefix("bloom_mixed_in_index_on_off_equivalent_").await; + let listener = Arc::new(IndexBuildListener::default()); + let engine = env + .create_engine_with( + async_build_mode_config(true), + None, + Some(listener.clone()), + None, + ) + .await; + let region_id = RegionId::new(1, 1); + + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let mut request = CreateRequestBuilder::new().build(); + let bloom_options = + SkippingIndexOptions::new_unchecked(1, 0.0001, SkippingIndexType::BloomFilter); + request.column_metadatas[0] + .column_schema + .set_skipping_options(&bloom_options) + .unwrap(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(0, 2), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + listener.wait_finish(1).await; + + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(scanner.num_files(), 1); + assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); + assert!( + engine + .all_index_metas() + .await + .iter() + .any(|meta| meta.index_type == "bloom_filter") + ); + + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(2, 4), + }, + ) + .await; + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(scanner.num_files(), 1); + assert_eq!(scanner.num_memtables(), 1); + + let mixed_in = Expr::InList(InList { + expr: Box::new(col("tag_0")), + list: vec![ + lit("definitely_absent_0"), + lit("definitely_absent_1"), + lit("definitely_absent_2"), + lit("definitely_absent_3"), + col("tag_0"), + ], + negated: false, + }); + let filtered_request = ScanRequest { + filters: vec![mixed_in], + ..Default::default() + }; + let cache_manager = bloom_filter_test_cache(); + + let indexed = scan_with_bloom_filter( + &engine, + region_id, + filtered_request.clone(), + cache_manager.clone(), + false, + ) + .await; + let unindexed = scan_with_bloom_filter( + &engine, + region_id, + filtered_request.clone(), + cache_manager, + true, + ) + .await; + let memtable_only = RecordBatches::try_collect( + engine + .scan_to_stream( + region_id, + ScanRequest { + skip_sst_files: true, + ..filtered_request.clone() + }, + ) + .await + .unwrap(), + ) + .await + .unwrap(); + let full_scan = RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap(); + + let indexed_rows = sort_batches_and_print(&indexed, &["ts"]); + let unindexed_rows = sort_batches_and_print(&unindexed, &["ts"]); + let memtable_rows = sort_batches_and_print(&memtable_only, &["ts"]); + let full_rows = sort_batches_and_print(&full_scan, &["ts"]); + assert_eq!( + memtable_rows, + "\ ++-------+---------+---------------------+ +| tag_0 | field_0 | ts | ++-------+---------+---------------------+ +| 2 | 2.0 | 1970-01-01T00:00:02 | +| 3 | 3.0 | 1970-01-01T00:00:03 | ++-------+---------+---------------------+" + ); + assert_eq!( + full_rows, + "\ ++-------+---------+---------------------+ +| tag_0 | field_0 | ts | ++-------+---------+---------------------+ +| 0 | 0.0 | 1970-01-01T00:00:00 | +| 1 | 1.0 | 1970-01-01T00:00:01 | +| 2 | 2.0 | 1970-01-01T00:00:02 | +| 3 | 3.0 | 1970-01-01T00:00:03 | ++-------+---------+---------------------+" + ); + assert_eq!( + unindexed_rows, full_rows, + "disabling the Bloom filter must retain all matching SST rows" + ); + assert_eq!( + indexed_rows, unindexed_rows, + "Bloom filtering must not prune rows matched by a nonliteral IN member" + ); +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum IndexPublicationPhase { BeforeManifestCommit, diff --git a/src/mito2/src/sst/index/bloom_filter/applier/builder.rs b/src/mito2/src/sst/index/bloom_filter/applier/builder.rs index beb26d1bf7c..e5968b5abb2 100644 --- a/src/mito2/src/sst/index/bloom_filter/applier/builder.rs +++ b/src/mito2/src/sst/index/bloom_filter/applier/builder.rs @@ -204,21 +204,13 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { return Ok(()); }; - // Convert all non-null literals to predicates - let predicates = in_list - .list - .iter() - .filter_map(Self::nonnull_lit) - .map(|lit| encode_lit(lit, data_type.clone())); - - // Collect successful conversions let mut valid_predicates = BTreeSet::new(); - for predicate in predicates { - match predicate { - Ok(p) => { - valid_predicates.insert(p); - } - Err(e) => warn!(e; "Failed to convert value in InList"), + for expr in &in_list.list { + let Expr::Literal(lit, _) = expr else { + return Ok(()); + }; + if !lit.is_null() { + valid_predicates.insert(encode_lit(lit, data_type.clone())?); } } @@ -311,14 +303,6 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { Ok(false) } - /// Helper function to get non-null literal value - fn nonnull_lit(expr: &Expr) -> Option<&ScalarValue> { - match expr { - Expr::Literal(lit, _) if !lit.is_null() => Some(lit), - _ => None, - } - } - /// Helper function to get the column and literal value from an equality expr (column = lit) fn eq_expr_col_lit<'b>( left: &'b Expr, @@ -435,6 +419,125 @@ mod tests { i.lit() } + fn build_bloom_predicates(exprs: &[Expr]) -> Option>> { + let (_d, factory) = PuffinManagerFactory::new_for_test_block("bloom_builder_"); + let metadata = test_region_metadata(); + BloomFilterIndexApplierBuilder::new( + "test".to_string(), + PathType::Bare, + test_object_store(), + &metadata, + factory, + ) + .build(exprs) + .unwrap() + .map(|applier| (*applier.default_predicates).clone()) + } + + fn int64_inlist_predicate(values: impl IntoIterator) -> InListPredicate { + InListPredicate { + list: values + .into_iter() + .map(|value| { + encode_lit( + &ScalarValue::Int64(Some(value)), + ConcreteDataType::int64_datatype(), + ) + .unwrap() + }) + .collect(), + } + } + + #[test] + fn bloom_pure_literal_in_extracts_exact_predicate() { + let expr = Expr::InList(InList { + expr: Box::new(column("column2")), + list: vec![int64_lit(1), int64_lit(2), int64_lit(3)], + negated: false, + }); + + assert_eq!( + build_bloom_predicates(&[expr]), + Some(BTreeMap::from([( + 2, + vec![int64_inlist_predicate([1, 2, 3])] + )])) + ); + } + + #[test] + fn bloom_pure_nonliteral_in_does_not_extract_predicate() { + let expr = Expr::InList(InList { + expr: Box::new(column("column1")), + list: vec![column("column1")], + negated: false, + }); + + assert_eq!(build_bloom_predicates(&[expr]), None); + } + + #[test] + fn bloom_mixed_literal_null_in_extracts_exact_predicate() { + let expr = Expr::InList(InList { + expr: Box::new(column("column2")), + list: vec![ + int64_lit(1), + Expr::Literal(ScalarValue::Int64(None), None), + int64_lit(3), + ], + negated: false, + }); + + assert_eq!( + build_bloom_predicates(&[expr]), + Some(BTreeMap::from([(2, vec![int64_inlist_predicate([1, 3])])])) + ); + } + + #[test] + fn bloom_all_literal_null_in_does_not_extract_predicate() { + let expr = Expr::InList(InList { + expr: Box::new(column("column2")), + list: vec![ + Expr::Literal(ScalarValue::Int64(None), None), + Expr::Literal(ScalarValue::Int64(None), None), + ], + negated: false, + }); + + assert_eq!(build_bloom_predicates(&[expr]), None); + } + + #[test] + fn bloom_mixed_nonliteral_in_keeps_only_independent_predicate() { + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::InList(InList { + expr: Box::new(column("column1")), + list: vec!["definitely_absent".lit(), column("column1")], + negated: false, + })), + op: Operator::And, + right: Box::new(column("column2").eq(int64_lit(42))), + }); + + assert_eq!( + build_bloom_predicates(&[expr]), + Some(BTreeMap::from([(2, vec![int64_inlist_predicate([42])])])) + ); + } + + #[test] + fn bloom_encoding_failure_in_does_not_extract_predicate() { + let expr = Expr::InList(InList { + expr: Box::new(column("column2")), + list: vec![int64_lit(1), "not_an_int64".lit()], + negated: false, + }); + + assert_eq!(build_bloom_predicates(&[expr]), None); + } + #[test] fn test_build_with_in_list() { let (_d, factory) = PuffinManagerFactory::new_for_test_block("test_build_with_in_list_"); From 1aadfa565c32332146e9628dd82455f0f59ce3da Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Thu, 6 Aug 2026 20:28:50 +0800 Subject: [PATCH 007/126] fix: cache physical table metadata lookups (#8777) fix: physical table dont cache Signed-off-by: shuiyisong --- src/catalog/src/kvbackend/manager.rs | 115 ++++++++++++++++++++++++--- 1 file changed, 106 insertions(+), 9 deletions(-) diff --git a/src/catalog/src/kvbackend/manager.rs b/src/catalog/src/kvbackend/manager.rs index dc86c703d88..a4d8214a55b 100644 --- a/src/catalog/src/kvbackend/manager.rs +++ b/src/catalog/src/kvbackend/manager.rs @@ -29,7 +29,7 @@ use common_meta::key::TableMetadataManagerRef; use common_meta::key::catalog_name::CatalogNameKey; use common_meta::key::flow::FlowMetadataManager; use common_meta::key::schema_name::SchemaNameKey; -use common_meta::key::table_info::{TableInfoManager, TableInfoValue}; +use common_meta::key::table_info::TableInfoValue; use common_meta::key::table_name::TableNameKey; use common_meta::kv_backend::KvBackendRef; use common_procedure::ProcedureManagerRef; @@ -112,7 +112,7 @@ impl KvBackendCatalogManager { // Override logical table's partition key indices with physical table's. async fn override_logical_table_partition_key_indices( table_route_cache: &TableRouteCacheRef, - table_info_manager: &TableInfoManager, + table_info_cache: &TableInfoCacheRef, table: TableRef, ) -> Result { // If the table is not a metric table, return the table directly. @@ -125,7 +125,7 @@ impl KvBackendCatalogManager { .await .context(TableMetadataManagerSnafu)? && let TableRoute::Logical(logical_route) = &*table_route_value - && let Some(physical_table_info_value) = table_info_manager + && let Some(physical_table_info) = table_info_cache .get(logical_route.physical_table_id()) .await .context(TableMetadataManagerSnafu)? @@ -135,15 +135,13 @@ impl KvBackendCatalogManager { let mut phy_part_cols_not_in_logical_table = vec![]; // Remap partition key indices from physical table to logical table - new_table_info.meta.partition_key_indices = physical_table_info_value - .table_info + new_table_info.meta.partition_key_indices = physical_table_info .meta .partition_key_indices .iter() .filter_map(|&physical_index| { // Get the column name from the physical table using the physical index - physical_table_info_value - .table_info + physical_table_info .meta .schema .column_schemas() @@ -321,9 +319,13 @@ impl CatalogManager for KvBackendCatalogManager { self.cache_registry.get().context(CacheNotFoundSnafu { name: "table_route_cache", })?; + let table_info_cache: TableInfoCacheRef = + self.cache_registry.get().context(CacheNotFoundSnafu { + name: "table_info_cache", + })?; return Self::override_logical_table_partition_key_indices( &table_route_cache, - self.table_metadata_manager.table_info_manager(), + &table_info_cache, table, ) .await @@ -453,6 +455,10 @@ impl CatalogManager for KvBackendCatalogManager { self.cache_registry.get().context(CacheNotFoundSnafu { name: "table_route_cache", }); + let table_info_cache: Result = + self.cache_registry.get().context(CacheNotFoundSnafu { + name: "table_info_cache", + }); common_runtime::spawn_global(async move { let table_route_cache = match table_route_cache { @@ -462,6 +468,13 @@ impl CatalogManager for KvBackendCatalogManager { return; } }; + let table_info_cache = match table_info_cache { + Ok(table_info_cache) => table_info_cache, + Err(e) => { + let _ = tx.send(Err(e)).await; + return; + } + }; let table_id_stream = metadata_manager .table_name_manager() @@ -490,6 +503,7 @@ impl CatalogManager for KvBackendCatalogManager { let tx = tx.clone(); let semaphore = semaphore.clone(); let table_route_cache = table_route_cache.clone(); + let table_info_cache = table_info_cache.clone(); common_runtime::spawn_global(async move { // we don't explicitly close the semaphore so just ignore the potential error. let _ = semaphore.acquire().await; @@ -509,7 +523,7 @@ impl CatalogManager for KvBackendCatalogManager { for table in table_info_values.into_values().map(build_table) { let table = Self::override_logical_table_partition_key_indices( &table_route_cache, - metadata_manager.table_info_manager(), + &table_info_cache, table, ) .await; @@ -645,3 +659,86 @@ impl SystemCatalog { } } } + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_meta::cache::{CacheContainer, Initializer}; + use common_meta::key::table_route::LogicalTableRouteValue; + use moka::future::Cache as FutureCache; + use table::metadata::TableInfo; + + use super::*; + + fn metric_table_info( + table_id: TableId, + table_name: &str, + partition_key_indices: Vec, + ) -> TableInfo { + let mut table_info = + common_meta::ddl::test_util::create_table::test_create_table_task(table_name, table_id) + .table_info; + table_info.meta.engine = METRIC_ENGINE_NAME.to_string(); + table_info.meta.partition_key_indices = partition_key_indices; + table_info + } + + #[tokio::test] + async fn test_logical_table_uses_cached_physical_table_info() { + const LOGICAL_TABLE_ID: TableId = 1; + const PHYSICAL_TABLE_ID: TableId = 2; + + let route = Arc::new(TableRoute::Logical(Arc::new(LogicalTableRouteValue::new( + PHYSICAL_TABLE_ID, + )))); + let route_initializer: Initializer> = Arc::new(move |_| { + let route = route.clone(); + Box::pin(async move { Ok(Some(route)) }) + }); + let table_route_cache = Arc::new(CacheContainer::new( + "test_table_route_cache".to_string(), + FutureCache::new(16), + Box::new(|_, _| Box::pin(async { Ok(()) })), + route_initializer, + |_| true, + )); + + let load_count = Arc::new(AtomicUsize::new(0)); + let physical_table_info = + Arc::new(metric_table_info(PHYSICAL_TABLE_ID, "physical", vec![1])); + let table_info_initializer: Initializer> = { + let load_count = load_count.clone(); + let table_info = physical_table_info.clone(); + Arc::new(move |_| { + load_count.fetch_add(1, Ordering::Relaxed); + let table_info = table_info.clone(); + Box::pin(async move { Ok(Some(table_info)) }) + }) + }; + let table_info_cache = Arc::new(CacheContainer::new( + "test_table_info_cache".to_string(), + FutureCache::new(16), + Box::new(|_, _| Box::pin(async { Ok(()) })), + table_info_initializer, + |_| true, + )); + + let logical_table = DistTable::table(Arc::new(metric_table_info( + LOGICAL_TABLE_ID, + "logical", + vec![], + ))); + for _ in 0..2 { + let table = KvBackendCatalogManager::override_logical_table_partition_key_indices( + &table_route_cache, + &table_info_cache, + logical_table.clone(), + ) + .await + .unwrap(); + assert_eq!(table.table_info().meta.partition_key_indices, vec![1]); + } + assert_eq!(load_count.load(Ordering::Relaxed), 1); + } +} From 1693b2727c8841db5e94c484ea7b31cc41df3b19 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 6 Aug 2026 21:01:32 +0800 Subject: [PATCH 008/126] refactor: port query regression runner to Rust (#8651) * refactor: port query regression runner to Rust Signed-off-by: discord9 * ci: remove optional OTLP report plotter Signed-off-by: discord9 * refactor: split query regression runner into modules Signed-off-by: discord9 * style: use crate-qualified imports in query regression runner Signed-off-by: discord9 * refactor: simplify query regression runner internals Signed-off-by: discord9 * feat: abstract inspect-footer storage access behind object store destination Add an optional --destination to inspect-footer (and --base-destination/--candidate-destination to finalize-remote) so the storage inspection reads DB data files through the opendal-backed object_store abstraction instead of bare std::fs. Local paths keep working unchanged via the --root shortcut (File backend); remote backends (S3/GCS/...) are described by a DestinationConfig TOML reusing the object-store crate's ObjectStoreConfig serde shape. - inspect_footer: list via ObjectStore::list + ObjectMeta filtering (parquet keys, non-zero size, metadata/ segment), read footers async via ParquetObjectReader + ParquetMetaDataReader with known file size (no extra HEAD); output JSON schema unchanged - finalize-remote: --base-data-home/--candidate-data-home become optional, mutually exclusive with the new --*-destination args - cmd deps: add object_store_opendal + datafusion_object_store - tests: fs-backend list+footer integration tests (metadata filtering, destination TOML mode, root/destination exclusivity) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * style: drop needless borrow in inspect footer test Fix clippy::needless_borrows_for_generic_args in the inspect-footer test (fs::create_dir_all(table.join("metadata"))). Missed by the earlier focused clippy run because it only covered --bin targets. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .github/scripts/query-regression-run.py | 527 +++++- .github/workflows/query-regression.yml | 10 +- Cargo.lock | 2 + src/cmd/Cargo.toml | 7 + .../bin/query_perf_fixture/inspect_footer.rs | 398 +++- src/cmd/src/bin/query_perf_fixture/mod.rs | 6 +- src/cmd/src/bin/query_regression_runner.rs | 23 + .../src/bin/query_regression_runner/direct.rs | 501 +++++ .../query_regression_runner/materialize.rs | 225 +++ .../bin/query_regression_runner/measure.rs | 423 +++++ .../src/bin/query_regression_runner/mod.rs | 204 ++ .../src/bin/query_regression_runner/model.rs | 235 +++ .../src/bin/query_regression_runner/otlp.rs | 618 ++++++ .../src/bin/query_regression_runner/plan.rs | 122 ++ .../query_regression_runner/remote/ingest.rs | 505 +++++ .../bin/query_regression_runner/remote/mod.rs | 183 ++ .../remote/read_bench.rs | 363 ++++ .../query_regression_runner/remote/storage.rs | 241 +++ .../src/bin/query_regression_runner/sql.rs | 238 +++ tests/perf/AGENTS.md | 15 +- tests/perf/README.md | 201 +- tests/perf/plot_otlp_trace_report.sh | 80 - tests/perf/query_regression_runner.py | 1663 ----------------- ...ery_regression_runner_compaction_toctou.py | 196 +- ...query_regression_runner_otlp_trace_load.py | 243 +-- 25 files changed, 4935 insertions(+), 2294 deletions(-) create mode 100644 src/cmd/src/bin/query_regression_runner.rs create mode 100644 src/cmd/src/bin/query_regression_runner/direct.rs create mode 100644 src/cmd/src/bin/query_regression_runner/materialize.rs create mode 100644 src/cmd/src/bin/query_regression_runner/measure.rs create mode 100644 src/cmd/src/bin/query_regression_runner/mod.rs create mode 100644 src/cmd/src/bin/query_regression_runner/model.rs create mode 100644 src/cmd/src/bin/query_regression_runner/otlp.rs create mode 100644 src/cmd/src/bin/query_regression_runner/plan.rs create mode 100644 src/cmd/src/bin/query_regression_runner/remote/ingest.rs create mode 100644 src/cmd/src/bin/query_regression_runner/remote/mod.rs create mode 100644 src/cmd/src/bin/query_regression_runner/remote/read_bench.rs create mode 100644 src/cmd/src/bin/query_regression_runner/remote/storage.rs create mode 100644 src/cmd/src/bin/query_regression_runner/sql.rs delete mode 100755 tests/perf/plot_otlp_trace_report.sh delete mode 100644 tests/perf/query_regression_runner.py diff --git a/.github/scripts/query-regression-run.py b/.github/scripts/query-regression-run.py index 3392a0658c9..d293aa3e0c9 100644 --- a/.github/scripts/query-regression-run.py +++ b/.github/scripts/query-regression-run.py @@ -18,10 +18,17 @@ from __future__ import annotations import argparse +import json import os import re +import signal +import socket import subprocess +import time +import urllib.request +from dataclasses import dataclass from pathlib import Path +from typing import Any DEFAULT_CASES = [ @@ -100,38 +107,479 @@ def append_step_summary(summary: Path) -> None: out.write(summary.read_text()) +@dataclass(frozen=True) +class RunTarget: + name: str + binary: Path + work_dir: Path + http_port: int + grpc_port: int + mysql_port: int + postgres_port: int + metasrv_rpc_port: int + metasrv_http_port: int + datanode_rpc_port: int + datanode_http_port: int + datanode_data_dir: Path + frontend_config: Path | None = None + + +def allocate_ports(n: int) -> list[int]: + socks = [] + try: + for _ in range(n): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("127.0.0.1", 0)) + socks.append(sock) + return [sock.getsockname()[1] for sock in socks] + finally: + for sock in socks: + sock.close() + + +def make_target( + name: str, + binary: Path, + root: Path, + ports: list[int], + frontend_config: Path | None = None, +) -> RunTarget: + work_dir = root / name + return RunTarget( + name=name, + binary=binary, + work_dir=work_dir, + http_port=ports[4], + grpc_port=ports[5], + mysql_port=ports[6], + postgres_port=ports[7], + metasrv_rpc_port=ports[0], + metasrv_http_port=ports[1], + datanode_rpc_port=ports[2], + datanode_http_port=ports[3], + datanode_data_dir=work_dir / "datanode-0" / "data", + frontend_config=frontend_config, + ) + + +def wait_health(port: int, timeout_s: float = 60.0) -> None: + deadline = time.monotonic() + timeout_s + last: Exception | None = None + while time.monotonic() < deadline: + try: + with urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=2) as response: + if response.status < 500: + return + except Exception as err: # noqa: BLE001 - retain the last health diagnostic + last = err + time.sleep(0.5) + raise TimeoutError(f"health check timed out on port {port}: {last}") + + +def component_log_dir(target: RunTarget, name: str) -> Path: + return target.work_dir / "logs" / ("datanode-0" if name == "datanode" else name) + + +def component_command(target: RunTarget, name: str) -> list[str]: + log_dir = component_log_dir(target, name) + if name == "metasrv": + return [ + str(target.binary), "metasrv", "start", + "--grpc-bind-addr", f"127.0.0.1:{target.metasrv_rpc_port}", + "--grpc-server-addr", f"127.0.0.1:{target.metasrv_rpc_port}", + "--http-addr", f"127.0.0.1:{target.metasrv_http_port}", + "--backend", "memory-store", "--enable-region-failover", "false", + "--log-dir", str(log_dir), + ] + if name == "datanode": + return [ + str(target.binary), "datanode", "start", + "--grpc-bind-addr", f"127.0.0.1:{target.datanode_rpc_port}", + "--grpc-server-addr", f"127.0.0.1:{target.datanode_rpc_port}", + "--http-addr", f"127.0.0.1:{target.datanode_http_port}", + "--data-home", str(target.datanode_data_dir), "--log-dir", str(log_dir), + "--node-id", "0", "--metasrv-addrs", f"127.0.0.1:{target.metasrv_rpc_port}", + ] + if name == "frontend": + command = [str(target.binary), "frontend", "start"] + if target.frontend_config is not None: + command.extend(["--config-file", str(target.frontend_config)]) + command.extend([ + "--metasrv-addrs", f"127.0.0.1:{target.metasrv_rpc_port}", + "--http-addr", f"127.0.0.1:{target.http_port}", + "--grpc-bind-addr", f"127.0.0.1:{target.grpc_port}", + "--grpc-server-addr", f"127.0.0.1:{target.grpc_port}", + "--mysql-addr", f"127.0.0.1:{target.mysql_port}", + "--postgres-addr", f"127.0.0.1:{target.postgres_port}", + "--log-dir", str(log_dir), + ]) + return command + raise ValueError(f"unknown component: {name}") + + +def start_component( + target: RunTarget, + name: str, + procs: dict[tuple[str, str], subprocess.Popen[bytes]], +) -> None: + if name != "metasrv": + metasrv = procs.get((target.name, "metasrv")) + if metasrv is None or metasrv.poll() is not None: + raise RuntimeError("metasrv exited; memory-store metadata is no longer valid") + if name == "datanode": + target.datanode_data_dir.mkdir(parents=True, exist_ok=True) + logs = component_log_dir(target, name) + logs.mkdir(parents=True, exist_ok=True) + with (logs / "stdout.log").open("ab") as out, (logs / "stderr.log").open("ab") as err: + procs[(target.name, name)] = subprocess.Popen( + component_command(target, name), + stdout=out, + stderr=err, + start_new_session=True, + ) + wait_health({ + "metasrv": target.metasrv_http_port, + "datanode": target.datanode_http_port, + "frontend": target.http_port, + }[name]) + + +def stop_component( + target: RunTarget, + name: str, + procs: dict[tuple[str, str], subprocess.Popen[bytes]], +) -> None: + proc = procs.pop((target.name, name), None) + if proc is None or proc.poll() is not None: + return + try: + os.killpg(proc.pid, signal.SIGTERM) + except ProcessLookupError: + return + try: + proc.wait(timeout=20) + except subprocess.TimeoutExpired: + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + return + proc.wait(timeout=20) + + +def restart_component( + target: RunTarget, + name: str, + procs: dict[tuple[str, str], subprocess.Popen[bytes]], +) -> None: + stop_component(target, name, procs) + start_component(target, name, procs) + + +def stop_all(targets: list[RunTarget], procs: dict[tuple[str, str], subprocess.Popen[bytes]]) -> None: + for target in reversed(targets): + for name in ("frontend", "datanode", "metasrv"): + stop_component(target, name, procs) + + +def load_plan(fixture_generator: Path, case_path: Path) -> dict[str, Any]: + result = subprocess.run( + [str(fixture_generator), "plan", "--case", str(case_path)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=False, + ) + if result.returncode != 0: + raise RuntimeError(f"query_perf_fixture plan failed: {result.stderr[:2000]}") + return json.loads(result.stdout) + + +def run_direct_case( + args: argparse.Namespace, + case_path: Path, + work_dir: Path, + base_bin: Path, + candidate_bin: Path, + fixture_generator: Path, + runner: Path, +) -> int: + ports = allocate_ports(16) + targets = [ + make_target("base", base_bin, work_dir, ports[:8]), + make_target("candidate", candidate_bin, work_dir, ports[8:]), + ] + for target in targets: + if target.work_dir.exists() and any(target.work_dir.iterdir()): + raise RuntimeError(f"target work_dir exists and is non-empty: {target.work_dir}") + target.work_dir.mkdir(parents=True, exist_ok=True) + + procs: dict[tuple[str, str], subprocess.Popen[bytes]] = {} + try: + for target in targets: + start_component(target, "metasrv", procs) + start_component(target, "datanode", procs) + start_component(target, "frontend", procs) + + prepare = [ + str(runner), "prepare-direct", "--case", str(case_path), + "--fixture-generator", str(fixture_generator), + "--base-http-port", str(targets[0].http_port), + "--candidate-http-port", str(targets[1].http_port), + "--fixture-dir", str(work_dir / "fixture"), + "--output", str(work_dir / "prepare-direct.json"), + "--http-timeout", str(args.http_timeout), + ] + if parse_bool(args.allow_large_fixture): + prepare.append("--allow-large-fixture") + prepare_status = subprocess.run(prepare, check=False).returncode + if prepare_status != 0: + return prepare_status + + for target in targets: + stop_component(target, "datanode", procs) + + prepared = json.loads((work_dir / "prepare-direct.json").read_text(encoding="utf-8")) + fixtures = prepared.get("fixtures") + if not isinstance(fixtures, list): + raise RuntimeError("prepare-direct report has no fixtures array") + for target in targets: + destination = target.work_dir / "materialize-destination.toml" + destination.write_text( + f"data_home = {json.dumps(str(target.datanode_data_dir))}\n" + "object_store = { type = \"File\" }\n", + encoding="utf-8", + ) + for fixture in fixtures: + fixture_dir = fixture.get("fixture_dir") if isinstance(fixture, dict) else None + if not isinstance(fixture_dir, str): + raise RuntimeError("prepare-direct fixture record has no fixture_dir") + materialize = [ + str(runner), "materialize", "--fixture-dir", fixture_dir, + "--destination", str(destination), + ] + materialize_status = subprocess.run(materialize, check=False).returncode + if materialize_status != 0: + return materialize_status + + for target in targets: + start_component(target, "datanode", procs) + + measure = [ + str(runner), "measure", "--case", str(case_path), + "--fixture-generator", str(fixture_generator), + "--base-http-port", str(targets[0].http_port), + "--candidate-http-port", str(targets[1].http_port), + "--output", str(work_dir / "query-regression-report.json"), + "--http-timeout", str(args.http_timeout), + ] + status = subprocess.run(measure, check=False).returncode + if status != 0: + for target in targets: + restart_component(target, "frontend", procs) + status = subprocess.run(measure, check=False).returncode + return status + finally: + stop_all(targets, procs) + + +def run_remote_case( + args: argparse.Namespace, + case_path: Path, + work_dir: Path, + base_bin: Path, + candidate_bin: Path, + fixture_generator: Path, + runner: Path, +) -> int: + ports = allocate_ports(16) + targets = [ + make_target( + "base", + base_bin, + work_dir, + ports[:8], + work_dir / "base" / "frontend-prom-store.toml", + ), + make_target( + "candidate", + candidate_bin, + work_dir, + ports[8:], + work_dir / "candidate" / "frontend-prom-store.toml", + ), + ] + for target in targets: + if target.work_dir.exists() and any(target.work_dir.iterdir()): + raise RuntimeError(f"target work_dir exists and is non-empty: {target.work_dir}") + target.work_dir.mkdir(parents=True, exist_ok=True) + if target.frontend_config is None: + raise RuntimeError(f"remote target has no frontend config path: {target.name}") + render = [ + str(runner), "render-remote-config", "--case", str(case_path), + "--fixture-generator", str(fixture_generator), + "--output", str(target.frontend_config), + ] + render_status = subprocess.run(render, check=False).returncode + if render_status != 0: + return render_status + + procs: dict[tuple[str, str], subprocess.Popen[bytes]] = {} + report = work_dir / "query-regression-report.json" + try: + for target in targets: + start_component(target, "metasrv", procs) + start_component(target, "datanode", procs) + start_component(target, "frontend", procs) + + prepare = [ + str(runner), "prepare-remote", "--case", str(case_path), + "--fixture-generator", str(fixture_generator), + "--base-http-port", str(targets[0].http_port), + "--candidate-http-port", str(targets[1].http_port), + "--output", str(work_dir / "prepare-remote.json"), + "--http-timeout", str(args.http_timeout), + ] + prepare_status = subprocess.run(prepare, check=False).returncode + if prepare_status != 0: + return prepare_status + + measure = [ + str(runner), "measure", "--case", str(case_path), + "--fixture-generator", str(fixture_generator), + "--base-http-port", str(targets[0].http_port), + "--candidate-http-port", str(targets[1].http_port), + "--output", str(report), "--http-timeout", str(args.http_timeout), + ] + measure_status = subprocess.run(measure, check=False).returncode + if measure_status != 0 and not report.exists(): + return measure_status + + for target in targets: + stop_component(target, "datanode", procs) + + finalize = [ + str(runner), "finalize-remote", "--case", str(case_path), + "--fixture-generator", str(fixture_generator), + "--candidate-bin", str(candidate_bin), + "--base-data-home", str(targets[0].datanode_data_dir), + "--candidate-data-home", str(targets[1].datanode_data_dir), + "--report", str(report), + ] + finalize_status = subprocess.run(finalize, check=False).returncode + if finalize_status != 0: + return finalize_status + final_report = json.loads(report.read_text(encoding="utf-8")) + return 1 if measure_status != 0 or final_report.get("status") == "failed" else 0 + finally: + stop_all(targets, procs) + + +def run_otlp_case( + args: argparse.Namespace, + case_path: Path, + work_dir: Path, + base_bin: Path, + candidate_bin: Path, + fixture_generator: Path, + runner: Path, +) -> int: + if args.otelgen_bin is None: + raise ValueError("--otelgen-bin is required for otlp_trace_load") + ports = allocate_ports(16) + targets = [ + make_target("base", base_bin, work_dir, ports[:8]), + make_target("candidate", candidate_bin, work_dir, ports[8:]), + ] + for target in targets: + if target.work_dir.exists() and any(target.work_dir.iterdir()): + raise RuntimeError(f"target work_dir exists and is non-empty: {target.work_dir}") + target.work_dir.mkdir(parents=True, exist_ok=True) + + procs: dict[tuple[str, str], subprocess.Popen[bytes]] = {} + try: + for target in targets: + try: + start_component(target, "metasrv", procs) + start_component(target, "datanode", procs) + start_component(target, "frontend", procs) + command = [ + str(runner), "run-otlp-target", "--case", str(case_path), + "--fixture-generator", str(fixture_generator), + "--otelgen-bin", str(args.otelgen_bin), + "--http-port", str(target.http_port), + "--target-name", target.name, + "--work-dir", str(target.work_dir), + "--output", str(target.work_dir / "report.json"), + "--http-timeout", str(args.http_timeout), + ] + status = subprocess.run(command, check=False).returncode + finally: + stop_all([target], procs) + if status != 0: + return status + + output = work_dir / "query-regression-report.json" + finalize = [ + str(runner), "finalize-otlp", "--case", str(case_path), + "--fixture-generator", str(fixture_generator), + "--base-result", str(targets[0].work_dir / "report.json"), + "--candidate-result", str(targets[1].work_dir / "report.json"), + "--output", str(output), + ] + status = subprocess.run(finalize, check=False).returncode + if status != 0: + return status + report = json.loads(output.read_text(encoding="utf-8")) + return 1 if report.get("status") == "failed" else 0 + finally: + stop_all(targets, procs) + + def run_case(args: argparse.Namespace, case_path: Path, work_dir: Path) -> int: target_dir = profile_dir(args.cargo_profile) base_bin = args.base_bin or args.base_src / "target" / target_dir / "greptime" candidate_bin = args.candidate_bin or args.candidate_src / "target" / target_dir / "greptime" fixture_generator = args.fixture_generator or args.candidate_src / "target" / target_dir / "query_perf_fixture" - cmd = [ - "uv", - "run", - "--no-project", - "python", - str(args.candidate_src / "tests/perf/query_regression_runner.py"), - "--case", - str(case_path), - "--base-bin", - str(base_bin), - "--candidate-bin", - str(candidate_bin), - "--fixture-generator", - str(fixture_generator), - "--work-dir", - str(work_dir), - "--http-timeout", - str(args.http_timeout), - ] - if parse_bool(args.allow_large_fixture): - cmd.append("--allow-large-fixture") - if args.otelgen_bin is not None: - cmd.extend(["--otelgen-bin", str(args.otelgen_bin)]) - + runner = args.runner or args.candidate_src / "target" / target_dir / "query_regression_runner" + plan = load_plan(fixture_generator, case_path) + scenario = plan.get("scenario") + kind = scenario.get("kind") if isinstance(scenario, dict) else None print(f"::group::Query regression case: {case_path}", flush=True) try: - return subprocess.run(cmd, check=False).returncode + if kind == "direct_readable_sst": + return run_direct_case( + args, + case_path, + work_dir, + base_bin, + candidate_bin, + fixture_generator, + runner, + ) + if kind == "prom_remote_write_then_query": + return run_remote_case( + args, + case_path, + work_dir, + base_bin, + candidate_bin, + fixture_generator, + runner, + ) + if kind == "otlp_trace_load": + return run_otlp_case( + args, + case_path, + work_dir, + base_bin, + candidate_bin, + fixture_generator, + runner, + ) + raise ValueError( + f"unsupported scenario kind {kind!r}; supported: " + "'direct_readable_sst', 'prom_remote_write_then_query', 'otlp_trace_load'" + ) finally: print("::endgroup::", flush=True) @@ -157,6 +605,27 @@ def write_summary(args: argparse.Namespace, reports: list[Path]) -> int: return subprocess.run(cmd, check=False).returncode +def write_failed_report(work_dir: Path, case_path: Path, error: Exception) -> None: + work_dir.mkdir(parents=True, exist_ok=True) + path = work_dir / "query-regression-report.json" + try: + report = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(report, dict) or not isinstance(report.get("targets"), list) or not isinstance(report.get("thresholds"), list): + raise ValueError("existing report is malformed") + except (OSError, ValueError, json.JSONDecodeError): + report = { + "case_path": str(case_path), + "targets": [], + "thresholds": [], + } + report["status"] = "failed" + report["error"] = repr(error) + path.write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--cases", action="append", help="'all', 'heavy', or comma/space separated case paths") @@ -170,6 +639,7 @@ def main() -> int: default=configured_path(os.environ.get("FIXTURE_GENERATOR")), ) parser.add_argument("--otelgen-bin", type=Path, default=configured_path(os.environ.get("OTELGEN_BIN"))) + parser.add_argument("--runner", type=Path, default=configured_path(os.environ.get("QUERY_REGRESSION_RUNNER"))) parser.add_argument("--cargo-profile", default=os.environ.get("CARGO_PROFILE", "nightly")) parser.add_argument("--work-dir", default=Path("query-regression-work"), type=Path) parser.add_argument("--http-timeout", default=os.environ.get("HTTP_TIMEOUT", "300")) @@ -199,7 +669,12 @@ def main() -> int: case_path = resolve_case_path(args.candidate_src, case) work_dir = args.work_dir / case_slug(case_path) reports.append(work_dir / "query-regression-report.json") - case_status = run_case(args, case_path, work_dir) + try: + case_status = run_case(args, case_path, work_dir) + except Exception as err: # noqa: BLE001 - preserve an aggregate report for the summary step + print(f"error: query regression case {case_path}: {err}", flush=True) + write_failed_report(work_dir, case_path, err) + case_status = 1 if case_status != 0: status = case_status or 1 diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index fd15e659604..0806d4d3e38 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -575,16 +575,17 @@ jobs: uv run --no-project python tests/perf/test_query_regression_summary_otlp.py uv run --no-project python tests/perf/test_query_regression_case_selection.py - - name: Build candidate greptime and fixture generators + - name: Build candidate greptime and query regression helpers working-directory: src run: | set -euo pipefail SECONDS=0 cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime cargo build --profile "${CARGO_PROFILE}" -p cmd --bin query_perf_fixture --features dev-tools + cargo build --profile "${CARGO_PROFILE}" -p cmd --bin query_regression_runner --features dev-tools candidate_build_elapsed="${SECONDS}" - printf 'Candidate greptime and fixture generator cargo builds elapsed: %s seconds\n' "${candidate_build_elapsed}" - printf -- '- Candidate greptime and fixture generator cargo builds: %s seconds\n' "${candidate_build_elapsed}" >> "${GITHUB_STEP_SUMMARY}" + printf 'Candidate greptime and query regression helper cargo builds elapsed: %s seconds\n' "${candidate_build_elapsed}" + printf -- '- Candidate greptime and query regression helper cargo builds: %s seconds\n' "${candidate_build_elapsed}" >> "${GITHUB_STEP_SUMMARY}" sccache --show-stats target_dir="${CARGO_PROFILE}" if [[ "${CARGO_PROFILE}" == "dev" ]]; then @@ -595,6 +596,8 @@ jobs: "${GITHUB_WORKSPACE}/query-regression-bins/candidate/greptime" cp "${CARGO_TARGET_DIR}/${target_dir}/query_perf_fixture" \ "${GITHUB_WORKSPACE}/query-regression-bins/candidate/query_perf_fixture" + cp "${CARGO_TARGET_DIR}/${target_dir}/query_regression_runner" \ + "${GITHUB_WORKSPACE}/query-regression-bins/candidate/query_regression_runner" - name: Run query regression id: run @@ -609,6 +612,7 @@ jobs: FIXTURE_GENERATOR: ${{ github.workspace }}/query-regression-bins/candidate/query_perf_fixture OTELGEN_BIN: /usr/local/bin/otelgen SUMMARY_SCRIPT: ${{ github.event_name == 'pull_request' && 'query-regression-trusted-scripts/query-regression-summary.py' || 'src/.github/scripts/query-regression-summary.py' }} + QUERY_REGRESSION_RUNNER: ${{ github.workspace }}/query-regression-bins/candidate/query_regression_runner run: >- uv run --no-project python src/.github/scripts/query-regression-run.py --base-src src diff --git a/Cargo.lock b/Cargo.lock index 77474ea66fd..6c610726ef8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2190,6 +2190,8 @@ dependencies = [ "mito2", "moka", "object-store", + "object_store", + "object_store_opendal", "parquet", "plugins", "pprof", diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 24edcd56338..10865b48a10 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -14,6 +14,11 @@ name = "query_perf_fixture" path = "src/bin/query_perf_fixture.rs" required-features = ["dev-tools"] +[[bin]] +name = "query_regression_runner" +path = "src/bin/query_regression_runner.rs" +required-features = ["dev-tools"] + [features] default = [ "servers/pprof", @@ -68,6 +73,7 @@ common-wal.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-physical-plan.workspace = true +datafusion_object_store.workspace = true datanode.workspace = true datatypes.workspace = true either = "1.15" @@ -87,6 +93,7 @@ mito-codec.workspace = true mito2.workspace = true moka = { workspace = true, features = ["future"] } object-store.workspace = true +object_store_opendal.workspace = true parquet = { workspace = true, features = ["object_store"] } plugins.workspace = true prometheus.workspace = true diff --git a/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs b/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs index a52005c4e30..4bc9845671c 100644 --- a/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs +++ b/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs @@ -13,17 +13,39 @@ // limitations under the License. use std::collections::BTreeSet; -use std::fs::{self, File}; +use std::fs; use std::path::{Path, PathBuf}; +use std::sync::Arc; use clap::Args as ClapArgs; -use parquet::file::reader::{FileReader, SerializedFileReader}; -use serde::Serialize; +use datafusion_object_store::path::Path as StorePath; +use datafusion_object_store::{ObjectMeta, ObjectStore}; +use futures::{StreamExt, TryStreamExt}; +use object_store::config::ObjectStoreConfig; +use object_store::factory::new_raw_object_store; +use object_store::services::Fs; +use parquet::arrow::async_reader::ParquetObjectReader; +use parquet::file::metadata::ParquetMetaDataReader; +use serde::{Deserialize, Serialize}; + +/// Same shape as `query_regression_runner::model::DestinationConfig`. The two +/// binaries are separate crate roots, so the serde-compatible struct is defined +/// here instead of being shared; it reuses the `object-store` crate's +/// `ObjectStoreConfig` deserialization. +#[derive(Debug, Deserialize)] +struct DestinationConfig { + data_home: String, + object_store: ObjectStoreConfig, +} #[derive(Debug, ClapArgs)] pub(super) struct InspectFooterArgs { + /// Local data home (convenience shortcut for a File destination). #[arg(long)] - root: PathBuf, + root: Option, + /// TOML file: data_home = "..." and object_store = { type = "File" | "S3" | ... }. + #[arg(long)] + destination: Option, #[arg(long, default_value = "greptime_value")] column: String, #[arg(long)] @@ -68,64 +90,140 @@ struct FooterColumnChunkReport { num_values: i64, } -pub(super) fn run_inspect_footer( +/// A parquet data file discovered by listing the store, with the metadata needed +/// to read its footer without any extra stat/head call. +#[derive(Debug)] +struct ListedFile { + location: StorePath, + size: u64, + path: String, + relative_path: String, +} + +pub(super) async fn run_inspect_footer( args: InspectFooterArgs, -) -> Result<(), Box> { - let files = collect_parquet_files(&args.root, args.include_metadata_files)?; - let reports = files - .iter() - .map(|p| inspect_file(&args.root, p, &args.column)) - .collect::, _>>()?; +) -> Result<(), Box> { + let (destination, root) = match (args.destination, args.root) { + (Some(_), Some(_)) => { + return Err("--destination and --root are mutually exclusive; pass exactly one".into()); + } + (None, None) => return Err("one of --destination or --root is required".into()), + (Some(path), None) => { + let destination: DestinationConfig = toml::from_str(&fs::read_to_string(path)?)?; + (Some(destination), None) + } + (None, root) => (None, root), + }; + let store = build_store(destination.as_ref(), root.as_deref()).await?; + let root_display = root + .as_deref() + .map(|root| root.display().to_string()) + .unwrap_or_else(|| { + destination + .as_ref() + .expect("destination is set when root is not") + .data_home + .clone() + }); + let files = + collect_parquet_files(store.as_ref(), root.as_deref(), args.include_metadata_files).await?; + let mut reports = futures::stream::iter(files.iter()) + .map(|file| inspect_file(Arc::clone(&store), file, &args.column)) + .buffer_unordered(8) + .try_collect::>() + .await?; + reports.sort_by(|a, b| a.relative_path.cmp(&b.relative_path)); println!( "{}", serde_json::to_string_pretty(&FooterReport { - root: args.root.display().to_string(), + root: root_display, summary: summarize_footer(&args.column, &reports), files: reports })? ); Ok(()) } -fn collect_parquet_files( - root: &Path, + +/// Builds an arrow `object_store::ObjectStore` for the given destination. The +/// File backend (and the `--root` shortcut) uses the opendal `Fs` builder +/// directly to avoid the create_dir_all / clean_temp_dir side effects of +/// `object_store::factory::new_fs_object_store`; every other backend goes +/// through `object_store::factory::new_raw_object_store`. +async fn build_store( + destination: Option<&DestinationConfig>, + root: Option<&Path>, +) -> Result, Box> { + let operator = match destination { + Some(destination) => match &destination.object_store { + ObjectStoreConfig::File(_) => { + object_store::ObjectStore::new(Fs::default().root(&destination.data_home))?.finish() + } + _ => new_raw_object_store(&destination.object_store, &destination.data_home).await?, + }, + None => { + let root = root.expect("root must be set when destination is None"); + object_store::ObjectStore::new(Fs::default().root(&root.to_string_lossy()))?.finish() + } + }; + Ok(Arc::new(object_store_opendal::OpendalStore::new(operator))) +} + +async fn collect_parquet_files( + store: &dyn ObjectStore, + root: Option<&Path>, include_metadata_files: bool, -) -> Result, Box> { +) -> Result, Box> { + let metas = store.list(None).try_collect::>().await?; let mut files = Vec::new(); - collect_parquet_files_inner(root, root, include_metadata_files, &mut files)?; - files.sort(); + for meta in metas { + if !is_parquet_data_file(&meta) { + continue; + } + let is_metadata = meta + .location + .parts() + .any(|part| part.as_ref() == "metadata"); + if !include_metadata_files && is_metadata { + continue; + } + let relative_path = meta.location.to_string(); + let path = match root { + Some(root) => root.join(&relative_path).display().to_string(), + None => relative_path.clone(), + }; + files.push(ListedFile { + location: meta.location, + size: meta.size, + path, + relative_path, + }); + } + files.sort_by(|a, b| a.relative_path.cmp(&b.relative_path)); Ok(files) } -fn collect_parquet_files_inner( - root: &Path, - path: &Path, - include_metadata_files: bool, - files: &mut Vec, -) -> Result<(), Box> { - if path.is_dir() { - for entry in fs::read_dir(path)? { - collect_parquet_files_inner(root, &entry?.path(), include_metadata_files, files)?; - } - } else if path.extension().is_some_and(|ext| ext == "parquet") - && (include_metadata_files - || !path - .strip_prefix(root) - .unwrap_or(path) - .components() - .any(|c| c.as_os_str() == "metadata")) - { - files.push(path.to_path_buf()); - } - Ok(()) + +/// Keeps only real parquet data files: `.parquet` keys with a non-zero size. +/// Zero-size keys (and keys ending in `/`) are directory markers emitted by +/// some backends and are skipped even when they happen to end in `.parquet`, +/// since an empty parquet file has no readable footer anyway. +fn is_parquet_data_file(meta: &ObjectMeta) -> bool { + meta.location.extension() == Some("parquet") + && meta.size > 0 + && !meta.location.as_ref().ends_with('/') } -fn inspect_file( - root: &Path, - path: &Path, + +async fn inspect_file( + store: Arc, + file: &ListedFile, column: &str, -) -> Result> { - let file_size = fs::metadata(path)?.len(); - let reader = SerializedFileReader::new(File::open(path)?)?; - let metadata = reader.metadata().file_metadata(); - let row_groups = reader.metadata().row_groups(); +) -> Result> { + let mut reader = + ParquetObjectReader::new(store, file.location.clone()).with_file_size(file.size); + let metadata = ParquetMetaDataReader::new() + .load_and_finish(&mut reader, file.size) + .await?; + let file_metadata = metadata.file_metadata(); + let row_groups = metadata.row_groups(); let mut columns = Vec::new(); for (row_group_index, rg) in row_groups.iter().enumerate() { for (column_index, chunk) in rg.columns().iter().enumerate() { @@ -145,14 +243,10 @@ fn inspect_file( } } Ok(FooterFileReport { - path: path.display().to_string(), - relative_path: path - .strip_prefix(root) - .unwrap_or(path) - .display() - .to_string(), - file_size, - num_rows: metadata.num_rows(), + path: file.path.clone(), + relative_path: file.relative_path.clone(), + file_size: file.size, + num_rows: file_metadata.num_rows(), num_row_groups: row_groups.len(), columns, }) @@ -179,3 +273,199 @@ fn summarize_footer(column: &str, files: &[FooterFileReport]) -> FooterSummary { s.unique_encodings = encodings.into_iter().collect(); s } + +#[cfg(test)] +mod tests { + use std::fs::{self, File}; + + use parquet::column::writer::ColumnWriter; + use parquet::file::properties::WriterProperties; + use parquet::file::writer::SerializedFileWriter; + use parquet::schema::parser::parse_message_type; + use tempfile::TempDir; + + use super::*; + + /// Writes a minimal valid parquet file with a single `value` Int64 column + /// and `rows` values, so the footer can be read back through the object + /// store path. + fn write_parquet_file(path: &Path, rows: i64) { + let schema = Arc::new( + parse_message_type("message schema { REQUIRED INT64 value; }").expect("valid schema"), + ); + let props = Arc::new(WriterProperties::builder().build()); + let file = File::create(path).expect("create parquet file"); + let mut writer = SerializedFileWriter::new(file, schema, props).expect("init writer"); + let mut row_group = writer.next_row_group().expect("next row group"); + let mut column = row_group + .next_column() + .expect("next column") + .expect("has column"); + match column.untyped() { + ColumnWriter::Int64ColumnWriter(typed) => { + typed + .write_batch(&(1..=rows).collect::>(), None, None) + .expect("write batch"); + } + _ => panic!("unexpected column writer"), + } + column.close().expect("close column"); + row_group.close().expect("close row group"); + writer.close().expect("close writer"); + } + + fn fixture_tree(dir: &Path) { + // Nested data tree with a `metadata` subdirectory, a non-parquet file + // and a zero-size fake file. + let table = dir.join("data/table"); + fs::create_dir_all(&table).expect("create table dir"); + fs::create_dir_all(table.join("metadata")).expect("create metadata dir"); + write_parquet_file(&table.join("0001.parquet"), 3); + write_parquet_file(&table.join("0002.parquet"), 5); + write_parquet_file(&table.join("metadata/0001.parquet"), 7); + fs::write(table.join("notes.txt"), b"not parquet").expect("write notes"); + File::create(table.join("empty.parquet")).expect("create empty file"); + } + + /// Runs the full list -> filter -> footer-read path against a File backend + /// rooted at `dir` and returns the collected file reports. + async fn inspect_fixture(dir: &Path, include_metadata_files: bool) -> Vec { + let root = Some(dir.to_path_buf()); + let store = build_store(None, root.as_deref()) + .await + .expect("build store"); + let files = collect_parquet_files(store.as_ref(), root.as_deref(), include_metadata_files) + .await + .expect("collect files"); + let mut reports = futures::stream::iter(files.iter()) + .map(|file| inspect_file(Arc::clone(&store), file, "value")) + .buffer_unordered(8) + .try_collect::>() + .await + .expect("inspect files"); + reports.sort_by(|a, b| a.relative_path.cmp(&b.relative_path)); + reports + } + + #[tokio::test] + async fn lists_and_inspects_footer_files_without_metadata_by_default() { + let dir = TempDir::new().expect("tempdir"); + fixture_tree(dir.path()); + + let reports = inspect_fixture(dir.path(), false).await; + + let relative_paths = reports + .iter() + .map(|report| report.relative_path.as_str()) + .collect::>(); + // metadata/0001.parquet excluded, notes.txt and empty.parquet excluded. + assert_eq!( + relative_paths, + vec!["data/table/0001.parquet", "data/table/0002.parquet"] + ); + assert!(reports.iter().all(|report| report.file_size > 0)); + assert!(reports.iter().all(|report| report.num_row_groups == 1)); + assert!(reports.iter().all(|report| report.num_rows > 0)); + assert!( + reports + .iter() + .all(|report| report.columns.len() == 1 && report.columns[0].column_path == "value") + ); + // Default writer properties use dictionary encoding, so the column chunk + // reports PLAIN (dictionary page) plus RLE_DICTIONARY (data pages). + assert!( + reports + .iter() + .all(|report| report.columns[0].encodings.contains(&"PLAIN".to_string())) + ); + assert!( + reports + .iter() + .all(|report| report.columns[0].compression == "UNCOMPRESSED") + ); + } + + #[tokio::test] + async fn includes_metadata_files_when_requested() { + let dir = TempDir::new().expect("tempdir"); + fixture_tree(dir.path()); + + let reports = inspect_fixture(dir.path(), true).await; + + let relative_paths = reports + .iter() + .map(|report| report.relative_path.as_str()) + .collect::>(); + assert_eq!( + relative_paths, + vec![ + "data/table/0001.parquet", + "data/table/0002.parquet", + "data/table/metadata/0001.parquet" + ] + ); + let metadata_report = reports + .iter() + .find(|report| report.relative_path == "data/table/metadata/0001.parquet") + .expect("metadata report"); + assert_eq!(metadata_report.num_rows, 7); + } + + #[tokio::test] + async fn accepts_a_file_destination_toml() { + let dir = TempDir::new().expect("tempdir"); + fixture_tree(dir.path()); + let destination = dir.path().join("destination.toml"); + fs::write( + &destination, + format!( + "data_home = \"{}\"\n[object_store]\ntype = \"File\"\n", + dir.path().display() + ), + ) + .expect("write destination toml"); + + let destination: DestinationConfig = + toml::from_str(&fs::read_to_string(&destination).expect("read toml")) + .expect("parse toml"); + let store = build_store(Some(&destination), None) + .await + .expect("build store"); + let files = collect_parquet_files(store.as_ref(), None, false) + .await + .expect("collect files"); + + // Destination mode has no root prefix: paths are store-relative keys. + assert_eq!( + files + .iter() + .map(|file| file.relative_path.as_str()) + .collect::>(), + vec!["data/table/0001.parquet", "data/table/0002.parquet"] + ); + assert!(files.iter().all(|file| file.path == file.relative_path)); + } + + #[test] + fn requires_exactly_one_of_root_or_destination() { + use futures::FutureExt; + + let both = run_inspect_footer(InspectFooterArgs { + root: Some(PathBuf::from("/tmp/root")), + destination: Some(PathBuf::from("/tmp/destination.toml")), + column: "value".to_string(), + include_metadata_files: false, + }) + .now_or_never(); + assert!(matches!(both, Some(Err(_)))); + + let neither = run_inspect_footer(InspectFooterArgs { + root: None, + destination: None, + column: "value".to_string(), + include_metadata_files: false, + }) + .now_or_never(); + assert!(matches!(neither, Some(Err(_)))); + } +} diff --git a/src/cmd/src/bin/query_perf_fixture/mod.rs b/src/cmd/src/bin/query_perf_fixture/mod.rs index 1a794c0c3b5..17f639bb54e 100644 --- a/src/cmd/src/bin/query_perf_fixture/mod.rs +++ b/src/cmd/src/bin/query_perf_fixture/mod.rs @@ -95,9 +95,9 @@ pub async fn run() { Some(Command::PromRemoteWrite(rw)) => run_prom_remote_write(rw) .await .expect("prom remote write failed"), - Some(Command::InspectFooter(inspect)) => { - run_inspect_footer(inspect).expect("inspect footer failed") - } + Some(Command::InspectFooter(inspect)) => run_inspect_footer(inspect) + .await + .expect("inspect footer failed"), Some(Command::Plan(plan)) => run_plan(plan).expect("plan failed"), None => run_direct_sst(args.legacy).await, } diff --git a/src/cmd/src/bin/query_regression_runner.rs b/src/cmd/src/bin/query_regression_runner.rs new file mode 100644 index 00000000000..adc7a2496d2 --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner.rs @@ -0,0 +1,23 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#![allow(clippy::print_stderr, clippy::print_stdout)] + +#[path = "query_regression_runner/mod.rs"] +mod query_regression_runner; + +#[tokio::main] +async fn main() { + query_regression_runner::run().await; +} diff --git a/src/cmd/src/bin/query_regression_runner/direct.rs b/src/cmd/src/bin/query_regression_runner/direct.rs new file mode 100644 index 00000000000..04c191c8c55 --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/direct.rs @@ -0,0 +1,501 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fs; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::time::{Duration, Instant}; + +use reqwest::Client; +use serde::Serialize; +use serde_json::{Value, json}; + +use crate::query_regression_runner::model::{Scenario, Table}; +use crate::query_regression_runner::plan::{load_plan, validate_direct_tables}; +use crate::query_regression_runner::sql::{ + extract_rows, http_post_sql, row_u64, row_value, sql_ident, sql_string, value_text, +}; +use crate::query_regression_runner::{PrepareDirectArgs, Result}; + +#[derive(Clone, Debug, Serialize)] +struct Discovery { + table: String, + catalog: String, + schema: String, + table_id: u64, + region_id: u64, + region_seq: u64, + table_dir: String, + region_dir: String, + peer_id: u64, + peer_addr: Value, + is_leader: Value, + status: Value, + discovery_queries: Value, +} + +#[derive(Debug)] +struct PreparedTarget { + creates: Vec, + discoveries: Vec, +} + +pub(super) async fn run_prepare_direct(args: PrepareDirectArgs) -> Result<()> { + if !args.http_timeout.is_finite() || args.http_timeout < 0.0 { + return Err("--http-timeout must be a non-negative finite number".into()); + } + let case_path = args.case.canonicalize()?; + let plan = load_plan(&args.fixture_generator, &case_path)?; + let scenario_value = plan + .get("scenario") + .cloned() + .ok_or("fixture plan has no scenario")?; + let scenario: Scenario = serde_json::from_value(scenario_value)?; + let tables = direct_tables(scenario)?; + let client = Client::builder() + .timeout(Duration::from_secs_f64(args.http_timeout)) + .build()?; + let base = prepare_target("base", args.base_http_port, &tables, &client).await?; + let candidate = prepare_target("candidate", args.candidate_http_port, &tables, &client).await?; + if base.discoveries.len() != candidate.discoveries.len() + || base + .discoveries + .iter() + .zip(&candidate.discoveries) + .any(|(base, candidate)| { + base.table != candidate.table + || base.region_id != candidate.region_id + || base.table_dir != candidate.table_dir + }) + { + return Err(format!( + "base/candidate metadata mismatch: base={:?}, candidate={:?}", + base.discoveries, candidate.discoveries + ) + .into()); + } + + let mut fixtures = Vec::with_capacity(tables.len()); + for (index, (table, discovery)) in tables.iter().zip(&base.discoveries).enumerate() { + let fixture_dir = table_fixture_dir(&args.fixture_dir, &tables, table, index); + fixtures.push(generate_direct_fixture( + &args.fixture_generator, + &case_path, + &fixture_dir, + table, + discovery, + tables.len() > 1, + args.allow_large_fixture, + )?); + } + + let report = json!({ + "case_path": case_path, + "scenario": "direct_readable_sst", + "base": { "create_table": base.creates, "discovery": base.discoveries }, + "candidate": { "create_table": candidate.creates, "discovery": candidate.discoveries }, + "fixtures": fixtures, + }); + let text = format!("{}\n", serde_json::to_string_pretty(&report)?); + if let Some(path) = args.output { + fs::write(path, &text)?; + } + print!("{text}"); + Ok(()) +} + +fn direct_tables(scenario: Scenario) -> Result> { + match scenario { + Scenario::DirectReadableSst { tables, layout, .. } => { + validate_direct_tables(tables, layout) + } + Scenario::PromRemoteWriteThenQuery { .. } => { + Err("prepare-direct requires scenario kind direct_readable_sst".into()) + } + Scenario::OtlpTraceLoad { .. } => { + Err("prepare-direct requires scenario kind direct_readable_sst".into()) + } + } +} + +async fn prepare_target( + name: &str, + port: u16, + tables: &[Table], + client: &Client, +) -> Result { + let mut creates = Vec::with_capacity(tables.len()); + let mut discoveries = Vec::with_capacity(tables.len()); + for table in tables { + let sql = create_table_sql(table)?; + let result = http_post_sql(client, port, &sql, &table.database).await; + if !result["ok"].as_bool().unwrap_or(false) { + return Err(format!("CREATE TABLE {} failed for {name}: {result}", table.name).into()); + } + creates.push(json!({ "table": table.name, "sql": sql, "result": result })); + discoveries.push(discover_region(client, port, table).await?); + } + Ok(PreparedTarget { + creates, + discoveries, + }) +} + +fn create_table_sql(table: &Table) -> Result { + if table.engine != "mito" { + return Err("prepare-direct requires table engine mito".into()); + } + let time_index = table + .time_index + .as_deref() + .ok_or("direct-SST table requires time_index")?; + if table.columns.is_empty() { + return Err("direct-SST table requires columns".into()); + } + let columns = table + .columns + .iter() + .map(|column| format!("{} {}", sql_ident(&column.name), column.data_type)) + .collect::>() + .join(",\n "); + let primary_key = table + .primary_key + .iter() + .map(|column| sql_ident(column)) + .collect::>() + .join(", "); + let mut options = Vec::new(); + if let Some(append_mode) = table.append_mode { + options.push(("append_mode", append_mode.to_string())); + } + if let Some(sst_format) = table + .sst_format + .as_deref() + .filter(|value| !value.is_empty()) + { + options.push(("sst_format", sst_format.to_string())); + } + let with = (!options.is_empty()).then(|| { + format!( + "\nWITH ({})", + options + .iter() + .map(|(key, value)| format!("'{key}'='{value}'")) + .collect::>() + .join(", ") + ) + }); + Ok(format!( + "CREATE TABLE {} (\n {columns},\n TIME INDEX ({}),\n PRIMARY KEY ({primary_key})\n) ENGINE=mito{};", + sql_ident(&table.name), + sql_ident(time_index), + with.unwrap_or_default() + )) +} + +async fn discover_region(client: &Client, port: u16, table: &Table) -> Result { + let schema = sql_string(&table.database); + let table_name = sql_string(&table.name); + let table_sql = format!( + "SELECT table_id FROM information_schema.tables WHERE table_schema = {schema} AND table_name = {table_name}" + ); + let table_result = http_post_sql(client, port, &table_sql, &table.database).await; + if !table_result["ok"].as_bool().unwrap_or(false) { + return Err(format!("table_id discovery failed: {table_result}").into()); + } + let table_rows = extract_rows(table_result.get("response").unwrap_or(&Value::Null)); + if table_rows.len() != 1 { + return Err(format!( + "expected one information_schema.tables row, got {}: {table_result}", + table_rows.len() + ) + .into()); + } + let table_id = row_u64(&table_rows[0], 0, "table_id")?; + + let region_sql = format!( + "SELECT region_id, peer_id, peer_addr, is_leader, status FROM information_schema.region_peers WHERE table_schema = {schema} AND table_name = {table_name}" + ); + let region_result = http_post_sql(client, port, ®ion_sql, &table.database).await; + if !region_result["ok"].as_bool().unwrap_or(false) { + return Err(format!("region_peers discovery failed: {region_result}").into()); + } + let region_rows = extract_rows(region_result.get("response").unwrap_or(&Value::Null)); + if region_rows.len() != 1 { + return Err(format!( + "expected one information_schema.region_peers row, got {}: {region_result}", + region_rows.len() + ) + .into()); + } + let row = ®ion_rows[0]; + let region_id = row_u64(row, 0, "region_id")?; + let peer_id = row_u64(row, 1, "peer_id")?; + let peer_addr = row_value(row, 2, "peer_addr") + .cloned() + .unwrap_or(Value::Null); + let is_leader = row_value(row, 3, "is_leader") + .cloned() + .unwrap_or(Value::Null); + let status = row_value(row, 4, "status").cloned().unwrap_or(Value::Null); + if !matches!( + value_text(&is_leader).to_ascii_lowercase().as_str(), + "yes" | "true" + ) { + return Err(format!( + "expected leader region peer, got is_leader={is_leader}: {region_result}" + ) + .into()); + } + if !value_text(&status).eq_ignore_ascii_case("ALIVE") { + return Err( + format!("expected ALIVE region peer, got status={status}: {region_result}").into(), + ); + } + if peer_id != 0 { + return Err(format!("expected region leader on datanode peer_id=0, got {peer_id}").into()); + } + if region_id >> 32 != table_id { + return Err(format!( + "table_id mismatch: tables={table_id}, region_id-derived={}", + region_id >> 32 + ) + .into()); + } + let region_seq = region_id & u32::MAX as u64; + let (table_dir, region_dir) = storage_paths(&table.database, table_id, region_seq); + Ok(Discovery { + table: table.name.clone(), + catalog: "greptime".to_string(), + schema: table.database.clone(), + table_id, + region_id, + region_seq, + table_dir: table_dir.clone(), + region_dir, + peer_id, + peer_addr, + is_leader, + status, + discovery_queries: json!({ "table": table_result, "region": region_result }), + }) +} + +fn storage_paths(database: &str, table_id: u64, region_seq: u64) -> (String, String) { + let table_dir = format!("data/greptime/{database}/{table_id}/"); + let region_dir = format!("{table_dir}{table_id}_{region_seq:010}"); + (table_dir, region_dir) +} + +fn table_fixture_dir(root: &Path, tables: &[Table], table: &Table, index: usize) -> PathBuf { + if tables.len() == 1 { + root.to_path_buf() + } else { + root.join(fixture_subdir(table, index)) + } +} + +fn fixture_subdir(table: &Table, index: usize) -> String { + let raw = format!("{index:02}_{}_{}", table.database, table.name); + let mut safe = String::new(); + let mut replaced = false; + for character in raw.chars() { + if character.is_ascii_alphanumeric() || matches!(character, '_' | '.' | '-') { + safe.push(character); + replaced = false; + } else if !replaced { + safe.push('_'); + replaced = true; + } + } + let safe = safe.trim_matches(['.', '_', '-']); + if safe.is_empty() { + format!("table_{index:02}") + } else { + safe.to_string() + } +} + +fn generate_direct_fixture( + generator: &Path, + case_path: &Path, + fixture_dir: &Path, + table: &Table, + discovery: &Discovery, + multi_table: bool, + allow_large_fixture: bool, +) -> Result { + if fixture_dir.exists() { + fs::remove_dir_all(fixture_dir)?; + } + fs::create_dir_all(fixture_dir)?; + let mut command = vec![ + generator.to_string_lossy().to_string(), + "direct-sst".to_string(), + "--case".to_string(), + case_path.to_string_lossy().to_string(), + "--out-dir".to_string(), + fixture_dir.to_string_lossy().to_string(), + ]; + if multi_table { + command.extend(["--table".to_string(), table.name.clone()]); + } + command.extend([ + "--region-id".to_string(), + discovery.region_id.to_string(), + "--table-dir".to_string(), + discovery.table_dir.clone(), + ]); + if allow_large_fixture { + command.push("--allow-large".to_string()); + } + let started = Instant::now(); + let output = Command::new(generator).args(&command[1..]).output()?; + let elapsed_seconds = started.elapsed().as_secs_f64(); + let stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + if !output.status.success() { + return Err(format!( + "fixture generator failed: command={command:?}, returncode={:?}, elapsed_seconds={elapsed_seconds:.3}, stderr={:.2000}", + output.status.code(), stderr + ) + .into()); + } + let summary: Value = serde_json::from_slice(&fs::read(fixture_dir.join("summary.json"))?)?; + assert_fixture_summary(&summary, table, discovery)?; + Ok(json!({ + "status": "ok", + "fixture_dir": fixture_dir, + "command": command, + "returncode": output.status.code(), + "elapsed_seconds": elapsed_seconds, + "stdout": stdout, + "stderr": stderr, + "summary": summary, + })) +} + +fn assert_fixture_summary(summary: &Value, table: &Table, discovery: &Discovery) -> Result<()> { + if summary.get("table").and_then(Value::as_str) != Some(&table.name) { + return Err(format!( + "fixture table mismatch: {:?} != {}", + summary.get("table"), + table.name + ) + .into()); + } + if summary.get("database").and_then(Value::as_str) != Some(&table.database) { + return Err(format!( + "fixture database mismatch: {:?} != {}", + summary.get("database"), + table.database + ) + .into()); + } + if summary.get("region_id").and_then(|value| { + value + .as_u64() + .or_else(|| value.as_str().and_then(|value| value.parse().ok())) + }) != Some(discovery.region_id) + { + return Err(format!( + "fixture region_id mismatch: {:?} != {}", + summary.get("region_id"), + discovery.region_id + ) + .into()); + } + if summary.get("table_dir").and_then(Value::as_str) != Some(&discovery.table_dir) { + return Err(format!( + "fixture table_dir mismatch: {:?} != {}", + summary.get("table_dir"), + discovery.table_dir + ) + .into()); + } + if summary + .get("region_dir") + .and_then(Value::as_str) + .map(|value| value.trim_matches('/')) + != Some(discovery.region_dir.trim_matches('/')) + { + return Err(format!( + "fixture region_dir mismatch: {:?} != {}", + summary.get("region_dir"), + discovery.region_dir + ) + .into()); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::query_regression_runner::model::Column; + + #[test] + fn creates_exact_direct_sst_table_sql() { + let table = Table { + database: "public".to_string(), + name: "metric\"name".to_string(), + engine: "mito".to_string(), + columns: vec![ + Column { + name: "host".to_string(), + data_type: "STRING".to_string(), + }, + Column { + name: "ts".to_string(), + data_type: "TIMESTAMP(9)".to_string(), + }, + ], + primary_key: vec!["host".to_string()], + time_index: Some("ts".to_string()), + append_mode: Some(true), + sst_format: Some("flat".to_string()), + validate_show_create_engine: true, + }; + assert_eq!( + create_table_sql(&table).unwrap(), + "CREATE TABLE \"metric\"\"name\" (\n \"host\" STRING,\n \"ts\" TIMESTAMP(9),\n TIME INDEX (\"ts\"),\n PRIMARY KEY (\"host\")\n) ENGINE=mito\nWITH ('append_mode'='true', 'sst_format'='flat');" + ); + } + + #[test] + fn extracts_named_and_positional_discovery_rows() { + let rows = extract_rows(&json!({"output": [{"data": [{"TABLE_ID": "7"}]}]})); + assert_eq!(rows, vec![json!({"TABLE_ID": "7"})]); + assert_eq!(row_u64(&rows[0], 0, "table_id").unwrap(), 7); + let rows = extract_rows(&json!({"data": [[30064771073_u64, 0, "addr", "YES", "ALIVE"]]})); + assert_eq!(rows.len(), 1); + assert_eq!(row_u64(&rows[0], 0, "region_id").unwrap(), 30_064_771_073); + assert_eq!(row_u64(&rows[0], 1, "peer_id").unwrap(), 0); + assert_eq!( + value_text(row_value(&rows[0], 3, "is_leader").unwrap()), + "YES" + ); + } + + #[test] + fn storage_paths_include_the_table_directory() { + assert_eq!( + storage_paths("public", 1024, 0), + ( + "data/greptime/public/1024/".to_string(), + "data/greptime/public/1024/1024_0000000000".to_string(), + ) + ); + } +} diff --git a/src/cmd/src/bin/query_regression_runner/materialize.rs b/src/cmd/src/bin/query_regression_runner/materialize.rs new file mode 100644 index 00000000000..3cc9e1aa07d --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/materialize.rs @@ -0,0 +1,225 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fs; +use std::path::Path; + +use futures::{AsyncWriteExt as _, TryStreamExt}; +use object_store::ObjectStore; +use object_store::factory::new_raw_object_store; +use object_store::services::Fs; + +use crate::query_regression_runner::model::{ + CopyCounts, DestinationConfig, FixtureSummary, MaterializeResult, +}; +use crate::query_regression_runner::{MaterializeArgs, Result}; + +pub(super) async fn run_materialize(args: MaterializeArgs) -> Result<()> { + let destination: DestinationConfig = toml::from_str(&fs::read_to_string(args.destination)?)?; + let result = materialize(&args.fixture_dir, destination).await?; + println!("{}", serde_json::to_string(&result)?); + Ok(()) +} + +async fn materialize( + fixture_dir: &Path, + destination: DestinationConfig, +) -> Result { + let summary: FixtureSummary = + serde_json::from_slice(&fs::read(fixture_dir.join("summary.json"))?)?; + let region_dir = validate_region_dir(&summary.region_dir)?; + let object_source = fs_operator(&fixture_dir.join("object-store"))?; + let manifest_source = fs_operator(&fixture_dir.join("manifest"))?; + let destination = + new_raw_object_store(&destination.object_store, &destination.data_home).await?; + + let region_prefix = format!("{region_dir}/"); + destination + .delete_with(®ion_prefix) + .recursive(true) + .await?; + + let object_store = copy_tree(&object_source, &destination, "/", "").await?; + let manifest = copy_tree( + &manifest_source, + &destination, + "/", + &format!("{region_dir}/manifest/"), + ) + .await?; + Ok(MaterializeResult { + region_dir, + object_store, + manifest, + }) +} + +fn fs_operator(root: &Path) -> Result { + Ok(ObjectStore::new(Fs::default().root(&root.to_string_lossy()))?.finish()) +} + +fn validate_region_dir(region_dir: &str) -> Result { + let region_dir = region_dir.trim_end_matches('/'); + if region_dir.is_empty() + || region_dir.starts_with('/') + || region_dir.contains('\\') + || region_dir + .split('/') + .any(|component| component.is_empty() || matches!(component, "." | "..")) + { + return Err( + "region_dir must be a non-empty relative OpenDAL key without dot components".into(), + ); + } + Ok(region_dir.to_string()) +} + +async fn copy_tree( + source: &ObjectStore, + destination: &ObjectStore, + source_prefix: &str, + destination_prefix: &str, +) -> Result { + let mut lister = source.lister_with(source_prefix).recursive(true).await?; + let mut counts = CopyCounts::default(); + while let Some(entry) = lister.try_next().await? { + if entry.metadata().is_dir() { + continue; + } + let source_path = entry.path().to_string(); + let reader = source + .reader(&source_path) + .await? + .into_futures_async_read(0..entry.metadata().content_length()) + .await?; + let mut writer = destination + .writer(&format!("{destination_prefix}{source_path}")) + .await? + .into_futures_async_write(); + let bytes = futures::io::copy(reader, &mut writer).await?; + writer.close().await?; + counts.files += 1; + counts.bytes += bytes; + } + Ok(counts) +} + +#[cfg(test)] +mod tests { + use std::fs; + + use object_store::config::ObjectStoreConfig; + + use super::*; + + #[test] + fn region_dir_must_be_a_relative_opendal_key() { + assert_eq!( + validate_region_dir("data/db/region/").unwrap(), + "data/db/region" + ); + for region_dir in [ + "", + "/region", + ".", + "a/../region", + "a/./region", + "a\\region", + "a//region", + ] { + assert!(validate_region_dir(region_dir).is_err(), "{region_dir}"); + } + } + + #[tokio::test] + async fn materializes_fixture_from_fs_to_fs() { + let fixture = tempfile::tempdir().unwrap(); + let destination = tempfile::tempdir().unwrap(); + let region_dir = "data/public/metrics/00000000000000000001"; + let sst = fixture + .path() + .join("object-store") + .join(region_dir) + .join("00000000000000000001.parquet"); + fs::create_dir_all(sst.parent().unwrap()).unwrap(); + fs::write(&sst, b"sst").unwrap(); + let checkpoint = fixture + .path() + .join("manifest/00000000000000000001.checkpoint"); + fs::create_dir_all(checkpoint.parent().unwrap()).unwrap(); + fs::write(&checkpoint, b"checkpoint").unwrap(); + fs::write(fixture.path().join("manifest/_last_checkpoint"), b"last").unwrap(); + fs::write( + fixture.path().join("summary.json"), + format!(r#"{{"region_dir":"{region_dir}"}}"#), + ) + .unwrap(); + fs::write(fixture.path().join("files.jsonl"), "metadata").unwrap(); + fs::create_dir_all(destination.path().join(region_dir)).unwrap(); + fs::write(destination.path().join(region_dir).join("stale"), "stale").unwrap(); + fs::write(destination.path().join("unrelated"), "keep").unwrap(); + + let result = materialize( + fixture.path(), + DestinationConfig { + data_home: destination.path().to_string_lossy().to_string(), + object_store: ObjectStoreConfig::default(), + }, + ) + .await + .unwrap(); + + assert_eq!(result.object_store.files, 1); + assert_eq!(result.object_store.bytes, 3); + assert_eq!(result.manifest.files, 2); + assert_eq!(result.manifest.bytes, 14); + assert_eq!( + fs::read( + destination + .path() + .join(region_dir) + .join("00000000000000000001.parquet") + ) + .unwrap(), + b"sst" + ); + assert_eq!( + fs::read( + destination + .path() + .join(region_dir) + .join("manifest/00000000000000000001.checkpoint") + ) + .unwrap(), + b"checkpoint" + ); + assert_eq!( + fs::read( + destination + .path() + .join(region_dir) + .join("manifest/_last_checkpoint") + ) + .unwrap(), + b"last" + ); + assert!(!destination.path().join(region_dir).join("stale").exists()); + assert_eq!( + fs::read(destination.path().join("unrelated")).unwrap(), + b"keep" + ); + assert!(!destination.path().join("summary.json").exists()); + assert!(!destination.path().join("files.jsonl").exists()); + } +} diff --git a/src/cmd/src/bin/query_regression_runner/measure.rs b/src/cmd/src/bin/query_regression_runner/measure.rs new file mode 100644 index 00000000000..2b389f7d782 --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/measure.rs @@ -0,0 +1,423 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::time::Duration; +use std::{fs, io}; + +use reqwest::Client; +use serde_json::{Map, Value, json}; + +use crate::query_regression_runner::model::{Measurement, Query, QueryResult, Scenario, Table}; +use crate::query_regression_runner::plan::{load_plan, normalize_scenario}; +use crate::query_regression_runner::sql::{http_post_sql, sql_ident}; +use crate::query_regression_runner::{MeasureArgs, Result}; + +pub(super) async fn run_measure(args: MeasureArgs) -> Result<()> { + if !args.http_timeout.is_finite() || args.http_timeout < 0.0 { + return Err("--http-timeout must be a non-negative finite number".into()); + } + + let case_path = args.case.canonicalize()?; + let plan = load_plan(&args.fixture_generator, &case_path)?; + let case_text = fs::read_to_string(&case_path)?; + let raw_case: toml::Value = toml::from_str(&case_text)?; + let case_metadata = raw_case + .get("case") + .cloned() + .unwrap_or(toml::Value::Table(toml::map::Map::new())); + let case_metadata = serde_json::to_value(case_metadata)?; + + let scenario_value = plan + .get("scenario") + .cloned() + .ok_or("fixture plan has no scenario")?; + let scenario: Scenario = serde_json::from_value(scenario_value.clone())?; + let (tables, configured_queries) = normalize_scenario(scenario)?; + + let client = Client::builder() + .timeout(Duration::from_secs_f64(args.http_timeout)) + .build()?; + let base = run_target(args.base_http_port, &tables, &configured_queries, &client).await; + let candidate = run_target( + args.candidate_http_port, + &tables, + &configured_queries, + &client, + ) + .await; + let thresholds = enforce_thresholds(&configured_queries, &base, &candidate)?; + let status = if base.status == "failed" + || candidate.status == "failed" + || thresholds + .iter() + .any(|threshold| threshold["status"] == "failed") + { + "failed" + } else { + "ok" + }; + let report = json!({ + "case_path": case_path, + "case": case_metadata, + "scenario": scenario_value, + "queries": configured_queries, + "query_mode": "endpoint", + "http_timeout": args.http_timeout, + "targets": [target_report("base", args.base_http_port, base), target_report("candidate", args.candidate_http_port, candidate)], + "thresholds": thresholds, + "status": status, + }); + let text = format!("{}\n", serde_json::to_string_pretty(&report)?); + if let Some(path) = args.output { + fs::write(path, &text)?; + } + print!("{text}"); + if status == "failed" { + std::process::exit(1); + } + Ok(()) +} + +fn target_report(name: &str, http_port: u16, result: QueryResult) -> Value { + let status = if result.status == "ok" { + "measured" + } else { + "failed" + }; + json!({ + "name": name, + "http_port": http_port, + "validation": result.validation, + "validation_errors": result.validation_errors, + "measurements": result.measurements, + "status": status, + }) +} + +async fn run_target( + port: u16, + tables: &[Table], + configured_queries: &[Query], + client: &Client, +) -> QueryResult { + let mut queries = configured_queries.to_vec(); + if queries.is_empty() { + queries.push(Query { + name: Some("count_all".to_string()), + kind: Some("sql".to_string()), + query: format!("SELECT count(*) FROM {}", sql_ident(&tables[0].name)), + warmup: 0, + iterations: 1, + thresholds: Map::new(), + }); + } + + let db = &tables[0].database; + let mut validation = Vec::new(); + let mut validation_errors = Vec::new(); + for table in tables { + let sql = format!("SHOW CREATE TABLE {}", sql_ident(&table.name)); + let sample = http_post_sql(client, port, &sql, &table.database).await; + if !sample["ok"].as_bool().unwrap_or(false) { + validation_errors.push(json!({ + "sql": sql, + "error": sample.get("error"), + "response": sample.get("response"), + })); + } else { + for error in validate_show_create(&sample, table) { + validation_errors.push(json!({ + "sql": sql, + "error": error, + "response": sample.get("response"), + })); + } + } + validation.push(sample); + } + let first = http_post_sql(client, port, &queries[0].query, db).await; + if !first["ok"].as_bool().unwrap_or(false) { + validation_errors.push(json!({ + "sql": queries[0].query, + "error": first.get("error"), + "response": first.get("response"), + })); + } + validation.push(first); + + let mut measurements = Vec::with_capacity(queries.len()); + for query in &queries { + for _ in 0..query.warmup { + let warmup = http_post_sql(client, port, &query.query, db).await; + if !warmup["ok"].as_bool().unwrap_or(false) { + validation_errors.push(json!({ + "sql": query.query, + "phase": "warmup", + "error": warmup.get("error"), + "response": warmup.get("response"), + })); + } + } + let mut samples = Vec::with_capacity(query.iterations); + let mut good_latencies = Vec::with_capacity(query.iterations); + for _ in 0..query.iterations { + let mut sample = http_post_sql(client, port, &query.query, db).await; + let execution_time = sample + .get("response") + .and_then(extract_execution_time) + .cloned() + .unwrap_or(Value::Null); + sample + .as_object_mut() + .expect("HTTP samples are objects") + .insert("execution_time_ms".to_string(), execution_time); + if sample["ok"].as_bool().unwrap_or(false) { + good_latencies.push(sample["latency_ms"].as_f64().unwrap_or_default()); + } + samples.push(sample); + } + let median = (!good_latencies.is_empty()).then(|| median(&good_latencies)); + let p95 = (!good_latencies.is_empty()).then(|| percentile(&good_latencies, 95.0)); + let status = if good_latencies.len() == samples.len() { + "ok" + } else { + "failed" + }; + measurements.push(Measurement { + name: query.name.clone(), + kind: query.kind.clone(), + iterations: samples.len(), + samples, + latency_ms_median: median, + latency_ms_p95: p95, + status: status.to_string(), + }); + } + let failed = !validation_errors.is_empty() || measurements.iter().any(|m| m.status == "failed"); + QueryResult { + validation, + validation_errors, + measurements, + status: if failed { "failed" } else { "ok" }.to_string(), + } +} + +fn validate_show_create(result: &Value, table: &Table) -> Vec<&'static str> { + let text = result + .get("response") + .map(response_text) + .unwrap_or_default() + .to_lowercase(); + let mut errors = Vec::new(); + if !text.contains(&table.name.to_lowercase()) { + errors.push("SHOW CREATE output does not contain table name"); + } + if table.validate_show_create_engine && (!text.contains("engine") || !text.contains("mito")) { + errors.push("SHOW CREATE output does not mention ENGINE=mito"); + } + if table.append_mode.is_some() && !text.contains("append_mode") { + errors.push("SHOW CREATE output does not mention append_mode"); + } + if table.sst_format.is_some() && !text.contains("sst_format") { + errors.push("SHOW CREATE output does not mention sst_format"); + } + errors +} + +fn response_text(body: &Value) -> String { + body.as_str() + .map(ToOwned::to_owned) + .unwrap_or_else(|| serde_json::to_string(body).unwrap_or_default()) +} + +fn extract_execution_time(body: &Value) -> Option<&Value> { + match body { + Value::Object(map) => { + for key in ["execution_time_ms", "execution_time", "elapsed"] { + if let Some(value) = map.get(key) { + return Some(value); + } + } + map.values().find_map(extract_execution_time) + } + Value::Array(values) => values.iter().find_map(extract_execution_time), + _ => None, + } +} + +pub(super) fn median(values: &[f64]) -> f64 { + let mut ordered = values.to_vec(); + ordered.sort_by(f64::total_cmp); + let middle = ordered.len() / 2; + if ordered.len().is_multiple_of(2) { + (ordered[middle - 1] + ordered[middle]) / 2.0 + } else { + ordered[middle] + } +} + +fn percentile(values: &[f64], pct: f64) -> f64 { + if values.is_empty() { + return 0.0; + } + let mut ordered = values.to_vec(); + ordered.sort_by(f64::total_cmp); + let index = round_ties_even((pct / 100.0) * (ordered.len() - 1) as f64) + .clamp(0, ordered.len() as isize - 1) as usize; + ordered[index] +} + +pub(super) fn round_ties_even(value: f64) -> isize { + let floor = value.floor(); + let fraction = value - floor; + if fraction < 0.5 { + floor as isize + } else if fraction > 0.5 { + floor as isize + 1 + } else if (floor as isize) % 2 == 0 { + floor as isize + } else { + floor as isize + 1 + } +} + +fn enforce_thresholds( + queries: &[Query], + base: &QueryResult, + candidate: &QueryResult, +) -> Result> { + let base_by_name: HashMap<_, _> = base + .measurements + .iter() + .map(|measurement| (measurement.name.as_deref(), measurement)) + .collect(); + let mut results = Vec::new(); + for candidate_measurement in &candidate.measurements { + let query = queries + .iter() + .find(|query| query.name == candidate_measurement.name); + let thresholds = query.map_or_else(Map::new, |query| query.thresholds.clone()); + let base_measurement = base_by_name.get(&candidate_measurement.name.as_deref()); + if let Some(limit) = thresholds + .get("max_candidate_latency_regression_pct") + .filter(|value| !value.is_null()) + .map(value_as_f64) + .transpose()? + { + let result = match base_measurement { + None => { + json!({"query": candidate_measurement.name, "threshold": "max_candidate_latency_regression_pct", "status": "failed", "reason": "missing base measurement"}) + } + Some(base) if matches!(base.latency_ms_median, None | Some(0.0)) => { + json!({"query": candidate_measurement.name, "threshold": "max_candidate_latency_regression_pct", "status": "failed", "reason": "base median latency is missing or zero", "base_latency_ms_median": base.latency_ms_median}) + } + Some(_) if candidate_measurement.latency_ms_median.is_none() => { + json!({"query": candidate_measurement.name, "threshold": "max_candidate_latency_regression_pct", "status": "failed", "reason": "missing candidate measurement"}) + } + Some(base) => { + let actual = (candidate_measurement.latency_ms_median.unwrap() + - base.latency_ms_median.unwrap()) + / base.latency_ms_median.unwrap() + * 100.0; + json!({"query": candidate_measurement.name, "threshold": "max_candidate_latency_regression_pct", "status": if actual <= limit { "passed" } else { "failed" }, "actual_pct": actual, "limit_pct": limit}) + } + }; + results.push(result); + } + for key in thresholds + .keys() + .filter(|key| *key != "max_candidate_latency_regression_pct") + { + results.push(json!({"query": candidate_measurement.name, "threshold": key, "status": "failed", "reason": "unsupported threshold"})); + } + } + Ok(results) +} + +fn value_as_f64(value: &Value) -> Result { + value.as_f64().ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "threshold must be numeric").into() + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn median_and_p95_match_python() { + assert_eq!(median(&[1.0, 8.0, 3.0, 4.0]), 3.5); + assert_eq!( + percentile( + &[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], + 95.0 + ), + 10.0 + ); + assert_eq!(percentile(&[], 95.0), 0.0); + } + + #[test] + fn finds_nested_execution_time_in_priority_order() { + let body = json!({"output": [{"elapsed": 3}], "execution_time": 2}); + assert_eq!(extract_execution_time(&body), Some(&json!(2))); + assert_eq!( + extract_execution_time(&json!({"output": [{"elapsed": 3}]})), + Some(&json!(3)) + ); + } + + #[test] + fn threshold_rejects_unknown_keys_and_zero_base() { + let query = Query { + name: Some("q".to_string()), + kind: None, + query: "SELECT 1".to_string(), + warmup: 0, + iterations: 1, + thresholds: Map::from_iter([ + ("max_candidate_latency_regression_pct".to_string(), json!(0)), + ("other".to_string(), json!(1)), + ]), + }; + let measurement = |median| Measurement { + name: Some("q".to_string()), + kind: None, + iterations: 1, + samples: vec![], + latency_ms_median: median, + latency_ms_p95: median, + status: "ok".to_string(), + }; + let base = QueryResult { + validation: vec![], + validation_errors: vec![], + measurements: vec![measurement(Some(0.0))], + status: "ok".to_string(), + }; + let candidate = QueryResult { + validation: vec![], + validation_errors: vec![], + measurements: vec![measurement(Some(1.0))], + status: "ok".to_string(), + }; + let results = enforce_thresholds(&[query], &base, &candidate).unwrap(); + assert_eq!( + results[0]["reason"], + "base median latency is missing or zero" + ); + assert_eq!(results[1]["reason"], "unsupported threshold"); + } +} diff --git a/src/cmd/src/bin/query_regression_runner/mod.rs b/src/cmd/src/bin/query_regression_runner/mod.rs new file mode 100644 index 00000000000..9ed903f5e07 --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/mod.rs @@ -0,0 +1,204 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::error::Error; +use std::path::PathBuf; + +use clap::{Parser, Subcommand}; + +mod direct; +mod materialize; +mod measure; +mod model; +mod otlp; +mod plan; +mod remote; +mod sql; + +pub(super) type Result = std::result::Result>; + +#[derive(Debug, Parser)] +#[command(about = "Measure query endpoints or materialize a direct-SST fixture")] +struct Cli { + #[command(subcommand)] + command: RunnerCommand, +} + +#[derive(Debug, Subcommand)] +enum RunnerCommand { + /// Run normalized query regression cases against two HTTP endpoints. + Measure(MeasureArgs), + /// Create direct-SST tables, discover their regions, and generate fixtures. + PrepareDirect(PrepareDirectArgs), + /// Render the frontend Prometheus store configuration for a remote-write case. + RenderRemoteConfig(RenderRemoteConfigArgs), + /// Create a database and ingest a normalized Prometheus remote-write case. + PrepareRemote(PrepareRemoteArgs), + /// Inspect stopped remote-write storage and finalize the aggregate report. + FinalizeRemote(FinalizeRemoteArgs), + /// Run one externally managed OTLP trace-load target. + RunOtlpTarget(RunOtlpTargetArgs), + /// Combine externally managed OTLP trace-load target results. + FinalizeOtlp(FinalizeOtlpArgs), + /// Copy a direct-SST fixture into an object store destination. + Materialize(MaterializeArgs), +} + +#[derive(Debug, Parser)] +struct MeasureArgs { + #[arg(long, value_name = "PATH")] + case: PathBuf, + #[arg(long, value_name = "PATH")] + fixture_generator: PathBuf, + #[arg(long)] + base_http_port: u16, + #[arg(long)] + candidate_http_port: u16, + #[arg(long, value_name = "PATH")] + output: Option, + #[arg(long, default_value_t = 120.0)] + http_timeout: f64, +} + +#[derive(Debug, Parser)] +struct PrepareDirectArgs { + #[arg(long, value_name = "PATH")] + case: PathBuf, + #[arg(long, value_name = "PATH")] + fixture_generator: PathBuf, + #[arg(long)] + base_http_port: u16, + #[arg(long)] + candidate_http_port: u16, + #[arg(long, value_name = "PATH")] + fixture_dir: PathBuf, + #[arg(long, value_name = "PATH")] + output: Option, + #[arg(long, default_value_t = 120.0)] + http_timeout: f64, + #[arg(long)] + allow_large_fixture: bool, +} + +#[derive(Debug, Parser)] +struct RenderRemoteConfigArgs { + #[arg(long, value_name = "PATH")] + case: PathBuf, + #[arg(long, value_name = "PATH")] + fixture_generator: PathBuf, + #[arg(long, value_name = "PATH")] + output: PathBuf, +} + +#[derive(Debug, Parser)] +struct PrepareRemoteArgs { + #[arg(long, value_name = "PATH")] + case: PathBuf, + #[arg(long, value_name = "PATH")] + fixture_generator: PathBuf, + #[arg(long)] + base_http_port: u16, + #[arg(long)] + candidate_http_port: u16, + #[arg(long, value_name = "PATH")] + output: Option, + #[arg(long, default_value_t = 120.0)] + http_timeout: f64, +} + +#[derive(Debug, Parser)] +struct FinalizeRemoteArgs { + #[arg(long, value_name = "PATH")] + case: PathBuf, + #[arg(long, value_name = "PATH")] + fixture_generator: PathBuf, + #[arg(long, value_name = "PATH")] + candidate_bin: PathBuf, + #[arg(long, value_name = "PATH")] + base_data_home: Option, + #[arg(long, value_name = "PATH")] + candidate_data_home: Option, + #[arg(long, value_name = "PATH")] + base_destination: Option, + #[arg(long, value_name = "PATH")] + candidate_destination: Option, + #[arg(long, value_name = "PATH")] + report: PathBuf, +} + +#[derive(Debug, Parser)] +struct RunOtlpTargetArgs { + #[arg(long, value_name = "PATH")] + case: PathBuf, + #[arg(long, value_name = "PATH")] + fixture_generator: PathBuf, + #[arg(long, value_name = "PATH")] + otelgen_bin: PathBuf, + #[arg(long)] + http_port: u16, + #[arg(long)] + target_name: String, + #[arg(long, value_name = "PATH")] + work_dir: PathBuf, + #[arg(long, value_name = "PATH")] + output: PathBuf, + #[arg(long, default_value_t = 120.0)] + http_timeout: f64, +} + +#[derive(Debug, Parser)] +struct FinalizeOtlpArgs { + #[arg(long, value_name = "PATH")] + case: PathBuf, + #[arg(long, value_name = "PATH")] + fixture_generator: PathBuf, + #[arg(long, value_name = "PATH")] + base_result: PathBuf, + #[arg(long, value_name = "PATH")] + candidate_result: PathBuf, + #[arg(long, value_name = "PATH")] + output: PathBuf, +} + +#[derive(Debug, Parser)] +struct MaterializeArgs { + #[arg(long, value_name = "PATH")] + fixture_dir: PathBuf, + #[arg( + long, + value_name = "PATH", + help = "TOML file: data_home = \"...\" and object_store = { type = \"File\" }; object_store uses object_store::config::ObjectStoreConfig" + )] + destination: PathBuf, +} + +pub(super) async fn run() { + if let Err(error) = run_inner().await { + eprintln!("query_regression_runner: {error}"); + std::process::exit(1); + } +} + +async fn run_inner() -> Result<()> { + match Cli::parse().command { + RunnerCommand::Measure(args) => measure::run_measure(args).await, + RunnerCommand::PrepareDirect(args) => direct::run_prepare_direct(args).await, + RunnerCommand::RenderRemoteConfig(args) => remote::run_render_remote_config(args).await, + RunnerCommand::PrepareRemote(args) => remote::run_prepare_remote(args).await, + RunnerCommand::FinalizeRemote(args) => remote::run_finalize_remote(args).await, + RunnerCommand::RunOtlpTarget(args) => otlp::run_otlp_target(args).await, + RunnerCommand::FinalizeOtlp(args) => otlp::run_finalize_otlp(args).await, + RunnerCommand::Materialize(args) => materialize::run_materialize(args).await, + } +} diff --git a/src/cmd/src/bin/query_regression_runner/model.rs b/src/cmd/src/bin/query_regression_runner/model.rs new file mode 100644 index 00000000000..50d41b3f85f --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/model.rs @@ -0,0 +1,235 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use object_store::config::ObjectStoreConfig; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +#[derive(Debug, Deserialize)] +pub(super) struct DestinationConfig { + pub(super) data_home: String, + pub(super) object_store: ObjectStoreConfig, +} + +#[derive(Debug, Deserialize)] +pub(super) struct FixtureSummary { + pub(super) region_dir: String, +} + +#[derive(Debug, Default, Serialize)] +pub(super) struct CopyCounts { + pub(super) files: u64, + pub(super) bytes: u64, +} + +#[derive(Debug, Serialize)] +pub(super) struct MaterializeResult { + pub(super) region_dir: String, + pub(super) object_store: CopyCounts, + pub(super) manifest: CopyCounts, +} + +#[derive(Debug, Deserialize)] +#[serde(tag = "kind")] +pub(super) enum Scenario { + #[serde(rename = "direct_readable_sst")] + DirectReadableSst { + tables: Vec, + layout: Layout, + #[serde(default)] + queries: Vec, + }, + #[serde(rename = "prom_remote_write_then_query")] + PromRemoteWriteThenQuery { + remote_write: RemoteWrite, + #[serde(default)] + queries: Vec, + }, + #[serde(rename = "otlp_trace_load")] + OtlpTraceLoad { load: OtlpTraceLoad }, +} + +#[derive(Debug, Deserialize)] +pub(super) struct Layout { + pub(super) regions: usize, +} + +#[derive(Clone, Debug, Deserialize)] +pub(super) struct Table { + pub(super) database: String, + pub(super) name: String, + #[serde(default)] + pub(super) engine: String, + #[serde(default)] + pub(super) columns: Vec, + #[serde(default)] + pub(super) primary_key: Vec, + #[serde(default)] + pub(super) time_index: Option, + #[serde(default)] + pub(super) append_mode: Option, + #[serde(default)] + pub(super) sst_format: Option, + #[serde(default = "default_show_create_engine")] + pub(super) validate_show_create_engine: bool, +} + +#[derive(Clone, Debug, Deserialize)] +pub(super) struct Column { + pub(super) name: String, + #[serde(rename = "type")] + pub(super) data_type: String, +} + +const fn default_show_create_engine() -> bool { + true +} + +#[derive(Clone, Debug, Deserialize)] +pub(super) struct RemoteWrite { + pub(super) database: String, + pub(super) metric: String, + pub(super) physical_table: String, + pub(super) series_count: u64, + pub(super) samples_per_series: u64, + pub(super) start_unix_millis: i64, + pub(super) step_millis: i64, + pub(super) chunk_series_count: u64, + pub(super) timeout_seconds: u64, + pub(super) sample_chunk_size: Option, + pub(super) flush_every_sample_chunks: u64, + pub(super) visibility_timeout_seconds: u64, + pub(super) prom_store: PromStore, + pub(super) value: RemoteValue, + pub(super) storage: Option, + pub(super) read_bench: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub(super) struct PromStore { + pub(super) pending_rows_flush_interval: String, + pub(super) max_batch_rows: u64, + pub(super) max_concurrent_flushes: u64, + pub(super) worker_channel_capacity: u64, + pub(super) max_inflight_requests: u64, +} + +#[derive(Clone, Debug, Deserialize)] +pub(super) struct RemoteValue { + pub(super) pattern: String, + pub(super) base: f64, + pub(super) step: f64, + pub(super) cardinality: u64, + pub(super) seed: u64, + pub(super) run_length: u64, + pub(super) stall_every: u64, + pub(super) stall_length: u64, + pub(super) mixed_every: u64, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(super) struct OtlpTraceLoad { + pub(super) load: OtlpLoad, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(super) struct OtlpLoad { + pub(super) database: String, + pub(super) table: String, + pub(super) pipeline: String, + pub(super) duration_seconds: u64, + pub(super) warmup_seconds: u64, + pub(super) rate: u64, + pub(super) workers: usize, + pub(super) exporter_shards: usize, + pub(super) workload: String, + pub(super) visibility_timeout_seconds: u64, + pub(super) thresholds: OtlpThresholds, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(super) struct OtlpThresholds { + pub(super) max_candidate_throughput_regression_pct: f64, + pub(super) max_candidate_mean_latency_regression_pct: f64, + pub(super) max_failure_count: u64, +} + +#[derive(Clone, Debug, Deserialize)] +pub(super) struct StorageConfig { + pub(super) inspect: bool, + pub(super) column: String, + pub(super) root_suffix: Option, + pub(super) include_metadata_files: bool, + pub(super) min_files: u64, + pub(super) min_files_with_column: u64, + pub(super) require_encodings: Vec, + pub(super) forbid_encodings: Vec, + pub(super) max_total_file_size_bytes: Option, + pub(super) max_column_compressed_size_bytes: Option, + pub(super) max_column_uncompressed_size_bytes: Option, + pub(super) max_candidate_total_file_size_regression_pct: Option, + pub(super) max_candidate_column_compressed_size_regression_pct: Option, + pub(super) max_candidate_column_uncompressed_size_regression_pct: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub(super) struct ReadBenchConfig { + pub(super) enabled: bool, + pub(super) parquetbench: bool, + pub(super) scanbench: bool, + pub(super) iterations: u64, + pub(super) projection: Vec, + pub(super) parquet_reader: String, + pub(super) scan_scanner: String, + pub(super) parallelism: u64, + pub(super) max_files: Option, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub(super) struct Query { + #[serde(default)] + pub(super) name: Option, + #[serde(default)] + pub(super) kind: Option, + pub(super) query: String, + #[serde(default)] + pub(super) warmup: usize, + #[serde(default = "one")] + pub(super) iterations: usize, + #[serde(default)] + pub(super) thresholds: Map, +} + +const fn one() -> usize { + 1 +} + +#[derive(Debug, Serialize)] +pub(super) struct QueryResult { + pub(super) validation: Vec, + pub(super) validation_errors: Vec, + pub(super) measurements: Vec, + pub(super) status: String, +} + +#[derive(Debug, Serialize)] +pub(super) struct Measurement { + pub(super) name: Option, + pub(super) kind: Option, + pub(super) iterations: usize, + pub(super) samples: Vec, + pub(super) latency_ms_median: Option, + pub(super) latency_ms_p95: Option, + pub(super) status: String, +} diff --git a/src/cmd/src/bin/query_regression_runner/otlp.rs b/src/cmd/src/bin/query_regression_runner/otlp.rs new file mode 100644 index 00000000000..ceccad0ffcc --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/otlp.rs @@ -0,0 +1,618 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::fs; +use std::path::Path; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; + +use regex::Regex; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; + +use crate::query_regression_runner::measure::round_ties_even; +use crate::query_regression_runner::model::{OtlpLoad, OtlpThresholds}; +use crate::query_regression_runner::plan::normalized_otlp_load; +use crate::query_regression_runner::sql::{ + extract_count_value, http_post_sql, sql_ident, sql_string, value_f64, value_u64, +}; +use crate::query_regression_runner::{FinalizeOtlpArgs, Result, RunOtlpTargetArgs}; + +const OTLP_ROWS: &str = "greptime_frontend_otlp_traces_rows"; +const OTLP_FAILURES: &str = "greptime_frontend_otlp_traces_failure_count"; +const OTLP_ELAPSED_SUM: &str = "greptime_servers_http_otlp_traces_elapsed_sum"; +const OTLP_ELAPSED_COUNT: &str = "greptime_servers_http_otlp_traces_elapsed_count"; + +#[derive(Clone, Debug, Deserialize, Serialize)] +struct OtlpSnapshot { + captured_monotonic_seconds: f64, + values: HashMap, +} + +pub(super) async fn run_otlp_target(args: RunOtlpTargetArgs) -> Result<()> { + if !args.http_timeout.is_finite() || args.http_timeout < 0.0 { + return Err("--http-timeout must be a non-negative finite number".into()); + } + let (_, load) = normalized_otlp_load(&args.fixture_generator, &args.case)?; + let client = Client::builder() + .timeout(Duration::from_secs_f64(args.http_timeout)) + .build()?; + let result = match run_otlp_target_inner(&args, &load, &client).await { + Ok(result) => result, + Err(error) => { + json!({ "name": args.target_name, "status": "failed", "error": error.to_string() }) + } + }; + let text = format!("{}\n", serde_json::to_string_pretty(&result)?); + fs::write(&args.output, &text)?; + print!("{text}"); + if result["status"] == "failed" { + std::process::exit(1); + } + Ok(()) +} + +async fn run_otlp_target_inner( + args: &RunOtlpTargetArgs, + load: &OtlpLoad, + client: &Client, +) -> Result { + let create_database = http_post_sql( + client, + args.http_port, + &format!( + "CREATE DATABASE IF NOT EXISTS {}", + sql_ident(&load.database) + ), + "public", + ) + .await; + if !create_database["ok"].as_bool().unwrap_or(false) { + return Ok( + json!({ "name": args.target_name, "create_database": create_database, "status": "failed" }), + ); + } + let otelgen = run_otelgen_load( + &args.otelgen_bin, + args.http_port, + &args.work_dir, + load, + client, + ) + .await?; + let metrics = summarize_otlp_metrics(&otelgen)?; + let flush = http_post_sql( + client, + args.http_port, + &format!("ADMIN FLUSH_TABLE({})", sql_string(&load.table)), + &load.database, + ) + .await; + let visibility = poll_otlp_visibility( + client, + args.http_port, + &load.table, + &load.database, + metrics["accepted_spans"].as_u64().unwrap_or_default(), + load.visibility_timeout_seconds, + ) + .await?; + let checks_ok = otelgen["status"] == "ok" + && metrics["missing_metrics"] + .as_array() + .is_some_and(Vec::is_empty) + && metrics["accepted_spans"].as_u64().unwrap_or_default() > 0 + && metrics["http_requests"].as_u64().unwrap_or_default() > 0 + && flush["ok"].as_bool().unwrap_or(false) + && visibility["ok"].as_bool().unwrap_or(false) + && visibility["row_count_ok"].as_bool().unwrap_or(false); + Ok(json!({ + "name": args.target_name, + "otelgen": otelgen, + "metrics": metrics, + "create_database": create_database, + "flush": flush, + "visibility": visibility, + "status": if checks_ok { "measured" } else { "failed" }, + })) +} + +fn otelgen_command(otelgen_bin: &Path, http_port: u16, load: &OtlpLoad) -> Vec { + vec![ + otelgen_bin.to_string_lossy().to_string(), + "--protocol".to_string(), + "http".to_string(), + "--otel-exporter-otlp-endpoint".to_string(), + format!("127.0.0.1:{http_port}"), + "--otel-exporter-otlp-url-path".to_string(), + "/v1/otlp/v1/traces".to_string(), + "--header".to_string(), + format!("x-greptime-pipeline-name={}", load.pipeline), + "--header".to_string(), + format!("x-greptime-db-name={}", load.database), + "--header".to_string(), + format!("x-greptime-trace-table-name={}", load.table), + "--log-level".to_string(), + "error".to_string(), + "--insecure".to_string(), + "--duration".to_string(), + load.duration_seconds.to_string(), + "--rate".to_string(), + load.rate.to_string(), + "traces".to_string(), + "multi".to_string(), + "--workers".to_string(), + load.workers.to_string(), + "--scenarios".to_string(), + load.workload.clone(), + "--exporter-shards".to_string(), + load.exporter_shards.to_string(), + ] +} + +async fn run_otelgen_load( + otelgen_bin: &Path, + http_port: u16, + work_dir: &Path, + load: &OtlpLoad, + client: &Client, +) -> Result { + let command = otelgen_command(otelgen_bin, http_port, load); + let clock = Instant::now(); + let initial = fetch_otlp_metrics(client, http_port, &clock).await?; + let log_dir = work_dir.join("otelgen"); + fs::create_dir_all(&log_dir)?; + let stdout_path = log_dir.join("stdout.log"); + let stderr_path = log_dir.join("stderr.log"); + let mut child = Command::new(otelgen_bin) + .args(&command[1..]) + .stdout(Stdio::from(fs::File::create(&stdout_path)?)) + .stderr(Stdio::from(fs::File::create(&stderr_path)?)) + .spawn()?; + let started = Instant::now(); + if load.warmup_seconds > 0 { + let _ = wait_for_child(&mut child, Duration::from_secs(load.warmup_seconds)).await?; + } + let warmed = fetch_otlp_metrics(client, http_port, &clock).await?; + let mut timed_out = false; + if child.try_wait()?.is_none() { + let remaining = load + .duration_seconds + .saturating_sub(load.warmup_seconds) + .saturating_add(60) + .max(60); + if !wait_for_child(&mut child, Duration::from_secs(remaining)).await? { + timed_out = true; + child.kill()?; + let _ = child.wait()?; + } + } + let final_snapshot = fetch_otlp_metrics(client, http_port, &clock).await?; + let returncode = match child.try_wait()? { + Some(status) => status.code(), + None => { + child.kill()?; + child.wait()?.code() + } + }; + let elapsed_seconds = started.elapsed().as_secs_f64(); + Ok(json!({ + "status": if returncode == Some(0) && !timed_out && elapsed_seconds + 1.0 >= load.duration_seconds as f64 { "ok" } else { "failed" }, + "cmd": command, + "returncode": returncode, + "timed_out": timed_out, + "elapsed_seconds": elapsed_seconds, + "stdout_path": stdout_path, + "stderr_path": stderr_path, + "snapshots": { "initial": initial, "warmup": warmed, "final": final_snapshot }, + })) +} + +async fn wait_for_child(child: &mut std::process::Child, timeout: Duration) -> Result { + let deadline = Instant::now() + timeout; + loop { + if child.try_wait()?.is_some() { + return Ok(true); + } + if Instant::now() >= deadline { + return Ok(false); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } +} + +async fn fetch_otlp_metrics( + client: &Client, + http_port: u16, + clock: &Instant, +) -> Result { + let text = client + .get(format!("http://127.0.0.1:{http_port}/metrics")) + .send() + .await? + .error_for_status()? + .text() + .await?; + Ok(OtlpSnapshot { + captured_monotonic_seconds: clock.elapsed().as_secs_f64(), + values: parse_prometheus_metrics(&text)?, + }) +} + +fn parse_prometheus_metrics(text: &str) -> Result> { + let sample = Regex::new(r"^([a-zA-Z_:][a-zA-Z0-9_:]*)(?:\{[^}]*\})?\s+(\S+)")?; + let mut values = HashMap::new(); + for line in text.lines() { + let Some(captures) = sample.captures(line.trim()) else { + continue; + }; + let (Some(name), Some(value)) = (captures.get(1), captures.get(2)) else { + continue; + }; + let name = name.as_str(); + if !matches!( + name, + OTLP_ROWS | OTLP_FAILURES | OTLP_ELAPSED_SUM | OTLP_ELAPSED_COUNT + ) { + continue; + } + let value: f64 = value.as_str().parse()?; + if !value.is_finite() { + return Err(format!("non-finite Prometheus sample for {name}").into()); + } + *values.entry(name.to_string()).or_default() += value; + } + Ok(values) +} + +fn metric_delta(after: &OtlpSnapshot, before: &OtlpSnapshot, name: &str) -> Result { + let delta = after.values.get(name).copied().unwrap_or_default() + - before.values.get(name).copied().unwrap_or_default(); + if delta < 0.0 { + return Err(format!("metric {name} decreased by {}", -delta).into()); + } + Ok(delta) +} + +fn summarize_otlp_metrics(run: &Value) -> Result { + let snapshots = run + .get("snapshots") + .ok_or("otelgen result has no snapshots")?; + let initial: OtlpSnapshot = serde_json::from_value(snapshots["initial"].clone())?; + let warmed: OtlpSnapshot = serde_json::from_value(snapshots["warmup"].clone())?; + let final_snapshot: OtlpSnapshot = serde_json::from_value(snapshots["final"].clone())?; + let missing_metrics = [OTLP_ROWS, OTLP_ELAPSED_SUM, OTLP_ELAPSED_COUNT] + .into_iter() + .filter(|name| !final_snapshot.values.contains_key(*name)) + .collect::>(); + let accepted_spans = + round_ties_even(metric_delta(&final_snapshot, &initial, OTLP_ROWS)?) as u64; + let measurement_accepted_spans = + round_ties_even(metric_delta(&final_snapshot, &warmed, OTLP_ROWS)?) as u64; + let http_requests = + round_ties_even(metric_delta(&final_snapshot, &warmed, OTLP_ELAPSED_COUNT)?) as u64; + let latency_seconds = metric_delta(&final_snapshot, &warmed, OTLP_ELAPSED_SUM)?; + let measurement_seconds = + final_snapshot.captured_monotonic_seconds - warmed.captured_monotonic_seconds; + Ok(json!({ + "accepted_spans": accepted_spans, + "measurement_accepted_spans": measurement_accepted_spans, + "accepted_spans_per_second": if measurement_seconds > 0.0 { Some(measurement_accepted_spans as f64 / measurement_seconds) } else { None }, + "http_requests": http_requests, + "mean_http_latency_ms": if http_requests > 0 { Some(latency_seconds / http_requests as f64 * 1000.0) } else { None }, + "failure_count": round_ties_even(metric_delta(&final_snapshot, &initial, OTLP_FAILURES)?) as u64, + "measurement_seconds": measurement_seconds, + "missing_metrics": missing_metrics, + })) +} + +async fn poll_otlp_visibility( + client: &Client, + port: u16, + table_name: &str, + database: &str, + expected_rows: u64, + timeout_seconds: u64, +) -> Result { + let sql = format!("SELECT count(*) FROM {}", sql_ident(table_name)); + let deadline = Instant::now() + Duration::from_secs(timeout_seconds); + let mut attempts = 0; + loop { + attempts += 1; + let mut result = http_post_sql(client, port, &sql, database).await; + let observed_rows = extract_count_value(&result); + let row_count_ok = + result["ok"].as_bool().unwrap_or(false) && observed_rows == Some(expected_rows); + result + .as_object_mut() + .ok_or("count result must be an object")? + .extend([ + ("expected_rows".to_string(), json!(expected_rows)), + ("observed_rows".to_string(), json!(observed_rows)), + ("attempts".to_string(), json!(attempts)), + ("row_count_ok".to_string(), json!(row_count_ok)), + ]); + if row_count_ok { + return Ok(result); + } + if Instant::now() >= deadline { + result["ok"] = json!(false); + result["row_count_ok"] = json!(false); + result["error"] = json!(format!( + "expected {expected_rows} rows but observed {observed_rows:?} after {attempts} attempts" + )); + return Ok(result); + } + tokio::time::sleep(Duration::from_millis(500)).await; + } +} + +pub(super) async fn run_finalize_otlp(args: FinalizeOtlpArgs) -> Result<()> { + let (case_path, load) = normalized_otlp_load(&args.fixture_generator, &args.case)?; + let raw_case: toml::Value = toml::from_str(&fs::read_to_string(&case_path)?)?; + let case_metadata = serde_json::to_value( + raw_case + .get("case") + .cloned() + .unwrap_or(toml::Value::Table(toml::map::Map::new())), + )?; + let base: Value = serde_json::from_slice(&fs::read(&args.base_result)?)?; + let candidate: Value = serde_json::from_slice(&fs::read(&args.candidate_result)?)?; + let thresholds = enforce_otlp_thresholds( + &load.thresholds, + base.get("metrics").unwrap_or(&Value::Null), + candidate.get("metrics").unwrap_or(&Value::Null), + ); + let failed = thresholds + .iter() + .any(|threshold| threshold["status"] == "failed") + || [base.get("status"), candidate.get("status")] + .into_iter() + .any(|status| status.and_then(Value::as_str) == Some("failed")); + let report = json!({ + "case_path": case_path, + "case": case_metadata, + "scenario": { "kind": "otlp_trace_load", "load": load }, + "targets": [base, candidate], + "thresholds": thresholds, + "status": if failed { "failed" } else { "ok" }, + }); + let text = format!("{}\n", serde_json::to_string_pretty(&report)?); + fs::write(&args.output, &text)?; + print!("{text}"); + if failed { + std::process::exit(1); + } + Ok(()) +} + +fn enforce_otlp_thresholds( + thresholds: &OtlpThresholds, + base: &Value, + candidate: &Value, +) -> Vec { + let mut results = Vec::new(); + for (target, metrics) in [("base", base), ("candidate", candidate)] { + let failures = value_u64(metrics.get("failure_count")); + results.push(json!({ + "target": target, + "threshold": "max_failure_count", + "status": if failures.is_some_and(|failures| failures <= thresholds.max_failure_count) { "passed" } else { "failed" }, + "actual": failures, + "limit": thresholds.max_failure_count, + })); + } + for (name, base_value, candidate_value, limit, inverse) in [ + ( + "max_candidate_throughput_regression_pct", + value_f64(base.get("accepted_spans_per_second")), + value_f64(candidate.get("accepted_spans_per_second")), + thresholds.max_candidate_throughput_regression_pct, + true, + ), + ( + "max_candidate_mean_latency_regression_pct", + value_f64(base.get("mean_http_latency_ms")), + value_f64(candidate.get("mean_http_latency_ms")), + thresholds.max_candidate_mean_latency_regression_pct, + false, + ), + ] { + let (Some(base_value), Some(candidate_value)) = + (base_value.filter(|value| *value != 0.0), candidate_value) + else { + results.push(json!({ + "threshold": name, + "status": "failed", + "reason": if inverse { "missing or zero throughput" } else { "missing or zero mean latency" }, + "base": base_value, + "candidate": candidate_value, + })); + continue; + }; + let actual = if inverse { + (base_value - candidate_value) / base_value * 100.0 + } else { + (candidate_value - base_value) / base_value * 100.0 + }; + results.push(json!({ + "threshold": name, + "status": if actual <= limit { "passed" } else { "failed" }, + "actual_pct": actual, + "limit_pct": limit, + "base": base_value, + "candidate": candidate_value, + })); + } + results +} + +#[cfg(test)] +mod tests { + use super::*; + + fn otlp_load_for_test() -> OtlpLoad { + serde_json::from_value(json!({ + "database": "public", + "table": "opentelemetry_traces", + "pipeline": "greptime_trace_v1", + "duration_seconds": 120, + "warmup_seconds": 60, + "rate": 50000, + "workers": 4, + "exporter_shards": 4, + "workload": "microservices", + "visibility_timeout_seconds": 30, + "thresholds": { + "max_candidate_throughput_regression_pct": 20.0, + "max_candidate_mean_latency_regression_pct": 20.0, + "max_failure_count": 0 + } + })) + .unwrap() + } + + #[test] + fn parses_labeled_otlp_metrics_and_rejects_non_finite_samples() { + let metrics = parse_prometheus_metrics( + "greptime_frontend_otlp_traces_rows 10\n\ + greptime_frontend_otlp_traces_failure_count{kind=\"a\"} 1\n\ + greptime_frontend_otlp_traces_failure_count{kind=\"b\"} 2\n", + ) + .unwrap(); + assert_eq!(metrics[OTLP_ROWS], 10.0); + assert_eq!(metrics[OTLP_FAILURES], 3.0); + assert!(parse_prometheus_metrics("greptime_frontend_otlp_traces_rows NaN").is_err()); + let before = OtlpSnapshot { + captured_monotonic_seconds: 0.0, + values: HashMap::from([(OTLP_ROWS.to_string(), 2.0)]), + }; + let after = OtlpSnapshot { + captured_monotonic_seconds: 1.0, + values: HashMap::from([(OTLP_ROWS.to_string(), 1.0)]), + }; + assert!(metric_delta(&after, &before, OTLP_ROWS).is_err()); + } + + #[test] + fn summarizes_otlp_deltas_and_builds_exact_command() { + let snapshot = |captured: f64, values: Vec<(String, f64)>| OtlpSnapshot { + captured_monotonic_seconds: captured, + values: values.into_iter().collect(), + }; + let initial = snapshot(0.0, vec![(OTLP_ROWS.to_string(), 10.0)]); + let warmed = snapshot( + 5.0, + vec![ + (OTLP_ROWS.to_string(), 110.0), + (OTLP_ELAPSED_SUM.to_string(), 1.0), + (OTLP_ELAPSED_COUNT.to_string(), 10.0), + ], + ); + let final_snapshot = snapshot( + 15.0, + vec![ + (OTLP_ROWS.to_string(), 310.0), + (OTLP_FAILURES.to_string(), 3.0), + (OTLP_ELAPSED_SUM.to_string(), 3.0), + (OTLP_ELAPSED_COUNT.to_string(), 30.0), + ], + ); + let metrics = summarize_otlp_metrics(&json!({ + "snapshots": { "initial": initial, "warmup": warmed, "final": final_snapshot } + })) + .unwrap(); + assert_eq!(metrics["accepted_spans"], 300); + assert_eq!(metrics["measurement_accepted_spans"], 200); + assert_eq!(metrics["accepted_spans_per_second"], 20.0); + assert_eq!(metrics["http_requests"], 20); + assert_eq!(metrics["mean_http_latency_ms"], 100.0); + assert_eq!(metrics["failure_count"], 3); + let missing = summarize_otlp_metrics(&json!({ + "snapshots": { + "initial": snapshot(0.0, vec![(OTLP_ROWS.to_string(), 0.0)]), + "warmup": snapshot(1.0, vec![(OTLP_ROWS.to_string(), 1.0)]), + "final": snapshot(2.0, vec![(OTLP_ROWS.to_string(), 2.0)]) + } + })) + .unwrap(); + assert_eq!( + missing["missing_metrics"], + json!([OTLP_ELAPSED_SUM, OTLP_ELAPSED_COUNT]) + ); + let command = otelgen_command(Path::new("/bin/otelgen"), 4000, &otlp_load_for_test()); + assert_eq!( + command, + vec![ + "/bin/otelgen", + "--protocol", + "http", + "--otel-exporter-otlp-endpoint", + "127.0.0.1:4000", + "--otel-exporter-otlp-url-path", + "/v1/otlp/v1/traces", + "--header", + "x-greptime-pipeline-name=greptime_trace_v1", + "--header", + "x-greptime-db-name=public", + "--header", + "x-greptime-trace-table-name=opentelemetry_traces", + "--log-level", + "error", + "--insecure", + "--duration", + "120", + "--rate", + "50000", + "traces", + "multi", + "--workers", + "4", + "--scenarios", + "microservices", + "--exporter-shards", + "4", + ] + ); + } + + #[test] + fn otlp_thresholds_handle_zero_metrics_and_pass_fail_cases() { + let load = otlp_load_for_test(); + let base = json!({ "failure_count": 0, "accepted_spans_per_second": 20.0, "mean_http_latency_ms": 100.0 }); + let candidate = json!({ "failure_count": 0, "accepted_spans_per_second": 18.0, "mean_http_latency_ms": 110.0 }); + let results = enforce_otlp_thresholds(&load.thresholds, &base, &candidate); + assert!(results.iter().all(|result| result["status"] == "passed")); + let failed = enforce_otlp_thresholds( + &load.thresholds, + &json!({ "failure_count": 0, "accepted_spans_per_second": 0.0, "mean_http_latency_ms": 100.0 }), + &json!({ "failure_count": 3, "accepted_spans_per_second": 20.0, "mean_http_latency_ms": 130.0 }), + ); + assert!( + failed + .iter() + .any(|result| result["threshold"] == "max_failure_count" + && result["target"] == "candidate" + && result["status"] == "failed") + ); + assert!(failed.iter().any(|result| result["threshold"] + == "max_candidate_throughput_regression_pct" + && result["reason"] == "missing or zero throughput")); + assert!(failed.iter().any(|result| result["threshold"] + == "max_candidate_mean_latency_regression_pct" + && result["status"] == "failed")); + } +} diff --git a/src/cmd/src/bin/query_regression_runner/plan.rs b/src/cmd/src/bin/query_regression_runner/plan.rs new file mode 100644 index 00000000000..4fcc5135cd8 --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/plan.rs @@ -0,0 +1,122 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashSet; +use std::path::{Path, PathBuf}; +use std::process::Command; + +use serde_json::Value; + +use crate::query_regression_runner::Result; +use crate::query_regression_runner::model::{ + Layout, OtlpLoad, Query, RemoteWrite, Scenario, Table, +}; + +pub(super) fn load_plan(generator: &PathBuf, case_path: &PathBuf) -> Result { + let output = Command::new(generator) + .args(["plan", "--case"]) + .arg(case_path) + .output()?; + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + return Err(format!("query_perf_fixture plan failed: {:.2000}", stderr).into()); + } + Ok(serde_json::from_slice(&output.stdout)?) +} + +pub(super) fn normalize_scenario(scenario: Scenario) -> Result<(Vec
, Vec)> { + match scenario { + Scenario::DirectReadableSst { + tables, + layout, + queries, + } => Ok((validate_direct_tables(tables, layout)?, queries)), + Scenario::PromRemoteWriteThenQuery { + remote_write, + queries, + } => Ok(( + vec![Table { + database: remote_write.database, + name: remote_write.metric, + engine: "metric".to_string(), + columns: vec![], + primary_key: vec![], + time_index: None, + append_mode: None, + sst_format: None, + validate_show_create_engine: false, + }], + queries, + )), + Scenario::OtlpTraceLoad { .. } => { + Err("measure requires a query scenario, not otlp_trace_load".into()) + } + } +} + +pub(super) fn validate_direct_tables(tables: Vec
, layout: Layout) -> Result> { + if tables.is_empty() || layout.regions != 1 { + return Err("runner supports one or more tables and exactly one region per table".into()); + } + let mut pairs = HashSet::new(); + let mut names = HashSet::new(); + for table in &tables { + if !pairs.insert((&table.database, &table.name)) { + return Err("duplicate (database, name) table entries are not supported".into()); + } + if !names.insert(&table.name) { + return Err("duplicate table names are not supported".into()); + } + } + Ok(tables) +} + +pub(super) fn normalized_remote_write( + generator: &PathBuf, + case: &Path, +) -> Result<(PathBuf, RemoteWrite)> { + let case_path = case.canonicalize()?; + let plan = load_plan(generator, &case_path)?; + let scenario = plan + .get("scenario") + .cloned() + .ok_or("fixture plan has no scenario")?; + match serde_json::from_value(scenario)? { + Scenario::PromRemoteWriteThenQuery { remote_write, .. } => Ok((case_path, remote_write)), + Scenario::DirectReadableSst { .. } => { + Err("remote command requires scenario kind prom_remote_write_then_query".into()) + } + Scenario::OtlpTraceLoad { .. } => { + Err("remote command requires scenario kind prom_remote_write_then_query".into()) + } + } +} + +pub(super) fn normalized_otlp_load( + generator: &PathBuf, + case: &Path, +) -> Result<(PathBuf, OtlpLoad)> { + let case_path = case.canonicalize()?; + let plan = load_plan(generator, &case_path)?; + let scenario = plan + .get("scenario") + .cloned() + .ok_or("fixture plan has no scenario")?; + match serde_json::from_value(scenario)? { + Scenario::OtlpTraceLoad { load } => Ok((case_path, load.load)), + Scenario::DirectReadableSst { .. } | Scenario::PromRemoteWriteThenQuery { .. } => { + Err("OTLP command requires scenario kind otlp_trace_load".into()) + } + } +} diff --git a/src/cmd/src/bin/query_regression_runner/remote/ingest.rs b/src/cmd/src/bin/query_regression_runner/remote/ingest.rs new file mode 100644 index 00000000000..f563402afcd --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/remote/ingest.rs @@ -0,0 +1,505 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fs; +use std::path::Path; +use std::process::Command; +use std::time::{Duration, Instant}; + +use reqwest::Client; +use serde_json::{Value, json}; + +use crate::query_regression_runner::model::{PromStore, RemoteWrite}; +use crate::query_regression_runner::plan::normalized_remote_write; +use crate::query_regression_runner::sql::{ + extract_count_value, http_post_sql, sql_ident, sql_string, value_f64, value_u64, +}; +use crate::query_regression_runner::{PrepareRemoteArgs, RenderRemoteConfigArgs, Result}; + +pub(super) async fn run_render_remote_config(args: RenderRemoteConfigArgs) -> Result<()> { + let (_, remote) = normalized_remote_write(&args.fixture_generator, &args.case)?; + fs::write(args.output, frontend_prom_config(&remote.prom_store)?)?; + Ok(()) +} + +fn frontend_prom_config(prom: &PromStore) -> Result { + Ok(format!( + "[prom_store]\nenable = true\nwith_metric_engine = true\npending_rows_flush_interval = {}\nmax_batch_rows = {}\nmax_concurrent_flushes = {}\nworker_channel_capacity = {}\nmax_inflight_requests = {}\n", + serde_json::to_string(&prom.pending_rows_flush_interval)?, + prom.max_batch_rows, + prom.max_concurrent_flushes, + prom.worker_channel_capacity, + prom.max_inflight_requests, + )) +} + +pub(super) async fn run_prepare_remote(args: PrepareRemoteArgs) -> Result<()> { + if !args.http_timeout.is_finite() || args.http_timeout < 0.0 { + return Err("--http-timeout must be a non-negative finite number".into()); + } + let (case_path, remote) = normalized_remote_write(&args.fixture_generator, &args.case)?; + let client = Client::builder() + .timeout(Duration::from_secs_f64(args.http_timeout)) + .build()?; + let base = prepare_remote_target( + "base", + args.base_http_port, + &args.fixture_generator, + &remote, + &client, + ) + .await?; + let candidate = prepare_remote_target( + "candidate", + args.candidate_http_port, + &args.fixture_generator, + &remote, + &client, + ) + .await?; + let report = json!({ + "case_path": case_path, + "scenario": "prom_remote_write_then_query", + "base": base, + "candidate": candidate, + "status": "ok", + }); + let text = format!("{}\n", serde_json::to_string_pretty(&report)?); + if let Some(path) = args.output { + fs::write(path, &text)?; + } + print!("{text}"); + Ok(()) +} + +async fn prepare_remote_target( + name: &str, + port: u16, + generator: &Path, + remote: &RemoteWrite, + client: &Client, +) -> Result { + let create_database = http_post_sql( + client, + port, + &format!( + "CREATE DATABASE IF NOT EXISTS {}", + sql_ident(&remote.database) + ), + "public", + ) + .await; + if !create_database["ok"].as_bool().unwrap_or(false) { + return Err(format!( + "CREATE DATABASE {} failed for {name}: {create_database}", + remote.database + ) + .into()); + } + let (remote_write, flushes) = ingest_remote_write(generator, port, remote, client).await?; + let expected_rows = remote + .series_count + .checked_mul(remote.samples_per_series) + .ok_or("expected remote-write row count overflows u64")?; + let visibility = poll_expected_count( + client, + port, + &remote.metric, + &remote.database, + expected_rows, + remote.visibility_timeout_seconds, + ) + .await?; + Ok(json!({ + "name": name, + "create_database": create_database, + "remote_write": remote_write, + "flushes": flushes, + "visibility": visibility, + "status": "ok", + })) +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct SampleChunk { + index: u64, + offset: u64, + samples_per_series: u64, + start_unix_millis: i64, +} + +async fn ingest_remote_write( + generator: &Path, + port: u16, + remote: &RemoteWrite, + client: &Client, +) -> Result<(Value, Vec)> { + let Some(chunks) = sample_chunks(remote)? else { + let remote_write = run_remote_write(generator, port, remote, None)?; + let flush = flush_remote_table(client, port, remote, "final", None).await?; + return Ok((remote_write, vec![flush])); + }; + let mut results = Vec::with_capacity(chunks.len()); + let mut flushes = Vec::new(); + let scheduled_flushes = scheduled_flushes(&chunks, remote.flush_every_sample_chunks); + for chunk in &chunks { + let mut result = run_remote_write(generator, port, remote, Some(chunk))?; + result + .as_object_mut() + .ok_or("remote-write result must be an object")? + .extend([ + ("sample_offset".to_string(), json!(chunk.offset)), + ( + "samples_per_series".to_string(), + json!(chunk.samples_per_series), + ), + ("chunk_index".to_string(), json!(chunk.index)), + ]); + results.push(result); + if let Some((_, reason)) = scheduled_flushes + .iter() + .find(|(index, _)| *index == chunk.index) + { + flushes + .push(flush_remote_table(client, port, remote, reason, Some(chunk.index)).await?); + } + } + Ok(( + json!({ + "status": "ok", + "mode": "sample-chunked", + "sample_chunk_size": remote.sample_chunk_size, + "flush_every_sample_chunks": remote.flush_every_sample_chunks, + "chunks": results, + "aggregate": summarize_remote_chunks(&results), + }), + flushes, + )) +} + +fn sample_chunks(remote: &RemoteWrite) -> Result>> { + let Some(chunk_samples) = remote.sample_chunk_size else { + return Ok(None); + }; + if chunk_samples == 0 { + return Err("scenario.remote_write.sample_chunk_size must be positive".into()); + } + if remote.flush_every_sample_chunks == 0 { + return Err("scenario.remote_write.flush_every_sample_chunks must be positive".into()); + } + let mut chunks = Vec::new(); + let mut offset = 0; + while offset < remote.samples_per_series { + let samples_per_series = chunk_samples.min(remote.samples_per_series - offset); + chunks.push(SampleChunk { + index: chunks.len() as u64 + 1, + offset, + samples_per_series, + start_unix_millis: remote.start_unix_millis + offset as i64 * remote.step_millis, + }); + offset += samples_per_series; + } + Ok(Some(chunks)) +} + +fn scheduled_flushes(chunks: &[SampleChunk], flush_every: u64) -> Vec<(u64, &'static str)> { + let mut scheduled = chunks + .iter() + .filter(|chunk| chunk.index % flush_every == 0) + .map(|chunk| (chunk.index, "periodic")) + .collect::>(); + if let Some(chunk) = chunks.last() + && chunk.index % flush_every != 0 + { + scheduled.push((chunk.index, "final")); + } + scheduled +} + +fn summarize_remote_chunks(chunks: &[Value]) -> Value { + let mut rows = 0; + let mut samples_written = 0; + let mut batches = 0; + let mut elapsed_seconds = 0.0; + for chunk in chunks { + let summary = chunk.get("summary").unwrap_or(&Value::Null); + let chunk_rows = value_u64(summary.get("rows")).unwrap_or(0); + rows += chunk_rows; + samples_written += value_u64(summary.get("samples_written")).unwrap_or(chunk_rows); + batches += value_u64(summary.get("batches")).unwrap_or(0); + elapsed_seconds += value_f64(summary.get("elapsed_seconds")) + .or_else(|| value_f64(chunk.get("elapsed_seconds"))) + .unwrap_or(0.0); + } + json!({ "rows": rows, "samples_written": samples_written, "batches": batches, "elapsed_seconds": elapsed_seconds }) +} + +fn run_remote_write( + generator: &Path, + port: u16, + remote: &RemoteWrite, + chunk: Option<&SampleChunk>, +) -> Result { + let command = remote_write_command(generator, port, remote, chunk); + let started = Instant::now(); + let output = Command::new(generator).args(&command[1..]).output()?; + let elapsed_seconds = started.elapsed().as_secs_f64(); + let stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + if !output.status.success() { + return Err(format!( + "remote-write generator failed: command={command:?}, returncode={:?}, elapsed_seconds={elapsed_seconds:.3}, stderr={:.2000}", + output.status.code(), stderr + ) + .into()); + } + let mut result = json!({ + "status": "ok", + "command": command, + "returncode": output.status.code(), + "elapsed_seconds": elapsed_seconds, + "stdout": stdout, + "stderr": stderr, + }); + match serde_json::from_str(result["stdout"].as_str().unwrap_or_default()) { + Ok(summary) => result["summary"] = summary, + Err(_) => result["summary_parse_error"] = result["stdout"].clone(), + } + Ok(result) +} + +fn remote_write_command( + generator: &Path, + port: u16, + remote: &RemoteWrite, + chunk: Option<&SampleChunk>, +) -> Vec { + let (samples_per_series, start_unix_millis, sample_offset, total_samples_per_series) = chunk + .map(|chunk| { + ( + chunk.samples_per_series, + chunk.start_unix_millis, + Some(chunk.offset), + Some(remote.samples_per_series), + ) + }) + .unwrap_or(( + remote.samples_per_series, + remote.start_unix_millis, + None, + None, + )); + let mut command = vec![ + generator.to_string_lossy().to_string(), + "prom-remote-write".to_string(), + "--endpoint".to_string(), + format!("http://127.0.0.1:{port}/v1/prometheus/write"), + "--database".to_string(), + remote.database.clone(), + "--metric".to_string(), + remote.metric.clone(), + "--physical-table".to_string(), + remote.physical_table.clone(), + "--series-count".to_string(), + remote.series_count.to_string(), + "--samples-per-series".to_string(), + samples_per_series.to_string(), + "--start-unix-millis".to_string(), + start_unix_millis.to_string(), + "--step-millis".to_string(), + remote.step_millis.to_string(), + "--chunk-series-count".to_string(), + remote.chunk_series_count.to_string(), + "--timeout-seconds".to_string(), + remote.timeout_seconds.to_string(), + "--value-pattern".to_string(), + remote.value.pattern.clone(), + "--value-base".to_string(), + json_number(remote.value.base), + "--value-step".to_string(), + json_number(remote.value.step), + "--value-cardinality".to_string(), + remote.value.cardinality.to_string(), + "--value-seed".to_string(), + remote.value.seed.to_string(), + "--value-run-length".to_string(), + remote.value.run_length.to_string(), + "--value-stall-every".to_string(), + remote.value.stall_every.to_string(), + "--value-stall-length".to_string(), + remote.value.stall_length.to_string(), + "--value-mixed-every".to_string(), + remote.value.mixed_every.to_string(), + ]; + if let Some(sample_offset) = sample_offset { + command.extend([ + "--value-sample-offset".to_string(), + sample_offset.to_string(), + ]); + } + if let Some(total_samples_per_series) = total_samples_per_series { + command.extend([ + "--value-total-samples-per-series".to_string(), + total_samples_per_series.to_string(), + ]); + } + command +} + +async fn flush_remote_table( + client: &Client, + port: u16, + remote: &RemoteWrite, + reason: &str, + chunk_index: Option, +) -> Result { + let mut result = http_post_sql( + client, + port, + &format!("ADMIN FLUSH_TABLE({})", sql_string(&remote.physical_table)), + &remote.database, + ) + .await; + if !result["ok"].as_bool().unwrap_or(false) { + return Err(format!( + "ADMIN FLUSH_TABLE {} failed: {result}", + remote.physical_table + ) + .into()); + } + result + .as_object_mut() + .ok_or("flush result must be an object")? + .extend([ + ("physical_table".to_string(), json!(remote.physical_table)), + ("reason".to_string(), json!(reason)), + ("chunk_index".to_string(), json!(chunk_index)), + ]); + Ok(result) +} + +async fn poll_expected_count( + client: &Client, + port: u16, + table_name: &str, + database: &str, + expected_rows: u64, + visibility_timeout_seconds: u64, +) -> Result { + let sql = format!("SELECT count(*) FROM {}", sql_ident(table_name)); + let deadline = Instant::now() + Duration::from_secs(visibility_timeout_seconds); + let mut attempts = 0; + loop { + attempts += 1; + let mut result = http_post_sql(client, port, &sql, database).await; + let observed_rows = extract_count_value(&result); + let row_count_ok = + result["ok"].as_bool().unwrap_or(false) && observed_rows == Some(expected_rows); + result + .as_object_mut() + .ok_or("count result must be an object")? + .extend([ + ("expected_rows".to_string(), json!(expected_rows)), + ("observed_rows".to_string(), json!(observed_rows)), + ("attempts".to_string(), json!(attempts)), + ("row_count_ok".to_string(), json!(row_count_ok)), + ]); + if row_count_ok { + return Ok(result); + } + if Instant::now() >= deadline { + return Err(format!( + "expected {expected_rows} rows but observed {observed_rows:?} after {attempts} attempts" + ) + .into()); + } + tokio::time::sleep(Duration::from_millis(500)).await; + } +} + +fn json_number(value: f64) -> String { + serde_json::to_string(&value).unwrap_or_else(|_| value.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::query_regression_runner::model::{PromStore, RemoteValue}; + + #[test] + fn schedules_remote_sample_chunks_and_flushes() { + let remote = RemoteWrite { + database: "public".to_string(), + metric: "metric".to_string(), + physical_table: "physical".to_string(), + series_count: 2, + samples_per_series: 5, + start_unix_millis: 100, + step_millis: 10, + chunk_series_count: 1, + timeout_seconds: 60, + sample_chunk_size: Some(2), + flush_every_sample_chunks: 2, + visibility_timeout_seconds: 30, + prom_store: PromStore { + pending_rows_flush_interval: "1s".to_string(), + max_batch_rows: 1, + max_concurrent_flushes: 1, + worker_channel_capacity: 1, + max_inflight_requests: 1, + }, + value: RemoteValue { + pattern: "linear".to_string(), + base: 0.0, + step: 1.0, + cardinality: 1, + seed: 0, + run_length: 1, + stall_every: 0, + stall_length: 0, + mixed_every: 0, + }, + storage: None, + read_bench: None, + }; + let chunks = sample_chunks(&remote).unwrap().unwrap(); + assert_eq!( + chunks, + vec![ + SampleChunk { + index: 1, + offset: 0, + samples_per_series: 2, + start_unix_millis: 100, + }, + SampleChunk { + index: 2, + offset: 2, + samples_per_series: 2, + start_unix_millis: 120, + }, + SampleChunk { + index: 3, + offset: 4, + samples_per_series: 1, + start_unix_millis: 140, + }, + ] + ); + assert_eq!( + scheduled_flushes(&chunks, remote.flush_every_sample_chunks), + vec![(2, "periodic"), (3, "final")] + ); + } +} diff --git a/src/cmd/src/bin/query_regression_runner/remote/mod.rs b/src/cmd/src/bin/query_regression_runner/remote/mod.rs new file mode 100644 index 00000000000..946b00b61de --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/remote/mod.rs @@ -0,0 +1,183 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fs; +use std::path::{Path, PathBuf}; + +use serde_json::{Value, json}; + +use self::read_bench::run_read_bench; +use self::storage::{enforce_storage_thresholds, run_storage_inspection}; +use crate::query_regression_runner::model::DestinationConfig; +use crate::query_regression_runner::plan::normalized_remote_write; +use crate::query_regression_runner::{ + FinalizeRemoteArgs, PrepareRemoteArgs, RenderRemoteConfigArgs, Result, +}; + +mod ingest; +mod read_bench; +mod storage; + +pub(super) async fn run_render_remote_config(args: RenderRemoteConfigArgs) -> Result<()> { + ingest::run_render_remote_config(args).await +} + +pub(super) async fn run_prepare_remote(args: PrepareRemoteArgs) -> Result<()> { + ingest::run_prepare_remote(args).await +} + +pub(super) async fn run_finalize_remote(args: FinalizeRemoteArgs) -> Result<()> { + let (_, remote) = normalized_remote_write(&args.fixture_generator, &args.case)?; + let storage = remote.storage.as_ref().filter(|storage| storage.inspect); + let mut report: Value = serde_json::from_slice(&fs::read(&args.report)?)?; + let bench_root = args.report.parent().unwrap_or_else(|| Path::new(".")); + let mut inspections = Vec::new(); + let target_failed = { + let targets = report + .get_mut("targets") + .and_then(Value::as_array_mut) + .ok_or("report has no targets array")?; + for (name, data_home, destination) in [ + ( + "base", + args.base_data_home.as_deref(), + args.base_destination.as_deref(), + ), + ( + "candidate", + args.candidate_data_home.as_deref(), + args.candidate_destination.as_deref(), + ), + ] { + let (data_home, destination) = match (data_home, destination) { + (Some(data_home), None) => (data_home.to_path_buf(), None), + (None, Some(path)) => { + let destination: DestinationConfig = + toml::from_str(&fs::read_to_string(path)?)?; + ( + PathBuf::from(destination.data_home), + Some(path.to_path_buf()), + ) + } + (Some(_), Some(_)) => { + return Err(format!( + "{name}: --{name}-data-home and --{name}-destination are mutually exclusive" + ) + .into()); + } + (None, None) => { + return Err(format!( + "{name}: one of --{name}-data-home or --{name}-destination is required" + ) + .into()); + } + }; + let target = targets + .iter_mut() + .find(|target| target.get("name").and_then(Value::as_str) == Some(name)) + .ok_or_else(|| format!("report has no {name} target"))?; + let Some(storage) = storage else { + target + .as_object_mut() + .ok_or("report target must be an object")? + .insert( + "read_bench".to_string(), + json!({"status": "skipped", "reason": "storage inspection disabled"}), + ); + continue; + }; + let inspection = run_storage_inspection( + &args.fixture_generator, + &data_home, + destination.as_deref(), + storage, + )?; + let bench_dir = bench_root.join(name).join("read_bench"); + let read_bench = run_read_bench( + &args.candidate_bin, + &data_home, + &bench_dir, + remote.read_bench.as_ref(), + &inspection, + )?; + let inspection_failed = inspection["status"] == "failed"; + let bench_failed = read_bench["status"] == "failed"; + let target = target + .as_object_mut() + .ok_or("report target must be an object")?; + target.insert("storage_inspection".to_string(), inspection.clone()); + target.insert("read_bench".to_string(), read_bench); + if inspection_failed || bench_failed { + target.insert("status".to_string(), json!("failed")); + } + inspections.push(inspection); + } + targets.iter().any(|target| { + target["status"] == "failed" + || target["storage_inspection"]["status"] == "failed" + || target["read_bench"]["status"] == "failed" + }) + }; + let storage_thresholds = storage + .map(|storage| { + enforce_storage_thresholds( + storage, + inspections.first().unwrap_or(&Value::Null), + inspections.get(1).unwrap_or(&Value::Null), + ) + }) + .unwrap_or_default(); + let threshold_failed = { + let thresholds = report + .get_mut("thresholds") + .and_then(Value::as_array_mut) + .ok_or("report has no thresholds array")?; + thresholds.retain(|threshold| !is_storage_threshold_entry(threshold)); + thresholds.extend(storage_thresholds); + thresholds + .iter() + .any(|threshold| threshold["status"] == "failed") + }; + let failed = target_failed || threshold_failed; + report["status"] = json!(if failed { "failed" } else { "ok" }); + let text = format!("{}\n", serde_json::to_string_pretty(&report)?); + fs::write(&args.report, &text)?; + print!("{text}"); + if failed { + std::process::exit(1); + } + Ok(()) +} + +fn is_storage_threshold_entry(threshold: &Value) -> bool { + let Some(name) = threshold.get("threshold").and_then(Value::as_str) else { + return false; + }; + matches!( + name, + "max_candidate_total_file_size_regression_pct" + | "max_candidate_column_compressed_size_regression_pct" + | "max_candidate_column_uncompressed_size_regression_pct" + ) || (threshold.get("target").is_some() + && matches!( + name, + "min_files" + | "min_files_with_column" + | "max_total_file_size_bytes" + | "max_column_compressed_size_bytes" + | "max_column_uncompressed_size_bytes" + | "require_encodings" + | "forbid_encodings" + )) +} diff --git a/src/cmd/src/bin/query_regression_runner/remote/read_bench.rs b/src/cmd/src/bin/query_regression_runner/remote/read_bench.rs new file mode 100644 index 00000000000..6b20a31c181 --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/remote/read_bench.rs @@ -0,0 +1,363 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fs; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::time::Instant; + +use regex::Regex; +use serde_json::{Value, json}; + +use crate::query_regression_runner::Result; +use crate::query_regression_runner::measure::median; +use crate::query_regression_runner::model::ReadBenchConfig; + +#[derive(Clone, Debug)] +struct BenchTarget { + relative_path: String, + table_dir: String, + region_id: String, + path_type: String, + file_id: String, +} + +pub(super) fn run_read_bench( + candidate_bin: &Path, + data_home: &Path, + bench_dir: &Path, + read_bench: Option<&ReadBenchConfig>, + inspection: &Value, +) -> Result { + let Some(read_bench) = read_bench.filter(|config| config.enabled) else { + return Ok(json!({ "status": "skipped", "reason": "read_bench disabled" })); + }; + if !read_bench.parquetbench && !read_bench.scanbench { + return Ok(json!({ "status": "skipped", "reason": "parquetbench and scanbench disabled" })); + } + let report = inspection + .get("summary") + .filter(|summary| summary.is_object()); + let files = report + .and_then(|report| report.get("files")) + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let root = report + .and_then(|report| report.get("root")) + .and_then(Value::as_str) + .map(PathBuf::from) + .or_else(|| { + inspection + .get("root") + .and_then(Value::as_str) + .map(PathBuf::from) + }) + .ok_or("storage inspection has no root")?; + let mut targets = files + .iter() + .filter(|file| { + file.get("relative_path") + .and_then(Value::as_str) + .is_some_and(|path| path.ends_with(".parquet")) + && file + .get("columns") + .and_then(Value::as_array) + .is_some_and(|columns| !columns.is_empty()) + }) + .filter_map(|file| inspected_bench_target(data_home, &root, file).transpose()) + .collect::>>()?; + if let Some(max_files) = read_bench.max_files { + targets.truncate(max_files); + } + if targets.is_empty() { + return Ok( + json!({ "status": "failed", "reason": "no inspected data SST files available for read_bench" }), + ); + } + fs::create_dir_all(bench_dir)?; + let config_toml = bench_dir.join("bench.toml"); + let scan_json = bench_dir.join("scan.json"); + fs::write( + &config_toml, + format!( + "[storage]\ndata_home = \"{}\"\ntype = \"File\"\n\n[[region_engine]]\n[region_engine.mito]\n", + data_home.display() + ), + )?; + fs::write( + &scan_json, + format!( + "{}\n", + serde_json::to_string_pretty(&json!({ "projection_names": read_bench.projection }))? + ), + )?; + let mut parquet_runs = Vec::new(); + if read_bench.parquetbench { + for target in &targets { + let command = vec![ + candidate_bin.to_string_lossy().to_string(), + "datanode".to_string(), + "parquetbench".to_string(), + "--config".to_string(), + config_toml.to_string_lossy().to_string(), + "--region-id".to_string(), + target.region_id.clone(), + "--table-dir".to_string(), + target.table_dir.clone(), + "--file-id".to_string(), + target.file_id.clone(), + "--scan-config".to_string(), + scan_json.to_string_lossy().to_string(), + "--path-type".to_string(), + target.path_type.clone(), + "--iterations".to_string(), + read_bench.iterations.to_string(), + "--reader".to_string(), + read_bench.parquet_reader.clone(), + ]; + parquet_runs.push(run_bench_command(command, bench_target_value(target))?); + } + } + let mut scan_runs = Vec::new(); + if read_bench.scanbench { + for (table_dir, region_id, path_type, files) in group_scan_paths(&targets) { + let command = vec![ + candidate_bin.to_string_lossy().to_string(), + "datanode".to_string(), + "scanbench".to_string(), + "--config".to_string(), + config_toml.to_string_lossy().to_string(), + "--region-id".to_string(), + region_id.clone(), + "--table-dir".to_string(), + table_dir.clone(), + "--scan-config".to_string(), + scan_json.to_string_lossy().to_string(), + "--path-type".to_string(), + path_type.clone(), + "--scanner".to_string(), + read_bench.scan_scanner.clone(), + "--parallelism".to_string(), + read_bench.parallelism.to_string(), + "--iterations".to_string(), + read_bench.iterations.to_string(), + ]; + scan_runs.push(run_bench_command(command, json!({ "table_dir": table_dir, "region_id": region_id, "path_type": path_type, "files": files }))?); + } + } + let failed = parquet_runs + .iter() + .chain(&scan_runs) + .any(|run| run["status"] == "failed"); + Ok(json!({ + "status": if failed { "failed" } else { "ok" }, + "config_path": config_toml, + "scan_config_path": scan_json, + "parquetbench": parquet_runs, + "scanbench": scan_runs, + "aggregate": { + "parquetbench_median_average_ms": bench_median(&parquet_runs), + "scanbench_median_average_ms": bench_median(&scan_runs), + }, + })) +} + +fn inspected_bench_target( + data_home: &Path, + root: &Path, + file: &Value, +) -> Result> { + let relative_path = file + .get("relative_path") + .and_then(Value::as_str) + .unwrap_or_default(); + if relative_path.is_empty() { + return Ok(None); + } + let root_suffix = root.strip_prefix(data_home).map_err(|_| { + format!( + "storage inspection root {} is not under datanode data home {}", + root.display(), + data_home.display() + ) + })?; + let full_relative = root_suffix.join(relative_path); + let parts = full_relative + .components() + .map(|component| component.as_os_str().to_string_lossy().to_string()) + .collect::>(); + if !parts.last().is_some_and(|part| part.ends_with(".parquet")) { + return Ok(None); + } + let Some((region_index, table_id, region_seq)) = + parts.iter().enumerate().find_map(|(index, part)| { + parse_region_dir(part).map(|(table_id, region_seq)| (index, table_id, region_seq)) + }) + else { + return Ok(None); + }; + if region_index == 0 { + return Ok(None); + } + let file_index = parts.len() - 1; + let path_type = if region_index + 1 < file_index + && matches!(parts[region_index + 1].as_str(), "data" | "metadata") + { + parts[region_index + 1].clone() + } else { + "bare".to_string() + }; + if path_type == "metadata" { + return Ok(None); + } + let file_id = parts[file_index].trim_end_matches(".parquet").to_string(); + Ok(Some(BenchTarget { + relative_path: full_relative.to_string_lossy().to_string(), + table_dir: format!("{}/", parts[..region_index].join("/")), + region_id: format!("{table_id}:{region_seq}"), + path_type, + file_id, + })) +} + +fn parse_region_dir(part: &str) -> Option<(u64, u64)> { + let (table_id, region_seq) = part.split_once('_')?; + (region_seq.len() == 10 + && table_id.chars().all(|character| character.is_ascii_digit()) + && region_seq + .chars() + .all(|character| character.is_ascii_digit())) + .then(|| Some((table_id.parse().ok()?, region_seq.parse().ok()?)))? +} + +fn bench_target_value(target: &BenchTarget) -> Value { + json!({ + "relative_path": target.relative_path, + "table_dir": target.table_dir, + "region_id": target.region_id, + "path_type": target.path_type, + "file_id": target.file_id, + }) +} + +fn group_scan_paths(targets: &[BenchTarget]) -> Vec<(String, String, String, Vec)> { + let mut groups: Vec<(String, String, String, Vec)> = Vec::new(); + for target in targets { + if let Some((_, _, _, paths)) = + groups + .iter_mut() + .find(|(table_dir, region_id, path_type, _)| { + table_dir == &target.table_dir + && region_id == &target.region_id + && path_type == &target.path_type + }) + { + paths.push(target.relative_path.clone()); + } else { + groups.push(( + target.table_dir.clone(), + target.region_id.clone(), + target.path_type.clone(), + vec![target.relative_path.clone()], + )); + } + } + groups +} + +fn run_bench_command(command: Vec, mut run: Value) -> Result { + let started = Instant::now(); + let output = Command::new(&command[0]).args(&command[1..]).output()?; + let elapsed_seconds = started.elapsed().as_secs_f64(); + let stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + let run = run.as_object_mut().ok_or("bench run must be an object")?; + run.extend([ + ("command".to_string(), json!(command)), + ("returncode".to_string(), json!(output.status.code())), + ("elapsed_seconds".to_string(), json!(elapsed_seconds)), + ("stdout".to_string(), json!(stdout)), + ("stderr".to_string(), json!(stderr)), + ( + "status".to_string(), + json!(if output.status.success() { + "ok" + } else { + "failed" + }), + ), + ]); + if let Some(average_ms) = parse_average_duration(run["stdout"].as_str().unwrap_or_default()) { + run.insert("average_ms".to_string(), json!(average_ms)); + } + Ok(Value::Object(run.clone())) +} + +fn parse_average_duration(stdout: &str) -> Option { + Regex::new(r"(?i)Average duration[^0-9]*([0-9.]+)\s*ms") + .ok()? + .captures(stdout)? + .get(1)? + .as_str() + .parse() + .ok() +} + +fn bench_median(runs: &[Value]) -> Option { + let values = runs + .iter() + .filter_map(|run| run.get("average_ms").and_then(Value::as_f64)) + .collect::>(); + (!values.is_empty()).then(|| median(&values)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_average_duration_and_groups_scan_paths() { + assert_eq!( + parse_average_duration("work\nAverage duration: 12.5 ms\n"), + Some(12.5) + ); + assert_eq!(parse_average_duration("completed"), None); + let targets = vec![ + BenchTarget { + relative_path: "data/a.parquet".to_string(), + table_dir: "data/".to_string(), + region_id: "1:2".to_string(), + path_type: "data".to_string(), + file_id: "a".to_string(), + }, + BenchTarget { + relative_path: "data/b.parquet".to_string(), + table_dir: "data/".to_string(), + region_id: "1:2".to_string(), + path_type: "data".to_string(), + file_id: "b".to_string(), + }, + ]; + assert_eq!( + group_scan_paths(&targets), + vec![( + "data/".to_string(), + "1:2".to_string(), + "data".to_string(), + vec!["data/a.parquet".to_string(), "data/b.parquet".to_string()], + )] + ); + } +} diff --git a/src/cmd/src/bin/query_regression_runner/remote/storage.rs b/src/cmd/src/bin/query_regression_runner/remote/storage.rs new file mode 100644 index 00000000000..f6bf920b780 --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/remote/storage.rs @@ -0,0 +1,241 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::path::Path; +use std::process::Command; +use std::time::Instant; + +use serde_json::{Value, json}; + +use crate::query_regression_runner::Result; +use crate::query_regression_runner::model::StorageConfig; +use crate::query_regression_runner::sql::{value_f64, value_u64}; + +pub(super) fn run_storage_inspection( + generator: &Path, + data_home: &Path, + destination: Option<&Path>, + storage: &StorageConfig, +) -> Result { + let mut command = vec![ + generator.to_string_lossy().to_string(), + "inspect-footer".to_string(), + "--column".to_string(), + storage.column.clone(), + ]; + let root = if let Some(destination) = destination { + command.push("--destination".to_string()); + command.push(destination.to_string_lossy().to_string()); + data_home.to_path_buf() + } else { + let root = storage + .root_suffix + .as_deref() + .map_or_else(|| data_home.to_path_buf(), |suffix| data_home.join(suffix)); + command.push("--root".to_string()); + command.push(root.to_string_lossy().to_string()); + root + }; + if storage.include_metadata_files { + command.push("--include-metadata-files".to_string()); + } + let started = Instant::now(); + let output = Command::new(generator).args(&command[1..]).output()?; + let elapsed_seconds = started.elapsed().as_secs_f64(); + let stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + if !output.status.success() { + return Err(format!( + "storage inspector failed: command={command:?}, returncode={:?}, elapsed_seconds={elapsed_seconds:.3}, stderr={:.2000}", + output.status.code(), stderr + ) + .into()); + } + let mut result = json!({ + "status": "ok", + "command": command, + "root": root, + "returncode": output.status.code(), + "elapsed_seconds": elapsed_seconds, + "stdout": stdout, + "stderr": stderr, + }); + match serde_json::from_str(result["stdout"].as_str().unwrap_or_default()) { + Ok(summary) => result["summary"] = summary, + Err(_) => { + result["summary_parse_error"] = result["stdout"].clone(); + result["status"] = json!("failed"); + } + } + Ok(result) +} + +fn storage_summary(inspection: &Value) -> Value { + inspection + .get("summary") + .and_then(|summary| summary.get("summary")) + .filter(|summary| summary.is_object()) + .cloned() + .unwrap_or_else(|| json!({})) +} + +pub(super) fn enforce_storage_thresholds( + storage: &StorageConfig, + base_inspection: &Value, + candidate_inspection: &Value, +) -> Vec { + let base = storage_summary(base_inspection); + let candidate = storage_summary(candidate_inspection); + let targets = [("base", &base), ("candidate", &candidate)]; + let mut results = Vec::new(); + for (threshold, field, limit) in [ + ("min_files", "file_count", storage.min_files), + ( + "min_files_with_column", + "files_with_column", + storage.min_files_with_column, + ), + ] { + for (target, summary) in &targets { + let actual = summary.get(field).cloned().unwrap_or(Value::Null); + let ok = value_u64(Some(&actual)).is_some_and(|actual| actual >= limit); + results.push(json!({ "target": target, "threshold": threshold, "status": if ok { "passed" } else { "failed" }, "actual": actual, "limit": limit })); + } + } + for (threshold, field, limit) in [ + ( + "max_total_file_size_bytes", + "total_file_size", + storage.max_total_file_size_bytes, + ), + ( + "max_column_compressed_size_bytes", + "column_compressed_size", + storage.max_column_compressed_size_bytes, + ), + ( + "max_column_uncompressed_size_bytes", + "column_uncompressed_size", + storage.max_column_uncompressed_size_bytes, + ), + ] { + let Some(limit) = limit else { continue }; + for (target, summary) in &targets { + let actual = summary.get(field).cloned().unwrap_or(Value::Null); + let ok = value_f64(Some(&actual)).is_some_and(|actual| actual <= limit as f64); + results.push(json!({ "target": target, "threshold": threshold, "status": if ok { "passed" } else { "failed" }, "actual": actual, "limit": limit })); + } + } + for (target, summary) in &targets { + let encodings = summary + .get("unique_encodings") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + for encoding in &storage.require_encodings { + results.push(json!({ "target": target, "threshold": "require_encodings", "encoding": encoding, "status": if encodings.iter().any(|value| value.as_str() == Some(encoding)) { "passed" } else { "failed" } })); + } + for encoding in &storage.forbid_encodings { + results.push(json!({ "target": target, "threshold": "forbid_encodings", "encoding": encoding, "status": if encodings.iter().all(|value| value.as_str() != Some(encoding)) { "passed" } else { "failed" } })); + } + } + for (threshold, field, limit) in [ + ( + "max_candidate_total_file_size_regression_pct", + "total_file_size", + storage.max_candidate_total_file_size_regression_pct, + ), + ( + "max_candidate_column_compressed_size_regression_pct", + "column_compressed_size", + storage.max_candidate_column_compressed_size_regression_pct, + ), + ( + "max_candidate_column_uncompressed_size_regression_pct", + "column_uncompressed_size", + storage.max_candidate_column_uncompressed_size_regression_pct, + ), + ] { + let Some(limit) = limit else { continue }; + let base_value = base.get(field).cloned().unwrap_or(Value::Null); + let candidate_value = candidate.get(field).cloned().unwrap_or(Value::Null); + let (Some(base_number), Some(candidate_number)) = ( + value_f64(Some(&base_value)).filter(|value| *value != 0.0), + value_f64(Some(&candidate_value)), + ) else { + results.push(json!({ "threshold": threshold, "status": "failed", "reason": "missing or zero base/candidate value", "base": base_value, "candidate": candidate_value })); + continue; + }; + let actual = (candidate_number - base_number) / base_number * 100.0; + results.push(json!({ "threshold": threshold, "status": if actual <= limit { "passed" } else { "failed" }, "actual_pct": actual, "limit_pct": limit, "base": base_value, "candidate": candidate_value })); + } + results +} + +#[cfg(test)] +mod tests { + use super::super::is_storage_threshold_entry; + use super::*; + + #[test] + fn enforces_storage_thresholds_for_each_target_and_regression() { + let storage: StorageConfig = serde_json::from_value(json!({ + "inspect": true, + "column": "value", + "root_suffix": null, + "include_metadata_files": false, + "min_files": 2, + "min_files_with_column": 1, + "require_encodings": ["PLAIN"], + "forbid_encodings": ["DELTA"], + "max_total_file_size_bytes": 105, + "max_column_compressed_size_bytes": null, + "max_column_uncompressed_size_bytes": null, + "max_candidate_total_file_size_regression_pct": 5.0, + "max_candidate_column_compressed_size_regression_pct": null, + "max_candidate_column_uncompressed_size_regression_pct": null + })) + .unwrap(); + let base = json!({"summary": {"summary": {"file_count": 2, "files_with_column": 1, "total_file_size": 100, "unique_encodings": ["PLAIN"]}}}); + let candidate = json!({"summary": {"summary": {"file_count": 2, "files_with_column": 1, "total_file_size": 110, "unique_encodings": ["PLAIN"]}}}); + let results = enforce_storage_thresholds(&storage, &base, &candidate); + assert!( + results + .iter() + .any(|result| result["threshold"] == "max_total_file_size_bytes" + && result["target"] == "candidate" + && result["status"] == "failed") + ); + assert!(results.iter().any(|result| result["threshold"] + == "max_candidate_total_file_size_regression_pct" + && result["actual_pct"] == 10.0 + && result["status"] == "failed")); + assert!(results.iter().any( + |result| result["threshold"] == "require_encodings" && result["status"] == "passed" + )); + } + + #[test] + fn identifies_prior_storage_threshold_entries() { + assert!(is_storage_threshold_entry( + &json!({"target": "base", "threshold": "min_files"}) + )); + assert!(is_storage_threshold_entry(&json!({ + "threshold": "max_candidate_total_file_size_regression_pct" + }))); + assert!(!is_storage_threshold_entry(&json!({ + "threshold": "max_candidate_latency_regression_pct" + }))); + } +} diff --git a/src/cmd/src/bin/query_regression_runner/sql.rs b/src/cmd/src/bin/query_regression_runner/sql.rs new file mode 100644 index 00000000000..4d8baab2c8c --- /dev/null +++ b/src/cmd/src/bin/query_regression_runner/sql.rs @@ -0,0 +1,238 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::time::Instant; + +use reqwest::Client; +use serde_json::{Value, json}; + +use crate::query_regression_runner::Result; + +pub(super) fn sql_string(value: &str) -> String { + format!("'{}'", value.replace('\'', "''")) +} + +pub(super) fn extract_rows(body: &Value) -> Vec { + fn visit(body: &Value, rows: &mut Vec) { + match body { + Value::Object(object) => { + for key in ["data", "rows", "records", "output"] { + let Some(value) = object.get(key) else { + continue; + }; + if matches!(key, "data" | "rows") && value.is_array() { + rows.extend(value.as_array().unwrap().iter().cloned()); + } else { + visit(value, rows); + } + } + } + Value::Array(values) => { + if !values.is_empty() + && values + .iter() + .all(|value| !value.is_object() && !value.is_array()) + { + rows.push(body.clone()); + } else { + for value in values { + visit(value, rows); + } + } + } + _ => {} + } + } + + let mut rows = Vec::new(); + visit(body, &mut rows); + rows +} + +pub(super) fn row_value<'a>(row: &'a Value, index: usize, name: &str) -> Option<&'a Value> { + match row { + Value::Object(values) => [ + name.to_string(), + name.to_ascii_uppercase(), + name.to_ascii_lowercase(), + ] + .into_iter() + .find_map(|key| values.get(&key)), + Value::Array(values) => values.get(index), + _ => Some(row), + } +} + +pub(super) fn row_u64(row: &Value, index: usize, name: &str) -> Result { + let value = + row_value(row, index, name).ok_or_else(|| format!("missing {name} in row {row}"))?; + value + .as_u64() + .or_else(|| value.as_str().and_then(|value| value.parse().ok())) + .ok_or_else(|| format!("invalid {name} in row {row}").into()) +} + +pub(super) fn value_text(value: &Value) -> String { + match value { + Value::String(value) => value.clone(), + Value::Bool(value) => value.to_string(), + Value::Number(value) => value.to_string(), + Value::Null => "None".to_string(), + _ => value.to_string(), + } +} + +pub(super) fn extract_count_value(result: &Value) -> Option { + let row = result + .get("response")? + .get("data")? + .as_array()? + .first()? + .as_object()?; + row.iter() + .find(|(key, _)| { + let key = key.to_ascii_lowercase(); + key == "count(*)" || key.starts_with("count(") + }) + .and_then(|(_, value)| value_u64(Some(value))) +} + +pub(super) fn value_u64(value: Option<&Value>) -> Option { + value?.as_u64().or_else(|| value?.as_str()?.parse().ok()) +} + +pub(super) fn value_f64(value: Option<&Value>) -> Option { + value?.as_f64().or_else(|| value?.as_str()?.parse().ok()) +} + +pub(super) async fn http_post_sql(client: &Client, port: u16, sql: &str, db: &str) -> Value { + let started = Instant::now(); + let request = client + .post(format!("http://127.0.0.1:{port}/v1/sql")) + .form(&[("sql", sql), ("db", db), ("format", "json")]); + match request.send().await { + Ok(response) => { + let status = response.status().as_u16(); + match response.text().await { + Ok(raw) => { + let body = serde_json::from_str(&raw).unwrap_or_else(|_| json!({"raw": raw})); + let ok = status < 400 && !response_has_error(&body); + let mut sample = json!({ + "ok": ok, + "status": status, + "latency_ms": started.elapsed().as_secs_f64() * 1000.0, + "response": body, + "sql": sql, + }); + if status >= 400 { + sample + .as_object_mut() + .expect("HTTP samples are objects") + .insert("error".to_string(), Value::String(format!("HTTP {status}"))); + } + sample + } + Err(error) => json!({ + "ok": false, + "status": status, + "latency_ms": started.elapsed().as_secs_f64() * 1000.0, + "error": error.to_string(), + "sql": sql, + }), + } + } + Err(error) => json!({ + "ok": false, + "status": Value::Null, + "latency_ms": started.elapsed().as_secs_f64() * 1000.0, + "error": error.to_string(), + "sql": sql, + }), + } +} + +fn response_has_error(body: &Value) -> bool { + let Some(body) = body.as_object() else { + return false; + }; + ["error", "err_msg", "error_msg"] + .into_iter() + .any(|key| body.get(key).is_some_and(is_truthy)) + || body + .get("error_code") + .is_some_and(|value| !is_success_code(value)) + || (!body.contains_key("output") + && body + .get("code") + .is_some_and(|value| !is_success_code(value))) +} + +fn is_truthy(value: &Value) -> bool { + match value { + Value::Null => false, + Value::Bool(value) => *value, + Value::Number(value) => value.as_f64().is_none_or(|value| value != 0.0), + Value::String(value) => !value.is_empty(), + Value::Array(value) => !value.is_empty(), + Value::Object(value) => !value.is_empty(), + } +} + +fn is_success_code(value: &Value) -> bool { + let code = match value { + Value::String(value) => value.clone(), + Value::Bool(value) => value.to_string(), + Value::Number(value) => value.to_string(), + Value::Null => "None".to_string(), + _ => value.to_string(), + }; + matches!(code.to_lowercase().as_str(), "" | "0" | "success") +} + +pub(super) fn sql_ident(name: &str) -> String { + format!("\"{}\"", name.replace('"', "\"\"")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn top_level_errors_do_not_inspect_rows() { + assert!(response_has_error(&json!({"error_code": 7}))); + assert!(response_has_error(&json!({"code": "bad"}))); + assert!(!response_has_error( + &json!({"output": [{"code": 7, "error": "row value"}]}) + )); + assert!(!response_has_error( + &json!({"error_code": "success", "output": []}) + )); + } + + #[test] + fn extracts_remote_write_count_from_data_map() { + assert_eq!( + extract_count_value(&json!({"response": {"data": [{"COUNT(*)": "12"}]}})), + Some(12) + ); + assert_eq!( + extract_count_value(&json!({"response": {"data": [{"count(value)": 7}]}})), + Some(7) + ); + assert_eq!( + extract_count_value(&json!({"response": {"data": [[]]}})), + None + ); + } +} diff --git a/tests/perf/AGENTS.md b/tests/perf/AGENTS.md index 7f64b283bdc..e7c24f16e62 100644 --- a/tests/perf/AGENTS.md +++ b/tests/perf/AGENTS.md @@ -7,18 +7,21 @@ the default case set. Do not hard-code a single case such as `promql_pushdown_7913` into the workflow path. - The case DSL is not required to keep compatibility inside this PR. When the - DSL changes, update TOML cases, the Python runner, Rust fixture generator, and + DSL changes, update TOML cases, the outer lifecycle script, Rust helpers, and docs together. - `[case]` is report metadata only. `[scenario]` is the executable regression configuration and must include `kind`, data layout, tables, queries, and thresholds. Rust owns case schema, defaults, validation, and normalized plan - output through `query_perf_fixture plan`; `tests/perf/query_regression_runner.py` - should only orchestrate from that normalized JSON. + output through `query_perf_fixture plan`. `.github/scripts/query-regression-run.py` + owns process lifecycle; `query_regression_runner` consumes normalized plans, + frontend endpoints, direct-SST materialization requests, and OTLP target/finalize + requests. - Keep the direct-SST generator generic. Issue-specific behavior belongs in case files and thresholds, not in Rust generator logic. - Before pushing perf harness changes, run at least: - - `uv run --no-project python -m py_compile .github/scripts/query-regression-run.py .github/scripts/query-regression-summary.py .github/scripts/query-regression-pr-metadata.py tests/perf/query_regression_runner.py` + - `uv run --no-project python -m py_compile .github/scripts/query-regression-run.py .github/scripts/query-regression-summary.py .github/scripts/query-regression-pr-metadata.py tests/perf/test_query_regression_runner_compaction_toctou.py tests/perf/test_query_regression_runner_otlp_trace_load.py` + - `uv run --no-project python tests/perf/test_query_regression_runner_compaction_toctou.py && uv run --no-project python tests/perf/test_query_regression_runner_otlp_trace_load.py` - `cargo fmt --all -- --check` - `cargo build -p cmd --bin query_perf_fixture --features dev-tools` - - dry-run the Python runner and Rust fixture generator against all built-in - cases when the DSL or workflow case selection changes. + - exercise the outer lifecycle script and Rust fixture generator against all + built-in cases when the DSL or workflow case selection changes. diff --git a/tests/perf/README.md b/tests/perf/README.md index 8746d1ad375..cbe03b51794 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -41,7 +41,7 @@ release-level realism. ## Prometheus remote-write scenario -The runner also supports `scenario.kind = "prom_remote_write_then_query"` for a +The outer driver supports `scenario.kind = "prom_remote_write_then_query"` for a bounded write-path smoke/regression flow. This path is explicit and separate from the direct-SST fixture path: it starts the base and candidate distributed clusters, writes deterministic Prometheus remote-write v1 samples through @@ -111,11 +111,10 @@ helper so non-linear value patterns use a stable global/per-series ordinal acros chunks. Case schema, value distribution defaults, storage defaults, and read-bench -defaults are owned by Rust. The Python runner calls -`query_perf_fixture plan --case ` and orchestrates the normalized JSON. -The same helper exposes `direct-sst`, `prom-remote-write`, and `inspect-footer` -subcommands; the old direct invocation (`query_perf_fixture --case ... --out-dir -...`) remains compatible. +defaults are owned by Rust. The outer CI driver calls +`query_perf_fixture plan --case ` once, then uses the normalized plan +to select the Rust runner lifecycle. The fixture helper exposes `direct-sst`, +`prom-remote-write`, and `inspect-footer` subcommands. Remote-write cases that need to validate storage output can add `[scenario.remote_write.storage]`. When present, the scenario body becomes: @@ -163,17 +162,15 @@ measures region scan cost, and query measurements still exercise the SQL/TQL frontend path. Treat all performance conclusions as release-only; debug builds are suitable only for command wiring and correctness checks. -The runner creates the configured database if needed, writes a per-target +`prepare-remote` creates the configured database if needed. The outer driver writes a per-target frontend config enabling `[prom_store]` with metric engine storage and a non-zero `pending_rows_flush_interval`, and validates that the logical metric table reaches `series_count * samples_per_series` rows before trusting the query measurements. -Use `--fixture-generator /path/to/query_perf_fixture` to provide the Rust helper. -Deprecated `--remote-write-generator` and `--storage-inspector` options are kept -only for CLI compatibility. `--fixture-only` is rejected for remote-write cases; -use `--dry-run` for planning. +Use `--fixture-generator /path/to/query_perf_fixture` to provide the Rust helper +to the outer driver. Large manual remote-write cases can set `sample_chunk_size` to split ingestion by -time. For each chunk, the runner invokes `query_perf_fixture prom-remote-write` with the +time. For each chunk, `prepare-remote` invokes `query_perf_fixture prom-remote-write` with the same series cardinality but a shorter `--samples-per-series` and an advanced `--start-unix-millis`. `flush_every_sample_chunks` controls periodic `ADMIN FLUSH_TABLE('')` calls; with `flush_every_sample_chunks = 1`, @@ -213,34 +210,27 @@ workflow dispatch accepts the `heavy` token to select this case. ## OTLP trace load scenario `scenario.kind = "otlp_trace_load"` runs a bounded native `otelgen` process -against each local distributed cluster. The runner excludes the configured -warmup window, derives throughput and mean request latency from GreptimeDB's -OTLP-specific metrics, flushes the trace table, and verifies its row count -against the accepted-span counter. The case is intentionally outside the -default set until its variance is known. +against each local distributed cluster. The outer driver runs the base target +alone with Rust `run-otlp-target`, fully stops it, then runs the candidate alone +and calls Rust `finalize-otlp` to aggregate metrics, thresholds, and the final +report. The case is intentionally outside the default set until its variance is +known. -Build base/candidate `greptime` binaries with the same profile, build the -candidate `query_perf_fixture`, then run the case explicitly: +Build base/candidate `greptime` binaries with the same profile and build the +candidate `query_perf_fixture` and `query_regression_runner`, then run the outer +driver explicitly: ```bash WORK_DIR="$(mktemp -d /tmp/query-perf-otlp.XXXXXX)" -REPORT="$WORK_DIR/trace-report.json" -python3 tests/perf/query_regression_runner.py \ - --case tests/perf/query_cases/otlp_trace_load/case.toml \ +uv run --no-project python .github/scripts/query-regression-run.py \ + --cases tests/perf/query_cases/otlp_trace_load/case.toml \ --base-bin /path/to/base/target/nightly/greptime \ --candidate-bin /path/to/candidate/target/nightly/greptime \ --fixture-generator /path/to/candidate/target/nightly/query_perf_fixture \ + --runner /path/to/candidate/target/nightly/query_regression_runner \ --otelgen-bin /path/to/otelgen \ - --work-dir "$WORK_DIR" \ - --output "$REPORT" -``` - -Install [YouPlot](https://github.com/red-data-tools/YouPlot) once, then render -all comparison metrics and threshold results in the terminal: - -```bash -brew install youplot -tests/perf/plot_otlp_trace_report.sh "$REPORT" + --work-dir "$WORK_DIR" +REPORT="$WORK_DIR/otlp_trace_load/query-regression-report.json" ``` For each target, `accepted_spans` should equal `table_rows`, and `failures` @@ -323,8 +313,8 @@ seed = 12345 # query, warmups, iterations, thresholds ``` -The runner currently supports `direct_readable_sst` and -`prom_remote_write_then_query`. +The outer driver currently supports `direct_readable_sst`, +`prom_remote_write_then_query`, and `otlp_trace_load`. ## Metrics @@ -340,120 +330,49 @@ Primary gates should compare query work rather than plan text: Plan details such as pushed filters are useful diagnostics, but should not be the main pass/fail signal. -## Runner MVP +## Runner lifecycle -`query_regression_runner.py` is the base-vs-candidate orchestration layer. The -current MVP parses a case, creates per-target work directories, and in real query -mode starts a local distributed cluster for each target: metasrv (memory-store, -region failover disabled), one datanode (`node_id=0`), and one frontend. It -creates the configured Mito table(s) through frontend HTTP SQL, discovers the -real one-region-per-table metadata via `information_schema`, stops only the -owning datanode, generates one shared direct-SST fixture per table using the -discovered `--region-id`, `--table-dir`, and `--table`, injects those region -subtrees into the datanode data home, restarts the datanode, then validates and -measures through frontend. Reports are written as JSON under the work directory. +`.github/scripts/query-regression-run.py` is the outer CI driver. It resolves the +base and candidate `greptime` binaries, candidate `query_perf_fixture`, and +candidate `query_regression_runner`; allocates disjoint localhost +meta/datanode/frontend HTTP, gRPC, MySQL, and Postgres ports for each target; and +writes component logs below `//logs/`. -The runner intentionally keeps metasrv alive for the whole target run because -memory-store metadata would otherwise be lost. It replaces only the discovered -datanode region directory under `data/greptime///...` with -generated SST files and a manifest checkpoint. For multi-table cases this is -repeated per table, enabling true JOIN fixtures while still requiring exactly one -region per table. Base and candidate must discover identical per-table -`table_dir` and `region_id`; otherwise the run fails. +The Rust runner consumes the normalized plan and endpoint ports. For +`direct_readable_sst`, the driver starts both clusters, runs `prepare-direct`, +stops only both datanodes, writes File-object-store destination TOMLs, invokes +`materialize` for every fixture and target, restarts datanodes, then runs +`measure` (with one frontend restart retry). `materialize` uses OpenDAL and only +replaces the fixture's exact region prefix in each target data home. -Multi-table direct-SST cases must use unique table names as well as unique -`(database, name)` pairs because the generator currently selects a table with -`--table `. Per-table fixture directories are derived from table index, -database, and table name with path-unsafe characters sanitized. +For `prom_remote_write_then_query`, it renders a target frontend configuration +with `render-remote-config`, starts both clusters, runs `prepare-remote` and +`measure`, stops both datanodes, then runs `finalize-remote`. Finalization owns +storage inspection and read-bench against the quiescent data homes. Both paths +write `query-regression-report.json` in the case work directory and always clean +up the remaining components. -Currently enforced threshold: - -- `max_candidate_latency_regression_pct`, based on client-side median latency. - -Server-side scan thresholds such as file ranges and scanned rows are planned for -a follow-up PR that extracts them from structured `EXPLAIN ANALYZE VERBOSE` -output. Do not add those threshold keys until the runner enforces them. - -Dry-run example: +For local orchestration after building the three binaries, invoke the outer +driver directly: ```bash -uv run --no-project python tests/perf/query_regression_runner.py \ - --case tests/perf/query_cases/promql_pushdown_7913/case.toml \ - --base-bin /path/to/base/greptime \ - --candidate-bin /path/to/candidate/greptime \ - --work-dir /tmp/query-perf-work \ - --dry-run -``` - -With a fixture generator: - -```bash -uv run --no-project python tests/perf/query_regression_runner.py \ - --case tests/perf/query_cases/promql_pushdown_7913/case.toml \ +uv run --no-project python .github/scripts/query-regression-run.py \ + --cases tests/perf/query_cases/smoke_direct_sst/case.toml \ --base-bin /path/to/base/greptime \ --candidate-bin /path/to/candidate/greptime \ --fixture-generator /path/to/query_perf_fixture \ - --fixture-cache-dir /mnt/query-regression-fixtures \ - --allow-large-fixture \ - --work-dir /tmp/query-perf-work + --runner /path/to/query_regression_runner \ + --work-dir /tmp/query-regression-work ``` -This mode launches metasrv, datanode, and frontend for each target with explicit -localhost HTTP/gRPC/MySQL/Postgres ports and writes component stdout/stderr under -each target's `logs/` directory. - -By default query mode requires fresh base/candidate work directories and fails if -either target directory already exists with contents. Use `--reuse-work-dir` only -when intentionally debugging an existing run directory. SQL HTTP requests default -to a 120 second timeout; override with `--http-timeout ` for slow lab -runs. - -For large direct-SST fixtures, pass `--fixture-cache-dir ` to store generated -fixtures in a persistent content-addressed cache keyed by case name and fixture -data configuration. Query and threshold edits reuse the same cached data as long -as the scenario layout and table definitions do not change. Cached fixtures are -reused automatically when their `summary.json` -matches the discovered table/region metadata; incompatible entries are -regenerated instead of reused. Fixture materialization keeps base and candidate -data directories isolated, but copies files efficiently by trying filesystem -reflinks first, hardlinks for immutable SST/object files next, and normal copies -as a fallback. Manifest files are reflinked or copied, not hardlinked. - -Fixture generator smoke test: +The Rust runner subcommands are also useful for focused diagnostics: ```bash -cargo run -p cmd --bin query_perf_fixture --features dev-tools -- \ - direct-sst \ - --case tests/perf/query_cases/smoke_direct_sst/case.toml \ - --out-dir /tmp/query-perf-smoke -``` - -Runner smoke test with fixture generation only: - -```bash -uv run --no-project python tests/perf/query_regression_runner.py \ - --case tests/perf/query_cases/smoke_direct_sst/case.toml \ - --base-bin /path/to/query_perf_fixture \ - --candidate-bin /path/to/query_perf_fixture \ - --fixture-generator /path/to/query_perf_fixture \ - --work-dir /tmp/query-perf-runner-smoke \ - --fixture-only -``` - -`--fixture-only` preserves the earlier smoke behavior: it does not start -standalone servers, and it materializes the generated fixture into base and -candidate data directories for plumbing validation. - -Remote-write runner dry-run: - -```bash -uv run --no-project python tests/perf/query_regression_runner.py \ - --case tests/perf/query_cases/prom_remote_write_seeded_random/case.toml \ - --base-bin /path/to/base/greptime \ - --candidate-bin /path/to/candidate/greptime \ - --fixture-generator /path/to/query_perf_fixture \ - --work-dir /tmp/query-perf-remote-write \ - --dry-run +query_regression_runner prepare-direct --case --fixture-generator \ + --base-http-port --candidate-http-port --fixture-dir --output +query_regression_runner materialize --fixture-dir --destination +query_regression_runner measure --case --fixture-generator \ + --base-http-port --candidate-http-port --output ``` ## GitHub Actions @@ -462,14 +381,14 @@ uv run --no-project python tests/perf/query_regression_runner.py \ query regression runs. It builds its own binaries for now: - base `greptime` from the PR base commit, or `workflow_dispatch` `base_ref` -- candidate `greptime` and `query_perf_fixture` from the PR merge ref/current - candidate checkout -- runner and summary formatter from the candidate checkout +- candidate `greptime`, `query_perf_fixture`, and `query_regression_runner` from + the PR merge ref/current candidate checkout +- outer driver and summary formatter from the candidate checkout The workflow builds base and candidate `greptime` as normal release-equivalent -binaries. Candidate `query_perf_fixture` is the extra head-side helper binary; -the runner uses candidate `greptime datanode parquetbench/scanbench` as the -read-bench tool against each target's data directory. +binaries. Candidate `query_perf_fixture` and `query_regression_runner` are the +extra head-side helpers; `finalize-remote` uses candidate `greptime datanode +parquetbench/scanbench` as the read-bench tool against each target's data directory. The workflow runs automatically only when `query-regression` or `heavy-regression` is added to a non-draft PR; it does not rerun on pushes, ready-for-review, or diff --git a/tests/perf/plot_otlp_trace_report.sh b/tests/perf/plot_otlp_trace_report.sh deleted file mode 100755 index fa1e5ed796d..00000000000 --- a/tests/perf/plot_otlp_trace_report.sh +++ /dev/null @@ -1,80 +0,0 @@ -#!/usr/bin/env bash - -set -euo pipefail - -if [[ $# -ne 1 ]]; then - printf 'Usage: %s REPORT.json\n' "${0##*/}" >&2 - exit 2 -fi - -report="$1" - -[[ -f "${report}" ]] || { - printf 'Report does not exist: %s\n' "${report}" >&2 - exit 1 -} -command -v jq >/dev/null || { - printf 'jq is required\n' >&2 - exit 1 -} -command -v uplot >/dev/null || { - printf 'YouPlot is required; install it with: brew install youplot\n' >&2 - exit 1 -} - -jq -e ' - (.status | type == "string") and - (.targets | type == "array" and length > 0) and - all(.targets[]; - (.name | type == "string") and - (.metrics.accepted_spans | type == "number") and - (.visibility.observed_rows | type == "number") and - (.metrics.accepted_spans_per_second | type == "number") and - (.metrics.mean_http_latency_ms | type == "number") and - (.metrics.failure_count | type == "number") - ) -' "${report}" >/dev/null || { - printf 'Report is not a completed OTLP trace comparison: %s\n' "${report}" >&2 - exit 1 -} - -plot_metric() { - local title="$1" - local filter="$2" - local color="$3" - - printf '\n' - jq -r ".targets[] | [.name, ${filter}] | @tsv" "${report}" | - uplot bar -t "${title}" -w 80 -c "${color}" -o - -} - -jq -r '"Status: \(.status)"' "${report}" -row_count_status=0 -jq -r ' - "\nAccepted spans vs visible table rows", - (.targets[] | - " [\(if .metrics.accepted_spans == .visibility.observed_rows then "passed" else "failed" end)] \(.name): accepted=\(.metrics.accepted_spans), visible=\(.visibility.observed_rows)" - ) -' "${report}" -jq -e 'all(.targets[]; .metrics.accepted_spans == .visibility.observed_rows)' "${report}" >/dev/null || row_count_status=1 - -plot_metric "Accepted spans" ".metrics.accepted_spans" blue -plot_metric "Visible table rows" ".visibility.observed_rows" green -plot_metric "Throughput (spans/s, higher is better)" ".metrics.accepted_spans_per_second" cyan -plot_metric "Mean HTTP latency (ms, lower is better)" ".metrics.mean_http_latency_ms" yellow -plot_metric "Failures (lower is better)" ".metrics.failure_count" red - -jq -r ' - "\nThresholds", - (.thresholds[] | - if has("actual_pct") then - " [\(.status)] \(.threshold): \(.actual_pct)% (limit \(.limit_pct)%; base \(.base), candidate \(.candidate))" - elif has("actual") then - " [\(.status)] \(.threshold) [\(.target)]: \(.actual) (limit \(.limit))" - else - " [\(.status)] \(.threshold): \(.reason // "no measured value")" - end - ) -' "${report}" - -exit "${row_count_status}" diff --git a/tests/perf/query_regression_runner.py b/tests/perf/query_regression_runner.py deleted file mode 100644 index 724c38ebf17..00000000000 --- a/tests/perf/query_regression_runner.py +++ /dev/null @@ -1,1663 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2023 Greptime Team -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Base-vs-candidate query performance regression runner.""" - -from __future__ import annotations - -import argparse -import fcntl -import hashlib -import json -import math -import os -import re -import shutil -import socket -import statistics -import subprocess -import sys -import time -import tomllib -import urllib.error -import urllib.parse -import urllib.request -from dataclasses import dataclass -from pathlib import Path -from typing import Any - - -FICLONE = 0x40049409 -OTLP_TRACE_METRICS = { - "greptime_frontend_otlp_traces_rows", - "greptime_frontend_otlp_traces_failure_count", - "greptime_servers_http_otlp_traces_elapsed_sum", - "greptime_servers_http_otlp_traces_elapsed_count", -} -PROMETHEUS_SAMPLE_RE = re.compile(r"^([A-Za-z_:][A-Za-z0-9_:]*)(?:\{.*\})?\s+([^\s]+)(?:\s+.*)?$") - - -@dataclass(frozen=True) -class RunTarget: - name: str - binary: Path - work_dir: Path - data_dir: Path - fixture_dir: Path - report_path: Path - http_port: int - grpc_port: int - mysql_port: int - postgres_port: int - metasrv_rpc_port: int - metasrv_http_port: int - datanode_rpc_port: int - datanode_http_port: int - datanode_data_dir: Path - - -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description="Run a query perf case against base and candidate binaries.") - p.add_argument("--case", required=True, type=Path) - p.add_argument("--base-bin", required=True, type=Path) - p.add_argument("--candidate-bin", required=True, type=Path) - p.add_argument("--fixture-generator", type=Path) - p.add_argument("--otelgen-bin", type=Path) - p.add_argument("--remote-write-generator", type=Path, help="deprecated: use --fixture-generator query_perf_fixture") - p.add_argument("--storage-inspector", type=Path, help="deprecated: use --fixture-generator query_perf_fixture") - p.add_argument("--work-dir", required=True, type=Path) - p.add_argument("--output", type=Path, help="write the final JSON report to this file instead of stdout") - p.add_argument("--fixture-cache-dir", type=Path, help="persistent directory for generated fixtures, keyed by case content") - p.add_argument("--reuse-fixture", action="store_true") - p.add_argument("--allow-large-fixture", action="store_true") - p.add_argument("--dry-run", action="store_true") - p.add_argument("--fixture-only", action="store_true", help="old smoke mode: generate/materialize fixture only") - p.add_argument("--reuse-work-dir", action="store_true", help="allow non-empty base/candidate work dirs") - p.add_argument("--http-timeout", type=float, default=120.0, help="HTTP SQL timeout seconds") - return p.parse_args() - - -def load_case(path: Path) -> dict[str, Any]: - with path.open("rb") as f: - return tomllib.load(f) - - -def load_normalized_case(case_path: Path, fixture_generator: Path | None) -> dict[str, Any]: - if fixture_generator is None: - raise ValueError("--fixture-generator query_perf_fixture is required for Rust-owned case planning") - result = run_command([str(fixture_generator), "plan", "--case", str(case_path)]) - if result["returncode"] != 0: - raise RuntimeError(f"query_perf_fixture plan failed: {result['stderr'][:2000]}") - plan = json.loads(result["stdout"]) - return {"case": load_case(case_path).get("case", {}), "scenario": plan["scenario"], "schema_version": plan.get("schema_version")} - - -def require_binary(path: Path, name: str, *, dry_run: bool) -> None: - if dry_run: - return - if not path.exists(): - raise FileNotFoundError(f"{name} binary does not exist: {path}") - if not os.access(path, os.X_OK): - raise PermissionError(f"{name} binary is not executable: {path}") - - -def allocate_ports(n: int) -> list[int]: - socks = [] - try: - for _ in range(n): - s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - s.bind(("127.0.0.1", 0)) - socks.append(s) - return [s.getsockname()[1] for s in socks] - finally: - for s in socks: - s.close() - - -def make_target(name: str, binary: Path, root: Path, ports: list[int], fixture_dir: Path | None = None) -> RunTarget: - work_dir = root / name - return RunTarget( - name, - binary, - work_dir, - work_dir / "cluster_data", - fixture_dir or root / "fixture", - work_dir / "report.json", - ports[4], - ports[5], - ports[6], - ports[7], - ports[0], - ports[1], - ports[2], - ports[3], - work_dir / "datanode-0" / "data", - ) - - -def run_command(cmd: list[str], cwd: Path | None = None) -> dict[str, Any]: - started = time.monotonic() - proc = subprocess.run(cmd, cwd=cwd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False) - return {"cmd": cmd, "cwd": str(cwd) if cwd else None, "returncode": proc.returncode, "elapsed_seconds": time.monotonic() - started, "stdout": proc.stdout, "stderr": proc.stderr} - - -def sql_ident(name: str) -> str: - return '"' + name.replace('"', '""') + '"' - - -def sql_string(value: str) -> str: - return "'" + value.replace("'", "''") + "'" - - -def column_sql(col: dict[str, Any]) -> str: - return f"{sql_ident(col['name'])} {col['type']}" - - -def scenario(case: dict[str, Any]) -> dict[str, Any]: - value = case.get("scenario") - if not isinstance(value, dict): - raise ValueError("case requires [scenario] with a supported kind") - kind = value.get("kind") - supported = ("direct_readable_sst", "prom_remote_write_then_query", "otlp_trace_load") - if kind not in supported: - raise ValueError(f"unsupported scenario kind {kind!r}; supported: {', '.join(supported)}") - return value - - -def case_tables(case: dict[str, Any]) -> list[dict[str, Any]]: - value = scenario(case) - if value.get("kind") == "otlp_trace_load": - load = value["load"] - return [{"database": load["database"], "name": load["table"], "engine": "trace", "validate_show_create_engine": False}] - if value.get("kind") == "prom_remote_write_then_query": - remote = value["remote_write"] - metric = remote["metric"] - database = remote["database"] - return [{"database": database, "name": metric, "engine": "metric", "validate_show_create_engine": False}] - tables = value.get("tables") or [] - if not tables or value.get("layout", {}).get("regions") != 1: - raise ValueError("runner supports one or more tables and exactly one region per table") - pairs = [(table.get("database"), table.get("name")) for table in tables] - if len(set(pairs)) != len(pairs): - raise ValueError("duplicate (database, name) table entries are not supported") - names = [table.get("name") for table in tables] - if len(set(names)) != len(names): - raise ValueError("duplicate table names are not supported because fixture generator --table selects by name") - return list(tables) - - -def fixture_subdir(table: dict[str, Any], index: int) -> str: - raw = f"{index:02d}_{table['database']}_{table['name']}" - safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", raw).strip("._-") - return safe or f"table_{index:02d}" - - -def safe_path_component(raw: str) -> str: - safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", raw).strip("._-") - return safe or "query_perf_case" - - -def fixture_root(work_root: Path, case_path: Path, case: dict[str, Any], fixture_cache_dir: Path | None) -> Path: - if fixture_cache_dir is None: - return work_root / "fixture" - case_name = case.get("case", {}).get("name") or case_path.parent.name or case_path.stem - fixture_config = dict(scenario(case)) - fixture_config.pop("queries", None) - digest = hashlib.sha256(json.dumps(fixture_config, sort_keys=True).encode()).hexdigest()[:16] - return fixture_cache_dir.resolve() / f"{safe_path_component(case_name)}-{digest}" - - -def table_fixture_dir(root: Path, tables: list[dict[str, Any]], table: dict[str, Any], index: int) -> Path: - return root if len(tables) == 1 else root / fixture_subdir(table, index) - - -def create_table_sql(table: dict[str, Any]) -> str: - cols = ",\n ".join(column_sql(c) for c in table["columns"]) - pk = ", ".join(sql_ident(c) for c in table.get("primary_key", [])) - opts = [] - if "append_mode" in table: - opts.append(("append_mode", str(table["append_mode"]).lower())) - if table.get("sst_format"): - opts.append(("sst_format", table["sst_format"])) - with_sql = "" - if opts: - with_sql = "\nWITH (" + ", ".join(f"'{k}'='{v}'" for k, v in opts) + ")" - return f"CREATE TABLE {sql_ident(table['name'])} (\n {cols},\n TIME INDEX ({sql_ident(table['time_index'])}),\n PRIMARY KEY ({pk})\n) ENGINE=mito{with_sql};" - - -def planned_queries(case: dict[str, Any]) -> list[dict[str, Any]]: - queries = scenario(case).get("queries", []) - if isinstance(queries, dict): - return [{"name": name, **value} for name, value in queries.items()] - return list(queries) - - -def http_post_sql(port: int, sql: str, db: str, timeout: float) -> dict[str, Any]: - data = urllib.parse.urlencode({"sql": sql, "db": db, "format": "json"}).encode() - req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/sql", data=data, method="POST") - started = time.monotonic() - elapsed_ms = (time.monotonic() - started) * 1000.0 - try: - with urllib.request.urlopen(req, timeout=timeout) as resp: - raw = resp.read().decode() - status = resp.status - elapsed_ms = (time.monotonic() - started) * 1000.0 - try: - body: Any = json.loads(raw) - except json.JSONDecodeError: - body = {"raw": raw} - ok = status < 400 and not response_has_error(body) - return {"ok": ok, "status": status, "latency_ms": elapsed_ms, "response": body, "sql": sql} - except urllib.error.HTTPError as e: - elapsed_ms = (time.monotonic() - started) * 1000.0 - raw = e.read().decode(errors="replace") - try: - body = json.loads(raw) - except json.JSONDecodeError: - body = {"raw": raw} - return { - "ok": False, - "status": e.code, - "latency_ms": elapsed_ms, - "response": body, - "error": repr(e), - "sql": sql, - } - except Exception as e: # noqa: BLE001 - report HTTP/query failures as JSON - elapsed_ms = (time.monotonic() - started) * 1000.0 - return {"ok": False, "status": None, "latency_ms": elapsed_ms, "error": repr(e), "sql": sql} - - -def response_has_error(body: Any) -> bool: - """Detect top-level GreptimeDB HTTP error envelopes without inspecting rows. - - Query outputs may legitimately contain columns named `code` or `error`, so - this intentionally avoids recursive checks through result rows. - """ - if isinstance(body, dict): - if body.get("error") or body.get("err_msg") or body.get("error_msg"): - return True - if "error_code" in body and str(body.get("error_code", "")).lower() not in ("", "0", "success"): - return True - if "code" in body and "output" not in body and str(body.get("code", "")).lower() not in ("", "0", "success"): - return True - return False - - -def wait_health(port: int, timeout_s: float = 60.0) -> None: - deadline = time.monotonic() + timeout_s - last = None - while time.monotonic() < deadline: - try: - with urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=2) as r: - if r.status < 500: - return - except Exception as e: # noqa: BLE001 - diagnostic loop - last = e - time.sleep(0.5) - raise TimeoutError(f"health check timed out on port {port}: {last}") - - -class DistributedCluster: - """Local metasrv + one datanode + frontend cluster for query-mode runs.""" - - def __init__(self, target: RunTarget): - self.target = target - self.procs: dict[str, subprocess.Popen[bytes]] = {} - - def component_report(self) -> dict[str, Any]: - return { - "metasrv": { - "grpc": f"127.0.0.1:{self.target.metasrv_rpc_port}", - "http": f"127.0.0.1:{self.target.metasrv_http_port}", - "logs": str(self.target.work_dir / "logs" / "metasrv"), - }, - "datanode_0": { - "node_id": 0, - "grpc": f"127.0.0.1:{self.target.datanode_rpc_port}", - "http": f"127.0.0.1:{self.target.datanode_http_port}", - "data_home": str(self.target.datanode_data_dir), - "logs": str(self.target.work_dir / "logs" / "datanode-0"), - }, - "frontend": { - "http": f"127.0.0.1:{self.target.http_port}", - "grpc": f"127.0.0.1:{self.target.grpc_port}", - "mysql": f"127.0.0.1:{self.target.mysql_port}", - "postgres": f"127.0.0.1:{self.target.postgres_port}", - "logs": str(self.target.work_dir / "logs" / "frontend"), - }, - } - - def _spawn(self, name: str, args: list[str]) -> None: - logs = self.target.work_dir / "logs" / name - logs.mkdir(parents=True, exist_ok=True) - with (logs / "stdout.log").open("ab") as out, (logs / "stderr.log").open("ab") as err: - self.procs[name] = subprocess.Popen(args, stdout=out, stderr=err) - - def _ensure_metasrv_alive(self) -> None: - proc = self.procs.get("metasrv") - if proc is None or proc.poll() is not None: - raise RuntimeError("metasrv exited; memory-store metadata is no longer valid") - - def start_metasrv(self) -> None: - log_dir = self.target.work_dir / "logs" / "metasrv" - self._spawn( - "metasrv", - [ - str(self.target.binary), - "metasrv", - "start", - "--grpc-bind-addr", - f"127.0.0.1:{self.target.metasrv_rpc_port}", - "--grpc-server-addr", - f"127.0.0.1:{self.target.metasrv_rpc_port}", - "--http-addr", - f"127.0.0.1:{self.target.metasrv_http_port}", - "--backend", - "memory-store", - "--enable-region-failover", - "false", - "--log-dir", - str(log_dir), - ], - ) - wait_health(self.target.metasrv_http_port) - - def start_datanode(self) -> None: - self._ensure_metasrv_alive() - log_dir = self.target.work_dir / "logs" / "datanode-0" - self.target.datanode_data_dir.mkdir(parents=True, exist_ok=True) - self._spawn( - "datanode", - [ - str(self.target.binary), - "datanode", - "start", - "--grpc-bind-addr", - f"127.0.0.1:{self.target.datanode_rpc_port}", - "--grpc-server-addr", - f"127.0.0.1:{self.target.datanode_rpc_port}", - "--http-addr", - f"127.0.0.1:{self.target.datanode_http_port}", - "--data-home", - str(self.target.datanode_data_dir), - "--log-dir", - str(log_dir), - "--node-id", - "0", - "--metasrv-addrs", - f"127.0.0.1:{self.target.metasrv_rpc_port}", - ], - ) - wait_health(self.target.datanode_http_port) - - def start_frontend(self, config_file: Path | None = None) -> None: - self._ensure_metasrv_alive() - log_dir = self.target.work_dir / "logs" / "frontend" - cmd = [ - str(self.target.binary), - "frontend", - "start", - ] - if config_file is not None: - cmd += ["--config-file", str(config_file)] - cmd += [ - "--metasrv-addrs", - f"127.0.0.1:{self.target.metasrv_rpc_port}", - "--http-addr", - f"127.0.0.1:{self.target.http_port}", - "--grpc-bind-addr", - f"127.0.0.1:{self.target.grpc_port}", - "--grpc-server-addr", - f"127.0.0.1:{self.target.grpc_port}", - "--mysql-addr", - f"127.0.0.1:{self.target.mysql_port}", - "--postgres-addr", - f"127.0.0.1:{self.target.postgres_port}", - "--log-dir", - str(log_dir), - ] - self._spawn("frontend", cmd) - wait_health(self.target.http_port) - - def start_all(self, frontend_config: Path | None = None) -> None: - self.start_metasrv() - self.start_datanode() - self.start_frontend(frontend_config) - - def stop_component(self, name: str) -> None: - proc = self.procs.pop(name, None) - if proc is None: - return - proc.terminate() - try: - proc.wait(timeout=20) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait(timeout=20) - - def stop_all(self) -> None: - for name in ("frontend", "datanode", "metasrv"): - self.stop_component(name) - - def restart_frontend(self) -> None: - self.stop_component("frontend") - self.start_frontend() - - -def extract_rows(body: Any) -> list[Any]: - """Extract rows from GreptimeDB HTTP JSON in a tolerant way.""" - rows: list[Any] = [] - if isinstance(body, dict): - for key in ("data", "rows", "records", "output"): - if key in body: - value = body[key] - if key in ("data", "rows") and isinstance(value, list): - rows.extend(value) - else: - rows.extend(extract_rows(value)) - elif isinstance(body, list): - if body and all(not isinstance(v, (dict, list)) for v in body): - rows.append(body) - else: - for item in body: - rows.extend(extract_rows(item)) - return rows - - -def row_value(row: Any, idx: int, name: str) -> Any: - if isinstance(row, dict): - for key in (name, name.upper(), name.lower()): - if key in row: - return row[key] - return None - if isinstance(row, list): - return row[idx] - return row - - -def discover_region_via_frontend(target: RunTarget, table: dict[str, Any], http_timeout: float) -> dict[str, Any]: - schema = table["database"] - table_name = table["name"].replace("'", "''") - schema_name = schema.replace("'", "''") - table_sql = f"SELECT table_id FROM information_schema.tables WHERE table_schema = '{schema_name}' AND table_name = '{table_name}'" - table_result = http_post_sql(target.http_port, table_sql, schema, http_timeout) - if not table_result["ok"]: - raise RuntimeError(f"table_id discovery failed: {table_result}") - table_rows = extract_rows(table_result.get("response")) - if len(table_rows) != 1: - raise RuntimeError(f"expected one information_schema.tables row, got {len(table_rows)}: {table_result}") - table_id = int(row_value(table_rows[0], 0, "table_id")) - - region_sql = f"SELECT region_id, peer_id, peer_addr, is_leader, status FROM information_schema.region_peers WHERE table_schema = '{schema_name}' AND table_name = '{table_name}'" - region_result = http_post_sql(target.http_port, region_sql, schema, http_timeout) - if not region_result["ok"]: - raise RuntimeError(f"region_peers discovery failed: {region_result}") - region_rows = extract_rows(region_result.get("response")) - if len(region_rows) != 1: - raise RuntimeError(f"expected one information_schema.region_peers row, got {len(region_rows)}: {region_result}") - row = region_rows[0] - region_id = int(row_value(row, 0, "region_id")) - peer_id = int(row_value(row, 1, "peer_id")) - peer_addr = row_value(row, 2, "peer_addr") - is_leader = row_value(row, 3, "is_leader") - status = row_value(row, 4, "status") - if str(is_leader).lower() not in ("yes", "true"): - raise RuntimeError(f"expected leader region peer, got is_leader={is_leader}: {region_result}") - if str(status).upper() != "ALIVE": - raise RuntimeError(f"expected ALIVE region peer, got status={status}: {region_result}") - if peer_id != 0: - raise RuntimeError(f"expected region leader on datanode peer_id=0, got {peer_id}") - computed_table_id = region_id >> 32 - region_seq = region_id & 0xFFFFFFFF - if computed_table_id != table_id: - raise RuntimeError(f"table_id mismatch: tables={table_id}, region_id-derived={computed_table_id}") - table_dir = f"data/greptime/{schema}/{table_id}/" - region_dir = f"data/greptime/{schema}/{table_id}/{table_id}_{region_seq:010}" - return { - "catalog": "greptime", - "schema": schema, - "table_id": table_id, - "region_id": region_id, - "region_seq": region_seq, - "table_dir": table_dir, - "region_dir": region_dir, - "peer_id": peer_id, - "peer_addr": peer_addr, - "is_leader": is_leader, - "status": status, - "discovery_queries": {"table": table_result, "region": region_result}, - } - - -def assert_fixture_summary(summary: dict[str, Any], *, region_id: int | None, table_dir: str | None, region_dir: str | None = None, table_name: str | None = None, database: str | None = None) -> None: - if table_name is not None and summary.get("table") != table_name: - raise RuntimeError(f"fixture table mismatch: {summary.get('table')} != {table_name}") - if database is not None and summary.get("database") != database: - raise RuntimeError(f"fixture database mismatch: {summary.get('database')} != {database}") - if region_id is not None and int(summary["region_id"]) != region_id: - raise RuntimeError(f"fixture region_id mismatch: {summary['region_id']} != {region_id}") - if table_dir is not None and summary["table_dir"] != table_dir: - raise RuntimeError(f"fixture table_dir mismatch: {summary['table_dir']} != {table_dir}") - if region_dir is not None and summary["region_dir"].strip("/") != region_dir.strip("/"): - raise RuntimeError(f"fixture region_dir mismatch: {summary['region_dir']} != {region_dir}") - - -def generate_fixture(generator: Path | None, case_path: Path, fixture_dir: Path, *, dry_run: bool, reuse_fixture: bool, allow_large_fixture: bool, table_name: str | None = None, database: str | None = None, region_id: int | None = None, table_dir: str | None = None, region_dir: str | None = None) -> dict[str, Any]: - if generator is None: - return {"status": "skipped", "reason": "no fixture generator provided"} - summary_path = fixture_dir / "summary.json" - reuse_miss: str | None = None - if reuse_fixture and summary_path.exists(): - summary = json.loads(summary_path.read_text()) - try: - assert_fixture_summary(summary, region_id=region_id, table_dir=table_dir, region_dir=region_dir, table_name=table_name, database=database) - return {"status": "reused", "fixture_dir": str(fixture_dir), "summary": summary} - except RuntimeError as e: - reuse_miss = str(e) - shutil.rmtree(fixture_dir) - if fixture_dir.exists() and not reuse_fixture: - shutil.rmtree(fixture_dir) - fixture_dir.mkdir(parents=True, exist_ok=True) - cmd = [str(generator), "direct-sst", "--case", str(case_path), "--out-dir", str(fixture_dir)] - if table_name is not None: - cmd += ["--table", table_name] - if region_id is not None: - cmd += ["--region-id", str(region_id)] - if table_dir is not None: - cmd += ["--table-dir", table_dir] - if allow_large_fixture: - cmd.append("--allow-large") - if dry_run: - return {"status": "dry-run", "cmd": cmd} - result = run_command(cmd) - result["status"] = "ok" if result["returncode"] == 0 else "failed" - if result["returncode"] != 0: - raise RuntimeError(f"fixture generator failed: {result['stderr'][:2000]}") - summary = json.loads(summary_path.read_text()) - assert_fixture_summary(summary, region_id=region_id, table_dir=table_dir, region_dir=region_dir, table_name=table_name, database=database) - result["summary"] = summary - if reuse_miss is not None: - result["reuse_miss"] = reuse_miss - return result - - -def copy_stats() -> dict[str, int]: - return {"files": 0, "dirs": 0, "reflinked": 0, "hardlinked": 0, "copied": 0} - - -def merge_copy_stats(left: dict[str, int], right: dict[str, int]) -> None: - for key, value in right.items(): - left[key] = left.get(key, 0) + value - - -def reflink_file(src: Path, dst: Path) -> bool: - try: - dst.parent.mkdir(parents=True, exist_ok=True) - sfd = os.open(src, os.O_RDONLY) - try: - dfd = os.open(dst, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, src.stat().st_mode & 0o777) - try: - fcntl.ioctl(dfd, FICLONE, sfd) - shutil.copystat(src, dst, follow_symlinks=True) - return True - finally: - os.close(dfd) - finally: - os.close(sfd) - except OSError: - try: - dst.unlink() - except OSError: - pass - return False - - -def materialize_file(src: Path, dst: Path, *, allow_hardlink: bool) -> str: - if reflink_file(src, dst): - return "reflinked" - if allow_hardlink: - try: - dst.parent.mkdir(parents=True, exist_ok=True) - try: - dst.unlink() - except FileNotFoundError: - pass - os.link(src, dst) - return "hardlinked" - except OSError: - pass - dst.parent.mkdir(parents=True, exist_ok=True) - shutil.copy2(src, dst) - return "copied" - - -def copy_tree_contents(src: Path, dst: Path, *, allow_hardlinks: bool = True) -> dict[str, int]: - if not src.exists(): - raise FileNotFoundError(f"fixture path does not exist: {src}") - stats = copy_stats() - dst.mkdir(parents=True, exist_ok=True) - for child in src.iterdir(): - target = dst / child.name - if child.is_dir(): - stats["dirs"] += 1 - merge_copy_stats(stats, copy_tree_contents(child, target, allow_hardlinks=allow_hardlinks)) - else: - mode = materialize_file(child, target, allow_hardlink=allow_hardlinks) - stats["files"] += 1 - stats[mode] += 1 - return stats - - -def materialize_fixture(target: RunTarget, *, dry_run: bool, preserve_state: bool, expected_region_dir: str | None = None, fixture_dir: Path | None = None, reset_data: bool = True) -> dict[str, Any]: - fixture_dir = fixture_dir or target.fixture_dir - summary_path = fixture_dir / "summary.json" - materialize_root = target.datanode_data_dir if preserve_state else target.data_dir - if dry_run: - return {"status": "dry-run", "summary_path": str(summary_path), "data_dir": str(materialize_root)} - summary = json.loads(summary_path.read_text()) - object_store_dir = fixture_dir / "object-store" - manifest_dir = fixture_dir / "manifest" - region_dir = summary["region_dir"].strip("/") - target_region_dir = materialize_root / region_dir - data_root = materialize_root.resolve() - resolved_region_dir = target_region_dir.resolve(strict=False) - if data_root != resolved_region_dir and data_root not in resolved_region_dir.parents: - raise RuntimeError(f"unsafe fixture region_dir escapes data_dir: {region_dir}") - if expected_region_dir is not None and region_dir != expected_region_dir.strip("/"): - raise RuntimeError(f"fixture region_dir {region_dir} does not equal discovered {expected_region_dir}") - target_manifest_dir = target_region_dir / "manifest" - if not preserve_state and reset_data and materialize_root.exists(): - shutil.rmtree(materialize_root) - materialize_root.mkdir(parents=True, exist_ok=True) - if preserve_state and target_region_dir.exists(): - shutil.rmtree(target_region_dir) - object_stats: dict[str, int] - if preserve_state: - fixture_region_dir = object_store_dir / region_dir - object_stats = copy_tree_contents(fixture_region_dir, target_region_dir, allow_hardlinks=True) - else: - object_stats = copy_tree_contents(object_store_dir, materialize_root, allow_hardlinks=True) - if target_manifest_dir.exists(): - shutil.rmtree(target_manifest_dir) - manifest_stats = copy_tree_contents(manifest_dir, target_manifest_dir, allow_hardlinks=False) - return {"status": "ok", "summary_path": str(summary_path), "data_dir": str(materialize_root), "region_dir": region_dir, "manifest_dir": str(target_manifest_dir), "preserve_state": preserve_state, "copy_stats": {"object_store": object_stats, "manifest": manifest_stats}} - - -def response_text(body: Any) -> str: - return json.dumps(body, sort_keys=True) if not isinstance(body, str) else body - - -def validate_show_create(result: dict[str, Any], table: dict[str, Any]) -> list[str]: - text = response_text(result.get("response", {})).lower() - errors = [] - if table["name"].lower() not in text: - errors.append("SHOW CREATE output does not contain table name") - if table.get("validate_show_create_engine", True) and ("engine" not in text or "mito" not in text): - errors.append("SHOW CREATE output does not mention ENGINE=mito") - if "append_mode" in table and "append_mode" not in text: - errors.append("SHOW CREATE output does not mention append_mode") - if table.get("sst_format") and "sst_format" not in text: - errors.append("SHOW CREATE output does not mention sst_format") - return errors - - -def run_queries(target: RunTarget, case: dict[str, Any], tables: list[dict[str, Any]], http_timeout: float) -> dict[str, Any]: - table = tables[0] - db = table["database"] - queries = planned_queries(case) - if not queries: - queries = [{"name": "count_all", "kind": "sql", "query": f"SELECT count(*) FROM {sql_ident(table['name'])}", "warmup": 0, "iterations": 1}] - validations = [] - validation_errors = [] - for table in tables: - for sql in [f"SHOW CREATE TABLE {sql_ident(table['name'])}"]: - result = http_post_sql(target.http_port, sql, table["database"], http_timeout) - if not result["ok"]: - validation_errors.append({"sql": sql, "error": result.get("error"), "response": result.get("response")}) - else: - for error in validate_show_create(result, table): - validation_errors.append({"sql": sql, "error": error, "response": result.get("response")}) - validations.append(result) - if queries: - result = http_post_sql(target.http_port, queries[0]["query"], db, http_timeout) - if not result["ok"]: - validation_errors.append({"sql": queries[0]["query"], "error": result.get("error"), "response": result.get("response")}) - validations.append(result) - measurements = [] - for q in queries: - for _ in range(int(q.get("warmup", 0))): - warmup = http_post_sql(target.http_port, q["query"], db, http_timeout) - if not warmup["ok"]: - validation_errors.append({"sql": q["query"], "phase": "warmup", "error": warmup.get("error"), "response": warmup.get("response")}) - samples = [] - for _ in range(int(q.get("iterations", 1))): - result = http_post_sql(target.http_port, q["query"], db, http_timeout) - result["execution_time_ms"] = extract_execution_time(result.get("response")) - samples.append(result) - good_lats = [s["latency_ms"] for s in samples if s["ok"]] - measurements.append({"name": q.get("name"), "kind": q.get("kind"), "iterations": len(samples), "samples": samples, "latency_ms_median": statistics.median(good_lats) if good_lats else None, "latency_ms_p95": percentile(good_lats, 95) if good_lats else None, "status": "ok" if len(good_lats) == len(samples) else "failed"}) - return {"validation": validations, "validation_errors": validation_errors, "measurements": measurements, "status": "failed" if validation_errors or any(m["status"] == "failed" for m in measurements) else "ok"} - - -def extract_execution_time(body: Any) -> Any: - if isinstance(body, dict): - for key in ("execution_time_ms", "execution_time", "elapsed"): - if key in body: - return body[key] - for value in body.values(): - found = extract_execution_time(value) - if found is not None: - return found - if isinstance(body, list): - for value in body: - found = extract_execution_time(value) - if found is not None: - return found - return None - - -def percentile(values: list[float], pct: float) -> float: - if not values: - return 0.0 - ordered = sorted(values) - idx = min(len(ordered) - 1, max(0, round((pct / 100.0) * (len(ordered) - 1)))) - return ordered[idx] - - -def enforce_thresholds(case: dict[str, Any], base: dict[str, Any], candidate: dict[str, Any]) -> list[dict[str, Any]]: - results = [] - base_by_name = {m["name"]: m for m in base.get("measurements", [])} - for cm in candidate.get("measurements", []): - bm = base_by_name.get(cm["name"]) - qcfg = next((q for q in planned_queries(case) if q.get("name") == cm["name"]), {}) - th = qcfg.get("thresholds") or {} - max_reg = th.get("max_candidate_latency_regression_pct") - if max_reg is not None and not bm: - results.append({"query": cm["name"], "threshold": "max_candidate_latency_regression_pct", "status": "failed", "reason": "missing base measurement"}) - elif max_reg is not None and bm and bm.get("latency_ms_median") in (None, 0): - results.append({"query": cm["name"], "threshold": "max_candidate_latency_regression_pct", "status": "failed", "reason": "base median latency is missing or zero", "base_latency_ms_median": bm.get("latency_ms_median")}) - elif max_reg is not None and bm and cm.get("latency_ms_median") is None: - results.append({"query": cm["name"], "threshold": "max_candidate_latency_regression_pct", "status": "failed", "reason": "missing candidate measurement"}) - elif bm and max_reg is not None: - ratio = (cm["latency_ms_median"] - bm["latency_ms_median"]) / bm["latency_ms_median"] * 100.0 - results.append({"query": cm["name"], "threshold": "max_candidate_latency_regression_pct", "status": "passed" if ratio <= max_reg else "failed", "actual_pct": ratio, "limit_pct": max_reg}) - for key in th: - if key != "max_candidate_latency_regression_pct": - results.append({"query": cm["name"], "threshold": key, "status": "failed", "reason": "unsupported threshold"}) - return results - - -def storage_config(remote: dict[str, Any]) -> dict[str, Any] | None: - value = remote["storage"] - if not isinstance(value, dict): - return None - if value["inspect"] is False: - return None - return value - - -def run_storage_inspection(helper: Path | None, target: RunTarget, storage: dict[str, Any], *, dry_run: bool) -> dict[str, Any]: - if helper is None: - if not dry_run: - raise ValueError("--fixture-generator query_perf_fixture is required when storage inspection is enabled") - helper = Path("query_perf_fixture") - root = target.datanode_data_dir - if storage.get("root_suffix"): - root = root / str(storage["root_suffix"]) - cmd = [ - str(helper), - "inspect-footer", - "--root", str(root), - "--column", str(storage["column"]), - ] - if storage["include_metadata_files"]: - cmd.append("--include-metadata-files") - if dry_run: - return {"status": "dry-run", "cmd": cmd, "root": str(root)} - result = run_command(cmd) - result["status"] = "ok" if result["returncode"] == 0 else "failed" - if result["returncode"] != 0: - raise RuntimeError(f"storage inspector failed for {target.name}: {result['stderr'][:2000]}") - try: - result["summary"] = json.loads(result["stdout"]) - except json.JSONDecodeError: - result["summary_parse_error"] = result["stdout"] - result["status"] = "failed" - return result - - -def storage_summary(inspection: dict[str, Any]) -> dict[str, Any]: - summary = inspection.get("summary") - if isinstance(summary, dict): - nested = summary.get("summary") - if isinstance(nested, dict): - return nested - return {} - - -def enforce_storage_thresholds(storage: dict[str, Any] | None, base: dict[str, Any] | None, candidate: dict[str, Any] | None) -> list[dict[str, Any]]: - if not storage: - return [] - results: list[dict[str, Any]] = [] - target_summaries = [("base", storage_summary(base or {})), ("candidate", storage_summary(candidate or {}))] - cand = target_summaries[1][1] - base_summary = storage_summary(base or {}) - - def add_abs(name: str, field: str) -> None: - limit = storage.get(name) - if limit is None: - return - for target_name, summary in target_summaries: - actual = summary.get(field) - ok = actual is not None and float(actual) <= float(limit) - results.append({"target": target_name, "threshold": name, "status": "passed" if ok else "failed", "actual": actual, "limit": limit}) - - def add_min(name: str, field: str) -> None: - limit = storage[name] - if limit is None: - return - for target_name, summary in target_summaries: - actual = summary.get(field) - ok = actual is not None and int(actual) >= int(limit) - results.append({"target": target_name, "threshold": name, "status": "passed" if ok else "failed", "actual": actual, "limit": limit}) - - def add_cmp(name: str, field: str) -> None: - limit = storage.get(name) - if limit is None: - return - base_value = base_summary.get(field) - candidate_value = cand.get(field) - if base_value in (None, 0) or candidate_value is None: - results.append({"threshold": name, "status": "failed", "reason": "missing or zero base/candidate value", "base": base_value, "candidate": candidate_value}) - return - actual = (float(candidate_value) - float(base_value)) / float(base_value) * 100.0 - results.append({"threshold": name, "status": "passed" if actual <= float(limit) else "failed", "actual_pct": actual, "limit_pct": limit, "base": base_value, "candidate": candidate_value}) - - add_min("min_files", "file_count") - add_min("min_files_with_column", "files_with_column") - add_abs("max_total_file_size_bytes", "total_file_size") - add_abs("max_column_compressed_size_bytes", "column_compressed_size") - add_abs("max_column_uncompressed_size_bytes", "column_uncompressed_size") - for target_name, summary in target_summaries: - encodings = set(summary.get("unique_encodings") or []) - for required in storage["require_encodings"]: - results.append({"target": target_name, "threshold": "require_encodings", "encoding": required, "status": "passed" if required in encodings else "failed"}) - for forbidden in storage["forbid_encodings"]: - results.append({"target": target_name, "threshold": "forbid_encodings", "encoding": forbidden, "status": "passed" if forbidden not in encodings else "failed"}) - add_cmp("max_candidate_total_file_size_regression_pct", "total_file_size") - add_cmp("max_candidate_column_compressed_size_regression_pct", "column_compressed_size") - add_cmp("max_candidate_column_uncompressed_size_regression_pct", "column_uncompressed_size") - return results - - -def planned_storage_thresholds(storage: dict[str, Any] | None) -> list[dict[str, Any]]: - if not storage: - return [] - return list(storage["planned_thresholds"]) - - -REGION_DIR_RE = re.compile(r"^(\d+)_(\d{10})$") - - -def inspection_report(inspection: dict[str, Any] | None) -> dict[str, Any]: - if not inspection: - return {} - report = inspection.get("summary") - return report if isinstance(report, dict) else {} - - -def inspected_relative_path(target: RunTarget, inspection: dict[str, Any], relative_path: str) -> Path: - report = inspection_report(inspection) - root_text = report.get("root") or inspection.get("root") - if not root_text: - return Path(relative_path) - root = Path(root_text).resolve(strict=False) - data_home = target.datanode_data_dir.resolve(strict=False) - try: - root_suffix = root.relative_to(data_home) - except ValueError as e: - raise RuntimeError(f"storage inspection root {root} is not under datanode data home {data_home}") from e - return root_suffix / relative_path - - -def parse_sst_bench_target(target: RunTarget, inspection: dict[str, Any], file: dict[str, Any], *, dry_run: bool) -> dict[str, str] | None: - relative_path = str(file.get("relative_path") or "") - if dry_run and (not relative_path or "<" in relative_path): - return { - "relative_path": relative_path or "/_/data/.parquet", - "table_dir": "/", - "region_id": ":", - "path_type": "data", - "file_id": "", - } - if not relative_path: - return None - full_relative = inspected_relative_path(target, inspection, relative_path) - parts = full_relative.parts - if not parts or not parts[-1].endswith(".parquet"): - return None - region_index = next((idx for idx, part in enumerate(parts) if REGION_DIR_RE.match(part)), None) - if region_index is None or region_index == 0: - if dry_run: - return { - "relative_path": relative_path, - "table_dir": "/", - "region_id": ":", - "path_type": "data", - "file_id": Path(parts[-1]).stem, - } - return None - region_match = REGION_DIR_RE.match(parts[region_index]) - assert region_match is not None - file_index = len(parts) - 1 - path_type = "bare" - if region_index + 1 < file_index and parts[region_index + 1] in ("data", "metadata"): - path_type = parts[region_index + 1] - if path_type == "metadata": - return None - table_dir = "/".join(parts[:region_index]).rstrip("/") + "/" - return { - "relative_path": str(full_relative), - "table_dir": table_dir, - "region_id": f"{int(region_match.group(1))}:{int(region_match.group(2))}", - "path_type": path_type, - "file_id": Path(parts[-1]).stem, - } - - -def run_read_bench(bench_binary: Path, target: RunTarget, read_bench: dict[str, Any] | None, storage_inspection: dict[str, Any] | None, *, dry_run: bool) -> dict[str, Any]: - if not read_bench or read_bench["enabled"] is False: - return {"status": "skipped", "reason": "read_bench disabled"} - run_parquetbench = read_bench["parquetbench"] - run_scanbench = read_bench["scanbench"] - if not run_parquetbench and not run_scanbench: - return {"status": "skipped", "reason": "parquetbench and scanbench disabled"} - if storage_inspection is None: - raise ValueError("read_bench requires storage inspection") - bench_dir = target.work_dir / "read_bench" - config_toml = bench_dir / "bench.toml" - scan_json = bench_dir / "scan.json" - files = inspection_report(storage_inspection).get("files") or [] - ssts = [f for f in files if f.get("relative_path", "").endswith(".parquet") and f.get("columns")] - if dry_run and not ssts: - ssts = [{"relative_path": ".parquet"}] - if read_bench["max_files"] is not None: - ssts = ssts[: int(read_bench["max_files"])] - bench_targets = [parsed for f in ssts if (parsed := parse_sst_bench_target(target, storage_inspection, f, dry_run=dry_run)) is not None] - if not bench_targets: - return {"status": "failed", "reason": "no inspected data SST files available for read_bench"} - config_text = f'[storage]\ndata_home = "{target.datanode_data_dir}"\ntype = "File"\n\n[[region_engine]]\n[region_engine.mito]\n' - commands = [] - parquet_runs = [] - scan_runs = [] - iterations = str(int(read_bench["iterations"])) - if run_parquetbench: - for bench_target in bench_targets: - cmd = [str(bench_binary), "datanode", "parquetbench", "--config", str(config_toml), "--region-id", bench_target["region_id"], "--table-dir", bench_target["table_dir"], "--file-id", bench_target["file_id"], "--scan-config", str(scan_json), "--path-type", bench_target["path_type"], "--iterations", iterations, "--reader", str(read_bench["parquet_reader"])] - commands.append(cmd) - parquet_runs.append({**bench_target, "cmd": cmd, "status": "dry-run" if dry_run else "planned"}) - regions: dict[tuple[str, str, str], list[str]] = {} - if run_scanbench: - for bench_target in bench_targets: - regions.setdefault((bench_target["table_dir"], bench_target["region_id"], bench_target["path_type"]), []).append(bench_target["relative_path"]) - for (table_dir, region_id, path_type), region_files in regions.items(): - cmd = [str(bench_binary), "datanode", "scanbench", "--config", str(config_toml), "--region-id", region_id, "--table-dir", table_dir, "--scan-config", str(scan_json), "--path-type", path_type, "--scanner", str(read_bench["scan_scanner"]), "--parallelism", str(int(read_bench["parallelism"])), "--iterations", iterations] - commands.append(cmd) - scan_runs.append({"table_dir": table_dir, "region_id": region_id, "path_type": path_type, "files": region_files, "cmd": cmd, "status": "dry-run" if dry_run else "planned"}) - if dry_run: - return {"status": "dry-run", "config_path": str(config_toml), "scan_config_path": str(scan_json), "commands": commands, "parquetbench": parquet_runs, "scanbench": scan_runs} - bench_dir.mkdir(parents=True, exist_ok=True) - config_toml.write_text(config_text) - scan_json.write_text(json.dumps({"projection_names": read_bench["projection"]}, indent=2) + "\n") - avg_re = re.compile(r"Average duration[^0-9]*([0-9.]+)\s*ms", re.I) - for runs in (parquet_runs, scan_runs): - for run in runs: - result = run_command(run["cmd"]) - run.update(result) - run["status"] = "ok" if result["returncode"] == 0 else "failed" - m = avg_re.search(result.get("stdout", "")) - if m: - run["average_ms"] = float(m.group(1)) - def median(runs: list[dict[str, Any]]) -> float | None: - vals = [r["average_ms"] for r in runs if "average_ms" in r] - return statistics.median(vals) if vals else None - return {"status": "ok" if all(r.get("status") == "ok" for r in parquet_runs + scan_runs) else "failed", "config_path": str(config_toml), "scan_config_path": str(scan_json), "parquetbench": parquet_runs, "scanbench": scan_runs, "aggregate": {"parquetbench_median_average_ms": median(parquet_runs), "scanbench_median_average_ms": median(scan_runs)}} - - -def write_json(path: Path, payload: dict[str, Any]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") - - -def output_report(report: dict[str, Any], output: Path | None) -> None: - if output is None: - print(json.dumps(report, indent=2, sort_keys=True)) - else: - write_json(output, report) - - -def parse_prometheus_metrics(text: str, names: set[str] = OTLP_TRACE_METRICS) -> dict[str, float]: - values: dict[str, float] = {} - for line in text.splitlines(): - match = PROMETHEUS_SAMPLE_RE.match(line.strip()) - if not match or match.group(1) not in names: - continue - value = float(match.group(2)) - if not math.isfinite(value): - raise ValueError(f"non-finite Prometheus sample for {match.group(1)}") - values[match.group(1)] = values.get(match.group(1), 0.0) + value - return values - - -def fetch_otlp_metrics(target: RunTarget, http_timeout: float) -> dict[str, Any]: - with urllib.request.urlopen(f"http://127.0.0.1:{target.http_port}/metrics", timeout=http_timeout) as response: - text = response.read().decode() - return {"captured_monotonic_seconds": time.monotonic(), "values": parse_prometheus_metrics(text)} - - -def metric_delta(after: dict[str, Any], before: dict[str, Any], name: str) -> float: - delta = float(after["values"].get(name, 0.0)) - float(before["values"].get(name, 0.0)) - if delta < 0: - raise RuntimeError(f"metric {name} decreased by {-delta}") - return delta - - -def otelgen_command(otelgen_bin: Path, target: RunTarget, load: dict[str, Any]) -> list[str]: - return [ - str(otelgen_bin), - "--protocol", "http", - "--otel-exporter-otlp-endpoint", f"127.0.0.1:{target.http_port}", - "--otel-exporter-otlp-url-path", "/v1/otlp/v1/traces", - "--header", f"x-greptime-pipeline-name={load['pipeline']}", - "--header", f"x-greptime-db-name={load['database']}", - "--header", f"x-greptime-trace-table-name={load['table']}", - "--log-level", "error", - "--insecure", - "--duration", str(int(load["duration_seconds"])), - "--rate", str(int(load["rate"])), - "traces", "multi", - "--workers", str(int(load["workers"])), - "--scenarios", str(load["workload"]), - "--exporter-shards", str(int(load["exporter_shards"])), - ] - - -def run_otelgen_load(otelgen_bin: Path | None, target: RunTarget, load: dict[str, Any], http_timeout: float, *, dry_run: bool) -> dict[str, Any]: - binary = otelgen_bin or Path("otelgen") - cmd = otelgen_command(binary, target, load) - if dry_run: - return {"status": "dry-run", "cmd": cmd} - - initial = fetch_otlp_metrics(target, http_timeout) - log_dir = target.work_dir / "otelgen" - log_dir.mkdir(parents=True, exist_ok=True) - stdout_path = log_dir / "stdout.log" - stderr_path = log_dir / "stderr.log" - duration = int(load["duration_seconds"]) - warmup = int(load["warmup_seconds"]) - timed_out = False - started = time.monotonic() - with stdout_path.open("wb") as stdout, stderr_path.open("wb") as stderr: - proc = subprocess.Popen(cmd, stdout=stdout, stderr=stderr) - try: - if warmup: - try: - proc.wait(timeout=warmup) - except subprocess.TimeoutExpired: - pass - warmed = fetch_otlp_metrics(target, http_timeout) - if proc.poll() is None: - try: - proc.wait(timeout=max(60, duration - warmup + 60)) - except subprocess.TimeoutExpired: - timed_out = True - proc.kill() - proc.wait(timeout=20) - final = fetch_otlp_metrics(target, http_timeout) - finally: - if proc.poll() is None: - proc.kill() - proc.wait(timeout=20) - elapsed = time.monotonic() - started - return { - "status": "ok" if proc.returncode == 0 and not timed_out and elapsed + 1 >= duration else "failed", - "cmd": cmd, - "returncode": proc.returncode, - "timed_out": timed_out, - "elapsed_seconds": elapsed, - "stdout_path": str(stdout_path), - "stderr_path": str(stderr_path), - "snapshots": {"initial": initial, "warmup": warmed, "final": final}, - } - - -def summarize_otlp_metrics(run: dict[str, Any]) -> dict[str, Any]: - initial = run["snapshots"]["initial"] - warmed = run["snapshots"]["warmup"] - final = run["snapshots"]["final"] - rows = "greptime_frontend_otlp_traces_rows" - failures = "greptime_frontend_otlp_traces_failure_count" - elapsed_sum = "greptime_servers_http_otlp_traces_elapsed_sum" - elapsed_count = "greptime_servers_http_otlp_traces_elapsed_count" - missing = sorted(name for name in (rows, elapsed_sum, elapsed_count) if name not in final["values"]) - accepted_spans = int(round(metric_delta(final, initial, rows))) - measurement_accepted_spans = int(round(metric_delta(final, warmed, rows))) - http_requests = int(round(metric_delta(final, warmed, elapsed_count))) - latency_seconds = metric_delta(final, warmed, elapsed_sum) - measurement_seconds = final["captured_monotonic_seconds"] - warmed["captured_monotonic_seconds"] - return { - "accepted_spans": accepted_spans, - "measurement_accepted_spans": measurement_accepted_spans, - "accepted_spans_per_second": measurement_accepted_spans / measurement_seconds if measurement_seconds > 0 else None, - "http_requests": http_requests, - "mean_http_latency_ms": latency_seconds / http_requests * 1000.0 if http_requests else None, - "failure_count": int(round(metric_delta(final, initial, failures))), - "measurement_seconds": measurement_seconds, - "missing_metrics": missing, - } - - -def planned_otlp_thresholds(load: dict[str, Any]) -> list[dict[str, Any]]: - thresholds = load["thresholds"] - return [ - {"threshold": "max_candidate_throughput_regression_pct", "status": "planned", "limit_pct": thresholds["max_candidate_throughput_regression_pct"]}, - {"threshold": "max_candidate_mean_latency_regression_pct", "status": "planned", "limit_pct": thresholds["max_candidate_mean_latency_regression_pct"]}, - {"target": "each", "threshold": "max_failure_count", "status": "planned", "limit": thresholds["max_failure_count"]}, - ] - - -def enforce_otlp_thresholds(load: dict[str, Any], base: dict[str, Any], candidate: dict[str, Any]) -> list[dict[str, Any]]: - thresholds = load["thresholds"] - results: list[dict[str, Any]] = [] - for target_name, metrics in (("base", base), ("candidate", candidate)): - failures = metrics.get("failure_count") - limit = int(thresholds["max_failure_count"]) - results.append({"target": target_name, "threshold": "max_failure_count", "status": "passed" if failures is not None and failures <= limit else "failed", "actual": failures, "limit": limit}) - - base_rate = base.get("accepted_spans_per_second") - candidate_rate = candidate.get("accepted_spans_per_second") - throughput_limit = float(thresholds["max_candidate_throughput_regression_pct"]) - if base_rate in (None, 0) or candidate_rate is None: - results.append({"threshold": "max_candidate_throughput_regression_pct", "status": "failed", "reason": "missing or zero throughput", "base": base_rate, "candidate": candidate_rate}) - else: - actual = (base_rate - candidate_rate) / base_rate * 100.0 - results.append({"threshold": "max_candidate_throughput_regression_pct", "status": "passed" if actual <= throughput_limit else "failed", "actual_pct": actual, "limit_pct": throughput_limit, "base": base_rate, "candidate": candidate_rate}) - - base_latency = base.get("mean_http_latency_ms") - candidate_latency = candidate.get("mean_http_latency_ms") - latency_limit = float(thresholds["max_candidate_mean_latency_regression_pct"]) - if base_latency in (None, 0) or candidate_latency is None: - results.append({"threshold": "max_candidate_mean_latency_regression_pct", "status": "failed", "reason": "missing or zero mean latency", "base": base_latency, "candidate": candidate_latency}) - else: - actual = (candidate_latency - base_latency) / base_latency * 100.0 - results.append({"threshold": "max_candidate_mean_latency_regression_pct", "status": "passed" if actual <= latency_limit else "failed", "actual_pct": actual, "limit_pct": latency_limit, "base": base_latency, "candidate": candidate_latency}) - return results - - -def require_fresh_work_dirs(targets: list[RunTarget], *, reuse_work_dir: bool, dry_run: bool, fixture_only: bool) -> None: - if dry_run or reuse_work_dir or fixture_only: - return - for target in targets: - if target.work_dir.exists() and any(target.work_dir.iterdir()): - raise RuntimeError(f"target work_dir exists and is non-empty: {target.work_dir}; pass --reuse-work-dir to override") - - -def write_frontend_prom_config(target: RunTarget, remote: dict[str, Any]) -> Path: - prom = remote["prom_store"] - path = target.work_dir / "frontend-prom-store.toml" - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text( - "[prom_store]\n" - "enable = true\n" - "with_metric_engine = true\n" - f"pending_rows_flush_interval = {json.dumps(str(prom['pending_rows_flush_interval']))}\n" - f"max_batch_rows = {int(prom['max_batch_rows'])}\n" - f"max_concurrent_flushes = {int(prom['max_concurrent_flushes'])}\n" - f"worker_channel_capacity = {int(prom['worker_channel_capacity'])}\n" - f"max_inflight_requests = {int(prom['max_inflight_requests'])}\n" - ) - return path - - -def run_remote_write(generator: Path | None, target: RunTarget, remote: dict[str, Any], *, dry_run: bool) -> dict[str, Any]: - if generator is None: - if not dry_run: - raise ValueError("--fixture-generator query_perf_fixture is required for prom_remote_write_then_query") - generator = Path("query_perf_fixture") - metric = remote["metric"] - physical_table = remote["physical_table"] - cmd = [ - str(generator), - "prom-remote-write", - "--endpoint", f"http://127.0.0.1:{target.http_port}/v1/prometheus/write", - "--database", remote["database"], - "--metric", metric, - "--physical-table", physical_table, - "--series-count", str(int(remote["series_count"])), - "--samples-per-series", str(int(remote["samples_per_series"])), - "--start-unix-millis", str(int(remote["start_unix_millis"])), - "--step-millis", str(int(remote["step_millis"])), - "--chunk-series-count", str(int(remote["chunk_series_count"])), - "--timeout-seconds", str(int(remote["timeout_seconds"])), - ] - value = remote["value"] - cmd.extend(["--value-pattern", str(value["pattern"])]) - cmd.extend(["--value-base", str(value["base"])]) - cmd.extend(["--value-step", str(value["step"])]) - cmd.extend(["--value-cardinality", str(int(value["cardinality"]))]) - cmd.extend(["--value-seed", str(int(value["seed"]))]) - cmd.extend(["--value-run-length", str(int(value["run_length"]))]) - cmd.extend(["--value-stall-every", str(int(value["stall_every"]))]) - cmd.extend(["--value-stall-length", str(int(value["stall_length"]))]) - cmd.extend(["--value-mixed-every", str(int(value["mixed_every"]))]) - if "sample_offset" in remote: - cmd.extend(["--value-sample-offset", str(int(remote["sample_offset"]))]) - if "total_samples_per_series" in remote: - cmd.extend(["--value-total-samples-per-series", str(int(remote["total_samples_per_series"]))]) - if dry_run: - return {"status": "dry-run", "cmd": cmd} - result = run_command(cmd) - result["status"] = "ok" if result["returncode"] == 0 else "failed" - if result["returncode"] != 0: - raise RuntimeError(f"remote-write generator failed for {target.name}: {result['stderr'][:2000]}") - try: - result["summary"] = json.loads(result["stdout"]) - except json.JSONDecodeError: - result["summary_parse_error"] = result["stdout"] - return result - - -def summarize_remote_write_chunks(chunks: list[dict[str, Any]]) -> dict[str, Any]: - summary = {"rows": 0, "samples_written": 0, "batches": 0, "elapsed_seconds": 0.0} - for chunk in chunks: - chunk_summary = chunk.get("summary") or {} - summary["rows"] += int(chunk_summary.get("rows", 0)) - summary["samples_written"] += int(chunk_summary.get("samples_written", chunk_summary.get("rows", 0))) - summary["batches"] += int(chunk_summary.get("batches", 0)) - summary["elapsed_seconds"] += float(chunk_summary.get("elapsed_seconds", chunk.get("elapsed_seconds", 0.0))) - return summary - - -def flush_remote_physical_table(target: RunTarget, db: str, physical_table: str, args: argparse.Namespace, *, dry_run: bool, reason: str, chunk_index: int | None = None) -> dict[str, Any]: - if dry_run: - return {"status": "dry-run", "physical_table": physical_table, "reason": reason, "chunk_index": chunk_index} - result = http_post_sql(target.http_port, f"ADMIN FLUSH_TABLE({sql_string(physical_table)})", db, args.http_timeout) - result["physical_table"] = physical_table - result["reason"] = reason - result["chunk_index"] = chunk_index - return result - - -def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote: dict[str, Any], args: argparse.Namespace, *, dry_run: bool) -> tuple[dict[str, Any], list[dict[str, Any]]]: - sample_chunk_size = remote["sample_chunk_size"] - db = remote["database"] - physical_table = remote["physical_table"] - if sample_chunk_size is None: - rw = run_remote_write(generator, target, remote, dry_run=dry_run) - flush = flush_remote_physical_table(target, db, physical_table, args, dry_run=dry_run, reason="final") - return rw, [flush] - - total_samples = int(remote["samples_per_series"]) - chunk_samples = int(sample_chunk_size) - if chunk_samples <= 0: - raise ValueError("scenario.remote_write.sample_chunk_size must be positive") - flush_every = int(remote["flush_every_sample_chunks"]) - if flush_every <= 0: - raise ValueError("scenario.remote_write.flush_every_sample_chunks must be positive") - start = int(remote["start_unix_millis"]) - step = int(remote["step_millis"]) - chunks: list[dict[str, Any]] = [] - flushes: list[dict[str, Any]] = [] - chunk_index = 0 - last_flushed_chunk = 0 - for offset in range(0, total_samples, chunk_samples): - current = min(chunk_samples, total_samples - offset) - chunk_index += 1 - chunk_remote = dict(remote) - chunk_remote["samples_per_series"] = current - chunk_remote["start_unix_millis"] = start + offset * step - chunk_remote["sample_offset"] = offset - chunk_remote["total_samples_per_series"] = total_samples - result = run_remote_write(generator, target, chunk_remote, dry_run=dry_run) - result["sample_offset"] = offset - result["samples_per_series"] = current - result["chunk_index"] = chunk_index - chunks.append(result) - if chunk_index % flush_every == 0: - flushes.append(flush_remote_physical_table(target, db, physical_table, args, dry_run=dry_run, reason="periodic", chunk_index=chunk_index)) - last_flushed_chunk = chunk_index - if last_flushed_chunk != chunk_index: - flushes.append(flush_remote_physical_table(target, db, physical_table, args, dry_run=dry_run, reason="final", chunk_index=chunk_index)) - aggregate = summarize_remote_write_chunks(chunks) - if dry_run: - series_count = int(remote["series_count"]) - chunk_series_count = int(remote["chunk_series_count"]) - aggregate["rows"] = series_count * total_samples - aggregate["samples_written"] = aggregate["rows"] - aggregate["batches"] = chunk_index * ((series_count + chunk_series_count - 1) // chunk_series_count) - return { - "status": "dry-run" if dry_run else ("ok" if all(chunk.get("status") == "ok" for chunk in chunks) else "failed"), - "mode": "sample-chunked", - "sample_chunk_size": chunk_samples, - "flush_every_sample_chunks": flush_every, - "chunks": chunks, - "aggregate": aggregate, - }, flushes - - -def expected_remote_write_rows(remote: dict[str, Any]) -> int: - return int(remote["series_count"]) * int(remote["samples_per_series"]) - - -def extract_count_value(result: dict[str, Any]) -> int | None: - body = result.get("response") - if not isinstance(body, dict): - return None - data = body.get("data") - if not isinstance(data, list) or not data: - return None - row = data[0] - if not isinstance(row, dict): - return None - for key, value in row.items(): - if key.lower() == "count(*)" or key.lower().startswith("count("): - try: - return int(value) - except (TypeError, ValueError): - return None - return None - - -def poll_expected_count(target: RunTarget, table_name: str, db: str, expected_rows: int, timeout_s: float, http_timeout: float) -> dict[str, Any]: - sql = f"SELECT count(*) FROM {sql_ident(table_name)}" - deadline = time.monotonic() + timeout_s - attempts = 0 - last: dict[str, Any] | None = None - while True: - attempts += 1 - result = http_post_sql(target.http_port, sql, db, http_timeout) - observed_rows = extract_count_value(result) - result["expected_rows"] = expected_rows - result["observed_rows"] = observed_rows - result["attempts"] = attempts - result["row_count_ok"] = result.get("ok") and observed_rows == expected_rows - if result["row_count_ok"]: - return result - last = result - if time.monotonic() >= deadline: - break - time.sleep(0.5) - assert last is not None - last["ok"] = False - last["row_count_ok"] = False - last["error"] = f"expected {expected_rows} rows but observed {last.get('observed_rows')} after {attempts} attempts" - return last - - -def run_remote_write_scenario(args: argparse.Namespace, case: dict[str, Any], case_path: Path, targets: list[RunTarget], report: dict[str, Any]) -> None: - remote = scenario(case)["remote_write"] - tables = case_tables(case) - if args.fixture_only: - raise ValueError("--fixture-only is not supported for prom_remote_write_then_query; use --dry-run for planning") - helper = args.fixture_generator or args.remote_write_generator - if helper is not None: - require_binary(helper, "query_perf_fixture", dry_run=args.dry_run) - elif not args.dry_run: - raise ValueError("--fixture-generator is required for prom_remote_write_then_query") - storage = storage_config(remote) - if storage and helper is None and not args.dry_run: - raise ValueError("--fixture-generator is required when scenario.remote_write.storage.inspect is enabled") - clusters: list[DistributedCluster] = [] - target_results = [] - storage_results = [] - try: - for target in targets: - target.work_dir.mkdir(parents=True, exist_ok=True) - config_path = write_frontend_prom_config(target, remote) - cluster = DistributedCluster(target) - clusters.append(cluster) - if not args.dry_run: - cluster.start_all(config_path) - db = remote["database"] - create_database = {"status": "dry-run", "database": db} - if not args.dry_run: - create_database = http_post_sql(target.http_port, f"CREATE DATABASE IF NOT EXISTS {sql_ident(db)}", "public", args.http_timeout) - rw, flushes = run_remote_write_ingestion(helper, target, remote, args, dry_run=args.dry_run) - visibility = {"status": "dry-run"} - query_result = {"validation": [], "validation_errors": [], "measurements": [], "status": "planned"} - if not args.dry_run: - visibility = poll_expected_count(target, tables[0]["name"], db, expected_remote_write_rows(remote), float(remote["visibility_timeout_seconds"]), args.http_timeout) - query_result = run_queries(target, case, tables, args.http_timeout) - cluster.stop_component("datanode") - storage_inspection = None - if storage: - storage_inspection = run_storage_inspection(helper or args.storage_inspector, target, storage, dry_run=args.dry_run) - storage_results.append(storage_inspection) - read_bench_result = run_read_bench(args.candidate_bin, target, remote["read_bench"], storage_inspection, dry_run=args.dry_run) - tr = {"name": target.name, "binary": str(target.binary), "work_dir": str(target.work_dir), "components": cluster.component_report(), "frontend_config": str(config_path), "create_database": create_database, "remote_write": rw, "flushes": flushes, "flush": flushes[-1] if flushes else None, "visibility": visibility, **query_result} - if storage_inspection is not None: - tr["storage_inspection"] = storage_inspection - tr["read_bench"] = read_bench_result - flushes_ok = all(flush.get("ok") for flush in flushes) - storage_ok = storage_inspection is None or args.dry_run or storage_inspection.get("status") == "ok" - read_bench_ok = args.dry_run or read_bench_result.get("status") in ("ok", "skipped") - remote_checks_ok = args.dry_run or (create_database.get("ok") and flushes_ok and visibility.get("ok") and visibility.get("row_count_ok") and storage_ok and read_bench_ok) - if not remote_checks_ok: - tr.setdefault("validation_errors", []).append({"phase": "remote_write_visibility_storage", "create_database_ok": create_database.get("ok"), "flushes_ok": flushes_ok, "visibility_ok": visibility.get("ok"), "row_count_ok": visibility.get("row_count_ok"), "storage_ok": storage_ok, "read_bench_ok": read_bench_ok, "create_database": create_database, "flushes": flushes, "visibility": visibility, "read_bench": read_bench_result}) - tr["status"] = "planned" if args.dry_run else ("measured" if remote_checks_ok and query_result["status"] == "ok" else "failed") - write_json(target.report_path, tr) - report["targets"].append(tr) - target_results.append(query_result) - if not args.dry_run: - report["thresholds"] = enforce_thresholds(case, target_results[0], target_results[1]) + enforce_storage_thresholds(storage, storage_results[0] if storage_results else None, storage_results[1] if len(storage_results) > 1 else None) - elif storage: - report["thresholds"] = planned_storage_thresholds(storage) - report["status"] = "planned" if args.dry_run else ("failed" if any(t["status"] == "failed" for t in report["thresholds"]) or any(t.get("status") == "failed" for t in report["targets"]) else "ok") - finally: - for cluster in reversed(clusters): - cluster.stop_all() - - -def run_otlp_trace_load_scenario(args: argparse.Namespace, case: dict[str, Any], targets: list[RunTarget], report: dict[str, Any]) -> None: - load = scenario(case)["load"] - if args.fixture_only: - raise ValueError("--fixture-only is not supported for otlp_trace_load; use --dry-run for planning") - if args.otelgen_bin is not None: - require_binary(args.otelgen_bin, "otelgen", dry_run=args.dry_run) - elif not args.dry_run: - raise ValueError("--otelgen-bin is required for otlp_trace_load") - - clusters: list[DistributedCluster] = [] - metrics_results: list[dict[str, Any]] = [] - try: - for target in targets: - target.work_dir.mkdir(parents=True, exist_ok=True) - cluster = DistributedCluster(target) - clusters.append(cluster) - if not args.dry_run: - cluster.start_all() - create_database: dict[str, Any] = {"status": "dry-run", "database": load["database"]} - if not args.dry_run: - create_database = http_post_sql(target.http_port, f"CREATE DATABASE IF NOT EXISTS {sql_ident(load['database'])}", "public", args.http_timeout) - otelgen = run_otelgen_load(args.otelgen_bin, target, load, args.http_timeout, dry_run=args.dry_run) - metrics: dict[str, Any] = {"status": "dry-run"} - flush: dict[str, Any] = {"status": "dry-run", "table": load["table"]} - visibility: dict[str, Any] = {"status": "dry-run"} - if not args.dry_run: - metrics = summarize_otlp_metrics(otelgen) - flush = http_post_sql(target.http_port, f"ADMIN FLUSH_TABLE({sql_string(load['table'])})", load["database"], args.http_timeout) - visibility = poll_expected_count(target, load["table"], load["database"], int(metrics["accepted_spans"]), float(load["visibility_timeout_seconds"]), args.http_timeout) - tr = { - "name": target.name, - "binary": str(target.binary), - "work_dir": str(target.work_dir), - "components": cluster.component_report(), - "create_database": create_database, - "otelgen": otelgen, - "metrics": metrics, - "flush": flush, - "visibility": visibility, - } - checks_ok = args.dry_run or ( - create_database.get("ok") - and otelgen.get("status") == "ok" - and not metrics.get("missing_metrics") - and metrics.get("accepted_spans", 0) > 0 - and metrics.get("http_requests", 0) > 0 - and flush.get("ok") - and visibility.get("ok") - and visibility.get("row_count_ok") - ) - tr["status"] = "planned" if args.dry_run else ("measured" if checks_ok else "failed") - write_json(target.report_path, tr) - report["targets"].append(tr) - metrics_results.append(metrics) - cluster.stop_all() - report["thresholds"] = planned_otlp_thresholds(load) if args.dry_run else enforce_otlp_thresholds(load, metrics_results[0], metrics_results[1]) - report["status"] = "planned" if args.dry_run else ("failed" if any(t["status"] == "failed" for t in report["thresholds"]) or any(t["status"] == "failed" for t in report["targets"]) else "ok") - finally: - for cluster in reversed(clusters): - cluster.stop_all() - - -def main() -> int: - args = parse_args() - case_path = args.case.resolve() - if args.fixture_generator is not None: - require_binary(args.fixture_generator, "fixture generator", dry_run=False) - case = load_normalized_case(case_path, args.fixture_generator) - scenario_config = scenario(case) - scenario_kind = scenario_config.get("kind") - tables = case_tables(case) - work_root = args.work_dir.resolve() - work_root.mkdir(parents=True, exist_ok=True) - require_binary(args.base_bin, "base", dry_run=args.dry_run or args.fixture_only) - require_binary(args.candidate_bin, "candidate", dry_run=args.dry_run or args.fixture_only) - if scenario_kind == "direct_readable_sst" and not args.dry_run and args.fixture_generator is None: - raise ValueError("--fixture-generator is required unless --dry-run is set") - ports = allocate_ports(16) - targets = [make_target("base", args.base_bin.resolve(), work_root, ports[:8]), make_target("candidate", args.candidate_bin.resolve(), work_root, ports[8:])] - require_fresh_work_dirs(targets, reuse_work_dir=args.reuse_work_dir, dry_run=args.dry_run, fixture_only=args.fixture_only) - fixture_dir = fixture_root(work_root, case_path, case, args.fixture_cache_dir) - reuse_fixture = args.reuse_fixture or args.fixture_cache_dir is not None - report: dict[str, Any] = {"case_path": str(case_path), "case": case.get("case", {}), "scenario": scenario_config, "queries": planned_queries(case), "dry_run": args.dry_run, "fixture_only": args.fixture_only, "query_mode": "fixture-only" if args.fixture_only else "distributed", "reuse_work_dir": args.reuse_work_dir, "reuse_fixture": reuse_fixture, "fixture_cache_dir": str(args.fixture_cache_dir.resolve()) if args.fixture_cache_dir is not None else None, "fixture_dir": str(fixture_dir), "http_timeout": args.http_timeout, "targets": [], "thresholds": [], "status": "planned" if args.dry_run else "running"} - - if scenario_kind == "otlp_trace_load": - try: - run_otlp_trace_load_scenario(args, case, targets, report) - except Exception as e: # noqa: BLE001 - write machine-readable failure report - report["status"] = "failed" - report["error"] = repr(e) - write_json(work_root / "query-regression-report.json", report) - output_report(report, args.output) - return 1 if report["status"] == "failed" else 0 - - if scenario_kind == "prom_remote_write_then_query": - try: - run_remote_write_scenario(args, case, case_path, targets, report) - except Exception as e: # noqa: BLE001 - write machine-readable failure report - report["status"] = "failed" - report["error"] = repr(e) - write_json(work_root / "query-regression-report.json", report) - output_report(report, args.output) - return 1 if report["status"] == "failed" else 0 - - if args.fixture_only or args.dry_run: - generations = [] - for table_idx, table in enumerate(tables): - per_table_fixture_dir = table_fixture_dir(fixture_dir, tables, table, table_idx) - generations.append(generate_fixture(args.fixture_generator, case_path, per_table_fixture_dir, dry_run=args.dry_run, reuse_fixture=reuse_fixture, allow_large_fixture=args.allow_large_fixture, table_name=table["name"] if len(tables) > 1 else None, database=table["database"] if len(tables) > 1 else None)) - report["fixture_generation"] = generations[0] if len(generations) == 1 else generations - for target in targets: - target.work_dir.mkdir(parents=True, exist_ok=True) - mats = [] - for table_idx, table in enumerate(tables): - per_table_fixture_dir = table_fixture_dir(fixture_dir, tables, table, table_idx) - mats.append(materialize_fixture(target, dry_run=args.dry_run, preserve_state=False, fixture_dir=per_table_fixture_dir, reset_data=not mats)) - tr = {"name": target.name, "binary": str(target.binary), "work_dir": str(target.work_dir), "data_dir": str(target.data_dir), "fixture_dir": str(fixture_dir), "fixture_materialization": mats[0] if len(mats) == 1 else mats, "measurements": [], "status": "planned" if args.dry_run else "fixture-ready"} - write_json(target.report_path, tr) - report["targets"].append(tr) - report["status"] = "planned" if args.dry_run else "fixture-ready" - write_json(work_root / "query-regression-report.json", report) - output_report(report, args.output) - return 0 - - clusters: list[DistributedCluster] = [] - try: - discovered = [] - for target in targets: - target.work_dir.mkdir(parents=True, exist_ok=True) - cluster = DistributedCluster(target) - clusters.append(cluster) - cluster.start_all() - create_results = [] - metas = [] - for table in tables: - create_result = http_post_sql(target.http_port, create_table_sql(table), table["database"], args.http_timeout) - if not create_result["ok"]: - raise RuntimeError(f"CREATE TABLE {table['name']} failed for {target.name}: {create_result}") - create_results.append({"table": table["name"], "result": create_result}) - metas.append({"table": table["name"], **discover_region_via_frontend(target, table, args.http_timeout)}) - cluster.stop_component("datanode") - cluster._ensure_metasrv_alive() - discovered.append(metas) - report["targets"].append({"name": target.name, "binary": str(target.binary), "work_dir": str(target.work_dir), "data_dir": str(target.data_dir), "datanode_data_home": str(target.datanode_data_dir), "components": cluster.component_report(), "create_table": create_results[0]["result"] if len(create_results) == 1 else create_results, "discovered": metas[0] if len(metas) == 1 else metas, "discovered_tables": metas}) - - if len(discovered[0]) != len(discovered[1]) or any(a["table"] != b["table"] or a["region_id"] != b["region_id"] or a["table_dir"] != b["table_dir"] for a, b in zip(discovered[0], discovered[1])): - raise RuntimeError(f"base/candidate metadata mismatch: {discovered}") - generations = [] - for table_idx, meta in enumerate(discovered[0]): - per_table_fixture_dir = table_fixture_dir(fixture_dir, tables, tables[table_idx], table_idx) - generations.append(generate_fixture(args.fixture_generator, case_path, per_table_fixture_dir, dry_run=False, reuse_fixture=reuse_fixture, allow_large_fixture=args.allow_large_fixture, table_name=meta["table"] if len(tables) > 1 else None, database=meta["schema"] if len(tables) > 1 else None, region_id=meta["region_id"], table_dir=meta["table_dir"], region_dir=meta["region_dir"])) - report["fixture_generation"] = generations[0] if len(generations) == 1 else generations - - target_results = [] - for idx, target in enumerate(targets): - cluster = clusters[idx] - materialize = [] - for table_idx, meta in enumerate(discovered[idx]): - per_table_fixture_dir = table_fixture_dir(fixture_dir, tables, tables[table_idx], table_idx) - materialize.append(materialize_fixture(target, dry_run=False, preserve_state=True, expected_region_dir=meta["region_dir"], fixture_dir=per_table_fixture_dir)) - cluster.start_datanode() - query_result = run_queries(target, case, tables, args.http_timeout) - if query_result["status"] == "failed": - cluster.restart_frontend() - retry_result = run_queries(target, case, tables, args.http_timeout) - query_result["frontend_restart_retry"] = retry_result - if retry_result["status"] == "ok": - query_result = retry_result - report["targets"][idx]["fixture_dir"] = str(fixture_dir) - report["targets"][idx]["fixture_materialization"] = materialize[0] if len(materialize) == 1 else materialize - report["targets"][idx].update(query_result) - report["targets"][idx]["status"] = "measured" if query_result["status"] == "ok" else "failed" - write_json(target.report_path, report["targets"][idx]) - target_results.append(query_result) - - report["thresholds"] = enforce_thresholds(case, target_results[0], target_results[1]) - report["status"] = "failed" if any(t["status"] == "failed" for t in report["thresholds"]) or any(r["status"] == "failed" for r in target_results) else "ok" - except Exception as e: # noqa: BLE001 - write machine-readable failure report - report["status"] = "failed" - report["error"] = repr(e) - finally: - for cluster in reversed(clusters): - cluster.stop_all() - write_json(work_root / "query-regression-report.json", report) - output_report(report, args.output) - return 1 if report["status"] == "failed" else 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/perf/test_query_regression_runner_compaction_toctou.py b/tests/perf/test_query_regression_runner_compaction_toctou.py index 0e79cf245de..58f7f31b7fa 100644 --- a/tests/perf/test_query_regression_runner_compaction_toctou.py +++ b/tests/perf/test_query_regression_runner_compaction_toctou.py @@ -13,9 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Regression coverage for quiescing remote-write storage before read benches.""" +"""Regression coverage for the outer remote-write lifecycle.""" -import argparse import importlib.util import json import sys @@ -25,151 +24,90 @@ from pathlib import Path from unittest.mock import patch -RUNNER_PATH = Path(__file__).with_name("query_regression_runner.py") -SPEC = importlib.util.spec_from_file_location("query_regression_runner_under_test", RUNNER_PATH) +RUNNER_PATH = Path(__file__).resolve().parents[2] / ".github/scripts/query-regression-run.py" +SPEC = importlib.util.spec_from_file_location("query_regression_outer_under_test", RUNNER_PATH) assert SPEC is not None and SPEC.loader is not None runner = importlib.util.module_from_spec(SPEC) sys.modules[SPEC.name] = runner SPEC.loader.exec_module(runner) -class ReportOutputTest(unittest.TestCase): - def test_final_report_output_defaults_to_stdout_or_writes_to_file(self) -> None: - report = {"status": "ok"} - expected = json.dumps(report, indent=2, sort_keys=True) +class RemoteWriteLifecycleTest(unittest.TestCase): + def setUp(self) -> None: + self.args = type("Args", (), {"http_timeout": "1"})() + self.ports = list(range(10000, 10016)) - with patch("builtins.print") as print_mock: - runner.output_report(report, None) - print_mock.assert_called_once_with(expected) + def test_measure_report_is_finalized_after_datanodes_stop(self) -> None: + events: list[str] = [] - with tempfile.TemporaryDirectory() as tmpdir: - output = Path(tmpdir) / "nested" / "report.json" - with patch("builtins.print") as print_mock: - runner.output_report(report, output) - print_mock.assert_not_called() - self.assertEqual(output.read_text(), expected + "\n") + def run(command, **_kwargs): + phase = command[1] + events.append(phase) + if phase == "measure": + output = Path(command[command.index("--output") + 1]) + output.write_text(json.dumps({"targets": [], "thresholds": [], "status": "ok"})) + return type("Result", (), {"returncode": 0})() + def start(target, component, _procs): + events.append(f"start:{target.name}:{component}") -class RemoteWriteCompactionToctouTest(unittest.TestCase): - def test_quiesces_each_target_before_storage_inspection_and_read_bench(self) -> None: - lifecycle: dict[str, list[str]] = {"base": [], "candidate": []} - clusters = [] + def stop(target, component, _procs): + events.append(f"stop:{target.name}:{component}") - class RecordingCluster: - def __init__(self, target): - self.target = target - self.active = set() - self.stop_all_calls = 0 - clusters.append(self) + with ( + tempfile.TemporaryDirectory() as tmpdir, + patch.object(runner, "allocate_ports", return_value=self.ports), + patch.object(runner, "start_component", side_effect=start), + patch.object(runner, "stop_component", side_effect=stop), + patch.object(runner.subprocess, "run", side_effect=run), + ): + self.assertEqual( + runner.run_remote_case( + self.args, Path("case.toml"), Path(tmpdir), Path("base"), Path("candidate"), Path("fixture"), Path("runner") + ), + 0, + ) + self.assertTrue((Path(tmpdir) / "query-regression-report.json").exists()) - def start_all(self, _config_path): - self.active.update(("metasrv", "datanode", "frontend")) + measure = events.index("measure") + finalize = events.index("finalize-remote") + self.assertLess(measure, events.index("stop:base:datanode")) + self.assertLess(events.index("stop:base:datanode"), finalize) + self.assertLess(events.index("stop:candidate:datanode"), finalize) + self.assertIn("stop:base:frontend", events[finalize + 1:]) + self.assertIn("stop:candidate:metasrv", events[finalize + 1:]) - def stop_component(self, name): - if name in self.active: - lifecycle[self.target.name].append(f'cluster.stop_component("{name}")') - self.active.remove(name) + def test_phase_exception_still_cleans_up_without_processes_or_network(self) -> None: + events: list[str] = [] - def stop_all(self): - self.stop_all_calls += 1 - for name in ("frontend", "datanode", "metasrv"): - self.stop_component(name) + def run(command, **_kwargs): + phase = command[1] + events.append(phase) + if phase == "prepare-remote": + raise RuntimeError("prepare exploded") + return type("Result", (), {"returncode": 0})() - def component_report(self): - return {} + def start(target, component, _procs): + events.append(f"start:{target.name}:{component}") - def run_queries(target, _case, _tables, _http_timeout): - lifecycle[target.name].append("query measurement completes") - return {"validation": [], "validation_errors": [], "measurements": [], "status": "ok"} + def stop(target, component, _procs): + events.append(f"stop:{target.name}:{component}") - def inspect_storage(_helper, target, _storage, *, dry_run): - self.assertFalse(dry_run) - lifecycle[target.name].append("storage inspection") - return {"status": "ok", "summary": {"summary": {}}} + with ( + tempfile.TemporaryDirectory() as tmpdir, + patch.object(runner, "allocate_ports", return_value=self.ports), + patch.object(runner, "start_component", side_effect=start), + patch.object(runner, "stop_component", side_effect=stop), + patch.object(runner.subprocess, "run", side_effect=run), + ): + with self.assertRaisesRegex(RuntimeError, "prepare exploded"): + runner.run_remote_case( + self.args, Path("case.toml"), Path(tmpdir), Path("base"), Path("candidate"), Path("fixture"), Path("runner") + ) - def read_bench(_binary, target, _read_bench, _inspection, *, dry_run): - self.assertFalse(dry_run) - lifecycle[target.name].append("read bench") - return {"status": "ok"} - - remote_write = { - "database": "public", - "metric": "cpu_usage", - "physical_table": "cpu_usage_physical", - "series_count": 1, - "samples_per_series": 1, - "visibility_timeout_seconds": 1, - "storage": { - "inspect": True, - "min_files": None, - "min_files_with_column": None, - "max_total_file_size_bytes": None, - "max_column_compressed_size_bytes": None, - "max_column_uncompressed_size_bytes": None, - "require_encodings": [], - "forbid_encodings": [], - }, - "read_bench": {"enabled": True}, - "prom_store": { - "pending_rows_flush_interval": "1s", - "max_batch_rows": 1, - "max_concurrent_flushes": 1, - "worker_channel_capacity": 1, - "max_inflight_requests": 1, - }, - } - case = { - "scenario": { - "kind": "prom_remote_write_then_query", - "remote_write": remote_write, - "queries": [], - } - } - args = argparse.Namespace( - fixture_only=False, - fixture_generator=Path("/bin/true"), - remote_write_generator=None, - storage_inspector=None, - candidate_bin=Path("/bin/true"), - dry_run=False, - http_timeout=1.0, - ) - - with tempfile.TemporaryDirectory() as tmpdir: - root = Path(tmpdir) - targets = [ - runner.make_target("base", Path("/bin/true"), root, list(range(10000, 10008))), - runner.make_target("candidate", Path("/bin/true"), root, list(range(10008, 10016))), - ] - report = {"targets": [], "thresholds": []} - with ( - patch.object(runner, "DistributedCluster", RecordingCluster), - patch.object(runner, "run_remote_write_ingestion", return_value=({"status": "ok"}, [{"ok": True}])), - patch.object(runner, "http_post_sql", return_value={"ok": True}), - patch.object(runner, "poll_expected_count", return_value={"ok": True, "row_count_ok": True}), - patch.object(runner, "run_queries", side_effect=run_queries), - patch.object(runner, "run_storage_inspection", side_effect=inspect_storage), - patch.object(runner, "run_read_bench", side_effect=read_bench), - ): - runner.run_remote_write_scenario(args, case, root / "case.toml", targets, report) - - for cluster in clusters: - self.assertEqual(cluster.stop_all_calls, 1) - self.assertEqual(cluster.active, set()) - cluster.stop_all() - self.assertEqual(cluster.stop_all_calls, 2) - self.assertEqual(cluster.active, set()) - - expected = [ - "query measurement completes", - 'cluster.stop_component("datanode")', - "storage inspection", - "read bench", - ] - relevant = {name: [event for event in events if event in expected] for name, events in lifecycle.items()} - for target_name in ("base", "candidate"): - with self.subTest(target=target_name): - self.assertEqual(relevant[target_name], expected) + self.assertIn("prepare-remote", events) + self.assertIn("stop:base:metasrv", events) + self.assertIn("stop:candidate:frontend", events) if __name__ == "__main__": diff --git a/tests/perf/test_query_regression_runner_otlp_trace_load.py b/tests/perf/test_query_regression_runner_otlp_trace_load.py index ab14ceef24c..8b2ca0a0ba1 100644 --- a/tests/perf/test_query_regression_runner_otlp_trace_load.py +++ b/tests/perf/test_query_regression_runner_otlp_trace_load.py @@ -13,215 +13,80 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Regression coverage for the local OTLP trace load lifecycle and metrics.""" +"""Regression coverage for the outer OTLP lifecycle.""" import importlib.util import json -import os -import subprocess import sys import tempfile import unittest from pathlib import Path +from types import SimpleNamespace from unittest.mock import patch -RUNNER_PATH = Path(__file__).with_name("query_regression_runner.py") -PLOTTER_PATH = Path(__file__).with_name("plot_otlp_trace_report.sh") -SPEC = importlib.util.spec_from_file_location("query_regression_runner_otlp_under_test", RUNNER_PATH) +RUNNER_PATH = Path(__file__).resolve().parents[2] / ".github/scripts/query-regression-run.py" +SPEC = importlib.util.spec_from_file_location("query_regression_outer_otlp_under_test", RUNNER_PATH) assert SPEC is not None and SPEC.loader is not None runner = importlib.util.module_from_spec(SPEC) sys.modules[SPEC.name] = runner SPEC.loader.exec_module(runner) -class OtlpTraceReportPlotTest(unittest.TestCase): - def test_plots_all_comparison_metrics_and_thresholds(self) -> None: - report = { - "status": "ok", - "targets": [ - { - "name": name, - "metrics": { - "accepted_spans": spans, - "accepted_spans_per_second": rate, - "mean_http_latency_ms": latency, - "failure_count": 0, - }, - "visibility": {"observed_rows": spans}, - } - for name, spans, rate, latency in (("base", 100, 10.0, 3.0), ("candidate", 120, 12.0, 2.5)) - ], - "thresholds": [ - {"target": "base", "threshold": "max_failure_count", "status": "passed", "actual": 0, "limit": 0}, - { - "threshold": "max_candidate_throughput_regression_pct", - "status": "passed", - "actual_pct": -20.0, - "limit_pct": 20.0, - "base": 10.0, - "candidate": 12.0, - }, - ], - } +class OtlpTraceOuterLifecycleTest(unittest.TestCase): + def test_stops_base_before_candidate_and_wires_rust_commands(self) -> None: + events: list[str] = [] + commands: list[list[str]] = [] - with tempfile.TemporaryDirectory() as tmpdir: + def run(command, **_kwargs): + commands.append(command) + phase = command[1] + events.append(phase) + output = Path(command[command.index("--output") + 1]) + output.write_text(json.dumps({"status": "ok", "targets": []})) + return SimpleNamespace(returncode=0) + + def start(target, component, _procs): + events.append(f"start:{target.name}:{component}") + + def stop(target, component, _procs): + events.append(f"stop:{target.name}:{component}") + + args = SimpleNamespace(http_timeout="1", otelgen_bin=Path("otelgen")) + ports = list(range(10000, 10016)) + with ( + tempfile.TemporaryDirectory() as tmpdir, + patch.object(runner, "allocate_ports", return_value=ports), + patch.object(runner, "start_component", side_effect=start), + patch.object(runner, "stop_component", side_effect=stop), + patch.object(runner.subprocess, "run", side_effect=run), + ): root = Path(tmpdir) - report_path = root / "report.json" - report_path.write_text(json.dumps(report)) - fake_bin = root / "bin" - fake_bin.mkdir() - uplot = fake_bin / "uplot" - uplot.write_text('#!/usr/bin/env bash\nprintf "CALL %s\\n" "$*" >>"$UPLOT_LOG"\ncat >>"$UPLOT_LOG"\n') - uplot.chmod(0o755) - log_path = root / "uplot.log" - env = os.environ.copy() - env["PATH"] = f"{fake_bin}{os.pathsep}{env['PATH']}" - env["UPLOT_LOG"] = str(log_path) - - result = subprocess.run([str(PLOTTER_PATH), str(report_path)], env=env, text=True, capture_output=True, check=False) - log = log_path.read_text() - report["targets"][1]["visibility"]["observed_rows"] = 119 - report_path.write_text(json.dumps(report)) - mismatch = subprocess.run([str(PLOTTER_PATH), str(report_path)], env=env, text=True, capture_output=True, check=False) - - self.assertEqual(result.returncode, 0, result.stderr) - self.assertEqual(log.count("CALL "), 5) - for metric in ("Accepted spans", "Visible table rows", "Throughput", "Mean HTTP latency", "Failures"): - self.assertIn(metric, log) - self.assertIn("base\t100", log) - self.assertIn("candidate\t12", log) - self.assertIn("Status: ok", result.stdout) - self.assertIn("[passed] candidate: accepted=120, visible=120", result.stdout) - self.assertIn("max_failure_count", result.stdout) - self.assertIn("max_candidate_throughput_regression_pct", result.stdout) - self.assertEqual(mismatch.returncode, 1, mismatch.stderr) - self.assertIn("[failed] candidate: accepted=120, visible=119", mismatch.stdout) - - -class OtlpTraceLoadTest(unittest.TestCase): - def test_stops_base_cluster_before_creating_candidate_cluster(self) -> None: - events = [] - - class FakeCluster: - def __init__(self, target): - self.target = target - self.stopped = False - events.append(f"create:{target.name}") - - def component_report(self): - return {} - - def stop_all(self): - if not self.stopped: - self.stopped = True - events.append(f"stop:{self.target.name}") - - load = { - "database": "public", - "table": "opentelemetry_traces", - "pipeline": "greptime_trace_v1", - "duration_seconds": 120, - "warmup_seconds": 60, - "rate": 50_000, - "workers": 4, - "workload": "microservices", - "exporter_shards": 4, - "visibility_timeout_seconds": 1, - "thresholds": { - "max_candidate_throughput_regression_pct": 20, - "max_candidate_mean_latency_regression_pct": 20, - "max_failure_count": 0, - }, - } - args = runner.argparse.Namespace(fixture_only=False, otelgen_bin=None, dry_run=True, http_timeout=1.0) - - with tempfile.TemporaryDirectory() as tmpdir: - root = Path(tmpdir) - targets = [ - runner.make_target("base", Path("/bin/true"), root, list(range(10_000, 10_008))), - runner.make_target("candidate", Path("/bin/true"), root, list(range(10_008, 10_016))), + self.assertEqual( + runner.run_otlp_case(args, Path("case.toml"), root, Path("base-bin"), Path("candidate-bin"), Path("fixture"), Path("runner")), + 0, + ) + expected = [ + [ + "runner", "run-otlp-target", "--case", "case.toml", "--fixture-generator", "fixture", + "--otelgen-bin", "otelgen", "--http-port", "10004", "--target-name", "base", + "--work-dir", str(root / "base"), "--output", str(root / "base" / "report.json"), "--http-timeout", "1", + ], + [ + "runner", "run-otlp-target", "--case", "case.toml", "--fixture-generator", "fixture", + "--otelgen-bin", "otelgen", "--http-port", "10012", "--target-name", "candidate", + "--work-dir", str(root / "candidate"), "--output", str(root / "candidate" / "report.json"), "--http-timeout", "1", + ], + [ + "runner", "finalize-otlp", "--case", "case.toml", "--fixture-generator", "fixture", + "--base-result", str(root / "base" / "report.json"), + "--candidate-result", str(root / "candidate" / "report.json"), + "--output", str(root / "query-regression-report.json"), + ], ] - with patch.object(runner, "DistributedCluster", FakeCluster): - runner.run_otlp_trace_load_scenario( - args, - {"scenario": {"kind": "otlp_trace_load", "load": load}}, - targets, - {"targets": []}, - ) + self.assertEqual(commands, expected) - self.assertEqual(events, ["create:base", "stop:base", "create:candidate", "stop:candidate"]) - - def test_warmup_lifecycle_and_labeled_metric_deltas(self) -> None: - class FakeProcess: - def __init__(self, *_args, **_kwargs): - self.returncode = None - self.wait_timeouts = [] - - def wait(self, timeout): - self.wait_timeouts.append(timeout) - if len(self.wait_timeouts) == 1: - raise subprocess.TimeoutExpired("otelgen", timeout) - self.returncode = 0 - return 0 - - def poll(self): - return self.returncode - - def kill(self): - self.returncode = -9 - - def snapshot(text: str, captured: float): - return {"captured_monotonic_seconds": captured, "values": runner.parse_prometheus_metrics(text)} - - snapshots = [ - snapshot("greptime_frontend_otlp_traces_rows 10\n", 0.0), - snapshot( - 'greptime_frontend_otlp_traces_rows 110\n' - 'greptime_servers_http_otlp_traces_elapsed_sum{db="public"} 1\n' - 'greptime_servers_http_otlp_traces_elapsed_count{db="public"} 10\n', - 5.0, - ), - snapshot( - 'greptime_frontend_otlp_traces_rows 310\n' - 'greptime_frontend_otlp_traces_failure_count{label="decode"} 1\n' - 'greptime_frontend_otlp_traces_failure_count{label="write"} 2\n' - 'greptime_servers_http_otlp_traces_elapsed_sum{db="public"} 3\n' - 'greptime_servers_http_otlp_traces_elapsed_count{db="public"} 30\n', - 15.0, - ), - ] - load = { - "database": "public", - "table": "opentelemetry_traces", - "pipeline": "greptime_trace_v1", - "duration_seconds": 120, - "warmup_seconds": 60, - "rate": 50_000, - "workers": 4, - "workload": "microservices", - "exporter_shards": 4, - } - - with tempfile.TemporaryDirectory() as tmpdir: - target = runner.make_target("base", Path("/bin/true"), Path(tmpdir), list(range(10_000, 10_008))) - process = FakeProcess() - with ( - patch.object(runner, "fetch_otlp_metrics", side_effect=snapshots), - patch.object(runner.subprocess, "Popen", return_value=process), - patch.object(runner.time, "monotonic", side_effect=[0.0, 120.0]), - ): - result = runner.run_otelgen_load(Path("/bin/otelgen"), target, load, 1.0, dry_run=False) - - self.assertEqual(result["status"], "ok") - self.assertEqual(process.wait_timeouts, [60, 120]) - metrics = runner.summarize_otlp_metrics(result) - self.assertEqual(metrics["accepted_spans"], 300) - self.assertEqual(metrics["measurement_accepted_spans"], 200) - self.assertEqual(metrics["accepted_spans_per_second"], 20.0) - self.assertEqual(metrics["http_requests"], 20) - self.assertEqual(metrics["mean_http_latency_ms"], 100.0) - self.assertEqual(metrics["failure_count"], 3) + self.assertLess(events.index("stop:base:metasrv"), events.index("start:candidate:metasrv")) if __name__ == "__main__": From ca1c0b6c9b533afbdd16f5d07765f13fd80e1d5e Mon Sep 17 00:00:00 2001 From: Yingwen Date: Thu, 6 Aug 2026 21:22:25 +0800 Subject: [PATCH 009/126] feat(mito2): add range-based metric series reader (#8703) * feat(mito2): add range-based metric series reader Signed-off-by: evenyag * fix(mito2): address series reader review feedback Signed-off-by: evenyag * fix(mito2): align series predicate filtering Signed-off-by: evenyag * refactor: update semaphore usage and move prefilter flag Signed-off-by: evenyag * fix(mito2): enforce candidate pruner invariant Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/src/memtable/bulk/part_reader.rs | 1 + src/mito2/src/read/prune.rs | 6 +- src/mito2/src/read/pruner.rs | 68 +- src/mito2/src/read/range_cache.rs | 1 - src/mito2/src/read/scan_region.rs | 4 +- src/mito2/src/read/series_candidate.rs | 79 ++- src/mito2/src/read/series_reader.rs | 697 ++++++++++++++++++++- src/mito2/src/sst/parquet/file_range.rs | 122 +++- src/mito2/src/sst/parquet/prefilter.rs | 56 +- src/mito2/src/sst/parquet/reader.rs | 27 +- 10 files changed, 1006 insertions(+), 55 deletions(-) diff --git a/src/mito2/src/memtable/bulk/part_reader.rs b/src/mito2/src/memtable/bulk/part_reader.rs index 062d8fc977e..b6da514e625 100644 --- a/src/mito2/src/memtable/bulk/part_reader.rs +++ b/src/mito2/src/memtable/bulk/part_reader.rs @@ -419,6 +419,7 @@ fn apply_combined_filters( let predicate_mask = context.base.compute_filter_mask_flat( &record_batch, skip_fields, + false, &mut tag_decode_state, )?; // If predicate filters out the entire batch, return None early diff --git a/src/mito2/src/read/prune.rs b/src/mito2/src/read/prune.rs index 52c36e6e08c..70d5f505c72 100644 --- a/src/mito2/src/read/prune.rs +++ b/src/mito2/src/read/prune.rs @@ -223,9 +223,9 @@ impl FlatPruneReader { } let num_rows_before_filter = record_batch.num_rows(); - let Some(filtered_batch) = self - .context - .precise_filter_flat(record_batch, self.skip_fields)? + let Some(filtered_batch) = + self.context + .precise_filter_flat(record_batch, self.skip_fields, false)? else { // the entire batch is filtered out self.metrics.filter_metrics.rows_precise_filtered += num_rows_before_filter; diff --git a/src/mito2/src/read/pruner.rs b/src/mito2/src/read/pruner.rs index 9b462a4566b..b1fff5e81c2 100644 --- a/src/mito2/src/read/pruner.rs +++ b/src/mito2/src/read/pruner.rs @@ -178,6 +178,30 @@ impl PartitionPruner { } } +/// Options to create a [`Pruner`]. +#[derive(Debug, Clone, Copy)] +pub struct PrunerOptions { + /// Keeps file range builders until the query-scoped pruner is dropped. + pub retain_builders: bool, + /// Whether [`FileRange`]s execute the reduced-column predicate prefilter. + /// + /// The flag is baked into the cached [`FileRangeBuilder`], so it belongs to + /// the whole scan. All partitions sharing this pruner must agree on it. + /// Two-stage series scans disable this prefilter: candidate discovery applies + /// tag predicates, then the data phase calls [`FileRange::precise_filter_flat`] + /// with tag filtering skipped so the predicates are applied exactly once. + pub enable_predicate_prefilter: bool, +} + +impl Default for PrunerOptions { + fn default() -> Self { + Self { + retain_builders: false, + enable_predicate_prefilter: true, + } + } +} + /// A pruner that prunes files for all partitions of a scanner. pub struct Pruner { /// Channels to send requests to workers. @@ -200,6 +224,8 @@ struct PrunerInner { manifest_pruned_files: Vec, /// Keeps file range builders until the query-scoped pruner is dropped. retain_builders: bool, + /// Whether FileRanges should execute the reduced-column predicate prefilter. + enable_predicate_prefilter: bool, } impl Drop for PrunerInner { @@ -275,15 +301,19 @@ impl Pruner { /// Initially all file_entries have `remaining_ranges = 0`. /// Call `add_partition_ranges()` to initialize ref counts. pub fn new(stream_ctx: Arc, num_workers: usize) -> Self { - Self::new_with_retained_builders(stream_ctx, num_workers, false) + Self::new_with_options(stream_ctx, num_workers, PrunerOptions::default()) } - /// Creates a new pruner and optionally retains file range builders until drop. - pub fn new_with_retained_builders( + /// Creates a new pruner with the given options. + pub fn new_with_options( stream_ctx: Arc, num_workers: usize, - retain_builders: bool, + options: PrunerOptions, ) -> Self { + let PrunerOptions { + retain_builders, + enable_predicate_prefilter, + } = options; let num_files = stream_ctx.input.num_files(); let file_entries: Vec<_> = (0..num_files) .map(|_| { @@ -311,6 +341,7 @@ impl Pruner { stream_ctx, manifest_pruned_files, retain_builders, + enable_predicate_prefilter, }); // Spawn worker tasks with their receivers @@ -486,6 +517,12 @@ impl Pruner { let _ = self.worker_senders[worker_idx].try_send(request); } + /// Returns whether this pruner builds file ranges with the reduced-column + /// predicate prefilter enabled. + pub fn predicate_prefilter_enabled(&self) -> bool { + self.inner.enable_predicate_prefilter + } + fn get_worker_idx(&self, file_id: FileId) -> usize { let file_id_hash = Uuid::from(file_id).as_u128() as usize; file_id_hash % self.inner.num_workers @@ -516,7 +553,13 @@ impl Pruner { .inner .stream_ctx .input - .prune_file_after_manifest_check(file, pre_filter_mode, predicate, reader_metrics) + .prune_file_after_manifest_check( + file, + pre_filter_mode, + self.inner.enable_predicate_prefilter, + predicate, + reader_metrics, + ) .await?; let arc_builder = Arc::new(builder); @@ -599,7 +642,13 @@ impl Pruner { inner .stream_ctx .input - .prune_file_after_manifest_check(file, pre_filter_mode, predicate, &mut metrics) + .prune_file_after_manifest_check( + file, + pre_filter_mode, + inner.enable_predicate_prefilter, + predicate, + &mut metrics, + ) .await }; @@ -796,10 +845,13 @@ mod tests { .with_files(files) .with_append_mode(true); let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input)); - let pruner = Arc::new(Pruner::new_with_retained_builders( + let pruner = Arc::new(Pruner::new_with_options( stream_ctx, 1, - retain_builders, + PrunerOptions { + retain_builders, + ..Default::default() + }, )); (env, pruner) } diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index 327af3d325b..036a574c1f5 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -501,7 +501,6 @@ pub(crate) fn build_candidate_range_cache_key( } /// Builds a cache key for a two-phase series-data partition-range result. -#[allow(dead_code)] pub(crate) fn build_series_range_cache_key( stream_ctx: &StreamContext, part_range: &PartitionRange, diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 482c2d4a332..250f0d7fc27 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -1269,7 +1269,7 @@ impl ScanInput { return Ok(FileRangeBuilder::default()); } - self.prune_file_after_manifest_check(file, pre_filter_mode, predicate, reader_metrics) + self.prune_file_after_manifest_check(file, pre_filter_mode, true, predicate, reader_metrics) .await } @@ -1284,6 +1284,7 @@ impl ScanInput { &self, file: &FileHandle, pre_filter_mode: PreFilterMode, + enable_predicate_prefilter: bool, predicate: Option, reader_metrics: &mut ReaderMetrics, ) -> Result { @@ -1320,6 +1321,7 @@ impl ScanInput { .expected_metadata(Some(self.mapper.metadata().clone())) .compaction(self.compaction) .pre_filter_mode(pre_filter_mode) + .enable_predicate_prefilter(enable_predicate_prefilter) .decode_primary_key_values(decode_pk_values) .build_reader_input(reader_metrics) .await; diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index d2154493eee..03d76d2adf6 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -88,6 +88,10 @@ impl SeriesCandidateScanner { /// /// Callers must fall back to the legacy series-scan path when the scan contains /// extension ranges. Candidate-series discovery does not support other range types. + /// + /// `pruner` must be built with `PrunerOptions::enable_predicate_prefilter` set to + /// `false`: both scan phases share its file range builders, and the two-stage path + /// applies simple filters on the precise-filter path instead. pub(crate) fn try_new( stream_ctx: Arc, partitions: Vec>, @@ -106,6 +110,15 @@ impl SeriesCandidateScanner { reason: "candidate-series scan does not support extension ranges; use the legacy series-scan path", } ); + ensure!( + !pruner.predicate_prefilter_enabled(), + UnexpectedSnafu { + reason: format!( + "candidate-series scan for region {} requires a pruner without predicate prefiltering", + stream_ctx.input.region_metadata().region_id + ), + } + ); let all_ranges = partitions.iter().flatten().copied().collect::>(); pruner.add_partition_ranges(&all_ranges); let partition_pruner = Arc::new(PartitionPruner::new(pruner, &all_ranges)); @@ -253,12 +266,15 @@ impl SeriesCandidateRangeBuilder { ); let filter = build_primary_key_filter( &metadata, + None, self.stream_ctx.input.predicate_group().predicate(), ); return Ok(Some(candidate_primary_key_stream(Box::pin(raw), filter))); } if self.stream_ctx.is_file_range_index(index) { + let file = self.stream_ctx.input.file_from_index(index); + let predicate = self.stream_ctx.input.predicate_for_file(file); if self .partition_pruner .try_skip_manifest_pruned_file_range(index, &self.part_metrics) @@ -277,9 +293,15 @@ impl SeriesCandidateRangeBuilder { self.part_metrics .merge_reader_metrics(&reader_metrics, None); - // Reuse the exact encoded-PK filter selected by the file's reader plan, but - // execute it in `candidate_primary_key_stream` after the PK-only read. - let filter = ranges.first().and_then(|range| range.primary_key_filter()); + // Build a fresh encoded-PK filter from this SST's metadata so schema + // compatibility is evaluated independently for each file. + let filter = ranges.first().and_then(|range| { + build_primary_key_filter( + range.region_metadata(), + Some(metadata.as_ref()), + predicate.as_ref(), + ) + }); let part_metrics = self.part_metrics.clone(); let raw = Box::pin(try_stream! { let fetch_metrics = part_metrics @@ -547,6 +569,8 @@ fn decode_metric_series( #[cfg(test)] mod tests { + use std::time::Instant; + use datafusion::execution::memory_pool::UnboundedMemoryPool; use datafusion_expr::{col, lit}; use datatypes::arrow::array::{ArrayRef, DictionaryArray, UInt32Array}; @@ -556,8 +580,55 @@ mod tests { use table::predicate::Predicate; use super::*; + use crate::read::flat_projection::FlatProjectionMapper; + use crate::read::scan_region::ScanInput; + use crate::read::scan_util::PartitionMetrics; + use crate::test_util::scheduler_util::SchedulerEnv; use crate::test_util::sst_util::sst_region_metadata_with_encoding; + #[tokio::test] + async fn candidate_scanner_rejects_predicate_prefilter_pruner() { + let env = SchedulerEnv::new().await; + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let mapper = + FlatProjectionMapper::new(&metadata, 0..metadata.column_metadatas.len()).unwrap(); + let stream_ctx = Arc::new(StreamContext::seq_scan_ctx(ScanInput::new( + env.access_layer.clone(), + mapper, + ))); + let pruner = Arc::new(Pruner::new(stream_ctx.clone(), 1)); + let metrics_set = ExecutionPlanMetricsSet::new(); + let part_metrics = PartitionMetrics::new( + metadata.region_id, + 0, + "candidate-test", + Instant::now(), + false, + &metrics_set, + ); + + let error = SeriesCandidateScanner::try_new( + stream_ctx, + Vec::new(), + pruner, + Arc::new(Semaphore::new(1)), + Arc::new(UnboundedMemoryPool::default()), + metrics_set, + part_metrics, + ) + .err() + .unwrap(); + + assert!(matches!(error, crate::error::Error::Unexpected { .. })); + assert!( + error + .to_string() + .contains("requires a pruner without predicate prefiltering") + ); + } + fn binary_batch(values: &[&[u8]]) -> RecordBatch { RecordBatch::try_new( primary_key_schema(), @@ -621,7 +692,7 @@ mod tests { let predicate = Predicate::new(vec![ col(store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME).eq(lit(1_u32)), ]); - let filter = build_primary_key_filter(&metadata, Some(&predicate)); + let filter = build_primary_key_filter(&metadata, None, Some(&predicate)); let input = Box::pin(futures::stream::iter(vec![Ok(dictionary_batch( &[table_1.as_slice(), table_2.as_slice()], &[0, 1], diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index 4e1a86e1bd1..d16d13f2dac 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -12,9 +12,41 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Shared types for range-based metric series reads. +//! Reads selected metric series from partition ranges. + +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Instant; + +use async_stream::try_stream; +use futures::TryStreamExt; +use mito_codec::row_converter::{PrimaryKeyFilter, SparsePrimaryKeyCodec}; +use snafu::ResultExt; +use store_api::region_engine::PartitionRange; +use tokio::sync::Semaphore; + +#[cfg(feature = "enterprise")] +use crate::error::InvalidRequestSnafu; +use crate::error::{JoinSnafu, Result, UnexpectedSnafu}; +use crate::read::BoxedRecordBatchStream; +use crate::read::pruner::PartitionPruner; +use crate::read::range_cache::{ + build_series_range_cache_key, cache_flat_range_stream, cached_flat_range_stream, +}; +use crate::read::scan_region::StreamContext; +use crate::read::scan_util::{ + PartitionMetrics, SplitRecordBatchStream, compute_average_batch_size, + compute_parallel_channel_size, new_filter_metrics, scan_flat_mem_ranges, + should_split_flat_batches_for_merge, +}; +use crate::read::seq_scan::SeqScan; +use crate::read::series_candidate::{MetricSeriesId, validate_metric_metadata}; +use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; +use crate::sst::parquet::flat_format::primary_key_column_index; +use crate::sst::parquet::prefilter::prefilter_flat_batch_by_primary_key; +use crate::sst::parquet::reader::ReaderMetrics; +use crate::sst::parquet::row_group::ParquetFetchMetrics; -#[allow(dead_code)] const TSID_DOMAIN_END: u128 = 1u128 << u64::BITS; /// A stable partition of the TSID integer domain. @@ -24,7 +56,6 @@ pub(crate) struct SeriesRange { end: u128, } -#[allow(dead_code)] impl SeriesRange { pub(crate) fn new(partition: usize, partitions: usize) -> Option { if partitions == 0 || partition >= partitions { @@ -35,28 +66,684 @@ impl SeriesRange { let partition = partition as u128; let boundary = |partition: u128| { let numerator = partition * TSID_DOMAIN_END; - numerator / partitions + u128::from(!numerator.is_multiple_of(partitions)) + numerator.div_ceil(partitions) }; Some(Self { start: boundary(partition), end: boundary(partition + 1), }) } + + fn partition_for(tsid: u64, partitions: usize) -> usize { + ((tsid as u128 * partitions as u128) >> u64::BITS) as usize + } +} + +/// All series assigned to one data-reader partition. +#[derive(Debug)] +pub(crate) struct AssignedSeriesBatch { + range: SeriesRange, + series: Vec, +} + +#[allow(dead_code)] +impl AssignedSeriesBatch { + fn new(range: SeriesRange, series: Vec) -> Self { + Self { range, series } + } + + pub(crate) fn range(&self) -> SeriesRange { + self.range + } + + pub(crate) fn series(&self) -> &[MetricSeriesId] { + &self.series + } +} + +/// Collects candidate batches and assigns every TSID by its integer range. +#[allow(dead_code)] +pub(crate) struct SeriesBatchCollector { + assignments: Vec>, +} + +#[allow(dead_code)] +impl SeriesBatchCollector { + pub(crate) fn new(partitions: usize) -> Option { + (partitions > 0).then(|| Self { + assignments: (0..partitions).map(|_| Vec::new()).collect(), + }) + } + + pub(crate) fn push(&mut self, batch: Vec) { + let partitions = self.assignments.len(); + for series in batch { + let partition = SeriesRange::partition_for(series.tsid, partitions); + self.assignments[partition].push(series); + } + } + + pub(crate) fn finish(self) -> Vec { + let partitions = self.assignments.len(); + self.assignments + .into_iter() + .enumerate() + .map(|(partition, series)| { + AssignedSeriesBatch::new(SeriesRange::new(partition, partitions).unwrap(), series) + }) + .collect() + } +} + +/// Immutable allow-list for one partition's metric series. +#[derive(Clone, Debug)] +struct MetricSeriesFilter { + range: SeriesRange, + series: Arc>, +} + +impl MetricSeriesFilter { + fn new(assigned: &AssignedSeriesBatch) -> Self { + let series = assigned.series.iter().copied().collect(); + Self { + range: assigned.range, + series: Arc::new(series), + } + } + + fn primary_key_filter(&self, codec: SparsePrimaryKeyCodec) -> Box { + Box::new(MetricSeriesPrimaryKeyFilter { + codec, + series: self.series.clone(), + last_primary_key: Vec::new(), + last_match: None, + }) + } + + fn overlaps_encoded_bounds( + &self, + codec: &SparsePrimaryKeyCodec, + encoded_min: &[u8], + encoded_max: &[u8], + ) -> Option { + let (min_table_id, min_tsid) = codec.decode_ids(encoded_min).ok()?; + let (max_table_id, max_tsid) = codec.decode_ids(encoded_max).ok()?; + let min = MetricSeriesId { + table_id: min_table_id, + tsid: min_tsid, + }; + let max = MetricSeriesId { + table_id: max_table_id, + tsid: max_tsid, + }; + if min > max { + return None; + } + if min_table_id != max_table_id { + return Some(true); + } + + Some(u128::from(min_tsid) < self.range.end && u128::from(max_tsid) >= self.range.start) + } +} + +struct MetricSeriesPrimaryKeyFilter { + codec: SparsePrimaryKeyCodec, + series: Arc>, + last_primary_key: Vec, + last_match: Option, +} + +impl PrimaryKeyFilter for MetricSeriesPrimaryKeyFilter { + fn matches(&mut self, primary_key: &[u8]) -> mito_codec::error::Result { + if let Some(last_match) = self.last_match + && self.last_primary_key == primary_key + { + return Ok(last_match); + } + + let (table_id, tsid) = self.codec.decode_ids(primary_key)?; + let matched = self.series.contains(&MetricSeriesId { table_id, tsid }); + self.last_primary_key.clear(); + self.last_primary_key.extend_from_slice(primary_key); + self.last_match = Some(matched); + Ok(matched) + } +} + +fn filter_flat_stream_by_series( + mut input: BoxedRecordBatchStream, + codec: SparsePrimaryKeyCodec, + filter: MetricSeriesFilter, +) -> BoxedRecordBatchStream { + Box::pin(try_stream! { + let mut primary_key_filter = filter.primary_key_filter(codec); + while let Some(batch) = input.try_next().await? { + let pk_idx = primary_key_column_index(batch.num_columns()); + if let Some(batch) = prefilter_flat_batch_by_primary_key( + batch, + pk_idx, + primary_key_filter.as_mut(), + )? { + yield batch; + } + } + }) +} + +/// Reads all collected metric series assigned to one partition. +#[allow(dead_code)] +pub(crate) struct SeriesReader { + stream_ctx: Arc, + partition_ranges: Vec, + range: SeriesRange, + filter: MetricSeriesFilter, + codec: SparsePrimaryKeyCodec, + partition_pruner: Arc, + range_semaphore: Arc, + part_metrics: PartitionMetrics, +} + +#[allow(dead_code)] +impl SeriesReader { + /// Creates a reader for the series assigned to one data partition. + /// + /// `partition_pruner` must come from the candidate scanner's pruner, which + /// has predicate prefiltering disabled. The file read path applies precise + /// filters with tag filtering skipped because candidate discovery has already + /// enforced tag predicates through the exact assigned-series set. + /// + /// A single `range_semaphore` covers both the range-build phase and the final + /// merge. + pub(crate) fn try_new( + stream_ctx: Arc, + partition_ranges: Vec, + assigned_series: AssignedSeriesBatch, + partition_pruner: Arc, + range_semaphore: Arc, + part_metrics: PartitionMetrics, + ) -> Result { + validate_metric_metadata(&stream_ctx)?; + #[cfg(feature = "enterprise")] + snafu::ensure!( + stream_ctx.input.extension_ranges().is_empty(), + InvalidRequestSnafu { + region_id: stream_ctx.input.region_metadata().region_id, + reason: "series reader does not support extension ranges", + } + ); + + let range = assigned_series.range(); + let filter = MetricSeriesFilter::new(&assigned_series); + let codec = SparsePrimaryKeyCodec::new(stream_ctx.input.region_metadata()); + Ok(Self { + stream_ctx, + partition_ranges, + range, + filter, + codec, + partition_pruner, + range_semaphore, + part_metrics, + }) + } + + pub(crate) async fn build_stream(&self) -> Result { + if self.partition_ranges.is_empty() || self.filter.series.is_empty() { + return Ok(Box::pin(futures::stream::empty())); + } + + let mut tasks = Vec::with_capacity(self.partition_ranges.len()); + for part_range in self.partition_ranges.iter().copied() { + let stream_ctx = self.stream_ctx.clone(); + let filter = self.filter.clone(); + let codec = self.codec.clone(); + let partition_pruner = self.partition_pruner.clone(); + let range_semaphore = self.range_semaphore.clone(); + let part_metrics = self.part_metrics.clone(); + let range = self.range; + tasks.push(common_runtime::spawn_query(async move { + let _permit = range_semaphore.acquire().await.map_err(|error| { + UnexpectedSnafu { + reason: format!("failed to acquire series range permit: {error}"), + } + .build() + })?; + build_series_partition_range( + stream_ctx, + part_range, + range, + filter, + codec, + partition_pruner, + part_metrics, + ) + .await + })); + } + + let mut range_streams = Vec::with_capacity(tasks.len()); + let mut estimated_batch_sizes = Vec::with_capacity(tasks.len()); + for task in tasks { + let (stream, estimated_batch_size) = task.await.context(JoinSnafu)??; + range_streams.push(stream); + estimated_batch_sizes.push(estimated_batch_size); + } + + // Every range task above has finished, so all build permits are released + // and the final merge can reuse the same semaphore. + let estimated_batch_size = compute_average_batch_size(estimated_batch_sizes); + SeqScan::build_flat_reader_from_sources( + &self.stream_ctx, + range_streams, + Some(self.range_semaphore.clone()), + Some(&self.part_metrics), + true, + compute_parallel_channel_size(estimated_batch_size), + ) + .await + } +} + +async fn build_series_partition_range( + stream_ctx: Arc, + part_range: PartitionRange, + range: SeriesRange, + filter: MetricSeriesFilter, + codec: SparsePrimaryKeyCodec, + partition_pruner: Arc, + part_metrics: PartitionMetrics, +) -> Result<(BoxedRecordBatchStream, usize)> { + let cache_key = build_series_range_cache_key(&stream_ctx, &part_range, range); + if let Some(key) = cache_key.as_ref() { + if let Some(value) = stream_ctx.input.cache_strategy.get_range_result(key) { + part_metrics.inc_range_cache_hit(); + return Ok((cached_flat_range_stream(value), DEFAULT_READ_BATCH_SIZE)); + } + part_metrics.inc_range_cache_miss(); + } + + let range_meta = &stream_ctx.ranges[part_range.identifier]; + let split_batch_size = should_split_flat_batches_for_merge(&stream_ctx, range_meta); + let mut sources = Vec::with_capacity(range_meta.row_group_indices.len()); + + for index in range_meta.row_group_indices.iter().copied() { + if stream_ctx.is_mem_range_index(index) { + let stream = Box::pin(scan_flat_mem_ranges( + stream_ctx.clone(), + part_metrics.clone(), + index, + range_meta.time_range, + )); + sources.push(filter_flat_stream_by_series( + stream, + codec.clone(), + filter.clone(), + )); + continue; + } + + if stream_ctx.is_file_range_index(index) { + let file = stream_ctx.input.file_from_index(index); + if matches!( + file.primary_key_range() + .and_then(|(min, max)| { filter.overlaps_encoded_bounds(&codec, &min, &max) }), + Some(false) + ) { + continue; + } + + if partition_pruner.try_skip_manifest_pruned_file_range(index, &part_metrics) { + continue; + } + let mut reader_metrics = ReaderMetrics { + filter_metrics: new_filter_metrics(part_metrics.explain_verbose()), + ..Default::default() + }; + let file_ranges = partition_pruner + .build_file_ranges(index, &part_metrics, &mut reader_metrics) + .await?; + part_metrics.inc_num_file_ranges(file_ranges.len()); + part_metrics.merge_reader_metrics(&reader_metrics, None); + let ranges = file_ranges + .iter() + .filter(|file_range| { + !matches!( + file_range.primary_key_range().and_then(|(min, max)| { + filter.overlaps_encoded_bounds(&codec, min, max) + }), + Some(false) + ) + }) + .cloned() + .collect::>(); + if ranges.is_empty() { + continue; + } + + let stream = scan_series_file_ranges( + part_metrics.clone(), + ranges, + filter.clone(), + codec.clone(), + ); + sources.push(Box::pin(stream) as BoxedRecordBatchStream); + continue; + } + + return UnexpectedSnafu { + reason: format!( + "series reader received unsupported range index {}", + index.index + ), + } + .fail(); + } + + if split_batch_size.is_some() { + sources = sources + .into_iter() + .map(|stream| Box::pin(SplitRecordBatchStream::new(stream)) as BoxedRecordBatchStream) + .collect(); + } + let estimated_batch_size = split_batch_size.unwrap_or(DEFAULT_READ_BATCH_SIZE); + let stream = SeqScan::build_flat_reader_from_sources( + &stream_ctx, + sources, + None, + Some(&part_metrics), + false, + compute_parallel_channel_size(estimated_batch_size), + ) + .await?; + let stream = match cache_key { + Some(key) => cache_flat_range_stream( + stream, + stream_ctx.input.cache_strategy.clone(), + key, + part_metrics, + ), + None => stream, + }; + Ok((stream, estimated_batch_size)) +} + +fn scan_series_file_ranges( + part_metrics: PartitionMetrics, + ranges: smallvec::SmallVec<[crate::sst::parquet::file_range::FileRange; 2]>, + filter: MetricSeriesFilter, + codec: SparsePrimaryKeyCodec, +) -> impl futures::Stream> { + try_stream! { + let fetch_metrics = part_metrics + .explain_verbose() + .then(|| Arc::new(ParquetFetchMetrics::default())); + let mut reader_metrics = ReaderMetrics { + fetch_metrics: fetch_metrics.clone(), + ..Default::default() + }; + let mut primary_key_filter = filter.primary_key_filter(codec); + + for range in ranges { + let build_start = Instant::now(); + let Some(mut reader) = range + .reader_by_primary_key( + primary_key_filter.as_mut(), + fetch_metrics.as_deref(), + ) + .await? + else { + continue; + }; + let build_cost = build_start.elapsed(); + reader_metrics.build_cost += build_cost; + part_metrics.inc_build_reader_cost(build_cost); + + let scan_start = Instant::now(); + while let Some(record_batch) = reader.next_batch().await? { + reader_metrics.num_record_batches += 1; + reader_metrics.num_batches += 1; + reader_metrics.num_rows += record_batch.num_rows(); + + let num_rows_before_filter = record_batch.num_rows(); + let Some(record_batch) = range.precise_filter_flat( + record_batch, + range.pre_filter_mode().skip_fields(), + true, + )? else { + reader_metrics.filter_metrics.rows_precise_filtered += + num_rows_before_filter; + continue; + }; + reader_metrics.filter_metrics.rows_precise_filtered += + num_rows_before_filter - record_batch.num_rows(); + + let record_batch = if let Some(mapper) = range.compaction_projection_mapper() { + mapper.project(record_batch)? + } else { + record_batch + }; + if let Some(compat) = range.compat_batch() { + yield compat.compat(record_batch)?; + } else { + yield record_batch; + } + } + reader_metrics.scan_cost += scan_start.elapsed(); + } + + reader_metrics.observe_rows("series_data"); + reader_metrics.filter_metrics.observe(); + part_metrics.merge_reader_metrics(&reader_metrics, None); + } } #[cfg(test)] mod tests { + use store_api::codec::PrimaryKeyEncoding; + use super::*; + use crate::error::DecodeSnafu; + use crate::test_util::sst_util::sst_region_metadata_with_encoding; + + fn series(table_id: u32, tsid: u64) -> MetricSeriesId { + MetricSeriesId { table_id, tsid } + } + + fn assigned_batch( + partitions: usize, + partition: usize, + series: Vec, + ) -> AssignedSeriesBatch { + let mut collector = SeriesBatchCollector::new(partitions).unwrap(); + collector.push(series); + collector.finish().remove(partition) + } + + #[test] + fn series_range_assignment_is_stable_across_batch_boundaries() { + let input = vec![ + series(1, 0), + series(2, 1u64 << 62), + series(1, 1u64 << 63), + series(2, 3u64 << 62), + series(1, u64::MAX), + series(3, 0), + ]; + let mut first = SeriesBatchCollector::new(4).unwrap(); + first.push(input.clone()); + let first = first.finish(); + + let mut second = SeriesBatchCollector::new(4).unwrap(); + for chunk in input.chunks(2) { + second.push(chunk.to_vec()); + } + let second = second.finish(); + + assert_eq!( + first + .iter() + .map(AssignedSeriesBatch::series) + .collect::>(), + second + .iter() + .map(AssignedSeriesBatch::series) + .collect::>() + ); + assert_eq!(&[series(1, 0), series(3, 0)], first[0].series()); + assert_eq!(&[series(2, 1u64 << 62)], first[1].series()); + assert_eq!( + input.len(), + first + .iter() + .map(|batch| batch.series().len()) + .sum::() + ); + } #[test] fn series_ranges_cover_the_tsid_domain() { let ranges = (0..3) .map(|partition| SeriesRange::new(partition, 3).unwrap()) .collect::>(); - assert_eq!(0, ranges[0].start); assert_eq!(TSID_DOMAIN_END, ranges[2].end); assert_eq!(ranges[0].end, ranges[1].start); assert_eq!(ranges[1].end, ranges[2].start); + assert_eq!(0, SeriesRange::partition_for(0, 3)); + assert_eq!(2, SeriesRange::partition_for(u64::MAX, 3)); + + for (partition, range) in ranges.iter().enumerate() { + assert_eq!(partition, SeriesRange::partition_for(range.start as u64, 3)); + if range.end < TSID_DOMAIN_END { + assert_eq!( + partition + 1, + SeriesRange::partition_for(range.end as u64, 3) + ); + } + } + } + + #[test] + fn metric_series_filter_matches_encoded_primary_key() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let assigned = assigned_batch(1, 0, vec![series(1, 10), series(2, 20)]); + let filter = MetricSeriesFilter::new(&assigned); + let codec = SparsePrimaryKeyCodec::new(&metadata); + let mut primary_key_filter = filter.primary_key_filter(codec); + + for selected in [series(1, 10), series(2, 20)] { + assert!( + primary_key_filter + .matches(&encode_series(&metadata, selected.table_id, selected.tsid)) + .context(DecodeSnafu) + .unwrap() + ); + } + for unselected in [series(2, 10), series(1, 20)] { + assert!( + !primary_key_filter + .matches(&encode_series( + &metadata, + unselected.table_id, + unselected.tsid, + )) + .context(DecodeSnafu) + .unwrap() + ); + } + } + + fn encode_series( + metadata: &store_api::metadata::RegionMetadataRef, + table_id: u32, + tsid: u64, + ) -> Vec { + let codec = SparsePrimaryKeyCodec::new(metadata); + let mut primary_key = Vec::new(); + codec + .encode_internal(table_id, tsid, &mut primary_key) + .unwrap(); + primary_key + } + + #[test] + fn series_range_overlaps_single_table_primary_key_bounds() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let assigned = assigned_batch(2, 0, vec![series(1, 10), series(2, 20)]); + let filter = MetricSeriesFilter::new(&assigned); + let codec = SparsePrimaryKeyCodec::new(&metadata); + + // The full assigned range overlaps even though the actual candidate bounds do not. + assert_eq!( + Some(true), + filter.overlaps_encoded_bounds( + &codec, + &encode_series(&metadata, 1, 100), + &encode_series(&metadata, 1, 200), + ) + ); + assert_eq!( + Some(false), + filter.overlaps_encoded_bounds( + &codec, + &encode_series(&metadata, 1, 1u64 << 63), + &encode_series(&metadata, 1, u64::MAX), + ) + ); + assert_eq!( + Some(true), + filter.overlaps_encoded_bounds( + &codec, + &encode_series(&metadata, 3, 10), + &encode_series(&metadata, 3, 20), + ) + ); + } + + #[test] + fn series_range_keeps_multiple_table_primary_key_bounds() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let assigned = assigned_batch(2, 0, vec![series(1, 10), series(2, 20)]); + let filter = MetricSeriesFilter::new(&assigned); + let codec = SparsePrimaryKeyCodec::new(&metadata); + + assert_eq!( + Some(true), + filter.overlaps_encoded_bounds( + &codec, + &encode_series(&metadata, 1, 1u64 << 63), + &encode_series(&metadata, 2, u64::MAX), + ) + ); + } + + #[test] + fn series_statistics_keep_invalid_or_inverted_bounds() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let assigned = assigned_batch(2, 0, vec![series(1, 10)]); + let filter = MetricSeriesFilter::new(&assigned); + let codec = SparsePrimaryKeyCodec::new(&metadata); + + assert_eq!( + None, + filter.overlaps_encoded_bounds(&codec, b"invalid", b"bounds") + ); + assert_eq!( + None, + filter.overlaps_encoded_bounds( + &codec, + &encode_series(&metadata, 2, 0), + &encode_series(&metadata, 1, 0), + ) + ); } } diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index 7dae5d0a792..2b9e800fb23 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -53,7 +53,7 @@ use crate::sst::parquet::flat_format::{ time_index_column_index, }; use crate::sst::parquet::json_align::ProjectedRecordBatchStream; -use crate::sst::parquet::prefilter::{CachedPrimaryKeyFilter, primary_key_filter_mask}; +use crate::sst::parquet::prefilter::primary_key_filter_mask; use crate::sst::parquet::reader::{ FlatRowGroupReader, MaybeFilter, RowGroupBuildContext, RowGroupReaderBuilder, SimpleFilterContext, @@ -100,13 +100,12 @@ pub struct FileRange { } impl FileRange { - /// Builds the encoded-primary-key filter selected for this file. - pub(crate) fn primary_key_filter(&self) -> Option { - self.context.reader_builder.primary_key_filter() + /// Returns the region metadata stored in this SST. + pub(crate) fn region_metadata(&self) -> &RegionMetadataRef { + self.context.read_format().metadata() } /// Returns encoded primary-key min/max statistics for this row group. - #[allow(dead_code)] pub(crate) fn primary_key_range(&self) -> Option<(&[u8], &[u8])> { let metadata = self.context.reader_builder.parquet_metadata(); let num_columns = metadata.file_metadata().schema_descr().num_columns(); @@ -301,9 +300,8 @@ impl FileRange { /// filter and this range's existing row selection. /// /// This deliberately bypasses generic predicate prefiltering. The series - /// pruner selected the row group independently, and non-tag simple predicates - /// are applied after merge and dedup by the series reader. - #[allow(dead_code)] + /// pruner selected the row group independently, and simple predicates retained + /// by the disabled prefilter plan are applied precisely before merge. pub(crate) async fn reader_by_primary_key( &self, primary_key_filter: &mut dyn mito_codec::row_converter::PrimaryKeyFilter, @@ -348,13 +346,28 @@ impl FileRange { self.context.compaction_projection_mapper() } + /// Filters a full-projection batch using this range's precise filters. + pub(crate) fn precise_filter_flat( + &self, + input: RecordBatch, + skip_fields: bool, + skip_tags: bool, + ) -> Result> { + self.context + .precise_filter_flat(input, skip_fields, skip_tags) + } + + /// Returns the precise-filter mode configured for this range. + pub(crate) fn pre_filter_mode(&self) -> PreFilterMode { + self.context.pre_filter_mode() + } + /// Returns the file handle of the file range. pub(crate) fn file_handle(&self) -> &FileHandle { self.context.reader_builder.file_handle() } } -#[allow(dead_code)] fn refine_primary_key_selection( masks: &[BooleanArray], original: &Option, @@ -431,8 +444,9 @@ impl FileRangeContext { &self, input: RecordBatch, skip_fields: bool, + skip_tags: bool, ) -> Result> { - self.base.precise_filter_flat(input, skip_fields) + self.base.precise_filter_flat(input, skip_fields, skip_tags) } pub(crate) fn pre_filter_mode(&self) -> PreFilterMode { @@ -534,13 +548,16 @@ impl RangeBase { /// # Arguments /// * `input` - The RecordBatch to filter /// * `skip_fields` - Whether to skip field filters based on PreFilterMode + /// * `skip_tags` - Whether to skip tag filters that were applied in an earlier phase pub(crate) fn precise_filter_flat( &self, input: RecordBatch, skip_fields: bool, + skip_tags: bool, ) -> Result> { let mut tag_decode_state = TagDecodeState::new(); - let mask = self.compute_filter_mask_flat(&input, skip_fields, &mut tag_decode_state)?; + let mask = + self.compute_filter_mask_flat(&input, skip_fields, skip_tags, &mut tag_decode_state)?; // If mask is None, the entire batch is filtered out let Some(mut mask) = mask else { @@ -558,9 +575,15 @@ impl RangeBase { mask = mask.bitand(&partition_mask); } - if mask.count_set_bits() == 0 { + let num_selected = mask.count_set_bits(); + if num_selected == 0 { return Ok(None); } + if num_selected == input.num_rows() { + // Nothing was filtered out, e.g. all filters were skipped by + // `skip_fields`/`skip_tags`. Avoid copying the whole batch. + return Ok(Some(input)); + } let filtered_batch = datatypes::arrow::compute::filter_record_batch(&input, &BooleanArray::from(mask)) @@ -582,10 +605,12 @@ impl RangeBase { /// # Arguments /// * `input` - The RecordBatch to compute mask for /// * `skip_fields` - Whether to skip field filters based on PreFilterMode + /// * `skip_tags` - Whether to skip tag filters that were applied in an earlier phase pub(crate) fn compute_filter_mask_flat( &self, input: &RecordBatch, skip_fields: bool, + skip_tags: bool, tag_decode_state: &mut TagDecodeState, ) -> Result> { let mut mask = BooleanBuffer::new_set(input.num_rows()); @@ -606,6 +631,9 @@ impl RangeBase { if skip_fields && filter_ctx.semantic_type() == SemanticType::Field { continue; } + if skip_tags && filter_ctx.semantic_type() == SemanticType::Tag { + continue; + } // Get the column directly by its projected index. // If the column is missing and it's not a tag/time column, this filter is skipped. @@ -783,7 +811,11 @@ mod tests { use std::sync::Arc; use datafusion_expr::{col, lit}; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::ColumnSchema; + use datatypes::value::Value; use parquet::arrow::arrow_reader::RowSelector; + use partition::expr::col as partition_col; use super::*; use crate::read::read_columns::ReadColumns; @@ -829,10 +861,33 @@ mod tests { let batch = new_record_batch_with_custom_sequence(&["b", "x"], 0, 4, 1); let mask = base - .compute_filter_mask_flat(&batch, false, &mut TagDecodeState::new()) + .compute_filter_mask_flat(&batch, false, false, &mut TagDecodeState::new()) .unwrap() .unwrap(); assert_eq!(mask.count_set_bits(), 0); + + let mask = base + .compute_filter_mask_flat(&batch, false, true, &mut TagDecodeState::new()) + .unwrap() + .unwrap(); + assert_eq!(mask.count_set_bits(), 2); + } + + #[test] + fn test_precise_filter_flat_returns_input_when_nothing_is_filtered() { + let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); + let tag_filter = + SimpleFilterContext::new_opt(&metadata, None, &col("tag_0").eq(lit("z"))).unwrap(); + let base = new_test_range_base(vec![tag_filter]); + let batch = new_record_batch_with_custom_sequence(&["b", "x"], 0, 4, 1); + + // The only filter is a tag filter, and it is skipped, so the batch must + // come back untouched rather than being copied through `filter_record_batch`. + let filtered = base + .precise_filter_flat(batch.clone(), false, true) + .unwrap() + .unwrap(); + assert_eq!(batch, filtered); } #[test] @@ -859,12 +914,51 @@ mod tests { let batch = new_record_batch_with_custom_sequence(&["b", "x"], 0, 4, 1); let mask = base - .compute_filter_mask_flat(&batch, false, &mut TagDecodeState::new()) + .compute_filter_mask_flat(&batch, false, false, &mut TagDecodeState::new()) .unwrap() .unwrap(); assert_eq!(mask.count_set_bits(), 4); } + #[test] + fn test_precise_filter_flat_applies_partition_filter_when_skipping_tags() { + let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); + let tag_filter = + SimpleFilterContext::new_opt(&metadata, None, &col("tag_0").eq(lit("z"))).unwrap(); + let mut base = new_test_range_base(vec![tag_filter]); + let batch = new_record_batch_with_custom_sequence(&["b", "x"], 0, 4, 1); + + let batch_schema = batch.schema(); + let tag_field = batch_schema.field(0); + let partition_schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::from_arrow_type(tag_field.data_type()), + tag_field.is_nullable(), + )])); + let partition_expr = partition_col("tag_0") + .gt_eq(Value::String("a".into())) + .and(partition_col("tag_0").lt(Value::String("c".into()))); + base.partition_filter = Some(PartitionFilterContext { + region_partition_physical_expr: partition_expr + .try_as_physical_expr(partition_schema.arrow_schema()) + .unwrap(), + partition_schema, + }); + + let filtered = base + .precise_filter_flat(batch, false, true) + .unwrap() + .unwrap(); + assert_eq!(filtered.num_rows(), 4); + + let out_of_partition = new_record_batch_with_custom_sequence(&["z", "x"], 0, 4, 1); + assert!( + base.precise_filter_flat(out_of_partition, false, true) + .unwrap() + .is_none() + ); + } + #[test] fn test_refine_primary_key_selection_intersects_original_selection() { let original = Some(RowSelection::from(vec![ diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index ada8f1860ef..5b08a0e4957 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -268,13 +268,14 @@ pub(crate) struct BulkFilterPlan { /// omitted. Callers use this as a pruning filter and must preserve the full predicate for /// authoritative filtering later in the scan. pub(crate) fn build_primary_key_filter( - metadata: &RegionMetadataRef, + sst_metadata: &RegionMetadataRef, + expected_metadata: Option<&RegionMetadata>, predicate: Option<&Predicate>, ) -> Option { let filters = predicate .into_iter() .flat_map(|predicate| predicate.exprs()) - .filter_map(|expr| SimpleFilterContext::new_opt(metadata, None, expr)) + .filter_map(|expr| SimpleFilterContext::new_opt(sst_metadata, expected_metadata, expr)) .filter_map(|filter_ctx| { (filter_ctx.semantic_type() == SemanticType::Tag) .then(|| filter_ctx.filter().as_filter().cloned()) @@ -285,8 +286,8 @@ pub(crate) fn build_primary_key_filter( return None; } - let codec = build_primary_key_codec(metadata.as_ref()); - let filter = codec.primary_key_filter(metadata, Arc::new(filters)); + let codec = build_primary_key_codec(sst_metadata.as_ref()); + let filter = codec.primary_key_filter(sst_metadata, Arc::new(filters)); Some(CachedPrimaryKeyFilter::new(filter)) } @@ -370,13 +371,15 @@ pub(crate) fn build_bulk_filter_plan( /// reader always re-applies the original predicate, so the prefilter pass is purely a /// pruning hint. /// -/// Tag and timestamp predicates that lower to [`SimpleFilterEvaluator`] are an -/// exception — the engine enforces them precisely, so the prefilter pass is the only -/// place they execute. They are never silently dropped. +/// With predicate prefiltering enabled, tag and timestamp predicates that lower to +/// [`SimpleFilterEvaluator`] are an exception — the engine enforces them precisely in +/// the prefilter pass. When it is disabled, all simple filters remain on the normal +/// precise-filter path instead. pub(crate) fn build_reader_filter_plan( predicate: Option<&Predicate>, expected_metadata: Option<&RegionMetadata>, pre_filter_mode: PreFilterMode, + enable_predicate_prefilter: bool, read_format: &FlatReadFormat, codec: &Arc, ) -> ReaderFilterPlan { @@ -417,6 +420,11 @@ pub(crate) fn build_reader_filter_plan( // Prefer cheap simple filters first. They also preserve `Matched` / // `Pruned` states for columns that only exist in expected metadata. if let Some(filter_ctx) = SimpleFilterContext::new_opt(metadata, expected_metadata, expr) { + if !enable_predicate_prefilter { + remaining_simple_filters.push(filter_ctx); + continue; + } + // `Matched` and `Pruned` come from expected-metadata compatibility // and must stay in the main filter list so later phases keep that // outcome. @@ -452,6 +460,10 @@ pub(crate) fn build_reader_filter_plan( continue; } + if !enable_predicate_prefilter { + continue; + } + // Best-effort physical-filter prefilter (see fn-level doc): `new_opt` // returning `None` means the column is not in the projected arrow // schema, and dropping the predicate is safe because the upper @@ -464,6 +476,13 @@ pub(crate) fn build_reader_filter_plan( } } + if !enable_predicate_prefilter { + return ReaderFilterPlan { + remaining_simple_filters, + prefilter_builder: None, + }; + } + let pk_filter_expr_strs = (!pk_filter_contexts.is_empty()).then(|| { let mut expr_strs = pk_filter_contexts .iter() @@ -1153,6 +1172,7 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use common_recordbatch::filter::SimpleFilterEvaluator; + use datafusion_common::ScalarValue; use datafusion_expr::{col, lit}; use datatypes::arrow::array::{ ArrayRef, DictionaryArray, TimestampMillisecondArray, UInt8Array, UInt32Array, UInt64Array, @@ -1557,6 +1577,7 @@ mod tests { ])), None, PreFilterMode::SkipFields, + true, &full_read_format, &codec, ); @@ -1584,6 +1605,7 @@ mod tests { Some(&Predicate::new(vec![col("tag_0").eq(lit("a"))])), None, PreFilterMode::All, + true, &projected_read_format, &metric_codec, ); @@ -1597,6 +1619,24 @@ mod tests { .is_some() ); assert!(pk_prefilter_plan.remaining_simple_filters.is_empty()); + + let disabled_plan = build_reader_filter_plan( + Some(&Predicate::new(vec![ + col("tag_0").eq(lit("a")), + col("field_0").gt(lit(1_u64)), + col("ts").gt_eq(lit(ScalarValue::TimestampMillisecond(Some(1), None))), + ])), + None, + PreFilterMode::All, + false, + &projected_read_format, + &metric_codec, + ); + assert!(disabled_plan.prefilter_builder.is_none()); + assert_eq!( + remaining_simple_filter_columns(&disabled_plan.remaining_simple_filters), + vec!["tag_0", "field_0", "ts"] + ); } #[test] @@ -1622,6 +1662,7 @@ mod tests { Some(&Predicate::new(vec![expr_a.clone(), expr_b.clone()])), None, PreFilterMode::All, + true, &read_format, &codec, ); @@ -1629,6 +1670,7 @@ mod tests { Some(&Predicate::new(vec![expr_b, expr_a])), None, PreFilterMode::All, + true, &read_format, &codec, ); diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index b4f0accf9fa..97a4abfac02 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -85,7 +85,7 @@ use crate::sst::parquet::format::{INTERNAL_COLUMN_NUM, need_override_sequence}; use crate::sst::parquet::json_align::{NestedSchemaAligner, ProjectedRecordBatchStream}; use crate::sst::parquet::metadata::MetadataLoader; use crate::sst::parquet::prefilter::{ - CachedPrimaryKeyFilter, PrefilterContextBuilder, build_reader_filter_plan, execute_prefilter, + PrefilterContextBuilder, build_reader_filter_plan, execute_prefilter, }; use crate::sst::parquet::push_decoder::{ SstParquetRangeFetcher, build_sst_parquet_record_batch_stream, @@ -183,6 +183,8 @@ pub struct ParquetReaderBuilder { compaction: bool, /// Mode to pre-filter columns. pre_filter_mode: PreFilterMode, + /// Whether to run the reduced-column predicate prefilter pass. + enable_predicate_prefilter: bool, /// Whether to decode primary key values eagerly when reading primary key format SSTs. decode_primary_key_values: bool, page_index_policy: PageIndexPolicy, @@ -217,6 +219,7 @@ impl ParquetReaderBuilder { expected_metadata: None, compaction: false, pre_filter_mode: PreFilterMode::All, + enable_predicate_prefilter: true, decode_primary_key_values: false, page_index_policy: Default::default(), defer_optional_page_index: false, @@ -318,6 +321,13 @@ impl ParquetReaderBuilder { self } + /// Sets whether to run the reduced-column predicate prefilter pass. + #[must_use] + pub(crate) fn enable_predicate_prefilter(mut self, enable: bool) -> Self { + self.enable_predicate_prefilter = enable; + self + } + /// Decodes primary key values eagerly when reading primary key format SSTs. #[must_use] pub(crate) fn decode_primary_key_values(mut self, decode: bool) -> Self { @@ -497,6 +507,7 @@ impl ParquetReaderBuilder { self.predicate.as_ref(), self.expected_metadata.as_deref(), self.pre_filter_mode, + self.enable_predicate_prefilter, &read_format, &codec, ); @@ -1838,13 +1849,6 @@ impl RowGroupReaderBuilder { self.prefilter_builder.is_some() } - /// Builds the encoded-primary-key filter selected by the reader filter plan. - pub(crate) fn primary_key_filter(&self) -> Option { - self.prefilter_builder - .as_ref() - .and_then(PrefilterContextBuilder::build_primary_key_filter) - } - /// Builds a parquet record batch stream to read the row group at `row_group_idx`. /// /// If prefiltering is applicable (based on `build_ctx`), this performs a two-phase read: @@ -1855,9 +1859,9 @@ impl RowGroupReaderBuilder { /// Predicates that cannot be lowered to prefilter columns (column not projected, /// expression not supported, etc.) are silently skipped. Correctness rests on the /// DataFusion `FilterExec` above this reader, which always re-applies the original - /// predicate. Tag and timestamp predicates that flow through [`SimpleFilterEvaluator`] - /// are an exception — the engine enforces them precisely, so the prefilter pass is the - /// only place they execute. See [`build_reader_filter_plan`] for the bucketing rules. + /// predicate. With predicate prefiltering enabled, tag and timestamp predicates that + /// flow through [`SimpleFilterEvaluator`] are enforced precisely in this pass. See + /// [`build_reader_filter_plan`] for the bucketing rules and disabled mode. /// /// When the prefilter result selects no rows, the second read still issues but /// parquet-rs short-circuits before any column-chunk IO: the row-group state machine @@ -1907,7 +1911,6 @@ impl RowGroupReaderBuilder { /// /// The series reader uses this after computing its own primary-key-only row /// selection. - #[allow(dead_code)] pub(crate) async fn build_without_prefilter( &self, build_ctx: RowGroupBuildContext<'_>, From 9b55715329a3931646586060391b0c19bb57cf62 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 7 Aug 2026 10:52:54 +0800 Subject: [PATCH 010/126] fix(ci): render query regression bot comment as compact table plus threshold details (#8774) Move case/storage thresholds out of the main summary table into a collapsible
block below it, so measured cases render one row per query instead of an extra Query=N/A storage row that stretched the table. Drop the backtick code span around the status cell (the escaping in text() turned it into a visible backtick) and unify multi-threshold separators on '; '. - query-regression-comment.cjs: classifyThresholds shared helper, syntheticThresholdDetail for case-level thresholds, kind marker on synthetic rows, renderThresholdDetails emitting unescaped details tags with text()-escaped entry content, plain-text status cell - query-regression-comment.test.cjs: updated separator/kind assertions, new test covering details-block rendering and plain-text status Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .github/scripts/query-regression-comment.cjs | 53 +++++++++++++++++-- .../scripts/query-regression-comment.test.cjs | 47 ++++++++++++++-- 2 files changed, 91 insertions(+), 9 deletions(-) diff --git a/.github/scripts/query-regression-comment.cjs b/.github/scripts/query-regression-comment.cjs index 130b1c7efa3..2d6c9c40dc8 100644 --- a/.github/scripts/query-regression-comment.cjs +++ b/.github/scripts/query-regression-comment.cjs @@ -92,7 +92,7 @@ function thresholdStatus(thresholds, query) { const hits = (Array.isArray(thresholds) ? thresholds : []) .filter(item => query === undefined || (hasScopedQuery(item) && String(item.query) === query)) .map(formatThreshold); - return hits.length > 0 ? hits.join(', ') : 'N/A'; + return hits.length > 0 ? hits.join('; ') : 'N/A'; } function hasScopedQuery(threshold) { @@ -115,7 +115,7 @@ function formatThreshold(threshold) { return `${name}${scope.length > 0 ? ` [${scope.join(', ')}]` : ''}: ${status}${reason}`; } -function syntheticThresholdStatus(thresholds, measurementNames) { +function classifyThresholds(thresholds, measurementNames) { const unscoped = []; const unmatched = new Map(); for (const threshold of Array.isArray(thresholds) ? thresholds : []) { @@ -130,7 +130,11 @@ function syntheticThresholdStatus(thresholds, measurementNames) { unmatched.set(query, entries); } } + return { unscoped, unmatched }; +} +function syntheticThresholdStatus(thresholds, measurementNames) { + const { unscoped, unmatched } = classifyThresholds(thresholds, measurementNames); const parts = []; if (unscoped.length > 0) { parts.push(`case/storage threshold: ${thresholdStatus(unscoped)}`); @@ -141,6 +145,21 @@ function syntheticThresholdStatus(thresholds, measurementNames) { return parts.length > 0 ? parts.join('; ') : 'N/A'; } +// Case-level (query-unscoped) thresholds formatted for the collapsible details +// block below the summary table: individual items, `; `-separated, without the +// `case/storage threshold:` label (the section heading already says that). +function syntheticThresholdDetail(thresholds, measurementNames) { + const { unscoped, unmatched } = classifyThresholds(thresholds, measurementNames); + const items = []; + if (unscoped.length > 0) { + items.push(thresholdStatus(unscoped)); + } + for (const query of Array.from(unmatched.keys()).sort()) { + items.push(`unmatched query ${query}: ${thresholdStatus(unmatched.get(query))}`); + } + return items.length > 0 ? items.join('; ') : 'N/A'; +} + function joinDetails(...details) { const present = details.filter(detail => detail && detail !== 'N/A'); return present.length > 0 ? present.join('; ') : 'N/A'; @@ -229,7 +248,7 @@ function collectReportRows(report, reportPath) { ), }; }); - const syntheticThresholds = syntheticThresholdStatus(thresholds, measurementNames); + const syntheticThresholds = syntheticThresholdDetail(thresholds, measurementNames); if (syntheticThresholds !== 'N/A') { rows.push({ caseName: name, @@ -239,22 +258,46 @@ function collectReportRows(report, reportPath) { candidateMedian: 'N/A', regression: 'N/A', threshold: syntheticThresholds, + kind: 'case-thresholds', }); } return rows; } +function renderThresholdDetails(rows) { + const entries = rows.filter(row => row.kind === 'case-thresholds'); + if (entries.length === 0) { + return ''; + } + // The
/ tags are emitted literally; only the entry content + // goes through text() (so `[ ] ( ) |` etc. stay escaped while the tags render). + const lines = [ + '
Case / storage thresholds', + '', + ]; + for (const row of entries) { + lines.push(`- ${text(row.caseName)}: ${text(row.threshold)}`); + } + lines.push('
'); + return lines.join('\n'); +} + function renderSummaryTable(rows) { const lines = [ '| Case | Query | Case status | Base median ms | Candidate median ms | Regression | Threshold |', '| --- | --- | --- | ---: | ---: | ---: | --- |', ]; for (const row of rows) { + if (row.kind === 'case-thresholds') { + continue; + } lines.push( - `| ${text(row.caseName)} | ${text(row.query)} | ${statusEmoji(row.status)} \`${text(row.status)}\` | ${text(row.baseMedian)} | ${text(row.candidateMedian)} | ${text(row.regression)} | ${text(row.threshold)} |` + `| ${text(row.caseName)} | ${text(row.query)} | ${statusEmoji(row.status)} ${text(row.status)} | ${text(row.baseMedian)} | ${text(row.candidateMedian)} | ${text(row.regression)} | ${text(row.threshold)} |` ); } - return lines.join('\n'); + const details = renderThresholdDetails(rows); + const table = lines.join('\n'); + return details === '' ? table : `${table}\n\n${details}`; } module.exports = async function validateQueryRegressionComment({ github, context, core }) { diff --git a/.github/scripts/query-regression-comment.test.cjs b/.github/scripts/query-regression-comment.test.cjs index 7a74a0d88c1..4b3d1a5cbfd 100644 --- a/.github/scripts/query-regression-comment.test.cjs +++ b/.github/scripts/query-regression-comment.test.cjs @@ -187,11 +187,11 @@ test('sorts the base and candidate query union and aggregates scoped thresholds' assert.equal(rows[0].threshold, 'base measurement missing'); assert.equal( rows[1].threshold, - 'candidate measurement missing; p95 [target=base]: warn, absolute [target=candidate, encoding=plain]: pass' + 'candidate measurement missing; p95 [target=base]: warn; absolute [target=candidate, encoding=plain]: pass' ); }); -test('preserves unscoped and unmatched thresholds in a synthetic N/A row', () => { +test('keeps unscoped and unmatched thresholds in a case-thresholds detail row', () => { const rows = collectReportRows(report('thresholds', { base: [{ name: 'measured', latency_ms_median: 10 }], candidate: [{ name: 'measured', latency_ms_median: 11 }], @@ -214,9 +214,10 @@ test('preserves unscoped and unmatched thresholds in a synthetic N/A row', () => assert.equal(rows[0].query, 'measured'); assert.equal(rows[0].threshold, 'query limit [target=base]: passed'); assert.equal(rows[1].query, 'N/A'); + assert.equal(rows[1].kind, 'case-thresholds'); assert.equal( rows[1].threshold, - 'case/storage threshold: min_files [target=base]: passed, min_files [target=candidate]: failed, encoding limit [target=candidate, encoding=plain]: failed; unmatched query not-measured: orphaned limit [target=base, encoding=json]: failed (reason: measurement unavailable)' + 'min_files [target=base]: passed; min_files [target=candidate]: failed; encoding limit [target=candidate, encoding=plain]: failed; unmatched query not-measured: orphaned limit [target=base, encoding=json]: failed (reason: measurement unavailable)' ); }); @@ -239,12 +240,50 @@ test('keeps unscoped thresholds out of undefined and null query rows', () => { assert.equal(byQuery.get('undefined').threshold, 'N/A'); assert.equal(byQuery.get('null').threshold, 'N/A'); assert.equal(rows.filter(row => row.query === 'N/A').length, 1); + assert.equal(byQuery.get('N/A').kind, 'case-thresholds'); assert.equal( byQuery.get('N/A').threshold, - 'case/storage threshold: min_files [target=base]: passed, min_files [target=candidate]: failed' + 'min_files [target=base]: passed; min_files [target=candidate]: failed' ); }); +test('renders case/storage thresholds in a collapsible details block below the table', () => { + const rows = [ + ...collectReportRows(report('first', { + base: [{ name: 'q1', latency_ms_median: 10 }], + candidate: [{ name: 'q1', latency_ms_median: 11 }], + }, [ + { query: 'q1', threshold: 'query limit', target: 'base', status: 'passed' }, + { threshold: 'min_files', target: 'base', status: 'passed' }, + { threshold: 'min_files', target: 'candidate', status: 'failed' }, + ]), '/reports/first/query-regression-report.json'), + ...collectReportRows({ + case: { name: 'broken' }, + status: 'failed', + error: 'connection refused', + }, '/reports/broken/query-regression-report.json'), + ]; + + const rendered = renderSummaryTable(rows); + const detailsIndex = rendered.indexOf('
'); + const table = rendered.slice(0, detailsIndex); + const details = rendered.slice(detailsIndex); + + // Query rows and abnormal N/A rows stay in the main table. + assert.match(table, /\| first \| q1 \| ✅ ok \|/); + assert.match(table, /\| broken \| N\/A \| ❌ failed \|/); + // Storage thresholds are not main-table rows. + assert.doesNotMatch(table, /^\| first \| N\/A \|/m); + assert.doesNotMatch(table, /min_files/); + // Status column is plain text without code-span backticks. + assert.doesNotMatch(table, /`/); + // Details block keeps its tags unescaped and its items '; '-separated. + assert.ok(details.startsWith('
Case / storage thresholds')); + assert.ok(details.includes('- first: min_files \\[target=base\\]: passed; min_files \\[target=candidate\\]: failed')); + assert.ok(details.endsWith('
')); + assert.ok(!rendered.includes('<details>')); +}); + test('escapes Markdown table content, including bare carriage returns', () => { const table = renderSummaryTable([{ caseName: 'safe\r| injected |\n@user ', From 4e43c279ff29e9071394f85b90b874906115d29f Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 7 Aug 2026 11:07:33 +0800 Subject: [PATCH 011/126] fix(object-store): fix unused import on Windows after #8735 (#8752) fix(object-store): make removed-entry lister test portable On Windows, DirEntry is a FindFirstFileW snapshot: file_type() and metadata() keep returning cached data after the file is removed, so read_list_entry() may yield the stale entry instead of None. The strong assertion is Unix-specific; keep the portable checks (call succeeds, path is correct) on all platforms and gate only the is_none() assertion. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/object-store/src/secure_fs.rs | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/src/object-store/src/secure_fs.rs b/src/object-store/src/secure_fs.rs index 496ae34d278..d793dee6e3b 100644 --- a/src/object-store/src/secure_fs.rs +++ b/src/object-store/src/secure_fs.rs @@ -625,11 +625,11 @@ mod tests { // On Windows, `std::fs::DirEntry` is a snapshot taken by // `FindFirstFileW`: `file_type()` and `metadata()` keep returning the // cached data even after the file is removed, so `read_list_entry` - // cannot observe the deletion and returns `Some` instead of `None`. - // This test only applies to platforms where metadata is fetched from - // the live filesystem (e.g. Unix `lstat` returns `ENOENT` after - // removal, which `read_list_entry` turns into `None`). - #[cfg(not(windows))] + // may return `Some` with the stale entry. On platforms where metadata + // is fetched from the live filesystem (e.g. Unix `lstat` returns + // `ENOENT` after removal), `read_list_entry` returns `None` instead. + // The strong assertion is Unix-specific; the portable part still + // verifies the call succeeds and yields the expected path. #[test] fn test_lister_skips_entry_removed_during_iteration() { let temp_dir = create_temp_dir("secure_fs_lister_removed_entry"); @@ -640,7 +640,14 @@ mod tests { let entry = read_dir.next().unwrap().unwrap(); std::fs::remove_file(path).unwrap(); - assert!(read_list_entry(entry, "").unwrap().is_none()); + let result = read_list_entry(entry, "").unwrap(); + + if let Some(entry) = &result { + assert_eq!("removed", entry.path()); + } + + #[cfg(not(windows))] + assert!(result.is_none()); } #[tokio::test] From 3d12273c84c50a0e5e9e7149000489f3a82aa113 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 7 Aug 2026 11:37:14 +0800 Subject: [PATCH 012/126] feat: read-time entity relationships graph over telemetry (M0+M1) (#8614) * feat(table): add entity semantic declarations Define open-ended greptime.semantic.entity.* options, validate entity columns at DDL time, and stamp OTLP trace tables with the service entity declaration. Signed-off-by: Dennis Zhuang * feat: add read-time entity relationships graph Add computed semantic graph tables, typed DataFusion derivation plans for entity registry and trace calls edges, and streaming read-time execution. Signed-off-by: Dennis Zhuang * test: exclude semantic graph tables from table constraints Signed-off-by: Dennis Zhuang * refactor(operator): name the plan-builder source groupings Review feedback: build_registry_plan / build_calls_plan took anonymous (declarations, DataFrame) tuples while the caller already grouped the same fields. Introduce RegistrySource { declarations, scan } and CallsSource { service, scan } next to the builders and flow them through the frontend caller and tests. The frontend-side EntitySource keeps holding a TableRef (the operator builders stay pure over already-built scans), so the named structs live in operator rather than reusing that type. No behavior change. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/catalog/src/kvbackend/builder.rs | 1 + src/catalog/src/kvbackend/manager.rs | 38 +- src/catalog/src/system_schema.rs | 1 + .../information_schema/table_semantics.rs | 17 + .../src/system_schema/semantic_graph.rs | 383 ++++++ src/catalog/src/system_schema/utils.rs | 17 + src/common/catalog/src/consts.rs | 41 + src/frontend/src/instance.rs | 1 + src/frontend/src/instance/builder.rs | 19 + src/frontend/src/instance/entity_graph.rs | 457 +++++++ src/frontend/src/instance/jaeger.rs | 11 +- .../src/instance/otlp/trace_ingest.rs | 11 +- src/operator/src/error.rs | 18 +- .../src/req_convert/insert/stmt_to_region.rs | 8 +- src/operator/src/statement.rs | 1 + src/operator/src/statement/ddl.rs | 171 ++- src/operator/src/statement/semantic_graph.rs | 1090 +++++++++++++++++ src/query/src/datafusion.rs | 6 +- src/servers/src/otlp/trace.rs | 14 +- src/table/src/requests.rs | 13 + src/table/src/requests/semantic.rs | 162 ++- tests-integration/tests/http.rs | 10 +- .../table_constraints.result | 10 +- .../information_schema/table_constraints.sql | 10 +- .../common/information_schema/tables.result | 2 +- .../common/information_schema/tables.sql | 2 +- .../common/system/information_schema.result | 22 +- .../common/system/information_schema.sql | 22 +- .../common/system/pg_catalog.result | 5 + .../standalone/common/system/pg_catalog.sql | 5 + .../common/system/semantic_graph.result | 155 +++ .../common/system/semantic_graph.sql | 94 ++ .../standalone/common/view/create.result | 5 +- tests/cases/standalone/common/view/create.sql | 5 +- tests/runner/src/formatter.rs | 6 +- 35 files changed, 2722 insertions(+), 111 deletions(-) create mode 100644 src/catalog/src/system_schema/semantic_graph.rs create mode 100644 src/frontend/src/instance/entity_graph.rs create mode 100644 src/operator/src/statement/semantic_graph.rs create mode 100644 tests/cases/standalone/common/system/semantic_graph.result create mode 100644 tests/cases/standalone/common/system/semantic_graph.sql diff --git a/src/catalog/src/kvbackend/builder.rs b/src/catalog/src/kvbackend/builder.rs index a71b309e69b..87e180c9420 100644 --- a/src/catalog/src/kvbackend/builder.rs +++ b/src/catalog/src/kvbackend/builder.rs @@ -101,6 +101,7 @@ impl KvBackendCatalogManagerBuilder { } = self; Arc::new_cyclic(|me| KvBackendCatalogManager { information_extension, + entity_graph_provider: Arc::new(std::sync::OnceLock::new()), partition_manager: Arc::new(PartitionRuleManager::new( backend.clone(), cache_registry diff --git a/src/catalog/src/kvbackend/manager.rs b/src/catalog/src/kvbackend/manager.rs index a4d8214a55b..db7361615d0 100644 --- a/src/catalog/src/kvbackend/manager.rs +++ b/src/catalog/src/kvbackend/manager.rs @@ -14,11 +14,12 @@ use std::any::Any; use std::collections::BTreeSet; -use std::sync::{Arc, Weak}; +use std::sync::{Arc, OnceLock, Weak}; use async_stream::try_stream; use common_catalog::consts::{ - DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, INFORMATION_SCHEMA_NAME, PG_CATALOG_NAME, + DEFAULT_CATALOG_NAME, DEFAULT_PRIVATE_SCHEMA_NAME, DEFAULT_SCHEMA_NAME, + INFORMATION_SCHEMA_NAME, PG_CATALOG_NAME, }; use common_error::ext::BoxedError; use common_meta::cache::{ @@ -61,6 +62,7 @@ use crate::process_manager::ProcessManagerRef; use crate::system_schema::SystemSchemaProvider; use crate::system_schema::numbers_table_provider::NumbersTableProvider; use crate::system_schema::pg_catalog::PGCatalogProvider; +use crate::system_schema::semantic_graph::{EntityGraphProviderRef, SemanticGraphTableProvider}; /// Access all existing catalog, schema and tables. /// @@ -71,6 +73,12 @@ use crate::system_schema::pg_catalog::PGCatalogProvider; pub struct KvBackendCatalogManager { /// Provides the extension methods for the `information_schema` tables pub(super) information_extension: InformationExtensionRef, + /// Backs the computed entity-graph tables (`semantic_entities` / + /// `semantic_relationships`). Set once, after the query engine is built, to + /// break the `catalog -> query` cycle; `None` until then. `Arc` so the cell + /// stays shared across `Clone`s of the manager (a plain `OnceLock` would + /// fork on clone). + pub(super) entity_graph_provider: Arc>, /// Manages partition rules. pub(super) partition_manager: PartitionRuleManagerRef, /// Manages table metadata. @@ -97,6 +105,19 @@ impl KvBackendCatalogManager { self.information_extension.clone() } + /// Returns the entity-graph provider, or `None` if it has not been injected + /// yet (before the query engine is built). The computed graph tables stream + /// empty until it is set. + pub fn entity_graph_provider(&self) -> Option { + self.entity_graph_provider.get().cloned() + } + + /// Injects the entity-graph provider once the query engine exists. A second + /// call is a no-op (the first binding wins). + pub fn set_entity_graph_provider(&self, provider: EntityGraphProviderRef) { + let _ = self.entity_graph_provider.set(provider); + } + pub fn partition_manager(&self) -> PartitionRuleManagerRef { self.partition_manager.clone() } @@ -591,6 +612,9 @@ impl SystemCatalog { self.pg_catalog_provider.table_names() } DEFAULT_SCHEMA_NAME => self.numbers_table_provider.table_names(), + // Computed entity-graph tables overlay the physical tables of + // `greptime_private` (the caller appends these to the physical list). + DEFAULT_PRIVATE_SCHEMA_NAME => SemanticGraphTableProvider::table_names(), _ => vec![], } } @@ -611,6 +635,8 @@ impl SystemCatalog { self.numbers_table_provider.table_exists(table) } else if schema == PG_CATALOG_NAME && channel == Channel::Postgres { self.pg_catalog_provider.table(table).is_some() + } else if schema == DEFAULT_PRIVATE_SCHEMA_NAME { + SemanticGraphTableProvider::table_exists(table) } else { false } @@ -654,6 +680,14 @@ impl SystemCatalog { } } else if schema == DEFAULT_SCHEMA_NAME { self.numbers_table_provider.table(table_name) + } else if schema == DEFAULT_PRIVATE_SCHEMA_NAME + && SemanticGraphTableProvider::table_exists(table_name) + { + // Constructed on demand (the provider is a name + a weak ref); the + // system catalog is consulted before physical resolution, so the + // computed tables shadow same-named physical tables by design. + SemanticGraphTableProvider::new(catalog.to_string(), self.catalog_manager.clone()) + .table(table_name) } else { None } diff --git a/src/catalog/src/system_schema.rs b/src/catalog/src/system_schema.rs index 6582704a88f..add91cc75f4 100644 --- a/src/catalog/src/system_schema.rs +++ b/src/catalog/src/system_schema.rs @@ -17,6 +17,7 @@ mod memory_table; pub mod numbers_table_provider; pub mod pg_catalog; pub mod predicate; +pub mod semantic_graph; mod utils; use std::collections::HashMap; diff --git a/src/catalog/src/system_schema/information_schema/table_semantics.rs b/src/catalog/src/system_schema/information_schema/table_semantics.rs index b74885968eb..62587e80c4e 100644 --- a/src/catalog/src/system_schema/information_schema/table_semantics.rs +++ b/src/catalog/src/system_schema/information_schema/table_semantics.rs @@ -453,4 +453,21 @@ mod tests { assert!(row.source.is_none()); assert!(row.options_json.is_none()); } + + #[test] + fn extract_folds_entity_keys_into_json_tail() { + // Entity keys are not promoted columns; they surface verbatim (prefix- + // stripped, sorted) in the JSON tail with no code change to this table. + let info = table_info(&[ + (SEMANTIC_SIGNAL_TYPE, "trace"), + ("greptime.semantic.entity.service.id", "service_name"), + ("greptime.semantic.entity.host.id", "host_id"), + ]); + let row = SemanticRow::extract(&info).unwrap(); + assert_eq!(row.signal_type, Some("trace")); + assert_eq!( + row.options_json.as_deref(), + Some(r#"{"entity.host.id":"host_id","entity.service.id":"service_name"}"#) + ); + } } diff --git a/src/catalog/src/system_schema/semantic_graph.rs b/src/catalog/src/system_schema/semantic_graph.rs new file mode 100644 index 00000000000..049377f5a38 --- /dev/null +++ b/src/catalog/src/system_schema/semantic_graph.rs @@ -0,0 +1,383 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The computed entity-graph tables `greptime_private.semantic_entities` and +//! `greptime_private.semantic_relationships`. +//! +//! They live in `greptime_private`, not `information_schema`: scanning them +//! triggers read-time derivation over telemetry tables (trace self-joins, ...), +//! which breaks the "cheap, metadata-only" expectation users have of +//! `information_schema`. `greptime_private` already signals "system-managed, +//! computed data objects", and it is also where the physical declared-edge +//! table will live (a follow-up), so derived and declared edges will share one +//! schema. +//! +//! These are thin forwarders: their rows are derived at read time by the injected +//! [`EntityGraphProvider`], which enumerates the `table_semantics` declarations, +//! builds typed DataFusion derivation plans, and executes them via the query +//! engine. When no provider is injected (e.g. before the engine is up, or on a +//! non-frontend node) they stream empty. The fixed schemas here must match the +//! columns the provider's plans project; JSON columns are `json` (JSONB), whose +//! Arrow storage type is `Binary` — the derived batches are rebuilt against the +//! declared Arrow schema (which carries the `json` extension metadata) in +//! [`SystemTable::to_stream`]. + +use std::sync::{Arc, LazyLock, Weak}; + +use common_catalog::consts::{ + DEFAULT_PRIVATE_SCHEMA_NAME, SEMANTIC_ENTITIES_TABLE_ID, + SEMANTIC_ENTITIES_TABLE_NAME as SEMANTIC_ENTITIES, SEMANTIC_RELATIONSHIPS_TABLE_ID, + SEMANTIC_RELATIONSHIPS_TABLE_NAME as SEMANTIC_RELATIONSHIPS, +}; +use common_error::ext::BoxedError; +use common_recordbatch::adapter::AsyncRecordBatchStreamAdapter; +use common_recordbatch::{ + DfRecordBatch, EmptyRecordBatchStream, RecordBatch, RecordBatchStreamWrapper, + SendableRecordBatchStream, +}; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; +use futures::StreamExt; +use snafu::ResultExt; +use store_api::storage::{ScanRequest, TableId}; +use table::TableRef; + +use crate::CatalogManager; +use crate::error::{InternalSnafu, Result}; +use crate::system_schema::{SystemSchemaProviderInner, SystemTable, SystemTableRef, utils}; + +pub type EntityGraphProviderRef = Arc; + +/// Produces the rows of the computed entity-graph tables +/// (`semantic_entities` / `semantic_relationships`) at read time. +/// +/// The computed tables are thin forwarders to this provider, which is +/// implemented above the query engine (in the frontend): it enumerates the entity +/// declarations from `table_semantics`, builds the typed derivation plans, and +/// executes them. It is injected into the catalog manager *after* construction — +/// the provider needs the engine, which needs the catalog manager — so this late +/// binding breaks the `catalog -> query` dependency cycle. Keeping derivation out +/// of `catalog` also respects the crate layering (the plan builders live in +/// `operator`). See `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. +#[async_trait::async_trait] +pub trait EntityGraphProvider: Send + Sync { + /// Produces the entity registry (`semantic_entities`) rows for `catalog`. + /// `None` means no source table declared an entity. + async fn scan_entities( + &self, + catalog: &str, + request: ScanRequest, + ) -> std::result::Result, BoxedError>; + + /// Produces the relationship set (`semantic_relationships`) rows for `catalog`. + async fn scan_relationships( + &self, + catalog: &str, + request: ScanRequest, + ) -> std::result::Result, BoxedError>; +} + +/// Serves the computed graph tables under `greptime_private`, overlaid on the +/// schema's physical tables the same way the `numbers` table overlays `public` +/// (the system catalog is consulted before physical table resolution). +pub(crate) struct SemanticGraphTableProvider { + catalog_name: String, + catalog_manager: Weak, +} + +impl SemanticGraphTableProvider { + pub(crate) fn new(catalog_name: String, catalog_manager: Weak) -> Self { + Self { + catalog_name, + catalog_manager, + } + } + + pub(crate) fn table_names() -> Vec { + vec![ + SEMANTIC_ENTITIES.to_string(), + SEMANTIC_RELATIONSHIPS.to_string(), + ] + } + + pub(crate) fn table_exists(name: &str) -> bool { + name == SEMANTIC_ENTITIES || name == SEMANTIC_RELATIONSHIPS + } + + pub(crate) fn table(&self, name: &str) -> Option { + self.build_table(name) + } +} + +impl SystemSchemaProviderInner for SemanticGraphTableProvider { + fn catalog_name(&self) -> &str { + &self.catalog_name + } + + fn schema_name() -> &'static str { + DEFAULT_PRIVATE_SCHEMA_NAME + } + + fn system_table(&self, name: &str) -> Option { + let kind = match name { + SEMANTIC_ENTITIES => GraphTableKind::Entities, + SEMANTIC_RELATIONSHIPS => GraphTableKind::Relationships, + _ => return None, + }; + Some(Arc::new(SemanticGraphTable::new( + kind, + self.catalog_name.clone(), + self.catalog_manager.clone(), + )) as _) + } +} + +fn ts() -> ConcreteDataType { + ConcreteDataType::timestamp_millisecond_datatype() +} + +fn string() -> ConcreteDataType { + ConcreteDataType::string_datatype() +} + +fn json() -> ConcreteDataType { + ConcreteDataType::json_datatype() +} + +/// Schema of `semantic_entities` — the node set of the graph, one row per entity +/// observed in a time window. Must match the registry derivation projection. +/// +/// Columns: +/// - `observed_at` — TIME INDEX; the 60s time bucket the entity was observed in. +/// - `window_start` — start of that observation window. +/// - `window_end` — end of the window (`window_start` + 60s). +/// - `fresh_until` — time up to which the entity is considered present; equals +/// `window_end` for derived rows (the graph is a sliding window, not a +/// current-state table: an entity exists in a query window only if it has +/// observed evidence there). +/// - `entity_type` — the entity's type, e.g. `service`, `host`, `k8s.pod`, +/// `process`, `service.instance` (the OTel-style, possibly dotted, type). +/// - `entity_id` — canonical identifier: the value verbatim for a +/// single-attribute identity, or a sorted `k=v,k=v` rendering for a composite. +/// - `entity_id_attrs` — JSON object of the identifying attributes (the +/// escaping-safe source of truth for composite ids); NULL for single-attribute ids. +/// - `scope` — namespace/environment the id is scoped to; empty when none. +/// - `descriptive` — JSON snapshot of the entity's descriptive (non-identifying) +/// attributes; NULL when no descriptive columns were declared. +/// - `source_tables` — JSON array of the telemetry tables that contributed this entity. +static ENTITIES_SCHEMA: LazyLock = LazyLock::new(|| { + Arc::new(Schema::new(vec![ + ColumnSchema::new("observed_at", ts(), false).with_time_index(true), + ColumnSchema::new("window_start", ts(), true), + ColumnSchema::new("window_end", ts(), true), + ColumnSchema::new("fresh_until", ts(), true), + ColumnSchema::new("entity_type", string(), false), + ColumnSchema::new("entity_id", string(), false), + ColumnSchema::new("entity_id_attrs", json(), true), + ColumnSchema::new("scope", string(), true), + ColumnSchema::new("descriptive", json(), true), + ColumnSchema::new("source_tables", json(), true), + ])) +}); + +/// Schema of `semantic_relationships` — the edge set of the graph, one row per +/// edge observed in a time window. This is the 16-column contract every derived +/// branch and the declared-edge table must project for the top-level `UNION ALL`. +/// +/// Columns: +/// - `observed_at` — TIME INDEX; the 60s time bucket the edge was observed in. +/// - `window_start` / `window_end` — the observation window (`window_start` + 60s). +/// - `fresh_until` — time up to which the edge is considered live; equals +/// `window_end` for derived edges (from `valid_until` for declared edges). +/// - `src_type` / `src_id` — type and canonical id of the source endpoint. +/// - `dst_type` / `dst_id` — type and canonical id of the destination endpoint. +/// - `rel_type` — relationship kind, e.g. `calls`, `runs_on`, `contains`, +/// `part_of`, `depends_on` (direction is src → dst; the inverse is a query concern). +/// - `provenance` — how the edge was obtained: `trace` (derived from spans), +/// `attribute` (shared-identity join), `declared` (hand-inserted), or `agent` +/// (agent-inferred). Part of the edge identity, so edges of different provenance +/// for the same pair coexist. +/// - `confidence` — derivation certainty in `[0, 1]`: `1.0` for paired or +/// declared edges, lower for virtual-node or agent-inferred edges. It does +/// not correct for trace sampling. +/// - `request_count` — RED: number of requests over the window (`calls` edges). +/// - `error_count` — RED: number of errored requests over the window. +/// - `duration_sum` — RED: sum of request durations, in seconds, over the window. +/// - `duration_count`— RED: number of durations summed (pair with `duration_sum` +/// to get an average). +/// - `attributes` — JSON of edge attributes, e.g. `connection_type`, +/// `db.system`, `peer.service`. +static RELATIONSHIPS_SCHEMA: LazyLock = LazyLock::new(|| { + Arc::new(Schema::new(vec![ + ColumnSchema::new("observed_at", ts(), false).with_time_index(true), + ColumnSchema::new("window_start", ts(), true), + ColumnSchema::new("window_end", ts(), true), + ColumnSchema::new("fresh_until", ts(), true), + ColumnSchema::new("src_type", string(), false), + ColumnSchema::new("src_id", string(), false), + ColumnSchema::new("dst_type", string(), false), + ColumnSchema::new("dst_id", string(), false), + ColumnSchema::new("rel_type", string(), false), + ColumnSchema::new("provenance", string(), false), + ColumnSchema::new("confidence", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new("request_count", ConcreteDataType::int64_datatype(), true), + ColumnSchema::new("error_count", ConcreteDataType::int64_datatype(), true), + ColumnSchema::new("duration_sum", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new("duration_count", ConcreteDataType::int64_datatype(), true), + ColumnSchema::new("attributes", json(), true), + ])) +}); + +/// Which computed table this shell represents, so the two share one forwarder. +#[derive(Clone, Copy)] +enum GraphTableKind { + Entities, + Relationships, +} + +/// Forwarder for a computed entity-graph table. +struct SemanticGraphTable { + kind: GraphTableKind, + schema: SchemaRef, + catalog_name: String, + catalog_manager: Weak, +} + +impl SemanticGraphTable { + fn new( + kind: GraphTableKind, + catalog_name: String, + catalog_manager: Weak, + ) -> Self { + let schema = match kind { + GraphTableKind::Entities => ENTITIES_SCHEMA.clone(), + GraphTableKind::Relationships => RELATIONSHIPS_SCHEMA.clone(), + }; + Self { + kind, + schema, + catalog_name, + catalog_manager, + } + } + + async fn derive( + kind: GraphTableKind, + catalog: String, + catalog_manager: Weak, + request: ScanRequest, + ) -> Result> { + let provider = utils::entity_graph_provider(&catalog_manager)?; + // No provider (engine not up / non-frontend node): stream empty. + let Some(provider) = provider else { + return Ok(None); + }; + match kind { + GraphTableKind::Entities => provider.scan_entities(&catalog, request).await, + GraphTableKind::Relationships => provider.scan_relationships(&catalog, request).await, + } + .context(InternalSnafu) + } + + fn align_schema( + stream: SendableRecordBatchStream, + schema: SchemaRef, + ) -> SendableRecordBatchStream { + let batch_schema = schema.clone(); + let arrow_schema = schema.arrow_schema().clone(); + let batches = stream.map(move |batch| { + let batch = batch?; + // The derivation output is structurally identical, but its JSON + // fields are plain Binary without the declared extension metadata. + let batch = DfRecordBatch::try_new( + arrow_schema.clone(), + batch.into_df_record_batch().columns().to_vec(), + ) + .context(common_recordbatch::error::NewDfRecordBatchSnafu)?; + Ok(RecordBatch::from_df_record_batch( + batch_schema.clone(), + batch, + )) + }); + Box::pin(RecordBatchStreamWrapper::new(schema, Box::pin(batches))) + } +} + +impl SystemTable for SemanticGraphTable { + fn table_id(&self) -> TableId { + match self.kind { + GraphTableKind::Entities => SEMANTIC_ENTITIES_TABLE_ID, + GraphTableKind::Relationships => SEMANTIC_RELATIONSHIPS_TABLE_ID, + } + } + + fn table_name(&self) -> &'static str { + match self.kind { + GraphTableKind::Entities => SEMANTIC_ENTITIES, + GraphTableKind::Relationships => SEMANTIC_RELATIONSHIPS, + } + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn to_stream(&self, request: ScanRequest) -> Result { + let schema = self.schema.clone(); + let kind = self.kind; + let catalog = self.catalog_name.clone(); + let catalog_manager = self.catalog_manager.clone(); + + let stream_schema = schema.clone(); + let stream = async move { + let stream = Self::derive(kind, catalog, catalog_manager, request) + .await + .map_err(BoxedError::new) + .context(common_recordbatch::error::ExternalSnafu)?; + Ok(match stream { + Some(stream) => Self::align_schema(stream, stream_schema.clone()), + None => Box::pin(EmptyRecordBatchStream::new(stream_schema.clone())), + }) + }; + + Ok(Box::pin(AsyncRecordBatchStreamAdapter::new( + schema, + Box::pin(stream), + ))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn graph_tables_use_observed_at_as_time_index() { + for schema in [&*ENTITIES_SCHEMA, &*RELATIONSHIPS_SCHEMA] { + assert_eq!( + schema.timestamp_column().map(|column| column.name.as_str()), + Some("observed_at") + ); + } + } + + #[test] + fn relationship_schema_does_not_expose_generation_id() { + assert!( + RELATIONSHIPS_SCHEMA + .column_schema_by_name("generation_id") + .is_none() + ); + } +} diff --git a/src/catalog/src/system_schema/utils.rs b/src/catalog/src/system_schema/utils.rs index 9231a48e981..b989bf91228 100644 --- a/src/catalog/src/system_schema/utils.rs +++ b/src/catalog/src/system_schema/utils.rs @@ -21,9 +21,26 @@ use crate::CatalogManager; use crate::error::{GetInformationExtensionSnafu, Result, UpgradeWeakCatalogManagerRefSnafu}; use crate::information_schema::InformationExtensionRef; use crate::kvbackend::KvBackendCatalogManager; +use crate::system_schema::semantic_graph::EntityGraphProviderRef; pub mod tables; +/// Try to get the entity-graph provider from a `[CatalogManager]` weak reference. +/// Returns `None` when the manager is not the kv-backed one or the provider has +/// not been injected yet (the computed graph tables then stream empty). +pub fn entity_graph_provider( + catalog_manager: &Weak, +) -> Result> { + let catalog_manager = catalog_manager + .upgrade() + .context(UpgradeWeakCatalogManagerRefSnafu)?; + + Ok(catalog_manager + .as_any() + .downcast_ref::() + .and_then(|manager| manager.entity_graph_provider())) +} + /// Try to get the `[InformationExtension]` from `[CatalogManager]` weak reference. pub fn information_extension( catalog_manager: &Weak, diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index 3149f6ccd4e..64bbbcbde86 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -125,6 +125,13 @@ pub const INFORMATION_SCHEMA_FLOW_STATISTICS_TABLE_ID: u32 = 45; // ----- End of information_schema tables ----- +// Computed system tables under `greptime_private` (same virtual-table id space +// as the information_schema tables above). +/// id for greptime_private.semantic_entities (computed entity registry) +pub const SEMANTIC_ENTITIES_TABLE_ID: u32 = 45; +/// id for greptime_private.semantic_relationships (computed edge set) +pub const SEMANTIC_RELATIONSHIPS_TABLE_ID: u32 = 46; + /// ----- Begin of pg_catalog tables ----- pub const PG_CATALOG_TABLE_ID_START: u32 = 256; // Please leave at 128 table ids for Postgres @@ -148,12 +155,34 @@ pub fn is_readonly_schema(schema: &str) -> bool { matches!(schema, INFORMATION_SCHEMA_NAME) } +/// Returns true for tables that must reject DDL/DML: everything in a read-only +/// schema, plus computed system tables overlaid on a writable schema (the +/// entity-graph tables under `greptime_private`). +pub fn is_readonly_table(schema: &str, table: &str) -> bool { + is_readonly_schema(schema) + || (schema == DEFAULT_PRIVATE_SCHEMA_NAME + && matches!( + table, + SEMANTIC_ENTITIES_TABLE_NAME | SEMANTIC_RELATIONSHIPS_TABLE_NAME + )) +} + // ---- special table and fields ---- pub const TRACE_ID_COLUMN: &str = "trace_id"; pub const SPAN_ID_COLUMN: &str = "span_id"; pub const SPAN_NAME_COLUMN: &str = "span_name"; pub const SERVICE_NAME_COLUMN: &str = "service_name"; pub const PARENT_SPAN_ID_COLUMN: &str = "parent_span_id"; +// More fixed columns/values of the `greptime_trace_v1` data model, shared by +// the ingest side (`servers::otlp::trace` re-exports them) and the read-time +// graph derivation so the two cannot silently drift. +pub const TRACE_TIMESTAMP_COLUMN: &str = "timestamp"; +pub const SPAN_KIND_COLUMN: &str = "span_kind"; +pub const SPAN_STATUS_CODE_COLUMN: &str = "span_status_code"; +pub const DURATION_NANO_COLUMN: &str = "duration_nano"; +pub const SPAN_KIND_CLIENT: &str = "SPAN_KIND_CLIENT"; +pub const SPAN_KIND_SERVER: &str = "SPAN_KIND_SERVER"; +pub const SPAN_STATUS_ERROR: &str = "STATUS_CODE_ERROR"; pub const TRACE_TABLE_NAME: &str = "opentelemetry_traces"; pub const TRACE_TABLE_NAME_SESSION_KEY: &str = "trace_table_name"; // ---- End of special table and fields ---- @@ -168,3 +197,15 @@ pub fn trace_operations_table_name(trace_table_name: &str) -> String { format!("{}_operations", trace_table_name) } // ---- End of special table and fields ---- + +// ---- Entity relationship graph tables (live in `greptime_private`) ---- +// The `semantic_` prefix ties them to the semantic layer they derive from +// (`greptime.semantic.entity.*` declarations, `information_schema.table_semantics`); +// no `__` prefix is needed since `greptime_private` already scopes them. +/// Computed entity registry: the node set of the observability graph, derived at +/// read time from tables that declared `greptime.semantic.entity.*` identities. +pub const SEMANTIC_ENTITIES_TABLE_NAME: &str = "semantic_entities"; +/// Computed relationship set: the edge set of the observability graph, derived at +/// read time (`calls`/`runs_on`/... ) and unioned with the declared-edge table. +pub const SEMANTIC_RELATIONSHIPS_TABLE_NAME: &str = "semantic_relationships"; +// ---- End of entity relationship graph tables ---- diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 706e651462b..1b9710073d0 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -14,6 +14,7 @@ pub mod builder; mod dashboard; +mod entity_graph; mod grpc; mod influxdb; mod jaeger; diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 42d234fae41..8e69c5951b8 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -17,7 +17,9 @@ use std::sync::atomic::AtomicBool; use cache::{PARTITION_INFO_CACHE_NAME, TABLE_FLOWNODE_SET_CACHE_NAME, TABLE_ROUTE_CACHE_NAME}; use catalog::CatalogManagerRef; +use catalog::kvbackend::KvBackendCatalogManager; use catalog::process_manager::ProcessManagerRef; +use catalog::system_schema::semantic_graph::EntityGraphProviderRef; use common_base::Plugins; use common_datasource::object_store::LocalFileAccess; use common_event_recorder::EventRecorderImpl; @@ -53,6 +55,7 @@ use crate::events::EventHandlerImpl; use crate::frontend::FrontendOptions; use crate::heartbeat::frontend_peer_addr; use crate::instance::Instance; +use crate::instance::entity_graph::EntityGraphProviderImpl; use crate::instance::region_query::FrontendRegionQueryHandler; /// The frontend [`Instance`] builder. @@ -268,6 +271,22 @@ impl FrontendBuilder { .context(DataFusionSnafu)? .query_engine(); + // Inject the entity-graph provider now that the query engine exists, so the + // computed `greptime_private.semantic_entities` / `semantic_relationships` + // tables can derive rows. Late binding here breaks the `catalog -> query` + // dependency cycle. + if let Some(kv_catalog) = self + .catalog_manager + .as_any() + .downcast_ref::() + { + let provider: EntityGraphProviderRef = Arc::new(EntityGraphProviderImpl::new( + query_engine.clone(), + Arc::downgrade(&self.catalog_manager), + )); + kv_catalog.set_entity_graph_provider(provider); + } + let frontend_peer_addr = frontend_peer_addr(&self.options); let statement_executor = StatementExecutor::new( self.catalog_manager.clone(), diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs new file mode 100644 index 00000000000..1d012b7fc60 --- /dev/null +++ b/src/frontend/src/instance/entity_graph.rs @@ -0,0 +1,457 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Frontend implementation of the [`EntityGraphProvider`]: the live connector that +//! makes the computed `greptime_private.semantic_entities` / +//! `semantic_relationships` tables produce rows. +//! +//! It enumerates the entity-identity declarations by iterating the catalog's +//! `TableInfo` options (`greptime.semantic.entity.*`), builds the read-time +//! derivation plans as typed DataFusion `Expr`s over the declaring tables' +//! DataFrames (`operator::statement::semantic_graph`), and executes them through +//! the query engine. Injected into the catalog manager after the engine is built, +//! breaking the `catalog -> query` cycle. + +use std::collections::HashMap; +use std::sync::Weak; + +use async_trait::async_trait; +use catalog::CatalogManager; +use catalog::system_schema::semantic_graph::EntityGraphProvider; +use common_catalog::consts::{ + DEFAULT_PRIVATE_SCHEMA_NAME, DEFAULT_SCHEMA_NAME, INFORMATION_SCHEMA_NAME, PG_CATALOG_NAME, + SERVICE_NAME_COLUMN, +}; +use common_error::ext::BoxedError; +use common_query::OutputData; +use common_recordbatch::SendableRecordBatchStream; +use common_telemetry::warn; +use datafusion::dataframe::DataFrame; +use datafusion_expr::LogicalPlan; +use futures::TryStreamExt; +use operator::statement::semantic_graph::{ + CallsSource, EntityDeclaration, GraphWindow, RegistrySource, build_calls_plan, + build_registry_plan, +}; +use query::QueryEngineRef; +use session::context::{QueryContextBuilder, QueryContextRef}; +use snafu::ResultExt; +use store_api::storage::ScanRequest; +use table::TableRef; +use table::metadata::TableInfo; +use table::requests::{ + EntityRole, is_trace_v1_table, parse_entity_columns, parse_entity_option_key, +}; + +use crate::error; + +/// The live [`EntityGraphProvider`], backed by the query engine. +pub struct EntityGraphProviderImpl { + query_engine: QueryEngineRef, + catalog_manager: Weak, +} + +struct EntitySource { + declarations: Vec, + table: TableRef, +} + +impl EntityGraphProviderImpl { + pub fn new(query_engine: QueryEngineRef, catalog_manager: Weak) -> Self { + Self { + query_engine, + catalog_manager, + } + } + + /// Parses `greptime.semantic.entity..{id|descriptive|scope}` options of + /// one table into per-type declarations. A type with no `id` columns is skipped. + fn parse_declarations(table_info: &TableInfo) -> Vec { + let Some(time_index) = table_info + .meta + .schema + .timestamp_column() + .map(|c| c.name.clone()) + else { + return vec![]; + }; + + // entity_type -> (id_columns, descriptive_columns, scope_columns) + type RoleColumns = (Vec, Vec, Vec); + let mut by_type: HashMap = HashMap::new(); + for (key, value) in &table_info.meta.options.extra_options { + let Some((entity_type, role)) = parse_entity_option_key(key) else { + continue; + }; + let cols = parse_entity_columns(value); + let entry = by_type.entry(entity_type.to_string()).or_default(); + match role { + EntityRole::Id => entry.0 = cols, + EntityRole::Descriptive => entry.1 = cols, + EntityRole::Scope => entry.2 = cols, + } + } + + by_type + .into_iter() + .filter(|(_, (id, _, _))| !id.is_empty()) + .filter_map( + |(entity_type, (id_columns, descriptive_columns, scope_columns))| { + // A stale declaration (e.g. its column was dropped later) + // must not poison every graph scan; skip it. + let schema = &table_info.meta.schema; + if let Some(missing) = id_columns + .iter() + .chain(&descriptive_columns) + .chain(&scope_columns) + .find(|c| schema.column_schema_by_name(c).is_none()) + { + warn!( + "Skipping entity declaration `{}` of table `{}`: column `{}` not found", + entity_type, table_info.name, missing + ); + return None; + } + Some(EntityDeclaration { + schema: table_info.schema_name.clone(), + table: table_info.name.clone(), + time_index: time_index.clone(), + entity_type, + id_columns, + descriptive_columns, + scope_columns, + }) + }, + ) + .collect() + } + + /// All entity declarations of one table. Trace-v1 tables created before the + /// ingest-side auto-stamp carry no `entity.service.id` option; synthesize it + /// (their schema is fixed), so their `calls` edges have endpoint entities. + /// A table with an explicit service declaration never gets the synthesized + /// one — not even when the explicit declaration is invalid and skipped: + /// silently falling back would change entity identity behind the user's back. + fn declarations_for(table_info: &TableInfo) -> Vec { + let mut declarations = Self::parse_declarations(table_info); + let has_explicit_service = table_info.meta.options.extra_options.keys().any(|key| { + parse_entity_option_key(key) + .is_some_and(|(ty, role)| ty == "service" && role == EntityRole::Id) + }); + if is_trace_v1_table(table_info) + && !has_explicit_service + && table_info + .meta + .schema + .column_schema_by_name(SERVICE_NAME_COLUMN) + .is_some() + && let Some(time_index) = table_info + .meta + .schema + .timestamp_column() + .map(|c| c.name.clone()) + { + declarations.push(EntityDeclaration { + schema: table_info.schema_name.clone(), + table: table_info.name.clone(), + time_index, + entity_type: "service".to_string(), + id_columns: vec![SERVICE_NAME_COLUMN.to_string()], + descriptive_columns: vec![], + scope_columns: vec![], + }); + } + declarations + } + + /// Enumerates entity declarations and trace tables across a catalog. Trace + /// tables are keyed off the engine-native `table_data_model` option (same + /// check as the Jaeger query path) so pre-existing trace tables without the + /// newer `greptime.semantic.*` stamps are recognized too. + async fn enumerate( + &self, + catalog: &str, + ) -> Result<(Vec, Vec<(EntityDeclaration, TableRef)>), BoxedError> { + let Some(catalog_manager) = self.catalog_manager.upgrade() else { + return Ok((vec![], vec![])); + }; + + let mut declarations = vec![]; + let mut traces = vec![]; + let schemas = catalog_manager + .schema_names(catalog, None) + .await + .map_err(BoxedError::new)?; + for schema in schemas { + // User telemetry never lives in the system schemas; skip them to avoid + // scanning information_schema (including the computed graph tables) etc. + if schema == INFORMATION_SCHEMA_NAME + || schema == PG_CATALOG_NAME + || schema == DEFAULT_PRIVATE_SCHEMA_NAME + { + continue; + } + let mut tables = catalog_manager.tables(catalog, &schema, None); + while let Some(table) = tables.try_next().await.map_err(BoxedError::new)? { + let table_info = table.table_info(); + let table_declarations = Self::declarations_for(&table_info); + if is_trace_v1_table(&table_info) { + // TODO(entity-graph): validate the complete fixed trace-v1 + // schema before adding the table to calls derivation, and + // skip malformed/stale tables instead of letting one poison + // the whole relationship scan. + match table_declarations + .iter() + .find(|d| d.entity_type == "service") + { + Some(service) => traces.push((service.clone(), table.clone())), + // No usable service identity: the table cannot + // contribute calls edges (see declarations_for). + None => warn!( + "Trace table `{}` has no usable service declaration; skipping calls derivation", + table_info.name + ), + } + } + if !table_declarations.is_empty() { + declarations.push(EntitySource { + declarations: table_declarations, + table, + }); + } + } + } + Ok((declarations, traces)) + } + + /// The time window to derive over, taken from the scan's time predicate. + /// + /// TODO(entity-graph): read the `observed_at` bounds out of `request.filters` + /// and build the window from them. Until then this always returns the default + /// (last hour), so a query whose time filter falls entirely outside the last + /// hour derives nothing and returns empty — a known limitation tracked for the + /// scan-time-predicate-pushdown follow-up. + fn query_window(_request: &ScanRequest) -> GraphWindow { + GraphWindow::default_last_hour() + } + + fn read_table(&self, table: TableRef) -> Result { + self.query_engine.read_table(table).map_err(BoxedError::new) + } + + /// Executes a derivation plan and returns its live result stream. + /// + /// TODO(entity-graph): the QueryContext must come from the outer query + /// instead of being built here, so the derivation inherits the caller's + /// permissions, cancellation and deadline. Requires threading the context + /// through the computed-table scan path (planned next PR). + async fn execute_plan( + &self, + catalog: &str, + plan: LogicalPlan, + ) -> Result, BoxedError> { + let query_ctx: QueryContextRef = QueryContextBuilder::default() + .current_catalog(catalog.to_string()) + .current_schema(DEFAULT_SCHEMA_NAME.to_string()) + .build() + .into(); + let output = self + .query_engine + .execute(plan, query_ctx) + .await + .map_err(BoxedError::new)?; + let stream = match output.data { + OutputData::Stream(stream) => stream, + OutputData::RecordBatches(batches) => batches.as_stream(), + OutputData::AffectedRows(_) => return Ok(None), + }; + Ok(Some(stream)) + } +} + +#[async_trait] +impl EntityGraphProvider for EntityGraphProviderImpl { + async fn scan_entities( + &self, + catalog: &str, + request: ScanRequest, + ) -> Result, BoxedError> { + let (sources, _) = self.enumerate(catalog).await?; + let mut plans = Vec::with_capacity(sources.len()); + for source in sources { + plans.push(RegistrySource { + declarations: source.declarations, + scan: self.read_table(source.table)?, + }); + } + let window = Self::query_window(&request); + let Some(plan) = build_registry_plan(plans, &window) + .context(error::DataFusionSnafu) + .map_err(BoxedError::new)? + else { + return Ok(None); + }; + self.execute_plan(catalog, plan).await + } + + async fn scan_relationships( + &self, + catalog: &str, + request: ScanRequest, + ) -> Result, BoxedError> { + let (_, traces) = self.enumerate(catalog).await?; + let mut scans = Vec::with_capacity(traces.len()); + for (service, trace) in traces { + scans.push(CallsSource { + service, + scan: self.read_table(trace)?, + }); + } + let window = Self::query_window(&request); + let Some(plan) = build_calls_plan(scans, &window) + .context(error::DataFusionSnafu) + .map_err(BoxedError::new)? + else { + return Ok(None); + }; + self.execute_plan(catalog, plan).await + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use common_catalog::consts::{DEFAULT_CATALOG_NAME, MITO_ENGINE}; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::{ColumnSchema, SchemaBuilder}; + use table::metadata::{TableInfoBuilder, TableMeta, TableType}; + use table::requests::{TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1, TableOptions}; + + use super::*; + + fn table_info(columns: &[&str], extra: &[(&str, &str)]) -> TableInfo { + let mut column_schemas = vec![ + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ]; + column_schemas.extend( + columns + .iter() + .map(|c| ColumnSchema::new(*c, ConcreteDataType::string_datatype(), true)), + ); + let schema = Arc::new( + SchemaBuilder::try_from_columns(column_schemas) + .unwrap() + .build() + .unwrap(), + ); + let options = TableOptions { + extra_options: extra + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect::>(), + ..Default::default() + }; + let meta = TableMeta { + schema, + primary_key_indices: vec![], + value_indices: vec![], + engine: MITO_ENGINE.to_string(), + next_column_id: 1, + options, + created_on: Default::default(), + updated_on: Default::default(), + partition_key_indices: vec![], + column_ids: vec![], + }; + TableInfoBuilder::default() + .table_id(1) + .name("t1") + .catalog_name(DEFAULT_CATALOG_NAME) + .schema_name(DEFAULT_SCHEMA_NAME) + .table_version(0) + .table_type(TableType::Base) + .meta(meta) + .build() + .unwrap() + } + + #[test] + fn declaration_referencing_missing_column_is_skipped() { + let info = table_info( + &["service_name"], + &[ + ("greptime.semantic.entity.service.id", "service_name"), + ("greptime.semantic.entity.host.id", "gone"), + ], + ); + let declarations = EntityGraphProviderImpl::declarations_for(&info); + assert_eq!(declarations.len(), 1); + assert_eq!(declarations[0].entity_type, "service"); + + let info = table_info( + &["service_name"], + &[ + ("greptime.semantic.entity.service.id", "service_name"), + ("greptime.semantic.entity.service.descriptive", "gone"), + ], + ); + assert!(EntityGraphProviderImpl::declarations_for(&info).is_empty()); + } + + #[test] + fn trace_v1_table_gets_implicit_service_declaration() { + let info = table_info( + &["service_name"], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + let declarations = EntityGraphProviderImpl::declarations_for(&info); + assert_eq!(declarations.len(), 1); + let decl = &declarations[0]; + assert_eq!(decl.entity_type, "service"); + assert_eq!(decl.id_columns, vec![SERVICE_NAME_COLUMN.to_string()]); + assert_eq!(decl.time_index, "ts"); + + // An explicit service declaration wins; no duplicate is synthesized. + let info = table_info( + &["service_name"], + &[ + (TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1), + ("greptime.semantic.entity.service.id", "service_name"), + ], + ); + assert_eq!(EntityGraphProviderImpl::declarations_for(&info).len(), 1); + + // A trace-model table without the fixed column synthesizes nothing. + let info = table_info(&[], &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)]); + assert!(EntityGraphProviderImpl::declarations_for(&info).is_empty()); + + // An explicit but invalid declaration must not silently fall back to + // the synthesized one: identity must not change behind the user's back. + let info = table_info( + &["service_name"], + &[ + (TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1), + ("greptime.semantic.entity.service.id", "gone"), + ], + ); + assert!(EntityGraphProviderImpl::declarations_for(&info).is_empty()); + } +} diff --git a/src/frontend/src/instance/jaeger.rs b/src/frontend/src/instance/jaeger.rs index bc9ae07b2aa..4d3b137f1c9 100644 --- a/src/frontend/src/instance/jaeger.rs +++ b/src/frontend/src/instance/jaeger.rs @@ -51,7 +51,6 @@ use servers::query_handler::JaegerQueryHandler; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt}; use table::TableRef; -use table::requests::{TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1}; use table::table::adapter::DfTableProviderAdapter; use crate::instance::Instance; @@ -368,15 +367,7 @@ async fn query_trace_table( schema: ctx.current_schema(), })?; - let is_data_model_v1 = table - .clone() - .table_info() - .meta - .options - .extra_options - .get(TABLE_DATA_MODEL) - .map(|s| s.as_str()) - == Some(TABLE_DATA_MODEL_TRACE_V1); + let is_data_model_v1 = table::requests::is_trace_v1_table(&table.table_info()); // collect to set let col_names = table diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs index 0866189b571..41db13661a8 100644 --- a/src/frontend/src/instance/otlp/trace_ingest.rs +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -31,16 +31,16 @@ use pipeline::{GreptimePipelineParams, PipelineWay}; use servers::error::{self, Result as ServerResult}; use servers::otlp; use servers::otlp::coerce::{coerce_value_data, trace_value_datatype}; -use servers::otlp::trace::TraceAuxData; use servers::otlp::trace::span::{TraceSpan, TraceSpanGroup}; use servers::otlp::trace::v1::{TraceBatchSchema, TraceBinaryType, TraceRetryColumns}; +use servers::otlp::trace::{SERVICE_NAME_COLUMN, TraceAuxData}; use servers::query_handler::{PipelineHandlerRef, TraceIngestOutcome}; use session::context::QueryContextRef; use snafu::{IntoError, ResultExt}; use table::requests::{ - SEMANTIC_PIPELINE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SEMANTIC_TRACE_CONVENTIONS, - SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, SIGNAL_TYPE_TRACE, SOURCE_OPENTELEMETRY, - TABLE_DATA_MODEL_TRACE_V1, + SEMANTIC_ENTITY_SERVICE_ID, SEMANTIC_PIPELINE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, + SEMANTIC_TRACE_CONVENTIONS, SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, SIGNAL_TYPE_TRACE, + SOURCE_OPENTELEMETRY, TABLE_DATA_MODEL_TRACE_V1, }; use crate::instance::Instance; @@ -842,6 +842,9 @@ impl Instance { let mut c = (*ctx).clone(); c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_TRACE); c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY); + // `service_name` is a tag column in both trace models, so the main span + // table declares the logical `service` entity (Layer 1 auto-stamp). + c.set_extension(SEMANTIC_ENTITY_SERVICE_ID, SERVICE_NAME_COLUMN); if is_trace_v1_model { c.set_extension(SEMANTIC_PIPELINE, TABLE_DATA_MODEL_TRACE_V1); c.set_extension(SEMANTIC_TRACE_CONVENTIONS, conventions); diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index eabd6477737..324825d4589 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -232,6 +232,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid entity semantic option, reason: {}", reason))] + InvalidEntitySemanticOption { + reason: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Table not found: {}", table_name))] TableNotFound { table_name: String }, @@ -342,6 +349,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Table `{name}` is read-only"))] + TableReadOnly { + name: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Table occurs error"))] Table { #[snafu(implicit)] @@ -931,6 +945,7 @@ impl ErrorExt for Error { | Error::InvalidConfigValue { .. } | Error::InvalidInsertRequest { .. } | Error::InvalidDeleteRequest { .. } + | Error::InvalidEntitySemanticOption { .. } | Error::IllegalPrimaryKeysDef { .. } | Error::SchemaNotFound { .. } | Error::SchemaExists { .. } @@ -970,7 +985,8 @@ impl ErrorExt for Error { } Error::NotSupported { .. } | Error::ShowCreateTableBaseOnly { .. } - | Error::SchemaReadOnly { .. } => StatusCode::Unsupported, + | Error::SchemaReadOnly { .. } + | Error::TableReadOnly { .. } => StatusCode::Unsupported, Error::TableMetadataManager { source, .. } => source.status_code(), Error::ParseSql { source, .. } => source.status_code(), Error::InvalidateTableCache { source, .. } => source.status_code(), diff --git a/src/operator/src/req_convert/insert/stmt_to_region.rs b/src/operator/src/req_convert/insert/stmt_to_region.rs index 39a30c61b1a..7bc0999d29a 100644 --- a/src/operator/src/req_convert/insert/stmt_to_region.rs +++ b/src/operator/src/req_convert/insert/stmt_to_region.rs @@ -32,7 +32,7 @@ use table::metadata::TableInfoRef; use crate::error::{ CatalogSnafu, ColumnDataTypeSnafu, ColumnDefaultValueSnafu, ColumnNoneDefaultValueSnafu, ColumnNotFoundSnafu, InvalidInsertRequestSnafu, InvalidSqlSnafu, MissingInsertBodySnafu, - ParseSqlSnafu, Result, SchemaReadOnlySnafu, TableNotFoundSnafu, + ParseSqlSnafu, Result, SchemaReadOnlySnafu, TableNotFoundSnafu, TableReadOnlySnafu, }; use crate::insert::InstantAndNormalInsertRequests; use crate::req_convert::common::partitioner::Partitioner; @@ -73,6 +73,12 @@ impl<'a> StatementToRegion<'a> { !common_catalog::consts::is_readonly_schema(&schema), SchemaReadOnlySnafu { name: schema } ); + // The computed entity-graph tables overlay a writable schema, so the + // schema-level check above does not cover them. + ensure!( + !common_catalog::consts::is_readonly_table(&schema, &table_name), + TableReadOnlySnafu { name: table_name } + ); let column_names = column_names(stmt, &table_schema); let column_count = column_names.len(); diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 320e7c2e642..fea54bb27c3 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -23,6 +23,7 @@ pub mod ddl; mod describe; mod dml; mod kill; +pub mod semantic_graph; mod set; mod show; mod tql; diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index b0d1ec8dcda..b0986de59f0 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -31,7 +31,9 @@ use api::v1::{ use catalog::CatalogManagerRef; use chrono::Utc; use common_base::regex_pattern::NAME_PATTERN_REG; -use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, is_readonly_schema}; +use common_catalog::consts::{ + DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, is_readonly_schema, is_readonly_table, +}; use common_catalog::{format_full_flow_name, format_full_table_name}; use common_error::ext::BoxedError; #[cfg(feature = "enterprise")] @@ -96,8 +98,8 @@ use table::TableRef; use table::dist_table::DistTable; use table::metadata::{self, TableId, TableInfo, TableMeta, TableType}; use table::requests::{ - AlterKind, AlterTableRequest, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, REPARTITION_COLUMN_HINT_KEY, - TableOptions, + AlterKind, AlterTableRequest, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, EntityRole, + REPARTITION_COLUMN_HINT_KEY, TableOptions, parse_entity_columns, parse_entity_option_key, }; use table::table_name::TableName; use table::table_reference::TableReference; @@ -106,12 +108,12 @@ use crate::error::{ self, AlterExprToRequestSnafu, BuildDfLogicalPlanSnafu, CatalogSnafu, ColumnDataTypeSnafu, ColumnNotFoundSnafu, ConvertSchemaSnafu, CreateLogicalTablesSnafu, DeserializePartitionExprSnafu, EmptyDdlExprSnafu, ExternalSnafu, ExtractTableNamesSnafu, - FlowNotFoundSnafu, InvalidPartitionRuleSnafu, InvalidPartitionSnafu, InvalidSqlSnafu, - InvalidTableNameSnafu, InvalidViewNameSnafu, InvalidViewStmtSnafu, NotSupportedSnafu, - PartitionExprToPbSnafu, Result, SchemaInUseSnafu, SchemaNotFoundSnafu, SchemaReadOnlySnafu, - SerializePartitionExprSnafu, SubstraitCodecSnafu, TableAlreadyExistsSnafu, - TableMetadataManagerSnafu, TableNotFoundSnafu, UnrecognizedTableOptionSnafu, - ViewAlreadyExistsSnafu, + FlowNotFoundSnafu, InvalidEntitySemanticOptionSnafu, InvalidPartitionRuleSnafu, + InvalidPartitionSnafu, InvalidSqlSnafu, InvalidTableNameSnafu, InvalidViewNameSnafu, + InvalidViewStmtSnafu, NotSupportedSnafu, PartitionExprToPbSnafu, Result, SchemaInUseSnafu, + SchemaNotFoundSnafu, SchemaReadOnlySnafu, SerializePartitionExprSnafu, SubstraitCodecSnafu, + TableAlreadyExistsSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu, TableReadOnlySnafu, + UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu, }; use crate::expr_helper::{self, RepartitionRequest, RepartitionSource}; use crate::statement::StatementExecutor; @@ -379,12 +381,7 @@ impl StatementExecutor { partitions: Option, query_ctx: QueryContextRef, ) -> Result { - ensure!( - !is_readonly_schema(&create_table.schema_name), - SchemaReadOnlySnafu { - name: create_table.schema_name.clone() - } - ); + ensure_table_writable(&create_table.schema_name, &create_table.table_name)?; if create_table.engine == METRIC_ENGINE_NAME && create_table @@ -1374,12 +1371,7 @@ impl StatementExecutor { ) -> Result { let mut tables = Vec::with_capacity(table_names.len()); for table_name in table_names { - ensure!( - !is_readonly_schema(&table_name.schema_name), - SchemaReadOnlySnafu { - name: table_name.schema_name.clone() - } - ); + ensure_table_writable(&table_name.schema_name, &table_name.table_name)?; if let Some(table) = self .catalog_manager @@ -1485,12 +1477,7 @@ impl StatementExecutor { time_ranges: Vec<(Timestamp, Timestamp)>, query_context: QueryContextRef, ) -> Result { - ensure!( - !is_readonly_schema(&table_name.schema_name), - SchemaReadOnlySnafu { - name: table_name.schema_name.clone() - } - ); + ensure_table_writable(&table_name.schema_name, &table_name.table_name)?; let table = self .catalog_manager @@ -1537,12 +1524,7 @@ impl StatementExecutor { query_context: &QueryContextRef, ) -> Result { // Check if the schema is read-only. - ensure!( - !is_readonly_schema(&request.schema_name), - SchemaReadOnlySnafu { - name: request.schema_name.clone() - } - ); + ensure_table_writable(&request.schema_name, &request.table_name)?; let table_ref = TableReference::full( &request.catalog_name, @@ -1822,12 +1804,7 @@ impl StatementExecutor { expr: AlterTableExpr, query_context: QueryContextRef, ) -> Result { - ensure!( - !is_readonly_schema(&expr.schema_name), - SchemaReadOnlySnafu { - name: expr.schema_name.clone() - } - ); + ensure_table_writable(&expr.schema_name, &expr.table_name)?; let catalog_name = if expr.catalog_name.is_empty() { DEFAULT_CATALOG_NAME.to_string() @@ -2315,6 +2292,9 @@ pub fn verify_alter( violated: format!("Invalid table name: {}", new_table_name) } ); + // Renaming INTO a computed graph table's name would let the overlay + // shadow the renamed physical table, orphaning its data. + ensure_table_writable(&table_info.schema_name, new_table_name)?; } else if let AlterKind::AddColumns { columns } = alter_kind { // If all the columns are marked as add_if_not_exists and they already exist in the table, // there is no need to perform the alter. @@ -2398,6 +2378,12 @@ pub fn create_table_info( &create_table.time_index, )?; + validate_entity_semantic_options( + &table_options, + &column_name_to_index_map, + &primary_key_indices, + )?; + let meta = TableMeta { schema, primary_key_indices, @@ -2511,6 +2497,54 @@ fn validate_repartition_column_hint( Ok(()) } +/// Rejects DDL against read-only schemas and computed entity-graph tables. +fn ensure_table_writable(schema: &str, table: &str) -> Result<()> { + ensure!( + !is_readonly_schema(schema), + SchemaReadOnlySnafu { + name: schema.to_string() + } + ); + ensure!( + !is_readonly_table(schema, table), + TableReadOnlySnafu { + name: table.to_string() + } + ); + Ok(()) +} + +/// Validates `greptime.semantic.entity..{id|descriptive|scope}` options +/// against the table schema: every named column must exist, and `id` columns must +/// additionally be tag (primary-key) columns so entity identity stays indexable +/// and joinable. See `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. +fn validate_entity_semantic_options( + table_options: &TableOptions, + column_name_to_index_map: &HashMap, + primary_key_indices: &[usize], +) -> Result<()> { + for (key, value) in &table_options.extra_options { + let Some((_, role)) = parse_entity_option_key(key) else { + continue; + }; + let is_id_role = role == EntityRole::Id; + for col in &parse_entity_columns(value) { + let idx = column_name_to_index_map + .get(col) + .context(ColumnNotFoundSnafu { msg: col })?; + ensure!( + !is_id_role || primary_key_indices.contains(idx), + InvalidEntitySemanticOptionSnafu { + reason: format!( + "entity id column `{col}` (option `{key}`) must be a tag/primary-key column" + ), + } + ); + } + } + Ok(()) +} + fn find_partition_columns(partitions: &Option) -> Result> { let columns = if let Some(partitions) = partitions { partitions @@ -2732,12 +2766,7 @@ async fn execute_undrop_table( table_name: TableName, query_context: QueryContextRef, ) -> Result { - ensure!( - !is_readonly_schema(&table_name.schema_name), - SchemaReadOnlySnafu { - name: table_name.schema_name.clone() - } - ); + ensure_table_writable(&table_name.schema_name, &table_name.table_name)?; let dropped = table_metadata_manager .get_dropped_table(&table_name) @@ -3031,6 +3060,58 @@ mod test { assert_eq!(Duration::from_secs(300), ddl_options.timeout); } + #[test] + fn test_validate_entity_semantic_options() { + let col_map = HashMap::from([ + ("service_name".to_string(), 0usize), + ("host_id".to_string(), 1usize), + ("value".to_string(), 2usize), + ]); + // service_name (0) and host_id (1) are tags; value (2) is a field. + let pk = vec![0usize, 1usize]; + let opts = |pairs: &[(&str, &str)]| TableOptions { + extra_options: pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(), + ..Default::default() + }; + + // id on a tag column, composite id on tags, and descriptive on a field + // column are all accepted. + for pairs in [ + &[("greptime.semantic.entity.service.id", "service_name")][..], + &[( + "greptime.semantic.entity.process.id", + "service_name,host_id", + )][..], + &[("greptime.semantic.entity.service.descriptive", "value")][..], + ] { + assert!(validate_entity_semantic_options(&opts(pairs), &col_map, &pk).is_ok()); + } + + // id pointing at a field column is rejected. + let err = validate_entity_semantic_options( + &opts(&[("greptime.semantic.entity.service.id", "value")]), + &col_map, + &pk, + ) + .unwrap_err(); + assert!(matches!( + err, + error::Error::InvalidEntitySemanticOption { .. } + )); + + // id pointing at a missing column is rejected. + let err = validate_entity_semantic_options( + &opts(&[("greptime.semantic.entity.service.id", "nope")]), + &col_map, + &pk, + ) + .unwrap_err(); + assert!(matches!(err, error::Error::ColumnNotFound { .. })); + } + #[test] fn test_validate_and_normalize_flow_options_empty() { assert!( diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs new file mode 100644 index 00000000000..b845dce3b75 --- /dev/null +++ b/src/operator/src/statement/semantic_graph.rs @@ -0,0 +1,1090 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Typed DataFusion plan builders for the read-time derivation behind the +//! computed `semantic_entities` / `semantic_relationships` tables. +//! +//! The derivation is built as typed [`Expr`]s over [`DataFrame`]s (never as SQL +//! text), so user-controlled identifiers are plain values — no quoting or SQL +//! injection surface — and the plans compose with DataFusion's optimizer, +//! including filter pushdown into the source table scans. See +//! `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. + +use std::sync::{Arc, LazyLock}; + +use common_catalog::consts::{ + DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, + SPAN_KIND_COLUMN, SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, + TRACE_ID_COLUMN, TRACE_TIMESTAMP_COLUMN, +}; +use common_function::function::FunctionContext; +use common_function::function_registry::FUNCTION_REGISTRY; +use datafusion::arrow::datatypes::{DataType, TimeUnit}; +use datafusion::dataframe::DataFrame; +use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns}; +use datafusion::functions_aggregate::expr_fn::{count, sum}; +use datafusion::functions_nested::expr_fn::make_array; +use datafusion_common::{Column, Result as DfResult, ScalarValue}; +use datafusion_expr::{Expr, ExprFunctionExt, JoinType, LogicalPlan, ScalarUDF, cast, ident, lit}; + +/// Bin width for the temporal window of derived rows: 60s buckets, matching the +/// service-graph convention. +const BIN_NANOS: i64 = 60 * 1_000_000_000; + +/// A single table's entity-identity declaration, projected from +/// `information_schema.table_semantics` (`greptime.semantic.entity..*`). +#[derive(Debug, Clone)] +pub struct EntityDeclaration { + /// The declaring table's schema, for the qualified `source_tables` lineage. + pub schema: String, + /// The declaring table's name, recorded in `source_tables`. + pub table: String, + /// The table's time index column, used for the temporal window filter. + pub time_index: String, + pub entity_type: String, + /// Identifying columns (>= 1). One column → id verbatim; several → composite. + pub id_columns: Vec, + /// Descriptive columns snapshotted into the `descriptive` JSON (may be empty). + pub descriptive_columns: Vec, + /// Scope columns (namespace/environment). One column → scope verbatim; + /// several → sorted `k=v,k=v`, mirroring the composite id rendering. + pub scope_columns: Vec, +} + +/// The read-time query window as scalar [`Expr`]s. Kept as expressions so the +/// provider can fill them from a scan's time predicate or the product default. +#[derive(Debug, Clone)] +pub struct GraphWindow { + pub start: Expr, + pub end: Expr, +} + +impl GraphWindow { + /// Conservative default when a query carries no explicit time predicate: the + /// last hour, so a bare `SELECT * FROM semantic_entities` never scans every + /// declaring table's full history. This is a product default, not a cap. + /// + /// The bounds are timestamp literals taken from the wall clock at plan-build + /// time — the same snapshot semantics `now()` would give, but as plain + /// constants they prune the source table scans without depending on + /// constant-folding. + pub fn default_last_hour() -> Self { + let end_ms = common_time::util::current_time_millis(); + Self { + start: lit(ScalarValue::TimestampMillisecond( + Some(end_ms - 60 * 60 * 1000), + None, + )), + end: lit(ScalarValue::TimestampMillisecond(Some(end_ms), None)), + } + } +} + +/// An `INTERVAL` literal of `nanos` nanoseconds. +fn interval(nanos: i64) -> Expr { + lit(ScalarValue::new_interval_mdn(0, 0, nanos)) +} + +fn bin_interval() -> Expr { + interval(BIN_NANOS) +} + +/// `date_bin(60s, ts)` cast to millisecond precision, so the output schema is +/// deterministic regardless of the source column's precision (trace tables are +/// nanosecond, metric tables millisecond). +fn bin_ms(ts: Expr) -> Expr { + cast( + datetime_fns::date_bin().call(vec![bin_interval(), ts]), + DataType::Timestamp(TimeUnit::Millisecond, None), + ) +} + +/// Folds union branches without requiring a non-empty input. +fn union_all(acc: Option, branch: DataFrame) -> DfResult> { + Ok(Some(match acc { + Some(acc) => acc.union(branch)?, + None => branch, + })) +} + +/// A column reference qualified by a join-side alias, built without string +/// parsing (so column names containing `.` or `"` stay verbatim). +fn qcol(relation: &str, name: &str) -> Expr { + Expr::Column(Column::new(Some(relation), name)) +} + +fn concat_expr(parts: Vec) -> Expr { + string_fns::concat().call(parts) +} + +/// `coalesce(CAST(column AS STRING), '')`: renders a nullable column for string +/// concatenation without collapsing the result to NULL. +fn cast_string_or_empty(column: &str) -> Expr { + core_fns::coalesce().call(vec![cast(ident(column), DataType::Utf8), lit("")]) +} + +/// The canonical entity-id expression for `id_columns`: the value verbatim for +/// a single column, the sorted `k=v,k=v` rendering for a composite. `col` +/// constructs the column reference (unqualified for registry branches, +/// join-side-qualified for the calls derivation). +fn entity_id_expr(id_columns: &[String], col: &dyn Fn(&str) -> Expr) -> Expr { + if let [id] = id_columns { + cast(col(id), DataType::Utf8) + } else { + let mut cols = id_columns.to_vec(); + cols.sort(); + sorted_kv_expr_with(&cols, false, col) + } +} + +/// The `parse_json` UDF, shared by all derivation plans. Resolved from the +/// global registry once: the UDF is stateless (its `FunctionContext` is unused). +static PARSE_JSON_UDF: LazyLock> = LazyLock::new(|| { + Arc::new( + FUNCTION_REGISTRY + .get_function("parse_json") + .expect("parse_json must be registered") + .provide(FunctionContext::default()), + ) +}); + +/// Parses a JSON text expression into a JSONB value, cast from the UDF's +/// `BinaryView` output to `Binary` — the storage type the computed tables' +/// declared `json` columns map to in Arrow. +fn parse_json_expr(json_text: Expr) -> Expr { + cast(PARSE_JSON_UDF.call(vec![json_text]), DataType::Binary) +} + +/// A NULL literal typed as JSONB storage (`Binary`), so branches without a JSON +/// value union-align with branches that produce one. +fn null_json() -> Expr { + lit(ScalarValue::Binary(None)) +} + +/// Renders a compile-time-known string as JSON text (quoted, fully escaped — +/// including control characters, unlike the runtime value escaping). +fn json_quote(value: &str) -> String { + serde_json::Value::from(value).to_string() +} + +/// Wraps a column reference in `replace` calls so its runtime value is +/// JSON-escaped (`\` then `"`); NULL becomes `''` so one NULL column does not +/// invalidate the whole JSON text (descriptive columns are nullable). +fn json_escaped_value_expr(column: &str) -> Expr { + let escaped_backslash = + string_fns::replace().call(vec![cast_string_or_empty(column), lit("\\"), lit("\\\\")]); + string_fns::replace().call(vec![escaped_backslash, lit("\""), lit("\\\"")]) +} + +/// Builds a JSONB object from `columns` by concatenating a JSON text and parsing +/// it — GreptimeDB has no struct→json function. Keys are JSON-escaped in Rust; +/// values are JSON-escaped at runtime via [`json_escaped_value_expr`]. +/// +/// TODO(entity-graph): replace the text round-trip with a UDF that assembles +/// JSONB directly from the value columns (`jsonb::ObjectBuilder`, keys baked +/// in), dropping the escaping helpers and the per-row parse cost. +fn json_object_expr(columns: &[String]) -> Expr { + if columns.is_empty() { + return parse_json_expr(lit("{}")); + } + let mut parts = vec![lit("{")]; + for (i, column) in columns.iter().enumerate() { + if i > 0 { + parts.push(lit(",")); + } + parts.push(lit(format!("{}:\"", json_quote(column)))); + parts.push(json_escaped_value_expr(column)); + parts.push(lit("\"")); + } + parts.push(lit("}")); + parse_json_expr(concat_expr(parts)) +} + +/// Renders pre-sorted columns as a `k=v,k=v` concatenation. `nullable` +/// coalesces each value to `''` (id columns are tags and non-null; scope +/// columns carry no such guarantee). +fn sorted_kv_expr_with(sorted_cols: &[String], nullable: bool, col: &dyn Fn(&str) -> Expr) -> Expr { + let mut parts = Vec::with_capacity(sorted_cols.len() * 3); + for (i, column) in sorted_cols.iter().enumerate() { + if i > 0 { + parts.push(lit(",")); + } + parts.push(lit(format!("{column}="))); + parts.push(if nullable { + core_fns::coalesce().call(vec![cast(col(column), DataType::Utf8), lit("")]) + } else { + cast(col(column), DataType::Utf8) + }); + } + concat_expr(parts) +} + +const REGISTRY_COLUMNS: [&str; 10] = [ + "observed_at", + "window_start", + "window_end", + "fresh_until", + "entity_type", + "entity_id", + "entity_id_attrs", + "scope", + "descriptive", + "source_tables", +]; + +const REGISTRY_VALID_COLUMN: &str = "__entity_valid"; + +/// Projects all entity declarations of one source table with one source scan. +/// +/// Each output field is first built as an array whose entries correspond to the +/// table's declarations. Unnesting the arrays expands one source row into one +/// row per declared entity without cloning the table scan under `UNION ALL`. +fn registry_source( + first: &EntityDeclaration, + rest: &[EntityDeclaration], + df: DataFrame, + window: &GraphWindow, +) -> DfResult { + let ts = ident(&first.time_index); + let bin = bin_ms(ts.clone()); + let window_predicate = ts + .clone() + .gt_eq(window.start.clone()) + .and(ts.lt(window.end.clone())); + + let mut arrays = vec![Vec::new(); REGISTRY_COLUMNS.len() + 1]; + for decl in std::iter::once(first).chain(rest) { + // CAST even a single-column id: id columns must be tags but not + // necessarily strings, and the computed table declares entity_id + // STRING. Composite ids additionally carry a JSON object of the id + // columns in entity_id_attrs. + let entity_id = entity_id_expr(&decl.id_columns, &|c| ident(c)); + let entity_id_attrs = if decl.id_columns.len() == 1 { + null_json() + } else { + let mut cols = decl.id_columns.clone(); + cols.sort(); + json_object_expr(&cols) + }; + + let scope = match decl.scope_columns.as_slice() { + [] => lit(""), + // Scope columns are not required to be tags, so guard against NULL. + [single] => cast_string_or_empty(single), + _ => { + let mut cols = decl.scope_columns.clone(); + cols.sort(); + sorted_kv_expr_with(&cols, true, &|c| ident(c)) + } + }; + + let descriptive = if decl.descriptive_columns.is_empty() { + null_json() + } else { + json_object_expr(&decl.descriptive_columns) + }; + + let source_tables = parse_json_expr(lit(format!( + "[{}]", + json_quote(&format!("{}.{}", decl.schema, decl.table)) + ))); + + // Tag columns may still be nullable; a NULL identity component + // identifies nothing. Keep this predicate per declaration so a NULL + // identity for one entity does not remove other entities on the row. + let valid = decl.id_columns.iter().fold(lit(true), |predicate, id| { + predicate.and(ident(id).is_not_null()) + }); + + let row = [ + valid, + bin.clone(), + bin.clone(), + bin.clone() + bin_interval(), + bin.clone() + bin_interval(), + lit(decl.entity_type.as_str()), + entity_id, + entity_id_attrs, + scope, + descriptive, + source_tables, + ]; + for (array, value) in arrays.iter_mut().zip(row) { + array.push(value); + } + } + + let array_names = std::iter::once(REGISTRY_VALID_COLUMN) + .chain(REGISTRY_COLUMNS) + .collect::>(); + let array_projection = array_names + .iter() + .zip(arrays) + .map(|(name, values)| make_array(values).alias(*name)) + .collect::>(); + + df.filter(window_predicate)? + .select(array_projection)? + .unnest_columns(&array_names)? + .filter(ident(REGISTRY_VALID_COLUMN))? + .select(REGISTRY_COLUMNS.into_iter().map(ident).collect::>())? + .distinct() +} + +/// A declaring table's scan paired with the entity declarations it carries — +/// the unit `build_registry_plan` derives registry rows from. +pub struct RegistrySource { + pub declarations: Vec, + pub scan: DataFrame, +} + +/// A trace table's scan paired with the `service` declaration it derives +/// `calls` edges for — the unit `build_calls_plan` consumes. +pub struct CallsSource { + pub service: EntityDeclaration, + pub scan: DataFrame, +} + +/// Builds the `semantic_entities` registry plan: one branch and source scan per +/// declaring table, filtered to `window`, then `UNION ALL` across source tables. +/// Returns `None` when nothing declared an entity, so the computed table streams +/// empty. +pub fn build_registry_plan( + sources: Vec, + window: &GraphWindow, +) -> DfResult> { + let mut union_df: Option = None; + for source in sources { + let Some((first, rest)) = source.declarations.split_first() else { + continue; + }; + union_df = union_all(union_df, registry_source(first, rest, source.scan, window)?)?; + } + Ok(union_df.map(DataFrame::into_unoptimized_plan)) +} + +/// The projected columns of `semantic_relationships`, in order. Every derived +/// branch and the declared-edge scan must project exactly these so the top-level +/// `UNION ALL` type-aligns. (The physical declared table additionally stores +/// `valid_from`/`valid_until`, which are applied as a filter, not projected.) +/// Test-only until the declared-edge union branch lands and enforces it in code. +#[cfg(test)] +const RELATIONSHIP_COLUMNS: [&str; 16] = [ + "observed_at", + "window_start", + "window_end", + "fresh_until", + "src_type", + "src_id", + "dst_type", + "dst_id", + "rel_type", + "provenance", + "confidence", + "request_count", + "error_count", + "duration_sum", + "duration_count", + "attributes", +]; + +/// A child server span starts no earlier than 5 minutes before its client span +/// (clock-skew allowance) and no later than 1 hour after it; the bounds keep the +/// join windowed instead of pairing arbitrarily distant spans of a long-lived +/// trace. +const CHILD_SPAN_EARLY_NANOS: i64 = 5 * 60 * 1_000_000_000; +const CHILD_SPAN_LATE_NANOS: i64 = 60 * 60 * 1_000_000_000; + +/// Builds the `calls` derivation (RFC §3a) over `traces`: pair each client span +/// with its child server span on `trace_id` + `parent_span_id`, project to +/// `service`, union the pairs of all trace tables, and aggregate to RED metrics +/// per 60s window in one pass — so an edge observed across several trace tables +/// yields one row, not per-table fragments. This is the plan form of the Tempo +/// servicegraph connector. Virtual-node edges (uninstrumented peers) are a +/// separate branch, added on top of this. Column names are the fixed +/// `greptime_trace_v1` schema (the reason `table_data_model = greptime_trace_v1` +/// is required); `span_status_code` is a string column (`STATUS_CODE_ERROR`), +/// verified against the trace ingest path. Returns `None` when there is no +/// trace table. +pub fn build_calls_plan( + traces: Vec, + window: &GraphWindow, +) -> DfResult> { + let mut union_df: Option = None; + for trace in traces { + union_df = union_all(union_df, calls_pairs(&trace.service, trace.scan, window)?)?; + } + let Some(pairs) = union_df else { + return Ok(None); + }; + + let plan = pairs + .aggregate( + vec![ident("observed_at"), ident("src_id"), ident("dst_id")], + vec![ + count(lit(1)).alias("request_count"), + count(lit(1)) + .filter(ident("status_code").eq(lit(SPAN_STATUS_ERROR))) + .build()? + .alias("error_count"), + sum(ident("duration_nano")).alias("duration_nano_sum"), + ], + )? + .select(vec![ + ident("observed_at"), + ident("observed_at").alias("window_start"), + (ident("observed_at") + bin_interval()).alias("window_end"), + (ident("observed_at") + bin_interval()).alias("fresh_until"), + lit("service").alias("src_type"), + ident("src_id"), + lit("service").alias("dst_type"), + ident("dst_id"), + lit("calls").alias("rel_type"), + lit("trace").alias("provenance"), + lit(1.0_f64).alias("confidence"), + ident("request_count"), + ident("error_count"), + // duration_nano sums in nanoseconds; the contract column is seconds. + (cast(ident("duration_nano_sum"), DataType::Float64) / lit(1e9_f64)) + .alias("duration_sum"), + ident("request_count").alias("duration_count"), + null_json().alias("attributes"), + ])? + .into_unoptimized_plan(); + Ok(Some(plan)) +} + +/// One trace table's client/server span pairs, projected to the aggregation +/// inputs `(observed_at, src_id, dst_id, status_code, duration_nano)`. +/// Endpoint ids are built from the table's `service` entity declaration, so +/// edges land on exactly the entity ids the registry emits (a composite +/// service identity renders the same sorted `k=v` form on both sides). +fn calls_pairs( + service: &EntityDeclaration, + trace: DataFrame, + window: &GraphWindow, +) -> DfResult { + let mut client_pred = ident(SPAN_KIND_COLUMN) + .eq(lit(SPAN_KIND_CLIENT)) + .and(ident(TRACE_TIMESTAMP_COLUMN).gt_eq(window.start.clone())) + .and(ident(TRACE_TIMESTAMP_COLUMN).lt(window.end.clone())); + let mut server_pred = ident(SPAN_KIND_COLUMN) + .eq(lit(SPAN_KIND_SERVER)) + // Static bounds implied by the window and the join's time-proximity + // conditions below; the join bounds reference client.timestamp and + // cannot prune the server-side scan. + .and( + ident(TRACE_TIMESTAMP_COLUMN) + .gt_eq(window.start.clone() - interval(CHILD_SPAN_EARLY_NANOS)), + ) + .and( + ident(TRACE_TIMESTAMP_COLUMN).lt(window.end.clone() + interval(CHILD_SPAN_LATE_NANOS)), + ); + // A NULL identity component identifies nothing, on either endpoint. + for id in &service.id_columns { + client_pred = client_pred.and(ident(id).is_not_null()); + server_pred = server_pred.and(ident(id).is_not_null()); + } + + let client = trace.clone().filter(client_pred)?.alias("client")?; + let server = trace.filter(server_pred)?.alias("server")?; + + let join_conditions = vec![ + qcol("client", TRACE_ID_COLUMN).eq(qcol("server", TRACE_ID_COLUMN)), + qcol("server", PARENT_SPAN_ID_COLUMN).eq(qcol("client", SPAN_ID_COLUMN)), + qcol("server", TRACE_TIMESTAMP_COLUMN) + .gt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) - interval(CHILD_SPAN_EARLY_NANOS)), + qcol("server", TRACE_TIMESTAMP_COLUMN) + .lt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) + interval(CHILD_SPAN_LATE_NANOS)), + ]; + + client + .join_on(server, JoinType::Inner, join_conditions)? + .select(vec![ + bin_ms(qcol("client", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), + // The cast inside entity_id_expr also normalizes tag columns, which + // come out of the storage engine dictionary-encoded. + entity_id_expr(&service.id_columns, &|c| qcol("client", c)).alias("src_id"), + entity_id_expr(&service.id_columns, &|c| qcol("server", c)).alias("dst_id"), + qcol("server", SPAN_STATUS_CODE_COLUMN).alias("status_code"), + qcol("server", DURATION_NANO_COLUMN).alias("duration_nano"), + ])? + // Exclude self-calls on the composed identity: an edge needs two + // distinct service entities. + .filter(ident("src_id").not_eq(ident("dst_id"))) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_catalog::consts::SERVICE_NAME_COLUMN; + use datafusion::arrow::array::{ + Array, ArrayRef, BinaryArray, Float64Array, Int64Array, StringArray, + TimestampMillisecondArray, TimestampNanosecondArray, UInt64Array, + }; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::MemTable; + use datafusion::prelude::SessionContext; + + use super::*; + + fn test_window() -> GraphWindow { + GraphWindow { + start: lit(ScalarValue::TimestampMillisecond(Some(0), None)), + end: lit(ScalarValue::TimestampMillisecond( + Some(10 * 60 * 1000), + None, + )), + } + } + + async fn collect(ctx: &SessionContext, plan: LogicalPlan) -> Vec { + ctx.execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap() + } + + fn json_texts(batch: &RecordBatch, column: usize) -> Vec> { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()) + .map(|i| { + array + .is_valid(i) + .then(|| jsonb::from_slice(array.value(i)).unwrap().to_string()) + }) + .collect() + } + + fn strings(batch: &RecordBatch, column: usize) -> Vec { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()) + .map(|i| array.value(i).to_string()) + .collect() + } + + /// A metric-like table: ms timestamps, service/pid identity, nullable + /// descriptive column with JSON-hostile characters. + fn metric_table_ctx() -> SessionContext { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("service_name", DataType::Utf8, false), + Field::new("pid", DataType::Int64, false), + Field::new("host", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 2_000, 61_000])) as ArrayRef, + Arc::new(StringArray::from(vec!["cart", "cart", "cart"])), + Arc::new(Int64Array::from(vec![42, 42, 42])), + Arc::new(StringArray::from(vec![ + Some(r#"we"ird\host"#), + None, + Some("h2"), + ])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "app_latency", + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + ctx + } + + fn decl(entity_type: &str, id_columns: &[&str]) -> EntityDeclaration { + EntityDeclaration { + schema: "public".to_string(), + table: "app_latency".to_string(), + time_index: "ts".to_string(), + entity_type: entity_type.to_string(), + id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + descriptive_columns: vec![], + scope_columns: vec![], + } + } + + #[tokio::test] + async fn registry_single_column_identity() { + let ctx = metric_table_ctx(); + let df = ctx.table("app_latency").await.unwrap(); + let plan = build_registry_plan( + vec![RegistrySource { + declarations: vec![decl("service", &["service_name"])], + scan: df, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let names = plan + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(); + assert_eq!( + names, + [ + "observed_at", + "window_start", + "window_end", + "fresh_until", + "entity_type", + "entity_id", + "entity_id_attrs", + "scope", + "descriptive", + "source_tables", + ] + ); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + // 3 rows in 2 distinct 60s bins, all the same entity -> 2 rows. + assert_eq!(total, 2); + let batch = &batches[0]; + assert_eq!(strings(batch, 4), vec!["service"; batch.num_rows()]); + assert_eq!(strings(batch, 5), vec!["cart"; batch.num_rows()]); + // Single-column id -> entity_id_attrs and descriptive are typed-JSON NULLs. + assert!(json_texts(batch, 6).iter().all(Option::is_none)); + assert!(json_texts(batch, 8).iter().all(Option::is_none)); + assert_eq!( + json_texts(batch, 9), + vec![Some(r#"["public.app_latency"]"#.to_string()); batch.num_rows()] + ); + } + + #[tokio::test] + async fn registry_composite_identity_and_descriptive_escaping() { + let ctx = metric_table_ctx(); + let df = ctx.table("app_latency").await.unwrap(); + let mut declaration = decl("process", &["service_name", "pid"]); + declaration.descriptive_columns = vec!["host".to_string()]; + let plan = build_registry_plan( + vec![RegistrySource { + declarations: vec![declaration], + scan: df, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let mut rows: Vec<(String, Option, Option)> = batches + .iter() + .flat_map(|batch| { + let ids = strings(batch, 5); + let id_attrs = json_texts(batch, 6); + let descriptives = json_texts(batch, 8); + ids.into_iter() + .zip(id_attrs) + .zip(descriptives) + .map(|((id, attrs), descriptive)| (id, attrs, descriptive)) + .collect::>() + }) + .collect(); + rows.sort(); + + // Composite id -> sorted `k=v,k=v` plus a JSON object of the id columns; + // descriptive JSON escapes `\` and `"` in runtime values, NULL -> "". + assert_eq!( + rows, + vec![ + ( + "pid=42,service_name=cart".to_string(), + Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), + Some(r#"{"host":""}"#.to_string()), + ), + ( + "pid=42,service_name=cart".to_string(), + Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), + Some(r#"{"host":"h2"}"#.to_string()), + ), + ( + "pid=42,service_name=cart".to_string(), + Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), + Some(r#"{"host":"we\"ird\\host"}"#.to_string()), + ), + ] + ); + } + + #[tokio::test] + async fn registry_scope_variants() { + let ctx = metric_table_ctx(); + + // Single scope column: its (NULL-safe) value verbatim. + let mut single = decl("service", &["service_name"]); + single.scope_columns = vec!["host".to_string()]; + let df = ctx.table("app_latency").await.unwrap(); + let plan = build_registry_plan( + vec![RegistrySource { + declarations: vec![single], + scan: df, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + let batches = collect(&ctx, plan).await; + let mut scopes: Vec = batches.iter().flat_map(|b| strings(b, 7)).collect(); + scopes.sort(); + assert_eq!(scopes, vec!["", "h2", r#"we"ird\host"#]); + + // Multiple scope columns: sorted `k=v,k=v`. + let mut multi = decl("service", &["service_name"]); + multi.scope_columns = vec!["pid".to_string(), "host".to_string()]; + let df = ctx.table("app_latency").await.unwrap(); + let plan = build_registry_plan( + vec![RegistrySource { + declarations: vec![multi], + scan: df, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + let batches = collect(&ctx, plan).await; + let mut scopes: Vec = batches.iter().flat_map(|b| strings(b, 7)).collect(); + scopes.sort(); + assert_eq!( + scopes, + vec![ + "host=,pid=42", + "host=h2,pid=42", + r#"host=we"ird\host,pid=42"# + ] + ); + } + + #[tokio::test] + async fn registry_skips_null_identity_rows() { + let ctx = metric_table_ctx(); + let df = ctx.table("app_latency").await.unwrap(); + let plan = build_registry_plan( + vec![RegistrySource { + declarations: vec![decl("host", &["host"])], + scan: df, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + let batches = collect(&ctx, plan).await; + let mut ids: Vec = batches.iter().flat_map(|b| strings(b, 5)).collect(); + ids.sort(); + assert_eq!(ids, vec!["h2", r#"we"ird\host"#]); + } + + #[tokio::test] + async fn registry_expands_declarations_with_one_source_scan() { + let ctx = metric_table_ctx(); + let df = ctx.table("app_latency").await.unwrap(); + let plan = build_registry_plan( + vec![RegistrySource { + declarations: vec![decl("service", &["service_name"]), decl("host", &["pid"])], + scan: df, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + assert_eq!( + plan.display_indent() + .to_string() + .matches("TableScan: app_latency") + .count(), + 1 + ); + let batches = collect(&ctx, plan).await; + let mut types: Vec = batches.iter().flat_map(|b| strings(b, 4)).collect(); + types.sort(); + assert_eq!(types, vec!["host", "host", "service", "service"]); + + // No declarations -> no plan. + assert!( + build_registry_plan(vec![], &test_window()) + .unwrap() + .is_none() + ); + } + + /// A trace-like table in the fixed `greptime_trace_v1` shape (ns timestamps). + fn trace_table_ctx() -> SessionContext { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TRACE_TIMESTAMP_COLUMN, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new(TRACE_ID_COLUMN, DataType::Utf8, false), + Field::new(SPAN_ID_COLUMN, DataType::Utf8, false), + Field::new(PARENT_SPAN_ID_COLUMN, DataType::Utf8, true), + Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false), + Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false), + Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false), + Field::new("service_namespace", DataType::Utf8, false), + Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false), + ])); + const MS: i64 = 1_000_000; + // Two client->server pairs frontend->cart (one errored), one pair + // cart->cart (self-call, excluded), one unmatched client span. + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampNanosecondArray::from(vec![ + 1_000 * MS, // client frontend->cart + 1_010 * MS, // server cart + 2_000 * MS, // client frontend->cart (error) + 2_010 * MS, // server cart (error) + 3_000 * MS, // client cart->cart (self-call) + 3_010 * MS, // server cart + 4_000 * MS, // client with no matching server + ])) as ArrayRef, + Arc::new(StringArray::from(vec![ + "t1", "t1", "t2", "t2", "t3", "t3", "t4", + ])), + Arc::new(StringArray::from(vec![ + "c1", "s1", "c2", "s2", "c3", "s3", "c4", + ])), + Arc::new(StringArray::from(vec![ + None, + Some("c1"), + None, + Some("c2"), + None, + Some("c3"), + None, + ])), + Arc::new(StringArray::from(vec![ + "SPAN_KIND_CLIENT", + "SPAN_KIND_SERVER", + "SPAN_KIND_CLIENT", + "SPAN_KIND_SERVER", + "SPAN_KIND_CLIENT", + "SPAN_KIND_SERVER", + "SPAN_KIND_CLIENT", + ])), + Arc::new(StringArray::from(vec![ + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + "STATUS_CODE_ERROR", + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + ])), + Arc::new(StringArray::from(vec![ + "frontend", "cart", "frontend", "cart", "cart", "cart", "frontend", + ])), + Arc::new(StringArray::from(vec!["ns1"; 7])), + Arc::new(UInt64Array::from(vec![ + 0, + 500_000_000, // 0.5s + 0, + 1_500_000_000, // 1.5s + 0, + 100, + 0, + ])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "opentelemetry_traces", + Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch.clone()]]).unwrap()), + ) + .unwrap(); + ctx.register_table( + "opentelemetry_traces_2", + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + ctx + } + + fn trace_service_decl(id_columns: &[&str]) -> EntityDeclaration { + EntityDeclaration { + schema: "public".to_string(), + table: "opentelemetry_traces".to_string(), + time_index: TRACE_TIMESTAMP_COLUMN.to_string(), + entity_type: "service".to_string(), + id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + descriptive_columns: vec![], + scope_columns: vec![], + } + } + + #[tokio::test] + async fn calls_plan_aggregates_red_metrics() { + let ctx = trace_table_ctx(); + let trace = ctx.table("opentelemetry_traces").await.unwrap(); + let plan = build_calls_plan( + vec![CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: trace, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let names = plan + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(); + assert_eq!(names, RELATIONSHIP_COLUMNS); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + // frontend->cart only: the self-call and the unmatched client drop out; + // both pairs land in the same 60s bin. + assert_eq!(total, 1); + let batch = &batches[0]; + assert_eq!(strings(batch, 5), vec!["frontend"]); + assert_eq!(strings(batch, 7), vec!["cart"]); + assert_eq!(strings(batch, 8), vec!["calls"]); + assert_eq!(strings(batch, 9), vec!["trace"]); + + let request_count = batch + .column(11) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(request_count.value(0), 2); + let error_count = batch + .column(12) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(error_count.value(0), 1); + let duration_sum = batch + .column(13) + .as_any() + .downcast_ref::() + .unwrap(); + assert!((duration_sum.value(0) - 2.0).abs() < 1e-9); + let duration_count = batch + .column(14) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(duration_count.value(0), 2); + // Derived calls edges carry no attributes: a typed-JSON NULL. + assert!(json_texts(batch, 15)[0].is_none()); + } + + #[tokio::test] + async fn calls_plan_merges_edges_across_trace_tables() { + let ctx = trace_table_ctx(); + let t1 = ctx.table("opentelemetry_traces").await.unwrap(); + let t2 = ctx.table("opentelemetry_traces_2").await.unwrap(); + let plan = build_calls_plan( + vec![ + CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: t1, + }, + CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: t2, + }, + ], + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + // The same frontend->cart edge from both tables folds into one row with + // summed RED metrics. + assert_eq!(total, 1); + let batch = &batches[0]; + let request_count = batch + .column(11) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(request_count.value(0), 4); + let error_count = batch + .column(12) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(error_count.value(0), 2); + let duration_sum = batch + .column(13) + .as_any() + .downcast_ref::() + .unwrap(); + assert!((duration_sum.value(0) - 4.0).abs() < 1e-9); + } + + #[tokio::test] + async fn calls_endpoints_follow_service_declaration() { + let ctx = trace_table_ctx(); + let trace = ctx.table("opentelemetry_traces").await.unwrap(); + let plan = build_calls_plan( + vec![CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN, "service_namespace"]), + scan: trace, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 1); + let batch = &batches[0]; + // Composite service identity renders the same sorted `k=v` form the + // registry emits, so edges land on registry entity ids. + assert_eq!( + strings(batch, 5), + vec!["service_name=frontend,service_namespace=ns1"] + ); + assert_eq!( + strings(batch, 7), + vec!["service_name=cart,service_namespace=ns1"] + ); + } +} diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 2d4f237dd63..754199a78fc 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -24,7 +24,7 @@ use std::sync::Arc; use async_trait::async_trait; use common_base::Plugins; -use common_catalog::consts::is_readonly_schema; +use common_catalog::consts::is_readonly_table; use common_error::ext::BoxedError; use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; @@ -259,7 +259,7 @@ impl DatafusionQueryEngine { let table_schema = table.schema(); ensure!( - !is_readonly_schema(&schema_name), + !is_readonly_table(&schema_name, &table_name), TableReadOnlySnafu { table: table_name } ); @@ -307,7 +307,7 @@ impl DatafusionQueryEngine { let table_name = table_name.table.to_string(); ensure!( - !is_readonly_schema(&schema_name), + !is_readonly_table(&schema_name, &table_name), TableReadOnlySnafu { table: table_name } ); diff --git a/src/servers/src/otlp/trace.rs b/src/servers/src/otlp/trace.rs index f9a0b5e4528..a8439cdb2c4 100644 --- a/src/servers/src/otlp/trace.rs +++ b/src/servers/src/otlp/trace.rs @@ -20,6 +20,12 @@ pub mod v1; use std::collections::HashSet; use api::v1::RowInsertRequests; +// Column names. The `greptime_trace_v1` fixed columns are canonically defined +// in `common_catalog::consts` (shared with the read-time graph derivation). +pub use common_catalog::consts::{ + DURATION_NANO_COLUMN, SERVICE_NAME_COLUMN, SPAN_KIND_COLUMN, + SPAN_STATUS_CODE_COLUMN as SPAN_STATUS_CODE, TRACE_TIMESTAMP_COLUMN as TIMESTAMP_COLUMN, +}; pub use common_catalog::consts::{ PARENT_SPAN_ID_COLUMN, SPAN_ID_COLUMN, SPAN_NAME_COLUMN, TRACE_ID_COLUMN, }; @@ -30,12 +36,6 @@ use crate::error::{NotSupportedSnafu, Result}; use crate::otlp::trace::span::TraceSpan; use crate::query_handler::PipelineHandlerRef; -// column names -pub const SERVICE_NAME_COLUMN: &str = "service_name"; -pub const TIMESTAMP_COLUMN: &str = "timestamp"; -pub const DURATION_NANO_COLUMN: &str = "duration_nano"; -pub const SPAN_KIND_COLUMN: &str = "span_kind"; -pub const SPAN_STATUS_CODE: &str = "span_status_code"; pub const SPAN_STATUS_MESSAGE_COLUMN: &str = "span_status_message"; pub const SPAN_ATTRIBUTES_COLUMN: &str = "span_attributes"; pub const SPAN_EVENTS_COLUMN: &str = "span_events"; @@ -64,7 +64,7 @@ pub const SPAN_KIND_PREFIX: &str = "SPAN_KIND_"; // The span status code prefix in the database. pub const SPAN_STATUS_PREFIX: &str = "STATUS_CODE_"; pub const SPAN_STATUS_UNSET: &str = "STATUS_CODE_UNSET"; -pub const SPAN_STATUS_ERROR: &str = "STATUS_CODE_ERROR"; +pub use common_catalog::consts::SPAN_STATUS_ERROR; /// Deduplicated auxiliary trace entities derived from successfully ingested /// spans. diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 8c736e26749..b4a7d5e052b 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -59,6 +59,19 @@ pub const FILE_TABLE_FORMAT_KEY: &str = "format"; pub const TABLE_DATA_MODEL: &str = "table_data_model"; pub const TABLE_DATA_MODEL_TRACE_V1: &str = "greptime_trace_v1"; +/// Returns true if the table stores spans in the `greptime_trace_v1` data model +/// (fixed span columns), the shape the Jaeger query path and the entity-graph +/// derivation rely on. +pub fn is_trace_v1_table(table_info: &crate::metadata::TableInfo) -> bool { + table_info + .meta + .options + .extra_options + .get(TABLE_DATA_MODEL) + .map(|v| v == TABLE_DATA_MODEL_TRACE_V1) + .unwrap_or(false) +} + pub const OTLP_METRIC_COMPAT_KEY: &str = "otlp_metric_compat"; pub const OTLP_METRIC_COMPAT_PROM: &str = "prom"; diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index 33c0d53576f..cb9d6aa7476 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -76,6 +76,40 @@ pub const SEMANTIC_METRIC_METADATA_QUALITY: &str = "greptime.semantic.metric.met /// consumer uses to look the metric up in the OTel semantic conventions. pub const SEMANTIC_METRIC_ORIGINAL_NAME: &str = "greptime.semantic.metric.original_name"; +// ---- Entity keys (open sub-namespace) ---- + +/// Reserved prefix for the entity-identity sub-namespace: +/// `greptime.semantic.entity..{id|descriptive|scope}`. Unlike the rest of +/// the vocabulary (a closed whitelist), entity types are open-ended (`service`, +/// `host`, `k8s.pod`, `process`, `agent`, custom, ...), so keys here are validated +/// by prefix + shape rather than membership. See +/// `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. +pub const SEMANTIC_ENTITY_PREFIX: &str = "greptime.semantic.entity."; + +/// The well-known entity-identity key auto-stamped on OTLP trace tables; its value +/// is the `service_name` tag column, declaring the logical `service` entity. +pub const SEMANTIC_ENTITY_SERVICE_ID: &str = "greptime.semantic.entity.service.id"; + +/// The role a set of columns plays for an entity: `id` (identifying attributes, +/// must be tag columns — enforced at DDL time), `descriptive`, or `scope`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EntityRole { + Id, + Descriptive, + Scope, +} + +impl EntityRole { + fn parse(role: &str) -> Option { + match role { + "id" => Some(EntityRole::Id), + "descriptive" => Some(EntityRole::Descriptive), + "scope" => Some(EntityRole::Scope), + _ => None, + } + } +} + // ---- Value constants ---- pub const SIGNAL_TYPE_TRACE: &str = "trace"; @@ -115,13 +149,58 @@ pub const SEMANTIC_OPTION_KEYS: &[&str] = &[ SEMANTIC_METRIC_ORIGINAL_NAME, ]; -/// Returns true if `key` is a recognised semantic table-option key (whitelist). +/// Returns true if `ty` is a syntactically valid entity type, e.g. `service`, +/// `host`, `k8s.pod`, `service.instance`. An entity type is one or more +/// dot-separated segments, each a non-empty `[a-z0-9_]+` token. The dotted form +/// carries the two-entity-layer convention (`service` vs `service.instance`). +fn is_valid_entity_type(ty: &str) -> bool { + !ty.is_empty() + && ty.split('.').all(|seg| { + !seg.is_empty() + && seg + .bytes() + .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'_') + }) +} + +/// Parses a well-formed entity-identity option key of the shape +/// `greptime.semantic.entity..{id|descriptive|scope}` into +/// `(entity_type, role)`. The `` may itself contain dots; the role is the +/// final dot-separated segment. This is the single parser of the key format — +/// DDL validation and the read-time derivation both go through it. +pub fn parse_entity_option_key(key: &str) -> Option<(&str, EntityRole)> { + let rest = key.strip_prefix(SEMANTIC_ENTITY_PREFIX)?; + let (ty, role) = rest.rsplit_once('.')?; + if !is_valid_entity_type(ty) { + return None; + } + Some((ty, EntityRole::parse(role)?)) +} + +/// Returns true if `key` is a well-formed entity-identity option key. +pub fn is_entity_option_key(key: &str) -> bool { + parse_entity_option_key(key).is_some() +} + +/// Tokenizes an entity option's comma-separated column list (trimmed, empty +/// tokens dropped). [`validate_semantic_option`] rejects empty tokens at DDL +/// time, so readers only ever drop what validation already refused. +pub fn parse_entity_columns(value: &str) -> Vec { + value + .split(',') + .map(|c| c.trim().to_string()) + .filter(|c| !c.is_empty()) + .collect() +} + +/// Returns true if `key` is a recognised semantic table-option key. /// -/// Note this is membership, not a prefix test: unknown keys under -/// [`SEMANTIC_PREFIX`] are rejected, and the internal -/// [`SEMANTIC_PER_TABLE_INDEX_KEY`] (outside the prefix) never matches. +/// Two acceptance rules: membership in the closed [`SEMANTIC_OPTION_KEYS`] +/// whitelist, OR the open entity sub-namespace ([`is_entity_option_key`], validated +/// by prefix + shape). Everything else under [`SEMANTIC_PREFIX`] is rejected, and +/// the internal [`SEMANTIC_PER_TABLE_INDEX_KEY`] (outside the prefix) never matches. pub fn is_semantic_option_key(key: &str) -> bool { - SEMANTIC_OPTION_KEYS.contains(&key) + SEMANTIC_OPTION_KEYS.contains(&key) || is_entity_option_key(key) } /// Validates a `greptime.semantic.*` option's `value` against its allowed domain. @@ -129,9 +208,15 @@ pub fn is_semantic_option_key(key: &str) -> bool { /// Open-value keys (unit, original_name, pipeline, conventions) accept any /// non-empty string. Closed-domain keys accept a fixed set, plus the `unknown` /// sentinel, plus `mixed` for the keys where one long-lived table can -/// legitimately see multiple values. Keys not in [`SEMANTIC_OPTION_KEYS`] are -/// rejected. +/// legitimately see multiple values. Entity keys ([`is_entity_option_key`]) take a +/// comma-separated column-name list (each token non-empty); the "id columns must +/// be tag columns" invariant is enforced later against the table schema at DDL +/// time, not here. Keys that are neither whitelisted nor a well-formed entity key +/// are rejected. pub fn validate_semantic_option(key: &str, value: &str) -> bool { + if is_entity_option_key(key) { + return !value.is_empty() && value.split(',').all(|col| !col.trim().is_empty()); + } match key { SEMANTIC_PIPELINE | SEMANTIC_SOURCE_VERSION @@ -273,4 +358,67 @@ mod tests { ); } } + + #[test] + fn test_entity_option_key() { + // Well-formed entity keys are accepted by the open sub-namespace, and are + // therefore recognised as semantic option keys. + for key in [ + "greptime.semantic.entity.service.id", + "greptime.semantic.entity.k8s.pod.id", + "greptime.semantic.entity.service.instance.descriptive", + "greptime.semantic.entity.host.scope", + "greptime.semantic.entity.agent.id", + ] { + assert!(is_entity_option_key(key), "should accept {key}"); + assert!(is_semantic_option_key(key), "should recognise {key}"); + } + + // Malformed: empty type segment, bogus role, missing role, bare prefix, + // invalid (uppercase) charset in the type. + for key in [ + "greptime.semantic.entity..id", + "greptime.semantic.entity.service.bogusrole", + "greptime.semantic.entity.service", + "greptime.semantic.entity.", + "greptime.semantic.entity.Service.id", + ] { + assert!(!is_entity_option_key(key), "should reject {key}"); + assert!(!is_semantic_option_key(key), "should not recognise {key}"); + } + + // A non-entity semantic key is not an entity key. + assert!(!is_entity_option_key(SEMANTIC_SIGNAL_TYPE)); + + // Drift guard: the auto-stamped constant is a well-formed entity key. + assert!(is_entity_option_key(SEMANTIC_ENTITY_SERVICE_ID)); + assert!(is_semantic_option_key(SEMANTIC_ENTITY_SERVICE_ID)); + } + + #[test] + fn test_validate_entity_option() { + // Single- and composite-id column lists validate. + assert!(validate_semantic_option( + "greptime.semantic.entity.service.id", + "service_name" + )); + assert!(validate_semantic_option( + "greptime.semantic.entity.process.id", + "pid,start_time" + )); + // Empty value and blank/empty tokens do not. + assert!(!validate_semantic_option( + "greptime.semantic.entity.service.id", + "" + )); + assert!(!validate_semantic_option( + "greptime.semantic.entity.process.id", + "pid," + )); + // A malformed entity key validates to false regardless of value. + assert!(!validate_semantic_option( + "greptime.semantic.entity.service.bogusrole", + "service_name" + )); + } } diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 330f6b83bca..d7912d0a975 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -6656,7 +6656,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; - let expected_ddl = r#"[["mytable","CREATE TABLE IF NOT EXISTS \"mytable\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '1',\n trace_id >= '1' AND trace_id < '2',\n trace_id >= '2' AND trace_id < '3',\n trace_id >= '3' AND trace_id < '4',\n trace_id >= '4' AND trace_id < '5',\n trace_id >= '5' AND trace_id < '6',\n trace_id >= '6' AND trace_id < '7',\n trace_id >= '7' AND trace_id < '8',\n trace_id >= '8' AND trace_id < '9',\n trace_id >= '9' AND trace_id < 'a',\n trace_id >= 'a' AND trace_id < 'b',\n trace_id >= 'b' AND trace_id < 'c',\n trace_id >= 'c' AND trace_id < 'd',\n trace_id >= 'd' AND trace_id < 'e',\n trace_id >= 'e' AND trace_id < 'f',\n trace_id >= 'f'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["mytable","CREATE TABLE IF NOT EXISTS \"mytable\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '1',\n trace_id >= '1' AND trace_id < '2',\n trace_id >= '2' AND trace_id < '3',\n trace_id >= '3' AND trace_id < '4',\n trace_id >= '4' AND trace_id < '5',\n trace_id >= '5' AND trace_id < '6',\n trace_id >= '6' AND trace_id < '7',\n trace_id >= '7' AND trace_id < '8',\n trace_id >= '8' AND trace_id < '9',\n trace_id >= '9' AND trace_id < 'a',\n trace_id >= 'a' AND trace_id < 'b',\n trace_id >= 'b' AND trace_id < 'c',\n trace_id >= 'c' AND trace_id < 'd',\n trace_id >= 'd' AND trace_id < 'e',\n trace_id >= 'e' AND trace_id < 'f',\n trace_id >= 'f'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -6709,7 +6709,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part1","CREATE TABLE IF NOT EXISTS \"trace_table_part1\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\n\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part1","CREATE TABLE IF NOT EXISTS \"trace_table_part1\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\n\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -6746,7 +6746,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part4","CREATE TABLE IF NOT EXISTS \"trace_table_part4\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '4',\n trace_id >= '4' AND trace_id < '8',\n trace_id >= '8' AND trace_id < 'c',\n trace_id >= 'c'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part4","CREATE TABLE IF NOT EXISTS \"trace_table_part4\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '4',\n trace_id >= '4' AND trace_id < '8',\n trace_id >= '8' AND trace_id < 'c',\n trace_id >= 'c'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -6783,7 +6783,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part32","CREATE TABLE IF NOT EXISTS \"trace_table_part32\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '08',\n trace_id >= '08' AND trace_id < '10',\n trace_id >= '10' AND trace_id < '18',\n trace_id >= '18' AND trace_id < '20',\n trace_id >= '20' AND trace_id < '28',\n trace_id >= '28' AND trace_id < '30',\n trace_id >= '30' AND trace_id < '38',\n trace_id >= '38' AND trace_id < '40',\n trace_id >= '40' AND trace_id < '48',\n trace_id >= '48' AND trace_id < '50',\n trace_id >= '50' AND trace_id < '58',\n trace_id >= '58' AND trace_id < '60',\n trace_id >= '60' AND trace_id < '68',\n trace_id >= '68' AND trace_id < '70',\n trace_id >= '70' AND trace_id < '78',\n trace_id >= '78' AND trace_id < '80',\n trace_id >= '80' AND trace_id < '88',\n trace_id >= '88' AND trace_id < '90',\n trace_id >= '90' AND trace_id < '98',\n trace_id >= '98' AND trace_id < 'a0',\n trace_id >= 'a0' AND trace_id < 'a8',\n trace_id >= 'a8' AND trace_id < 'b0',\n trace_id >= 'b0' AND trace_id < 'b8',\n trace_id >= 'b8' AND trace_id < 'c0',\n trace_id >= 'c0' AND trace_id < 'c8',\n trace_id >= 'c8' AND trace_id < 'd0',\n trace_id >= 'd0' AND trace_id < 'd8',\n trace_id >= 'd8' AND trace_id < 'e0',\n trace_id >= 'e0' AND trace_id < 'e8',\n trace_id >= 'e8' AND trace_id < 'f0',\n trace_id >= 'f0' AND trace_id < 'f8',\n trace_id >= 'f8'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part32","CREATE TABLE IF NOT EXISTS \"trace_table_part32\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '08',\n trace_id >= '08' AND trace_id < '10',\n trace_id >= '10' AND trace_id < '18',\n trace_id >= '18' AND trace_id < '20',\n trace_id >= '20' AND trace_id < '28',\n trace_id >= '28' AND trace_id < '30',\n trace_id >= '30' AND trace_id < '38',\n trace_id >= '38' AND trace_id < '40',\n trace_id >= '40' AND trace_id < '48',\n trace_id >= '48' AND trace_id < '50',\n trace_id >= '50' AND trace_id < '58',\n trace_id >= '58' AND trace_id < '60',\n trace_id >= '60' AND trace_id < '68',\n trace_id >= '68' AND trace_id < '70',\n trace_id >= '70' AND trace_id < '78',\n trace_id >= '78' AND trace_id < '80',\n trace_id >= '80' AND trace_id < '88',\n trace_id >= '88' AND trace_id < '90',\n trace_id >= '90' AND trace_id < '98',\n trace_id >= '98' AND trace_id < 'a0',\n trace_id >= 'a0' AND trace_id < 'a8',\n trace_id >= 'a8' AND trace_id < 'b0',\n trace_id >= 'b0' AND trace_id < 'b8',\n trace_id >= 'b8' AND trace_id < 'c0',\n trace_id >= 'c0' AND trace_id < 'c8',\n trace_id >= 'c8' AND trace_id < 'd0',\n trace_id >= 'd0' AND trace_id < 'd8',\n trace_id >= 'd8' AND trace_id < 'e0',\n trace_id >= 'e0' AND trace_id < 'e8',\n trace_id >= 'e8' AND trace_id < 'f0',\n trace_id >= 'f0' AND trace_id < 'f8',\n trace_id >= 'f8'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -9164,7 +9164,7 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { .await; assert_eq!(StatusCode::OK, res.status()); - let trace_table_sql = "[[\"mytable\",\"CREATE TABLE IF NOT EXISTS \\\"mytable\\\" (\\n \\\"timestamp\\\" TIMESTAMP(9) NOT NULL,\\n \\\"timestamp_end\\\" TIMESTAMP(9) NULL,\\n \\\"duration_nano\\\" BIGINT UNSIGNED NULL,\\n \\\"parent_span_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"trace_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_id\\\" STRING NULL,\\n \\\"span_kind\\\" STRING NULL,\\n \\\"span_name\\\" STRING NULL,\\n \\\"span_status_code\\\" STRING NULL,\\n \\\"span_status_message\\\" STRING NULL,\\n \\\"trace_state\\\" STRING NULL,\\n \\\"scope_name\\\" STRING NULL,\\n \\\"scope_version\\\" STRING NULL,\\n \\\"service_name\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_attributes.operation.type\\\" STRING NULL,\\n \\\"span_attributes.net.peer.ip\\\" STRING NULL,\\n \\\"span_attributes.peer.service\\\" STRING NULL,\\n \\\"span_events\\\" JSON NULL,\\n \\\"span_links\\\" JSON NULL,\\n TIME INDEX (\\\"timestamp\\\"),\\n PRIMARY KEY (\\\"service_name\\\")\\n)\\nPARTITION ON COLUMNS (\\\"trace_id\\\") (\\n trace_id < '1',\\n trace_id >= '1' AND trace_id < '2',\\n trace_id >= '2' AND trace_id < '3',\\n trace_id >= '3' AND trace_id < '4',\\n trace_id >= '4' AND trace_id < '5',\\n trace_id >= '5' AND trace_id < '6',\\n trace_id >= '6' AND trace_id < '7',\\n trace_id >= '7' AND trace_id < '8',\\n trace_id >= '8' AND trace_id < '9',\\n trace_id >= '9' AND trace_id < 'a',\\n trace_id >= 'a' AND trace_id < 'b',\\n trace_id >= 'b' AND trace_id < 'c',\\n trace_id >= 'c' AND trace_id < 'd',\\n trace_id >= 'd' AND trace_id < 'e',\\n trace_id >= 'e' AND trace_id < 'f',\\n trace_id >= 'f'\\n)\\nENGINE=mito\\nWITH(\\n 'comment' = 'Created on insertion',\\n append_mode = 'true',\\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\\n 'greptime.semantic.signal_type' = 'trace',\\n 'greptime.semantic.source' = 'opentelemetry',\\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\\n table_data_model = 'greptime_trace_v1',\\n ttl = '7days'\\n)\"]]"; + let trace_table_sql = "[[\"mytable\",\"CREATE TABLE IF NOT EXISTS \\\"mytable\\\" (\\n \\\"timestamp\\\" TIMESTAMP(9) NOT NULL,\\n \\\"timestamp_end\\\" TIMESTAMP(9) NULL,\\n \\\"duration_nano\\\" BIGINT UNSIGNED NULL,\\n \\\"parent_span_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"trace_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_id\\\" STRING NULL,\\n \\\"span_kind\\\" STRING NULL,\\n \\\"span_name\\\" STRING NULL,\\n \\\"span_status_code\\\" STRING NULL,\\n \\\"span_status_message\\\" STRING NULL,\\n \\\"trace_state\\\" STRING NULL,\\n \\\"scope_name\\\" STRING NULL,\\n \\\"scope_version\\\" STRING NULL,\\n \\\"service_name\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_attributes.operation.type\\\" STRING NULL,\\n \\\"span_attributes.net.peer.ip\\\" STRING NULL,\\n \\\"span_attributes.peer.service\\\" STRING NULL,\\n \\\"span_events\\\" JSON NULL,\\n \\\"span_links\\\" JSON NULL,\\n TIME INDEX (\\\"timestamp\\\"),\\n PRIMARY KEY (\\\"service_name\\\")\\n)\\nPARTITION ON COLUMNS (\\\"trace_id\\\") (\\n trace_id < '1',\\n trace_id >= '1' AND trace_id < '2',\\n trace_id >= '2' AND trace_id < '3',\\n trace_id >= '3' AND trace_id < '4',\\n trace_id >= '4' AND trace_id < '5',\\n trace_id >= '5' AND trace_id < '6',\\n trace_id >= '6' AND trace_id < '7',\\n trace_id >= '7' AND trace_id < '8',\\n trace_id >= '8' AND trace_id < '9',\\n trace_id >= '9' AND trace_id < 'a',\\n trace_id >= 'a' AND trace_id < 'b',\\n trace_id >= 'b' AND trace_id < 'c',\\n trace_id >= 'c' AND trace_id < 'd',\\n trace_id >= 'd' AND trace_id < 'e',\\n trace_id >= 'e' AND trace_id < 'f',\\n trace_id >= 'f'\\n)\\nENGINE=mito\\nWITH(\\n 'comment' = 'Created on insertion',\\n append_mode = 'true',\\n 'greptime.semantic.entity.service.id' = 'service_name',\\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\\n 'greptime.semantic.signal_type' = 'trace',\\n 'greptime.semantic.source' = 'opentelemetry',\\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\\n table_data_model = 'greptime_trace_v1',\\n ttl = '7days'\\n)\"]]"; validate_data( "trace_v1_create_table", &client, diff --git a/tests/cases/standalone/common/information_schema/table_constraints.result b/tests/cases/standalone/common/information_schema/table_constraints.result index 57333645877..3fc0bec64e4 100644 --- a/tests/cases/standalone/common/information_schema/table_constraints.result +++ b/tests/cases/standalone/common/information_schema/table_constraints.result @@ -17,7 +17,10 @@ DESC TABLE TABLE_CONSTRAINTS; | enforced | String | | NO | | FIELD | +--------------------+--------+-----+------+---------+---------------+ -SELECT * FROM TABLE_CONSTRAINTS ORDER BY TABLE_NAME, CONSTRAINT_NAME; +SELECT * +FROM TABLE_CONSTRAINTS +WHERE TABLE_SCHEMA != 'greptime_private' +ORDER BY TABLE_NAME, CONSTRAINT_NAME; +--------------------+-------------------+-----------------+--------------+------------+-----------------+----------+ | constraint_catalog | constraint_schema | constraint_name | table_schema | table_name | constraint_type | enforced | @@ -37,7 +40,10 @@ use INFORMATION_SCHEMA; Affected Rows: 0 -SELECT * FROM TABLE_CONSTRAINTS ORDER BY TABLE_NAME, CONSTRAINT_NAME; +SELECT * +FROM TABLE_CONSTRAINTS +WHERE TABLE_SCHEMA != 'greptime_private' +ORDER BY TABLE_NAME, CONSTRAINT_NAME; +--------------------+-------------------+-----------------+--------------+------------+-----------------+----------+ | constraint_catalog | constraint_schema | constraint_name | table_schema | table_name | constraint_type | enforced | diff --git a/tests/cases/standalone/common/information_schema/table_constraints.sql b/tests/cases/standalone/common/information_schema/table_constraints.sql index f5e0542696d..557826c8d56 100644 --- a/tests/cases/standalone/common/information_schema/table_constraints.sql +++ b/tests/cases/standalone/common/information_schema/table_constraints.sql @@ -4,7 +4,10 @@ USE INFORMATION_SCHEMA; DESC TABLE TABLE_CONSTRAINTS; -SELECT * FROM TABLE_CONSTRAINTS ORDER BY TABLE_NAME, CONSTRAINT_NAME; +SELECT * +FROM TABLE_CONSTRAINTS +WHERE TABLE_SCHEMA != 'greptime_private' +ORDER BY TABLE_NAME, CONSTRAINT_NAME; use public; @@ -12,7 +15,10 @@ CREATE TABLE test(i double, j string, ts timestamp time index, primary key(j)); use INFORMATION_SCHEMA; -SELECT * FROM TABLE_CONSTRAINTS ORDER BY TABLE_NAME, CONSTRAINT_NAME; +SELECT * +FROM TABLE_CONSTRAINTS +WHERE TABLE_SCHEMA != 'greptime_private' +ORDER BY TABLE_NAME, CONSTRAINT_NAME; SELECT * FROM TABLE_CONSTRAINTS WHERE TABLE_NAME = 'test' ORDER BY TABLE_NAME, CONSTRAINT_NAME; diff --git a/tests/cases/standalone/common/information_schema/tables.result b/tests/cases/standalone/common/information_schema/tables.result index 93a93a9c980..7c9e7c84d5a 100644 --- a/tests/cases/standalone/common/information_schema/tables.result +++ b/tests/cases/standalone/common/information_schema/tables.result @@ -22,7 +22,7 @@ create table t (ts timestamp time index); Affected Rows: 0 -select table_catalog, table_schema, table_name from information_schema.tables where table_schema != 'information_schema'; +select table_catalog, table_schema, table_name from information_schema.tables where table_schema not in ('information_schema', 'greptime_private'); +---------------+--------------+------------+ | table_catalog | table_schema | table_name | diff --git a/tests/cases/standalone/common/information_schema/tables.sql b/tests/cases/standalone/common/information_schema/tables.sql index ae2bcf04306..1316170c236 100644 --- a/tests/cases/standalone/common/information_schema/tables.sql +++ b/tests/cases/standalone/common/information_schema/tables.sql @@ -10,7 +10,7 @@ use abcde; create table t (ts timestamp time index); -select table_catalog, table_schema, table_name from information_schema.tables where table_schema != 'information_schema'; +select table_catalog, table_schema, table_name from information_schema.tables where table_schema not in ('information_schema', 'greptime_private'); use public; diff --git a/tests/cases/standalone/common/system/information_schema.result b/tests/cases/standalone/common/system/information_schema.result index a6c44229a25..e05cd064e1d 100644 --- a/tests/cases/standalone/common/system/information_schema.result +++ b/tests/cases/standalone/common/system/information_schema.result @@ -9,6 +9,7 @@ Error: 1004(InvalidArguments), Schema information_schema already exists select * from information_schema.tables where table_name != 'scripts' + and table_schema != 'greptime_private' order by table_schema, table_name; +++++++++++++++++++++++++ @@ -56,7 +57,10 @@ order by table_schema, table_name; |greptime|public|numbers|LOCALTEMPORARY|2|0|0|0|0|0|test_engine|11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|||Y| +++++++++++++++++++++++++ -select * from information_schema.columns order by table_schema, table_name, column_name; +select * +from information_schema.columns +where table_schema != 'greptime_private' +order by table_schema, table_name, column_name; +---------------+--------------------+---------------------------------------+-----------------------------------+------------------+--------------------------+------------------------+-------------------+---------------+--------------------+--------------------+----------------+------------+-------+---------------+-----------------------+----------------------+---------------------+---------------+----------------+-------------+---------------------+----------------+--------+ | table_catalog | table_schema | table_name | column_name | ordinal_position | character_maximum_length | character_octet_length | numeric_precision | numeric_scale | datetime_precision | character_set_name | collation_name | column_key | extra | privileges | generation_expression | greptime_data_type | data_type | semantic_type | column_default | is_nullable | column_type | column_comment | srs_id | @@ -623,7 +627,7 @@ order by table_name; select table_name from information_schema.tables -where table_schema not in ('my_db', 'information_schema', 'pg_catalog') +where table_schema not in ('my_db', 'information_schema', 'pg_catalog', 'greptime_private') order by table_name; +------------+ @@ -638,6 +642,7 @@ where table_catalog = 'greptime' and table_schema != 'public' and table_schema != 'information_schema' and table_schema != 'pg_catalog' + and table_schema != 'greptime_private' order by table_schema, table_name; +---------------+--------------+------------+------------+--------+ @@ -652,6 +657,7 @@ where table_catalog = 'greptime' and table_schema != 'public' and table_schema != 'information_schema' and table_schema != 'pg_catalog' + and table_schema != 'greptime_private' order by table_schema, table_name, column_name; +---------------+--------------+------------+-------------+--------------+---------------+ @@ -689,12 +695,12 @@ use information_schema; Affected Rows: 0 -- test query filter for key_column_usage -- -select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME = 'TIME INDEX'; +select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME = 'TIME INDEX' and TABLE_SCHEMA != 'greptime_private'; ++ ++ -select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME != 'TIME INDEX'; +select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME != 'TIME INDEX' and TABLE_SCHEMA != 'greptime_private'; +--------------------+-------------------+-----------------+---------------+--------------------+--------------+------------+-------------+------------------+-------------------------------+-------------------------+-----------------------+------------------------+---------------------+ | constraint_catalog | constraint_schema | constraint_name | table_catalog | real_table_catalog | table_schema | table_name | column_name | ordinal_position | position_in_unique_constraint | referenced_table_schema | referenced_table_name | referenced_column_name | greptime_index_type | @@ -702,12 +708,12 @@ select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME != 'TIME INDEX'; | def | public | PRIMARY | def | greptime | public | numbers | number | 1 | | | | | dense | +--------------------+-------------------+-----------------+---------------+--------------------+--------------+------------+-------------+------------------+-------------------------------+-------------------------+-----------------------+------------------------+---------------------+ -select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME LIKE '%INDEX'; +select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME LIKE '%INDEX' and TABLE_SCHEMA != 'greptime_private'; ++ ++ -select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME NOT LIKE '%INDEX'; +select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME NOT LIKE '%INDEX' and TABLE_SCHEMA != 'greptime_private'; +--------------------+-------------------+-----------------+---------------+--------------------+--------------+------------+-------------+------------------+-------------------------------+-------------------------+-----------------------+------------------------+---------------------+ | constraint_catalog | constraint_schema | constraint_name | table_catalog | real_table_catalog | table_schema | table_name | column_name | ordinal_position | position_in_unique_constraint | referenced_table_schema | referenced_table_name | referenced_column_name | greptime_index_type | @@ -715,7 +721,7 @@ select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME NOT LIKE '%INDEX'; | def | public | PRIMARY | def | greptime | public | numbers | number | 1 | | | | | dense | +--------------------+-------------------+-----------------+---------------+--------------------+--------------+------------+-------------+------------------+-------------------------------+-------------------------+-----------------------+------------------------+---------------------+ -select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME == 'TIME INDEX' AND CONSTRAINT_SCHEMA != 'my_db'; +select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME == 'TIME INDEX' AND CONSTRAINT_SCHEMA != 'my_db' and TABLE_SCHEMA != 'greptime_private'; ++ ++ @@ -794,7 +800,7 @@ desc table key_column_usage; | greptime_index_type | String | | YES | | FIELD | +-------------------------------+--------+-----+------+---------+---------------+ -select * from key_column_usage; +select * from key_column_usage where table_schema != 'greptime_private'; +--------------------+-------------------+-----------------+---------------+--------------------+--------------+------------+-------------+------------------+-------------------------------+-------------------------+-----------------------+------------------------+---------------------+ | constraint_catalog | constraint_schema | constraint_name | table_catalog | real_table_catalog | table_schema | table_name | column_name | ordinal_position | position_in_unique_constraint | referenced_table_schema | referenced_table_name | referenced_column_name | greptime_index_type | diff --git a/tests/cases/standalone/common/system/information_schema.sql b/tests/cases/standalone/common/system/information_schema.sql index 992ed8cd14c..0b74978269a 100644 --- a/tests/cases/standalone/common/system/information_schema.sql +++ b/tests/cases/standalone/common/system/information_schema.sql @@ -7,9 +7,13 @@ create database information_schema; select * from information_schema.tables where table_name != 'scripts' + and table_schema != 'greptime_private' order by table_schema, table_name; -select * from information_schema.columns order by table_schema, table_name, column_name; +select * +from information_schema.columns +where table_schema != 'greptime_private' +order by table_schema, table_name, column_name; select table_schema, table_name from information_schema.tables order by table_name limit 5; @@ -49,7 +53,7 @@ order by table_name; select table_name from information_schema.tables -where table_schema not in ('my_db', 'information_schema', 'pg_catalog') +where table_schema not in ('my_db', 'information_schema', 'pg_catalog', 'greptime_private') order by table_name; select table_catalog, table_schema, table_name, table_type, engine @@ -58,6 +62,7 @@ where table_catalog = 'greptime' and table_schema != 'public' and table_schema != 'information_schema' and table_schema != 'pg_catalog' + and table_schema != 'greptime_private' order by table_schema, table_name; select table_catalog, table_schema, table_name, column_name, data_type, semantic_type @@ -66,6 +71,7 @@ where table_catalog = 'greptime' and table_schema != 'public' and table_schema != 'information_schema' and table_schema != 'pg_catalog' + and table_schema != 'greptime_private' order by table_schema, table_name, column_name; -- test query filter for columns -- @@ -84,15 +90,15 @@ drop schema my_db; use information_schema; -- test query filter for key_column_usage -- -select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME = 'TIME INDEX'; +select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME = 'TIME INDEX' and TABLE_SCHEMA != 'greptime_private'; -select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME != 'TIME INDEX'; +select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME != 'TIME INDEX' and TABLE_SCHEMA != 'greptime_private'; -select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME LIKE '%INDEX'; +select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME LIKE '%INDEX' and TABLE_SCHEMA != 'greptime_private'; -select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME NOT LIKE '%INDEX'; +select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME NOT LIKE '%INDEX' and TABLE_SCHEMA != 'greptime_private'; -select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME == 'TIME INDEX' AND CONSTRAINT_SCHEMA != 'my_db'; +select * from KEY_COLUMN_USAGE where CONSTRAINT_NAME == 'TIME INDEX' AND CONSTRAINT_SCHEMA != 'my_db' and TABLE_SCHEMA != 'greptime_private'; -- schemata -- @@ -109,7 +115,7 @@ select count(*) from build_info; desc table key_column_usage; -select * from key_column_usage; +select * from key_column_usage where table_schema != 'greptime_private'; -- tables not implemented DESC TABLE COLUMN_PRIVILEGES; diff --git a/tests/cases/standalone/common/system/pg_catalog.result b/tests/cases/standalone/common/system/pg_catalog.result index 7a6b65c7c89..a36f0fe606b 100644 --- a/tests/cases/standalone/common/system/pg_catalog.result +++ b/tests/cases/standalone/common/system/pg_catalog.result @@ -727,6 +727,7 @@ WHERE c.relkind IN ('r','p','v','m','S','f','') AND n.nspname <> 'pg_catalog' AND n.nspname !~ '^pg_toast' AND n.nspname <> 'information_schema' + AND n.nspname <> 'greptime_private' AND pg_catalog.pg_table_is_visible(c.oid) ORDER BY 1,2; @@ -748,6 +749,7 @@ WHERE c.relkind IN ('r','p','') AND n.nspname <> 'pg_catalog' AND n.nspname !~ '^pg_toast' AND n.nspname <> 'information_schema' + AND n.nspname <> 'greptime_private' AND pg_catalog.pg_table_is_visible(c.oid) ORDER BY 1,2; @@ -816,6 +818,7 @@ WHERE c.relkind IN ('r','p','') AND n.nspname <> 'pg_catalog' AND n.nspname !~ '^pg_toast' AND n.nspname <> 'information_schema' + AND n.nspname <> 'greptime_private' AND pg_catalog.pg_table_is_visible(c.oid) ORDER BY 1,2; @@ -878,6 +881,7 @@ SELECT FROM information_schema.tables WHERE quote_ident(table_schema) NOT IN ('information_schema', 'pg_catalog', + 'greptime_private', '_timescaledb_cache', '_timescaledb_catalog', '_timescaledb_internal', @@ -940,6 +944,7 @@ where relnamespace in ( where nspname <> 'public' and nspname <> 'information_schema' and nspname <> 'pg_catalog' + and nspname <> 'greptime_private' ) order by relnamespace, relname; diff --git a/tests/cases/standalone/common/system/pg_catalog.sql b/tests/cases/standalone/common/system/pg_catalog.sql index 2e84ecd7ce3..1dd97c9dc38 100644 --- a/tests/cases/standalone/common/system/pg_catalog.sql +++ b/tests/cases/standalone/common/system/pg_catalog.sql @@ -53,6 +53,7 @@ WHERE c.relkind IN ('r','p','v','m','S','f','') AND n.nspname <> 'pg_catalog' AND n.nspname !~ '^pg_toast' AND n.nspname <> 'information_schema' + AND n.nspname <> 'greptime_private' AND pg_catalog.pg_table_is_visible(c.oid) ORDER BY 1,2; @@ -68,6 +69,7 @@ WHERE c.relkind IN ('r','p','') AND n.nspname <> 'pg_catalog' AND n.nspname !~ '^pg_toast' AND n.nspname <> 'information_schema' + AND n.nspname <> 'greptime_private' AND pg_catalog.pg_table_is_visible(c.oid) ORDER BY 1,2; @@ -110,6 +112,7 @@ WHERE c.relkind IN ('r','p','') AND n.nspname <> 'pg_catalog' AND n.nspname !~ '^pg_toast' AND n.nspname <> 'information_schema' + AND n.nspname <> 'greptime_private' AND pg_catalog.pg_table_is_visible(c.oid) ORDER BY 1,2; @@ -152,6 +155,7 @@ SELECT FROM information_schema.tables WHERE quote_ident(table_schema) NOT IN ('information_schema', 'pg_catalog', + 'greptime_private', '_timescaledb_cache', '_timescaledb_catalog', '_timescaledb_internal', @@ -201,6 +205,7 @@ where relnamespace in ( where nspname <> 'public' and nspname <> 'information_schema' and nspname <> 'pg_catalog' + and nspname <> 'greptime_private' ) order by relnamespace, relname; diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result new file mode 100644 index 00000000000..bda28755f1e --- /dev/null +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -0,0 +1,155 @@ +-- The computed entity-graph tables under greptime_private are read-only +-- virtual tables: readable, but rejecting every DDL/DML path. +select observed_at, entity_type, entity_id, scope from greptime_private.semantic_entities; + +++ +++ + +select observed_at, src_id, dst_id, rel_type from greptime_private.semantic_relationships; + +++ +++ + +insert into greptime_private.semantic_entities (observed_at, entity_type, entity_id) values (now(), 'service', 'svc-a'); + +Error: 1001(Unsupported), Cannot change read-only table: semantic_entities + +-- Plain-literal VALUES takes the direct region-insert path, not the query +-- engine; the table-level guard must hold there too. +insert into greptime_private.semantic_entities (observed_at, entity_type, entity_id) values (0, 'service', 'svc-b'); + +Error: 1001(Unsupported), Table `semantic_entities` is read-only + +create table greptime_private.semantic_entities (ts timestamp time index); + +Error: 1001(Unsupported), Table `semantic_entities` is read-only + +create table greptime_private.semantic_relationships (ts timestamp time index); + +Error: 1001(Unsupported), Table `semantic_relationships` is read-only + +alter table greptime_private.semantic_entities add column extra string; + +Error: 1001(Unsupported), Table `semantic_entities` is read-only + +truncate table greptime_private.semantic_entities; + +Error: 1001(Unsupported), Table `semantic_entities` is read-only + +drop table greptime_private.semantic_entities; + +Error: 1001(Unsupported), Table `semantic_entities` is read-only + +drop table greptime_private.semantic_relationships; + +Error: 1001(Unsupported), Table `semantic_relationships` is read-only + +-- Read-time derivation: declare entity identities on a user table and observe +-- them in the computed registry (single-column and composite ids, scope and +-- descriptive columns, JSON output columns). +create table graph_app_latency ( + ts timestamp time index, + service_name string, + host string, + env string, + latency double, + primary key (service_name, host, env) +) with ( + 'greptime.semantic.entity.service.id' = 'service_name', + 'greptime.semantic.entity.service.scope' = 'env', + 'greptime.semantic.entity.process.id' = 'service_name,host', + 'greptime.semantic.entity.process.descriptive' = 'env' +); + +Affected Rows: 0 + +insert into graph_app_latency values (now(), 'cart', 'h1', 'us-east', 3.5); + +Affected Rows: 1 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id; + ++-------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ +| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | ++-------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ +| process | host=h1,service_name=cart | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_latency"] | +| service | cart | | us-east | | ["public.graph_app_latency"] | ++-------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ + +drop table graph_app_latency; + +Affected Rows: 0 + +-- Renaming a physical table INTO a reserved computed-table name would let the +-- overlay shadow it, orphaning its data. +create table greptime_private.graph_rename_probe (ts timestamp time index); + +Affected Rows: 0 + +alter table greptime_private.graph_rename_probe rename semantic_entities; + +Error: 1001(Unsupported), Table `semantic_entities` is read-only + +drop table greptime_private.graph_rename_probe; + +Affected Rows: 0 + +-- Calls derivation: a trace-v1-model table with client/server span pairs. +-- Trace tables need no entity declaration: the service entities are implicit. +create table graph_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +-- The epoch-timestamped pair falls outside the default one-hour derivation +-- window and must not appear in the results. +insert into graph_traces values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), + (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000), + (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), + (now(), 't2', 's2', 'c2', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'cart', 500000000), + (0, 't0', 'c0', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'stale-src', 0), + (0, 't0', 's0', 'c0', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'stale-dst', 100); + +Affected Rows: 6 + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence, + request_count, error_count, duration_sum, duration_count, attributes +from greptime_private.semantic_relationships +order by src_id; + ++----------+----------+----------+--------+----------+------------+------------+---------------+-------------+--------------+----------------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | request_count | error_count | duration_sum | duration_count | attributes | ++----------+----------+----------+--------+----------+------------+------------+---------------+-------------+--------------+----------------+------------+ +| service | frontend | service | cart | calls | trace | 1 | 2 | 1 | 2 | 2 | | ++----------+----------+----------+--------+----------+------------+------------+---------------+-------------+--------------+----------------+------------+ + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, entity_id_attrs, scope, source_tables +from greptime_private.semantic_entities +order by entity_id; + ++-------------+-----------+-----------------+-------+-------------------------+ +| entity_type | entity_id | entity_id_attrs | scope | source_tables | ++-------------+-----------+-----------------+-------+-------------------------+ +| service | cart | | | ["public.graph_traces"] | +| service | frontend | | | ["public.graph_traces"] | ++-------------+-----------+-----------------+-------+-------------------------+ + +drop table graph_traces; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql new file mode 100644 index 00000000000..000401e1666 --- /dev/null +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -0,0 +1,94 @@ +-- The computed entity-graph tables under greptime_private are read-only +-- virtual tables: readable, but rejecting every DDL/DML path. +select observed_at, entity_type, entity_id, scope from greptime_private.semantic_entities; + +select observed_at, src_id, dst_id, rel_type from greptime_private.semantic_relationships; + +insert into greptime_private.semantic_entities (observed_at, entity_type, entity_id) values (now(), 'service', 'svc-a'); + +-- Plain-literal VALUES takes the direct region-insert path, not the query +-- engine; the table-level guard must hold there too. +insert into greptime_private.semantic_entities (observed_at, entity_type, entity_id) values (0, 'service', 'svc-b'); + +create table greptime_private.semantic_entities (ts timestamp time index); + +create table greptime_private.semantic_relationships (ts timestamp time index); + +alter table greptime_private.semantic_entities add column extra string; + +truncate table greptime_private.semantic_entities; + +drop table greptime_private.semantic_entities; + +drop table greptime_private.semantic_relationships; + +-- Read-time derivation: declare entity identities on a user table and observe +-- them in the computed registry (single-column and composite ids, scope and +-- descriptive columns, JSON output columns). +create table graph_app_latency ( + ts timestamp time index, + service_name string, + host string, + env string, + latency double, + primary key (service_name, host, env) +) with ( + 'greptime.semantic.entity.service.id' = 'service_name', + 'greptime.semantic.entity.service.scope' = 'env', + 'greptime.semantic.entity.process.id' = 'service_name,host', + 'greptime.semantic.entity.process.descriptive' = 'env' +); + +insert into graph_app_latency values (now(), 'cart', 'h1', 'us-east', 3.5); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id; + +drop table graph_app_latency; + +-- Renaming a physical table INTO a reserved computed-table name would let the +-- overlay shadow it, orphaning its data. +create table greptime_private.graph_rename_probe (ts timestamp time index); + +alter table greptime_private.graph_rename_probe rename semantic_entities; + +drop table greptime_private.graph_rename_probe; + +-- Calls derivation: a trace-v1-model table with client/server span pairs. +-- Trace tables need no entity declaration: the service entities are implicit. +create table graph_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +-- The epoch-timestamped pair falls outside the default one-hour derivation +-- window and must not appear in the results. +insert into graph_traces values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), + (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000), + (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), + (now(), 't2', 's2', 'c2', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'cart', 500000000), + (0, 't0', 'c0', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'stale-src', 0), + (0, 't0', 's0', 'c0', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'stale-dst', 100); + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence, + request_count, error_count, duration_sum, duration_count, attributes +from greptime_private.semantic_relationships +order by src_id; + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, entity_id_attrs, scope, source_tables +from greptime_private.semantic_entities +order by entity_id; + +drop table graph_traces; diff --git a/tests/cases/standalone/common/view/create.result b/tests/cases/standalone/common/view/create.result index a4451bfed15..76b5ed78935 100644 --- a/tests/cases/standalone/common/view/create.result +++ b/tests/cases/standalone/common/view/create.result @@ -88,7 +88,10 @@ Error: 4001(TableNotFound), Failed to plan SQL: Table not found: greptime.pg_cat -- SQLNESS REPLACE (\s\d+\s) ID -- SQLNESS REPLACE (\s[\-0-9T:\.]{15,}) DATETIME -- SQLNESS REPLACE [\u0020\-]+ -SELECT * FROM INFORMATION_SCHEMA.TABLES ORDER BY TABLE_NAME, TABLE_TYPE; +SELECT * +FROM INFORMATION_SCHEMA.TABLES +WHERE TABLE_SCHEMA != 'greptime_private' +ORDER BY TABLE_NAME, TABLE_TYPE; +++++++++++++++++++++++++ |table_catalog|table_schema|table_name|table_type|table_id|data_length|max_data_length|index_length|max_index_length|avg_row_length|engine|version|row_format|table_rows|data_free|auto_increment|create_time|update_time|check_time|table_collation|checksum|create_options|table_comment|temporary| diff --git a/tests/cases/standalone/common/view/create.sql b/tests/cases/standalone/common/view/create.sql index 91149f44f46..c1f63c897fd 100644 --- a/tests/cases/standalone/common/view/create.sql +++ b/tests/cases/standalone/common/view/create.sql @@ -46,7 +46,10 @@ ORDER BY 1,2; -- SQLNESS REPLACE (\s\d+\s) ID -- SQLNESS REPLACE (\s[\-0-9T:\.]{15,}) DATETIME -- SQLNESS REPLACE [\u0020\-]+ -SELECT * FROM INFORMATION_SCHEMA.TABLES ORDER BY TABLE_NAME, TABLE_TYPE; +SELECT * +FROM INFORMATION_SCHEMA.TABLES +WHERE TABLE_SCHEMA != 'greptime_private' +ORDER BY TABLE_NAME, TABLE_TYPE; -- SQLNESS REPLACE (\s\d+\s) ID -- SQLNESS REPLACE (\s[\-0-9T:\.]{15,}) DATETIME diff --git a/tests/runner/src/formatter.rs b/tests/runner/src/formatter.rs index 0ae3b5ab207..99a7d1656b9 100644 --- a/tests/runner/src/formatter.rs +++ b/tests/runner/src/formatter.rs @@ -203,7 +203,7 @@ pub fn build_recordbatches_from_mysql_rows(rows: &[MySqlRow]) -> RecordBatches { names .iter() .map(|name| { - ColumnSchema::new(name.to_string(), ConcreteDataType::string_datatype(), false) + ColumnSchema::new(name.to_string(), ConcreteDataType::string_datatype(), true) }) .collect(), )); @@ -214,7 +214,9 @@ pub fn build_recordbatches_from_mysql_rows(rows: &[MySqlRow]) -> RecordBatches { .collect(); for row in rows.iter() { for (i, name) in names.iter().enumerate() { - columns[i].push(row.get::(name).as_deref()); + // `get::` panics on NULL values, so fetch as Option. + let value: Option = row.get::, &str>(name).flatten(); + columns[i].push(value.as_deref()); } } let columns: Vec = columns From 9fab11271bf1273c9d243661d7fc4a2e0f797688 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Fri, 7 Aug 2026 12:07:34 +0800 Subject: [PATCH 013/126] perf: reduce ingestion and flat-merge overhead (#8778) chore: minor performance fix Signed-off-by: shuiyisong --- src/mito2/src/read/flat_merge.rs | 34 +++++++++-- src/servers/src/prom_remote_write/v2.rs | 75 ++++++++++++++++--------- src/servers/src/row_writer.rs | 5 +- 3 files changed, 80 insertions(+), 34 deletions(-) diff --git a/src/mito2/src/read/flat_merge.rs b/src/mito2/src/read/flat_merge.rs index b1c304f2441..2562f364309 100644 --- a/src/mito2/src/read/flat_merge.rs +++ b/src/mito2/src/read/flat_merge.rs @@ -512,6 +512,11 @@ impl RowCursor { self.offset >= self.columns.num_rows() } + /// Returns whether advancing this cursor will finish the current batch. + fn is_last_row(&self) -> bool { + self.offset.checked_add(1) == Some(self.columns.num_rows()) + } + fn advance(&mut self) { self.offset += 1; } @@ -819,12 +824,15 @@ impl FlatMergeReader { } } - let start = Instant::now(); - if let Some(next) = hottest.advance_row().await? { + // Only read the clock when advancing will attempt to fetch the next batch. + // Check first because `None` means either no fetch or EOF after a fetch attempt. + let start = hottest.current_cursor().is_last_row().then(Instant::now); + let next = hottest.advance_row().await?; + if let Some(start) = start { self.metrics.fetch_cost += start.elapsed(); + } + if let Some(next) = next { self.in_progress.push_batch(hottest.node_index, next); - } else { - self.metrics.fetch_cost += start.elapsed(); } self.algo.reheap(hottest); @@ -1066,6 +1074,24 @@ mod tests { Box::new(batches.into_iter().map(Ok)) } + #[test] + fn test_row_cursor_last_row() { + let batch = create_test_record_batch( + &[b"k1", b"k1"], + &[1000, 2000], + &[21, 22], + &[OpType::Put, OpType::Put], + &[11, 12], + ); + let mut cursor = RowCursor::new(SortColumns::new(&batch)); + + assert!(!cursor.is_last_row()); + cursor.advance(); + assert!(cursor.is_last_row()); + cursor.advance(); + assert!(!cursor.is_last_row()); + } + /// Helper function to check if two record batches are equivalent. fn assert_record_batches_eq(expected: &[RecordBatch], actual: &[RecordBatch]) { for (exp, act) in expected.iter().zip(actual.iter()) { diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 923f860be38..229fac01666 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -90,8 +90,7 @@ pub(crate) fn into_write_requests(request: Request) -> Result>::new(); let mut histogram_tables = HashMap::>::new(); - let mut sample_metrics = HashSet::<(PromCtx, String)>::new(); - let mut histogram_metrics = HashSet::<(PromCtx, String)>::new(); + let mut label_names = HashSet::new(); let mut sample_count_total = 0; let mut histogram_count_total = 0; @@ -103,31 +102,27 @@ pub(crate) fn into_write_requests(request: Request) -> Result 0 { - ensure!( - !histogram_metrics.contains(&metric_key), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 metric `{table_name}` contains both samples and native histograms" - ), - } - ); - sample_metrics.insert(metric_key.clone()); - } - if histogram_count > 0 { - ensure!( - !sample_metrics.contains(&metric_key), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 metric `{table_name}` contains both samples and native histograms" - ), - } - ); - histogram_metrics.insert(metric_key); - } + let has_other_value_type = if sample_count > 0 { + histogram_count > 0 + || histogram_tables + .get(&prom_ctx) + .is_some_and(|tables| tables.contains_key(&table_name)) + } else { + sample_tables + .get(&prom_ctx) + .is_some_and(|tables| tables.contains_key(&table_name)) + }; + ensure!( + !has_other_value_type, + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 metric `{table_name}` contains both samples and native histograms" + ), + } + ); if sample_count > 0 && histogram_count == 0 { // Fast path for regular sample-only series. Move the resolved labels into @@ -745,7 +740,11 @@ fn ensure_no_internal_histogram_labels(tags: &PromTags) -> Result<()> { Ok(()) } -fn resolve_series_labels(symbols: &[String], series: &TimeSeries) -> Result { +fn resolve_series_labels<'a>( + symbols: &'a [String], + series: &TimeSeries, + label_names: &mut HashSet<&'a str>, +) -> Result { ensure!( series.labels_refs.len().is_multiple_of(2), error::InvalidPromRemoteRequestSnafu { @@ -756,7 +755,7 @@ fn resolve_series_labels(symbols: &[String], series: &TimeSeries) -> Result, rows: Vec, - column_indexes: HashMap, + column_indexes: AHashMap, } impl TableData { @@ -45,7 +46,7 @@ impl TableData { Self { schema: Vec::with_capacity(num_columns), rows: Vec::with_capacity(num_rows), - column_indexes: HashMap::with_capacity(num_columns), + column_indexes: AHashMap::with_capacity(num_columns), } } From fbbe5af76744eb93067a3bed099834656d7da4d6 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 7 Aug 2026 14:46:46 +0800 Subject: [PATCH 014/126] fix(query): avoid unsafe count wildcard rewrites (#8522) * fix(query): avoid unsafe count wildcard rewrites Signed-off-by: discord9 * fix(query): preserve outer count alias Signed-off-by: discord9 * fix(query): address review comments on count wildcard rewrite - Remove the has_projection check: the row count is correct regardless of whether a projection exists (per review). - Explain why checking the first input is equivalent to checking all inputs (a plan with zero inputs falls back to count(1)). - Rename qa_ prefixed tests to follow the module convention. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(tql): update tql-cte expectations for count wildcard rewrite The QP-026 count-wildcard fix rewrites count(*) -> count(time_index), so the EXPLAIN output for the filtered/final CTE aggregates names the time-index column. Aligns tql-cte.result with the actual output (CI failure). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/optimizer/count_wildcard.rs | 241 +++++++++++++++++- .../standalone/common/tql/tql-cte.result | 4 +- 2 files changed, 231 insertions(+), 14 deletions(-) diff --git a/src/query/src/optimizer/count_wildcard.rs b/src/query/src/optimizer/count_wildcard.rs index affdaa18f77..1c6a5b814b3 100644 --- a/src/query/src/optimizer/count_wildcard.rs +++ b/src/query/src/optimizer/count_wildcard.rs @@ -85,20 +85,23 @@ impl CountWildcardToTimeIndexRule { plan.visit(&mut finder).unwrap(); let col = finder.into_column(); - // check if the time index is a valid column as for current plan + // The resolved time index must be present and non-nullable in the + // immediate input schema. Schema-changing nodes can otherwise expose + // a nullable field with the same name as the source time index, and + // `count(
)` would then count fewer rows than `count(*)`. if let Some(col) = &col { - let mut is_valid = false; // if more than one input, we give up and just use `count(1)` if plan.inputs().len() > 1 { return None; } - for input in plan.inputs() { - if input.schema().has_column(col) { - is_valid = true; - break; - } - } - if !is_valid { + // The guard above guarantees exactly one input here, so checking + // the first input is equivalent to checking all inputs as the rule + // used to: a plan with zero inputs also falls back to `count(1)`. + let input = plan.inputs().first().copied()?; + let Ok((_, field)) = input.schema().qualified_field_from_column(col) else { + return None; + }; + if field.is_nullable() { return None; } } @@ -147,7 +150,8 @@ impl TreeNodeVisitor<'_> for TimeIndexFinder { fn f_down(&mut self, node: &Self::Node) -> DataFusionResult { if let LogicalPlan::SubqueryAlias(subquery_alias) = node { - self.table_alias = Some(subquery_alias.alias.clone()); + self.table_alias + .get_or_insert_with(|| subquery_alias.alias.clone()); } if let LogicalPlan::TableScan(table_scan) = &node @@ -199,17 +203,20 @@ mod test { use common_catalog::consts::DEFAULT_CATALOG_NAME; use common_error::ext::{BoxedError, ErrorExt, StackError}; use common_error::status_code::StatusCode; - use common_recordbatch::SendableRecordBatchStream; + use common_recordbatch::{RecordBatch, SendableRecordBatchStream}; use datafusion::functions_aggregate::count::count_all; + use datafusion::functions_aggregate::min_max::max; use datafusion_common::Column; use datafusion_expr::LogicalPlanBuilder; use datafusion_sql::TableReference; use datatypes::data_type::ConcreteDataType; - use datatypes::schema::{ColumnSchema, SchemaBuilder}; + use datatypes::schema::{ColumnSchema, Schema, SchemaBuilder}; + use datatypes::vectors::{Int64Vector, TimestampMillisecondVector, VectorRef}; use store_api::data_source::DataSource; use store_api::storage::ScanRequest; use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder, TableType}; use table::table::numbers::NumbersTable; + use table::test_util::MemTable; use table::{Table, TableRef}; use super::*; @@ -319,6 +326,216 @@ mod test { ); } + #[test] + fn count_wildcard_shape_matrix() { + let config = datafusion::config::ConfigOptions::default(); + + let direct = CountWildcardToTimeIndexRule + .analyze(count_star(source_plan("source")), &config) + .unwrap(); + assert_count_argument_column(&direct, "source", "ts"); + + let simple_alias = count_star( + LogicalPlanBuilder::from(source_plan("source")) + .alias("projected") + .unwrap() + .build() + .unwrap(), + ); + let simple_alias = CountWildcardToTimeIndexRule + .analyze(simple_alias, &config) + .unwrap(); + assert_count_argument_column(&simple_alias, "projected", "ts"); + + let nested_alias = count_star( + LogicalPlanBuilder::from(source_plan("source")) + .alias("inner") + .unwrap() + .alias("outer") + .unwrap() + .build() + .unwrap(), + ); + let nested_alias = CountWildcardToTimeIndexRule + .analyze(nested_alias, &config) + .unwrap(); + assert_count_argument_column(&nested_alias, "outer", "ts"); + + let nested_rename = count_star( + LogicalPlanBuilder::from(source_plan("source")) + .project(vec![col("ts").alias("renamed")]) + .unwrap() + .alias("projected") + .unwrap() + .build() + .unwrap(), + ); + let nested_rename = CountWildcardToTimeIndexRule + .analyze(nested_rename, &config) + .unwrap(); + assert_count_argument_literal_one(&nested_rename); + + let nested_rename_with_payload_reorder = count_star( + LogicalPlanBuilder::from(source_plan("source")) + .project(vec![col("payload"), col("ts").alias("renamed")]) + .unwrap() + .alias("projected") + .unwrap() + .build() + .unwrap(), + ); + let nested_rename_with_payload_reorder = CountWildcardToTimeIndexRule + .analyze(nested_rename_with_payload_reorder, &config) + .unwrap(); + assert_count_argument_literal_one(&nested_rename_with_payload_reorder); + + let multi_input = count_star( + LogicalPlanBuilder::from(source_plan("left")) + .cross_join(source_plan("right")) + .unwrap() + .build() + .unwrap(), + ); + let multi_input = CountWildcardToTimeIndexRule + .analyze(multi_input, &config) + .unwrap(); + assert_count_argument_literal_one(&multi_input); + } + + #[test] + fn projection_name_collision_falls_back_to_literal_one() { + let before = count_star( + LogicalPlanBuilder::from(source_plan("source")) + .project(vec![col("payload").alias("ts")]) + .unwrap() + .alias("projected") + .unwrap() + .build() + .unwrap(), + ); + + let aggregate = aggregate_plan(&before); + let field = aggregate + .input + .schema() + .qualified_field_with_name(Some(&TableReference::bare("projected")), "ts") + .unwrap(); + assert!(field.1.is_nullable()); + + let after = CountWildcardToTimeIndexRule + .analyze(before, &datafusion::config::ConfigOptions::default()) + .unwrap(); + assert_count_argument_literal_one(&after); + } + + #[test] + fn inner_aggregate_nullable_time_index_name_falls_back_to_literal_one() { + let before = count_star( + LogicalPlanBuilder::from(source_plan("source")) + .aggregate(Vec::::new(), vec![max(col("payload")).alias("ts")]) + .unwrap() + .alias("aggregated") + .unwrap() + .build() + .unwrap(), + ); + + let aggregate = aggregate_plan(&before); + let field = aggregate + .input + .schema() + .qualified_field_with_name(Some(&TableReference::bare("aggregated")), "ts") + .unwrap(); + assert!(field.1.is_nullable()); + + let after = CountWildcardToTimeIndexRule + .analyze(before, &datafusion::config::ConfigOptions::default()) + .unwrap(); + assert_count_argument_literal_one(&after); + } + + fn source_plan(table_name: &str) -> LogicalPlan { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("payload", ConcreteDataType::int64_datatype(), true), + ])); + let columns: Vec = vec![ + Arc::new(TimestampMillisecondVector::from_slice([1, 2, 3])), + Arc::new(Int64Vector::from(vec![Some(10), None, Some(30)])), + ]; + let table = MemTable::table( + table_name, + RecordBatch::new(schema, columns).expect("test record batch must be valid"), + ); + let source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(table), + ))); + LogicalPlanBuilder::scan_with_filters(table_name, source, None, vec![]) + .unwrap() + .build() + .unwrap() + } + + fn count_star(input: LogicalPlan) -> LogicalPlan { + LogicalPlanBuilder::from(input) + .aggregate(Vec::::new(), vec![count_all()]) + .unwrap() + .build() + .unwrap() + } + + fn count_aggregate(plan: &LogicalPlan) -> &AggregateFunction { + let LogicalPlan::Aggregate(aggregate) = plan else { + panic!("expected aggregate plan, got {plan:?}"); + }; + assert_eq!(1, aggregate.aggr_expr.len()); + let expr = unwrap_aliases(&aggregate.aggr_expr[0]); + let Expr::AggregateFunction(count) = expr else { + panic!("expected count aggregate, got {:?}", aggregate.aggr_expr[0]); + }; + assert_eq!("count", count.func.name()); + count + } + + fn unwrap_aliases(expr: &Expr) -> &Expr { + match expr { + Expr::Alias(alias) => unwrap_aliases(alias.expr.as_ref()), + expr => expr, + } + } + + fn assert_count_argument_column(plan: &LogicalPlan, relation: &str, name: &str) { + let count = count_aggregate(plan); + let [Expr::Column(column)] = count.params.args.as_slice() else { + panic!( + "expected one column count argument, got {:?}", + count.params.args + ); + }; + assert_eq!(Some(TableReference::bare(relation)), column.relation); + assert_eq!(name, column.name); + } + + fn assert_count_argument_literal_one(plan: &LogicalPlan) { + let count = count_aggregate(plan); + assert!(matches!( + count.params.args.as_slice(), + [Expr::Literal(ScalarValue::Int64(Some(1)), _)] + )); + } + + fn aggregate_plan(plan: &LogicalPlan) -> &datafusion_expr::logical_plan::Aggregate { + let LogicalPlan::Aggregate(aggregate) = plan else { + panic!("expected aggregate plan, got {plan:?}"); + }; + aggregate + } + fn build_time_index_table(table_name: &str, schema_name: &str, catalog_name: &str) -> TableRef { let column_schemas = vec![ ColumnSchema::new( diff --git a/tests/cases/standalone/common/tql/tql-cte.result b/tests/cases/standalone/common/tql/tql-cte.result index 62a81804fae..0547754b483 100644 --- a/tests/cases/standalone/common/tql/tql-cte.result +++ b/tests/cases/standalone/common/tql/tql-cte.result @@ -155,7 +155,7 @@ SELECT count(*) FROM filtered; +---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------+ | logical_plan | MergeScan [is_placeholder=false, remote_input=[ | | | Projection: count(Int64(1)) AS count(*) | -| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] | +| | Aggregate: groupBy=[[]], aggr=[[count(filtered.ts) AS count(Int64(1))]] | | | SubqueryAlias: filtered | | | Projection: tql_data.ts, tql_data.val | | | SubqueryAlias: tql_data | @@ -642,7 +642,7 @@ SELECT count(*) as high_values FROM final; +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | logical_plan | MergeScan [is_placeholder=false, remote_input=[ | | | Projection: count(Int64(1)) AS count(*) AS high_values | -| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] | +| | Aggregate: groupBy=[[]], aggr=[[count(final.ts) AS count(Int64(1))]] | | | SubqueryAlias: final | | | Projection: processed.ts AS ts, processed.percent AS percent | | | Projection: processed.ts, processed.percent | From 97ca104129536075fcd7329a924db8c40ef88d7f Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 7 Aug 2026 14:48:15 +0800 Subject: [PATCH 015/126] fix(frontend): remove gRPC DDL panics for DropView and non-timestamp time index (#8739) * fix(frontend): remove gRPC DDL panics for DropView and non-timestamp time index Direct gRPC DDL bypasses the SQL parser, so two client-controlled DDL payloads could panic a request handler: - QX-152: DdlExpr::DropView hit todo!() (instance/grpc.rs:247-248). Wire it to the real drop-view implementation (drop_view was pub(crate); widened to pub) so a DropView DDL returns a structured error (e.g. TableNotFound) instead of panicking. - QX-153: a CreateTableExpr whose time_index column is not a timestamp reached Schema::new's unwrap (ddl.rs:2346 -> schema.rs:114-119). create_table_info now uses Schema::try_new with ConvertSchemaSnafu context (InvalidArguments), and the direct gRPC CreateTable arm validates the request via validate_create_expr (which now also checks the time-index column type is a timestamp) before any catalog work. SQL/HTTP paths were already protected by the parser; unchanged. Tests: qx_152_drop_view_via_grpc_ddl_returns_error_not_panic, qx_153_create_table_with_non_timestamp_time_index_via_grpc_returns_error (asserts InvalidArguments), test_create_table_info_rejects_non_timestamp_time_index. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(frontend): add gRPC DDL happy-path coverage for DropView and CreateTable Per review: the initial tests only asserted error paths. Add: - drop_if_exists=true on a missing view succeeds (no error) - dropping an existing view via gRPC DDL succeeds end-to-end - a valid CreateTableExpr with a timestamp time index still succeeds (guards validate_create_expr against rejecting good requests) - qx_152 test now asserts the TableNotFound status instead of is_err() Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/frontend/src/instance.rs | 527 +++++++++++++++++++++++++++++- src/frontend/src/instance/grpc.rs | 15 +- src/operator/src/expr_helper.rs | 31 +- src/operator/src/statement/ddl.rs | 36 +- 4 files changed, 596 insertions(+), 13 deletions(-) diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 1b9710073d0..14acf87787e 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1691,7 +1691,7 @@ mod tests { use common_meta::cache::LayeredCacheRegistryBuilder; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor}; - use common_meta::rpc::ddl::{SubmitDdlTaskRequest, SubmitDdlTaskResponse}; + use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; use common_meta::rpc::procedure::{ MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, }; @@ -1722,7 +1722,9 @@ mod tests { }; use store_api::storage::ScanRequest; use strfmt::Format; - use table::metadata::{FilterPushDownType, TableInfo, TableInfoBuilder, TableMetaBuilder}; + use table::metadata::{ + FilterPushDownType, TableInfo, TableInfoBuilder, TableMetaBuilder, TableType, + }; use table::table_name::TableName; use table::test_util::{EmptyTable, MemTable}; use table::{Table, TableRef}; @@ -2307,6 +2309,139 @@ mod tests { } } + /// A test [`ProcedureExecutor`] that completes create/drop DDL tasks against the + /// in-memory catalog, mimicking what the meta DDL procedures do in production. + /// This allows happy-path DDL requests (create/drop table/view) to be exercised + /// end to end through the gRPC ingress. + struct MockProcedureExecutor { + catalog_manager: Arc, + next_table_id: std::sync::atomic::AtomicU32, + submitted: std::sync::Mutex>, + } + + impl MockProcedureExecutor { + fn new(catalog_manager: Arc) -> Self { + Self { + catalog_manager, + next_table_id: std::sync::atomic::AtomicU32::new(1026), + submitted: std::sync::Mutex::new(Vec::new()), + } + } + } + + #[async_trait::async_trait] + impl ProcedureExecutor for MockProcedureExecutor { + async fn submit_ddl_task( + &self, + _ctx: &ExecutorContext, + request: SubmitDdlTaskRequest, + ) -> common_meta::error::Result { + self.submitted.lock().unwrap().push(request.task.clone()); + match request.task { + DdlTask::CreateTable(task) => { + let table_id = self + .next_table_id + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let mut table_info = task.table_info; + table_info.ident.table_id = table_id; + self.catalog_manager + .register_table_sync(catalog::RegisterTableRequest { + catalog: table_info.catalog_name.clone(), + schema: table_info.schema_name.clone(), + table_name: table_info.name.clone(), + table_id, + table: table::dist_table::DistTable::table(Arc::new(table_info)), + }) + .map_err(BoxedError::new) + .context(common_meta::error::ExternalSnafu)?; + Ok(SubmitDdlTaskResponse { + key: Vec::new(), + table_ids: vec![table_id], + }) + } + DdlTask::CreateView(task) => { + let view_id = self + .next_table_id + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let mut view_info = task.view_info; + view_info.ident.table_id = view_id; + self.catalog_manager + .register_table_sync(catalog::RegisterTableRequest { + catalog: task.create_view.catalog_name.clone(), + schema: task.create_view.schema_name.clone(), + table_name: task.create_view.view_name.clone(), + table_id: view_id, + table: table::dist_table::DistTable::table(Arc::new(view_info)), + }) + .map_err(BoxedError::new) + .context(common_meta::error::ExternalSnafu)?; + Ok(SubmitDdlTaskResponse { + key: Vec::new(), + table_ids: vec![view_id], + }) + } + DdlTask::DropView(task) => { + self.catalog_manager + .deregister_table_sync(catalog::DeregisterTableRequest { + catalog: task.catalog.clone(), + schema: task.schema.clone(), + table_name: task.view.clone(), + }) + .map_err(BoxedError::new) + .context(common_meta::error::ExternalSnafu)?; + Ok(SubmitDdlTaskResponse::default()) + } + other => common_meta::error::UnsupportedSnafu { + operation: format!("mock submit_ddl_task: {other:?}"), + } + .fail(), + } + } + + async fn migrate_region( + &self, + _ctx: &ExecutorContext, + _request: MigrateRegionRequest, + ) -> common_meta::error::Result { + common_meta::error::UnsupportedSnafu { + operation: "migrate_region", + } + .fail() + } + + async fn reconcile( + &self, + _ctx: &ExecutorContext, + _request: ReconcileRequest, + ) -> common_meta::error::Result { + common_meta::error::UnsupportedSnafu { + operation: "reconcile", + } + .fail() + } + + async fn query_procedure_state( + &self, + _ctx: &ExecutorContext, + _pid: &str, + ) -> common_meta::error::Result { + common_meta::error::UnsupportedSnafu { + operation: "query_procedure_state", + } + .fail() + } + + async fn list_procedures( + &self, + _ctx: &ExecutorContext, + ) -> common_meta::error::Result { + common_meta::error::UnsupportedSnafu { + operation: "list_procedures", + } + .fail() + } + } + fn test_cache_registry( kv_backend: common_meta::kv_backend::KvBackendRef, ) -> TestResult { @@ -2485,10 +2620,29 @@ mod tests { target_table: TableRef, plugins: Plugins, metric_names_table: Option, + ) -> TestResult { + let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table); + test_instance_with_catalog_manager( + catalog_manager, + target_table, + plugins, + metric_names_table, + Arc::new(NoopProcedureExecutor), + ) + .await + } + + /// Builds a test frontend `Instance` over the given (already source-registered) + /// catalog manager, completing DDL tasks through `procedure_executor`. + async fn test_instance_with_catalog_manager( + catalog_manager: Arc, + target_table: TableRef, + plugins: Plugins, + metric_names_table: Option, + procedure_executor: ProcedureExecutorRef, ) -> TestResult { let kv_backend = Arc::new(MemoryKvBackend::new()); let process_manager = Arc::new(ProcessManager::new("test-frontend".to_string(), None)); - let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table); let target_table_name = "target"; catalog_manager .register_table_sync(catalog::RegisterTableRequest { @@ -2529,7 +2683,7 @@ mod tests { cache_registry, catalog_manager, Arc::new(client::client_manager::NodeClients::default()), - Arc::new(NoopProcedureExecutor), + procedure_executor, process_manager, ) .with_plugin(plugins) @@ -3277,4 +3431,369 @@ mod tests { assert_eq!(result.is_ok(), is_ok); } } + + /// A `DropView` DDL sent through the direct gRPC ingress must return an error + /// (e.g. table not found) instead of panicking on `todo!()`. + #[tokio::test] + async fn qx_152_drop_view_via_grpc_ddl_returns_error_not_panic() -> TestResult<()> { + let instance = + test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?) + .await?; + + let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest { + expr: Some(api::v1::ddl_request::Expr::DropView( + api::v1::DropViewExpr { + catalog_name: String::new(), + schema_name: String::new(), + view_name: "non_existent_view".to_string(), + view_id: None, + drop_if_exists: false, + }, + )), + }); + + let result = servers::query_handler::grpc::GrpcQueryHandler::do_query( + &instance, + request, + QueryContext::arc(), + ) + .await; + + let err = match result { + Ok(_) => panic!("DropView DDL request must return an error instead of panicking"), + Err(err) => err, + }; + assert_eq!( + err.status_code(), + StatusCode::TableNotFound, + "dropping a non-existent view without IF EXISTS must report TableNotFound, got {err}" + ); + Ok(()) + } + + /// `DROP VIEW IF EXISTS` on a missing view through the direct gRPC ingress must + /// succeed with 0 affected rows (no error, no DDL task submitted), instead of + /// returning `TableNotFound`. + #[tokio::test] + async fn qx_152_drop_view_if_exists_missing_view_via_grpc_ddl_succeeds() -> TestResult<()> { + let catalog_manager = + catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?); + let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone())); + let instance = test_instance_with_catalog_manager( + catalog_manager, + test_table(1025, "target")?, + Plugins::new(), + None, + procedure_executor.clone() as ProcedureExecutorRef, + ) + .await?; + + let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest { + expr: Some(api::v1::ddl_request::Expr::DropView( + api::v1::DropViewExpr { + catalog_name: String::new(), + schema_name: String::new(), + view_name: "non_existent_view".to_string(), + view_id: None, + drop_if_exists: true, + }, + )), + }); + + let result = servers::query_handler::grpc::GrpcQueryHandler::do_query( + &instance, + request, + QueryContext::arc(), + ) + .await; + + let output = match result { + Ok(output) => output, + Err(err) => { + panic!("DROP VIEW IF EXISTS on a missing view must succeed, got error: {err}") + } + }; + assert!( + matches!(output.data, OutputData::AffectedRows(0)), + "DROP VIEW IF EXISTS on a missing view must report 0 affected rows" + ); + assert!( + procedure_executor.submitted.lock().unwrap().is_empty(), + "DROP VIEW IF EXISTS on a missing view must not submit a DDL task" + ); + Ok(()) + } + + /// A `CREATE VIEW` followed by `DROP VIEW` through the direct gRPC ingress must + /// succeed end to end: the view is registered in the catalog and then removed. + #[tokio::test] + async fn qx_152_drop_existing_view_via_grpc_ddl_succeeds() -> TestResult<()> { + let catalog_manager = + catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?); + let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone())); + let instance = test_instance_with_catalog_manager( + catalog_manager, + test_table(1025, "target")?, + Plugins::new(), + None, + procedure_executor.clone() as ProcedureExecutorRef, + ) + .await?; + + // The default "greptime.public" schema must be visible to the kv-backed table + // metadata manager for `CREATE VIEW`/`CREATE TABLE` to pass validation. + instance + .table_metadata_manager() + .schema_manager() + .create( + common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"), + None, + true, + ) + .await + .unwrap(); + + let create_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest { + expr: Some(api::v1::ddl_request::Expr::CreateView( + api::v1::CreateViewExpr { + catalog_name: String::new(), + schema_name: String::new(), + view_name: "my_view".to_string(), + logical_plan: vec![1, 2, 3], + create_if_not_exists: false, + or_replace: false, + table_names: vec![], + columns: vec![], + plan_columns: vec![], + definition: "CREATE VIEW my_view AS SELECT * FROM source".to_string(), + }, + )), + }); + + let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query( + &instance, + create_view_request, + QueryContext::arc(), + ) + .await + { + Ok(output) => output, + Err(err) => panic!("CREATE VIEW via gRPC DDL must succeed, got error: {err}"), + }; + assert!( + matches!(output.data, OutputData::AffectedRows(0)), + "CREATE VIEW via gRPC DDL must report 0 affected rows" + ); + + // The view is registered in the catalog as a view. + let view = instance + .catalog_manager() + .table("greptime", "public", "my_view", None) + .await + .unwrap() + .expect("view should exist after CREATE VIEW"); + assert_eq!(view.table_info().table_type, TableType::View); + + let drop_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest { + expr: Some(api::v1::ddl_request::Expr::DropView( + api::v1::DropViewExpr { + catalog_name: String::new(), + schema_name: String::new(), + view_name: "my_view".to_string(), + view_id: None, + drop_if_exists: false, + }, + )), + }); + + let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query( + &instance, + drop_view_request, + QueryContext::arc(), + ) + .await + { + Ok(output) => output, + Err(err) => panic!("DROP VIEW via gRPC DDL must succeed, got error: {err}"), + }; + assert!( + matches!(output.data, OutputData::AffectedRows(0)), + "DROP VIEW via gRPC DDL must report 0 affected rows" + ); + + // The view is gone after the drop. + assert!( + instance + .catalog_manager() + .table("greptime", "public", "my_view", None) + .await + .unwrap() + .is_none(), + "view should be removed after DROP VIEW" + ); + + let submitted = procedure_executor.submitted.lock().unwrap(); + assert_eq!( + submitted.len(), + 2, + "expected create and drop view tasks, got {submitted:?}" + ); + assert!(matches!(&submitted[0], DdlTask::CreateView(_))); + assert!(matches!(&submitted[1], DdlTask::DropView(_))); + Ok(()) + } + + /// A direct gRPC `CreateTable` whose time index column is not a timestamp must + /// be rejected with `InvalidArguments` instead of panicking while building the schema. + #[tokio::test] + async fn qx_153_create_table_with_non_timestamp_time_index_via_grpc_returns_error() + -> TestResult<()> { + let instance = + test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?) + .await?; + + let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest { + expr: Some(api::v1::ddl_request::Expr::CreateTable( + api::v1::CreateTableExpr { + catalog_name: String::new(), + schema_name: String::new(), + table_name: "demo".to_string(), + desc: String::new(), + column_defs: vec![api::v1::ColumnDef { + name: "host".to_string(), + data_type: api::v1::ColumnDataType::String as i32, + is_nullable: true, + default_constraint: vec![], + semantic_type: 0, + comment: String::new(), + datatype_extension: None, + options: None, + }], + time_index: "host".to_string(), + primary_keys: vec![], + create_if_not_exists: false, + table_options: HashMap::new(), + table_id: None, + engine: "mito".to_string(), + }, + )), + }); + + let result = servers::query_handler::grpc::GrpcQueryHandler::do_query( + &instance, + request, + QueryContext::arc(), + ) + .await; + + let err = match result { + Ok(_) => panic!("CreateTable with a non-timestamp time index must be rejected"), + Err(err) => err, + }; + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); + Ok(()) + } + + /// A valid `CREATE TABLE` (timestamp time index) through the direct gRPC ingress + /// must succeed, guarding that the `validate_create_expr` ingress check doesn't + /// accidentally reject good requests. + #[tokio::test] + async fn qx_153_create_table_with_timestamp_time_index_via_grpc_succeeds() -> TestResult<()> { + let catalog_manager = + catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?); + let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone())); + let instance = test_instance_with_catalog_manager( + catalog_manager, + test_table(1025, "target")?, + Plugins::new(), + None, + procedure_executor.clone() as ProcedureExecutorRef, + ) + .await?; + + // The default "greptime.public" schema must be visible to the kv-backed table + // metadata manager for `CREATE TABLE` to pass validation. + instance + .table_metadata_manager() + .schema_manager() + .create( + common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"), + None, + true, + ) + .await + .unwrap(); + + let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest { + expr: Some(api::v1::ddl_request::Expr::CreateTable( + api::v1::CreateTableExpr { + catalog_name: String::new(), + schema_name: String::new(), + table_name: "demo".to_string(), + desc: String::new(), + column_defs: vec![ + api::v1::ColumnDef { + name: "host".to_string(), + data_type: api::v1::ColumnDataType::String as i32, + is_nullable: true, + default_constraint: vec![], + semantic_type: 0, + comment: String::new(), + datatype_extension: None, + options: None, + }, + api::v1::ColumnDef { + name: "ts".to_string(), + data_type: api::v1::ColumnDataType::TimestampMillisecond as i32, + is_nullable: true, + default_constraint: vec![], + semantic_type: 0, + comment: String::new(), + datatype_extension: None, + options: None, + }, + ], + time_index: "ts".to_string(), + primary_keys: vec![], + create_if_not_exists: false, + table_options: HashMap::new(), + table_id: None, + engine: "mito".to_string(), + }, + )), + }); + + let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query( + &instance, + request, + QueryContext::arc(), + ) + .await + { + Ok(output) => output, + Err(err) => panic!("CREATE TABLE via gRPC DDL must succeed, got error: {err}"), + }; + assert!( + matches!(output.data, OutputData::AffectedRows(0)), + "CREATE TABLE via gRPC DDL must report 0 affected rows" + ); + + // The table is registered in the catalog. + let table = instance + .catalog_manager() + .table("greptime", "public", "demo", None) + .await + .unwrap() + .expect("table should exist after CREATE TABLE"); + assert_eq!(table.table_info().table_type, TableType::Base); + + let submitted = procedure_executor.submitted.lock().unwrap(); + assert_eq!( + submitted.len(), + 1, + "expected one create table task, got {submitted:?}" + ); + assert!(matches!(&submitted[0], DdlTask::CreateTable(_))); + Ok(()) + } } diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index 661afb7f43d..3caac10b841 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -176,6 +176,9 @@ impl GrpcQueryHandler for Instance { match expr { DdlExpr::CreateTable(mut expr) => { + // Direct gRPC DDL bypasses the SQL parser, so validate the + // request here (e.g. the time index must be a timestamp). + operator::expr_helper::validate_create_expr(&expr)?; let _ = self .statement_executor .create_table_inner(&mut expr, None, ctx.clone()) @@ -244,8 +247,16 @@ impl GrpcQueryHandler for Instance { Output::new_with_affected_rows(0) } - DdlExpr::DropView(_) => { - todo!("implemented in the following PR") + DdlExpr::DropView(expr) => { + self.statement_executor + .drop_view( + expr.catalog_name, + expr.schema_name, + expr.view_name, + expr.drop_if_exists, + ctx.clone(), + ) + .await? } DdlExpr::CommentOn(expr) => { self.statement_executor diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index 45e3e340a33..3221598cf24 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -36,6 +36,7 @@ use common_error::ext::BoxedError; use common_grpc_expr::util::ColumnExpr; use common_time::Timezone; use datafusion::sql::planner::object_name_to_table_reference; +use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, @@ -442,14 +443,34 @@ pub fn validate_create_expr(create: &CreateTableExpr) -> Result<()> { } // verify time_index exists - let _ = column_to_indices - .get(&create.time_index) - .with_context(|| InvalidSqlSnafu { + let time_index_idx = + column_to_indices + .get(&create.time_index) + .with_context(|| InvalidSqlSnafu { + err_msg: format!( + "column name `{}` is not found in column list", + create.time_index + ), + })?; + + // verify time_index is a timestamp column + let time_index_column = &create.column_defs[*time_index_idx]; + let data_type = ConcreteDataType::from( + ColumnDataTypeWrapper::try_new( + time_index_column.data_type, + time_index_column.datatype_extension.clone(), + ) + .context(ColumnDataTypeSnafu)?, + ); + ensure!( + data_type.is_timestamp(), + InvalidSqlSnafu { err_msg: format!( - "column name `{}` is not found in column list", + "column `{}` is not a timestamp type, it can't be used as time index", create.time_index ), - })?; + } + ); // verify primary_key exists for pk in &create.primary_keys { diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index b0986de59f0..96321b13a3f 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -1213,7 +1213,7 @@ impl StatementExecutor { /// Drop a view #[tracing::instrument(skip_all)] - pub(crate) async fn drop_view( + pub async fn drop_view( &self, catalog: String, schema: String, @@ -2343,7 +2343,7 @@ pub fn create_table_info( } let next_column_id = column_schemas.len() as u32; - let schema = Arc::new(Schema::new(column_schemas)); + let schema = Arc::new(Schema::try_new(column_schemas).context(ConvertSchemaSnafu)?); let primary_key_indices = create_table .primary_keys @@ -3443,6 +3443,38 @@ WITH ('repartition.column.hint' = ' host ')", ); } + #[test] + fn test_create_table_info_rejects_non_timestamp_time_index() { + let expr = CreateTableExpr { + catalog_name: "greptime".to_string(), + schema_name: "public".to_string(), + table_name: "demo".to_string(), + desc: String::new(), + column_defs: vec![api::v1::ColumnDef { + name: "host".to_string(), + data_type: api::v1::ColumnDataType::String as i32, + is_nullable: true, + default_constraint: vec![], + semantic_type: 0, + comment: String::new(), + datatype_extension: None, + options: None, + }], + time_index: "host".to_string(), + primary_keys: vec![], + create_if_not_exists: false, + table_options: HashMap::new(), + table_id: None, + engine: "mito".to_string(), + }; + + let err = create_table_info(&expr, vec![]).unwrap_err(); + assert_eq!( + common_error::ext::ErrorExt::status_code(&err), + common_error::status_code::StatusCode::InvalidArguments + ); + } + #[test] fn test_json2_requires_append_mode() { let cases = [ From ec51113ab6769fe9d2ecd34497e72fcc727d0e8a Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:23:43 +0800 Subject: [PATCH 016/126] feat(protocol): validate native histogram ingestion (#8775) Signed-off-by: shuiyisong --- src/operator/src/insert.rs | 62 +++++++- src/servers/src/otlp/metrics.rs | 9 +- src/servers/src/prom_remote_write/decode.rs | 156 ++++++++++++++++++-- src/servers/src/prom_remote_write/v2.rs | 126 ++++++++++++---- src/servers/src/prom_store.rs | 2 + 5 files changed, 313 insertions(+), 42 deletions(-) diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 74ecb6b3cd8..2bfcc617cb1 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -1496,7 +1496,7 @@ mod tests { } #[tokio::test] - async fn test_accommodate_existing_schema_logic() { + async fn test_accommodate_existing_schema_and_reject_kind_changes() { let ts_name = "my_ts"; let field_name = "my_field"; let table = @@ -1548,6 +1548,66 @@ mod tests { let req_schema = req.rows.as_ref().unwrap().schema.clone(); assert_eq!(req_schema[0].column_name, ts_name); assert_eq!(req_schema[1].column_name, field_name); + + let (datatype, datatype_extension) = + ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone()) + .unwrap() + .into_parts(); + let mut histogram_req = RowInsertRequest { + table_name: "test_table".to_string(), + rows: Some(Rows { + schema: vec![ + time_index_column_schema("ts", ColumnDataType::TimestampMillisecond), + api::v1::ColumnSchema { + column_name: greptime_native_histogram().to_string(), + datatype: datatype as i32, + semantic_type: SemanticType::Field as i32, + datatype_extension, + options: None, + }, + ], + rows: vec![], + }), + }; + let error = inserter + .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, true, true, true) + .unwrap_err(); + assert!( + error + .to_string() + .contains("cannot mix native histogram and float sample fields") + ); + + let histogram_table = make_table_ref_with_schema( + "ts", + greptime_native_histogram(), + native_histogram_value_type().clone(), + ); + let mut sample_req = RowInsertRequest { + table_name: "test_table".to_string(), + rows: Some(Rows { + schema: vec![ + time_index_column_schema("ts", ColumnDataType::TimestampMillisecond), + field_column_schema(greptime_value(), ColumnDataType::Float64), + ], + rows: vec![], + }), + }; + let error = inserter + .get_alter_table_expr_on_demand( + &mut sample_req, + &histogram_table, + &ctx, + true, + true, + true, + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("cannot mix native histogram and float sample fields") + ); } #[test] diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 30a10aebed9..56ab0530338 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -374,7 +374,8 @@ fn encode_metrics( metric_ctx, )?; } - // TODO(sunng87) leave ExponentialHistogram for next release + // TODO: Convert OTLP exponential histograms into the canonical native-histogram + // Struct. See docs/rfcs/2026-08-04-native-histograms.md. metric::Data::ExponentialHistogram(_hist) => {} } } @@ -714,12 +715,6 @@ fn encode_histogram( Ok(()) } -#[allow(dead_code)] -fn encode_exponential_histogram(_name: &str, _hist: &ExponentialHistogram) -> Result<()> { - // TODO(sunng87): implement this using a prometheus compatible way - Ok(()) -} - fn encode_summary( table_writer: &mut MultiTableData, name: &str, diff --git a/src/servers/src/prom_remote_write/decode.rs b/src/servers/src/prom_remote_write/decode.rs index 3f108cdc2fa..6f635c7caf1 100644 --- a/src/servers/src/prom_remote_write/decode.rs +++ b/src/servers/src/prom_remote_write/decode.rs @@ -56,6 +56,12 @@ pub(crate) struct PromTimeSeries { impl Clear for PromTimeSeries { fn clear(&mut self) { self.table_name.clear(); + self.schema.clear(); + self.physical_table.clear(); + // Labels borrow from the request buffer, which may be replaced after this reset. + for label in self.labels.iter_mut() { + label.clear(); + } self.labels.clear(); self.samples.clear(); } @@ -93,27 +99,31 @@ impl PromTimeSeries { } #[allow(deprecated)] - match label.name { + let is_special_label = match label.name { METRIC_NAME_LABEL_BYTES => { self.table_name = prom_validation_mode.decode_string(label.value)?; - self.labels.truncate(self.labels.len() - 1); + true } SCHEMA_LABEL_BYTES => { self.schema = Some(prom_validation_mode.decode_string(label.value)?); - self.labels.truncate(self.labels.len() - 1); + true } DATABASE_LABEL_BYTES | DATABASE_LABEL_ALT_BYTES => { if self.schema.is_none() { self.schema = Some(prom_validation_mode.decode_string(label.value)?); } - self.labels.truncate(self.labels.len() - 1); + true } PHYSICAL_TABLE_LABEL_BYTES | PHYSICAL_TABLE_LABEL_ALT_BYTES => { self.physical_table = Some(prom_validation_mode.decode_string(label.value)?); - self.labels.truncate(self.labels.len() - 1); + true } - _ => {} + _ => false, + }; + if is_special_label { + label.clear(); + self.labels.truncate(self.labels.len() - 1); } Ok(()) @@ -129,6 +139,9 @@ impl PromTimeSeries { Ok(()) } 3u32 => prost::encoding::skip_field(wire_type, tag, buf, Default::default()), + 4u32 => Err(DecodeError::new( + "remote write v1 native histogram ingestion is unsupported; use remote write v2", + )), _ => prost::encoding::skip_field(wire_type, tag, buf, Default::default()), } } @@ -170,6 +183,7 @@ pub struct PromWriteRequest<'a> { impl<'a> Clear for PromWriteRequest<'a> { fn clear(&mut self) { + self.series.clear(); self.table_data.clear(); } } @@ -186,6 +200,7 @@ impl<'a> PromWriteRequest<'a> { processor: &mut PromSeriesProcessor, ) -> Result<(), DecodeError> { const STRUCT_NAME: &str = "PromWriteRequest"; + self.clear(); self.table_data.set_raw_data(buf); let mut offset = 0; while offset < self.table_data.raw_data.len() { @@ -247,8 +262,7 @@ impl<'a> PromWriteRequest<'a> { } if decoded_timeseries { - self.series.labels.clear(); - self.series.samples.clear(); + self.series.clear(); } } @@ -371,7 +385,7 @@ impl PromSeriesProcessor { mod tests { use std::collections::HashMap; - use api::prom_store::remote::WriteRequest; + use api::prom_store::remote::{Histogram, Label, Sample, TimeSeries, WriteRequest}; use api::v1::{Row, RowInsertRequests, Rows}; use bytes::Bytes; use prost::Message; @@ -453,6 +467,130 @@ mod tests { } } + #[test] + fn test_decode_rejects_remote_write_v1_native_histograms() { + for samples in [ + Vec::new(), + vec![Sample { + value: 1.0, + timestamp: 1000, + }], + ] { + let request = WriteRequest { + timeseries: vec![TimeSeries { + samples, + histograms: vec![Histogram::default()], + ..Default::default() + }], + ..Default::default() + }; + let mut processor = PromSeriesProcessor::default_processor(); + let mut write_request = PromWriteRequest::default(); + let error = write_request + .decode( + request.encode_to_vec(), + PromValidationMode::Strict, + &mut processor, + ) + .unwrap_err(); + + assert!(error.to_string().contains( + "remote write v1 native histogram ingestion is unsupported; use remote write v2" + )); + assert_eq!( + write_request.as_row_insert_requests().ref_all_req().count(), + 0 + ); + } + } + + #[test] + fn test_decode_clears_state_after_error() { + let label = |name: &str, value: &str| Label { + name: name.to_string(), + value: value.to_string(), + }; + let sample = Sample { + value: 1.0, + timestamp: 1000, + }; + let failed_request = WriteRequest { + timeseries: vec![ + TimeSeries { + labels: vec![ + label("__name__", "stale_metric"), + label("stale_label", "stale_value"), + ], + samples: vec![sample], + ..Default::default() + }, + TimeSeries { + labels: vec![ + label("__name__", "rejected_metric"), + label("rejected_label", "rejected_value"), + label("__schema__", "rejected_schema"), + label("x_greptime_physical_table", "rejected_physical_table"), + ], + samples: vec![sample], + histograms: vec![Histogram::default()], + ..Default::default() + }, + ], + ..Default::default() + }; + let successful_request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + label("__name__", "fresh_metric"), + label("fresh_label", "fresh_value"), + ], + samples: vec![sample], + ..Default::default() + }], + ..Default::default() + }; + let mut processor = PromSeriesProcessor::default_processor(); + let mut write_request = PromWriteRequest::default(); + + write_request + .decode( + failed_request.encode_to_vec(), + PromValidationMode::Strict, + &mut processor, + ) + .unwrap_err(); + write_request + .decode( + successful_request.encode_to_vec(), + PromValidationMode::Strict, + &mut processor, + ) + .unwrap(); + + assert_eq!(write_request.table_data.tables.len(), 1); + let (prom_ctx, tables) = write_request.table_data.tables.iter().next().unwrap(); + assert_eq!(prom_ctx.schema, None); + assert_eq!(prom_ctx.physical_table, None); + assert_eq!(tables.len(), 1); + assert!(tables.contains_key("fresh_metric")); + + let requests = write_request + .as_row_insert_requests() + .all_req() + .collect::>(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].table_name, "fresh_metric"); + let rows = requests[0].rows.as_ref().unwrap(); + assert_eq!(rows.rows.len(), 1); + assert_eq!( + rows.schema + .iter() + .map(|column| column.column_name.as_str()) + .collect::>(), + vec![greptime_timestamp(), greptime_value(), "fresh_label"] + ); + } + #[test] fn test_decode_string_strict_mode_valid_utf8() { let valid_utf8 = Bytes::from("hello world"); diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 229fac01666..cfa772f6919 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -45,7 +45,6 @@ use crate::row_writer::{self, TableData}; type PromTags = Vec<(String, String)>; type ResolvedSeriesLabels = (PromCtx, String, PromTags); -const MIN_REMOTE_WRITE_V2_SCHEMA: i32 = -4; const MAX_REMOTE_WRITE_V2_SCHEMA: i32 = 8; const MAX_REDUCIBLE_REMOTE_WRITE_V2_SCHEMA: i32 = 52; @@ -346,7 +345,7 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { } fn validate_native_histogram(histogram: &Histogram, uses_float_counts: bool) -> Result<()> { - validate_native_histogram_schema(histogram.schema)?; + let exponential_overflow_index = validate_native_histogram_schema(histogram.schema)?; validate_native_histogram_custom_values(histogram)?; if histogram.schema == CUSTOM_BUCKETS_SCHEMA { @@ -379,8 +378,11 @@ fn validate_native_histogram(histogram: &Histogram, uses_float_counts: bool) -> histogram.negative_deltas.len(), ) }; - let custom_max_index = if histogram.schema == CUSTOM_BUCKETS_SCHEMA { - Some( + let bucket_index_range = if let Some(overflow_index) = exponential_overflow_index { + (i32::MIN, overflow_index) + } else { + ( + 0, i32::try_from(histogram.custom_values.len()).ok().context( error::InvalidPromRemoteRequestSnafu { msg: "remote write v2 custom native histogram has too many custom_values" @@ -388,30 +390,30 @@ fn validate_native_histogram(histogram: &Histogram, uses_float_counts: bool) -> }, )?, ) - } else { - None }; validate_native_histogram_spans( "positive", &histogram.positive_spans, positive_buckets, - custom_max_index, + bucket_index_range, )?; validate_native_histogram_spans( "negative", &histogram.negative_spans, negative_buckets, - None, + bucket_index_range, )?; Ok(()) } -fn validate_native_histogram_schema(schema: i32) -> Result<()> { - if schema == CUSTOM_BUCKETS_SCHEMA - || (MIN_REMOTE_WRITE_V2_SCHEMA..=MAX_REMOTE_WRITE_V2_SCHEMA).contains(&schema) - { - return Ok(()); +fn validate_native_histogram_schema(schema: i32) -> Result> { + if schema == CUSTOM_BUCKETS_SCHEMA { + return Ok(None); + } + + if let Some(overflow_index) = exponential_overflow_bucket_index(schema) { + return Ok(Some(overflow_index)); } if (MAX_REMOTE_WRITE_V2_SCHEMA + 1..=MAX_REDUCIBLE_REMOTE_WRITE_V2_SCHEMA).contains(&schema) { @@ -467,7 +469,7 @@ fn validate_native_histogram_spans( name: &str, spans: &[BucketSpan], bucket_count: usize, - custom_max_index: Option, + bucket_index_range: (i32, i32), ) -> Result<()> { let span_len = spans .iter() @@ -487,7 +489,7 @@ fn validate_native_histogram_spans( let mut current_index = 0i32; for (span_index, span) in spans.iter().enumerate() { ensure!( - span.offset >= 0 || (span_index == 0 && custom_max_index.is_none()), + span.offset >= 0 || (span_index == 0 && bucket_index_range.0 == i32::MIN), error::InvalidPromRemoteRequestSnafu { msg: format!( "remote write v2 native histogram {name} span {} has negative offset {}", @@ -509,16 +511,14 @@ fn validate_native_histogram_spans( }; for _ in 0..span.length { - if let Some(max_index) = custom_max_index { - ensure!( - (0..=max_index).contains(¤t_index), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 custom native histogram {name} bucket index {current_index} is out of range" - ), - } - ); - } + ensure!( + (bucket_index_range.0..=bucket_index_range.1).contains(¤t_index), + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram {name} bucket index {current_index} is out of range" + ), + } + ); current_index = current_index .checked_add(1) @@ -1473,6 +1473,82 @@ mod tests { ); } + #[test] + fn test_into_context_req_rejects_metric_kind_conflict_across_label_sets() { + let request = Request { + symbols: vec![ + "".to_string(), + METRIC_NAME_LABEL.to_string(), + "metric".to_string(), + "job".to_string(), + "api".to_string(), + "worker".to_string(), + ], + timeseries: vec![ + TimeSeries { + labels_refs: vec![1, 2, 3, 4], + samples: vec![Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + ..Default::default() + }, + TimeSeries { + labels_refs: vec![1, 2, 3, 5], + histograms: vec![Histogram::default()], + ..Default::default() + }, + ], + }; + + assert_invalid( + "same metric kind conflict across label sets", + request, + "contains both samples and native histograms", + ); + } + + #[test] + fn test_into_context_req_validates_exponential_overflow_bucket_index() { + for schema in [-4, 0, 8] { + let max_index = exponential_overflow_bucket_index(schema).unwrap(); + for positive in [true, false] { + let mut histogram = Histogram { + schema, + count: Some(Count::CountInt(1)), + ..Default::default() + }; + if positive { + histogram.positive_spans = vec![BucketSpan { + offset: max_index, + length: 1, + }]; + histogram.positive_deltas = vec![1]; + } else { + histogram.negative_spans = vec![BucketSpan { + offset: max_index, + length: 1, + }]; + histogram.negative_deltas = vec![1]; + } + into_write_requests(request_with_histogram(histogram.clone())).unwrap(); + + let beyond = max_index + 1; + if positive { + histogram.positive_spans[0].offset = beyond; + } else { + histogram.negative_spans[0].offset = beyond; + } + assert_invalid( + "exponential overflow bucket index", + request_with_histogram(histogram), + &format!("bucket index {beyond} is out of range"), + ); + } + } + } + #[test] fn test_into_context_req_converts_histograms_and_ignores_exemplars() { let request = Request { diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 9988f7051f4..772689f0e07 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -381,6 +381,8 @@ fn recordbatch_to_timeseries(table: &str, recordbatch: RecordBatch) -> Result Date: Fri, 7 Aug 2026 16:12:35 +0800 Subject: [PATCH 017/126] feat: support generic heartbeat response extension accumulation (#8786) Signed-off-by: jeremyhi --- src/meta-srv/src/handler.rs | 55 ++++++++++++++++++++++++++++++++++--- 1 file changed, 51 insertions(+), 4 deletions(-) diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs index 082f902ee12..50b53f4970a 100644 --- a/src/meta-srv/src/handler.rs +++ b/src/meta-srv/src/handler.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::cmp::Ordering; -use std::collections::{BTreeMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt::{Debug, Display}; use std::ops::Bound; use std::sync::{Arc, Mutex}; @@ -119,6 +119,8 @@ pub struct HeartbeatAccumulator { pub stat: Option, pub inactive_region_ids: HashSet, pub region_lease: Option, + /// Generic heartbeat response extensions accumulated by handlers. + pub extensions: HashMap>, } impl HeartbeatAccumulator { @@ -405,7 +407,7 @@ impl HeartbeatHandlerGroup { region_lease: acc.region_lease, mailbox_message, heartbeat_config, - extensions: Default::default(), + extensions: std::mem::take(&mut acc.extensions), }; Ok(res) } @@ -888,16 +890,21 @@ mod tests { use std::sync::Arc; use std::time::Duration; - use api::v1::meta::{MailboxMessage, Role}; + use api::v1::meta::{HeartbeatRequest, MailboxMessage, RequestHeader, Role}; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::sequence::SequenceBuilder; use tokio::sync::mpsc; - use super::{HeartbeatHandlerGroupBuilder, PusherId, Pushers}; + use super::{ + HandleControl, HeartbeatAccumulator, HeartbeatHandler, HeartbeatHandlerGroupBuilder, + PusherId, Pushers, + }; use crate::error; use crate::handler::collect_stats_handler::CollectStatsHandler; use crate::handler::response_header_handler::ResponseHeaderHandler; + use crate::handler::test_utils::TestEnv; use crate::handler::{HeartbeatHandlerGroup, HeartbeatMailbox, Pusher}; + use crate::metasrv::Context; use crate::service::mailbox::{Channel, MailboxReceiver, MailboxRef}; #[tokio::test] @@ -1316,4 +1323,44 @@ mod tests { drop(pusher); deregister_signal_tx.changed().await.unwrap(); } + + #[tokio::test] + async fn test_custom_handler_extension_reaches_response() { + struct CustomExtensionHandler; + + #[async_trait::async_trait] + impl HeartbeatHandler for CustomExtensionHandler { + fn is_acceptable(&self, role: Role) -> bool { + role == Role::Frontend + } + + async fn handle( + &self, + _req: &HeartbeatRequest, + _ctx: &mut Context, + acc: &mut HeartbeatAccumulator, + ) -> crate::error::Result { + acc.extensions + .insert("custom.key".to_string(), b"custom-value".to_vec()); + Ok(HandleControl::Continue) + } + } + + let mut builder = + HeartbeatHandlerGroupBuilder::new(Pushers::default()).add_default_handlers(); + builder.add_handler_last(CustomExtensionHandler); + + let group = builder.build().unwrap(); + + let req = HeartbeatRequest { + header: Some(RequestHeader::new(1, Role::Frontend, Default::default())), + ..Default::default() + }; + let ctx = TestEnv::new().ctx(); + + let res = group.handle(req, ctx).await.unwrap(); + + let extensions = res.extensions; + assert_eq!(extensions.get("custom.key").unwrap(), b"custom-value"); + } } From bb6d55a99f8fffdf4e20b3a809162d387ffc42b4 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 7 Aug 2026 16:40:46 +0800 Subject: [PATCH 018/126] feat: support old-stage datanode config overlays (#8647) * feat: support old-stage datanode config overlays Signed-off-by: discord9 * fix: derive compat overlay policy from WAL config Signed-off-by: discord9 --------- Signed-off-by: discord9 --- tests/compatibility/AGENTS.md | 37 +- tests/compatibility/README.md | 42 +- tests/runner/src/cmd.rs | 1 + tests/runner/src/cmd/compat.rs | 787 +++++++++++++++++++---- tests/runner/src/cmd/compat_case.rs | 93 +++ tests/runner/src/cmd/datanode_overlay.rs | 637 ++++++++++++++++++ tests/runner/src/env/bare.rs | 214 +++++- tests/runner/src/server_mode.rs | 176 ++++- tests/runner/src/util.rs | 62 +- 9 files changed, 1870 insertions(+), 179 deletions(-) create mode 100644 tests/runner/src/cmd/datanode_overlay.rs diff --git a/tests/compatibility/AGENTS.md b/tests/compatibility/AGENTS.md index 8ef9b6c0117..144f5af3322 100644 --- a/tests/compatibility/AGENTS.md +++ b/tests/compatibility/AGENTS.md @@ -28,12 +28,42 @@ This file is intended for AI agents editing compat cases or the compat runner. F - All required fields must be non-empty: `name`, `reason`, `introduced_by`, `topologies`, `from_range`, `to_range`, `features`, `owner`. +## Old-Stage Datanode Overlays + +- An old-stage datanode overlay is declared only as: + + ```toml + [old_config] + datanode = "old-datanode.overlay.toml" + ``` + + `datanode` is required if `[old_config]` exists. Empty tables and unknown + keys are parse errors. +- The sidecar path is relative to its case directory and must remain confined + there. It contains native datanode TOML and is loaded and preflighted before + services or state are created. +- The runner merges tables recursively only when both values are tables. + Scalars, type mismatches, arrays, and arrays of tables replace atomically; + `region_engine` has no special merge behavior. +- Do not use an overlay to set runner-owned fields: `mode`, `node_id`, + `storage.data_home`, `meta_client_options.metasrv_addrs`, `wal.provider`, or + `wal.dir` for Raft WAL / `wal.broker_endpoints` for Kafka WAL. The runner + restores or deletes them according to its baseline and warns without showing + values. + ## Phase Semantics - `setup.sql` runs on the **old (from)** binary. Only success is required; output is not compared to any file. - `verify.sql` runs on the **new (to)** binary. Output is compared against `verify.result`. +- Overlays apply only to old-stage datanodes and survive old-stage setup + restarts; the current stage uses a clean configuration. +- The baseline profile runs first. Cases with semantically equivalent datanode + TOML share one sequential, isolated profile, and every profile has an + independent state and etcd lifecycle. +- Without fail-fast, only cases with successful setup are verified. Fail-fast + cleans up the active profile before stopping. ## PostgreSQL Protocol Cases @@ -48,9 +78,10 @@ cargo run -p sqlness-runner -- compat --dry-run [--from-version vX.Y.Z] [--test- ``` The dry-run performs full discovery and filtering (name, topology, metadata -validation, namespace dedup, version-range matching) but starts no services, -creates no temp dirs, and mutates no files. Use it to check which cases -would be selected before a real run. +validation, namespace dedup, version-range matching) and displays selected +profiles, cases, and sidecar paths without configuration values. It starts no +services, creates no temp dirs, and mutates no files. Use it to check which +cases would be selected before a real run. ## CI Version Window diff --git a/tests/compatibility/README.md b/tests/compatibility/README.md index a88fddeaf79..542f10debfd 100644 --- a/tests/compatibility/README.md +++ b/tests/compatibility/README.md @@ -68,6 +68,34 @@ namespace = "my_explicit_namespace" # defaults to sanitized directory name **Required fields**: `name`, `reason`, `introduced_by`, `topologies`, `from_range`, `to_range`, `features`, `owner`. +### Old-Stage Datanode Configuration Overlay + +To apply a datanode configuration overlay while running the old stage, add this +strict optional table to `case.toml`: + +```toml +[old_config] +datanode = "old-datanode.overlay.toml" +``` + +`datanode` is required whenever `[old_config]` is present; empty tables and +unknown keys are rejected. The reference is relative to the case directory and +must remain confined to that directory. The sidecar is native datanode TOML, +which the runner loads and preflights before starting services or creating +state. + +The runner first applies the datanode baseline, then merges the sidecar. Tables +merge recursively only when both values are tables. Scalars, type mismatches, +arrays, and arrays of tables replace the baseline value atomically. In +particular, `region_engine` has no special merge behavior. + +Runner-owned fields cannot be changed by an overlay: `mode`, `node_id`, +`storage.data_home`, `meta_client_options.metasrv_addrs`, and `wal.provider`, +plus `wal.dir` for Raft WAL or `wal.broker_endpoints` for Kafka WAL. The runner +restores these fields to its baseline values, or deletes them when the baseline +has no value. It warns about protected-field overrides without printing their +values. + ### Version-Range Filtering `from_range` and `to_range` control which binary versions a case applies to: @@ -166,9 +194,19 @@ Each case runs in its own database namespace to prevent cross-case interference: ## Batch Behavior -- All cases in a run share one cluster lifecycle: start from-version cluster → run all setups → restart with to-version binary → run all verifies -- Cases run **serially** (no parallelism in PR1). Namespace state is session/protocol state and cannot be shared concurrently. +- The baseline (no-overlay) profile runs first. Cases whose old datanode TOML + is semantically equivalent share one profile; profiles run serially and in + isolation. +- Each profile has its own state and etcd lifecycle. Its overlay is applied + only to old-stage datanodes and remains in effect through old-stage setup + restarts. The current stage always uses a clean configuration. +- Cases run **serially** (no parallelism in PR1). Namespace state is + session/protocol state and cannot be shared concurrently. - Same namespace across cases is rejected. +- Without fail-fast, the runner verifies only cases whose setup succeeded. + With fail-fast, it cleans up the active profile before stopping. +- `--dry-run` displays the selected profiles, cases, and sidecar paths without + printing configuration values; it starts no services. ## xfail Policy (Future) diff --git a/tests/runner/src/cmd.rs b/tests/runner/src/cmd.rs index e17d910b855..c0fed551a33 100644 --- a/tests/runner/src/cmd.rs +++ b/tests/runner/src/cmd.rs @@ -15,6 +15,7 @@ pub(crate) mod bare; pub(crate) mod compat; pub(crate) mod compat_case; +pub(crate) mod datanode_overlay; pub(crate) mod kube; use std::path::PathBuf; diff --git a/tests/runner/src/cmd/compat.rs b/tests/runner/src/cmd/compat.rs index ce33673a4a2..b7473131ef2 100644 --- a/tests/runner/src/cmd/compat.rs +++ b/tests/runner/src/cmd/compat.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::BTreeMap; use std::path::PathBuf; use std::sync::Arc; @@ -22,7 +23,10 @@ use sqlness::interceptor::{InterceptorRef, Registry}; use crate::cmd::bare::ServerAddr; use crate::cmd::compat_case::{self, CompatCase, try_infer_version, version_matches_range}; -use crate::env::bare::{Env, StoreConfig, WalConfig}; +use crate::cmd::datanode_overlay::{ + DatanodeOverlay, DatanodeProtectionPolicy, PreparedDatanodeOverlay, +}; +use crate::env::bare::{Env, GreptimeDB, StoreConfig, WalConfig}; use crate::protocol_interceptor::{self, POSTGRES, PROTOCOL_KEY}; use crate::util; @@ -326,16 +330,12 @@ impl CompatCommand { ); } - if let Some(expected_cases) = self.expect_cases { - assert_eq!( - cases.len(), - expected_cases, - "Expected {expected_cases} compatibility cases after all filtering, found {}", - cases.len() - ); - } - if cases.is_empty() { + if let Some(expected_cases) = self.expect_cases { + panic!( + "Expected {expected_cases} compatibility cases after all filtering, found 0" + ); + } if dry_run { println!("DRY-RUN: no compat cases would run after version-range filtering"); } else { @@ -344,8 +344,25 @@ impl CompatCommand { return; } + let wal_config = WalConfig::RaftEngine; + let profiles = + prepare_compat_profiles(cases, &wal_config).unwrap_or_else(|error| panic!("{error}")); + emit_protected_path_warnings(&profiles); + + if let Some(expected_cases) = self.expect_cases { + assert_eq!( + profiles.case_count(), + expected_cases, + "Expected {expected_cases} compatibility cases after all filtering, found {}", + profiles.case_count() + ); + } + if dry_run { - println!("DRY-RUN: would run {} compat case(s)", cases.len()); + println!( + "DRY-RUN: would run {} compat case(s)", + profiles.case_count() + ); println!(" topology: {}", topology.as_str()); println!( " from version: {}", @@ -359,22 +376,30 @@ impl CompatCommand { .as_deref() .unwrap_or("unknown (use --to-bins-dir or build debug binary)") ); - if pre_filter_count != cases.len() { + if pre_filter_count != profiles.case_count() { println!(); println!( "Version-range filtering reduced {} → {} cases (see 'Skipping case' messages above)", pre_filter_count, - cases.len() + profiles.case_count() ); } println!(); - for c in &cases { + for c in profiles.cases() { println!(" case: {}", c.metadata.name); println!(" namespace: {}", c.namespace); println!(" from_range: {:?}", c.metadata.from_range); println!(" to_range: {:?}", c.metadata.to_range); println!(" features: {:?}", c.metadata.features); } + println!("DRY-RUN: compatibility profiles:"); + for profile in profiles.iter() { + println!(" profile: {}", profile.profile_id()); + println!(" cases: {}", profile.case_names().join(", ")); + for source in profile.sources() { + println!(" sidecar: {}", source.display()); + } + } println!(); println!("Dry run complete. Remove --dry-run to execute."); return; @@ -382,133 +407,499 @@ impl CompatCommand { println!( "Running {} compat case(s) with topology {}:", - cases.len(), + profiles.case_count(), topology.as_str() ); - for c in &cases { + for c in profiles.cases() { println!( " - {} (namespace: {}, topologies: {:?})", c.metadata.name, c.namespace, c.metadata.topologies ); } - // ---- 6. Create temp directory (after filtering so early exits don't leave empty dirs) ---- - let temp_dir = tempfile::Builder::new() - .prefix("sqlness-compat") - .tempdir() - .unwrap(); - let sqlness_home = temp_dir.keep(); - unsafe { - std::env::set_var( - "SQLNESS_HOME", - sqlness_home.join("copy").display().to_string(), - ); - } - - // ---- 7. Build interceptor registry ---- + // ---- 6. Build interceptor registry ---- let interceptor_registry = create_interceptor_registry(); - - // ---- 7b. Create Env for the selected topology ---- - let setup_etcd = topology == CompatTopology::Distributed && self.setup_etcd; - let store_config = StoreConfig { - store_addrs: if setup_etcd { - vec!["127.0.0.1:2379".to_string()] - } else { - vec![] - }, - setup_etcd, - setup_pg: None, - setup_mysql: None, - enable_flat_format: false, - enable_gc: false, - }; - - let env = Env::new( - sqlness_home.clone(), - ServerAddr::default(), - WalConfig::RaftEngine, - self.pull_version_on_need, + let to_bins_dir = to_bins_dir.expect("to_bins_dir must be resolved in non-dry-run mode"); + let profile_config = ProfileRunConfig { + interceptor_registry: &interceptor_registry, from_bins_dir, - store_config, - vec![], - ); - - // ---- 7c. Etcd cleanup guard ---- - // Arm this only immediately before starting the cluster. Earlier validation - // failures should not stop an unrelated local container named `etcd`. - let mut etcd_guard = if setup_etcd { - Some(EtcdGuard::new()) - } else { - None + to_bins_dir, + pull_version_on_need: self.pull_version_on_need, + setup_etcd: self.setup_etcd, + fail_fast: self.fail_fast, + topology, + preserve_state: self.preserve_state, + wal_config, }; - - // ---- 8. Run setup phase on the from-version cluster ---- - println!("Starting from-version {} cluster...", topology.as_str()); - let mut db = match topology { - CompatTopology::Distributed => env.compat_start_distributed(0).await, - CompatTopology::Standalone => env.compat_start_standalone(0).await, - }; - - println!("Running setup phase..."); - for case in &cases { - run_compat_phase(&db, case, &interceptor_registry, CompatPhase::Setup) - .await - .unwrap_or_else(|e| panic!("Setup failed for case '{}': {e}", case.metadata.name)); - println!(" Setup: {} - OK", case.metadata.name); + let mut failures = Vec::new(); + let mut preserved_states = Vec::new(); + for profile in profiles.iter() { + let outcome = run_profile(profile, &profile_config).await; + failures.extend(outcome.failures); + if let Some(state) = outcome.preserved_state { + preserved_states.push(state); + } + if outcome.stop_remaining_profiles { + break; + } } - // ---- 9. Switch to the to-version binary and restart cluster ---- - // to_bins_dir was already resolved during version-range filtering - println!("Restarting cluster with to-version binary on preserved state..."); - let to_bins_dir = to_bins_dir.expect("to_bins_dir must be resolved in non-dry-run mode"); - env.compat_restart(&db, to_bins_dir).await; + for state in &preserved_states { + println!("Preserved compat profile state: {}", state.display()); + } + + if failures.is_empty() { + println!("\n\x1b[32mAll compat tests passed!\x1b[0m"); + } else { + panic!("\n\x1b[31mFailed cases: {}\x1b[0m", failures.join(", ")); + } + } +} + +/// Cases sharing one old-stage datanode configuration lifecycle. +#[derive(Debug)] +enum CompatProfile { + Baseline { + cases: Vec, + }, + Overlay { + overlay: Arc, + cases: Vec, + sources: Vec, + }, +} + +impl CompatProfile { + fn profile_id(&self) -> &str { + match self { + Self::Baseline { .. } => "baseline", + Self::Overlay { overlay, .. } => overlay.profile_id(), + } + } + + fn cases(&self) -> &[CompatCase] { + match self { + Self::Baseline { cases } | Self::Overlay { cases, .. } => cases, + } + } + + fn case_names(&self) -> Vec<&str> { + self.cases() + .iter() + .map(|case| case.metadata.name.as_str()) + .collect() + } + + fn sources(&self) -> &[PathBuf] { + match self { + Self::Baseline { .. } => &[], + Self::Overlay { sources, .. } => sources, + } + } + + fn overlay(&self) -> Option> { + match self { + Self::Baseline { .. } => None, + Self::Overlay { overlay, .. } => Some(Arc::clone(overlay)), + } + } + + fn touched_protected_paths(&self) -> &[crate::cmd::datanode_overlay::DottedPath] { + match self { + Self::Baseline { .. } => &[], + Self::Overlay { overlay, .. } => overlay.touched_protected_paths(), + } + } +} + +/// Deterministically ordered compatibility profiles. +#[derive(Debug)] +struct CompatProfiles(Vec); + +impl CompatProfiles { + fn iter(&self) -> impl Iterator { + self.0.iter() + } + + fn cases(&self) -> impl Iterator { + self.0.iter().flat_map(CompatProfile::cases) + } + + fn case_count(&self) -> usize { + self.0.iter().map(|profile| profile.cases().len()).sum() + } +} + +#[derive(Debug)] +struct OverlayProfileGroup { + overlay: Arc, + cases: Vec, + sources: Vec, +} + +fn prepare_compat_profiles( + cases: Vec, + wal_config: &WalConfig, +) -> Result { + let protection = DatanodeProtectionPolicy::for_wal(wal_config); + let mut baseline_cases = Vec::new(); + let mut overlays: BTreeMap<[u8; 32], OverlayProfileGroup> = BTreeMap::new(); + + for case in cases { + let Some(reference) = case.metadata.old_datanode_overlay() else { + baseline_cases.push(case); + continue; + }; + let overlay = DatanodeOverlay::load(&case.dir, reference).map_err(|error| { + format!( + "Failed to load old datanode overlay for compatibility case '{}': {error}", + case.metadata.name + ) + })?; + let prepared = Arc::new(overlay.prepare(&protection).map_err(|error| { + format!( + "Failed to prepare old datanode overlay for compatibility case '{}': {error}", + case.metadata.name + ) + })?); + let key = *prepared.profile_key(); + let source = prepared.source().to_path_buf(); + let entry = overlays.entry(key).or_insert_with(|| OverlayProfileGroup { + overlay: Arc::clone(&prepared), + cases: Vec::new(), + sources: Vec::new(), + }); + entry.cases.push(case); + entry.sources.push(source); + } + + baseline_cases.sort_by(|left, right| left.metadata.name.cmp(&right.metadata.name)); + let mut profiles = Vec::new(); + if !baseline_cases.is_empty() { + profiles.push(CompatProfile::Baseline { + cases: baseline_cases, + }); + } + for ( + _, + OverlayProfileGroup { + overlay, + mut cases, + mut sources, + }, + ) in overlays + { + cases.sort_by(|left, right| left.metadata.name.cmp(&right.metadata.name)); + sources.sort(); + sources.dedup(); + profiles.push(CompatProfile::Overlay { + overlay, + cases, + sources, + }); + } + + Ok(CompatProfiles(profiles)) +} + +fn emit_protected_path_warnings(profiles: &CompatProfiles) { + for profile in profiles.iter() { + let paths: Vec<_> = profile + .touched_protected_paths() + .iter() + .map(ToString::to_string) + .collect(); + if paths.is_empty() { + continue; + } + println!( + "{}", + format_protected_path_warning(profile.profile_id(), profile.case_names(), paths) + ); + } +} + +fn format_protected_path_warning( + profile_id: &str, + mut case_names: Vec<&str>, + mut paths: Vec, +) -> String { + case_names.sort_unstable(); + paths.sort_unstable(); + format!( + "Warning: datanode overlay profile {profile_id} touches runner-owned paths [{}] for cases [{}]", + paths.join(", "), + case_names.join(", ") + ) +} + +#[derive(Debug)] +struct ProfileOutcome { + failures: Vec, + stop_remaining_profiles: bool, + preserved_state: Option, +} + +/// Pure compatibility profile progress policy, independent of process ownership. +#[derive(Default)] +struct ProfileProgress { + successful_setup_indexes: Vec, + failures: Vec, + fail_fast_setup_failure: bool, + fail_fast_verify_failure: bool, + cleanup_failed: bool, +} + +impl ProfileProgress { + fn record_setup_success(&mut self, case_index: usize) { + self.successful_setup_indexes.push(case_index); + } + + fn record_setup_failure(&mut self, failure: String, fail_fast: bool) -> bool { + self.failures.push(failure); + self.fail_fast_setup_failure = fail_fast; + fail_fast + } + + fn record_verify_failure(&mut self, failure: String, fail_fast: bool) -> bool { + self.failures.push(failure); + self.fail_fast_verify_failure = fail_fast; + fail_fast + } + + fn record_cleanup_failure(&mut self, failure: String) { + self.failures.push(failure); + self.cleanup_failed = true; + } + + fn should_transition_to_current(&self) -> bool { + !self.successful_setup_indexes.is_empty() && !self.fail_fast_setup_failure + } + + fn should_stop_remaining_profiles(&self) -> bool { + self.fail_fast_setup_failure || self.fail_fast_verify_failure || self.cleanup_failed + } +} + +struct ProfileRunConfig<'a> { + interceptor_registry: &'a Registry, + from_bins_dir: Option, + to_bins_dir: PathBuf, + pull_version_on_need: bool, + setup_etcd: bool, + fail_fast: bool, + topology: CompatTopology, + preserve_state: bool, + wal_config: WalConfig, +} + +async fn run_profile(profile: &CompatProfile, config: &ProfileRunConfig<'_>) -> ProfileOutcome { + println!("Running compatibility profile {}", profile.profile_id()); + let temp_dir = tempfile::Builder::new() + .prefix(&format!("sqlness-compat-{}-", profile.profile_id())) + .tempdir() + .unwrap(); + let profile_state = ProfileStateGuard::new(temp_dir, config.preserve_state); + let sqlness_home = profile_state.path().to_path_buf(); + // Standalone runs need no etcd. + let setup_etcd = config.topology == CompatTopology::Distributed && config.setup_etcd; + unsafe { + std::env::set_var( + "SQLNESS_HOME", + sqlness_home.join("copy").display().to_string(), + ); + } + + let store_config = StoreConfig { + store_addrs: setup_etcd + .then(|| "127.0.0.1:2379".to_string()) + .into_iter() + .collect(), + setup_etcd, + setup_pg: None, + setup_mysql: None, + enable_flat_format: false, + enable_gc: false, + }; + let env = Env::new( + sqlness_home.clone(), + ServerAddr::default(), + config.wal_config.clone(), + config.pull_version_on_need, + config.from_bins_dir.clone(), + store_config, + vec![], + ); + if let Some(overlay) = profile.overlay() { + env.activate_compat_old(overlay); + } + + // Arm immediately before starting the profile so earlier preflight failures + // cannot remove a developer-owned etcd container. + let etcd_guard = setup_etcd.then(EtcdGuard::new); + println!( + "Starting old-version {} cluster...", + config.topology.as_str() + ); + let db = match config.topology { + CompatTopology::Distributed => env.compat_start_distributed(0).await, + CompatTopology::Standalone => env.compat_start_standalone(0).await, + }; + + let mut progress = ProfileProgress::default(); + for (case_index, case) in profile.cases().iter().enumerate() { + match run_compat_phase(&db, case, config.interceptor_registry, CompatPhase::Setup).await { + Ok(()) => { + println!(" Setup: {} - OK", case.metadata.name); + progress.record_setup_success(case_index); + } + Err(error) => { + println!(" Setup: {} - FAILED: {error}", case.metadata.name); + if progress.record_setup_failure( + format!("{} (setup): {error}", case.metadata.name), + config.fail_fast, + ) { + break; + } + } + } + } + + if progress.should_transition_to_current() { + println!("Restarting cluster with new-version binary on preserved state..."); + env.activate_compat_current(); + env.compat_restart(&db, config.to_bins_dir.clone()).await; - // ---- 10. Run verify phase on the to-version cluster ---- println!("Running verify phase..."); - let mut failed = Vec::new(); - for case in &cases { - match run_compat_phase(&db, case, &interceptor_registry, CompatPhase::Verify).await { + for case_index in progress.successful_setup_indexes.clone() { + let case = &profile.cases()[case_index]; + match run_compat_phase(&db, case, config.interceptor_registry, CompatPhase::Verify) + .await + { Ok(()) => println!(" Verify: {} - PASSED", case.metadata.name), - Err(e) => { - println!(" Verify: {} - FAILED: {e}", case.metadata.name); - failed.push(case.metadata.name.clone()); - if self.fail_fast { + Err(error) => { + println!(" Verify: {} - FAILED: {error}", case.metadata.name); + if progress.record_verify_failure( + format!("{} (verify): {error}", case.metadata.name), + config.fail_fast, + ) { break; } } } } + } else if progress.successful_setup_indexes.is_empty() { + println!("Skipping current-version restart: no setup cases succeeded"); + } - // ---- 11. Stop cluster ---- - db.compat_stop(); + let cleanup = cleanup_profile(db, etcd_guard, setup_etcd, profile_state).await; + for failure in cleanup.failures { + progress.record_cleanup_failure(failure); + } + let stop_remaining_profiles = progress.should_stop_remaining_profiles(); + ProfileOutcome { + failures: progress.failures, + stop_remaining_profiles, + preserved_state: cleanup.preserved_state, + } +} - // ---- 12. Cleanup ---- - // Etcd is always cleaned up; --preserve-state only preserves sqlness_home. - if setup_etcd { - println!("Stopping etcd"); - util::stop_rm_etcd(); +struct CleanupOutcome { + failures: Vec, + preserved_state: Option, +} + +/// Owns one profile state directory until normal cleanup or unwind finalization. +struct ProfileStateGuard { + temp_dir: Option, + preserve_state: bool, +} + +impl ProfileStateGuard { + fn new(temp_dir: tempfile::TempDir, preserve_state: bool) -> Self { + Self { + temp_dir: Some(temp_dir), + preserve_state, + } + } + + fn path(&self) -> &std::path::Path { + self.temp_dir.as_ref().unwrap().path() + } + + fn finalize(mut self) -> CleanupOutcome { + let temp_dir = self.temp_dir.take().unwrap(); + let mut failures = Vec::new(); + if self.preserve_state { + let path = temp_dir.keep(); + if let Err(error) = std::fs::metadata(&path) { + failures.push(format!( + "failed to confirm preserved profile state {}: {error}", + path.display() + )); + } + return CleanupOutcome { + failures, + preserved_state: Some(path), + }; } - if !self.preserve_state { - println!("Removing state in {:?}", sqlness_home); - tokio::fs::remove_dir_all(sqlness_home) - .await - .unwrap_or_else(|e| println!("Warning: failed to clean up temp dir: {e}")); + println!("Removing state in {}", temp_dir.path().display()); + if let Err(error) = temp_dir.close() { + failures.push(format!("failed to remove profile state: {error}")); } + CleanupOutcome { + failures, + preserved_state: None, + } + } +} - // Disarm the etcd guard now that we've done normal cleanup. - if let Some(mut guard) = etcd_guard.take() { - guard.disarm(); +impl Drop for ProfileStateGuard { + fn drop(&mut self) { + let Some(temp_dir) = self.temp_dir.take() else { + return; + }; + if self.preserve_state { + let path = temp_dir.keep(); + println!( + "Preserved compat profile state after abnormal exit: {}", + path.display() + ); } + } +} - if failed.is_empty() { - println!("\n\x1b[32mAll compat tests passed!\x1b[0m"); - } else { - println!("\n\x1b[31mFailed cases: {}\x1b[0m", failed.join(", ")); - // Explicitly drop the guard before exit so it doesn't double-cleanup. - std::process::exit(1); +async fn cleanup_profile( + mut db: GreptimeDB, + mut etcd_guard: Option, + setup_etcd: bool, + profile_state: ProfileStateGuard, +) -> CleanupOutcome { + db.compat_stop(); + drop(db); + let mut failures = Vec::new(); + if setup_etcd { + println!("Stopping etcd"); + match util::stop_rm_etcd_checked() { + Ok(()) => { + if let Some(guard) = etcd_guard.as_mut() { + guard.disarm(); + } + } + Err(error) => failures.push(format!("profile etcd cleanup: {error}")), } } + // On failed checked cleanup this Drop performs one best-effort retry before + // profile state is finalized. + drop(etcd_guard); + + let mut state_outcome = profile_state.finalize(); + failures.append(&mut state_outcome.failures); + CleanupOutcome { + failures, + preserved_state: state_outcome.preserved_state, + } } /// Guard that stops/removes Docker etcd on drop (panic or early exit). @@ -546,13 +937,9 @@ impl Drop for EtcdGuard { fn drop(&mut self) { if self.active { println!("EtcdGuard: emergency etcd cleanup (panic or early exit)"); - // Best-effort: don't panic in Drop - let _ = std::process::Command::new("docker") - .args(["container", "stop", "etcd"]) - .status(); - let _ = std::process::Command::new("docker") - .args(["container", "rm", "etcd"]) - .status(); + if let Err(error) = util::stop_rm_etcd_checked() { + println!("EtcdGuard: emergency etcd cleanup failed: {error}"); + } } } } @@ -942,7 +1329,36 @@ fn simple_diff(expected: &str, actual: &str) -> String { #[cfg(test)] mod tests { - use super::trim_trailing_blank_lines; + use std::path::Path; + + use super::*; + use crate::cmd::compat_case::{CaseMetadata, OldConfigMetadata}; + + fn test_case(temp_dir: &Path, name: &str, overlay: Option<&str>) -> CompatCase { + let case_dir = temp_dir.join(name); + std::fs::create_dir_all(&case_dir).unwrap(); + if let Some(overlay) = overlay { + std::fs::write(case_dir.join("overlay.toml"), overlay).unwrap(); + } + CompatCase { + metadata: CaseMetadata { + name: name.to_string(), + reason: "test".to_string(), + introduced_by: "test".to_string(), + topologies: vec![CompatTopology::Distributed.as_str().to_string()], + from_range: vec!["*".to_string()], + to_range: vec!["*".to_string()], + features: vec!["table".to_string()], + owner: "test".to_string(), + namespace: None, + old_config: overlay.map(|_| OldConfigMetadata { + datanode: PathBuf::from("overlay.toml"), + }), + }, + dir: case_dir, + namespace: name.to_string(), + } + } #[test] fn test_trim_trailing_blank_lines_preserves_single_final_newline() { @@ -951,4 +1367,139 @@ mod tests { assert_eq!(output, "SELECT 1;\n\n+---+\n"); } + + #[test] + fn profiles_are_baseline_first_and_group_by_full_semantic_digest() { + let temp_dir = tempfile::tempdir().unwrap(); + let profiles = prepare_compat_profiles( + vec![ + test_case(temp_dir.path(), "overlay_b", Some("[x]\nb = 2\na = 1\n")), + test_case(temp_dir.path(), "baseline", None), + test_case(temp_dir.path(), "overlay_a", Some("[x]\na = 1\nb = 2\n")), + test_case(temp_dir.path(), "overlay_c", Some("value = 3\n")), + ], + &WalConfig::RaftEngine, + ) + .unwrap(); + let profiles: Vec<_> = profiles.iter().collect(); + + assert!(matches!(profiles[0], CompatProfile::Baseline { .. })); + assert_eq!(profiles[0].case_names(), ["baseline"]); + assert_eq!(profiles[1].case_names(), ["overlay_a", "overlay_b"]); + assert_eq!(profiles[1].sources().len(), 2); + let overlay_keys: Vec<_> = profiles[1..] + .iter() + .map(|profile| match profile { + CompatProfile::Overlay { overlay, .. } => *overlay.profile_key(), + CompatProfile::Baseline { .. } => unreachable!(), + }) + .collect(); + assert!(overlay_keys.windows(2).all(|keys| keys[0] < keys[1])); + } + + #[test] + fn non_fail_fast_setup_verifies_only_successful_cases() { + let mut progress = ProfileProgress::default(); + progress.record_setup_success(0); + assert!(!progress.record_setup_failure("case_b (setup): failed".to_string(), false)); + progress.record_setup_success(2); + + assert!(progress.should_transition_to_current()); + assert_eq!(progress.successful_setup_indexes, [0, 2]); + assert!(!progress.should_stop_remaining_profiles()); + } + + #[test] + fn fail_fast_setup_blocks_current_and_later_profiles() { + let mut progress = ProfileProgress::default(); + progress.record_setup_success(0); + + assert!(progress.record_setup_failure("case_b (setup): failed".to_string(), true)); + assert!(!progress.should_transition_to_current()); + assert!(progress.should_stop_remaining_profiles()); + } + + #[test] + fn zero_successful_setups_skips_current() { + let mut progress = ProfileProgress::default(); + progress.record_setup_failure("case_a (setup): failed".to_string(), false); + + assert!(!progress.should_transition_to_current()); + assert!(!progress.should_stop_remaining_profiles()); + } + + #[test] + fn verify_fail_fast_stops_later_profiles_after_cleanup() { + let mut progress = ProfileProgress::default(); + progress.record_setup_success(0); + + assert!(progress.record_verify_failure("case_a (verify): failed".to_string(), true)); + assert!(progress.should_stop_remaining_profiles()); + } + + #[test] + fn non_fail_fast_aggregates_failures_but_cleanup_failure_stops_profiles() { + let mut progress = ProfileProgress::default(); + progress.record_setup_failure("case_a (setup): failed".to_string(), false); + progress.record_setup_success(1); + progress.record_verify_failure("case_b (verify): failed".to_string(), false); + + assert!(!progress.should_stop_remaining_profiles()); + progress.record_cleanup_failure("failed to remove profile state".to_string()); + assert_eq!(progress.failures.len(), 3); + assert!(progress.should_stop_remaining_profiles()); + } + + #[test] + fn protected_path_warning_is_sorted_and_value_free() { + let warning = format_protected_path_warning( + "abcdef123456", + vec!["case_z", "case_a"], + vec!["wal.provider".to_string(), "mode".to_string()], + ); + + assert_eq!( + warning, + "Warning: datanode overlay profile abcdef123456 touches runner-owned paths [mode, wal.provider] for cases [case_a, case_z]" + ); + assert!(!warning.contains("secret-overlay-value")); + } + + #[test] + fn profile_state_guard_preserves_state_on_forced_unwind() { + let temp_dir = tempfile::tempdir().unwrap(); + let path = temp_dir.path().to_path_buf(); + + let unwind = std::panic::catch_unwind(|| { + let _state = ProfileStateGuard::new(temp_dir, true); + panic!("forced profile unwind"); + }); + + assert!(unwind.is_err()); + assert!(path.is_dir()); + std::fs::remove_dir_all(path).unwrap(); + } + + #[test] + fn profile_state_guard_removes_state_on_drop_without_preservation() { + let temp_dir = tempfile::tempdir().unwrap(); + let path = temp_dir.path().to_path_buf(); + + drop(ProfileStateGuard::new(temp_dir, false)); + + assert!(!path.exists()); + } + + #[test] + fn profile_state_guard_normal_finalization_transfers_ownership_once() { + let temp_dir = tempfile::tempdir().unwrap(); + let path = temp_dir.path().to_path_buf(); + + let cleanup = ProfileStateGuard::new(temp_dir, true).finalize(); + + assert!(cleanup.failures.is_empty()); + assert_eq!(cleanup.preserved_state.as_deref(), Some(path.as_path())); + assert!(path.is_dir()); + std::fs::remove_dir_all(path).unwrap(); + } } diff --git a/tests/runner/src/cmd/compat_case.rs b/tests/runner/src/cmd/compat_case.rs index bc9065c4492..7c6dbbfa740 100644 --- a/tests/runner/src/cmd/compat_case.rs +++ b/tests/runner/src/cmd/compat_case.rs @@ -43,6 +43,17 @@ pub struct CaseMetadata { /// Must match `[a-z0-9_]+`. #[serde(default)] pub namespace: Option, + /// Optional old-stage server configuration sidecars. + #[serde(default)] + pub old_config: Option, +} + +/// Optional configuration sidecars for the old compatibility stage. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct OldConfigMetadata { + /// Datanode TOML sidecar, resolved relative to the compatibility case directory. + pub datanode: PathBuf, } impl CaseMetadata { @@ -53,6 +64,13 @@ impl CaseMetadata { .clone() .unwrap_or_else(|| sanitize_namespace(case_dir_name)) } + + /// Returns the old-stage datanode sidecar reference, if configured. + pub fn old_datanode_overlay(&self) -> Option<&Path> { + self.old_config + .as_ref() + .map(|config| config.datanode.as_path()) + } } /// A loaded compatibility case (metadata + file paths). @@ -469,6 +487,74 @@ mod tests { assert_eq!(sanitize_namespace("a.b-c"), "a_b_c"); } + #[test] + fn test_case_metadata_parses_old_datanode_overlay() { + let metadata: CaseMetadata = toml::from_str( + r#" + name = "case" + reason = "reason" + introduced_by = "test" + topologies = ["distributed"] + from_range = ["*"] + to_range = ["*"] + features = ["table"] + owner = "team" + + [old_config] + datanode = "old-datanode.overlay.toml" + "#, + ) + .unwrap(); + + assert_eq!( + metadata.old_datanode_overlay(), + Some(Path::new("old-datanode.overlay.toml")) + ); + } + + #[test] + fn test_case_metadata_rejects_empty_old_config() { + let error = toml::from_str::( + r#" + name = "case" + reason = "reason" + introduced_by = "test" + topologies = ["distributed"] + from_range = ["*"] + to_range = ["*"] + features = ["table"] + owner = "team" + + [old_config] + "#, + ) + .unwrap_err(); + + assert!(error.to_string().contains("missing field `datanode`")); + } + + #[test] + fn test_case_metadata_rejects_unknown_old_config_fields() { + let error = toml::from_str::( + r#" + name = "case" + reason = "reason" + introduced_by = "test" + topologies = ["distributed"] + from_range = ["*"] + to_range = ["*"] + features = ["table"] + owner = "team" + + [old_config] + unsupported = "value" + "#, + ) + .unwrap_err(); + + assert!(error.to_string().contains("unknown field `unsupported`")); + } + #[test] fn test_validate_cases_metadata_rejects_empty_required_vectors() { let case = CompatCase { @@ -482,6 +568,7 @@ mod tests { features: vec!["table".to_string()], owner: "team".to_string(), namespace: None, + old_config: None, }, dir: PathBuf::from("case"), namespace: "case".to_string(), @@ -503,6 +590,7 @@ mod tests { features: vec!["table".to_string()], owner: "test".to_string(), namespace: None, + old_config: None, }, dir: PathBuf::from("bad_constraint"), namespace: "bad_constraint".to_string(), @@ -524,6 +612,7 @@ mod tests { features: vec!["table".to_string()], owner: "test".to_string(), namespace: None, + old_config: None, }, dir: PathBuf::from("case_a"), namespace: "shared_name".to_string(), @@ -539,6 +628,7 @@ mod tests { features: vec!["table".to_string()], owner: "test".to_string(), namespace: None, + old_config: None, }, dir: PathBuf::from("case_b"), namespace: "shared_name".to_string(), @@ -563,6 +653,7 @@ mod tests { features: vec!["table".to_string()], owner: "test".to_string(), namespace: Some("shared_name".to_string()), + old_config: None, }, dir: PathBuf::from("case_a"), namespace: "shared_name".to_string(), @@ -578,6 +669,7 @@ mod tests { features: vec!["table".to_string()], owner: "test".to_string(), namespace: Some("shared_name".to_string()), + old_config: None, }, dir: PathBuf::from("case_b"), namespace: "shared_name".to_string(), @@ -599,6 +691,7 @@ mod tests { features: vec!["table".to_string()], owner: "team".to_string(), namespace: None, + old_config: None, }, dir: PathBuf::from("case"), namespace: "case".to_string(), diff --git a/tests/runner/src/cmd/datanode_overlay.rs b/tests/runner/src/cmd/datanode_overlay.rs new file mode 100644 index 00000000000..21892b9a091 --- /dev/null +++ b/tests/runner/src/cmd/datanode_overlay.rs @@ -0,0 +1,637 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fs::File; +use std::io::Read; +use std::path::{Component, Path, PathBuf}; + +use sha2::{Digest, Sha256}; +use toml::value::Table; + +use crate::env::bare::WalConfig; +/// A dotted TOML path owned by the compatibility runner. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) struct DottedPath(Vec); + +impl DottedPath { + fn new(parts: &[&str]) -> Self { + Self(parts.iter().map(|part| (*part).to_string()).collect()) + } + + fn parts(&self) -> impl Iterator { + self.0.iter().map(String::as_str) + } +} + +impl std::fmt::Display for DottedPath { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.0.join(".")) + } +} + +/// Runner-owned datanode configuration paths that an overlay cannot override. +#[derive(Debug, Clone)] +pub(crate) struct DatanodeProtectionPolicy { + protected_paths: Vec, +} + +impl DatanodeProtectionPolicy { + /// Builds the policy for the runner-selected WAL configuration. + pub(crate) fn for_wal(wal: &WalConfig) -> Self { + let mut protected_paths = vec![ + DottedPath::new(&["mode"]), + DottedPath::new(&["node_id"]), + DottedPath::new(&["storage", "data_home"]), + DottedPath::new(&["meta_client_options", "metasrv_addrs"]), + DottedPath::new(&["wal", "provider"]), + ]; + protected_paths.push(match wal { + WalConfig::RaftEngine => DottedPath::new(&["wal", "dir"]), + WalConfig::Kafka { .. } => DottedPath::new(&["wal", "broker_endpoints"]), + }); + protected_paths.sort(); + + Self { protected_paths } + } +} + +/// A parsed old-stage datanode sidecar, loaded from a confined case directory. +#[derive(Debug, Clone)] +pub(crate) struct DatanodeOverlay { + source: PathBuf, + value: toml::Value, + profile_key: [u8; 32], + profile_id: String, +} + +impl DatanodeOverlay { + /// Loads and parses a datanode sidecar referenced relative to `case_dir`. + /// + /// This is intended to prevent accidental case-directory escapes; it is not + /// an adversarial concurrent-filesystem security boundary. + pub(crate) fn load(case_dir: &Path, relative_ref: &Path) -> Result { + validate_relative_reference(relative_ref)?; + + let canonical_case_dir = case_dir.canonicalize().map_err(|error| { + format!( + "Failed to canonicalize compatibility case directory {}: {error}", + case_dir.display() + ) + })?; + let requested_path = case_dir.join(relative_ref); + let canonical_target = requested_path.canonicalize().map_err(|error| { + format!( + "Failed to resolve datanode overlay {} relative to case directory {}: {error}", + relative_ref.display(), + case_dir.display() + ) + })?; + + if !canonical_target.starts_with(&canonical_case_dir) { + return Err(format!( + "Datanode overlay {} resolves outside compatibility case directory {}", + relative_ref.display(), + canonical_case_dir.display() + )); + } + + let mut file = File::open(&canonical_target).map_err(|error| { + format!( + "Failed to open datanode overlay {}: {error}", + canonical_target.display() + ) + })?; + let metadata = file.metadata().map_err(|error| { + format!( + "Failed to inspect datanode overlay {}: {error}", + canonical_target.display() + ) + })?; + if !metadata.is_file() { + return Err(format!( + "Datanode overlay {} must be a regular file", + canonical_target.display() + )); + } + + let mut content = String::new(); + file.read_to_string(&mut content).map_err(|error| { + format!( + "Failed to read datanode overlay {}: {error}", + canonical_target.display() + ) + })?; + let value: toml::Value = toml::from_str(&content).map_err(|error| { + format!( + "Failed to parse datanode overlay {}: {error}", + canonical_target.display() + ) + })?; + if !value.is_table() { + return Err(format!( + "Datanode overlay {} must have a TOML table at its root", + canonical_target.display() + )); + } + + let profile_key = semantic_profile_key(&value); + let profile_id = hex::encode(profile_key)[..12].to_string(); + + Ok(Self { + source: canonical_target, + value, + profile_key, + profile_id, + }) + } + + /// Validates runner-owned path conflicts and records overridden protected paths. + pub(crate) fn prepare( + self, + protection: &DatanodeProtectionPolicy, + ) -> Result { + let mut touched_protected_paths = Vec::new(); + for path in &protection.protected_paths { + validate_protected_ancestors(&self.value, path).map_err(|error| { + format!( + "Datanode overlay {} (profile {}): {error}", + self.source.display(), + self.profile_id + ) + })?; + if value_at_path(&self.value, path).is_some() { + touched_protected_paths.push(path.clone()); + } + } + touched_protected_paths.sort(); + + Ok(PreparedDatanodeOverlay { + source: self.source, + value: self.value, + profile_key: self.profile_key, + profile_id: self.profile_id, + protected_paths: protection.protected_paths.clone(), + touched_protected_paths, + }) + } +} + +/// A validated datanode overlay ready for application to a rendered baseline. +#[derive(Debug, Clone)] +pub(crate) struct PreparedDatanodeOverlay { + source: PathBuf, + value: toml::Value, + profile_key: [u8; 32], + profile_id: String, + protected_paths: Vec, + touched_protected_paths: Vec, +} + +impl PreparedDatanodeOverlay { + /// Returns the canonical full SHA-256 profile key used for grouping. + pub(crate) fn profile_key(&self) -> &[u8; 32] { + &self.profile_key + } + + /// Returns a truncated profile ID suitable only for diagnostics. + pub(crate) fn profile_id(&self) -> &str { + &self.profile_id + } + + /// Returns the canonical sidecar path used for diagnostics. + pub(crate) fn source(&self) -> &Path { + &self.source + } + + /// Returns protected paths declared by the sidecar, in dotted-path order. + pub(crate) fn touched_protected_paths(&self) -> &[DottedPath] { + &self.touched_protected_paths + } + + /// Merges the sidecar into a rendered baseline and restores runner-owned fields. + pub(crate) fn apply_to_rendered_baseline(&self, baseline: &str) -> Result { + let baseline_value: toml::Value = toml::from_str(baseline) + .map_err(|error| format!("Failed to parse rendered datanode baseline: {error}"))?; + if !baseline_value.is_table() { + return Err( + "Rendered datanode baseline must have a TOML table at its root".to_string(), + ); + } + + let mut merged = baseline_value.clone(); + merge_toml_values(&mut merged, &self.value); + + for path in &self.protected_paths { + restore_protected_path(&mut merged, &baseline_value, path)?; + } + + toml::to_string(&merged) + .map_err(|error| format!("Failed to serialize merged datanode configuration: {error}")) + } +} + +fn validate_relative_reference(relative_ref: &Path) -> Result<(), String> { + if relative_ref.as_os_str().is_empty() + || relative_ref + .components() + .any(|component| !matches!(component, Component::Normal(_))) + { + return Err(format!( + "Datanode overlay reference {} must be a non-empty relative path with normal components only", + relative_ref.display() + )); + } + Ok(()) +} + +fn validate_protected_ancestors(value: &toml::Value, path: &DottedPath) -> Result<(), String> { + let mut current = value; + let parts: Vec<_> = path.parts().collect(); + for (index, part) in parts.iter().enumerate().take(parts.len().saturating_sub(1)) { + let Some(table) = current.as_table() else { + return Err(format!( + "Datanode overlay makes ancestor {} of protected path {} non-table", + parts[..index].join("."), + path + )); + }; + let Some(next) = table.get(*part) else { + return Ok(()); + }; + if !next.is_table() { + return Err(format!( + "Datanode overlay makes ancestor {} of protected path {} non-table", + parts[..=index].join("."), + path + )); + } + current = next; + } + Ok(()) +} + +fn value_at_path<'a>(value: &'a toml::Value, path: &DottedPath) -> Option<&'a toml::Value> { + let mut current = value; + for part in path.parts() { + current = current.as_table()?.get(part)?; + } + Some(current) +} + +fn merge_toml_values(baseline: &mut toml::Value, overlay: &toml::Value) { + match (baseline, overlay) { + (toml::Value::Table(baseline), toml::Value::Table(overlay)) => { + for (key, overlay_value) in overlay { + match baseline.get_mut(key) { + Some(baseline_value) => merge_toml_values(baseline_value, overlay_value), + None => { + baseline.insert(key.clone(), overlay_value.clone()); + } + } + } + } + (baseline, overlay) => *baseline = overlay.clone(), + } +} + +fn restore_protected_path( + merged: &mut toml::Value, + baseline: &toml::Value, + path: &DottedPath, +) -> Result<(), String> { + let Some(merged_table) = merged.as_table_mut() else { + return Err("Merged datanode configuration must have a TOML table at its root".to_string()); + }; + let baseline_value = value_at_path(baseline, path).cloned(); + restore_path_in_table(merged_table, &path.0, baseline_value); + Ok(()) +} + +fn restore_path_in_table(table: &mut Table, parts: &[String], baseline_value: Option) { + let Some((part, remaining)) = parts.split_first() else { + return; + }; + if remaining.is_empty() { + match baseline_value { + Some(value) => { + table.insert(part.clone(), value); + } + None => { + table.remove(part); + } + } + return; + } + + let Some(value) = table.get_mut(part) else { + return; + }; + if let Some(table) = value.as_table_mut() { + restore_path_in_table(table, remaining, baseline_value); + } +} + +fn semantic_profile_key(value: &toml::Value) -> [u8; 32] { + let mut encoded = Vec::new(); + encode_semantic_value(&mut encoded, value); + Sha256::digest(encoded).into() +} + +fn encode_semantic_value(output: &mut Vec, value: &toml::Value) { + match value { + toml::Value::String(value) => encode_bytes(output, b"string", value.as_bytes()), + toml::Value::Integer(value) => { + encode_bytes(output, b"integer", &value.to_be_bytes()); + } + toml::Value::Float(value) => { + encode_bytes(output, b"float", &value.to_bits().to_be_bytes()); + } + toml::Value::Boolean(value) => { + encode_bytes(output, b"boolean", &[u8::from(*value)]); + } + toml::Value::Datetime(value) => { + encode_bytes(output, b"datetime", value.to_string().as_bytes()) + } + toml::Value::Array(values) => { + let mut payload = Vec::new(); + encode_length(&mut payload, values.len()); + for value in values { + encode_semantic_value(&mut payload, value); + } + encode_bytes(output, b"array", &payload); + } + toml::Value::Table(table) => { + let mut payload = Vec::new(); + encode_length(&mut payload, table.len()); + let mut entries: Vec<_> = table.iter().collect(); + entries.sort_unstable_by(|(left, _), (right, _)| left.as_bytes().cmp(right.as_bytes())); + for (key, value) in entries { + encode_bytes(&mut payload, b"key", key.as_bytes()); + encode_semantic_value(&mut payload, value); + } + encode_bytes(output, b"table", &payload); + } + } +} + +fn encode_bytes(output: &mut Vec, tag: &[u8], value: &[u8]) { + encode_length(output, tag.len()); + output.extend_from_slice(tag); + encode_length(output, value.len()); + output.extend_from_slice(value); +} + +fn encode_length(output: &mut Vec, length: usize) { + output.extend_from_slice(&(length as u64).to_be_bytes()); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn load_overlay(content: &str) -> DatanodeOverlay { + let temp_dir = tempfile::tempdir().unwrap(); + let path = temp_dir.path().join("overlay.toml"); + std::fs::write(&path, content).unwrap(); + let overlay = DatanodeOverlay::load(temp_dir.path(), Path::new("overlay.toml")).unwrap(); + // `load` owns the parsed TOML, so the temporary source need not outlive it. + overlay + } + + fn prepared_overlay(content: &str, wal: &WalConfig) -> PreparedDatanodeOverlay { + load_overlay(content) + .prepare(&DatanodeProtectionPolicy::for_wal(wal)) + .unwrap() + } + + #[test] + fn merges_recursive_tables_and_replaces_non_tables_atomically() { + let prepared = prepared_overlay( + r#" + scalar = "replacement" + primitive_array = [3, 2, 1] + array_of_tables = [{ name = "new" }] + [nested] + new_key = true + old_key = "replacement" + "#, + &WalConfig::RaftEngine, + ); + let merged = prepared + .apply_to_rendered_baseline( + r#" + scalar = 1 + primitive_array = [1, 2] + array_of_tables = [{ name = "old" }, { name = "older" }] + untouched = "kept" + [nested] + old_key = 1 + retained = "kept" + "#, + ) + .unwrap(); + let value: toml::Value = toml::from_str(&merged).unwrap(); + + assert_eq!(value["scalar"].as_str(), Some("replacement")); + assert_eq!(value["primitive_array"].as_array().unwrap().len(), 3); + assert_eq!(value["array_of_tables"].as_array().unwrap().len(), 1); + assert_eq!(value["nested"]["old_key"].as_str(), Some("replacement")); + assert_eq!(value["nested"]["new_key"].as_bool(), Some(true)); + assert_eq!(value["nested"]["retained"].as_str(), Some("kept")); + assert_eq!(value["untouched"].as_str(), Some("kept")); + } + + #[test] + fn restores_protected_paths_and_deletes_absent_raft_wal_dir() { + let prepared = prepared_overlay( + r#" + mode = "standalone" + node_id = 99 + [storage] + data_home = "/overlay/data" + [meta_client_options] + metasrv_addrs = ["overlay:3002"] + [wal] + provider = "kafka" + dir = "/overlay/wal" + tuning = 42 + "#, + &WalConfig::RaftEngine, + ); + let merged = prepared + .apply_to_rendered_baseline( + r#" + mode = "distributed" + node_id = 1 + [storage] + data_home = "/runner/data" + [meta_client_options] + metasrv_addrs = ["runner:3002"] + [wal] + provider = "raft_engine" + "#, + ) + .unwrap(); + let value: toml::Value = toml::from_str(&merged).unwrap(); + + assert_eq!(value["mode"].as_str(), Some("distributed")); + assert_eq!(value["node_id"].as_integer(), Some(1)); + assert_eq!(value["storage"]["data_home"].as_str(), Some("/runner/data")); + assert_eq!( + value["meta_client_options"]["metasrv_addrs"][0].as_str(), + Some("runner:3002") + ); + assert_eq!(value["wal"]["provider"].as_str(), Some("raft_engine")); + assert!(value["wal"].get("dir").is_none()); + assert_eq!(value["wal"]["tuning"].as_integer(), Some(42)); + } + + #[test] + fn kafka_policy_restores_broker_endpoints() { + let prepared = prepared_overlay( + r#" + [wal] + broker_endpoints = ["overlay:9092"] + provider = "raft_engine" + "#, + &WalConfig::Kafka { + needs_kafka_cluster: false, + broker_endpoints: vec![], + }, + ); + let merged = prepared + .apply_to_rendered_baseline( + r#" + [wal] + provider = "kafka" + broker_endpoints = ["runner:9092"] + "#, + ) + .unwrap(); + let value: toml::Value = toml::from_str(&merged).unwrap(); + + assert_eq!(value["wal"]["provider"].as_str(), Some("kafka")); + assert_eq!( + value["wal"]["broker_endpoints"][0].as_str(), + Some("runner:9092") + ); + } + + #[test] + fn rejects_non_table_protected_ancestor() { + let overlay = load_overlay("wal = \"not a table\""); + let error = overlay + .prepare(&DatanodeProtectionPolicy::for_wal(&WalConfig::RaftEngine)) + .unwrap_err(); + + assert!(error.contains("ancestor wal")); + assert!(error.contains("profile")); + assert!(error.contains("wal.dir") || error.contains("wal.provider")); + } + + #[test] + fn collects_sorted_touched_protected_paths() { + let prepared = prepared_overlay( + r#" + node_id = 4 + [wal] + dir = "/overlay/wal" + provider = "kafka" + "#, + &WalConfig::RaftEngine, + ); + + let paths: Vec<_> = prepared + .touched_protected_paths() + .iter() + .map(ToString::to_string) + .collect(); + assert_eq!(paths, ["node_id", "wal.dir", "wal.provider"]); + } + + #[test] + fn semantic_identity_ignores_formatting_and_table_order() { + let first = load_overlay("[settings]\nb = 2\na = \"value\"\n"); + let second = load_overlay("# comment\n[settings]\na = \"value\"\nb = 2\n"); + + assert_eq!(first.profile_key, second.profile_key); + assert_eq!(first.profile_id, second.profile_id); + } + + #[test] + fn semantic_identity_preserves_scalar_types_and_array_order() { + let integer = load_overlay("value = 1"); + let float = load_overlay("value = 1.0"); + let first_order = load_overlay("values = [1, 2]"); + let second_order = load_overlay("values = [2, 1]"); + + assert_ne!(integer.profile_key, float.profile_key); + assert_ne!(first_order.profile_key, second_order.profile_key); + } + + #[test] + fn prepared_overlay_exposes_profile_and_source_for_runner_grouping() { + let overlay = load_overlay("value = 1"); + let expected_key = overlay.profile_key; + let expected_id = overlay.profile_id.clone(); + let prepared = overlay + .prepare(&DatanodeProtectionPolicy::for_wal(&WalConfig::RaftEngine)) + .unwrap(); + + assert_eq!(prepared.profile_key(), &expected_key); + assert_eq!(prepared.profile_id(), expected_id); + assert!(prepared.source().ends_with("overlay.toml")); + } + + #[test] + fn rejects_absolute_and_traversal_references() { + let temp_dir = tempfile::tempdir().unwrap(); + let absolute = temp_dir.path().join("overlay.toml"); + std::fs::write(&absolute, "value = 1").unwrap(); + + assert!(DatanodeOverlay::load(temp_dir.path(), &absolute).is_err()); + assert!(DatanodeOverlay::load(temp_dir.path(), Path::new("../overlay.toml")).is_err()); + assert!(DatanodeOverlay::load(temp_dir.path(), Path::new("./overlay.toml")).is_err()); + } + + #[test] + fn rejects_directories_and_parse_errors() { + let temp_dir = tempfile::tempdir().unwrap(); + std::fs::create_dir(temp_dir.path().join("directory.toml")).unwrap(); + std::fs::write(temp_dir.path().join("broken.toml"), "[broken").unwrap(); + + let directory_error = + DatanodeOverlay::load(temp_dir.path(), Path::new("directory.toml")).unwrap_err(); + assert!(directory_error.contains("regular file")); + let parse_error = + DatanodeOverlay::load(temp_dir.path(), Path::new("broken.toml")).unwrap_err(); + assert!(parse_error.contains("Failed to parse datanode overlay")); + assert!(parse_error.contains("broken.toml")); + } + + #[cfg(unix)] + #[test] + fn rejects_symlink_escape() { + use std::os::unix::fs::symlink; + + let case_dir = tempfile::tempdir().unwrap(); + let outside_dir = tempfile::tempdir().unwrap(); + let outside_file = outside_dir.path().join("outside.toml"); + std::fs::write(&outside_file, "value = 1").unwrap(); + symlink(&outside_file, case_dir.path().join("escape.toml")).unwrap(); + + let error = DatanodeOverlay::load(case_dir.path(), Path::new("escape.toml")).unwrap_err(); + assert!(error.contains("outside compatibility case directory")); + } +} diff --git a/tests/runner/src/env/bare.rs b/tests/runner/src/env/bare.rs index ed248a1fe75..8fc3855edbb 100644 --- a/tests/runner/src/env/bare.rs +++ b/tests/runner/src/env/bare.rs @@ -31,6 +31,7 @@ use tokio::sync::Mutex as TokioMutex; use crate::client::MultiProtocolClient; use crate::cmd::bare::ServerAddr; use crate::cmd::compat_case::try_infer_version; +use crate::cmd::datanode_overlay::PreparedDatanodeOverlay; use crate::formatter::{ErrorFormatter, MysqlFormatter, OutputFormatter, PostgresqlFormatter}; use crate::protocol_interceptor::{MYSQL, PROTOCOL_KEY}; use crate::server_mode::{GrpcArgStyle, ServerMode}; @@ -104,6 +105,39 @@ pub struct Env { extra_args: Vec, /// Cache for the inferred gRPC argument style per `bins_dir`. grpc_arg_style_cache: Arc>>, + compat_config_stage: Arc>, +} + +/// Kills a process unless ownership has been transferred to [`GreptimeDB`]. +struct ChildGuard(Option); + +impl ChildGuard { + fn new(child: Child) -> Self { + Self(Some(child)) + } + + fn into_inner(mut self) -> Child { + self.0.take().unwrap() + } +} + +impl Drop for ChildGuard { + fn drop(&mut self) { + if let Some(child) = self.0.as_mut() { + Env::stop_server(child); + } + } +} + +/// Compatibility configuration selected for future server renders. +#[derive(Clone, Debug)] +pub(crate) enum CompatConfigStage { + /// Ordinary and baseline compatibility renders use the template unchanged. + Baseline, + /// Old compatibility renders apply this datanode overlay. + Old(Arc), + /// Current compatibility renders use the template unchanged. + Current, } #[async_trait] @@ -157,9 +191,25 @@ impl Env { store_config, extra_args, grpc_arg_style_cache: Arc::new(Mutex::new(HashMap::new())), + compat_config_stage: Arc::new(Mutex::new(CompatConfigStage::Baseline)), } } + /// Selects the old-stage overlay for subsequent compatibility renders. + pub(crate) fn activate_compat_old(&self, overlay: Arc) { + *self.compat_config_stage.lock().unwrap() = CompatConfigStage::Old(overlay); + } + + /// Selects clean current-stage rendering for subsequent compatibility renders. + pub(crate) fn activate_compat_current(&self) { + *self.compat_config_stage.lock().unwrap() = CompatConfigStage::Current; + } + + /// Takes a cheap compatibility-stage snapshot before rendering or spawning. + pub(crate) fn compat_config_stage(&self) -> CompatConfigStage { + self.compat_config_stage.lock().unwrap().clone() + } + async fn start_standalone(&self, id: usize) -> GreptimeDB { println!("Starting standalone instance id: {id}"); @@ -176,7 +226,8 @@ impl Env { let server_process = self.start_server(server_mode, &db_ctx, id, true).await; let mut greptimedb = self.connect_db(&server_addr, id).await; - greptimedb.server_processes = Some(Arc::new(Mutex::new(vec![server_process]))); + greptimedb.server_processes = + Some(Arc::new(Mutex::new(vec![server_process.into_inner()]))); greptimedb.is_standalone = true; greptimedb.ctx = db_ctx; @@ -237,10 +288,12 @@ impl Env { let mut greptimedb = self.connect_db(&server_addr, id).await; - greptimedb.metasrv_process = Some(meta_server).into(); - greptimedb.server_processes = Some(Arc::new(Mutex::new(datanodes))); - greptimedb.frontend_process = Some(frontend).into(); - greptimedb.flownode_process = Some(flownode).into(); + greptimedb.metasrv_process = Some(meta_server.into_inner()).into(); + greptimedb.server_processes = Some(Arc::new(Mutex::new( + datanodes.into_iter().map(ChildGuard::into_inner).collect(), + ))); + greptimedb.frontend_process = Some(frontend.into_inner()).into(); + greptimedb.flownode_process = Some(flownode.into_inner()).into(); greptimedb.is_standalone = false; greptimedb.ctx = db_ctx; @@ -311,7 +364,7 @@ impl Env { db_ctx: &GreptimeDBContext, id: usize, truncate_log: bool, - ) -> Child { + ) -> ChildGuard { let bins_dir = self.bins_dir.lock().unwrap().clone().expect( "GreptimeDB binary is not available. Please pass in the path to the directory that contains the pre-built GreptimeDB binary. Or you may call `self.build_db()` beforehand.", ); @@ -327,7 +380,7 @@ impl Env { id: usize, truncate_log: bool, bins_dir: PathBuf, - ) -> Child { + ) -> ChildGuard { let log_file_name = match mode { ServerMode::Datanode { node_id, .. } => { db_ctx.incr_datanode_id(); @@ -351,7 +404,15 @@ impl Env { .unwrap(); let arg_style = self.infer_grpc_arg_style(&bins_dir); - let args = mode.get_args(&self.sqlness_home, self, db_ctx, id, arg_style); + let compat_stage = self.compat_config_stage(); + let args = mode.get_args( + &self.sqlness_home, + self, + db_ctx, + id, + arg_style, + &compat_stage, + ); let check_ip_addrs = mode.check_addrs(); for check_ip_addr in &check_ip_addrs { @@ -369,7 +430,7 @@ impl Env { .canonicalize() .expect("Failed to canonicalize bins_dir"); - let mut process = Command::new(abs_bins_dir.join(program)) + let process = Command::new(abs_bins_dir.join(program)) .current_dir(bins_dir.clone()) .env("TZ", "UTC") .args(args) @@ -382,10 +443,10 @@ impl Env { bins_dir.join(program) ); }); + let process = ChildGuard::new(process); for check_ip_addr in &check_ip_addrs { if !util::check_port(check_ip_addr.parse().unwrap(), Duration::from_secs(30)).await { - Env::stop_server(&mut process); panic!( "{} doesn't up in 30 seconds, check {} for more details.", mode.name(), @@ -454,6 +515,7 @@ impl Env { vec![new_server_process] } else { db.ctx.reset_datanode_id(); + let mut new_metasrv = None; if is_full_restart { let metasrv_mode = db .ctx @@ -469,10 +531,7 @@ impl Env { bins_dir.clone(), ) .await; - db.metasrv_process - .lock() - .expect("lock poisoned") - .replace(metasrv); + new_metasrv = Some(metasrv); // wait for metasrv to start // since it seems older version of db might take longer to complete election @@ -498,6 +557,7 @@ impl Env { processes.push(new_server_process); } + let mut new_frontend = None; if is_full_restart { let frontend_mode = db .ctx @@ -514,10 +574,6 @@ impl Env { bins_dir.clone(), ) .await; - db.frontend_process - .lock() - .expect("lock poisoned") - .replace(frontend); // Reconnect protocol clients to the new frontend process // so that MySQL/Postgres queries use the restarted frontend, @@ -529,9 +585,11 @@ impl Env { client .reconnect_pg_client(server_addr.pg_server_addr.as_ref().unwrap()) .await; + new_frontend = Some(frontend); } // Restart flownode. + let mut new_flownode = None; if let Some(flownode_mode) = db.ctx.get_server_mode(SERVER_MODE_FLOWNODE_IDX).cloned() { let flownode = self .start_server_with_bins_dir( @@ -542,10 +600,20 @@ impl Env { bins_dir.clone(), ) .await; - db.flownode_process - .lock() - .expect("lock poisoned") - .replace(flownode); + new_flownode = Some(flownode); + } + + if let Some(metasrv) = new_metasrv { + let mut metasrv_process = db.metasrv_process.lock().expect("lock poisoned"); + metasrv_process.replace(metasrv.into_inner()); + } + if let Some(frontend) = new_frontend { + let mut frontend_process = db.frontend_process.lock().expect("lock poisoned"); + frontend_process.replace(frontend.into_inner()); + } + if let Some(flownode) = new_flownode { + let mut flownode_process = db.flownode_process.lock().expect("lock poisoned"); + flownode_process.replace(flownode.into_inner()); } processes @@ -553,7 +621,10 @@ impl Env { if let Some(server_processes) = db.server_processes.clone() { let mut server_processes = server_processes.lock().unwrap(); - *server_processes = new_server_processes; + *server_processes = new_server_processes + .into_iter() + .map(ChildGuard::into_inner) + .collect(); } } @@ -969,3 +1040,100 @@ impl GreptimeDBContext { self.server_modes.get(idx) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::cmd::bare::ServerAddr; + use crate::cmd::datanode_overlay::{DatanodeOverlay, DatanodeProtectionPolicy}; + + fn test_env(temp_dir: &Path) -> Env { + Env::new( + temp_dir.to_path_buf(), + ServerAddr::default(), + WalConfig::RaftEngine, + false, + None, + StoreConfig { + store_addrs: vec![], + setup_etcd: false, + setup_pg: None, + setup_mysql: None, + enable_flat_format: false, + enable_gc: false, + }, + vec![], + ) + } + + #[test] + fn compat_config_stage_is_shared_by_env_clones_and_transitions_cleanly() { + let temp_dir = tempfile::tempdir().unwrap(); + std::fs::write(temp_dir.path().join("overlay.toml"), "value = 1").unwrap(); + let overlay = DatanodeOverlay::load(temp_dir.path(), Path::new("overlay.toml")) + .unwrap() + .prepare(&DatanodeProtectionPolicy::for_wal(&WalConfig::RaftEngine)) + .unwrap(); + let env = test_env(temp_dir.path()); + let clone = env.clone(); + + assert!(matches!( + clone.compat_config_stage(), + CompatConfigStage::Baseline + )); + env.activate_compat_old(Arc::new(overlay)); + assert!(matches!( + clone.compat_config_stage(), + CompatConfigStage::Old(_) + )); + clone.activate_compat_current(); + assert!(matches!( + env.compat_config_stage(), + CompatConfigStage::Current + )); + } + + #[cfg(unix)] + #[test] + fn child_guard_kills_untransferred_processes() { + let child = std::process::Command::new("sleep") + .arg("60") + .spawn() + .unwrap(); + let pid = child.id().to_string(); + + drop(ChildGuard::new(child)); + + assert!( + !std::process::Command::new("kill") + .args(["-0", &pid]) + .status() + .unwrap() + .success() + ); + } + + #[cfg(unix)] + #[test] + fn child_guard_kills_process_on_forced_unwind() { + let child = std::process::Command::new("sleep") + .arg("60") + .spawn() + .unwrap(); + let pid = child.id().to_string(); + + let unwind = std::panic::catch_unwind(|| { + let _guard = ChildGuard::new(child); + panic!("forced startup unwind"); + }); + + assert!(unwind.is_err()); + assert!( + !std::process::Command::new("kill") + .args(["-0", &pid]) + .status() + .unwrap() + .success() + ); + } +} diff --git a/tests/runner/src/server_mode.rs b/tests/runner/src/server_mode.rs index 43c2d0cca88..1d59403af17 100644 --- a/tests/runner/src/server_mode.rs +++ b/tests/runner/src/server_mode.rs @@ -21,7 +21,7 @@ use tinytemplate::TinyTemplate; use crate::cmd::bare::ServerAddr; use crate::cmd::compat_case::Version; -use crate::env::bare::{Env, GreptimeDBContext, ServiceProvider}; +use crate::env::bare::{CompatConfigStage, Env, GreptimeDBContext, ServiceProvider}; use crate::util; const DEFAULT_LOG_LEVEL: &str = "--log-level=debug,hyper=warn,tower=warn,datafusion=warn,reqwest=warn,sqlparser=warn,h2=info,opendal=info"; @@ -305,6 +305,7 @@ impl ServerMode { sqlness_home: &Path, db_ctx: &GreptimeDBContext, id: usize, + compat_stage: &CompatConfigStage, ) -> String { let mut tt = TinyTemplate::new(); @@ -365,9 +366,34 @@ impl ServerMode { }; let rendered = tt.render(self.name(), &ctx).unwrap(); + let rendered = match (self, compat_stage) { + (ServerMode::Datanode { .. }, CompatConfigStage::Old(overlay)) => overlay + .apply_to_rendered_baseline(&rendered) + .unwrap_or_else(|error| { + panic!( + "Failed to apply old datanode overlay {}: {error}", + overlay.source().display() + ) + }), + _ => rendered, + }; + let stage_suffix = match compat_stage { + CompatConfigStage::Baseline => "baseline", + CompatConfigStage::Old(_) => "old", + CompatConfigStage::Current => "current", + }; let conf_file = data_home - .join(format!("{}-{}-{}.toml", self.name(), id, db_ctx.time())) + .join(if matches!(self, ServerMode::Datanode { .. }) { + format!( + "{}-{}-{}-{stage_suffix}.toml", + self.name(), + id, + db_ctx.time() + ) + } else { + format!("{}-{}-{}.toml", self.name(), id, db_ctx.time()) + }) .display() .to_string(); println!( @@ -387,6 +413,7 @@ impl ServerMode { db_ctx: &GreptimeDBContext, id: usize, arg_style: GrpcArgStyle, + compat_stage: &CompatConfigStage, ) -> Vec { let mut args = env .extra_args() @@ -410,7 +437,7 @@ impl ServerMode { id ), "-c".to_string(), - self.generate_config_file(sqlness_home, db_ctx, id), + self.generate_config_file(sqlness_home, db_ctx, id, compat_stage), format!("--http-addr={http_addr}"), format!("{}={rpc_bind_addr}", arg_style.bind_addr_arg()), format!("--mysql-addr={mysql_addr}"), @@ -439,7 +466,7 @@ impl ServerMode { id ), "-c".to_string(), - self.generate_config_file(sqlness_home, db_ctx, id), + self.generate_config_file(sqlness_home, db_ctx, id, compat_stage), ]); } ServerMode::Metasrv { @@ -461,7 +488,7 @@ impl ServerMode { id ), "-c".to_string(), - self.generate_config_file(sqlness_home, db_ctx, id), + self.generate_config_file(sqlness_home, db_ctx, id, compat_stage), ]); if matches!( @@ -540,7 +567,7 @@ impl ServerMode { format!("--log-dir={}/logs", data_home.display()), format!("--node-id={node_id}"), "-c".to_string(), - self.generate_config_file(sqlness_home, db_ctx, id), + self.generate_config_file(sqlness_home, db_ctx, id, compat_stage), format!("--metasrv-addrs={metasrv_addr}"), ]); } @@ -573,8 +600,10 @@ impl ServerMode { #[cfg(test)] mod tests { use std::path::PathBuf; + use std::sync::Arc; use super::*; + use crate::cmd::datanode_overlay::{DatanodeOverlay, DatanodeProtectionPolicy}; use crate::env::bare::{StoreConfig, WalConfig}; fn test_env(sqlness_home: &Path) -> (Env, GreptimeDBContext) { @@ -600,6 +629,19 @@ mod tests { (env, db_ctx) } + fn old_stage(temp_dir: &Path) -> CompatConfigStage { + std::fs::write( + temp_dir.join("old-datanode.toml"), + "old_only = \"applied\"\nmode = \"standalone\"\n", + ) + .unwrap(); + let overlay = DatanodeOverlay::load(temp_dir, Path::new("old-datanode.toml")) + .unwrap() + .prepare(&DatanodeProtectionPolicy::for_wal(&WalConfig::RaftEngine)) + .unwrap(); + CompatConfigStage::Old(Arc::new(overlay)) + } + fn has_arg(args: &[String], name: &str) -> bool { let prefix = format!("{name}="); args.iter() @@ -641,7 +683,14 @@ mod tests { postgres_addr: "127.0.0.1:4003".to_string(), }; assert_uses_style( - &standalone.get_args(temp_dir, env, db_ctx, 0, style), + &standalone.get_args( + temp_dir, + env, + db_ctx, + 0, + style, + &CompatConfigStage::Baseline, + ), style, false, ); @@ -654,7 +703,14 @@ mod tests { metasrv_addr: "127.0.0.1:4001".to_string(), }; assert_uses_style( - &frontend.get_args(temp_dir, env, db_ctx, 0, style), + &frontend.get_args( + temp_dir, + env, + db_ctx, + 0, + style, + &CompatConfigStage::Baseline, + ), style, true, ); @@ -665,7 +721,14 @@ mod tests { http_addr: "127.0.0.1:4200".to_string(), }; assert_uses_style( - &metasrv.get_args(temp_dir, env, db_ctx, 0, style), + &metasrv.get_args( + temp_dir, + env, + db_ctx, + 0, + style, + &CompatConfigStage::Baseline, + ), style, true, ); @@ -678,7 +741,14 @@ mod tests { node_id: 0, }; assert_uses_style( - &datanode.get_args(temp_dir, env, db_ctx, 0, style), + &datanode.get_args( + temp_dir, + env, + db_ctx, + 0, + style, + &CompatConfigStage::Baseline, + ), style, true, ); @@ -691,7 +761,14 @@ mod tests { node_id: 0, }; assert_uses_style( - &flownode.get_args(temp_dir, env, db_ctx, 0, style), + &flownode.get_args( + temp_dir, + env, + db_ctx, + 0, + style, + &CompatConfigStage::Baseline, + ), style, true, ); @@ -759,14 +836,79 @@ mod tests { node_id: 0, }; - let metasrv_config = - std::fs::read_to_string(metasrv.generate_config_file(temp_dir.path(), &db_ctx, 0)) - .unwrap(); - let datanode_config = - std::fs::read_to_string(datanode.generate_config_file(temp_dir.path(), &db_ctx, 0)) - .unwrap(); + let metasrv_config = std::fs::read_to_string(metasrv.generate_config_file( + temp_dir.path(), + &db_ctx, + 0, + &CompatConfigStage::Baseline, + )) + .unwrap(); + let datanode_config = std::fs::read_to_string(datanode.generate_config_file( + temp_dir.path(), + &db_ctx, + 0, + &CompatConfigStage::Baseline, + )) + .unwrap(); assert!(metasrv_config.contains("[gc]\nenable = true")); assert!(datanode_config.contains("[region_engine.mito.gc]\nenable = true")); } + + #[test] + fn test_datanode_rendering_applies_only_old_stage_with_distinct_filenames() { + let temp_dir = tempfile::tempdir().unwrap(); + let (_, db_ctx) = test_env(temp_dir.path()); + let datanode = ServerMode::Datanode { + rpc_bind_addr: "127.0.0.1:4301".to_string(), + rpc_server_addr: "127.0.0.1:4301".to_string(), + http_addr: "127.0.0.1:4300".to_string(), + metasrv_addr: "127.0.0.1:4201".to_string(), + node_id: 0, + }; + let baseline_path = datanode.generate_config_file( + temp_dir.path(), + &db_ctx, + 0, + &CompatConfigStage::Baseline, + ); + let old_path = + datanode.generate_config_file(temp_dir.path(), &db_ctx, 0, &old_stage(temp_dir.path())); + let current_path = + datanode.generate_config_file(temp_dir.path(), &db_ctx, 0, &CompatConfigStage::Current); + + let baseline = std::fs::read_to_string(&baseline_path).unwrap(); + let old = std::fs::read_to_string(&old_path).unwrap(); + let current = std::fs::read_to_string(¤t_path).unwrap(); + assert!(baseline_path.ends_with("-baseline.toml")); + assert!(old_path.ends_with("-old.toml")); + assert!(current_path.ends_with("-current.toml")); + assert!(!baseline.contains("old_only")); + assert!(old.contains("old_only = \"applied\"")); + assert!(!current.contains("old_only")); + assert_eq!( + toml::from_str::(&old).unwrap()["mode"].as_str(), + toml::from_str::(&baseline).unwrap()["mode"].as_str() + ); + } + + #[test] + fn test_non_datanode_ignores_old_overlay() { + let temp_dir = tempfile::tempdir().unwrap(); + let (_, db_ctx) = test_env(temp_dir.path()); + let metasrv = ServerMode::Metasrv { + rpc_bind_addr: "127.0.0.1:4201".to_string(), + rpc_server_addr: "127.0.0.1:4201".to_string(), + http_addr: "127.0.0.1:4200".to_string(), + }; + let config_path = + metasrv.generate_config_file(temp_dir.path(), &db_ctx, 0, &old_stage(temp_dir.path())); + + assert!(!config_path.ends_with("-old.toml")); + assert!( + !std::fs::read_to_string(config_path) + .unwrap() + .contains("old_only") + ); + } } diff --git a/tests/runner/src/util.rs b/tests/runner/src/util.rs index 6c00c8a9a65..1be4183b49a 100644 --- a/tests/runner/src/util.rs +++ b/tests/runner/src/util.rs @@ -285,25 +285,55 @@ pub fn setup_etcd(client_ports: Vec, peer_port: Option, etcd_version: } } -/// Stop and remove the etcd container -pub fn stop_rm_etcd() { +/// Stop and remove the etcd container, failing if it cannot be confirmed absent. +pub fn stop_rm_etcd_checked() -> Result<(), String> { let status = std::process::Command::new("docker") - .args(["container", "stop", "etcd"]) - .status(); - if status.is_err() { - panic!("Failed to stop etcd: {:?}", status); - } else { - println!("Stopped etcd"); - } - // rm the container - let status = std::process::Command::new("docker") - .args(["container", "rm", "etcd"]) - .status(); - if status.is_err() { - panic!("Failed to remove etcd container: {:?}", status); - } else { + .args(["container", "rm", "--force", "etcd"]) + .status() + .map_err(|error| format!("Failed to run Docker while removing etcd: {error}"))?; + if status.success() { println!("Removed etcd container"); + return Ok(()); } + + let listed = std::process::Command::new("docker") + .args([ + "container", + "ls", + "--all", + "--filter", + "name=^/etcd$", + "--format", + "{{.ID}}", + ]) + .output() + .map_err(|error| { + format!( + "Docker failed to remove etcd ({status}) and could not verify its absence: {error}" + ) + })?; + if !listed.status.success() { + return Err(format!( + "Docker failed to remove etcd ({status}) and listing containers failed: {}", + String::from_utf8_lossy(&listed.stderr).trim() + )); + } + if String::from_utf8_lossy(&listed.stdout).trim().is_empty() { + println!("Etcd container is already absent"); + return Ok(()); + } + + Err(format!( + "Docker failed to remove etcd container (status {status}); the container still exists" + )) +} + +/// Stop and remove the etcd container. +/// +/// Legacy callers retain panic-on-cleanup-failure behavior. Compatibility +/// profiles use [`stop_rm_etcd_checked`] to report the failure structurally. +pub fn stop_rm_etcd() { + stop_rm_etcd_checked().unwrap_or_else(|error| panic!("{error}")); } /// Set up a PostgreSQL server in docker. From 57b8239ff88920d244c2465d86e2b3eab1ad529e Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 7 Aug 2026 17:05:20 +0800 Subject: [PATCH 019/126] test: rename internal bug numbers in tests to semantic names (#8779) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- tests-integration/tests/http.rs | 71 +++++++------------ .../copy/copy_from_fs_parquet.result | 34 ++++----- .../standalone/copy/copy_from_fs_parquet.sql | 34 ++++----- 3 files changed, 60 insertions(+), 79 deletions(-) diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index d7912d0a975..49ae28dec0d 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -3017,12 +3017,12 @@ pub async fn test_prometheus_remote_write_with_pipeline(store_type: StorageType) guard.remove_all().await; } -fn qx_077_schema_series(schema: &str, value: f64) -> TimeSeries { +fn schema_series(schema: &str, value: f64) -> TimeSeries { TimeSeries { labels: vec![ Label { name: "__name__".to_string(), - value: "qx_077_schema_metric".to_string(), + value: "schema_metric".to_string(), }, Label { name: "__schema__".to_string(), @@ -3041,10 +3041,10 @@ fn qx_077_schema_series(schema: &str, value: f64) -> TimeSeries { } } -fn qx_077_schema_query(schema: Option<&str>) -> Query { +fn schema_query(schema: Option<&str>) -> Query { let mut matchers = vec![LabelMatcher { name: "__name__".to_string(), - value: "qx_077_schema_metric".to_string(), + value: "schema_metric".to_string(), r#type: MatcherType::Eq as i32, }]; if let Some(schema) = schema { @@ -3063,7 +3063,7 @@ fn qx_077_schema_query(schema: Option<&str>) -> Query { } } -async fn post_qx_077_remote_read(client: &TestClient, queries: Vec) -> ReadResponse { +async fn post_remote_read(client: &TestClient, queries: Vec) -> ReadResponse { let read_request = ReadRequest { queries, ..Default::default() @@ -3083,10 +3083,7 @@ async fn post_qx_077_remote_read(client: &TestClient, queries: Vec) -> Re ReadResponse::decode(decompressed_response.as_slice()).unwrap() } -async fn post_qx_077_remote_read_from_schema_a( - client: &TestClient, - queries: Vec, -) -> ReadResponse { +async fn post_remote_read_from_schema_a(client: &TestClient, queries: Vec) -> ReadResponse { let read_request = ReadRequest { queries, ..Default::default() @@ -3096,7 +3093,7 @@ async fn post_qx_077_remote_read_from_schema_a( let response = client .post("/v1/prometheus/read") - .header(GREPTIME_DB_HEADER_NAME.clone(), "qx_077_schema_a") + .header(GREPTIME_DB_HEADER_NAME.clone(), "schema_a") .body(compressed_request) .send() .await; @@ -3107,7 +3104,7 @@ async fn post_qx_077_remote_read_from_schema_a( ReadResponse::decode(decompressed_response.as_slice()).unwrap() } -fn qx_077_read_sample_values(read_response: ReadResponse) -> Vec { +fn read_sample_values(read_response: ReadResponse) -> Vec { read_response .results .into_iter() @@ -3127,8 +3124,8 @@ pub async fn test_prometheus_remote_schema_labels(store_type: StorageType) { setup_test_prom_app_with_frontend(store_type, "test_prometheus_remote_schema_labels").await; let client = TestClient::new(app).await; - const SCHEMA_A: &str = "qx_077_schema_a"; - const SCHEMA_B: &str = "qx_077_schema_b"; + const SCHEMA_A: &str = "schema_a"; + const SCHEMA_B: &str = "schema_b"; for schema in [SCHEMA_A, SCHEMA_B] { let res = client @@ -3142,8 +3139,8 @@ pub async fn test_prometheus_remote_schema_labels(store_type: StorageType) { // Both series have the same metric, non-schema labels, and timestamp. let write_request = WriteRequest { timeseries: vec![ - qx_077_schema_series(SCHEMA_A, 101.0), - qx_077_schema_series(SCHEMA_B, 202.0), + schema_series(SCHEMA_A, 101.0), + schema_series(SCHEMA_B, 202.0), ], ..Default::default() }; @@ -3160,36 +3157,26 @@ pub async fn test_prometheus_remote_schema_labels(store_type: StorageType) { // Independent controls establish that each selector resolves its own schema. assert_eq!( - qx_077_read_sample_values( - post_qx_077_remote_read(&client, vec![qx_077_schema_query(Some(SCHEMA_A))]).await, - ), + read_sample_values(post_remote_read(&client, vec![schema_query(Some(SCHEMA_A))]).await,), vec![101.0] ); assert_eq!( - qx_077_read_sample_values( - post_qx_077_remote_read(&client, vec![qx_077_schema_query(Some(SCHEMA_B))]).await, - ), + read_sample_values(post_remote_read(&client, vec![schema_query(Some(SCHEMA_B))]).await,), vec![202.0] ); // Results must stay aligned with the request query order. - let forward_values = qx_077_read_sample_values( - post_qx_077_remote_read( + let forward_values = read_sample_values( + post_remote_read( &client, - vec![ - qx_077_schema_query(Some(SCHEMA_A)), - qx_077_schema_query(Some(SCHEMA_B)), - ], + vec![schema_query(Some(SCHEMA_A)), schema_query(Some(SCHEMA_B))], ) .await, ); - let reverse_values = qx_077_read_sample_values( - post_qx_077_remote_read( + let reverse_values = read_sample_values( + post_remote_read( &client, - vec![ - qx_077_schema_query(Some(SCHEMA_B)), - qx_077_schema_query(Some(SCHEMA_A)), - ], + vec![schema_query(Some(SCHEMA_B)), schema_query(Some(SCHEMA_A))], ) .await, ); @@ -3199,23 +3186,17 @@ pub async fn test_prometheus_remote_schema_labels(store_type: StorageType) { ); // With request database A, unqualified queries must resolve to A independently. - let explicit_then_default_values = qx_077_read_sample_values( - post_qx_077_remote_read_from_schema_a( + let explicit_then_default_values = read_sample_values( + post_remote_read_from_schema_a( &client, - vec![ - qx_077_schema_query(Some(SCHEMA_B)), - qx_077_schema_query(None), - ], + vec![schema_query(Some(SCHEMA_B)), schema_query(None)], ) .await, ); - let default_then_explicit_values = qx_077_read_sample_values( - post_qx_077_remote_read_from_schema_a( + let default_then_explicit_values = read_sample_values( + post_remote_read_from_schema_a( &client, - vec![ - qx_077_schema_query(None), - qx_077_schema_query(Some(SCHEMA_B)), - ], + vec![schema_query(None), schema_query(Some(SCHEMA_B))], ) .await, ); diff --git a/tests/cases/standalone/copy/copy_from_fs_parquet.result b/tests/cases/standalone/copy/copy_from_fs_parquet.result index 6be39497c8d..ab58311433a 100644 --- a/tests/cases/standalone/copy/copy_from_fs_parquet.result +++ b/tests/cases/standalone/copy/copy_from_fs_parquet.result @@ -243,32 +243,32 @@ drop table demo_with_less_columns; Affected Rows: 0 --- QX-083: COPY FROM Parquet honors LIMIT for a single file. -CREATE TABLE qx_083_copy_limit_source(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); +-- COPY FROM Parquet honors LIMIT for a single file. +CREATE TABLE copy_limit_source(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); Affected Rows: 0 -INSERT INTO qx_083_copy_limit_source VALUES +INSERT INTO copy_limit_source VALUES (1, 1704067200000), (2, 1704067201000), (3, 1704067202000); Affected Rows: 3 -COPY qx_083_copy_limit_source TO '${SQLNESS_HOME}/qx_083_copy_limit.parquet' WITH (format='parquet'); +COPY copy_limit_source TO '${SQLNESS_HOME}/copy_limit.parquet' WITH (format='parquet'); Affected Rows: 3 -- LIMIT 0 control: affected rows 0; count 0. -CREATE TABLE qx_083_copy_limit_limit_0(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); +CREATE TABLE copy_limit_limit_0(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); Affected Rows: 0 -COPY qx_083_copy_limit_limit_0 FROM '${SQLNESS_HOME}/qx_083_copy_limit.parquet' WITH (format='parquet') LIMIT 0; +COPY copy_limit_limit_0 FROM '${SQLNESS_HOME}/copy_limit.parquet' WITH (format='parquet') LIMIT 0; Affected Rows: 0 -SELECT count(*) FROM qx_083_copy_limit_limit_0; +SELECT count(*) FROM copy_limit_limit_0; +----------+ | count(*) | @@ -277,15 +277,15 @@ SELECT count(*) FROM qx_083_copy_limit_limit_0; +----------+ -- LIMIT 1 target: affected rows 1; count 1. -CREATE TABLE qx_083_copy_limit_limit_1(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); +CREATE TABLE copy_limit_limit_1(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); Affected Rows: 0 -COPY qx_083_copy_limit_limit_1 FROM '${SQLNESS_HOME}/qx_083_copy_limit.parquet' WITH (format='parquet') LIMIT 1; +COPY copy_limit_limit_1 FROM '${SQLNESS_HOME}/copy_limit.parquet' WITH (format='parquet') LIMIT 1; Affected Rows: 1 -SELECT count(*) FROM qx_083_copy_limit_limit_1; +SELECT count(*) FROM copy_limit_limit_1; +----------+ | count(*) | @@ -294,15 +294,15 @@ SELECT count(*) FROM qx_083_copy_limit_limit_1; +----------+ -- Cross-file LIMIT 4 target: affected rows 4; count 4. -CREATE TABLE qx_083_copy_limit_cross_file(host string, cpu double, memory double, ts TIMESTAMP TIME INDEX); +CREATE TABLE copy_limit_cross_file(host string, cpu double, memory double, ts TIMESTAMP TIME INDEX); Affected Rows: 0 -COPY qx_083_copy_limit_cross_file FROM '${SQLNESS_HOME}/demo/export/parquet_files/' WITH (start_time='2022-06-15 07:02:36') LIMIT 4; +COPY copy_limit_cross_file FROM '${SQLNESS_HOME}/demo/export/parquet_files/' WITH (start_time='2022-06-15 07:02:36') LIMIT 4; Affected Rows: 4 -SELECT count(*) FROM qx_083_copy_limit_cross_file; +SELECT count(*) FROM copy_limit_cross_file; +----------+ | count(*) | @@ -310,19 +310,19 @@ SELECT count(*) FROM qx_083_copy_limit_cross_file; | 4 | +----------+ -DROP TABLE qx_083_copy_limit_source; +DROP TABLE copy_limit_source; Affected Rows: 0 -DROP TABLE qx_083_copy_limit_limit_0; +DROP TABLE copy_limit_limit_0; Affected Rows: 0 -DROP TABLE qx_083_copy_limit_limit_1; +DROP TABLE copy_limit_limit_1; Affected Rows: 0 -DROP TABLE qx_083_copy_limit_cross_file; +DROP TABLE copy_limit_cross_file; Affected Rows: 0 diff --git a/tests/cases/standalone/copy/copy_from_fs_parquet.sql b/tests/cases/standalone/copy/copy_from_fs_parquet.sql index 9cdfbb241df..4f2da8cc38c 100644 --- a/tests/cases/standalone/copy/copy_from_fs_parquet.sql +++ b/tests/cases/standalone/copy/copy_from_fs_parquet.sql @@ -100,41 +100,41 @@ drop table demo_with_external_column; drop table demo_with_less_columns; --- QX-083: COPY FROM Parquet honors LIMIT for a single file. -CREATE TABLE qx_083_copy_limit_source(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); +-- COPY FROM Parquet honors LIMIT for a single file. +CREATE TABLE copy_limit_source(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); -INSERT INTO qx_083_copy_limit_source VALUES +INSERT INTO copy_limit_source VALUES (1, 1704067200000), (2, 1704067201000), (3, 1704067202000); -COPY qx_083_copy_limit_source TO '${SQLNESS_HOME}/qx_083_copy_limit.parquet' WITH (format='parquet'); +COPY copy_limit_source TO '${SQLNESS_HOME}/copy_limit.parquet' WITH (format='parquet'); -- LIMIT 0 control: affected rows 0; count 0. -CREATE TABLE qx_083_copy_limit_limit_0(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); +CREATE TABLE copy_limit_limit_0(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); -COPY qx_083_copy_limit_limit_0 FROM '${SQLNESS_HOME}/qx_083_copy_limit.parquet' WITH (format='parquet') LIMIT 0; +COPY copy_limit_limit_0 FROM '${SQLNESS_HOME}/copy_limit.parquet' WITH (format='parquet') LIMIT 0; -SELECT count(*) FROM qx_083_copy_limit_limit_0; +SELECT count(*) FROM copy_limit_limit_0; -- LIMIT 1 target: affected rows 1; count 1. -CREATE TABLE qx_083_copy_limit_limit_1(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); +CREATE TABLE copy_limit_limit_1(pk_col INTEGER PRIMARY KEY, ts TIMESTAMP TIME INDEX); -COPY qx_083_copy_limit_limit_1 FROM '${SQLNESS_HOME}/qx_083_copy_limit.parquet' WITH (format='parquet') LIMIT 1; +COPY copy_limit_limit_1 FROM '${SQLNESS_HOME}/copy_limit.parquet' WITH (format='parquet') LIMIT 1; -SELECT count(*) FROM qx_083_copy_limit_limit_1; +SELECT count(*) FROM copy_limit_limit_1; -- Cross-file LIMIT 4 target: affected rows 4; count 4. -CREATE TABLE qx_083_copy_limit_cross_file(host string, cpu double, memory double, ts TIMESTAMP TIME INDEX); +CREATE TABLE copy_limit_cross_file(host string, cpu double, memory double, ts TIMESTAMP TIME INDEX); -COPY qx_083_copy_limit_cross_file FROM '${SQLNESS_HOME}/demo/export/parquet_files/' WITH (start_time='2022-06-15 07:02:36') LIMIT 4; +COPY copy_limit_cross_file FROM '${SQLNESS_HOME}/demo/export/parquet_files/' WITH (start_time='2022-06-15 07:02:36') LIMIT 4; -SELECT count(*) FROM qx_083_copy_limit_cross_file; +SELECT count(*) FROM copy_limit_cross_file; -DROP TABLE qx_083_copy_limit_source; +DROP TABLE copy_limit_source; -DROP TABLE qx_083_copy_limit_limit_0; +DROP TABLE copy_limit_limit_0; -DROP TABLE qx_083_copy_limit_limit_1; +DROP TABLE copy_limit_limit_1; -DROP TABLE qx_083_copy_limit_cross_file; +DROP TABLE copy_limit_cross_file; From 70470bafbe38df2d3fd21e70e1e7097ffd7317c6 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 7 Aug 2026 17:49:55 +0800 Subject: [PATCH 020/126] fix(flow): fix flow stats aggregation and df_plan_to_sql quoting (#8729) * fix(flow): fix flow stats aggregation and df_plan_to_sql quoting 1. Distributed-mode flow stats last-writer-wins overwrite: Each flownode heartbeat put its local flow state map into the single global __flow/state key, so reports from different nodes overwrote each other. Store per-flownode reports under __flow/state/node/{node_id} in the in-memory KV and aggregate on each heartbeat (last_exec_time_map/state_size/start_time_map take the max across nodes) into the global key. FlowStateHandler derives node identity from header.member_id (fallback peer.id) and ignores identity-less reports. Per-node keys clear automatically on leader change KV reset. Adapts to FlowStateValue.start_time_map added in #8392. 2. df_plan_to_sql unquoted special characters break flush/scheduled execution: ForceQuoteIdentifiers only quoted uppercase identifiers, so Prometheus-style table names with ':' (e.g. cpu_cores:sum) were left unquoted, producing invalid SQL ('keyword: :'). Quote any identifier with non-[a-z0-9_] chars using double quotes (dialect-neutral). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): address review comments on quoting and logging - df_plan_to_sql: also quote digit-leading identifiers (e.g. 123metrics) which would produce invalid SQL when re-parsed. SQL keywords are intentionally not checked (ALL_KEYWORDS would over-quote common column names like number; the unparse failure path has an InsertIntoPlan fallback). - flow_state_handler: downgrade identity-less report log from warn! to debug! to avoid an anomalous sender spamming warn every heartbeat. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/common/meta/src/key/flow/flow_state.rs | 391 +++++++++++++++++- src/flow/src/batching_mode/utils.rs | 24 +- src/flow/src/batching_mode/utils/test.rs | 77 +++- .../src/handler/flow_state_handler.rs | 153 ++++++- 4 files changed, 632 insertions(+), 13 deletions(-) diff --git a/src/common/meta/src/key/flow/flow_state.rs b/src/common/meta/src/key/flow/flow_state.rs index 77ccd6206ab..5423cef37b3 100644 --- a/src/common/meta/src/key/flow/flow_state.rs +++ b/src/common/meta/src/key/flow/flow_state.rs @@ -16,16 +16,20 @@ use std::collections::BTreeMap; use std::sync::Arc; use serde::{Deserialize, Serialize}; +use tokio::sync::Mutex; use crate::error::{self, Result}; use crate::key::flow::FlowScoped; use crate::key::{FlowId, MetadataKey, MetadataValue}; use crate::kv_backend::KvBackendRef; -use crate::rpc::store::PutRequest; +use crate::rpc::store::{PutRequest, RangeRequest}; /// The entire FlowId to Flow Size's Map is stored directly in the value part of the key. pub const FLOW_STATE_KEY: &str = "state"; +/// The inner prefix (under `state/`) of the per-flownode flow state keys. +pub const FLOW_STATE_NODE_KEY_PREFIX: &str = "node"; + /// The key of flow state. #[derive(Debug, Clone, Copy, PartialEq)] struct FlowStateKeyInner; @@ -92,6 +96,99 @@ impl<'a> MetadataKey<'a, FlowStateKey> for FlowStateKey { } } +/// The inner key of a per-flownode flow state entry: `state/node/{node_id}`. +/// +/// `node_id` is the operator-configured flownode id (unique within the +/// cluster; a flownode requires `node_id` in its config, see +/// `src/cmd/src/flownode.rs`). It is the same value reported as +/// `HeartbeatRequest.header.member_id` — the canonical identity metasrv uses +/// for flownodes, see `get_node_id` in `src/meta-srv/src/service/heartbeat.rs` +/// — and as `HeartbeatRequest.peer.id`. +#[derive(Debug, Clone, PartialEq)] +struct FlowStateNodeKeyInner { + node_id: u64, +} + +impl FlowStateNodeKeyInner { + pub fn new(node_id: u64) -> Self { + Self { node_id } + } +} + +impl<'a> MetadataKey<'a, FlowStateNodeKeyInner> for FlowStateNodeKeyInner { + fn to_bytes(&self) -> Vec { + format!( + "{FLOW_STATE_KEY}/{FLOW_STATE_NODE_KEY_PREFIX}/{}", + self.node_id + ) + .into_bytes() + } + + fn from_bytes(bytes: &'a [u8]) -> Result { + let key = std::str::from_utf8(bytes).map_err(|e| { + error::InvalidMetadataSnafu { + err_msg: format!( + "FlowStateNodeKeyInner '{}' is not a valid UTF8 string: {e}", + String::from_utf8_lossy(bytes) + ), + } + .build() + })?; + let prefix = format!("{FLOW_STATE_KEY}/{FLOW_STATE_NODE_KEY_PREFIX}/"); + let Some(node_id) = key.strip_prefix(&prefix) else { + return Err(error::InvalidMetadataSnafu { + err_msg: format!("Invalid FlowStateNodeKeyInner '{key}'"), + } + .build()); + }; + let node_id = node_id.parse::().map_err(|_| { + error::InvalidMetadataSnafu { + err_msg: format!("Invalid node id '{node_id}' in FlowStateNodeKeyInner '{key}'"), + } + .build() + })?; + Ok(FlowStateNodeKeyInner::new(node_id)) + } +} + +/// The key stores the per-flownode flow state report. +/// +/// The layout: `__flow/state/node/{node_id}`. +/// +/// Per-node keys live in the in-memory KV (same as the global `__flow/state` +/// key), so they are automatically cleared when metasrv resets the in-memory +/// KV on leader change; no separate cleanup is needed. +pub struct FlowStateNodeKey(FlowScoped); + +impl FlowStateNodeKey { + /// Returns the [FlowStateNodeKey] of the given node. + pub fn new(node_id: u64) -> FlowStateNodeKey { + FlowStateNodeKey(FlowScoped::new(FlowStateNodeKeyInner::new(node_id))) + } + + /// Returns the full key prefix of all per-node flow state keys: + /// `__flow/state/node/`. + pub fn prefix() -> Vec { + format!( + "{}{FLOW_STATE_KEY}/{FLOW_STATE_NODE_KEY_PREFIX}/", + FlowScoped::::PREFIX + ) + .into_bytes() + } +} + +impl<'a> MetadataKey<'a, FlowStateNodeKey> for FlowStateNodeKey { + fn to_bytes(&self) -> Vec { + self.0.to_bytes() + } + + fn from_bytes(bytes: &'a [u8]) -> Result { + Ok(FlowStateNodeKey( + FlowScoped::::from_bytes(bytes)?, + )) + } +} + /// The value of flow state size #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] pub struct FlowStateValue { @@ -125,13 +222,25 @@ pub type FlowStateManagerRef = Arc; /// /// This is only used in distributed mode. When meta-srv use heartbeat to update the flow stat report /// and frontned use get to get the latest flow stat report. +/// +/// Per-flownode reports are stored under `__flow/state/node/{node_id}` keys in +/// the in-memory KV (not in a separate in-process map), so a metasrv leader +/// change — which resets the in-memory KV — automatically clears all per-node +/// state without leaving stale entries behind. pub struct FlowStateManager { in_memory: KvBackendRef, + /// Serializes the critical section of [`FlowStateManager::merge`] + /// (write per-node key -> scan -> aggregate -> write global key). It holds + /// no long-lived state; per-node reports live in the in-memory KV. + merge_lock: Mutex<()>, } impl FlowStateManager { pub fn new(in_memory: KvBackendRef) -> Self { - Self { in_memory } + Self { + in_memory, + merge_lock: Mutex::new(()), + } } pub async fn get(&self) -> Result> { @@ -150,6 +259,78 @@ impl FlowStateManager { self.in_memory.put(req).await?; Ok(()) } + + /// Merges a flow state report from a single flownode into the global view. + /// + /// `node_id` is the operator-configured flownode id (unique within the + /// cluster; reported as `HeartbeatRequest.header.member_id`, the canonical + /// metasrv identity of a flownode, and equal to `peer.id`). + /// + /// Reports are tracked per node under `__flow/state/node/{node_id}` in the + /// in-memory KV. A new report from the same node unconditionally replaces + /// that node's previous entry: reports are processed strictly in arrival + /// order, so no epoch comparison is made and a wall-clock rollback after a + /// node restart cannot permanently drop the node's later reports. The + /// global `FlowStateValue` is then aggregated over all per-node entries: + /// for the same flow, `last_exec_time_map` takes the max reported + /// timestamp and `state_size` takes the max reported size across nodes (a + /// flow normally runs on a single active flownode, so max is a safe + /// approximation). The aggregated value is written into the in-memory KV + /// under the global key `__flow/state` via the same path as `put`, keeping + /// `get()` behavior unchanged. + /// + /// Per-node keys are cleared automatically when the in-memory KV is reset + /// on a leader change. Known limitation: entries of dropped flows are not + /// proactively removed here. Once a flow is dropped its metadata is gone, + /// so the flows table join simply can't see the stale entry; it becomes + /// user-invisible until the owning flownode reports again (or stops + /// heartbeating forever, in which case the stale flow id lingers in the + /// aggregate but is never joined against any flow metadata). + pub async fn merge(&self, node_id: u64, incoming: FlowStateValue) -> Result<()> { + let _guard = self.merge_lock.lock().await; + + // 1. Store this node's latest report under its per-node key. + let node_key = FlowStateNodeKey::new(node_id).to_bytes(); + let value = incoming.try_as_raw_value()?; + let req = PutRequest::new().with_key(node_key).with_value(value); + self.in_memory.put(req).await?; + + // 2. Read back every per-node report and aggregate them. + let req = RangeRequest::new().with_prefix(FlowStateNodeKey::prefix()); + let resp = self.in_memory.range(req).await?; + let mut state_size = BTreeMap::new(); + let mut last_exec_time_map = BTreeMap::new(); + let mut start_time_map = BTreeMap::new(); + for kv in resp.kvs { + let state = FlowStateValue::try_from_raw_value(&kv.value)?; + for (flow_id, size) in state.state_size { + state_size + .entry(flow_id) + .and_modify(|v: &mut usize| *v = (*v).max(size)) + .or_insert(size); + } + for (flow_id, ts) in state.last_exec_time_map { + last_exec_time_map + .entry(flow_id) + .and_modify(|v: &mut i64| *v = (*v).max(ts)) + .or_insert(ts); + } + for (flow_id, ts) in state.start_time_map { + start_time_map + .entry(flow_id) + .and_modify(|v: &mut i64| *v = (*v).max(ts)) + .or_insert(ts); + } + } + + // 3. Write the aggregated value to the global key. + let aggregated = FlowStateValue::new(state_size, last_exec_time_map, start_time_map); + let key = FlowStateKey::new().to_bytes(); + let value = aggregated.try_as_raw_value()?; + let req = PutRequest::new().with_key(key).with_value(value); + self.in_memory.put(req).await?; + Ok(()) + } } /// Flow's state report, send regularly through heartbeat message @@ -234,4 +415,210 @@ mod tests { let decoded: FlowStateValue = serde_json::from_str(&json).unwrap(); assert_eq!(decoded, value); } + + use std::sync::Arc; + + use super::*; + use crate::kv_backend::memory::MemoryKvBackend; + + fn state(last_exec_time_map: BTreeMap) -> FlowStateValue { + FlowStateValue::new(BTreeMap::new(), last_exec_time_map, BTreeMap::new()) + } + + #[tokio::test] + async fn test_merge_keeps_reports_from_different_nodes() { + let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default())); + + // Node A reports flow 1 executed at t1. + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + // Node B reports flow 2 executed at t2. Before the per-node merge this + // would have wiped out node A's flow 1 entry. + manager + .merge(2, state(BTreeMap::from([(2, 200)]))) + .await + .unwrap(); + // Node A reports flow 1 executed at t3. + manager + .merge(1, state(BTreeMap::from([(1, 300)]))) + .await + .unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + // flow 1 and flow 2 are both present, and flow 1 takes max(t1, t3). + assert_eq!(value.last_exec_time_map.get(&1), Some(&300)); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_merge_replaces_same_node_state() { + let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default())); + + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + // A new report from the same node replaces the previous one. There is + // no epoch comparison: arrival order alone decides, so this is also + // what a restarted node (new epoch) hits. + manager + .merge(1, state(BTreeMap::from([(2, 200)]))) + .await + .unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + assert!(!value.last_exec_time_map.contains_key(&1)); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_merge_accepts_clock_rollback_from_same_node() { + let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default())); + + // Simulates a flownode restart whose wall clock rolled back: the node + // first reports flow 1 at t1, then (after restart) reports a *smaller* + // timestamp t0. Because reports are processed strictly in arrival + // order (no epoch comparison), the later report must win instead of + // being permanently rejected. + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + manager + .merge(1, state(BTreeMap::from([(1, 50)]))) + .await + .unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + assert_eq!(value.last_exec_time_map.get(&1), Some(&50)); + } + + #[tokio::test] + async fn test_merge_aggregates_state_size_and_last_exec_time_by_max() { + let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default())); + + // Both nodes report the same flow; the aggregate must take the max of + // state_size and last_exec_time_map across nodes. + manager + .merge( + 1, + FlowStateValue::new( + BTreeMap::from([(1, 1024)]), + BTreeMap::from([(1, 100)]), + BTreeMap::new(), + ), + ) + .await + .unwrap(); + manager + .merge( + 2, + FlowStateValue::new( + BTreeMap::from([(1, 2048)]), + BTreeMap::from([(1, 50)]), + BTreeMap::new(), + ), + ) + .await + .unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + assert_eq!(value.state_size.get(&1), Some(&2048)); + assert_eq!(value.last_exec_time_map.get(&1), Some(&100)); + } + + #[tokio::test] + async fn test_merge_concurrent_reports_no_lost_update() { + let manager = Arc::new(FlowStateManager::new(Arc::new(MemoryKvBackend::default()))); + + // Two nodes report concurrently; the merge lock must serialize the + // read-modify-write so neither node's report is lost. + let m1 = manager.clone(); + let h1 = tokio::spawn(async move { + m1.merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + }); + let m2 = manager.clone(); + let h2 = tokio::spawn(async move { + m2.merge(2, state(BTreeMap::from([(2, 200)]))) + .await + .unwrap(); + }); + h1.await.unwrap(); + h2.await.unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + assert_eq!(value.last_exec_time_map.get(&1), Some(&100)); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_merge_empty_report_removes_own_flows_keeps_others() { + let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default())); + + // Node A reports flow 1, node B reports flow 2. + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + manager + .merge(2, state(BTreeMap::from([(2, 200)]))) + .await + .unwrap(); + + // Node A reports an empty map: its own flow 1 disappears from the + // aggregate while node B's flow 2 is retained. + manager.merge(1, state(BTreeMap::new())).await.unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + assert!(!value.last_exec_time_map.contains_key(&1)); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_merge_state_cleared_on_in_memory_kv_reset() { + let backend = Arc::new(MemoryKvBackend::default()); + let manager = FlowStateManager::new(backend.clone()); + + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + assert!(manager.get().await.unwrap().is_some()); + + // Simulate a metasrv leader change, which clears the in-memory KV + // (including the per-node keys, since they live in the same KV). + backend.clear(); + assert!(manager.get().await.unwrap().is_none()); + + // A fresh report works again after the reset. + manager + .merge(2, state(BTreeMap::from([(2, 200)]))) + .await + .unwrap(); + let value = manager.get().await.unwrap().unwrap(); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_merge_writes_global_key() { + let backend = Arc::new(MemoryKvBackend::default()); + let manager = FlowStateManager::new(backend.clone()); + + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + + // The global key is still `__flow/state` and holds the serialized + // FlowStateValue, so existing get()/client readers are unchanged. + let dump = backend.dump(); + let global_key = "__flow/state".as_bytes().to_vec(); + assert!(dump.contains_key(&global_key)); + let value = FlowStateValue::try_from_raw_value(dump.get(&global_key).unwrap()).unwrap(); + assert_eq!(value.last_exec_time_map.get(&1), Some(&100)); + } } diff --git a/src/flow/src/batching_mode/utils.rs b/src/flow/src/batching_mode/utils.rs index 3a4c86b30a9..a707c2d9eb6 100644 --- a/src/flow/src/batching_mode/utils.rs +++ b/src/flow/src/batching_mode/utils.rs @@ -969,15 +969,27 @@ pub(crate) async fn gen_plan_with_matching_schema( } pub fn df_plan_to_sql(plan: &LogicalPlan) -> Result { - /// A dialect that forces identifiers to be quoted when have uppercase + /// A dialect that forces identifiers to be quoted when they contain + /// anything other than lowercase alphanumerics and underscores, or start + /// with a digit. + /// + /// Unquoted identifiers are normalized to lowercase by the SQL parser, so + /// uppercase letters need quoting to preserve case. Special characters + /// (e.g. ':' in Prometheus-style table names like + /// `kube_pod_cpu_cores:sum`, '.', '-', spaces) and digit-leading names + /// (e.g. `123metrics`) would produce invalid SQL if left unquoted, so they + /// are quoted as well. SQL keywords are intentionally not checked here: + /// quoting every ALL_KEYWORDS member would also quote common column names + /// like `number`; the unparse failure path has an InsertIntoPlan fallback. struct ForceQuoteIdentifiers; impl datafusion::sql::unparser::dialect::Dialect for ForceQuoteIdentifiers { fn identifier_quote_style(&self, identifier: &str) -> Option { - if identifier.to_lowercase() != identifier { - Some('`') - } else { - None - } + let is_plain = !identifier.is_empty() + && !identifier.starts_with(|c: char| c.is_ascii_digit()) + && identifier + .chars() + .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_'); + if is_plain { None } else { Some('"') } } } let unparser = Unparser::new(&ForceQuoteIdentifiers); diff --git a/src/flow/src/batching_mode/utils/test.rs b/src/flow/src/batching_mode/utils/test.rs index 7fbd3300c4d..1356f87e8f6 100644 --- a/src/flow/src/batching_mode/utils/test.rs +++ b/src/flow/src/batching_mode/utils/test.rs @@ -268,7 +268,7 @@ async fn test_sql_plan_convert() { let new_sql = df_plan_to_sql(&new).unwrap(); assert_eq!( - r#"SELECT `UPPERCASE_NUMBERS_WITH_TS`.`NUMBER` FROM `UPPERCASE_NUMBERS_WITH_TS`"#, + r#"SELECT "UPPERCASE_NUMBERS_WITH_TS"."NUMBER" FROM "UPPERCASE_NUMBERS_WITH_TS""#, new_sql ); } @@ -1985,3 +1985,78 @@ async fn test_gen_plan_with_matching_schema_last_non_null_rejects_extra_flow_col "{err}" ); } + +#[test] +fn test_df_plan_to_sql_quotes_colon_table_name() { + // Prometheus-style table names contain ':' (e.g. + // `kube_pod_cpu_cores:sum`). The unparser dialect must quote them, + // otherwise the re-parsed SQL is invalid (`keyword: :`). + let table = single_row_u32_table("kube_pod_cpu_cores:sum", vec!["value"]); + let provider = Arc::new(DfTableProviderAdapter::new(table)); + let table_source = Arc::new(DefaultTableSource::new(provider)); + let table_ref = TableReference::full("catalog", "schema", "kube_pod_cpu_cores:sum"); + let plan = LogicalPlanBuilder::scan(table_ref, table_source, None) + .unwrap() + .build() + .unwrap(); + + let sql = df_plan_to_sql(&plan).unwrap(); + assert!( + sql.contains("\"kube_pod_cpu_cores:sum\""), + "expected quoted table name in {sql}" + ); + // The only occurrence of `cores:sum` must be inside the quoted identifier. + assert_eq!( + sql.matches("cores:sum").count(), + 1, + "colon should only appear inside quotes in {sql}" + ); +} + +#[test] +fn test_df_plan_to_sql_does_not_quote_plain_lowercase() { + let table = single_row_u32_table("plain_table", vec!["value"]); + let provider = Arc::new(DfTableProviderAdapter::new(table)); + let table_source = Arc::new(DefaultTableSource::new(provider)); + let table_ref = TableReference::full("catalog", "schema", "plain_table"); + let plan = LogicalPlanBuilder::scan(table_ref, table_source, None) + .unwrap() + .project(vec![datafusion_expr::col("value")]) + .unwrap() + .build() + .unwrap(); + + let sql = df_plan_to_sql(&plan).unwrap(); + assert!( + sql.contains("plain_table") && !sql.contains("\"plain_table\""), + "plain lowercase table should stay unquoted in {sql}" + ); + // `value` is not a reserved word, so the column stays unquoted. + assert!( + sql.contains("plain_table.value"), + "column unquoted in {sql}" + ); + assert!(!sql.contains('`'), "no backtick quoting in {sql}"); +} + +#[test] +fn test_df_plan_to_sql_quotes_digit_leading_table_name() { + // A table literally named `123metrics` starts with a digit and must be + // quoted, otherwise the re-parsed SQL is invalid. + let table = single_row_u32_table("123metrics", vec!["value"]); + let provider = Arc::new(DfTableProviderAdapter::new(table)); + let table_source = Arc::new(DefaultTableSource::new(provider)); + let table_ref = TableReference::full("catalog", "schema", "123metrics"); + let plan = LogicalPlanBuilder::scan(table_ref, table_source, None) + .unwrap() + .project(vec![datafusion_expr::col("value")]) + .unwrap() + .build() + .unwrap(); + + let sql = df_plan_to_sql(&plan).unwrap(); + assert!( + sql.contains("\"123metrics\""), + "expected digit-leading table name quoted in {sql}" + ); +} diff --git a/src/meta-srv/src/handler/flow_state_handler.rs b/src/meta-srv/src/handler/flow_state_handler.rs index a13c18a9e03..b4abbd7e376 100644 --- a/src/meta-srv/src/handler/flow_state_handler.rs +++ b/src/meta-srv/src/handler/flow_state_handler.rs @@ -14,12 +14,27 @@ use api::v1::meta::{FlowStat, HeartbeatRequest, Role}; use common_meta::key::flow::flow_state::{FlowStateManager, FlowStateValue}; +use common_telemetry::debug; use snafu::ResultExt; use crate::error::{FlowStateHandlerSnafu, Result}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::Context; +/// Extracts the flownode identity from a heartbeat request. +/// +/// Prefers `header.member_id` — the canonical identity metasrv uses for +/// flownodes (see `get_node_id` in `src/meta-srv/src/service/heartbeat.rs`). +/// Falls back to `peer.id`, the operator-configured node id, which is also +/// unique within the cluster (a flownode requires `node_id` in its config, see +/// `src/cmd/src/flownode.rs`). Returns `None` when neither is present. +fn node_identity(req: &HeartbeatRequest) -> Option { + req.header + .as_ref() + .map(|header| header.member_id) + .or_else(|| req.peer.as_ref().map(|peer| peer.id)) +} + pub struct FlowStateHandler { flow_state_manager: FlowStateManager, } @@ -60,11 +75,141 @@ impl HeartbeatHandler for FlowStateHandler { // mode until a follow-up PR adds heartbeat propagation. let value: FlowStateValue = FlowStateValue::new(state_size, last_exec_time_map, Default::default()); - self.flow_state_manager - .put(value) - .await - .context(FlowStateHandlerSnafu)?; + match node_identity(req) { + Some(node_id) => { + // Merge by node so that reports from different flownodes + // don't overwrite each other in the global state. + self.flow_state_manager + .merge(node_id, value) + .await + .context(FlowStateHandlerSnafu)?; + } + // No usable identity in the request: ignore the report instead + // of falling back to a whole-map replace, which would clobber + // other nodes' reports. + // Normal flownodes always carry header.member_id/peer.id; a + // report without either indicates an old or malformed client. + // Log at debug to avoid an anomalous sender spamming warn. + None => { + debug!( + "Ignore flow state report without node identity (no header.member_id and no peer.id): {value:?}" + ); + } + } } Ok(HandleControl::Continue) } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use api::v1::meta::{HeartbeatRequest, Peer, RequestHeader, Role}; + + use super::*; + use crate::handler::test_utils::TestEnv; + + #[test] + fn test_node_identity_prefers_header_member_id() { + let req = HeartbeatRequest { + header: Some(RequestHeader::new(42, Role::Flownode, HashMap::new())), + peer: Some(Peer { + id: 99, + addr: "127.0.0.1:4001".to_string(), + }), + ..Default::default() + }; + assert_eq!(node_identity(&req), Some(42)); + } + + #[test] + fn test_node_identity_falls_back_to_peer_id() { + let req = HeartbeatRequest { + header: None, + peer: Some(Peer { + id: 99, + addr: "127.0.0.1:4001".to_string(), + }), + ..Default::default() + }; + assert_eq!(node_identity(&req), Some(99)); + } + + #[test] + fn test_node_identity_none_without_header_and_peer() { + let req = HeartbeatRequest::default(); + assert_eq!(node_identity(&req), None); + } + + fn flow_stat() -> api::v1::meta::FlowStat { + api::v1::meta::FlowStat { + flow_stat_size: HashMap::from([(1, 1024)]), + flow_last_exec_time_map: HashMap::from([(1, 100)]), + } + } + + #[tokio::test] + async fn test_handle_merges_reports_from_different_nodes() { + let env = TestEnv::new(); + let ctx = env.ctx(); + let flow_state_manager = FlowStateManager::new(ctx.in_memory.clone().as_kv_backend_ref()); + let handler = FlowStateHandler::new(flow_state_manager); + + // Node 42 (header.member_id) reports flow 1. + let req_a = HeartbeatRequest { + header: Some(RequestHeader::new(42, Role::Flownode, HashMap::new())), + peer: Some(Peer { + id: 42, + addr: "127.0.0.1:4001".to_string(), + }), + flow_stat: Some(flow_stat()), + ..Default::default() + }; + let mut ctx = env.ctx(); + let mut acc = HeartbeatAccumulator::default(); + handler.handle(&req_a, &mut ctx, &mut acc).await.unwrap(); + + // Node 7 (only peer.id present) reports flow 2. + let req_b = HeartbeatRequest { + header: None, + peer: Some(Peer { + id: 7, + addr: "127.0.0.1:4007".to_string(), + }), + flow_stat: Some(api::v1::meta::FlowStat { + flow_stat_size: HashMap::from([(2, 2048)]), + flow_last_exec_time_map: HashMap::from([(2, 200)]), + }), + ..Default::default() + }; + let mut ctx = env.ctx(); + let mut acc = HeartbeatAccumulator::default(); + handler.handle(&req_b, &mut ctx, &mut acc).await.unwrap(); + + let value = handler.flow_state_manager.get().await.unwrap().unwrap(); + assert_eq!(value.last_exec_time_map.get(&1), Some(&100)); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_handle_ignores_report_without_identity() { + let env = TestEnv::new(); + let ctx = env.ctx(); + let flow_state_manager = FlowStateManager::new(ctx.in_memory.clone().as_kv_backend_ref()); + let handler = FlowStateHandler::new(flow_state_manager); + + // No header and no peer: the report must be ignored and no KV written. + let req = HeartbeatRequest { + header: None, + peer: None, + flow_stat: Some(flow_stat()), + ..Default::default() + }; + let mut ctx = env.ctx(); + let mut acc = HeartbeatAccumulator::default(); + handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); + + assert!(handler.flow_state_manager.get().await.unwrap().is_none()); + } +} From 6cd19e6eeaa4543cff965d49c92c6d55cdfb0ae0 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Fri, 7 Aug 2026 18:16:30 +0800 Subject: [PATCH 021/126] ci: stop running enterprise tests in OSS CI (#8795) Signed-off-by: Lei, HUANG --- .github/workflows/rust.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 1302f00201c..b5bfa1a2968 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -139,9 +139,7 @@ jobs: run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: | - cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index - cargo nextest run -p tests-integration --test main --features enterprise -E 'test(=table_ddl_event::test_table_ddl_procedure_events)' + run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 From 1f1c9270a806658d2c3d9700c8cd1445d687f89f Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 7 Aug 2026 20:37:58 +0800 Subject: [PATCH 022/126] feat(event): add event context to procedure events (#8734) * feat(event): add trigger context to procedure events Signed-off-by: WenyXu * test(event): fix trigger context event contracts Signed-off-by: WenyXu * fix(event): preserve trigger context origins Signed-off-by: WenyXu * fix(event): preserve lifecycle trigger contexts Signed-off-by: WenyXu * fix(meta): gate enterprise trigger imports Signed-off-by: WenyXu * test(event): assert trigger contexts exactly Signed-off-by: WenyXu * fix(test): restore migration test literals Signed-off-by: WenyXu * feat(event): add trigger context to non-table ddl Signed-off-by: WenyXu * refactor(event): simplify trigger context encoding Signed-off-by: WenyXu * fix(event): preserve auto alter trigger reason Signed-off-by: WenyXu * refactor(event): require explicit trigger context Signed-off-by: WenyXu * refactor(event): derive trigger context at ddl boundary Signed-off-by: WenyXu * refactor(event): propagate DDL trigger reasons Signed-off-by: WenyXu * fix(meta): resolve alter table rebase conflict Signed-off-by: WenyXu * test(meta): fix alter table trigger context setup Signed-off-by: WenyXu * refactor(meta): pass trigger context to region migration Signed-off-by: WenyXu * test(event): fix unknown trigger context assertions Signed-off-by: WenyXu * fix(meta): centralize trigger context protocol mapping Signed-off-by: WenyXu * test(event): fix migration trigger context assertion Signed-off-by: WenyXu * refactor(procedure): rename event runtime context Signed-off-by: WenyXu * refactor(event): rename trigger context Signed-off-by: WenyXu * test(event): fix migration event context assertion Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 2 +- src/common/event-recorder/src/event_table.rs | 3 + src/common/meta/Cargo.toml | 1 + src/common/meta/src/ddl/alter_database.rs | 23 +- .../meta/src/ddl/alter_logical_tables.rs | 14 +- src/common/meta/src/ddl/alter_table.rs | 33 +- .../src/ddl/alter_table/region_request.rs | 8 +- src/common/meta/src/ddl/create_database.rs | 19 +- src/common/meta/src/ddl/create_flow.rs | 22 +- .../meta/src/ddl/create_logical_tables.rs | 18 +- src/common/meta/src/ddl/create_table.rs | 34 +- src/common/meta/src/ddl/create_view.rs | 14 +- src/common/meta/src/ddl/drop_database.rs | 25 +- src/common/meta/src/ddl/drop_flow.rs | 15 +- src/common/meta/src/ddl/drop_table.rs | 21 +- src/common/meta/src/ddl/drop_view.rs | 12 +- src/common/meta/src/ddl/event/database.rs | 25 +- src/common/meta/src/ddl/event/flow.rs | 28 +- src/common/meta/src/ddl/event/table.rs | 59 ++- src/common/meta/src/ddl/event/view.rs | 25 +- .../meta/src/ddl/purge_dropped_table.rs | 39 +- src/common/meta/src/ddl/test_util.rs | 9 +- .../src/ddl/tests/alter_logical_tables.rs | 49 ++- src/common/meta/src/ddl/tests/alter_table.rs | 72 +++- src/common/meta/src/ddl/tests/create_flow.rs | 87 ++++- .../src/ddl/tests/create_logical_tables.rs | 36 +- src/common/meta/src/ddl/tests/create_table.rs | 74 +++- src/common/meta/src/ddl/tests/create_view.rs | 30 +- .../meta/src/ddl/tests/drop_database.rs | 4 + src/common/meta/src/ddl/tests/drop_flow.rs | 14 +- src/common/meta/src/ddl/tests/drop_table.rs | 225 +++++++++--- src/common/meta/src/ddl/tests/drop_view.rs | 14 +- .../meta/src/ddl/tests/event/database.rs | 48 ++- src/common/meta/src/ddl/tests/event/flow.rs | 26 +- src/common/meta/src/ddl/tests/event/table.rs | 159 ++++++-- .../meta/src/ddl/tests/event/test_util.rs | 20 +- src/common/meta/src/ddl/tests/event/view.rs | 60 ++- src/common/meta/src/ddl/truncate_table.rs | 22 +- src/common/meta/src/ddl/undrop_table.rs | 48 ++- src/common/meta/src/ddl_manager.rs | 203 +++++++---- src/common/meta/src/rpc/ddl.rs | 131 ++++++- src/common/procedure/src/lib.rs | 6 +- src/common/procedure/src/local.rs | 12 +- src/common/procedure/src/local/runner.rs | 6 +- src/common/procedure/src/procedure.rs | 8 +- src/common/session/src/lib.rs | 49 ++- src/flow/src/adapter/util.rs | 3 +- src/frontend/src/instance/dashboard.rs | 8 +- src/frontend/src/instance/grpc.rs | 10 +- .../src/instance/otlp/trace_ingest.rs | 3 +- src/frontend/src/instance/prom_store.rs | 9 +- src/meta-srv/Cargo.toml | 1 - src/meta-srv/src/event/gc.rs | 66 +++- src/meta-srv/src/event/region_migration.rs | 18 +- src/meta-srv/src/event/repartition.rs | 39 +- src/meta-srv/src/gc/ctx.rs | 6 + src/meta-srv/src/gc/handler.rs | 10 +- src/meta-srv/src/gc/mock.rs | 4 + src/meta-srv/src/gc/mock/basic.rs | 72 +++- src/meta-srv/src/gc/mock/concurrent.rs | 30 +- src/meta-srv/src/gc/mock/full_list.rs | 6 + src/meta-srv/src/gc/mock/misc.rs | 15 +- src/meta-srv/src/gc/procedure.rs | 45 ++- src/meta-srv/src/gc/scheduler.rs | 22 +- .../src/procedure/region_migration.rs | 98 ++++- .../downgrade_leader_region.rs | 4 +- .../src/procedure/region_migration/manager.rs | 5 +- .../procedure/region_migration/test_util.rs | 7 +- .../upgrade_candidate_region.rs | 4 +- src/meta-srv/src/procedure/repartition.rs | 42 ++- .../src/procedure/repartition/collect.rs | 3 + .../src/procedure/repartition/group.rs | 21 +- .../src/procedure/repartition/test_util.rs | 2 + src/meta-srv/src/procedure/tests.rs | 14 +- src/meta-srv/src/procedure/wal_prune.rs | 25 +- src/meta-srv/src/service/procedure.rs | 1 - src/operator/src/insert.rs | 23 +- src/operator/src/procedure.rs | 19 +- src/operator/src/statement/comment.rs | 22 +- src/operator/src/statement/ddl.rs | 121 ++++-- src/operator/src/utils.rs | 43 ++- src/pipeline/src/manager/pipeline_operator.rs | 3 +- src/session/src/context.rs | 19 +- src/standalone/src/procedure.rs | 2 + tests-integration/src/tests/gc.rs | 2 + tests-integration/tests/database_ddl_event.rs | 11 +- tests-integration/tests/flow_ddl_event.rs | 18 + tests-integration/tests/gc_event.rs | 18 +- tests-integration/tests/http.rs | 42 ++- tests-integration/tests/region_migration.rs | 16 +- tests-integration/tests/repartition.rs | 2 + tests-integration/tests/table_ddl_event.rs | 344 +++++++++++++++++- tests-integration/tests/view_ddl_event.rs | 18 + 93 files changed, 2546 insertions(+), 622 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6c610726ef8..476133db9e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2642,6 +2642,7 @@ dependencies = [ "common-query", "common-recordbatch", "common-runtime", + "common-session", "common-telemetry", "common-test-util", "common-time", @@ -8401,7 +8402,6 @@ dependencies = [ "serde", "serde_json", "servers", - "session", "snafu 0.8.6", "sqlx", "store-api", diff --git a/src/common/event-recorder/src/event_table.rs b/src/common/event-recorder/src/event_table.rs index 40d6625b5d9..0ec38c8eabc 100644 --- a/src/common/event-recorder/src/event_table.rs +++ b/src/common/event-recorder/src/event_table.rs @@ -79,6 +79,9 @@ pub const TYPE_COLUMN: EventTableColumn = /// The canonical event payload column. pub const PAYLOAD_COLUMN: EventTableColumn = EventTableColumn::json_binary("payload", ColumnDataType::Binary, SemanticType::Field); +/// The optional context describing why an event was triggered. +pub const EVENT_CONTEXT_COLUMN: EventTableColumn = + EventTableColumn::json_binary("event_context", ColumnDataType::Binary, SemanticType::Field); /// The canonical event timestamp column. pub const TIMESTAMP_COLUMN: EventTableColumn = EventTableColumn::new( "timestamp", diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 8e0cbbdf354..43714b7d599 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -41,6 +41,7 @@ common-procedure.workspace = true common-query.workspace = true common-recordbatch.workspace = true common-runtime.workspace = true +common-session.workspace = true common-telemetry.workspace = true common-time.workspace = true common-wal.workspace = true diff --git a/src/common/meta/src/ddl/alter_database.rs b/src/common/meta/src/ddl/alter_database.rs index dfa42f4009c..7e16275da3c 100644 --- a/src/common/meta/src/ddl/alter_database.rs +++ b/src/common/meta/src/ddl/alter_database.rs @@ -15,7 +15,7 @@ use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Status, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Status, }; use common_telemetry::tracing::info; use serde::{Deserialize, Serialize}; @@ -32,7 +32,7 @@ use crate::key::DeserializedValueWithBytes; use crate::key::schema_name::{SchemaName, SchemaNameKey, SchemaNameValue}; use crate::lock_key::{CatalogLock, SchemaLock}; use crate::rpc::ddl::UnsetDatabaseOption::{self}; -use crate::rpc::ddl::{AlterDatabaseKind, AlterDatabaseTask, SetDatabaseOption}; +use crate::rpc::ddl::{AlterDatabaseKind, AlterDatabaseTask, EventContext, SetDatabaseOption}; pub struct AlterDatabaseProcedure { pub context: DdlContext, @@ -73,10 +73,14 @@ fn build_new_schema_value( impl AlterDatabaseProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::AlterDatabase"; - pub fn new(task: AlterDatabaseTask, context: DdlContext) -> Result { + pub fn new( + task: AlterDatabaseTask, + event_context: EventContext, + context: DdlContext, + ) -> Result { Ok(Self { context, - data: AlterDatabaseData::new(task)?, + data: AlterDatabaseData::new(task, event_context)?, }) } @@ -176,7 +180,10 @@ impl Procedure for AlterDatabaseProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx.event_type_filter.allows(ALTER_DATABASE_EVENT_TYPE) { return None; } @@ -186,6 +193,7 @@ impl Procedure for AlterDatabaseProcedure { self.data.catalog(), self.data.schema(), &self.data.kind, + self.data.event_context.clone(), ) } else { DatabaseDdlEvent::alter_lifecycle() @@ -209,16 +217,19 @@ pub struct AlterDatabaseData { catalog_name: String, schema_name: String, schema_value: Option>, + #[serde(default)] + event_context: EventContext, } impl AlterDatabaseData { - pub fn new(task: AlterDatabaseTask) -> Result { + pub fn new(task: AlterDatabaseTask, event_context: EventContext) -> Result { Ok(Self { state: AlterDatabaseState::Prepare, kind: AlterDatabaseKind::try_from(task.alter_expr.kind.unwrap())?, catalog_name: task.alter_expr.catalog_name, schema_name: task.alter_expr.schema_name, schema_value: None, + event_context, }) } diff --git a/src/common/meta/src/ddl/alter_logical_tables.rs b/src/common/meta/src/ddl/alter_logical_tables.rs index 683248965ff..816db0a6afb 100644 --- a/src/common/meta/src/ddl/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/alter_logical_tables.rs @@ -20,7 +20,7 @@ use api::region::RegionResponse; use async_trait::async_trait; use common_catalog::format_full_table_name; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; -use common_procedure::{Context, EventContext, EventTrigger, LockKey, Procedure, Status}; +use common_procedure::{Context, EventRuntimeContext, EventTrigger, LockKey, Procedure, Status}; use common_telemetry::{debug, error, info, warn}; pub use executor::make_alter_region_request; use serde::{Deserialize, Serialize}; @@ -47,7 +47,7 @@ use crate::key::table_info::TableInfoValue; use crate::key::table_route::PhysicalTableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; use crate::metrics; -use crate::rpc::ddl::AlterTableTask; +use crate::rpc::ddl::{AlterTableTask, EventContext}; use crate::rpc::router::RegionRoute; pub struct AlterLogicalTablesProcedure { @@ -88,6 +88,7 @@ impl AlterLogicalTablesProcedure { pub fn new( tasks: Vec, physical_table_id: TableId, + event_context: EventContext, context: DdlContext, ) -> Self { Self { @@ -100,6 +101,7 @@ impl AlterLogicalTablesProcedure { physical_table_info: None, physical_columns: vec![], table_cache_keys_to_invalidate: vec![], + event_context, }, physical_table_route: None, } @@ -320,7 +322,10 @@ impl Procedure for AlterLogicalTablesProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::AlterLogicalTables.as_str()) @@ -348,6 +353,7 @@ impl Procedure for AlterLogicalTablesProcedure { locators, self.data.tasks.len(), kinds, + self.data.event_context.clone(), ))) } } @@ -364,6 +370,8 @@ pub struct AlterTablesData { physical_table_info: Option>, physical_columns: Vec, table_cache_keys_to_invalidate: Vec, + #[serde(default)] + event_context: EventContext, } impl AlterTablesData { diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index 59fd2af2985..9504e8672ca 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -27,7 +27,7 @@ use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE}; use common_error::ext::BoxedError; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, ContextProvider, Error as ProcedureError, EventContext, + Context as ProcedureContext, ContextProvider, Error as ProcedureError, EventRuntimeContext, EventTrigger, LockKey, PoisonKey, PoisonKeys, Procedure, ProcedureId, Status, StringKey, }; use common_telemetry::{error, info, warn}; @@ -58,7 +58,7 @@ use crate::key::{DeserializedValueWithBytes, RegionDistribution}; use crate::lock_key::{CatalogLock, RegionLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; use crate::poison_key::table_poison_key; -use crate::rpc::ddl::AlterTableTask; +use crate::rpc::ddl::{AlterTableTask, EventContext}; use crate::rpc::router::{RegionRoute, find_leaders, region_distribution}; /// The alter table procedure @@ -97,18 +97,24 @@ fn build_executor_from_alter_expr(alter_data: &AlterTableData) -> AlterTableExec impl AlterTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::AlterTable"; - pub fn new(table_id: TableId, task: AlterTableTask, context: DdlContext) -> Result { - Self::new_with_region_locks(table_id, task, vec![], context) + pub fn new( + table_id: TableId, + task: AlterTableTask, + event_context: EventContext, + context: DdlContext, + ) -> Result { + Self::new_with_region_locks(table_id, task, event_context, vec![], context) } pub(crate) fn new_with_region_locks( table_id: TableId, task: AlterTableTask, + event_context: EventContext, region_locks: Vec, context: DdlContext, ) -> Result { task.validate()?; - let data = AlterTableData::new(task, table_id, region_locks); + let data = AlterTableData::new(task, table_id, event_context, region_locks); let executor = build_executor_from_alter_expr(&data); Ok(Self { context, @@ -558,7 +564,10 @@ impl Procedure for AlterTableProcedure { PoisonKeys::new(vec![self.table_poison_key()]) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::AlterTable.as_str()) @@ -578,7 +587,7 @@ impl Procedure for AlterTableProcedure { .kind .as_ref() .and_then(alter_table_kind_name); - TableDdlEvent::alter_table_submitted(locator, kind) + TableDdlEvent::alter_table_submitted(locator, kind, self.data.event_context.clone()) } _ => TableDdlEvent::lifecycle(TableDdlEventType::AlterTable), }; @@ -606,6 +615,8 @@ pub struct AlterTableData { task: AlterTableTask, table_id: TableId, #[serde(default)] + event_context: EventContext, + #[serde(default)] column_metadatas: Vec, /// Table info value before alteration. table_info_value: Option>, @@ -617,11 +628,17 @@ pub struct AlterTableData { } impl AlterTableData { - pub fn new(task: AlterTableTask, table_id: TableId, region_locks: Vec) -> Self { + pub fn new( + task: AlterTableTask, + table_id: TableId, + event_context: EventContext, + region_locks: Vec, + ) -> Self { Self { state: AlterTableState::Prepare, task, table_id, + event_context, column_metadatas: vec![], table_info_value: None, region_distribution: None, diff --git a/src/common/meta/src/ddl/alter_table/region_request.rs b/src/common/meta/src/ddl/alter_table/region_request.rs index 73edd4c3615..ef8d1487c70 100644 --- a/src/common/meta/src/ddl/alter_table/region_request.rs +++ b/src/common/meta/src/ddl/alter_table/region_request.rs @@ -144,7 +144,7 @@ mod tests { }; use crate::key::table_route::TableRouteValue; use crate::peer::Peer; - use crate::rpc::ddl::AlterTableTask; + use crate::rpc::ddl::{AlterTableTask, EventContext}; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -240,7 +240,8 @@ mod tests { }, }; - let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context).unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let alter_kind = procedure.make_region_alter_kind().unwrap(); let Some(Body::Alter(alter_region_request)) = @@ -296,7 +297,8 @@ mod tests { }, }; - let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context).unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let alter_kind = procedure.make_region_alter_kind().unwrap(); let Some(Body::Alter(alter_region_request)) = diff --git a/src/common/meta/src/ddl/create_database.rs b/src/common/meta/src/ddl/create_database.rs index b5c50235c63..8f9c5bee1a5 100644 --- a/src/common/meta/src/ddl/create_database.rs +++ b/src/common/meta/src/ddl/create_database.rs @@ -19,8 +19,8 @@ use async_trait::async_trait; use common_error::ext::{BoxedError, ErrorExt}; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureId, - Status, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + ProcedureId, Status, }; use serde::{Deserialize, Serialize}; use serde_with::{DefaultOnNull, serde_as}; @@ -34,7 +34,7 @@ use crate::error::{self, Result}; use crate::instruction::{CacheIdent, UserCacheIdent}; use crate::key::schema_name::{SchemaNameKey, SchemaNameValue}; use crate::lock_key::{CatalogLock, SchemaLock}; -use crate::rpc::ddl::CreatorGrantIntent; +use crate::rpc::ddl::{CreatorGrantIntent, EventContext}; /// Describes the creator-access result of an atomic create. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -93,6 +93,7 @@ impl CreateDatabaseProcedure { create_if_not_exists: bool, options: HashMap, creator: Option, + event_context: EventContext, context: DdlContext, ) -> Self { Self { @@ -104,6 +105,7 @@ impl CreateDatabaseProcedure { create_if_not_exists, options, creator, + event_context, }, } } @@ -278,7 +280,10 @@ impl Procedure for CreateDatabaseProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx.event_type_filter.allows(CREATE_DATABASE_EVENT_TYPE) { return None; } @@ -289,6 +294,7 @@ impl Procedure for CreateDatabaseProcedure { &self.data.schema, self.data.create_if_not_exists, &self.data.options, + self.data.event_context.clone(), ) } else { DatabaseDdlEvent::create_lifecycle() @@ -321,6 +327,8 @@ pub struct CreateDatabaseData { /// Authenticated creator whose access is ensured, absent for legacy schema-only requests. #[serde(default)] pub creator: Option, + #[serde(default)] + pub event_context: EventContext, } #[cfg(test)] @@ -415,6 +423,7 @@ mod tests { username: "alice".to_string(), created_at_ns: 1, }), + EventContext::default(), context, ); (procedure, committer) @@ -574,6 +583,7 @@ mod tests { false, HashMap::new(), None, + EventContext::default(), context, ); procedure.data.state = CreateDatabaseState::CreateMetadata; @@ -609,6 +619,7 @@ mod tests { username: "alice".to_string(), created_at_ns: 1, }), + EventContext::default(), context, ); procedure.data.state = CreateDatabaseState::CreateMetadata; diff --git a/src/common/meta/src/ddl/create_flow.rs b/src/common/meta/src/ddl/create_flow.rs index 489e4f760ce..66083f5949c 100644 --- a/src/common/meta/src/ddl/create_flow.rs +++ b/src/common/meta/src/ddl/create_flow.rs @@ -24,8 +24,8 @@ use async_trait::async_trait; use common_catalog::format_full_flow_name; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureState, - Result as ProcedureResult, Status, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + ProcedureState, Result as ProcedureResult, Status, }; use common_telemetry::info; use common_telemetry::tracing_context::TracingContext; @@ -50,7 +50,7 @@ use crate::key::{DeserializedValueWithBytes, FlowId, FlowPartitionId}; use crate::lock_key::{CatalogLock, FlowNameLock}; use crate::metrics; use crate::peer::Peer; -use crate::rpc::ddl::{CreateFlowTask, FlowQueryContext, QueryContext}; +use crate::rpc::ddl::{CreateFlowTask, EventContext, FlowQueryContext, QueryContext}; /// The procedure of flow creation. pub struct CreateFlowProcedure { @@ -62,7 +62,12 @@ impl CreateFlowProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::CreateFlow"; /// Returns a new [CreateFlowProcedure]. - pub fn new(task: CreateFlowTask, query_context: QueryContext, context: DdlContext) -> Self { + pub fn new( + task: CreateFlowTask, + query_context: QueryContext, + event_context: EventContext, + context: DdlContext, + ) -> Self { Self { context, data: CreateFlowData { @@ -76,6 +81,7 @@ impl CreateFlowProcedure { prev_flow_info_value: None, did_replace: false, flow_type: None, + event_context, }, } } @@ -404,7 +410,10 @@ impl Procedure for CreateFlowProcedure { ]) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx.event_type_filter.allows(CREATE_FLOW_EVENT_TYPE) { return None; } @@ -419,6 +428,7 @@ impl Procedure for CreateFlowProcedure { expire_after: self.data.task.expire_after, eval_interval_secs: self.data.task.eval_interval_secs, }, + self.data.event_context.clone(), ), EventTrigger::Succeeded => { let flow_id = match ctx.lifecycle_state { @@ -701,6 +711,8 @@ pub struct CreateFlowData { #[serde(default)] pub(crate) did_replace: bool, pub(crate) flow_type: Option, + #[serde(default)] + pub(crate) event_context: EventContext, } impl CreateFlowData { diff --git a/src/common/meta/src/ddl/create_logical_tables.rs b/src/common/meta/src/ddl/create_logical_tables.rs index 08f6cb855c1..8b93ecfcc0f 100644 --- a/src/common/meta/src/ddl/create_logical_tables.rs +++ b/src/common/meta/src/ddl/create_logical_tables.rs @@ -24,8 +24,8 @@ use common_catalog::consts::METRIC_ENGINE; use common_event_recorder::Event; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureState, - Status, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + ProcedureState, Status, }; use common_telemetry::{debug, error, warn}; use futures::future; @@ -48,7 +48,7 @@ use crate::error::Result; use crate::key::table_route::TableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; -use crate::rpc::ddl::CreateTableTask; +use crate::rpc::ddl::{CreateTableTask, EventContext}; use crate::rpc::router::{RegionRoute, find_leaders}; pub struct CreateLogicalTablesProcedure { @@ -62,6 +62,7 @@ impl CreateLogicalTablesProcedure { pub fn new( tasks: Vec, physical_table_id: TableId, + event_context: EventContext, context: DdlContext, ) -> Self { Self { @@ -74,6 +75,7 @@ impl CreateLogicalTablesProcedure { physical_region_numbers: vec![], physical_columns: vec![], physical_partition_columns: vec![], + event_context, }, } } @@ -258,7 +260,7 @@ impl Procedure for CreateLogicalTablesProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::CreateLogicalTables.as_str()) @@ -275,7 +277,11 @@ impl Procedure for CreateLogicalTablesProcedure { ) .with_physical_table_id(self.data.physical_table_id) }); - TableDdlEvent::create_logical_tables_submitted(locators, self.data.tasks.len()) + TableDdlEvent::create_logical_tables_submitted( + locators, + self.data.tasks.len(), + self.data.event_context.clone(), + ) } EventTrigger::Succeeded => match ctx.lifecycle_state { ProcedureState::Done { @@ -321,6 +327,8 @@ pub struct CreateTablesData { physical_region_numbers: Vec, physical_columns: Vec, physical_partition_columns: Vec, + #[serde(default)] + event_context: EventContext, } impl CreateTablesData { diff --git a/src/common/meta/src/ddl/create_table.rs b/src/common/meta/src/ddl/create_table.rs index 298fa8e56a7..c1690416a2e 100644 --- a/src/common/meta/src/ddl/create_table.rs +++ b/src/common/meta/src/ddl/create_table.rs @@ -23,8 +23,8 @@ use common_procedure::error::{ }; use common_procedure::local::DynamicKeyLockGuard; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureId, - ProcedureState, Status, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + ProcedureId, ProcedureState, Status, }; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -47,7 +47,7 @@ use crate::lock_key::{CatalogLock, SchemaLock, TableNameLock}; use crate::metrics; use crate::peer::PeerAllocContext; use crate::region_keeper::OperatingRegionGuard; -use crate::rpc::ddl::{CreateTableTask, QueryContext}; +use crate::rpc::ddl::{CreateTableTask, EventContext, QueryContext}; use crate::rpc::router::{RegionRoute, operating_leader_region_roles}; use crate::wal_provider::{ RegionWalOptions, acquire_remote_wal_read_locks, optional_region_wal_options_serde, @@ -84,20 +84,17 @@ fn build_executor_from_create_table_data( impl CreateTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::CreateTable"; - pub fn new(task: CreateTableTask, context: DdlContext) -> Result { - Self::new_with_query_context(task, QueryContext::default(), context) - } - - pub fn new_with_query_context( + pub fn new( task: CreateTableTask, query_context: QueryContext, + event_context: EventContext, context: DdlContext, ) -> Result { let executor = build_executor_from_create_table_data(&task.create_table)?; Ok(Self { context, - data: CreateTableData::new(task, query_context), + data: CreateTableData::new_with_query_context(task, event_context, query_context), opening_regions: vec![], executor, remote_wal_lock_guards: vec![], @@ -399,7 +396,10 @@ impl Procedure for CreateTableProcedure { ]) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::CreateTable.as_str()) @@ -416,6 +416,7 @@ impl Procedure for CreateTableProcedure { locator, create_table.create_if_not_exists, &create_table.engine, + self.data.event_context.clone(), ) } EventTrigger::Succeeded => match ctx.lifecycle_state { @@ -452,6 +453,8 @@ pub struct CreateTableData { #[serde(default)] pub query_context: QueryContext, #[serde(default)] + pub event_context: EventContext, + #[serde(default)] pub column_metadatas: Vec, /// None stands for not allocated yet. pub(crate) table_route: Option, @@ -462,12 +465,21 @@ pub struct CreateTableData { } impl CreateTableData { - pub fn new(task: CreateTableTask, query_context: QueryContext) -> Self { + pub fn new(task: CreateTableTask, event_context: EventContext) -> Self { + Self::new_with_query_context(task, event_context, QueryContext::default()) + } + + pub fn new_with_query_context( + task: CreateTableTask, + event_context: EventContext, + query_context: QueryContext, + ) -> Self { CreateTableData { state: CreateTableState::Prepare, column_metadatas: vec![], task, query_context, + event_context, table_route: None, region_wal_options: None, } diff --git a/src/common/meta/src/ddl/create_view.rs b/src/common/meta/src/ddl/create_view.rs index 1a6b85d55aa..0cb622774f0 100644 --- a/src/common/meta/src/ddl/create_view.rs +++ b/src/common/meta/src/ddl/create_view.rs @@ -16,8 +16,8 @@ use async_trait::async_trait; use common_event_recorder::Event; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureState, - Status, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + ProcedureState, Status, }; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -35,7 +35,7 @@ use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; use crate::lock_key::{CatalogLock, SchemaLock, TableNameLock}; use crate::metrics; -use crate::rpc::ddl::CreateViewTask; +use crate::rpc::ddl::{CreateViewTask, EventContext}; // The procedure to execute `[CreateViewTask]`. pub struct CreateViewProcedure { @@ -46,13 +46,14 @@ pub struct CreateViewProcedure { impl CreateViewProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::CreateView"; - pub fn new(task: CreateViewTask, context: DdlContext) -> Self { + pub fn new(task: CreateViewTask, event_context: EventContext, context: DdlContext) -> Self { Self { context, data: CreateViewData { state: CreateViewState::Prepare, task, need_update: false, + event_context, }, } } @@ -271,7 +272,7 @@ impl Procedure for CreateViewProcedure { ]) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { if !ctx.event_type_filter.allows(CREATE_VIEW_EVENT_TYPE) { return None; } @@ -289,6 +290,7 @@ impl Procedure for CreateViewProcedure { referenced_table_count: self.data.task.table_names().len(), column_count: self.data.task.columns().len(), }, + self.data.event_context.clone(), ) } EventTrigger::Succeeded => match ctx.lifecycle_state { @@ -321,6 +323,8 @@ pub struct CreateViewData { pub task: CreateViewTask, /// Whether to update the view info. pub need_update: bool, + #[serde(default)] + pub event_context: EventContext, } impl CreateViewData { diff --git a/src/common/meta/src/ddl/drop_database.rs b/src/common/meta/src/ddl/drop_database.rs index 6b6a6852daa..d84a5e6dc86 100644 --- a/src/common/meta/src/ddl/drop_database.rs +++ b/src/common/meta/src/ddl/drop_database.rs @@ -23,7 +23,7 @@ use std::fmt::Debug; use common_error::ext::BoxedError; use common_procedure::error::{ExternalSnafu, FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use futures::stream::BoxStream; @@ -38,11 +38,13 @@ use crate::ddl::utils::map_to_procedure_error; use crate::error::Result; use crate::key::table_name::TableNameValue; use crate::lock_key::{CatalogLock, SchemaLock}; +use crate::rpc::ddl::EventContext; pub struct DropDatabaseProcedure { /// The context of procedure runtime. runtime_context: DdlContext, context: DropDatabaseContext, + event_context: EventContext, state: Box, } @@ -85,9 +87,16 @@ pub(crate) trait State: Send + Debug { impl DropDatabaseProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropDatabase"; - pub fn new(catalog: String, schema: String, drop_if_exists: bool, context: DdlContext) -> Self { + pub fn new( + catalog: String, + schema: String, + drop_if_exists: bool, + event_context: EventContext, + context: DdlContext, + ) -> Self { Self { runtime_context: context, + event_context, context: DropDatabaseContext { catalog, schema, @@ -104,11 +113,13 @@ impl DropDatabaseProcedure { catalog, schema, drop_if_exists, + event_context, state, } = serde_json::from_str(json).context(FromJsonSnafu)?; Ok(Self { runtime_context, + event_context, context: DropDatabaseContext { catalog, schema, @@ -159,6 +170,7 @@ impl Procedure for DropDatabaseProcedure { catalog: &self.context.catalog, schema: &self.context.schema, drop_if_exists: self.context.drop_if_exists, + event_context: &self.event_context, state: self.state.as_ref(), }; @@ -174,7 +186,10 @@ impl Procedure for DropDatabaseProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx.event_type_filter.allows(DROP_DATABASE_EVENT_TYPE) { return None; } @@ -184,6 +199,7 @@ impl Procedure for DropDatabaseProcedure { &self.context.catalog, &self.context.schema, self.context.drop_if_exists, + self.event_context.clone(), ) } else { DatabaseDdlEvent::drop_lifecycle() @@ -199,6 +215,7 @@ struct DropDatabaseData<'a> { // The schema name schema: &'a str, drop_if_exists: bool, + event_context: &'a EventContext, state: &'a dyn State, } @@ -209,5 +226,7 @@ struct DropDatabaseOwnedData { // The schema name schema: String, drop_if_exists: bool, + #[serde(default)] + event_context: EventContext, state: Box, } diff --git a/src/common/meta/src/ddl/drop_flow.rs b/src/common/meta/src/ddl/drop_flow.rs index 1725b125508..badded386ad 100644 --- a/src/common/meta/src/ddl/drop_flow.rs +++ b/src/common/meta/src/ddl/drop_flow.rs @@ -21,7 +21,7 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; @@ -41,7 +41,7 @@ use crate::key::flow::flow_info::FlowInfoValue; use crate::key::flow::flow_route::FlowRouteValue; use crate::lock_key::{CatalogLock, FlowLock}; use crate::metrics; -use crate::rpc::ddl::DropFlowTask; +use crate::rpc::ddl::{DropFlowTask, EventContext}; /// The procedure for dropping a flow. pub struct DropFlowProcedure { @@ -54,7 +54,7 @@ pub struct DropFlowProcedure { impl DropFlowProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropFlow"; - pub fn new(task: DropFlowTask, context: DdlContext) -> Self { + pub fn new(task: DropFlowTask, event_context: EventContext, context: DdlContext) -> Self { Self { context, data: DropFlowData { @@ -62,6 +62,7 @@ impl DropFlowProcedure { task, flow_info_value: None, flow_route_values: vec![], + event_context, }, } } @@ -222,7 +223,10 @@ impl Procedure for DropFlowProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx.event_type_filter.allows(DROP_FLOW_EVENT_TYPE) { return None; } @@ -233,6 +237,7 @@ impl Procedure for DropFlowProcedure { &self.data.task.flow_name, self.data.task.flow_id, self.data.task.drop_if_exists, + self.data.event_context.clone(), ), _ => FlowDdlEvent::drop_lifecycle(), }; @@ -248,6 +253,8 @@ pub(crate) struct DropFlowData { task: DropFlowTask, pub(crate) flow_info_value: Option, pub(crate) flow_route_values: Vec, + #[serde(default)] + pub(crate) event_context: EventContext, } /// The state of drop flow diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs index 182b563239e..2dcfa2f7d82 100644 --- a/src/common/meta/src/ddl/drop_table.rs +++ b/src/common/meta/src/ddl/drop_table.rs @@ -22,8 +22,8 @@ use common_error::ext::BoxedError; use common_event_recorder::Event; use common_procedure::error::{ExternalSnafu, FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, - Procedure, Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, + LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; use common_telemetry::tracing::warn; @@ -48,7 +48,7 @@ use crate::key::table_route::TableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; use crate::region_keeper::OperatingRegionGuard; -use crate::rpc::ddl::DropTableTask; +use crate::rpc::ddl::{DropTableTask, EventContext}; use crate::rpc::router::{RegionRoute, operating_leader_region_roles}; #[cfg(feature = "enterprise")] @@ -74,11 +74,12 @@ pub struct DropTableProcedure { impl DropTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropTable"; - pub fn new(task: DropTableTask, context: DdlContext) -> Self { + pub fn new(task: DropTableTask, context: DdlContext, event_context: EventContext) -> Self { let data = DropTableData::new( task, cfg!(feature = "enterprise") && context.soft_drop_enabled, context.soft_drop_retention, + event_context, ); let executor = data.build_executor(); Self { @@ -413,7 +414,7 @@ impl Procedure for DropTableProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::DropTable.as_str()) @@ -425,7 +426,11 @@ impl Procedure for DropTableProcedure { let task = &self.data.task; let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) .with_table_id(task.table_id); - TableDdlEvent::drop_table_submitted(locator, task.drop_if_exists) + TableDdlEvent::drop_table_submitted( + locator, + task.drop_if_exists, + self.data.event_context.clone(), + ) } _ => TableDdlEvent::lifecycle(TableDdlEventType::DropTable), }; @@ -483,6 +488,8 @@ pub struct DropTableData { pub soft_drop_retention_millis: Option, #[serde(default)] pub drop_generation: Option, + #[serde(default)] + pub event_context: EventContext, } impl DropTableData { @@ -490,6 +497,7 @@ impl DropTableData { task: DropTableTask, soft_drop_enabled: bool, soft_drop_retention: Option, + event_context: EventContext, ) -> Self { Self { state: DropTableState::Prepare, @@ -503,6 +511,7 @@ impl DropTableData { retention_expires_at: None, soft_drop_retention_millis: soft_drop_retention_millis(soft_drop_retention), drop_generation: None, + event_context, } } diff --git a/src/common/meta/src/ddl/drop_view.rs b/src/common/meta/src/ddl/drop_view.rs index 3804df6862c..53a62e76159 100644 --- a/src/common/meta/src/ddl/drop_view.rs +++ b/src/common/meta/src/ddl/drop_view.rs @@ -16,7 +16,7 @@ use async_trait::async_trait; use common_event_recorder::Event; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; @@ -35,7 +35,7 @@ use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; use crate::metrics; -use crate::rpc::ddl::DropViewTask; +use crate::rpc::ddl::{DropViewTask, EventContext}; /// The procedure for dropping a view. pub struct DropViewProcedure { @@ -48,12 +48,13 @@ pub struct DropViewProcedure { impl DropViewProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropView"; - pub fn new(task: DropViewTask, context: DdlContext) -> Self { + pub fn new(task: DropViewTask, event_context: EventContext, context: DdlContext) -> Self { Self { context, data: DropViewData { state: DropViewState::Prepare, task, + event_context, }, } } @@ -213,7 +214,7 @@ impl Procedure for DropViewProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { if !ctx.event_type_filter.allows(DROP_VIEW_EVENT_TYPE) { return None; } @@ -227,6 +228,7 @@ impl Procedure for DropViewProcedure { table_ref.table, self.data.view_id(), self.data.task.drop_if_exists, + self.data.event_context.clone(), ) } _ => ViewDdlEvent::drop_lifecycle(), @@ -241,6 +243,8 @@ impl Procedure for DropViewProcedure { pub(crate) struct DropViewData { state: DropViewState, task: DropViewTask, + #[serde(default)] + event_context: EventContext, } impl DropViewData { diff --git a/src/common/meta/src/ddl/event/database.rs b/src/common/meta/src/ddl/event/database.rs index a3db0ebab28..9d70ea35a5e 100644 --- a/src/common/meta/src/ddl/event/database.rs +++ b/src/common/meta/src/ddl/event/database.rs @@ -19,13 +19,14 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, - SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, column_schemas, nullable_string, + CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, + SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, column_schemas, nullable_json, + nullable_string, }; use serde::Serialize; use snafu::ResultExt; -use crate::rpc::ddl::{AlterDatabaseKind, SetDatabaseOption, UnsetDatabaseOption}; +use crate::rpc::ddl::{AlterDatabaseKind, EventContext, SetDatabaseOption, UnsetDatabaseOption}; pub(crate) const CREATE_DATABASE_EVENT_TYPE: &str = "create_database"; pub(crate) const ALTER_DATABASE_EVENT_TYPE: &str = "alter_database"; @@ -39,6 +40,7 @@ pub(crate) struct DatabaseDdlEvent { catalog_name: Option, schema_name: Option, payload: Option, + event_context: Option, } #[derive(Debug, Serialize)] @@ -88,6 +90,7 @@ impl DatabaseDdlEvent { schema_name: &str, create_if_not_exists: bool, options: &HashMap, + event_context: EventContext, ) -> Self { let mut options = options.iter().collect::>(); options.sort_unstable_by_key(|(left, _)| *left); @@ -107,6 +110,7 @@ impl DatabaseDdlEvent { create_if_not_exists, options, }), + event_context, ) } @@ -114,6 +118,7 @@ impl DatabaseDdlEvent { catalog_name: &str, schema_name: &str, kind: &AlterDatabaseKind, + event_context: EventContext, ) -> Self { let intent = match kind { AlterDatabaseKind::SetDatabaseOptions(options) => AlterDatabaseIntent::Set { @@ -152,6 +157,7 @@ impl DatabaseDdlEvent { version: PAYLOAD_VERSION, intent, }), + event_context, ) } @@ -159,6 +165,7 @@ impl DatabaseDdlEvent { catalog_name: &str, schema_name: &str, drop_if_exists: bool, + event_context: EventContext, ) -> Self { Self::submitted( DROP_DATABASE_EVENT_TYPE, @@ -168,6 +175,7 @@ impl DatabaseDdlEvent { version: PAYLOAD_VERSION, drop_if_exists, }), + event_context, ) } @@ -188,12 +196,14 @@ impl DatabaseDdlEvent { catalog_name: &str, schema_name: &str, payload: DatabaseDdlPayload, + event_context: EventContext, ) -> Self { Self { event_type, catalog_name: Some(catalog_name.to_string()), schema_name: Some(schema_name.to_string()), payload: Some(payload), + event_context: Some(event_context), } } @@ -203,6 +213,7 @@ impl DatabaseDdlEvent { catalog_name: None, schema_name: None, payload: None, + event_context: None, } } } @@ -223,14 +234,22 @@ impl Event for DatabaseDdlEvent { column_schemas([ &EVENT_TABLE_CATALOG_NAME_COLUMN, &EVENT_TABLE_SCHEMA_NAME_COLUMN, + &EVENT_CONTEXT_COLUMN, ]) } fn extra_rows(&self) -> Result> { + let event_context = self + .event_context + .as_ref() + .map(serde_json::to_value) + .transpose() + .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.schema_name.as_deref()), + nullable_json(event_context.as_ref()), ], }]) } diff --git a/src/common/meta/src/ddl/event/flow.rs b/src/common/meta/src/ddl/event/flow.rs index 68ee51a4fc6..8a52486ddfe 100644 --- a/src/common/meta/src/ddl/event/flow.rs +++ b/src/common/meta/src/ddl/event/flow.rs @@ -19,12 +19,14 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, column_schemas, nullable_string, - nullable_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, column_schemas, + nullable_json, nullable_string, nullable_value, }; use serde::Serialize; use snafu::ResultExt; +use crate::rpc::ddl::EventContext; + pub(crate) const CREATE_FLOW_EVENT_TYPE: &str = "create_flow"; pub(crate) const DROP_FLOW_EVENT_TYPE: &str = "drop_flow"; @@ -69,14 +71,15 @@ pub(crate) struct FlowDdlEvent { flow_name: Option, flow_id: Option, payload: Option, + event_context: Option, } impl FlowDdlEvent { - /// Builds the bounded event emitted when creating a Flow is submitted. pub(crate) fn create_submitted( catalog_name: &str, flow_name: &str, intent: CreateFlowEventIntent, + event_context: EventContext, ) -> Self { Self { event_type: CREATE_FLOW_EVENT_TYPE, @@ -90,15 +93,16 @@ impl FlowDdlEvent { expire_after: intent.expire_after, eval_interval_secs: intent.eval_interval_secs, })), + event_context: Some(event_context), } } - /// Builds the bounded event emitted when dropping a Flow is submitted. pub(crate) fn drop_submitted( catalog_name: &str, flow_name: &str, flow_id: u32, drop_if_exists: bool, + event_context: EventContext, ) -> Self { Self { event_type: DROP_FLOW_EVENT_TYPE, @@ -109,6 +113,7 @@ impl FlowDdlEvent { version: PAYLOAD_VERSION, drop_if_exists, })), + event_context: Some(event_context), } } @@ -137,6 +142,7 @@ impl FlowDdlEvent { flow_name: None, flow_id: None, payload: None, + event_context: None, } } } @@ -154,15 +160,27 @@ impl Event for FlowDdlEvent { } fn extra_schema(&self) -> Vec { - column_schemas([&CATALOG_NAME_COLUMN, &FLOW_NAME_COLUMN, &FLOW_ID_COLUMN]) + column_schemas([ + &CATALOG_NAME_COLUMN, + &FLOW_NAME_COLUMN, + &FLOW_ID_COLUMN, + &EVENT_CONTEXT_COLUMN, + ]) } fn extra_rows(&self) -> Result> { + let event_context = self + .event_context + .as_ref() + .map(serde_json::to_value) + .transpose() + .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.flow_name.as_deref()), nullable_value(self.flow_id.map(ValueData::U32Value)), + nullable_json(event_context.as_ref()), ], }]) } diff --git a/src/common/meta/src/ddl/event/table.rs b/src/common/meta/src/ddl/event/table.rs index f4b25b141b8..8608b3151a6 100644 --- a/src/common/meta/src/ddl/event/table.rs +++ b/src/common/meta/src/ddl/event/table.rs @@ -21,14 +21,17 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN, - TABLE_NAME_COLUMN, column_schemas, nullable_string, nullable_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, + TABLE_ID_COLUMN, TABLE_NAME_COLUMN, column_schemas, nullable_json, nullable_string, + nullable_value, }; use serde::Serialize; use serde_json::Value as JsonValue; use snafu::ResultExt; use store_api::storage::TableId; +use crate::rpc::ddl::EventContext; + /// Current version of table DDL event payloads. pub(crate) const TABLE_DDL_PAYLOAD_VERSION: u8 = 1; @@ -211,6 +214,7 @@ pub(crate) struct TableDdlEvent { event_type: TableDdlEventType, locators: Vec, payload: Option, + event_context: Option, } impl TableDdlEvent { @@ -219,6 +223,7 @@ impl TableDdlEvent { locator: TableDdlLocator, create_if_not_exists: bool, engine: &str, + event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::CreateTable, @@ -228,6 +233,7 @@ impl TableDdlEvent { create_if_not_exists, engine: engine.to_string(), }), + Some(event_context), ) } @@ -235,6 +241,7 @@ impl TableDdlEvent { pub(crate) fn create_logical_tables_submitted( locators: impl IntoIterator, table_count: usize, + event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::CreateLogicalTables, @@ -243,6 +250,7 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, table_count, }), + Some(event_context), ) } @@ -250,6 +258,7 @@ impl TableDdlEvent { pub(crate) fn alter_table_submitted( locator: TableDdlLocator, kind: Option<&'static str>, + event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::AlterTable, @@ -258,6 +267,7 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, kind, }), + Some(event_context), ) } @@ -266,6 +276,7 @@ impl TableDdlEvent { locators: impl IntoIterator, table_count: usize, kinds: impl IntoIterator, + event_context: EventContext, ) -> Self { let kinds = kinds .into_iter() @@ -280,11 +291,16 @@ impl TableDdlEvent { table_count, kinds, }), + Some(event_context), ) } /// Builds the bounded event emitted when dropping a table is submitted. - pub(crate) fn drop_table_submitted(locator: TableDdlLocator, drop_if_exists: bool) -> Self { + pub(crate) fn drop_table_submitted( + locator: TableDdlLocator, + drop_if_exists: bool, + event_context: EventContext, + ) -> Self { Self::submitted( TableDdlEventType::DropTable, [locator], @@ -292,30 +308,39 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, drop_if_exists, }), + Some(event_context), ) } /// Builds the bounded event emitted when restoring a dropped table is submitted. #[cfg(feature = "enterprise")] - pub(crate) fn undrop_table_submitted(locator: TableDdlLocator) -> Self { + pub(crate) fn undrop_table_submitted( + locator: TableDdlLocator, + event_context: EventContext, + ) -> Self { Self::submitted( TableDdlEventType::UndropTable, [locator], TableDdlPayload::UndropTable(UndropTablePayload { version: TABLE_DDL_PAYLOAD_VERSION, }), + Some(event_context), ) } /// Builds the bounded event emitted when purging a dropped table is submitted. #[cfg(feature = "enterprise")] - pub(crate) fn purge_dropped_table_submitted(locator: TableDdlLocator) -> Self { + pub(crate) fn purge_dropped_table_submitted( + locator: TableDdlLocator, + event_context: EventContext, + ) -> Self { Self::submitted( TableDdlEventType::PurgeDroppedTable, [locator], TableDdlPayload::PurgeDroppedTable(PurgeDroppedTablePayload { version: TABLE_DDL_PAYLOAD_VERSION, }), + Some(event_context), ) } @@ -323,6 +348,7 @@ impl TableDdlEvent { pub(crate) fn truncate_table_submitted( locator: TableDdlLocator, time_range_count: usize, + event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::TruncateTable, @@ -331,6 +357,7 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, time_range_count, }), + Some(event_context), ) } @@ -340,6 +367,7 @@ impl TableDdlEvent { event_type, locators: vec![TableDdlLocator::default()], payload: None, + event_context: None, } } @@ -349,6 +377,7 @@ impl TableDdlEvent { event_type: TableDdlEventType::CreateTable, locators: vec![TableDdlLocator::from_table_id(table_id)], payload: None, + event_context: None, } } @@ -360,6 +389,7 @@ impl TableDdlEvent { event_type: TableDdlEventType::CreateLogicalTables, locators: locators.into_iter().collect(), payload: None, + event_context: None, } } @@ -367,11 +397,13 @@ impl TableDdlEvent { event_type: TableDdlEventType, locators: impl IntoIterator, payload: TableDdlPayload, + event_context: Option, ) -> Self { Self { event_type, locators: locators.into_iter().collect(), payload: Some(payload), + event_context, } } @@ -384,7 +416,7 @@ impl TableDdlEvent { ]) } - fn locator_row(&self, locator: &TableDdlLocator) -> Row { + fn locator_row(&self, locator: &TableDdlLocator) -> Result { let mut values = vec![ nullable_string(locator.catalog_name.as_deref()), nullable_string(locator.schema_name.as_deref()), @@ -394,7 +426,14 @@ impl TableDdlEvent { if self.event_type.has_physical_table_id() { values.push(nullable_table_id(locator.physical_table_id)); } - Row { values } + let event_context = self + .event_context + .as_ref() + .map(serde_json::to_value) + .transpose() + .context(SerializeEventSnafu)?; + values.push(nullable_json(event_context.as_ref())); + Ok(Row { values }) } } @@ -415,15 +454,15 @@ impl Event for TableDdlEvent { if self.event_type.has_physical_table_id() { schema.push(PHYSICAL_TABLE_ID_COLUMN.column_schema()); } + schema.push(EVENT_CONTEXT_COLUMN.column_schema()); schema } fn extra_rows(&self) -> Result> { - Ok(self - .locators + self.locators .iter() .map(|locator| self.locator_row(locator)) - .collect()) + .collect() } fn as_any(&self) -> &dyn Any { diff --git a/src/common/meta/src/ddl/event/view.rs b/src/common/meta/src/ddl/event/view.rs index 651805f46d1..d6dcf2c5f17 100644 --- a/src/common/meta/src/ddl/event/view.rs +++ b/src/common/meta/src/ddl/event/view.rs @@ -19,12 +19,14 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, VIEW_NAME_COLUMN, column_schemas, - nullable_string, nullable_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, + VIEW_NAME_COLUMN, column_schemas, nullable_json, nullable_string, nullable_value, }; use serde::Serialize; use snafu::ResultExt; +use crate::rpc::ddl::EventContext; + pub(crate) const CREATE_VIEW_EVENT_TYPE: &str = "create_view"; pub(crate) const DROP_VIEW_EVENT_TYPE: &str = "drop_view"; @@ -62,6 +64,7 @@ pub(crate) struct ViewDdlEvent { view_name: Option, view_id: Option, payload: Option, + event_context: Option, } #[derive(Debug, Serialize)] @@ -72,12 +75,12 @@ enum ViewDdlPayload { } impl ViewDdlEvent { - /// Builds the bounded event emitted when creating a View is submitted. pub(crate) fn create_submitted( catalog_name: &str, schema_name: &str, view_name: &str, intent: CreateViewEventIntent, + event_context: EventContext, ) -> Self { Self::submitted( CREATE_VIEW_EVENT_TYPE, @@ -92,16 +95,17 @@ impl ViewDdlEvent { referenced_table_count: intent.referenced_table_count, column_count: intent.column_count, }), + event_context, ) } - /// Builds the bounded event emitted when dropping a View is submitted. pub(crate) fn drop_submitted( catalog_name: &str, schema_name: &str, view_name: &str, view_id: u32, drop_if_exists: bool, + event_context: EventContext, ) -> Self { Self::submitted( DROP_VIEW_EVENT_TYPE, @@ -113,6 +117,7 @@ impl ViewDdlEvent { version: PAYLOAD_VERSION, drop_if_exists, }), + event_context, ) } @@ -138,6 +143,7 @@ impl ViewDdlEvent { view_name: &str, view_id: Option, payload: ViewDdlPayload, + event_context: EventContext, ) -> Self { Self { event_type, @@ -146,6 +152,7 @@ impl ViewDdlEvent { view_name: Some(view_name.to_string()), view_id, payload: Some(payload), + event_context: Some(event_context), } } @@ -157,6 +164,7 @@ impl ViewDdlEvent { view_name: None, view_id: None, payload: None, + event_context: None, } } @@ -168,6 +176,7 @@ impl ViewDdlEvent { view_name: None, view_id: Some(view_id), payload: None, + event_context: None, } } } @@ -190,16 +199,24 @@ impl Event for ViewDdlEvent { &SCHEMA_NAME_COLUMN, &VIEW_NAME_COLUMN, &VIEW_ID_COLUMN, + &EVENT_CONTEXT_COLUMN, ]) } fn extra_rows(&self) -> Result> { + let event_context = self + .event_context + .as_ref() + .map(serde_json::to_value) + .transpose() + .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.schema_name.as_deref()), nullable_string(self.view_name.as_deref()), nullable_value(self.view_id.map(ValueData::U32Value)), + nullable_json(event_context.as_ref()), ], }]) } diff --git a/src/common/meta/src/ddl/purge_dropped_table.rs b/src/common/meta/src/ddl/purge_dropped_table.rs index 18c161a297f..111b79cabe8 100644 --- a/src/common/meta/src/ddl/purge_dropped_table.rs +++ b/src/common/meta/src/ddl/purge_dropped_table.rs @@ -14,7 +14,7 @@ use std::collections::HashMap; use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; @@ -38,7 +38,7 @@ use crate::error::{self, Result}; use crate::key::DroppedTableMetadata; use crate::key::table_route::TableRouteValue; use crate::lock_key::TableLock; -use crate::rpc::ddl::PurgeDroppedTableTask; +use crate::rpc::ddl::{EventContext, PurgeDroppedTableTask}; use crate::rpc::router::RegionRoute; pub struct PurgeDroppedTableProcedure { @@ -50,17 +50,25 @@ impl PurgeDroppedTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::PurgeDroppedTable"; pub const EXPIRED_TYPE_NAME: &'static str = "metasrv-procedure::PurgeExpiredDroppedTable"; - pub fn new(task: PurgeDroppedTableTask, context: DdlContext) -> Self { + pub fn new( + task: PurgeDroppedTableTask, + context: DdlContext, + event_context: EventContext, + ) -> Self { Self { context, - data: PurgeDroppedTableData::new(task), + data: PurgeDroppedTableData::new(task, event_context), } } - pub fn new_if_expired(task: PurgeDroppedTableTask, context: DdlContext) -> Self { + pub fn new_if_expired( + task: PurgeDroppedTableTask, + context: DdlContext, + event_context: EventContext, + ) -> Self { Self { context, - data: PurgeDroppedTableData::new_if_expired(task), + data: PurgeDroppedTableData::new_if_expired(task, event_context), } } @@ -256,7 +264,10 @@ impl Procedure for PurgeDroppedTableProcedure { LockKey::new(vec![TableLock::Write(self.data.task.table_id).into()]) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::PurgeDroppedTable.as_str()) @@ -266,7 +277,10 @@ impl Procedure for PurgeDroppedTableProcedure { let event = match &ctx.trigger { EventTrigger::Submitted => { let locator = TableDdlLocator::from_table_id(self.data.task.table_id); - TableDdlEvent::purge_dropped_table_submitted(locator) + TableDdlEvent::purge_dropped_table_submitted( + locator, + self.data.event_context.clone(), + ) } _ => TableDdlEvent::lifecycle(TableDdlEventType::PurgeDroppedTable), }; @@ -293,10 +307,12 @@ pub struct PurgeDroppedTableData { drop_generation_loaded: bool, #[serde(default)] purging_claimed: bool, + #[serde(default)] + event_context: EventContext, } impl PurgeDroppedTableData { - fn new(task: PurgeDroppedTableTask) -> Self { + fn new(task: PurgeDroppedTableTask, event_context: EventContext) -> Self { Self { state: PurgeDroppedTableState::Prepare, task, @@ -309,13 +325,14 @@ impl PurgeDroppedTableData { drop_generation: None, drop_generation_loaded: false, purging_claimed: false, + event_context, } } - fn new_if_expired(task: PurgeDroppedTableTask) -> Self { + fn new_if_expired(task: PurgeDroppedTableTask, event_context: EventContext) -> Self { Self { check_expired: true, - ..Self::new(task) + ..Self::new(task, event_context) } } diff --git a/src/common/meta/src/ddl/test_util.rs b/src/common/meta/src/ddl/test_util.rs index 7d8cdd895b2..3686ceb33d2 100644 --- a/src/common/meta/src/ddl/test_util.rs +++ b/src/common/meta/src/ddl/test_util.rs @@ -45,7 +45,7 @@ use crate::key::node_address::{NodeAddressKey, NodeAddressValue}; use crate::key::table_route::TableRouteValue; use crate::key::{MetadataKey, MetadataValue}; use crate::peer::Peer; -use crate::rpc::ddl::CreateTableTask; +use crate::rpc::ddl::{CreateTableTask, EventContext}; use crate::rpc::store::PutRequest; pub async fn create_physical_table_metadata( @@ -104,7 +104,12 @@ pub async fn create_logical_table( table_name: &str, ) -> TableId { let tasks = vec![test_create_logical_table_task(table_name)]; - let mut procedure = CreateLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); + let mut procedure = CreateLogicalTablesProcedure::new( + tasks, + physical_table_id, + EventContext::default(), + ddl_context, + ); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, diff --git a/src/common/meta/src/ddl/tests/alter_logical_tables.rs b/src/common/meta/src/ddl/tests/alter_logical_tables.rs index bef06cd8391..09662199389 100644 --- a/src/common/meta/src/ddl/tests/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/tests/alter_logical_tables.rs @@ -43,7 +43,7 @@ use crate::error::Error::{AlterLogicalTablesInvalidArguments, TableNotFound}; use crate::error::Result; use crate::key::table_name::TableNameKey; use crate::key::table_route::{PhysicalTableRouteValue, TableRouteValue}; -use crate::rpc::ddl::AlterTableTask; +use crate::rpc::ddl::{AlterTableTask, EventContext}; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -162,7 +162,12 @@ async fn test_on_prepare_check_schema() { ), ]; let physical_table_id = 1024u32; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new( + tasks, + physical_table_id, + EventContext::default(), + ddl_context, + ); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -177,7 +182,12 @@ async fn test_on_prepare_check_alter_kind() { "new_table1", )]; let physical_table_id = 1024u32; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new( + tasks, + physical_table_id, + EventContext::default(), + ddl_context, + ); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -197,7 +207,8 @@ async fn test_on_prepare_different_physical_table() { make_alter_logical_table_add_column_task(None, "table2", vec!["column2".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy1_id, ddl_context); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy1_id, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -218,7 +229,8 @@ async fn test_on_prepare_logical_table_not_exists() { make_alter_logical_table_add_column_task(None, "table2", vec!["column2".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, TableNotFound { .. }); } @@ -241,7 +253,8 @@ async fn test_on_prepare() { make_alter_logical_table_add_column_task(None, "table3", vec!["column3".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, EventContext::default(), ddl_context); let result = procedure.on_prepare().await; assert_matches!( result, @@ -277,7 +290,12 @@ async fn test_on_update_metadata() { make_alter_logical_table_add_column_task(None, "table3", vec!["new_col".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); + let mut procedure = AlterLogicalTablesProcedure::new( + tasks, + phy_id, + EventContext::default(), + ddl_context.clone(), + ); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -361,7 +379,12 @@ async fn test_on_part_duplicate_alter_request() { make_alter_logical_table_add_column_task(None, "table2", vec!["col_0".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); + let mut procedure = AlterLogicalTablesProcedure::new( + tasks, + phy_id, + EventContext::default(), + ddl_context.clone(), + ); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -446,7 +469,12 @@ async fn test_on_part_duplicate_alter_request() { ), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); + let mut procedure = AlterLogicalTablesProcedure::new( + tasks, + phy_id, + EventContext::default(), + ddl_context.clone(), + ); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -617,7 +645,8 @@ async fn test_on_submit_alter_region_request() { make_alter_logical_table_add_column_task(None, "table2", vec!["mew_col".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, EventContext::default(), ddl_context); procedure.on_prepare().await.unwrap(); procedure.on_submit_alter_region_requests().await.unwrap(); let mut results = Vec::new(); diff --git a/src/common/meta/src/ddl/tests/alter_table.rs b/src/common/meta/src/ddl/tests/alter_table.rs index 38fc14e760a..af72aaa65df 100644 --- a/src/common/meta/src/ddl/tests/alter_table.rs +++ b/src/common/meta/src/ddl/tests/alter_table.rs @@ -59,7 +59,7 @@ use crate::lock_key::RegionLock; use crate::node_manager::NodeManagerRef; use crate::peer::Peer; use crate::poison_key::table_poison_key; -use crate::rpc::ddl::AlterTableTask; +use crate::rpc::ddl::{AlterTableTask, EventContext}; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -121,7 +121,8 @@ async fn test_on_prepare_table_exists_err() { .unwrap(); let task = test_rename_alter_table_task("non-exists", "foo"); - let mut procedure = AlterTableProcedure::new(1024, task, ddl_context).unwrap(); + let mut procedure = + AlterTableProcedure::new(1024, task, EventContext::default(), ddl_context).unwrap(); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err.status_code(), StatusCode::TableAlreadyExists); } @@ -131,7 +132,8 @@ async fn test_on_prepare_table_not_exists_err() { let node_manager = Arc::new(MockDatanodeManager::new(())); let ddl_context = new_ddl_context(node_manager); let task = test_rename_alter_table_task("non-exists", "foo"); - let mut procedure = AlterTableProcedure::new(1024, task, ddl_context).unwrap(); + let mut procedure = + AlterTableProcedure::new(1024, task, EventContext::default(), ddl_context).unwrap(); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err.status_code(), StatusCode::TableNotFound); } @@ -245,8 +247,13 @@ async fn test_on_submit_alter_request() { let alter_table_task = test_alter_table_task(table_name); let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = - AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); + let mut procedure = AlterTableProcedure::new( + table_id, + alter_table_task, + EventContext::default(), + ddl_context.clone(), + ) + .unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -299,8 +306,13 @@ async fn test_on_submit_alter_request_without_sync_request() { let alter_table_task = test_alter_table_task(table_name); let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = - AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); + let mut procedure = AlterTableProcedure::new( + table_id, + alter_table_task, + EventContext::default(), + ddl_context.clone(), + ) + .unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -357,7 +369,13 @@ async fn test_on_submit_alter_request_with_outdated_request() { }; let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = AlterTableProcedure::new(table_id, alter_table_task, ddl_context).unwrap(); + let mut procedure = AlterTableProcedure::new( + table_id, + alter_table_task, + EventContext::default(), + ddl_context, + ) + .unwrap(); procedure.on_prepare().await.unwrap(); let err = procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -386,7 +404,9 @@ async fn test_on_update_metadata_rename() { .unwrap(); let task = test_rename_alter_table_task(table_name, new_table_name); - let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) + .unwrap(); procedure.on_prepare().await.unwrap(); procedure.on_update_metadata().await.unwrap(); @@ -466,7 +486,9 @@ async fn test_on_update_metadata_add_columns() { }; let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) + .unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -563,7 +585,9 @@ async fn test_on_update_table_options() { }; let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) + .unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -634,6 +658,7 @@ async fn test_skip_wal_rejects_mixed_table_options() { })), }, }, + EventContext::default(), ddl_context, ) .unwrap(); @@ -662,8 +687,13 @@ async fn test_skip_wal_rejects_logical_table() { }, }; - let mut procedure = - AlterTableProcedure::new(logical_table_id, task, ddl_context.clone()).unwrap(); + let mut procedure = AlterTableProcedure::new( + logical_table_id, + task, + EventContext::default(), + ddl_context.clone(), + ) + .unwrap(); let error = procedure.on_prepare().await.unwrap_err(); assert_matches!(error, Error::Unsupported { .. }); @@ -709,7 +739,9 @@ async fn test_skip_wal_rejects_file_engine_table() { }, }; - let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) + .unwrap(); let error = procedure.on_prepare().await.unwrap_err(); assert_matches!(error, Error::Unsupported { .. }); @@ -746,6 +778,7 @@ fn test_skip_wal_holds_region_locks() { let procedure = AlterTableProcedure::new_with_region_locks( table_id, task, + EventContext::default(), region_ids.clone(), context.clone(), ) @@ -797,6 +830,7 @@ async fn test_skip_wal_detects_region_route_change() { let mut procedure = AlterTableProcedure::new_with_region_locks( table_id, alter_task, + EventContext::default(), stale_region_locks, ddl_context.clone(), ) @@ -857,6 +891,7 @@ async fn test_skip_wal_updates_metadata_before_all_replicas() { let mut procedure = AlterTableProcedure::new_with_region_locks( table_id, alter_task, + EventContext::default(), region_locks, ddl_context.clone(), ) @@ -938,8 +973,13 @@ async fn prepare_alter_table_procedure( }, }; let procedure_id = ProcedureId::random(); - let mut procedure = - AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); + let mut procedure = AlterTableProcedure::new( + table_id, + alter_table_task, + EventContext::default(), + ddl_context.clone(), + ) + .unwrap(); procedure.on_prepare().await.unwrap(); (procedure, procedure_id) } diff --git a/src/common/meta/src/ddl/tests/create_flow.rs b/src/common/meta/src/ddl/tests/create_flow.rs index b526856bdac..280a2633e32 100644 --- a/src/common/meta/src/ddl/tests/create_flow.rs +++ b/src/common/meta/src/ddl/tests/create_flow.rs @@ -37,7 +37,7 @@ use crate::key::flow::flow_info::{ FlowInfoValue, FlowMissedTickPolicy, FlowScheduleConfig, FlowStatus, }; use crate::key::table_route::TableRouteValue; -use crate::rpc::ddl::{CreateFlowTask, FlowQueryContext, QueryContext}; +use crate::rpc::ddl::{CreateFlowTask, EventContext, FlowQueryContext, QueryContext}; use crate::test_util::{MockFlownodeManager, new_ddl_context}; pub(crate) fn test_query_context() -> QueryContext { @@ -120,7 +120,8 @@ async fn test_create_flow_source_table_not_found() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); + let mut procedure = + CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -143,7 +144,12 @@ async fn test_create_pending_flow_source_table_not_found_with_defer() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context.clone()); + let mut procedure = CreateFlowProcedure::new( + task, + query_ctx, + EventContext::default(), + ddl_context.clone(), + ); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Executing { persist: true, .. }); assert_eq!(procedure.data.unresolved_source_table_names.len(), 1); @@ -183,7 +189,8 @@ async fn test_create_pending_flow_source_table_not_found_with_defer_false() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); + let mut procedure = + CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -233,7 +240,12 @@ async fn test_create_pending_flow_records_partial_source_resolution() { ); enable_defer_on_missing_source(&mut task); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context.clone()); + let mut procedure = CreateFlowProcedure::new( + task, + query_ctx, + EventContext::default(), + ddl_context.clone(), + ); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Executing { persist: true, .. }); assert_eq!(procedure.data.source_table_ids, vec![existing_table_id]); @@ -428,6 +440,7 @@ fn test_resolved_schedule_defaults_in_create_request() { // Construct CreateFlowData and verify CreateRequest carries internal key. let data = CreateFlowData { + event_context: EventContext::default(), state: CreateFlowState::CreateFlows, task, flow_id: Some(1024), @@ -441,6 +454,7 @@ fn test_resolved_schedule_defaults_in_create_request() { }; let data2 = CreateFlowData { + event_context: EventContext::default(), state: CreateFlowState::CreateMetadata, task: data.task.clone(), flow_id: Some(1024), @@ -580,7 +594,8 @@ async fn test_create_flow_rejects_unknown_option_in_meta_task() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); + let mut procedure = + CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unexpected { .. }); @@ -601,6 +616,7 @@ fn test_create_request_strips_defer_on_missing_source_runtime_option() { enable_defer_on_missing_source(&mut task); let data = CreateFlowData { + event_context: EventContext::default(), state: CreateFlowState::CreateFlows, task, flow_id: Some(1024), @@ -650,7 +666,12 @@ pub(crate) async fn create_test_flow( false, ); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(task.clone(), query_ctx, ddl_context.clone()); + let mut procedure = CreateFlowProcedure::new( + task.clone(), + query_ctx, + EventContext::default(), + ddl_context.clone(), + ); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let flow_id = output.downcast_ref::().unwrap(); @@ -671,7 +692,12 @@ pub(crate) async fn create_test_pending_flow( ); enable_defer_on_missing_source(&mut task); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context.clone()); + let mut procedure = CreateFlowProcedure::new( + task, + query_ctx, + EventContext::default(), + ddl_context.clone(), + ); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let flow_id = output.downcast_ref::().unwrap(); @@ -718,7 +744,12 @@ async fn test_create_flow() { true, ); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(task.clone(), query_ctx, ddl_context.clone()); + let mut procedure = CreateFlowProcedure::new( + task.clone(), + query_ctx, + EventContext::default(), + ddl_context.clone(), + ); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let flow_id = output.downcast_ref::().unwrap(); assert_eq!(*flow_id, 1024); @@ -726,7 +757,12 @@ async fn test_create_flow() { // Creates again let task = test_create_flow_task("my_flow", source_table_names, sink_table_name, false); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(task.clone(), query_ctx, ddl_context); + let mut procedure = CreateFlowProcedure::new( + task.clone(), + query_ctx, + EventContext::default(), + ddl_context, + ); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowAlreadyExists { .. }); } @@ -783,7 +819,12 @@ async fn test_replace_pending_flow_with_active_flow_is_unsupported() { ); replace_task.or_replace = true; let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(replace_task, query_ctx, ddl_context.clone()); + let mut procedure = CreateFlowProcedure::new( + replace_task, + query_ctx, + EventContext::default(), + ddl_context.clone(), + ); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -841,7 +882,12 @@ async fn test_replace_active_flow_with_pending_flow_is_unsupported() { enable_defer_on_missing_source(&mut replace_task); replace_task.or_replace = true; let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(replace_task, query_ctx, ddl_context.clone()); + let mut procedure = CreateFlowProcedure::new( + replace_task, + query_ctx, + EventContext::default(), + ddl_context.clone(), + ); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -901,7 +947,12 @@ async fn test_replace_pending_flow_with_pending_flow_updates_metadata() { enable_defer_on_missing_source(&mut replace_task); replace_task.or_replace = true; let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(replace_task, query_ctx, ddl_context.clone()); + let mut procedure = CreateFlowProcedure::new( + replace_task, + query_ctx, + EventContext::default(), + ddl_context.clone(), + ); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let replaced_flow_id = *output.downcast_ref::().unwrap(); assert_eq!(replaced_flow_id, original_flow_id); @@ -953,7 +1004,8 @@ async fn test_create_flow_same_source_and_sink_table() { // Try to create a flow with same source and sink table - should fail let task = test_create_flow_task("my_flow", source_table_names, sink_table_name, false); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); + let mut procedure = + CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); @@ -986,7 +1038,7 @@ fn test_create_flow_lock_key_does_not_lock_sink_table_name() { let task = create_test_flow_task_for_serialization(); let query_ctx = test_query_context(); let ddl_context = new_ddl_context(Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler))); - let procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); + let procedure = CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); let lock_keys = procedure.lock_key().get_keys(); @@ -1061,6 +1113,7 @@ fn test_create_flow_data_new_format_serialization() { }; let data = CreateFlowData { + event_context: EventContext::default(), state: CreateFlowState::Prepare, task: create_test_flow_task_for_serialization(), flow_id: None, @@ -1129,7 +1182,7 @@ fn test_create_flow_procedure_strips_scheduled_time_extension() { .extensions .insert("flow.other".to_string(), "kept".to_string()); - let procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); + let procedure = CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); assert!( !procedure @@ -1162,6 +1215,7 @@ fn test_flow_info_conversion_with_flow_context() { }; let data = CreateFlowData { + event_context: EventContext::default(), state: CreateFlowState::CreateMetadata, task: create_test_flow_task_for_serialization(), flow_id: Some(123), @@ -1204,6 +1258,7 @@ fn test_flow_info_conversion_strips_scheduled_time_extension() { }; let data = CreateFlowData { + event_context: EventContext::default(), state: CreateFlowState::CreateMetadata, task: create_test_flow_task_for_serialization(), flow_id: Some(123), diff --git a/src/common/meta/src/ddl/tests/create_logical_tables.rs b/src/common/meta/src/ddl/tests/create_logical_tables.rs index e2927e8df65..d1cc2420e04 100644 --- a/src/common/meta/src/ddl/tests/create_logical_tables.rs +++ b/src/common/meta/src/ddl/tests/create_logical_tables.rs @@ -39,6 +39,7 @@ use crate::ddl::test_util::{ }; use crate::error::{Error, Result}; use crate::key::table_route::{PhysicalTableRouteValue, TableRouteValue}; +use crate::rpc::ddl::EventContext; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -96,7 +97,12 @@ async fn test_on_prepare_physical_table_not_found() { let ddl_context = new_ddl_context(node_manager); let tasks = vec![test_create_logical_table_task("foo")]; let physical_table_id = 1024u32; - let mut procedure = CreateLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); + let mut procedure = CreateLogicalTablesProcedure::new( + tasks, + physical_table_id, + EventContext::default(), + ddl_context, + ); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableRouteNotFound { .. }); } @@ -126,7 +132,12 @@ async fn test_on_prepare() { // The create logical table procedure. let tasks = vec![test_create_logical_table_task("foo")]; let physical_table_id = table_id; - let mut procedure = CreateLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); + let mut procedure = CreateLogicalTablesProcedure::new( + tasks, + physical_table_id, + EventContext::default(), + ddl_context, + ); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -172,8 +183,12 @@ async fn test_on_prepare_logical_table_exists_err() { .unwrap(); // The create logical table procedure. let physical_table_id = table_id; - let mut procedure = - CreateLogicalTablesProcedure::new(vec![task], physical_table_id, ddl_context); + let mut procedure = CreateLogicalTablesProcedure::new( + vec![task], + physical_table_id, + EventContext::default(), + ddl_context, + ); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -216,8 +231,12 @@ async fn test_on_prepare_with_create_if_table_exists() { let physical_table_id = table_id; // Sets `create_if_not_exists` task.create_table.create_if_not_exists = true; - let mut procedure = - CreateLogicalTablesProcedure::new(vec![task], physical_table_id, ddl_context); + let mut procedure = CreateLogicalTablesProcedure::new( + vec![task], + physical_table_id, + EventContext::default(), + ddl_context, + ); let status = procedure.on_prepare().await.unwrap(); let output = status.downcast_output_ref::>().unwrap(); assert_eq!(*output, vec![8192]); @@ -264,6 +283,7 @@ async fn test_on_prepare_part_logical_tables_exist() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, non_exist_task], physical_table_id, + EventContext::default(), ddl_context, ); let status = procedure.on_prepare().await.unwrap(); @@ -311,6 +331,7 @@ async fn test_on_create_metadata() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, yet_another_task], physical_table_id, + EventContext::default(), ddl_context.clone(), ); let status = procedure.on_prepare().await.unwrap(); @@ -404,6 +425,7 @@ async fn test_on_create_metadata_part_logical_tables_exist() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, non_exist_task], physical_table_id, + EventContext::default(), ddl_context.clone(), ); let status = procedure.on_prepare().await.unwrap(); @@ -476,6 +498,7 @@ async fn test_on_create_metadata_err() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task.clone(), yet_another_task], physical_table_id, + EventContext::default(), ddl_context.clone(), ); let status = procedure.on_prepare().await.unwrap(); @@ -536,6 +559,7 @@ async fn test_on_submit_create_request() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, yet_another_task], physical_table_id, + EventContext::default(), ddl_context, ); procedure.on_prepare().await.unwrap(); diff --git a/src/common/meta/src/ddl/tests/create_table.rs b/src/common/meta/src/ddl/tests/create_table.rs index db94a12f84f..2eee317dc14 100644 --- a/src/common/meta/src/ddl/tests/create_table.rs +++ b/src/common/meta/src/ddl/tests/create_table.rs @@ -48,7 +48,7 @@ use crate::ddl::test_util::{assert_column_name, get_raw_table_info, put_datanode use crate::error::{Error, Result}; use crate::key::table_route::TableRouteValue; use crate::kv_backend::memory::MemoryKvBackend; -use crate::rpc::ddl::CreateTableTask; +use crate::rpc::ddl::{CreateTableTask, EventContext, QueryContext}; use crate::test_util::{MockDatanodeManager, new_ddl_context, new_ddl_context_with_kv_backend}; fn create_request_handler(_peer: Peer, request: RegionRequest) -> Result { @@ -189,7 +189,13 @@ async fn test_on_prepare_table_exists_err() { ) .await .unwrap(); - let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); + let mut procedure = CreateTableProcedure::new( + task, + QueryContext::default(), + EventContext::default(), + ddl_context, + ) + .unwrap(); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -212,7 +218,13 @@ async fn test_on_prepare_with_create_if_table_exists() { ) .await .unwrap(); - let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); + let mut procedure = CreateTableProcedure::new( + task, + QueryContext::default(), + EventContext::default(), + ddl_context, + ) + .unwrap(); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Done { output: Some(..) }); let table_id = *status.downcast_output_ref::().unwrap(); @@ -225,7 +237,13 @@ async fn test_on_prepare_without_create_if_table_exists() { let ddl_context = new_ddl_context(node_manager); let mut task = test_create_table_task("foo"); task.create_table.create_if_not_exists = true; - let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); + let mut procedure = CreateTableProcedure::new( + task, + QueryContext::default(), + EventContext::default(), + ddl_context, + ) + .unwrap(); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -244,7 +262,13 @@ async fn test_on_datanode_create_regions_should_retry() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); + let mut procedure = CreateTableProcedure::new( + task, + QueryContext::default(), + EventContext::default(), + ddl_context, + ) + .unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), @@ -261,7 +285,13 @@ async fn test_on_datanode_create_regions_should_not_retry() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); + let mut procedure = CreateTableProcedure::new( + task, + QueryContext::default(), + EventContext::default(), + ddl_context, + ) + .unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), @@ -278,7 +308,13 @@ async fn test_on_datanode_create_regions_remaps_addresses_when_retrying() { let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); - let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); + let mut procedure = CreateTableProcedure::new( + task, + QueryContext::default(), + EventContext::default(), + ddl_context.clone(), + ) + .unwrap(); procedure.on_prepare().await.unwrap(); let table_route = procedure.data.table_route.as_mut().unwrap(); @@ -299,7 +335,13 @@ async fn test_on_create_metadata_error() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new(task.clone(), ddl_context.clone()).unwrap(); + let mut procedure = CreateTableProcedure::new( + task.clone(), + QueryContext::default(), + EventContext::default(), + ddl_context.clone(), + ) + .unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), @@ -332,7 +374,13 @@ async fn test_on_create_metadata() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); + let mut procedure = CreateTableProcedure::new( + task, + QueryContext::default(), + EventContext::default(), + ddl_context.clone(), + ) + .unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), @@ -360,7 +408,13 @@ async fn test_memory_region_keeper_guard_dropped_on_procedure_done() { let ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); let task = test_create_table_task("foo"); - let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); + let mut procedure = CreateTableProcedure::new( + task, + QueryContext::default(), + EventContext::default(), + ddl_context.clone(), + ) + .unwrap(); execute_procedure_until(&mut procedure, |p| { p.data.state == CreateTableState::CreateMetadata diff --git a/src/common/meta/src/ddl/tests/create_view.rs b/src/common/meta/src/ddl/tests/create_view.rs index cc98bb1baed..618249fd1cd 100644 --- a/src/common/meta/src/ddl/tests/create_view.rs +++ b/src/common/meta/src/ddl/tests/create_view.rs @@ -29,7 +29,7 @@ use crate::ddl::create_view::CreateViewProcedure; use crate::ddl::test_util::datanode_handler::NaiveDatanodeHandler; use crate::ddl::tests::create_table::test_create_table_task; use crate::error::Error; -use crate::rpc::ddl::CreateViewTask; +use crate::rpc::ddl::{CreateViewTask, EventContext, QueryContext}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; pub(crate) fn test_table_names() -> HashSet { @@ -112,7 +112,7 @@ async fn test_on_prepare_view_exists_err() { ) .await .unwrap(); - let mut procedure = CreateViewProcedure::new(task, ddl_context); + let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::ViewAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -138,7 +138,7 @@ async fn test_on_prepare_with_create_if_view_exists() { ) .await .unwrap(); - let mut procedure = CreateViewProcedure::new(task, ddl_context); + let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Done { output: Some(..) }); let table_id = *status.downcast_output_ref::().unwrap(); @@ -151,7 +151,7 @@ async fn test_on_prepare_without_create_if_table_exists() { let ddl_context = new_ddl_context(node_manager); let mut task = test_create_view_task("foo"); task.create_view.create_if_not_exists = true; - let mut procedure = CreateViewProcedure::new(task, ddl_context); + let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -170,7 +170,7 @@ async fn test_on_create_metadata() { let ddl_context = new_ddl_context(node_manager); let task = test_create_view_task("foo"); assert!(!task.create_view.create_if_not_exists); - let mut procedure = CreateViewProcedure::new(task, ddl_context); + let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), @@ -189,7 +189,8 @@ async fn test_replace_view_metadata() { let ddl_context = new_ddl_context(node_manager.clone()); let task = test_create_view_task("foo"); assert!(!task.create_view.create_if_not_exists); - let mut procedure = CreateViewProcedure::new(task.clone(), ddl_context.clone()); + let mut procedure = + CreateViewProcedure::new(task.clone(), EventContext::default(), ddl_context.clone()); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), @@ -214,7 +215,8 @@ async fn test_replace_view_metadata() { let mut task = test_create_view_task("foo"); // The view already exists, prepare should fail { - let mut procedure = CreateViewProcedure::new(task.clone(), ddl_context.clone()); + let mut procedure = + CreateViewProcedure::new(task.clone(), EventContext::default(), ddl_context.clone()); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::ViewAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -225,7 +227,8 @@ async fn test_replace_view_metadata() { task.create_view.logical_plan = vec![4, 5, 6]; task.create_view.definition = "new_definition".to_string(); - let mut procedure = CreateViewProcedure::new(task, ddl_context.clone()); + let mut procedure = + CreateViewProcedure::new(task, EventContext::default(), ddl_context.clone()); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), @@ -259,7 +262,13 @@ async fn test_replace_table() { { // Create a `foo` table. let task = test_create_table_task("foo"); - let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); + let mut procedure = CreateTableProcedure::new( + task, + QueryContext::default(), + EventContext::default(), + ddl_context.clone(), + ) + .unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), @@ -272,7 +281,8 @@ async fn test_replace_table() { // Try to replace a view named `foo` too. let mut task = test_create_view_task("foo"); task.create_view.or_replace = true; - let mut procedure = CreateViewProcedure::new(task.clone(), ddl_context.clone()); + let mut procedure = + CreateViewProcedure::new(task.clone(), EventContext::default(), ddl_context.clone()); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); diff --git a/src/common/meta/src/ddl/tests/drop_database.rs b/src/common/meta/src/ddl/tests/drop_database.rs index 81cc29a51ba..955625e4bb3 100644 --- a/src/common/meta/src/ddl/tests/drop_database.rs +++ b/src/common/meta/src/ddl/tests/drop_database.rs @@ -26,6 +26,7 @@ use crate::ddl::drop_database::executor::DropDatabaseExecutor; use crate::ddl::test_util::datanode_handler::{NaiveDatanodeHandler, RetryErrorDatanodeHandler}; use crate::ddl::test_util::{create_logical_table, create_physical_table}; use crate::key::schema_name::SchemaNameKey; +use crate::rpc::ddl::EventContext; use crate::test_util::{MockDatanodeManager, new_ddl_context}; #[tokio::test] @@ -54,6 +55,7 @@ async fn test_drop_database_with_logical_tables() { DEFAULT_CATALOG_NAME.to_string(), DEFAULT_SCHEMA_NAME.to_string(), false, + EventContext::default(), ddl_context.clone(), ); @@ -102,6 +104,7 @@ async fn test_drop_database_retryable_error() { DEFAULT_CATALOG_NAME.to_string(), DEFAULT_SCHEMA_NAME.to_string(), false, + EventContext::default(), ddl_context.clone(), ); @@ -146,6 +149,7 @@ async fn test_drop_database_recover() { DEFAULT_CATALOG_NAME.to_string(), DEFAULT_SCHEMA_NAME.to_string(), false, + EventContext::default(), ddl_context.clone(), ); let num_operating_regions = 1; diff --git a/src/common/meta/src/ddl/tests/drop_flow.rs b/src/common/meta/src/ddl/tests/drop_flow.rs index 6d36fc93ba3..0976aa572f7 100644 --- a/src/common/meta/src/ddl/tests/drop_flow.rs +++ b/src/common/meta/src/ddl/tests/drop_flow.rs @@ -26,7 +26,7 @@ use crate::ddl::test_util::flownode_handler::NaiveFlownodeHandler; use crate::ddl::tests::create_flow::{create_test_flow, create_test_pending_flow}; use crate::error; use crate::key::table_route::TableRouteValue; -use crate::rpc::ddl::DropFlowTask; +use crate::rpc::ddl::{DropFlowTask, EventContext}; use crate::test_util::{MockFlownodeManager, new_ddl_context}; pub(crate) fn test_drop_flow_task( @@ -48,7 +48,7 @@ async fn test_drop_flow_not_found() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let task = test_drop_flow_task("my_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, ddl_context); + let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowNotFound { .. }); } @@ -81,17 +81,17 @@ async fn test_drop_flow() { create_test_flow(&ddl_context, "my_flow", source_table_names, sink_table_name).await; // Drops the flows let task = test_drop_flow_task("my_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, ddl_context.clone()); + let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; // Drops if not exists let task = test_drop_flow_task("my_flow", flow_id, true); - let mut procedure = DropFlowProcedure::new(task, ddl_context.clone()); + let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; // Drops again let task = test_drop_flow_task("my_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, ddl_context); + let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowNotFound { .. }); } @@ -129,11 +129,11 @@ async fn test_drop_pending_flow_without_routes() { assert!(flow_info.flownode_ids().is_empty()); let task = test_drop_flow_task("drop_pending_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, ddl_context.clone()); + let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; let task = test_drop_flow_task("drop_pending_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, ddl_context); + let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowNotFound { .. }); } diff --git a/src/common/meta/src/ddl/tests/drop_table.rs b/src/common/meta/src/ddl/tests/drop_table.rs index a278b951dd5..176f847c428 100644 --- a/src/common/meta/src/ddl/tests/drop_table.rs +++ b/src/common/meta/src/ddl/tests/drop_table.rs @@ -65,9 +65,9 @@ use crate::key::table_route::TableRouteValue; use crate::kv_backend::KvBackend; use crate::kv_backend::memory::MemoryKvBackend; use crate::peer::Peer; -use crate::rpc::ddl::DropTableTask; +use crate::rpc::ddl::{DropTableTask, EventContext}; #[cfg(feature = "enterprise")] -use crate::rpc::ddl::{PurgeDroppedTableTask, UndropTableTask}; +use crate::rpc::ddl::{PurgeDroppedTableTask, QueryContext, UndropTableTask}; use crate::rpc::router::{Region, RegionRoute}; #[cfg(feature = "enterprise")] use crate::rpc::store::{BatchDeleteRequest, PutRequest}; @@ -128,7 +128,11 @@ async fn undrop_at_restore_metadata( table_id: TableId, ) -> UndropTableProcedure { create_dropped_table(context, table_id, Some(1), Some(i64::MAX)).await; - let mut undrop = UndropTableProcedure::new(new_undrop_table_task(table_id), context.clone()); + let mut undrop = UndropTableProcedure::new( + new_undrop_table_task(table_id), + context.clone(), + EventContext::default(), + ); let procedure_context = new_test_procedure_context(); undrop.execute(&procedure_context).await.unwrap(); undrop.execute(&procedure_context).await.unwrap(); @@ -157,6 +161,7 @@ fn test_old_drop_table_json_defaults_to_hard_drop() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", 1024, false), runtime_context.clone(), + EventContext::default(), ); let mut old_data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); old_data @@ -164,6 +169,7 @@ fn test_old_drop_table_json_defaults_to_hard_drop() { .unwrap() .remove("soft_drop_enabled"); old_data.as_object_mut().unwrap().remove("dropped_at"); + old_data.as_object_mut().unwrap().remove("event_context"); let recovered = DropTableProcedure::from_json(&old_data.to_string(), runtime_context).unwrap(); let recovered_data: serde_json::Value = @@ -171,12 +177,20 @@ fn test_old_drop_table_json_defaults_to_hard_drop() { assert_eq!(recovered_data["soft_drop_enabled"], false); assert_eq!(recovered_data["dropped_at"], serde_json::Value::Null); + assert_eq!( + recovered_data["event_context"], + serde_json::json!({"reason": "unknown"}) + ); } #[test] fn test_drop_table_lock_key_includes_table_name() { let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); - let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", 1024, false), + context, + EventContext::default(), + ); let keys = procedure .lock_key() @@ -205,7 +219,11 @@ fn test_new_drop_table_procedure_disables_soft_drop_without_enterprise() { let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); context.soft_drop_enabled = true; - let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", 1024, false), + context, + EventContext::default(), + ); assert!(!procedure.data.soft_drop_enabled); } @@ -216,7 +234,11 @@ fn test_new_drop_table_procedure_preserves_soft_drop_with_enterprise() { let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); context.soft_drop_enabled = true; - let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", 1024, false), + context, + EventContext::default(), + ); assert!(procedure.data.soft_drop_enabled); } @@ -251,6 +273,7 @@ async fn test_recovered_soft_drop_preserves_persisted_mode() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), soft_context.clone(), + EventContext::default(), ); let hard_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); @@ -290,9 +313,13 @@ async fn test_legacy_soft_drop_prepare_preserves_persisted_mode() { .unwrap(); let mut persisted: serde_json::Value = serde_json::from_str( - &DropTableProcedure::new(new_drop_table_task("foo", table_id, false), context.clone()) - .dump() - .unwrap(), + &DropTableProcedure::new( + new_drop_table_task("foo", table_id, false), + context.clone(), + EventContext::default(), + ) + .dump() + .unwrap(), ) .unwrap(); persisted["soft_drop_enabled"] = true.into(); @@ -322,7 +349,11 @@ async fn test_disabled_soft_drop_creates_hard_drop_procedure() { context.soft_drop_enabled = false; context.soft_drop_retention = Some(Duration::from_millis(100)); - let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", 1024, false), + context, + EventContext::default(), + ); assert!(!procedure.data.soft_drop_enabled); assert_eq!(None, procedure.data.dropped_at); @@ -356,6 +387,7 @@ async fn test_recovered_hard_drop_ignores_soft_runtime_context() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), hard_context.clone(), + EventContext::default(), ); let mut soft_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); @@ -397,7 +429,7 @@ async fn test_on_prepare_table_not_exists_err() { .unwrap(); let task = new_drop_table_task("bar", table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context); + let mut procedure = DropTableProcedure::new(task, ddl_context, EventContext::default()); assert_eq!(procedure.data.dropped_at, None); assert_eq!(procedure.data.retention_expires_at, None); let err = procedure.on_prepare().await.unwrap_err(); @@ -427,6 +459,7 @@ async fn test_soft_drop_prepare_assigns_stable_deadline_and_recovers_it() { let mut procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), ddl_context.clone(), + EventContext::default(), ); procedure.on_prepare().await.unwrap(); @@ -463,6 +496,7 @@ async fn test_soft_drop_prepare_rejects_deadline_overflow_before_metadata_delete let mut procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), ddl_context.clone(), + EventContext::default(), ); let error = procedure.on_prepare().await.unwrap_err(); @@ -502,13 +536,13 @@ async fn test_on_prepare_table() { let task = new_drop_table_task("bar", table_id, true); // Drop if exists - let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); procedure.on_prepare().await.unwrap(); assert!(!procedure.rollback_supported()); let task = new_drop_table_task(table_name, table_id, false); // Drop table - let mut procedure = DropTableProcedure::new(task, ddl_context); + let mut procedure = DropTableProcedure::new(task, ddl_context, EventContext::default()); procedure.on_prepare().await.unwrap(); } @@ -559,7 +593,7 @@ async fn test_on_datanode_drop_regions() { let task = new_drop_table_task(table_name, table_id, false); // Drop table - let mut procedure = DropTableProcedure::new(task, ddl_context); + let mut procedure = DropTableProcedure::new(task, ddl_context, EventContext::default()); procedure.on_prepare().await.unwrap(); procedure.on_datanode_drop_regions(false).await.unwrap(); @@ -628,7 +662,7 @@ async fn test_on_datanode_drop_regions_remaps_addresses_when_retrying() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); procedure.on_prepare().await.unwrap(); put_datanode_address(&ddl_context, 1, "new-leader").await; @@ -687,7 +721,7 @@ async fn test_soft_drop_closes_regions_and_keeps_tombstone() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); execute_procedure_until_done(&mut procedure).await; @@ -801,7 +835,7 @@ async fn test_soft_drop_keeps_metadata_live_until_regions_close() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); put_datanode_address(&ddl_context, 1, "new-leader").await; @@ -888,6 +922,7 @@ async fn test_soft_drop_timestamp_is_stable_across_retry_and_recovery() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), ddl_context.clone(), + EventContext::default(), ); assert_eq!(procedure.data.dropped_at, None); let mut procedure = procedure; @@ -958,7 +993,7 @@ async fn test_hard_drop_keeps_delete_tombstone_flow() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); execute_procedure_until(&mut procedure, |p| { p.data.state == DropTableState::DeleteTombstone @@ -1001,13 +1036,20 @@ async fn test_create_table_succeeds_while_tombstone_exists() { .unwrap(); let drop_task = new_drop_table_task(table_name, dropped_table_id, false); - let mut drop_procedure = DropTableProcedure::new(drop_task, ddl_context.clone()); + let mut drop_procedure = + DropTableProcedure::new(drop_task, ddl_context.clone(), EventContext::default()); execute_procedure_until_done(&mut drop_procedure).await; let mut create_task = test_create_table_task(table_name, 1025); create_task.create_table.table_id = None; create_task.table_info.ident.table_id = 0; - let mut create_procedure = CreateTableProcedure::new(create_task, ddl_context.clone()).unwrap(); + let mut create_procedure = CreateTableProcedure::new( + create_task, + QueryContext::default(), + EventContext::default(), + ddl_context.clone(), + ) + .unwrap(); execute_procedure_until_done(&mut create_procedure).await; let live_table = ddl_context @@ -1052,7 +1094,8 @@ async fn test_hard_drop_recreated_table_fails_when_soft_tombstone_exists() { .unwrap(); let drop_task = new_drop_table_task(table_name, original_table_id, false); - let mut drop_procedure = DropTableProcedure::new(drop_task, ddl_context.clone()); + let mut drop_procedure = + DropTableProcedure::new(drop_task, ddl_context.clone(), EventContext::default()); execute_procedure_until_done(&mut drop_procedure).await; ddl_context.soft_drop_enabled = false; @@ -1069,6 +1112,7 @@ async fn test_hard_drop_recreated_table_fails_when_soft_tombstone_exists() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context, + EventContext::default(), ); let err = procedure.on_prepare().await.unwrap_err(); @@ -1098,7 +1142,8 @@ async fn test_hard_drop_recreated_table_ignores_previous_orphan_tombstone() { .unwrap(); let drop_task = new_drop_table_task(table_name, original_table_id, false); - let mut drop_procedure = DropTableProcedure::new(drop_task, ddl_context.clone()); + let mut drop_procedure = + DropTableProcedure::new(drop_task, ddl_context.clone(), EventContext::default()); execute_procedure_until(&mut drop_procedure, |p| { p.data.state == DropTableState::DeleteTombstone }) @@ -1117,6 +1162,7 @@ async fn test_hard_drop_recreated_table_ignores_previous_orphan_tombstone() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context, + EventContext::default(), ); procedure.on_prepare().await.unwrap(); @@ -1156,6 +1202,7 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; assert!( @@ -1172,8 +1219,11 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { while rx.try_recv().is_ok() {} - let mut procedure = - UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(table_id), + ddl_context.clone(), + EventContext::default(), + ); execute_procedure_until_done(&mut procedure).await; let live_table = ddl_context @@ -1264,12 +1314,16 @@ async fn test_undrop_table_opens_regions_before_restoring_live_metadata() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} - let mut procedure = - UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(table_id), + ddl_context.clone(), + EventContext::default(), + ); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -1333,13 +1387,17 @@ async fn test_undrop_logical_table_skips_datanode_open() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, logical_table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} - let mut procedure = - UndropTableProcedure::new(new_undrop_table_task(logical_table_id), ddl_context.clone()); + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(logical_table_id), + ddl_context.clone(), + EventContext::default(), + ); execute_procedure_until_done(&mut procedure).await; let live_table = ddl_context @@ -1370,6 +1428,7 @@ async fn test_soft_drop_metric_logical_table_falls_back_to_hard_drop() { let mut procedure = DropTableProcedure::new( new_drop_table_task("foo", logical_table_id, false), ddl_context.clone(), + EventContext::default(), ); procedure.on_prepare().await.unwrap(); assert!(!procedure.data.soft_drop_enabled); @@ -1413,6 +1472,7 @@ async fn test_soft_drop_metric_physical_table_remains_enabled() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut procedure).await; @@ -1457,6 +1517,7 @@ async fn test_soft_drop_file_engine_table_falls_back_to_hard_drop() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), + EventContext::default(), ); procedure.on_prepare().await.unwrap(); assert!(!procedure.data.soft_drop_enabled); @@ -1493,6 +1554,7 @@ async fn test_file_engine_fallback_is_resolved_before_tombstone_conflict() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context, + EventContext::default(), ); procedure.on_prepare().await.unwrap(); @@ -1508,8 +1570,11 @@ async fn test_undrop_metric_logical_table_fails() { let logical_table_id = create_metric_logical_table_tombstone(&ddl_context, physical_table_id, "foo").await; - let mut procedure = - UndropTableProcedure::new(new_undrop_table_task(logical_table_id), ddl_context); + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(logical_table_id), + ddl_context, + EventContext::default(), + ); let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::Unsupported); @@ -1527,6 +1592,7 @@ async fn test_purge_metric_logical_table_fails() { let mut procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(logical_table_id), ddl_context, + EventContext::default(), ); let err = procedure .execute(&new_test_procedure_context()) @@ -1558,6 +1624,7 @@ async fn test_undrop_table_fails_when_live_name_exists() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; ddl_context @@ -1570,8 +1637,11 @@ async fn test_undrop_table_fails_when_live_name_exists() { .await .unwrap(); - let mut procedure = - UndropTableProcedure::new(new_undrop_table_task(dropped_table_id), ddl_context); + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(dropped_table_id), + ddl_context, + EventContext::default(), + ); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); @@ -1599,11 +1669,15 @@ async fn test_undrop_table_fails_when_live_name_is_created_after_prepare() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; - let mut procedure = - UndropTableProcedure::new(new_undrop_table_task(dropped_table_id), ddl_context.clone()); + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(dropped_table_id), + ddl_context.clone(), + EventContext::default(), + ); procedure.on_prepare().await.unwrap(); ddl_context .table_metadata_manager @@ -1667,13 +1741,17 @@ async fn test_undrop_table_closes_opened_regions_when_restore_metadata_races_wit let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} detector_controller.clear().await; - let mut procedure = - UndropTableProcedure::new(new_undrop_table_task(dropped_table_id), ddl_context.clone()); + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(dropped_table_id), + ddl_context.clone(), + EventContext::default(), + ); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -1767,6 +1845,7 @@ async fn test_undrop_table_lock_key_includes_original_table_name_before_prepare( let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; @@ -1781,6 +1860,7 @@ async fn test_undrop_table_lock_key_includes_original_table_name_before_prepare( new_undrop_table_task(table_id), ddl_context, Some(original_table_name), + EventContext::default(), ); let keys = procedure @@ -1848,10 +1928,12 @@ async fn test_undrop_rejects_stale_id_after_name_tombstone_is_consumed() { new_undrop_table_task(original_table_id), ddl_context.clone(), Some(original_table_name.clone()), + EventContext::default(), ); let mut hard_drop = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut hard_drop).await; assert!( @@ -1906,11 +1988,15 @@ async fn test_undrop_table_replayed_restore_metadata_is_idempotent() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; - let mut procedure = - UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(table_id), + ddl_context.clone(), + EventContext::default(), + ); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -2014,6 +2100,7 @@ async fn test_purge_dropped_table_cleans_regions_offline_and_deletes_tombstone() let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} @@ -2022,6 +2109,7 @@ async fn test_purge_dropped_table_cleans_regions_offline_and_deletes_tombstone() let mut procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut procedure).await; @@ -2100,6 +2188,7 @@ async fn test_automatic_purge_rechecks_unexpired_tombstone() { let procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), + EventContext::default(), ); let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); data["check_expired"] = true.into(); @@ -2156,6 +2245,7 @@ async fn test_recovered_automatic_purge_rechecks_unexpired_tombstone() { let procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), + EventContext::default(), ); let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); data["check_expired"] = true.into(); @@ -2233,6 +2323,7 @@ async fn test_recovered_automatic_purge_obeys_runtime_disable() { let procedure = PurgeDroppedTableProcedure::new_if_expired( new_purge_dropped_table_task(table_id), ddl_context.clone(), + EventContext::default(), ); let mut recovered = PurgeDroppedTableProcedure::from_json(&procedure.dump().unwrap(), ddl_context.clone()) @@ -2264,6 +2355,7 @@ async fn test_recovered_automatic_purge_finishes_after_cleanup_is_claimed() { let mut procedure = PurgeDroppedTableProcedure::new_if_expired( new_purge_dropped_table_task(table_id), enabled_context.clone(), + EventContext::default(), ); procedure .execute(&new_test_procedure_context()) @@ -2327,6 +2419,7 @@ async fn test_recovered_automatic_purge_ignores_foreign_generation_claim() { let mut procedure = PurgeDroppedTableProcedure::new_if_expired( new_purge_dropped_table_task(table_id), ddl_context.clone(), + EventContext::default(), ); procedure .execute(&new_test_procedure_context()) @@ -2369,8 +2462,11 @@ async fn test_undrop_rejects_table_claimed_for_purge() { .await .unwrap(); - let mut procedure = - UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(table_id), + ddl_context.clone(), + EventContext::default(), + ); let error = procedure.on_prepare().await.unwrap_err(); assert_eq!(StatusCode::TableNotFound, error.status_code()); @@ -2392,8 +2488,11 @@ async fn test_recovered_undrop_rejects_tombstone_deleted_by_purge() { let table_id = 1024; create_dropped_table(&ddl_context, table_id, Some(1), Some(i64::MAX)).await; - let mut undrop = - UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + let mut undrop = UndropTableProcedure::new( + new_undrop_table_task(table_id), + ddl_context.clone(), + EventContext::default(), + ); undrop.on_prepare().await.unwrap(); let open_regions_state = undrop.dump().unwrap(); undrop.execute(&new_test_procedure_context()).await.unwrap(); @@ -2402,6 +2501,7 @@ async fn test_recovered_undrop_rejects_tombstone_deleted_by_purge() { let mut purge = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut purge).await; @@ -2433,8 +2533,11 @@ async fn test_recovered_undrop_rejects_replacement_generation() { .await .unwrap(); - let mut undrop = - UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + let mut undrop = UndropTableProcedure::new( + new_undrop_table_task(table_id), + ddl_context.clone(), + EventContext::default(), + ); undrop.on_prepare().await.unwrap(); let persisted = undrop.dump().unwrap(); kv_backend @@ -2500,8 +2603,11 @@ async fn test_legacy_unmarked_tombstone_can_be_undropped_and_purged() { assert_eq!(dropped_tables.len(), 1); assert_eq!(dropped_tables[0].dropped_at, None); - let mut undrop = - UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + let mut undrop = UndropTableProcedure::new( + new_undrop_table_task(table_id), + ddl_context.clone(), + EventContext::default(), + ); execute_procedure_until_done(&mut undrop).await; ddl_context @@ -2518,6 +2624,7 @@ async fn test_legacy_unmarked_tombstone_can_be_undropped_and_purged() { let mut purge = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut purge).await; assert!( @@ -2560,6 +2667,7 @@ async fn test_purge_dropped_table_by_id_selects_tombstone_when_live_table_exists let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; ddl_context @@ -2576,6 +2684,7 @@ async fn test_purge_dropped_table_by_id_selects_tombstone_when_live_table_exists let mut procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(dropped_table_id), ddl_context.clone(), + EventContext::default(), ); execute_procedure_until_done(&mut procedure).await; @@ -2626,8 +2735,12 @@ async fn test_on_rollback() { let physical_table_id = table_id; // Creates the logical table metadata. let task = test_create_logical_table_task("foo"); - let mut procedure = - CreateLogicalTablesProcedure::new(vec![task], physical_table_id, ddl_context.clone()); + let mut procedure = CreateLogicalTablesProcedure::new( + vec![task], + physical_table_id, + EventContext::default(), + ddl_context.clone(), + ); procedure.on_prepare().await.unwrap(); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); @@ -2640,7 +2753,8 @@ async fn test_on_rollback() { // Drops the physical table { let task = new_drop_table_task("phy_table", physical_table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); + let mut procedure = + DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); procedure.on_prepare().await.unwrap(); assert!(procedure.rollback_supported()); procedure.on_delete_metadata().await.unwrap(); @@ -2651,7 +2765,8 @@ async fn test_on_rollback() { // The physical table can be dropped again after the metadata rollback. let retry_task = new_drop_table_task("phy_table", physical_table_id, false); - let mut retry = DropTableProcedure::new(retry_task, ddl_context.clone()); + let mut retry = + DropTableProcedure::new(retry_task, ddl_context.clone(), EventContext::default()); retry.on_prepare().await.unwrap(); retry.on_delete_metadata().await.unwrap(); retry.rollback(&ctx).await.unwrap(); @@ -2668,7 +2783,7 @@ async fn test_on_rollback() { // Drops the logical table ddl_context.soft_drop_enabled = false; let task = new_drop_table_task("foo", table_ids[0], false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); procedure.on_prepare().await.unwrap(); assert!(!procedure.rollback_supported()); } @@ -2765,7 +2880,8 @@ async fn test_memory_region_keeper_guard_dropped_on_procedure_done() { let logical_table_id = create_logical_table(ddl_context.clone(), physical_table_id, "s").await; let inner_test = |task: DropTableTask| async { - let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); + let mut procedure = + DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); execute_procedure_until(&mut procedure, |p| { p.data.state == DropTableState::InvalidateTableCache }) @@ -2817,7 +2933,8 @@ async fn test_from_json() { let physical_table_id = create_physical_table(&ddl_context, "t").await; let task = new_drop_table_task("t", physical_table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); + let mut procedure = + DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); execute_procedure_until(&mut procedure, |p| p.data.state == state).await; let data = procedure.dump().unwrap(); assert_eq!( @@ -2846,7 +2963,7 @@ async fn test_from_json() { let physical_table_id = create_physical_table(&ddl_context, "t").await; let task = new_drop_table_task("t", physical_table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); execute_procedure_until_done(&mut procedure).await; let data = procedure.dump().unwrap(); assert_eq!( diff --git a/src/common/meta/src/ddl/tests/drop_view.rs b/src/common/meta/src/ddl/tests/drop_view.rs index 412c5a78f0b..09b2a450c52 100644 --- a/src/common/meta/src/ddl/tests/drop_view.rs +++ b/src/common/meta/src/ddl/tests/drop_view.rs @@ -24,7 +24,7 @@ use crate::ddl::drop_view::{DropViewProcedure, DropViewState}; use crate::ddl::test_util::create_table::test_create_table_task; use crate::ddl::tests::create_view::{test_create_view_task, test_table_names}; use crate::key::table_route::TableRouteValue; -use crate::rpc::ddl::DropViewTask; +use crate::rpc::ddl::{DropViewTask, EventContext}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; pub(crate) fn new_drop_view_task( @@ -63,7 +63,7 @@ async fn test_on_prepare_view_not_exists_err() { .unwrap(); let task = new_drop_view_task("bar", view_id, false); - let mut procedure = DropViewProcedure::new(task, ddl_context); + let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::TableNotFound); } @@ -87,7 +87,7 @@ async fn test_on_prepare_not_view_err() { .unwrap(); let task = new_drop_view_task(view_name, view_id, false); - let mut procedure = DropViewProcedure::new(task, ddl_context); + let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); // It's not a view, expect error let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::InvalidArguments); @@ -117,12 +117,12 @@ async fn test_on_prepare_success() { let task = new_drop_view_task("bar", view_id, true); // Drop if exists - let mut procedure = DropViewProcedure::new(task, ddl_context.clone()); + let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context.clone()); procedure.on_prepare().await.unwrap(); let task = new_drop_view_task(view_name, view_id, false); // Prepare success - let mut procedure = DropViewProcedure::new(task, ddl_context); + let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); procedure.on_prepare().await.unwrap(); assert_eq!(DropViewState::DeleteMetadata, procedure.state()); } @@ -161,7 +161,7 @@ async fn test_drop_view_success() { let task = new_drop_view_task(view_name, view_id, false); // Prepare success - let mut procedure = DropViewProcedure::new(task, ddl_context.clone()); + let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; assert_eq!(DropViewState::InvalidateViewCache, procedure.state()); @@ -178,7 +178,7 @@ async fn test_drop_view_success() { // Drop again let task = new_drop_view_task(view_name, view_id, false); - let mut procedure = DropViewProcedure::new(task, ddl_context); + let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::TableNotFound); } diff --git a/src/common/meta/src/ddl/tests/event/database.rs b/src/common/meta/src/ddl/tests/event/database.rs index 39c001e9ffd..8aa85c2fbbd 100644 --- a/src/common/meta/src/ddl/tests/event/database.rs +++ b/src/common/meta/src/ddl/tests/event/database.rs @@ -20,17 +20,17 @@ use api::v1::value::ValueData; use api::v1::{AlterDatabaseExpr, Row, SetDatabaseOptions as PbSetDatabaseOptions, Value}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, + CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN as EVENT_TABLE_PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN as EVENT_TABLE_PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN as EVENT_TABLE_PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN as EVENT_TABLE_PROCEDURE_TRIGGER_COLUMN, - SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, jsonb_value, + SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; -use super::test_util::assert_event_filter; +use super::test_util::{assert_event_filter, default_event_context_value, procedure_trigger_value}; use crate::ddl::alter_database::AlterDatabaseProcedure; use crate::ddl::create_database::CreateDatabaseProcedure; use crate::ddl::drop_database::DropDatabaseProcedure; @@ -39,7 +39,7 @@ use crate::ddl::event::database::{ DatabaseDdlEvent, }; use crate::rpc::ddl::{ - AlterDatabaseKind, SetDatabaseOption, SetDatabaseOptions, UnsetDatabaseOption, + AlterDatabaseKind, EventContext, SetDatabaseOption, SetDatabaseOptions, UnsetDatabaseOption, UnsetDatabaseOptions, }; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -50,7 +50,13 @@ fn test_create_database_submitted_event_contract() { ("password".to_string(), "do-not-record".to_string()), ("compaction.type".to_string(), "twcs".to_string()), ]); - let event = DatabaseDdlEvent::create_submitted("greptime", "metrics", true, &options); + let event = DatabaseDdlEvent::create_submitted( + "greptime", + "metrics", + true, + &options, + EventContext::default(), + ); assert_event_locator( &event, @@ -87,6 +93,7 @@ fn test_alter_database_set_and_unset_event_contracts() { SetDatabaseOption::Other("secret_token".to_string(), "hidden".to_string()), SetDatabaseOption::Ttl(std::time::Duration::from_secs(3600).into()), ])), + EventContext::default(), ); assert_event_locator( &set, @@ -113,6 +120,7 @@ fn test_alter_database_set_and_unset_event_contracts() { UnsetDatabaseOption::Ttl, UnsetDatabaseOption::Other("compaction.type".to_string()), ])), + EventContext::default(), ); assert_event_locator( &unset, @@ -132,7 +140,8 @@ fn test_alter_database_set_and_unset_event_contracts() { #[test] fn test_drop_database_submitted_event_contract() { - let event = DatabaseDdlEvent::drop_submitted("greptime", "metrics", true); + let event = + DatabaseDdlEvent::drop_submitted("greptime", "metrics", true, EventContext::default()); assert_event_locator( &event, @@ -165,7 +174,14 @@ fn test_database_lifecycle_events_have_fixed_schema_and_null_intent() { assert_eq!( DatabaseDdlEvent::create_lifecycle().extra_schema(), - DatabaseDdlEvent::create_submitted("c", "s", false, &HashMap::new()).extra_schema() + DatabaseDdlEvent::create_submitted( + "c", + "s", + false, + &HashMap::new(), + EventContext::default() + ) + .extra_schema() ); } @@ -179,6 +195,7 @@ fn test_database_events_preserve_procedure_envelope_contract() { "metrics", false, &HashMap::new(), + EventContext::default(), )), ProcedureState::Running, EventTrigger::Submitted, @@ -217,6 +234,7 @@ fn test_create_database_event_filter() { false, HashMap::new(), None, + EventContext::default(), new_ddl_context(Arc::new(MockDatanodeManager::new(()))), ); @@ -237,6 +255,7 @@ fn test_alter_database_event_filter() { )), }, }, + EventContext::default(), new_ddl_context(Arc::new(MockDatanodeManager::new(()))), ) .unwrap(); @@ -250,6 +269,7 @@ fn test_drop_database_event_filter() { "greptime".to_string(), "metrics".to_string(), false, + EventContext::default(), new_ddl_context(Arc::new(MockDatanodeManager::new(()))), ); @@ -268,6 +288,7 @@ fn assert_event_locator( &[ EVENT_TABLE_CATALOG_NAME_COLUMN.column_schema(), EVENT_TABLE_SCHEMA_NAME_COLUMN.column_schema(), + EVENT_CONTEXT_COLUMN.column_schema(), ], &[Row { values: vec![ @@ -277,6 +298,11 @@ fn assert_event_locator( Value { value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, + if catalog_name.is_some() { + default_event_context_value() + } else { + Value { value_data: None } + }, ], }], ); @@ -300,6 +326,7 @@ fn assert_procedure_event_contract( EVENT_TABLE_PROCEDURE_TRIGGER_COLUMN.column_schema(), EVENT_TABLE_CATALOG_NAME_COLUMN.column_schema(), EVENT_TABLE_SCHEMA_NAME_COLUMN.column_schema(), + EVENT_CONTEXT_COLUMN.column_schema(), ], &[Row { values: vec![ @@ -314,13 +341,18 @@ fn assert_procedure_event_contract( Value { value_data: Some(ValueData::StringValue(String::new())), }, - jsonb_value(&serde_json::json!({"type": procedure_trigger})), + procedure_trigger_value(procedure_trigger), Value { value_data: catalog_name.map(|value| ValueData::StringValue(value.to_string())), }, Value { value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, + if catalog_name.is_some() { + default_event_context_value() + } else { + Value { value_data: None } + }, ], }], ); diff --git a/src/common/meta/src/ddl/tests/event/flow.rs b/src/common/meta/src/ddl/tests/event/flow.rs index db3d5e2ce11..b2758197917 100644 --- a/src/common/meta/src/ddl/tests/event/flow.rs +++ b/src/common/meta/src/ddl/tests/event/flow.rs @@ -19,14 +19,14 @@ use api::v1::{ColumnSchema, Row, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, PROCEDURE_ERROR_COLUMN, - PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, + PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; use table::table_name::TableName; -use super::test_util::assert_event_filter; +use super::test_util::{assert_event_filter, default_event_context_value, procedure_trigger_value}; use crate::ddl::create_flow::CreateFlowProcedure; use crate::ddl::drop_flow::DropFlowProcedure; use crate::ddl::event::flow::{ @@ -35,6 +35,7 @@ use crate::ddl::event::flow::{ use crate::ddl::test_util::flownode_handler::NaiveFlownodeHandler; use crate::ddl::tests::create_flow::{test_create_flow_task, test_query_context}; use crate::ddl::tests::drop_flow::test_drop_flow_task; +use crate::rpc::ddl::EventContext; use crate::test_util::{MockFlownodeManager, new_ddl_context}; #[test] @@ -48,6 +49,7 @@ fn test_flow_submitted_event_contracts() { expire_after: Some(300), eval_interval_secs: Some(60), }, + EventContext::default(), ); assert_event_contract( &create, @@ -58,6 +60,7 @@ fn test_flow_submitted_event_contracts() { ValueData::StringValue("greptime".to_string()).into(), ValueData::StringValue("metrics".to_string()).into(), Value { value_data: None }, + default_event_context_value(), ], }], ); @@ -72,7 +75,8 @@ fn test_flow_submitted_event_contracts() { }) ); - let drop = FlowDdlEvent::drop_submitted("greptime", "metrics", 42, true); + let drop = + FlowDdlEvent::drop_submitted("greptime", "metrics", 42, true, EventContext::default()); assert_event_contract( &drop, DROP_FLOW_EVENT_TYPE, @@ -82,6 +86,7 @@ fn test_flow_submitted_event_contracts() { ValueData::StringValue("greptime".to_string()).into(), ValueData::StringValue("metrics".to_string()).into(), ValueData::U32Value(42).into(), + default_event_context_value(), ], }], ); @@ -107,6 +112,7 @@ fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { Value { value_data: None }, Value { value_data: None }, Value { value_data: None }, + Value { value_data: None }, ], }], ); @@ -123,6 +129,7 @@ fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { Value { value_data: None }, Value { value_data: None }, ValueData::U32Value(42).into(), + Value { value_data: None }, ], }], ); @@ -142,6 +149,7 @@ fn test_flow_events_preserve_procedure_envelope_contract() { expire_after: None, eval_interval_secs: None, }, + EventContext::default(), )), ProcedureState::Running, EventTrigger::Submitted, @@ -187,6 +195,7 @@ fn test_create_flow_event_filter() { false, ), test_query_context(), + EventContext::default(), new_ddl_context(Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler))), ); assert_event_filter(&procedure, CREATE_FLOW_EVENT_TYPE); @@ -196,6 +205,7 @@ fn test_create_flow_event_filter() { fn test_drop_flow_event_filter() { let procedure = DropFlowProcedure::new( test_drop_flow_task("flow", 42, false), + EventContext::default(), new_ddl_context(Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler))), ); assert_event_filter(&procedure, DROP_FLOW_EVENT_TYPE); @@ -206,6 +216,7 @@ fn flow_schema() -> Vec { CATALOG_NAME_COLUMN.column_schema(), FLOW_NAME_COLUMN.column_schema(), FLOW_ID_COLUMN.column_schema(), + EVENT_CONTEXT_COLUMN.column_schema(), ] } @@ -238,7 +249,7 @@ fn assert_procedure_event_contract( ValueData::StringValue(event.procedure_id.to_string()).into(), ValueData::StringValue(state.to_string()).into(), ValueData::StringValue(String::new()).into(), - jsonb_value(&serde_json::json!({"type": trigger})), + procedure_trigger_value(trigger), optional_string(locator.catalog_name), optional_string(locator.flow_name), locator @@ -246,6 +257,11 @@ fn assert_procedure_event_contract( .map(ValueData::U32Value) .map(Into::into) .unwrap_or(Value { value_data: None }), + if locator.catalog_name.is_some() { + default_event_context_value() + } else { + Value { value_data: None } + }, ], }], ); diff --git a/src/common/meta/src/ddl/tests/event/table.rs b/src/common/meta/src/ddl/tests/event/table.rs index c0a731a20e6..90e7d83b3f2 100644 --- a/src/common/meta/src/ddl/tests/event/table.rs +++ b/src/common/meta/src/ddl/tests/event/table.rs @@ -19,13 +19,13 @@ use api::v1::value::ValueData; use api::v1::{ColumnDataType, Repartition, SemanticType, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN, - TABLE_NAME_COLUMN, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, + TABLE_ID_COLUMN, TABLE_NAME_COLUMN, jsonb_value, }; use common_event_recorder::{Event, EventTypeFilter}; use common_procedure::{ - ChildSubmissionOutcome, EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState, - RetryPhase, + ChildSubmissionOutcome, EventRuntimeContext, EventTrigger, Procedure, ProcedureId, + ProcedureState, RetryPhase, }; use common_time::Timestamp; use serde_json::{Value as JsonValue, json}; @@ -52,9 +52,9 @@ use crate::ddl::truncate_table::TruncateTableProcedure; use crate::ddl::undrop_table::UndropTableProcedure; use crate::key::DeserializedValueWithBytes; use crate::key::table_info::TableInfoValue; -use crate::rpc::ddl::{DropTableTask, TruncateTableTask}; +use crate::rpc::ddl::{DropTableTask, EventContext, QueryContext, TruncateTableTask}; #[cfg(feature = "enterprise")] -use crate::rpc::ddl::{PurgeDroppedTableTask, UndropTableTask}; +use crate::rpc::ddl::{PurgeDroppedTableTask, TriggerReason, UndropTableTask}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; struct EventCase { @@ -177,7 +177,13 @@ fn later_lifecycle_events_are_uniform() { fn create_success_events_keep_allocated_ids() { let mut task = test_create_table_task("create_success"); task.table_info.ident.table_id = 7; - let create_table = CreateTableProcedure::new(task, test_context()).unwrap(); + let create_table = CreateTableProcedure::new( + task, + QueryContext::default(), + EventContext::default(), + test_context(), + ) + .unwrap(); let state = ProcedureState::Done { output: Some(Arc::new(42_u32)), }; @@ -196,6 +202,7 @@ fn create_success_events_keep_allocated_ids() { test_create_logical_table_task("bar"), ], 1024, + EventContext::default(), test_context(), ); let state = ProcedureState::Done { @@ -234,13 +241,14 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "create"), true, "mito2", + EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "create_if_not_exists": true, "engine": "mito2", }), - rows: vec![table_locator_values(Some("create"), None)], + rows: vec![submitted_table_locator_values(Some("create"), None)], }, EventCase { event_type: TableDdlEventType::CreateLogicalTables, @@ -252,14 +260,15 @@ fn event_cases() -> Vec { .with_physical_table_id(10), ], 2, + EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "table_count": 2, }), rows: vec![ - logical_locator_values("logical1", None, 10), - logical_locator_values("logical2", None, 10), + submitted_logical_locator_values("logical1", None, 10), + submitted_logical_locator_values("logical2", None, 10), ], }, EventCase { @@ -268,12 +277,13 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "alter") .with_table_id(11), Some("drop_columns"), + EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "kind": "drop_columns", }), - rows: vec![table_locator_values(Some("alter"), Some(11))], + rows: vec![submitted_table_locator_values(Some("alter"), Some(11))], }, EventCase { event_type: TableDdlEventType::AlterLogicalTables, @@ -286,6 +296,7 @@ fn event_cases() -> Vec { ], 2, ["rename_table", "add_columns", "add_columns"], + EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, @@ -293,8 +304,8 @@ fn event_cases() -> Vec { "kinds": ["add_columns", "rename_table"], }), rows: vec![ - logical_locator_values("logical1", None, 10), - logical_locator_values("logical2", None, 10), + submitted_logical_locator_values("logical1", None, 10), + submitted_logical_locator_values("logical2", None, 10), ], }, EventCase { @@ -303,26 +314,33 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "drop") .with_table_id(12), true, + EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "drop_if_exists": true, }), - rows: vec![table_locator_values(Some("drop"), Some(12))], + rows: vec![submitted_table_locator_values(Some("drop"), Some(12))], }, #[cfg(feature = "enterprise")] EventCase { event_type: TableDdlEventType::UndropTable, - event: TableDdlEvent::undrop_table_submitted(TableDdlLocator::from_table_id(13)), + event: TableDdlEvent::undrop_table_submitted( + TableDdlLocator::from_table_id(13), + EventContext::default(), + ), payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![table_locator_values(None, Some(13))], + rows: vec![submitted_table_locator_values(None, Some(13))], }, #[cfg(feature = "enterprise")] EventCase { event_type: TableDdlEventType::PurgeDroppedTable, - event: TableDdlEvent::purge_dropped_table_submitted(TableDdlLocator::from_table_id(14)), + event: TableDdlEvent::purge_dropped_table_submitted( + TableDdlLocator::from_table_id(14), + EventContext::default(), + ), payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![table_locator_values(None, Some(14))], + rows: vec![submitted_table_locator_values(None, Some(14))], }, EventCase { event_type: TableDdlEventType::TruncateTable, @@ -330,29 +348,41 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "truncate") .with_table_id(15), 4, + EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "time_range_count": 4, }), - rows: vec![table_locator_values(Some("truncate"), Some(15))], + rows: vec![submitted_table_locator_values(Some("truncate"), Some(15))], }, ] } fn procedure_cases() -> Vec { - let create_table = - CreateTableProcedure::new(test_create_table_task("create"), test_context()).unwrap(); + let create_table = CreateTableProcedure::new( + test_create_table_task("create"), + QueryContext::default(), + EventContext::default(), + test_context(), + ) + .unwrap(); let create_logical_tables = CreateLogicalTablesProcedure::new( vec![ test_create_logical_table_task("logical1"), test_create_logical_table_task("logical2"), ], 41, + EventContext::default(), test_context(), ); - let alter_table = - AlterTableProcedure::new(42, test_alter_table_task("alter"), test_context()).unwrap(); + let alter_table = AlterTableProcedure::new( + 42, + test_alter_table_task("alter"), + EventContext::default(), + test_context(), + ) + .unwrap(); let alter_logical_tables = AlterLogicalTablesProcedure::new( vec![ make_alter_logical_table_add_column_task( @@ -367,6 +397,7 @@ fn procedure_cases() -> Vec { ), ], 43, + EventContext::default(), test_context(), ); let drop_table = DropTableProcedure::new( @@ -378,12 +409,20 @@ fn procedure_cases() -> Vec { drop_if_exists: true, }, test_context(), + EventContext::default(), ); #[cfg(feature = "enterprise")] - let undrop_table = UndropTableProcedure::new(UndropTableTask { table_id: 45 }, test_context()); + let undrop_table = UndropTableProcedure::new( + UndropTableTask { table_id: 45 }, + test_context(), + EventContext::default(), + ); #[cfg(feature = "enterprise")] - let purge_dropped_table = - PurgeDroppedTableProcedure::new(PurgeDroppedTableTask { table_id: 46 }, test_context()); + let purge_dropped_table = PurgeDroppedTableProcedure::new_if_expired( + PurgeDroppedTableTask { table_id: 46 }, + test_context(), + EventContext::new(TriggerReason::ScheduledGc), + ); let truncate_table = truncate_procedure(TruncateTableTask { catalog: DEFAULT_CATALOG_NAME.to_string(), schema: DEFAULT_SCHEMA_NAME.to_string(), @@ -404,7 +443,7 @@ fn procedure_cases() -> Vec { "create_if_not_exists": false, "engine": "mito2", }), - rows: vec![table_locator_values(Some("create"), None)], + rows: vec![submitted_table_locator_values(Some("create"), None)], }, ProcedureCase { procedure: Box::new(create_logical_tables), @@ -414,8 +453,8 @@ fn procedure_cases() -> Vec { "table_count": 2, }), rows: vec![ - logical_locator_values("logical1", None, 41), - logical_locator_values("logical2", None, 41), + submitted_logical_locator_values("logical1", None, 41), + submitted_logical_locator_values("logical2", None, 41), ], }, ProcedureCase { @@ -425,7 +464,7 @@ fn procedure_cases() -> Vec { "version": TABLE_DDL_PAYLOAD_VERSION, "kind": "drop_columns", }), - rows: vec![table_locator_values(Some("alter"), Some(42))], + rows: vec![submitted_table_locator_values(Some("alter"), Some(42))], }, ProcedureCase { procedure: Box::new(alter_logical_tables), @@ -436,8 +475,8 @@ fn procedure_cases() -> Vec { "kinds": ["add_columns"], }), rows: vec![ - logical_locator_values("logical1", None, 43), - logical_locator_values("logical2", None, 43), + submitted_logical_locator_values("logical1", None, 43), + submitted_logical_locator_values("logical2", None, 43), ], }, ProcedureCase { @@ -447,21 +486,25 @@ fn procedure_cases() -> Vec { "version": TABLE_DDL_PAYLOAD_VERSION, "drop_if_exists": true, }), - rows: vec![table_locator_values(Some("drop"), Some(44))], + rows: vec![submitted_table_locator_values(Some("drop"), Some(44))], }, #[cfg(feature = "enterprise")] ProcedureCase { procedure: Box::new(undrop_table), event_type: "undrop_table", payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![table_locator_values(None, Some(45))], + rows: vec![submitted_table_locator_values(None, Some(45))], }, #[cfg(feature = "enterprise")] ProcedureCase { procedure: Box::new(purge_dropped_table), event_type: "purge_dropped_table", payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![table_locator_values(None, Some(46))], + rows: vec![submitted_table_locator_values_with_event_context( + None, + Some(46), + EventContext::new(TriggerReason::ScheduledGc), + )], }, ProcedureCase { procedure: Box::new(truncate_table), @@ -470,7 +513,7 @@ fn procedure_cases() -> Vec { "version": TABLE_DDL_PAYLOAD_VERSION, "time_range_count": 1, }), - rows: vec![table_locator_values(Some("truncate"), Some(47))], + rows: vec![submitted_table_locator_values(Some("truncate"), Some(47))], }, ] } @@ -491,6 +534,7 @@ fn expected_schema(event_type: TableDdlEventType) -> Vec<(&'static str, i32)> { ColumnDataType::Uint32 as i32, )); } + schema.push((EVENT_CONTEXT_COLUMN.name(), ColumnDataType::Binary as i32)); schema } @@ -503,12 +547,14 @@ fn table_locator_values(table_name: Option<&str>, table_id: Option) -> Vec< } else { (Value::default(), Value::default()) }; - vec![ + let mut values = vec![ catalog_name, schema_name, table_name.map(string_value).unwrap_or_default(), table_id.map(table_id_value).unwrap_or_default(), - ] + ]; + values.push(Value::default()); + values } fn logical_locator_values( @@ -517,10 +563,44 @@ fn logical_locator_values( physical_table_id: u32, ) -> Vec { let mut values = table_locator_values(Some(table_name), table_id); + let event_context = values.pop().unwrap(); values.push(table_id_value(physical_table_id)); + values.push(event_context); values } +fn submitted_table_locator_values(table_name: Option<&str>, table_id: Option) -> Vec { + submitted_table_locator_values_with_event_context(table_name, table_id, EventContext::default()) +} + +fn submitted_table_locator_values_with_event_context( + table_name: Option<&str>, + table_id: Option, + event_context: EventContext, +) -> Vec { + let mut values = table_locator_values(table_name, table_id); + *values.last_mut().unwrap() = event_context_value_for(event_context); + values +} + +fn submitted_logical_locator_values( + table_name: &str, + table_id: Option, + physical_table_id: u32, +) -> Vec { + let mut values = logical_locator_values(table_name, table_id, physical_table_id); + *values.last_mut().unwrap() = event_context_value(); + values +} + +fn event_context_value() -> Value { + event_context_value_for(EventContext::default()) +} + +fn event_context_value_for(event_context: EventContext) -> Value { + jsonb_value(&serde_json::to_value(event_context).unwrap()) +} + fn string_value(value: &str) -> Value { ValueData::StringValue(value.to_string()).into() } @@ -535,6 +615,7 @@ fn truncate_procedure(task: TruncateTableTask) -> TruncateTableProcedure { task, DeserializedValueWithBytes::from_inner(TableInfoValue::new(table_info)), test_context(), + EventContext::default(), ) } @@ -552,7 +633,7 @@ fn event_for_state( lifecycle_state: &ProcedureState, ) -> Box { procedure - .event(&EventContext { + .event(&EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, diff --git a/src/common/meta/src/ddl/tests/event/test_util.rs b/src/common/meta/src/ddl/tests/event/test_util.rs index 884fa3ec4d4..281b5deb488 100644 --- a/src/common/meta/src/ddl/tests/event/test_util.rs +++ b/src/common/meta/src/ddl/tests/event/test_util.rs @@ -15,12 +15,24 @@ use std::collections::HashSet; use std::sync::Arc; +use api::v1::Value; use common_event_recorder::EventTypeFilter; -use common_procedure::{EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState}; +use common_event_recorder::event_table::jsonb_value; +use common_procedure::{EventRuntimeContext, EventTrigger, Procedure, ProcedureId, ProcedureState}; + +use crate::rpc::ddl::EventContext; + +pub(crate) fn default_event_context_value() -> Value { + jsonb_value(&serde_json::to_value(EventContext::default()).unwrap()) +} + +pub(crate) fn procedure_trigger_value(trigger: &str) -> Value { + jsonb_value(&serde_json::json!({"type": trigger})) +} pub(crate) fn assert_event_filter(procedure: &dyn Procedure, event_type: &str) { let state = ProcedureState::Running; - let event_context = |event_type_filter| EventContext { + let runtime_context = |event_type_filter| EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, @@ -28,7 +40,7 @@ pub(crate) fn assert_event_filter(procedure: &dyn Procedure, event_type: &str) { }; let allowed = procedure - .event(&event_context(EventTypeFilter::Only(HashSet::from([ + .event(&runtime_context(EventTypeFilter::Only(HashSet::from([ event_type.to_string(), ])))) .unwrap(); @@ -37,7 +49,7 @@ pub(crate) fn assert_event_filter(procedure: &dyn Procedure, event_type: &str) { for denied in [HashSet::from(["other_event".to_string()]), HashSet::new()] { assert!( procedure - .event(&event_context(EventTypeFilter::Only(denied))) + .event(&runtime_context(EventTypeFilter::Only(denied))) .is_none() ); } diff --git a/src/common/meta/src/ddl/tests/event/view.rs b/src/common/meta/src/ddl/tests/event/view.rs index 024aa023c87..10947e45ee8 100644 --- a/src/common/meta/src/ddl/tests/event/view.rs +++ b/src/common/meta/src/ddl/tests/event/view.rs @@ -17,17 +17,18 @@ use std::sync::Arc; use api::v1::value::ValueData; use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, - PROCEDURE_TRIGGER_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, VIEW_NAME_COLUMN, jsonb_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, + PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, + VIEW_NAME_COLUMN, }; use common_event_recorder::testing::assert_event_contract; use common_event_recorder::{Event, EventTypeFilter}; use common_procedure::{ - ChildSubmissionOutcome, EventContext, EventTrigger, Output, Procedure, ProcedureEvent, + ChildSubmissionOutcome, EventRuntimeContext, EventTrigger, Output, Procedure, ProcedureEvent, ProcedureId, ProcedureState, RetryPhase, }; -use super::test_util::assert_event_filter; +use super::test_util::{assert_event_filter, default_event_context_value, procedure_trigger_value}; use crate::ddl::create_view::CreateViewProcedure; use crate::ddl::drop_view::DropViewProcedure; use crate::ddl::event::view::{ @@ -35,6 +36,7 @@ use crate::ddl::event::view::{ }; use crate::ddl::tests::create_view::test_create_view_task; use crate::ddl::tests::drop_view::new_drop_view_task; +use crate::rpc::ddl::EventContext; use crate::test_util::{MockDatanodeManager, new_ddl_context}; #[test] @@ -42,7 +44,7 @@ fn test_view_submitted_event_contracts() { let mut task = test_create_view_task("v_metrics"); task.create_view.or_replace = true; task.create_view.create_if_not_exists = true; - let create = CreateViewProcedure::new(task, test_context()); + let create = CreateViewProcedure::new(task, EventContext::default(), test_context()); let event = event_for(&create, EventTrigger::Submitted); assert_view_event_contract( @@ -70,7 +72,11 @@ fn test_view_submitted_event_contracts() { assert!(!payload.contains(omitted)); } - let drop = DropViewProcedure::new(new_drop_view_task("view_name", 42, true), test_context()); + let drop = DropViewProcedure::new( + new_drop_view_task("view_name", 42, true), + EventContext::default(), + test_context(), + ); let event = event_for(&drop, EventTrigger::Submitted); assert_view_event_contract( @@ -116,8 +122,16 @@ fn test_view_lifecycle_event_contracts() { #[test] fn test_view_procedures_emit_lightweight_lifecycle_events() { - let create = CreateViewProcedure::new(test_create_view_task("view_name"), test_context()); - let drop = DropViewProcedure::new(new_drop_view_task("view_name", 42, false), test_context()); + let create = CreateViewProcedure::new( + test_create_view_task("view_name"), + EventContext::default(), + test_context(), + ); + let drop = DropViewProcedure::new( + new_drop_view_task("view_name", 42, false), + EventContext::default(), + test_context(), + ); let triggers = [ EventTrigger::Recovered, EventTrigger::ChildSubmitted { @@ -149,7 +163,11 @@ fn test_view_procedures_emit_lightweight_lifecycle_events() { #[test] fn test_create_view_succeeded_output_mapping() { - let procedure = CreateViewProcedure::new(test_create_view_task("view_name"), test_context()); + let procedure = CreateViewProcedure::new( + test_create_view_task("view_name"), + EventContext::default(), + test_context(), + ); let state = ProcedureState::Done { output: Some(Arc::new(84_u32)), }; @@ -175,14 +193,21 @@ fn test_create_view_succeeded_output_mapping() { #[test] fn test_create_view_event_filter() { - let procedure = CreateViewProcedure::new(test_create_view_task("view_name"), test_context()); + let procedure = CreateViewProcedure::new( + test_create_view_task("view_name"), + EventContext::default(), + test_context(), + ); assert_event_filter(&procedure, CREATE_VIEW_EVENT_TYPE); } #[test] fn test_drop_view_event_filter() { - let procedure = - DropViewProcedure::new(new_drop_view_task("view_name", 42, false), test_context()); + let procedure = DropViewProcedure::new( + new_drop_view_task("view_name", 42, false), + EventContext::default(), + test_context(), + ); assert_event_filter(&procedure, DROP_VIEW_EVENT_TYPE); } @@ -201,6 +226,7 @@ fn test_view_event_procedure_envelope_contract() { referenced_table_count: 1, column_count: 1, }, + EventContext::default(), )), ProcedureState::Running, EventTrigger::Submitted, @@ -254,6 +280,11 @@ impl ViewEventLocator<'_> { .map(ValueData::U32Value) .map(Into::into) .unwrap_or_default(), + if self.catalog_name.is_some() { + default_event_context_value() + } else { + Value { value_data: None } + }, ] } } @@ -264,6 +295,7 @@ fn view_schema() -> Vec { SCHEMA_NAME_COLUMN.column_schema(), VIEW_NAME_COLUMN.column_schema(), VIEW_ID_COLUMN.column_schema(), + EVENT_CONTEXT_COLUMN.column_schema(), ] } @@ -302,7 +334,7 @@ fn assert_procedure_event_contract( ValueData::StringValue(event.procedure_id.to_string()).into(), ValueData::StringValue(state.to_string()).into(), ValueData::StringValue(String::new()).into(), - jsonb_value(&serde_json::json!({"type": trigger})), + procedure_trigger_value(trigger), ]; values.extend(locator.values()); @@ -323,7 +355,7 @@ fn event_for_state( lifecycle_state: &ProcedureState, ) -> Box { procedure - .event(&EventContext { + .event(&EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, diff --git a/src/common/meta/src/ddl/truncate_table.rs b/src/common/meta/src/ddl/truncate_table.rs index 454b7e2de45..9d0ed09e8c8 100644 --- a/src/common/meta/src/ddl/truncate_table.rs +++ b/src/common/meta/src/ddl/truncate_table.rs @@ -20,7 +20,7 @@ use api::v1::region::{ use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::debug; @@ -43,7 +43,7 @@ use crate::key::table_info::TableInfoValue; use crate::key::table_name::TableNameKey; use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; use crate::metrics; -use crate::rpc::ddl::TruncateTableTask; +use crate::rpc::ddl::{EventContext, TruncateTableTask}; use crate::rpc::router::{find_leader_regions, find_leaders}; pub struct TruncateTableProcedure { @@ -89,7 +89,10 @@ impl Procedure for TruncateTableProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::TruncateTable.as_str()) @@ -101,7 +104,11 @@ impl Procedure for TruncateTableProcedure { let task = &self.data.task; let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) .with_table_id(task.table_id); - TableDdlEvent::truncate_table_submitted(locator, task.time_ranges.len()) + TableDdlEvent::truncate_table_submitted( + locator, + task.time_ranges.len(), + self.data.event_context.clone(), + ) } _ => TableDdlEvent::lifecycle(TableDdlEventType::TruncateTable), }; @@ -117,10 +124,11 @@ impl TruncateTableProcedure { task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, context: DdlContext, + event_context: EventContext, ) -> Self { Self { context, - data: TruncateTableData::new(task, table_info_value), + data: TruncateTableData::new(task, table_info_value, event_context), } } @@ -227,17 +235,21 @@ pub struct TruncateTableData { state: TruncateTableState, task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, + #[serde(default)] + event_context: EventContext, } impl TruncateTableData { pub fn new( task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, + event_context: EventContext, ) -> Self { Self { state: TruncateTableState::Prepare, task, table_info_value, + event_context, } } diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs index 816ba3cd9b3..635d15ee6a5 100644 --- a/src/common/meta/src/ddl/undrop_table.rs +++ b/src/common/meta/src/ddl/undrop_table.rs @@ -17,7 +17,7 @@ use api::v1::region::{ use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::tracing_context::TracingContext; @@ -43,7 +43,7 @@ use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; -use crate::rpc::ddl::UndropTableTask; +use crate::rpc::ddl::{EventContext, UndropTableTask}; use crate::rpc::router::{ RegionRoute, find_follower_regions, find_followers, find_leader_regions, find_leaders, }; @@ -56,16 +56,17 @@ pub struct UndropTableProcedure { impl UndropTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::UndropTable"; - pub fn new(task: UndropTableTask, context: DdlContext) -> Self { - Self::new_with_original_table_name(task, context, None) + pub fn new(task: UndropTableTask, context: DdlContext, event_context: EventContext) -> Self { + Self::new_with_original_table_name(task, context, None, event_context) } pub(crate) fn new_with_original_table_name( task: UndropTableTask, context: DdlContext, table_name: Option, + event_context: EventContext, ) -> Self { - let mut data = UndropTableData::new(task); + let mut data = UndropTableData::new(task, event_context); data.table_name = table_name; Self { context, data } } @@ -426,7 +427,10 @@ impl Procedure for UndropTableProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::UndropTable.as_str()) @@ -448,7 +452,7 @@ impl Procedure for UndropTableProcedure { }) .unwrap_or_default() .with_table_id(self.data.task.table_id); - TableDdlEvent::undrop_table_submitted(locator) + TableDdlEvent::undrop_table_submitted(locator, self.data.event_context.clone()) } _ => TableDdlEvent::lifecycle(TableDdlEventType::UndropTable), }; @@ -556,10 +560,12 @@ pub struct UndropTableData { drop_generation: Option, #[serde(default)] tombstone_identity_loaded: bool, + #[serde(default)] + event_context: EventContext, } impl UndropTableData { - fn new(task: UndropTableTask) -> Self { + fn new(task: UndropTableTask, event_context: EventContext) -> Self { Self { state: UndropTableState::Prepare, task, @@ -571,6 +577,7 @@ impl UndropTableData { retention_expires_at: None, drop_generation: None, tombstone_identity_loaded: false, + event_context, } } @@ -635,7 +642,11 @@ mod tests { #[test] fn test_map_restore_metadata_error_without_table_name() { let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); - let procedure = UndropTableProcedure::new(UndropTableTask { table_id: 42 }, context); + let procedure = UndropTableProcedure::new( + UndropTableTask { table_id: 42 }, + context, + EventContext::default(), + ); let err = procedure.map_restore_metadata_error( error::TombstoneTargetAlreadyExistsSnafu { @@ -650,8 +661,11 @@ mod tests { #[test] fn test_recovered_legacy_snapshot_reloads_tombstone_identity() { let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); - let mut procedure = - UndropTableProcedure::new(UndropTableTask { table_id: 42 }, context.clone()); + let mut procedure = UndropTableProcedure::new( + UndropTableTask { table_id: 42 }, + context.clone(), + EventContext::default(), + ); procedure.data.state = UndropTableState::OpenRegions; let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); let data = data.as_object_mut().unwrap(); @@ -676,7 +690,11 @@ mod tests { let table_id = 1024; let region_id = RegionId::new(table_id, 1); - let mut procedure = UndropTableProcedure::new(UndropTableTask { table_id }, context); + let mut procedure = UndropTableProcedure::new( + UndropTableTask { table_id }, + context, + EventContext::default(), + ); procedure.data.table_route_value = Some(TableRouteValue::physical(vec![RegionRoute { region: Region::new_test(region_id), leader_peer: Some(Peer::empty(1)), @@ -715,7 +733,11 @@ mod tests { let table_id = 1024; let region_id = RegionId::new(table_id, 1); - let mut procedure = UndropTableProcedure::new(UndropTableTask { table_id }, context); + let mut procedure = UndropTableProcedure::new( + UndropTableTask { table_id }, + context, + EventContext::default(), + ); procedure.data.table_name = Some(TableName::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index e333b3ccb7f..137a16793d7 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -70,13 +70,15 @@ use crate::rpc::ddl::DdlTask::{ DropTable, DropView, PurgeDroppedTable, TruncateTable, UndropTable, }; #[cfg(feature = "enterprise")] +use crate::rpc::ddl::TriggerReason; +#[cfg(feature = "enterprise")] use crate::rpc::ddl::trigger::CreateTriggerTask; #[cfg(feature = "enterprise")] use crate::rpc::ddl::trigger::DropTriggerTask; use crate::rpc::ddl::{ AlterDatabaseTask, AlterTableTask, CommentOnTask, CreateDatabaseTask, CreateFlowTask, CreateTableTask, CreateViewTask, DropDatabaseTask, DropFlowTask, DropTableTask, DropViewTask, - PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse, + EventContext, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TruncateTableTask, UndropTableTask, }; @@ -177,6 +179,7 @@ pub enum RepartitionSource { #[async_trait::async_trait] pub trait RepartitionProcedureFactory: Send + Sync { + #[allow(clippy::too_many_arguments)] fn create( &self, ddl_ctx: &DdlContext, @@ -185,6 +188,7 @@ pub trait RepartitionProcedureFactory: Send + Sync { source: RepartitionSource, to_exprs: Vec, timeout: Option, + event_context: EventContext, ) -> std::result::Result; fn register_loaders( @@ -338,6 +342,7 @@ impl DdlManager { repartition: Repartition, wait: bool, timeout: Duration, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -373,6 +378,7 @@ impl DdlManager { source, into_partition_exprs, Some(timeout), + event_context, ) .context(CreateRepartitionProcedureSnafu)?; self.repartition_procedure_factory @@ -395,6 +401,7 @@ impl DdlManager { &self, table_id: TableId, alter_table_task: AlterTableTask, + event_context: EventContext, ddl_options: DdlOptions, ) -> Result<(ProcedureId, Option)> { // make alter_table_task mutable so we can call .take() on its field @@ -415,6 +422,7 @@ impl DdlManager { repartition, ddl_options.wait, ddl_options.timeout, + event_context, ) .await; } @@ -448,6 +456,7 @@ impl DdlManager { let procedure = AlterTableProcedure::new_with_region_locks( table_id, alter_table_task.clone(), + event_context.clone(), region_ids_to_lock, context, )?; @@ -481,14 +490,12 @@ impl DdlManager { &self, create_table_task: CreateTableTask, query_context: QueryContext, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = CreateTableProcedure::new_with_query_context( - create_table_task, - query_context, - context, - )?; + let procedure = + CreateTableProcedure::new(create_table_task, query_context, event_context, context)?; let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); @@ -500,10 +507,11 @@ impl DdlManager { pub async fn submit_create_view_task( &self, create_view_task: CreateViewTask, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = CreateViewProcedure::new(create_view_task, context); + let procedure = CreateViewProcedure::new(create_view_task, event_context, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); @@ -516,11 +524,16 @@ impl DdlManager { &self, create_table_tasks: Vec, physical_table_id: TableId, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = - CreateLogicalTablesProcedure::new(create_table_tasks, physical_table_id, context); + let procedure = CreateLogicalTablesProcedure::new( + create_table_tasks, + physical_table_id, + event_context, + context, + ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); @@ -533,11 +546,16 @@ impl DdlManager { &self, alter_table_tasks: Vec, physical_table_id: TableId, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = - AlterLogicalTablesProcedure::new(alter_table_tasks, physical_table_id, context); + let procedure = AlterLogicalTablesProcedure::new( + alter_table_tasks, + physical_table_id, + event_context, + context, + ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); @@ -549,10 +567,11 @@ impl DdlManager { pub async fn submit_drop_table_task( &self, drop_table_task: DropTableTask, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = DropTableProcedure::new(drop_table_task, context); + let procedure = DropTableProcedure::new(drop_table_task, context, event_context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); @@ -565,6 +584,7 @@ impl DdlManager { pub async fn submit_undrop_table_task( &self, undrop_table_task: UndropTableTask, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -590,6 +610,7 @@ impl DdlManager { undrop_table_task, context, Some(original_table_name), + event_context, ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); @@ -603,6 +624,7 @@ impl DdlManager { pub async fn submit_purge_dropped_table_task( &self, purge_dropped_table_task: PurgeDroppedTableTask, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -616,7 +638,8 @@ impl DdlManager { #[cfg(feature = "enterprise")] { let context = self.create_context(); - let procedure = PurgeDroppedTableProcedure::new(purge_dropped_table_task, context); + let procedure = + PurgeDroppedTableProcedure::new(purge_dropped_table_task, context, event_context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); self.execute_procedure_and_wait(procedure_with_id).await @@ -643,8 +666,11 @@ impl DdlManager { #[cfg(feature = "enterprise")] { let context = self.create_context(); - let procedure = - PurgeDroppedTableProcedure::new_if_expired(purge_dropped_table_task, context); + let procedure = PurgeDroppedTableProcedure::new_if_expired( + purge_dropped_table_task, + context, + EventContext::new(TriggerReason::ScheduledGc), + ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); self.execute_procedure_and_wait(procedure_with_id).await @@ -662,6 +688,7 @@ impl DdlManager { options, creator, }: CreateDatabaseTask, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateDatabaseProcedure::new( @@ -670,6 +697,7 @@ impl DdlManager { create_if_not_exists, options, creator, + event_context, context, ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); @@ -686,9 +714,11 @@ impl DdlManager { schema, drop_if_exists, }: DropDatabaseTask, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = DropDatabaseProcedure::new(catalog, schema, drop_if_exists, context); + let procedure = + DropDatabaseProcedure::new(catalog, schema, drop_if_exists, event_context, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); self.execute_procedure_and_wait(procedure_with_id).await @@ -697,9 +727,10 @@ impl DdlManager { pub async fn submit_alter_database( &self, alter_database_task: AlterDatabaseTask, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = AlterDatabaseProcedure::new(alter_database_task, context)?; + let procedure = AlterDatabaseProcedure::new(alter_database_task, event_context, context)?; let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); self.execute_procedure_and_wait(procedure_with_id).await @@ -711,9 +742,11 @@ impl DdlManager { &self, create_flow: CreateFlowTask, query_context: QueryContext, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = CreateFlowProcedure::new(create_flow, query_context, context); + let procedure = + CreateFlowProcedure::new(create_flow, query_context, event_context, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); self.execute_procedure_and_wait(procedure_with_id).await @@ -724,9 +757,10 @@ impl DdlManager { pub async fn submit_drop_flow_task( &self, drop_flow: DropFlowTask, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = DropFlowProcedure::new(drop_flow, context); + let procedure = DropFlowProcedure::new(drop_flow, event_context, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); self.execute_procedure_and_wait(procedure_with_id).await @@ -737,9 +771,10 @@ impl DdlManager { pub async fn submit_drop_view_task( &self, drop_view: DropViewTask, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = DropViewProcedure::new(drop_view, context); + let procedure = DropViewProcedure::new(drop_view, event_context, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); self.execute_procedure_and_wait(procedure_with_id).await @@ -751,9 +786,15 @@ impl DdlManager { &self, truncate_table_task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, + event_context: EventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = TruncateTableProcedure::new(truncate_table_task, table_info_value, context); + let procedure = TruncateTableProcedure::new( + truncate_table_task, + table_info_value, + context, + event_context, + ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); @@ -822,62 +863,75 @@ impl DdlManager { .map(TracingContext::from_w3c) .unwrap_or_else(TracingContext::from_current_span) .attach(tracing::info_span!("DdlManager::submit_ddl_task")); - let ddl_options = DdlOptions { - wait: request.wait, - timeout: request.timeout, - }; + let SubmitDdlTaskRequest { + query_context, + wait, + timeout, + task, + } = request; + let event_context = EventContext::from_query_context(&query_context); + let ddl_options = DdlOptions { wait, timeout }; async move { - debug!("Submitting Ddl task: {:?}", request.task); - match request.task { + debug!("Submitting Ddl task: {:?}", task); + match task { CreateTable(create_table_task) => { - handle_create_table_task(self, create_table_task, request.query_context).await + handle_create_table_task(self, create_table_task, query_context, event_context) + .await + } + DropTable(drop_table_task) => { + handle_drop_table_task(self, drop_table_task, event_context).await } - DropTable(drop_table_task) => handle_drop_table_task(self, drop_table_task).await, UndropTable(undrop_table_task) => { - handle_undrop_table_task(self, undrop_table_task).await + handle_undrop_table_task(self, undrop_table_task, event_context).await } PurgeDroppedTable(purge_dropped_table_task) => { - handle_purge_dropped_table_task(self, purge_dropped_table_task).await + handle_purge_dropped_table_task(self, purge_dropped_table_task, event_context) + .await } AlterTable(alter_table_task) => { - handle_alter_table_task(self, alter_table_task, ddl_options).await + handle_alter_table_task(self, alter_table_task, ddl_options, event_context) + .await } TruncateTable(truncate_table_task) => { - handle_truncate_table_task(self, truncate_table_task).await + handle_truncate_table_task(self, truncate_table_task, event_context).await } CreateLogicalTables(create_table_tasks) => { - handle_create_logical_table_tasks(self, create_table_tasks).await + handle_create_logical_table_tasks(self, create_table_tasks, event_context).await } AlterLogicalTables(alter_table_tasks) => { - handle_alter_logical_table_tasks(self, alter_table_tasks).await + handle_alter_logical_table_tasks(self, alter_table_tasks, event_context).await } DropLogicalTables(_) => todo!(), CreateDatabase(create_database_task) => { - handle_create_database_task(self, create_database_task).await + handle_create_database_task(self, create_database_task, event_context).await } DropDatabase(drop_database_task) => { - handle_drop_database_task(self, drop_database_task).await + handle_drop_database_task(self, drop_database_task, event_context).await } AlterDatabase(alter_database_task) => { - handle_alter_database_task(self, alter_database_task).await + handle_alter_database_task(self, alter_database_task, event_context).await } CreateFlow(create_flow_task) => { - handle_create_flow_task(self, create_flow_task, request.query_context).await + handle_create_flow_task(self, create_flow_task, query_context, event_context) + .await + } + DropFlow(drop_flow_task) => { + handle_drop_flow_task(self, drop_flow_task, event_context).await } - DropFlow(drop_flow_task) => handle_drop_flow_task(self, drop_flow_task).await, CreateView(create_view_task) => { - handle_create_view_task(self, create_view_task).await + handle_create_view_task(self, create_view_task, event_context).await + } + DropView(drop_view_task) => { + handle_drop_view_task(self, drop_view_task, event_context).await } - DropView(drop_view_task) => handle_drop_view_task(self, drop_view_task).await, CommentOn(comment_on_task) => handle_comment_on_task(self, comment_on_task).await, #[cfg(feature = "enterprise")] CreateTrigger(create_trigger_task) => { - handle_create_trigger_task(self, create_trigger_task, request.query_context) - .await + handle_create_trigger_task(self, create_trigger_task, query_context).await } #[cfg(feature = "enterprise")] DropTrigger(drop_trigger_task) => { - handle_drop_trigger_task(self, drop_trigger_task, request.query_context).await + handle_drop_trigger_task(self, drop_trigger_task, query_context).await } } } @@ -889,6 +943,7 @@ impl DdlManager { async fn handle_truncate_table_task( ddl_manager: &DdlManager, truncate_table_task: TruncateTableTask, + event_context: EventContext, ) -> Result { let table_id = truncate_table_task.table_id; let table_metadata_manager = &ddl_manager.table_metadata_manager(); @@ -913,7 +968,7 @@ async fn handle_truncate_table_task( ); let (id, _) = ddl_manager - .submit_truncate_table_task(truncate_table_task, table_info_value) + .submit_truncate_table_task(truncate_table_task, table_info_value, event_context) .await?; info!("Table: {table_id} is truncated via procedure_id {id:?}"); @@ -928,6 +983,7 @@ async fn handle_alter_table_task( ddl_manager: &DdlManager, alter_table_task: AlterTableTask, ddl_options: DdlOptions, + event_context: EventContext, ) -> Result { let table_ref = alter_table_task.table_ref(); @@ -960,7 +1016,7 @@ async fn handle_alter_table_task( ); let (id, _) = ddl_manager - .submit_alter_table_task(table_id, alter_table_task, ddl_options) + .submit_alter_table_task(table_id, alter_table_task, event_context, ddl_options) .await?; info!("Table: {table_id} is altered via procedure_id {id:?}"); @@ -974,9 +1030,12 @@ async fn handle_alter_table_task( async fn handle_drop_table_task( ddl_manager: &DdlManager, drop_table_task: DropTableTask, + event_context: EventContext, ) -> Result { let table_id = drop_table_task.table_id; - let (id, _) = ddl_manager.submit_drop_table_task(drop_table_task).await?; + let (id, _) = ddl_manager + .submit_drop_table_task(drop_table_task, event_context) + .await?; info!("Table: {table_id} is dropped via procedure_id {id:?}"); @@ -989,10 +1048,11 @@ async fn handle_drop_table_task( async fn handle_undrop_table_task( ddl_manager: &DdlManager, undrop_table_task: UndropTableTask, + event_context: EventContext, ) -> Result { let table_id = undrop_table_task.table_id; let (id, _) = ddl_manager - .submit_undrop_table_task(undrop_table_task) + .submit_undrop_table_task(undrop_table_task, event_context) .await?; info!("Table: {table_id} is undropped via procedure_id {id:?}"); @@ -1006,9 +1066,10 @@ async fn handle_undrop_table_task( async fn handle_purge_dropped_table_task( ddl_manager: &DdlManager, purge_dropped_table_task: PurgeDroppedTableTask, + event_context: EventContext, ) -> Result { let (id, _) = ddl_manager - .submit_purge_dropped_table_task(purge_dropped_table_task) + .submit_purge_dropped_table_task(purge_dropped_table_task, event_context) .await?; info!("Dropped table is purged via procedure_id {id:?}"); @@ -1023,9 +1084,10 @@ async fn handle_create_table_task( ddl_manager: &DdlManager, create_table_task: CreateTableTask, query_context: QueryContext, + event_context: EventContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_table_task(create_table_task, query_context) + .submit_create_table_task(create_table_task, query_context, event_context) .await?; let procedure_id = id.to_string(); @@ -1048,6 +1110,7 @@ async fn handle_create_table_task( async fn handle_create_logical_table_tasks( ddl_manager: &DdlManager, create_table_tasks: Vec, + event_context: EventContext, ) -> Result { ensure!( !create_table_tasks.is_empty(), @@ -1063,7 +1126,7 @@ async fn handle_create_logical_table_tasks( let num_logical_tables = create_table_tasks.len(); let (id, output) = ddl_manager - .submit_create_logical_table_tasks(create_table_tasks, physical_table_id) + .submit_create_logical_table_tasks(create_table_tasks, physical_table_id, event_context) .await?; info!( @@ -1092,11 +1155,12 @@ async fn handle_create_logical_table_tasks( async fn handle_create_database_task( ddl_manager: &DdlManager, create_database_task: CreateDatabaseTask, + event_context: EventContext, ) -> Result { let catalog = create_database_task.catalog.clone(); let schema = create_database_task.schema.clone(); let (id, _) = ddl_manager - .submit_create_database(create_database_task) + .submit_create_database(create_database_task, event_context) .await?; let procedure_id = id.to_string(); @@ -1114,9 +1178,10 @@ async fn handle_create_database_task( async fn handle_drop_database_task( ddl_manager: &DdlManager, drop_database_task: DropDatabaseTask, + event_context: EventContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_database(drop_database_task.clone()) + .submit_drop_database(drop_database_task.clone(), event_context) .await?; let procedure_id = id.to_string(); @@ -1134,9 +1199,10 @@ async fn handle_drop_database_task( async fn handle_alter_database_task( ddl_manager: &DdlManager, alter_database_task: AlterDatabaseTask, + event_context: EventContext, ) -> Result { let (id, _) = ddl_manager - .submit_alter_database(alter_database_task.clone()) + .submit_alter_database(alter_database_task.clone(), event_context) .await?; let procedure_id = id.to_string(); @@ -1155,9 +1221,10 @@ async fn handle_alter_database_task( async fn handle_drop_flow_task( ddl_manager: &DdlManager, drop_flow_task: DropFlowTask, + event_context: EventContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_flow_task(drop_flow_task.clone()) + .submit_drop_flow_task(drop_flow_task.clone(), event_context) .await?; let procedure_id = id.to_string(); @@ -1199,9 +1266,10 @@ async fn handle_drop_trigger_task( async fn handle_drop_view_task( ddl_manager: &DdlManager, drop_view_task: DropViewTask, + event_context: EventContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_view_task(drop_view_task.clone()) + .submit_drop_view_task(drop_view_task.clone(), event_context) .await?; let procedure_id = id.to_string(); @@ -1221,9 +1289,10 @@ async fn handle_create_flow_task( ddl_manager: &DdlManager, create_flow_task: CreateFlowTask, query_context: QueryContext, + event_context: EventContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_flow_task(create_flow_task.clone(), query_context) + .submit_create_flow_task(create_flow_task.clone(), query_context, event_context) .await?; let procedure_id = id.to_string(); @@ -1280,6 +1349,7 @@ async fn handle_create_trigger_task( async fn handle_alter_logical_table_tasks( ddl_manager: &DdlManager, alter_table_tasks: Vec, + event_context: EventContext, ) -> Result { ensure!( !alter_table_tasks.is_empty(), @@ -1299,7 +1369,7 @@ async fn handle_alter_logical_table_tasks( let num_logical_tables = alter_table_tasks.len(); let (id, _) = ddl_manager - .submit_alter_logical_table_tasks(alter_table_tasks, physical_table_id) + .submit_alter_logical_table_tasks(alter_table_tasks, physical_table_id, event_context) .await?; info!( @@ -1318,9 +1388,10 @@ async fn handle_alter_logical_table_tasks( async fn handle_create_view_task( ddl_manager: &DdlManager, create_view_task: CreateViewTask, + event_context: EventContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_view_task(create_view_task) + .submit_create_view_task(create_view_task, event_context) .await?; let procedure_id = id.to_string(); @@ -1398,7 +1469,7 @@ mod tests { use crate::procedure_executor::ExecutorContext; use crate::region_keeper::MemoryRegionKeeper; use crate::region_registry::LeaderRegionRegistry; - use crate::rpc::ddl::{CreatorGrantIntent, UndropTableTask}; + use crate::rpc::ddl::{CreatorGrantIntent, EventContext, UndropTableTask}; #[cfg(not(feature = "enterprise"))] use crate::rpc::ddl::{DdlTask, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest}; use crate::sequence::SequenceBuilder; @@ -1435,6 +1506,7 @@ mod tests { _source: RepartitionSource, _to_exprs: Vec, _timeout: Option, + _event_context: EventContext, ) -> std::result::Result { unimplemented!() } @@ -1608,7 +1680,7 @@ mod tests { let ddl_manager = build_soft_drop_test_ddl_manager(); let err = ddl_manager - .submit_undrop_table_task(UndropTableTask { table_id: 1024 }) + .submit_undrop_table_task(UndropTableTask { table_id: 1024 }, EventContext::default()) .await .unwrap_err(); @@ -1622,13 +1694,16 @@ mod tests { let ddl_manager = build_soft_drop_test_ddl_manager(); let err = ddl_manager - .submit_undrop_table_task(UndropTableTask { table_id: 1024 }) + .submit_undrop_table_task(UndropTableTask { table_id: 1024 }, EventContext::default()) .await .unwrap_err(); assert!(matches!(err, crate::error::Error::Unsupported { .. })); let err = ddl_manager - .submit_purge_dropped_table_task(PurgeDroppedTableTask { table_id: 1024 }) + .submit_purge_dropped_table_task( + PurgeDroppedTableTask { table_id: 1024 }, + EventContext::default(), + ) .await .unwrap_err(); assert!(matches!(err, crate::error::Error::Unsupported { .. })); diff --git a/src/common/meta/src/rpc/ddl.rs b/src/common/meta/src/rpc/ddl.rs index 0b61d2b2a5e..9cb75aa29c0 100644 --- a/src/common/meta/src/rpc/ddl.rs +++ b/src/common/meta/src/rpc/ddl.rs @@ -44,11 +44,13 @@ use base64::Engine as _; use base64::engine::general_purpose; use common_catalog::{format_full_flow_name, format_full_table_name}; use common_error::ext::BoxedError; +use common_session::channel_protocol; use common_time::{DatabaseTimeToLive, Timestamp}; use prost::Message; use serde::{Deserialize, Serialize}; use serde_with::{DefaultOnNull, serde_as}; use snafu::{OptionExt, ResultExt}; +use strum::{AsRefStr, EnumString}; use table::metadata::{TableId, TableInfo}; use table::requests::validate_database_option; use table::table_name::TableName; @@ -66,7 +68,8 @@ use crate::key::table_name::{TableNameKey, TableNameManager}; /// Reserved query-context extension key for the frontend peer address that submitted a DDL request. pub const ORIGIN_FRONTEND_ADDR_EXTENSION_KEY: &str = "__greptime_origin_frontend.addr"; - +/// Reserved query-context extension key for the trigger reason supplied by frontend. +pub const TRIGGER_REASON_EXTENSION_KEY: &str = "__greptime_event.trigger_reason"; /// Reserved query-context extension key for the authenticated database creator. pub const CREATE_DATABASE_CREATOR_EXTENSION_KEY: &str = "__greptime_create_database.creator"; /// Internal gRPC metadata key for the authenticated database creator. @@ -1665,6 +1668,86 @@ pub struct QueryContext { pub sst_min_sequences: HashMap, } +impl QueryContext { + /// Returns the protocol name represented by the wire channel value. + pub fn channel_protocol(&self) -> Option<&'static str> { + channel_protocol(self.channel) + } +} + +/// The stable context recorded for a procedure event. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct EventContext { + pub reason: TriggerReason, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub protocol: Option, + #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")] + pub extensions: serde_json::Map, +} + +impl EventContext { + /// Creates an event context with no additional extensions. + pub fn new(reason: TriggerReason) -> Self { + Self { + reason, + protocol: None, + extensions: Default::default(), + } + } + + /// Adds the protocol that originated the operation. + pub fn with_protocol(mut self, protocol: impl Into) -> Self { + self.protocol = Some(protocol.into()); + self + } + + /// Builds an event context from frontend query context metadata. + pub fn from_query_context(query_context: &QueryContext) -> Self { + let reason = query_context + .extensions + .get(TRIGGER_REASON_EXTENSION_KEY) + .map(|reason| TriggerReason::from_extension(reason)) + .unwrap_or_default(); + let context = Self::new(reason); + if let Some(protocol) = query_context.channel_protocol() { + context.with_protocol(protocol) + } else { + context + } + } +} + +impl Default for EventContext { + fn default() -> Self { + Self::new(TriggerReason::default()) + } +} + +/// The stable classification of a procedure trigger. +#[derive( + Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, AsRefStr, EnumString, +)] +#[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +pub enum TriggerReason { + Manual, + AutoCreate, + AutoAlter, + AutoRepartition, + AutoRebalance, + RegionFailover, + ScheduledGc, + #[default] + #[serde(other)] + Unknown, +} + +impl TriggerReason { + pub fn from_extension(value: &str) -> Self { + value.parse().unwrap_or_default() + } +} + impl QueryContext { /// Get the current catalog pub fn current_catalog(&self) -> &str { @@ -2237,4 +2320,50 @@ mod tests { assert_eq!(pb_roundtrip.channel, pb.channel); assert_eq!(pb_roundtrip.snapshot_seqs, pb.snapshot_seqs); } + + #[test] + fn test_trigger_reason_deserializes_unknown_value() { + let reason: TriggerReason = serde_json::from_str("\"future_reason\"").unwrap(); + assert_eq!(TriggerReason::Unknown, reason); + } + + #[test] + fn test_event_context_serialization() { + let context = EventContext::new(TriggerReason::Manual).with_protocol("mysql"); + + assert_eq!( + serde_json::json!({ + "reason": "manual", + "protocol": "mysql", + }), + serde_json::to_value(context).unwrap() + ); + + assert_eq!( + serde_json::json!({ "reason": "manual" }), + serde_json::to_value(EventContext::new(TriggerReason::Manual)).unwrap() + ); + } + + #[test] + fn test_event_context_from_query_context_preserves_extensions() { + let mut query_context = QueryContext::default(); + query_context.extensions.insert( + TRIGGER_REASON_EXTENSION_KEY.to_string(), + TriggerReason::AutoCreate.as_ref().to_string(), + ); + query_context.channel = 4; + + let event_context = EventContext::from_query_context(&query_context); + + assert_eq!(event_context.reason, TriggerReason::AutoCreate); + assert_eq!(event_context.protocol.as_deref(), Some("prometheus")); + assert_eq!( + query_context + .extensions + .get(TRIGGER_REASON_EXTENSION_KEY) + .map(String::as_str), + Some("auto_create") + ); + } } diff --git a/src/common/procedure/src/lib.rs b/src/common/procedure/src/lib.rs index a1525ad242c..1d7a200ca58 100644 --- a/src/common/procedure/src/lib.rs +++ b/src/common/procedure/src/lib.rs @@ -30,8 +30,8 @@ pub use crate::error::{Error, Result}; pub use crate::event::ProcedureEvent; pub use crate::procedure::{ BoxedProcedure, BoxedProcedureLoader, ChildSubmissionOutcome, Context, ContextProvider, - ContextProviderRef, EventContext, EventTrigger, LockKey, Output, ParseIdError, PoisonKey, - PoisonKeys, Procedure, ProcedureId, ProcedureInfo, ProcedureManager, ProcedureManagerRef, - ProcedureState, ProcedureWithId, RetryPhase, Status, StringKey, + ContextProviderRef, EventRuntimeContext, EventTrigger, LockKey, Output, ParseIdError, + PoisonKey, PoisonKeys, Procedure, ProcedureId, ProcedureInfo, ProcedureManager, + ProcedureManagerRef, ProcedureState, ProcedureWithId, RetryPhase, Status, StringKey, }; pub use crate::watcher::Watcher; diff --git a/src/common/procedure/src/local.rs b/src/common/procedure/src/local.rs index fbcbc131b8d..59eaee2d789 100644 --- a/src/common/procedure/src/local.rs +++ b/src/common/procedure/src/local.rs @@ -1053,8 +1053,8 @@ mod tests { use crate::store::state_store::ObjectStateStore; use crate::test_util::InMemoryPoisonStore; use crate::{ - ChildSubmissionOutcome, Context, EventContext, EventTrigger, Procedure, ProcedureEvent, - Status, + ChildSubmissionOutcome, Context, EventRuntimeContext, EventTrigger, Procedure, + ProcedureEvent, Status, }; fn new_test_manager_context() -> ManagerContext { @@ -1296,7 +1296,7 @@ mod tests { self.poison_keys.clone() } - fn event(&self, _ctx: &EventContext<'_>) -> Option> { + fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { Some(Box::new(TestProcedureEvent)) } } @@ -1354,7 +1354,7 @@ mod tests { LockKey::default() } - fn event(&self, _ctx: &EventContext<'_>) -> Option> { + fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { Some(Box::new(TestProcedureEvent)) } } @@ -1397,7 +1397,7 @@ mod tests { LockKey::default() } - fn event(&self, _ctx: &EventContext<'_>) -> Option> { + fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { Some(Box::new(TestProcedureEvent)) } } @@ -1424,7 +1424,7 @@ mod tests { LockKey::default() } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { *self.captured_filter.lock().unwrap() = Some(ctx.event_type_filter.clone()); Some(Box::new(TestProcedureEvent)) } diff --git a/src/common/procedure/src/local/runner.rs b/src/common/procedure/src/local/runner.rs index eea95f711b7..a460b7376bf 100644 --- a/src/common/procedure/src/local/runner.rs +++ b/src/common/procedure/src/local/runner.rs @@ -34,8 +34,8 @@ use crate::procedure::{Output, StringKey}; use crate::rwlock::OwnedKeyRwLockGuard; use crate::store::{ProcedureMessage, ProcedureStore}; use crate::{ - BoxedProcedure, ChildSubmissionOutcome, Context, Error, EventContext, EventTrigger, Procedure, - ProcedureId, ProcedureState, ProcedureWithId, RetryPhase, Status, + BoxedProcedure, ChildSubmissionOutcome, Context, Error, EventRuntimeContext, EventTrigger, + Procedure, ProcedureId, ProcedureState, ProcedureWithId, RetryPhase, Status, }; /// A guard to cleanup procedure state. @@ -776,7 +776,7 @@ impl Runner { pub(crate) fn build_event(&self, trigger: EventTrigger) -> Option { let recorder = self.event_recorder.as_ref()?; let state = self.meta.state(); - let context = EventContext { + let context = EventRuntimeContext { procedure_id: self.meta.id, lifecycle_state: &state, trigger: trigger.clone(), diff --git a/src/common/procedure/src/procedure.rs b/src/common/procedure/src/procedure.rs index c37e1b58bcc..a8af03e0da0 100644 --- a/src/common/procedure/src/procedure.rs +++ b/src/common/procedure/src/procedure.rs @@ -237,13 +237,13 @@ pub trait Procedure: Send { /// [`Event::extra_schema`] values. The event recorder batches events by type /// and rejects incompatible schemas; use a distinct event type for a /// different schema. - fn event(&self, _ctx: &EventContext<'_>) -> Option> { + fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { None } } /// Framework-owned context supplied when a procedure builds a lifecycle event. -pub struct EventContext<'a> { +pub struct EventRuntimeContext<'a> { /// Id of the procedure associated with the event. pub procedure_id: ProcedureId, /// Current framework state of the procedure. @@ -336,7 +336,7 @@ impl Procedure for Box { (**self).poison_keys() } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { (**self).event(ctx) } } @@ -721,7 +721,7 @@ mod tests { #[test] fn test_default_procedure_event_hook() { let state = ProcedureState::Running; - let context = EventContext { + let context = EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Succeeded, diff --git a/src/common/session/src/lib.rs b/src/common/session/src/lib.rs index f09baf92866..418425df15e 100644 --- a/src/common/session/src/lib.rs +++ b/src/common/session/src/lib.rs @@ -15,6 +15,27 @@ use serde::{Deserialize, Serialize}; use strum::{AsRefStr, Display, EnumString}; +/// Returns the protocol name represented by a wire query channel value. +pub fn channel_protocol(channel: u8) -> Option<&'static str> { + match channel { + 1 => Some("mysql"), + 2 => Some("postgres"), + 3 => Some("httpsql"), + 4 => Some("prometheus"), + 5 => Some("otlp"), + 6 => Some("grpc"), + 7 => Some("influx"), + 8 => Some("opentsdb"), + 9 => Some("loki"), + 10 => Some("elasticsearch"), + 11 => Some("jaeger"), + 12 => Some("log"), + 13 => Some("promql"), + 14 => Some("splunk"), + _ => None, + } +} + /// Defines the read preference for frontend route operations, /// determining whether to read from the region leader or follower. #[derive( @@ -31,7 +52,33 @@ pub enum ReadPreference { mod tests { use std::str::FromStr; - use crate::ReadPreference; + use crate::{ReadPreference, channel_protocol}; + + #[test] + fn test_channel_protocol() { + let expected = [ + (1, "mysql"), + (2, "postgres"), + (3, "httpsql"), + (4, "prometheus"), + (5, "otlp"), + (6, "grpc"), + (7, "influx"), + (8, "opentsdb"), + (9, "loki"), + (10, "elasticsearch"), + (11, "jaeger"), + (12, "log"), + (13, "promql"), + (14, "splunk"), + ]; + + for (channel, protocol) in expected { + assert_eq!(Some(protocol), channel_protocol(channel)); + } + assert_eq!(None, channel_protocol(0)); + assert_eq!(None, channel_protocol(15)); + } #[test] fn test_read_preference() { diff --git a/src/flow/src/adapter/util.rs b/src/flow/src/adapter/util.rs index 80b3f5ae199..b9e73e296c3 100644 --- a/src/flow/src/adapter/util.rs +++ b/src/flow/src/adapter/util.rs @@ -21,6 +21,7 @@ use api::v1::column_def::options_from_column_schema; use api::v1::{ColumnDataType, ColumnDataTypeExtension, CreateTableExpr, SemanticType}; use common_error::ext::BoxedError; use common_meta::key::table_info::TableInfoValue; +use common_meta::rpc::ddl::TriggerReason; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema}; use itertools::Itertools; @@ -131,7 +132,7 @@ impl StreamingEngine { .build(), ); stmt_exec - .create_table_inner(&mut create_table, None, ctx) + .create_table_inner(&mut create_table, None, ctx, TriggerReason::AutoCreate) .await .map_err(BoxedError::new) .context(ExternalSnafu)?; diff --git a/src/frontend/src/instance/dashboard.rs b/src/frontend/src/instance/dashboard.rs index 068e3e102c1..33a9083cc8e 100644 --- a/src/frontend/src/instance/dashboard.rs +++ b/src/frontend/src/instance/dashboard.rs @@ -24,6 +24,7 @@ use async_trait::async_trait; use auth::{DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, PermissionReq}; use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, default_engine}; use common_error::ext::BoxedError; +use common_meta::rpc::ddl::TriggerReason; use common_query::OutputData; use common_recordbatch::util as record_util; use common_telemetry::info; @@ -165,7 +166,12 @@ impl Instance { }; self.statement_executor - .create_table_inner(&mut create_table_expr, None, ctx.clone()) + .create_table_inner( + &mut create_table_expr, + None, + ctx.clone(), + TriggerReason::AutoCreate, + ) .await .map_err(BoxedError::new) .context(ExecuteQuerySnafu)?; diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index 3caac10b841..acb3a79e049 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -31,6 +31,7 @@ use auth::{ }; use common_error::ext::BoxedError; use common_grpc::flight::do_put::DoPutResponse; +use common_meta::rpc::ddl::TriggerReason; use common_query::Output; use common_query::logical_plan::add_insert_to_logical_plan; use common_telemetry::tracing::{self}; @@ -181,7 +182,12 @@ impl GrpcQueryHandler for Instance { operator::expr_helper::validate_create_expr(&expr)?; let _ = self .statement_executor - .create_table_inner(&mut expr, None, ctx.clone()) + .create_table_inner( + &mut expr, + None, + ctx.clone(), + TriggerReason::Manual, + ) .await?; Output::new_with_affected_rows(0) } @@ -194,7 +200,7 @@ impl GrpcQueryHandler for Instance { } DdlExpr::AlterTable(expr) => { self.statement_executor - .alter_table_inner(expr, ctx.clone()) + .alter_table_inner(expr, ctx.clone(), TriggerReason::Manual) .await? } DdlExpr::CreateDatabase(expr) => { diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs index 41db13661a8..f6ce60b2d11 100644 --- a/src/frontend/src/instance/otlp/trace_ingest.rs +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -24,6 +24,7 @@ use api::v1::{ use client::Output; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; +use common_meta::rpc::ddl::TriggerReason; use common_telemetry::warn; use datatypes::prelude::ConcreteDataType; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; @@ -1681,7 +1682,7 @@ impl Instance { if let Err(err) = self .statement_executor - .alter_table_inner(alter_expr, ctx.clone()) + .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter) .await { let table = self diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index 3f3d2f88121..7c86f3316a5 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -30,6 +30,7 @@ use auth::{ use client::OutputData; use common_catalog::{format_full_table_name, parse_optional_catalog_and_schema_from_db_string}; use common_error::ext::BoxedError; +use common_meta::rpc::ddl::TriggerReason; use common_query::Output; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; use common_recordbatch::RecordBatches; @@ -279,7 +280,7 @@ impl PendingRowsSchemaAlterer for Instance { AutoCreateTableType::Logical(_) => { // Use the batch API for logical tables. self.statement_executor - .create_logical_tables(&create_exprs, ctx) + .create_logical_tables(&create_exprs, ctx, TriggerReason::AutoCreate) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu)?; @@ -290,7 +291,7 @@ impl PendingRowsSchemaAlterer for Instance { expr.table_options .insert(SST_FORMAT_KEY.to_string(), "flat".to_string()); self.statement_executor - .create_table_inner(&mut expr, None, ctx.clone()) + .create_table_inner(&mut expr, None, ctx.clone(), TriggerReason::AutoCreate) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu)?; @@ -357,7 +358,7 @@ impl PendingRowsSchemaAlterer for Instance { } self.statement_executor - .alter_logical_tables(alter_exprs, ctx) + .alter_logical_tables(alter_exprs, ctx, TriggerReason::AutoAlter) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu)?; @@ -609,7 +610,7 @@ impl Instance { fill_metric_physical_table_options(&mut create_table_expr.table_options); self.statement_executor - .create_table_inner(&mut create_table_expr, None, ctx) + .create_table_inner(&mut create_table_expr, None, ctx, TriggerReason::AutoCreate) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu)?; diff --git a/src/meta-srv/Cargo.toml b/src/meta-srv/Cargo.toml index 7e8347e83bf..e6585505db4 100644 --- a/src/meta-srv/Cargo.toml +++ b/src/meta-srv/Cargo.toml @@ -100,6 +100,5 @@ common-event-recorder = { workspace = true, features = ["testing"] } common-meta = { workspace = true, features = ["testing"] } common-procedure-test.workspace = true common-wal = { workspace = true, features = ["testing"] } -session.workspace = true tracing = "0.1" tracing-subscriber.workspace = true diff --git a/src/meta-srv/src/event/gc.rs b/src/meta-srv/src/event/gc.rs index dd417ca77fb..c467bc665c6 100644 --- a/src/meta-srv/src/event/gc.rs +++ b/src/meta-srv/src/event/gc.rs @@ -21,9 +21,10 @@ use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - GC_REPORT_COLUMN, REGION_ID_COLUMN, REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, column_schemas, - nullable_json, + EVENT_CONTEXT_COLUMN, GC_REPORT_COLUMN, REGION_ID_COLUMN, REGION_NUMBER_COLUMN, + TABLE_ID_COLUMN, column_schemas, nullable_json, }; +use common_meta::rpc::ddl::EventContext; use serde::Serialize; use snafu::ResultExt; use store_api::storage::{GcReport, IndexVersion, RegionId}; @@ -66,6 +67,7 @@ struct BatchGcRegionRow { #[derive(Debug)] pub(crate) struct BatchGcEvent { payload: Option, + event_context: Option, // None emits a procedure-level lifecycle row with null Region dimensions. regions: Option>, } @@ -75,6 +77,7 @@ impl BatchGcEvent { regions: &[RegionId], full_file_listing: bool, timeout: Duration, + event_context: Option, ) -> Self { Self { payload: Some(BatchGcPayload { @@ -83,6 +86,7 @@ impl BatchGcEvent { full_file_listing, timeout, }), + event_context, regions: None, } } @@ -107,6 +111,7 @@ impl BatchGcEvent { (!regions.is_empty()).then_some(Self { payload: None, + event_context: None, regions: Some(regions), }) } @@ -131,8 +136,22 @@ impl Event for BatchGcEvent { } fn extra_rows(&self) -> Result> { + let event_context = self + .event_context + .as_ref() + .map(serde_json::to_value) + .transpose() + .context(SerializeEventSnafu)?; let Some(regions) = &self.regions else { - return Ok(vec![null_row()]); + return Ok(vec![Row { + values: vec![ + Value { value_data: None }, + Value { value_data: None }, + Value { value_data: None }, + Value { value_data: None }, + nullable_json(event_context.as_ref()), + ], + }]); }; regions @@ -145,6 +164,7 @@ impl Event for BatchGcEvent { ValueData::U32Value(region.region_id.table_id()).into(), ValueData::U32Value(region.region_id.region_number()).into(), nullable_json(Some(&report)), + nullable_json(None), ], }) }) @@ -162,9 +182,11 @@ fn schema() -> Vec { &TABLE_ID_COLUMN, ®ION_NUMBER_COLUMN, &GC_REPORT_COLUMN, + &EVENT_CONTEXT_COLUMN, ]) } +#[cfg(test)] fn null_row() -> Row { Row { values: (0..schema().len()) @@ -221,7 +243,7 @@ mod tests { use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::sequence::SequenceBuilder; use common_procedure::{ - EventContext, EventTrigger, Procedure, ProcedureEvent, ProcedureId, ProcedureState, + EventRuntimeContext, EventTrigger, Procedure, ProcedureEvent, ProcedureId, ProcedureState, RetryPhase, }; use store_api::storage::FileId; @@ -234,8 +256,12 @@ mod tests { fn test_batch_gc_lifecycle_event_contract() { let first = RegionId::new(1024, 1); let second = RegionId::new(1024, 2); - let event = - BatchGcEvent::with_config(&[second, first, second], true, Duration::from_secs(10)); + let event = BatchGcEvent::with_config( + &[second, first, second], + true, + Duration::from_secs(10), + None, + ); assert_event_contract(&event, BATCH_GC_EVENT_TYPE, &schema(), &[null_row()]); assert_eq!( @@ -330,7 +356,7 @@ mod tests { fn test_batch_gc_event_filter() { let procedure = batch_gc_procedure(); let running = ProcedureState::Running; - let event_context = |trigger, lifecycle_state, event_type_filter| EventContext { + let runtime_context = |trigger, lifecycle_state, event_type_filter| EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, @@ -339,7 +365,7 @@ mod tests { assert!( procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Submitted, &running, EventTypeFilter::All, @@ -356,7 +382,7 @@ mod tests { }; assert!( procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Succeeded, &done, EventTypeFilter::All, @@ -366,7 +392,7 @@ mod tests { assert!( procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Recovered, &running, EventTypeFilter::All, @@ -375,7 +401,7 @@ mod tests { ); let retrying = procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Retrying { phase: RetryPhase::Execute, attempt: 1, @@ -397,7 +423,7 @@ mod tests { assert!( procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Recovered, &running, EventTypeFilter::Only(HashSet::from(["another_event".to_string()])), @@ -406,7 +432,7 @@ mod tests { ); assert!( procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Recovered, &running, EventTypeFilter::Only(HashSet::new()), @@ -417,7 +443,7 @@ mod tests { let missing = ProcedureState::Done { output: None }; assert!( procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Succeeded, &missing, EventTypeFilter::All, @@ -429,7 +455,7 @@ mod tests { }; assert!( procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Succeeded, &wrong, EventTypeFilter::All, @@ -442,7 +468,7 @@ mod tests { fn test_batch_gc_terminal_event_with_report() { let mut procedure = batch_gc_procedure(); let running = ProcedureState::Running; - let event_context = |trigger| EventContext { + let runtime_context = |trigger| EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state: &running, trigger, @@ -455,7 +481,7 @@ mod tests { procedure.set_gc_report_for_test(report); let retrying = procedure - .event(&event_context(EventTrigger::Retrying { + .event(&runtime_context(EventTrigger::Retrying { phase: RetryPhase::Execute, attempt: 2, })) @@ -472,7 +498,7 @@ mod tests { assert_eq!(retrying.extra_rows().unwrap(), vec![null_row()]); for trigger in [EventTrigger::Failed, EventTrigger::Poisoned] { - let event = procedure.event(&event_context(trigger)).unwrap(); + let event = procedure.event(&runtime_context(trigger)).unwrap(); assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); assert_eq!( event.extra_rows().unwrap(), @@ -487,7 +513,7 @@ mod tests { let procedure = batch_gc_procedure(); let failed_without_report = procedure - .event(&event_context(EventTrigger::Failed)) + .event(&runtime_context(EventTrigger::Failed)) .unwrap(); assert_eq!( failed_without_report.json_payload().unwrap(), @@ -563,6 +589,7 @@ mod tests { ValueData::U32Value(region_id.table_id()).into(), ValueData::U32Value(region_id.region_number()).into(), nullable_json(report.as_ref()), + nullable_json(None), ], } } @@ -580,6 +607,7 @@ mod tests { true, Duration::from_secs(10), HashMap::new(), + EventContext::default(), ) } } diff --git a/src/meta-srv/src/event/region_migration.rs b/src/meta-srv/src/event/region_migration.rs index f7b311c13ba..65b89df2515 100644 --- a/src/meta-srv/src/event/region_migration.rs +++ b/src/meta-srv/src/event/region_migration.rs @@ -20,9 +20,10 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - REGION_ID_COLUMN, REGION_MIGRATION_DST_NODE_ID_COLUMN, REGION_MIGRATION_DST_PEER_ADDR_COLUMN, - REGION_MIGRATION_SRC_NODE_ID_COLUMN, REGION_MIGRATION_SRC_PEER_ADDR_COLUMN, - REGION_MIGRATION_TRIGGER_REASON_COLUMN, REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, column_schemas, + EVENT_CONTEXT_COLUMN, REGION_ID_COLUMN, REGION_MIGRATION_DST_NODE_ID_COLUMN, + REGION_MIGRATION_DST_PEER_ADDR_COLUMN, REGION_MIGRATION_SRC_NODE_ID_COLUMN, + REGION_MIGRATION_SRC_PEER_ADDR_COLUMN, REGION_MIGRATION_TRIGGER_REASON_COLUMN, + REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, column_schemas, nullable_json, }; use serde::Serialize; use snafu::ResultExt; @@ -49,6 +50,7 @@ pub(crate) struct RegionMigrationEvent { dst_peer_addr: String, // The timeout of the region migration. timeout: Duration, + event_context: Option, } #[derive(Debug, Serialize)] @@ -58,7 +60,7 @@ struct Payload { } impl RegionMigrationEvent { - pub fn from_persistent_ctx(ctx: &PersistentContext) -> Self { + pub fn from_persistent_ctx(ctx: &PersistentContext, include_event_context: bool) -> Self { Self { region_ids: ctx.region_ids.clone(), trigger_reason: ctx.trigger_reason, @@ -67,6 +69,7 @@ impl RegionMigrationEvent { dst_node_id: ctx.to_peer.id, dst_peer_addr: ctx.to_peer.addr.clone(), timeout: ctx.timeout, + event_context: include_event_context.then(|| ctx.event_context.clone()), } } } @@ -86,6 +89,7 @@ impl Event for RegionMigrationEvent { ®ION_MIGRATION_SRC_PEER_ADDR_COLUMN, ®ION_MIGRATION_DST_NODE_ID_COLUMN, ®ION_MIGRATION_DST_PEER_ADDR_COLUMN, + &EVENT_CONTEXT_COLUMN, ]) } @@ -102,6 +106,12 @@ impl Event for RegionMigrationEvent { ValueData::StringValue(self.src_peer_addr.clone()).into(), ValueData::U64Value(self.dst_node_id).into(), ValueData::StringValue(self.dst_peer_addr.clone()).into(), + self.event_context + .as_ref() + .map(serde_json::to_value) + .transpose() + .context(SerializeEventSnafu) + .map(|value| nullable_json(value.as_ref()))?, ], }); } diff --git a/src/meta-srv/src/event/repartition.rs b/src/meta-srv/src/event/repartition.rs index d266b581f82..ad47fbf9d42 100644 --- a/src/meta-srv/src/event/repartition.rs +++ b/src/meta-srv/src/event/repartition.rs @@ -21,11 +21,11 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, PARENT_PROCEDURE_ID_COLUMN, REPARTITION_GROUP_ID_COLUMN, - SCHEMA_NAME_COLUMN, SOURCE_PARTITION_EXPR_COLUMN, SOURCE_REGION_ID_COLUMN, - SOURCE_REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, TABLE_NAME_COLUMN, TARGET_PARTITION_EXPR_COLUMN, - TARGET_REGION_ID_COLUMN, TARGET_REGION_NUMBER_COLUMN, column_schemas, nullable_string, - nullable_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PARENT_PROCEDURE_ID_COLUMN, + REPARTITION_GROUP_ID_COLUMN, SCHEMA_NAME_COLUMN, SOURCE_PARTITION_EXPR_COLUMN, + SOURCE_REGION_ID_COLUMN, SOURCE_REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, TABLE_NAME_COLUMN, + TARGET_PARTITION_EXPR_COLUMN, TARGET_REGION_ID_COLUMN, TARGET_REGION_NUMBER_COLUMN, + column_schemas, nullable_json, nullable_string, nullable_value, }; use serde::Serialize; use snafu::ResultExt; @@ -61,6 +61,7 @@ pub(crate) struct RepartitionEvent { table_name: Option, table_id: Option, payload: Option, + event_context: Option, } impl RepartitionEvent { @@ -90,6 +91,7 @@ impl RepartitionEvent { target_partition_columns: intent.target_partition_columns().map(ToOwned::to_owned), timeout: persistent_ctx.timeout, }), + event_context: Some(persistent_ctx.event_context.clone()), } } @@ -100,10 +102,17 @@ impl RepartitionEvent { table_name: None, table_id: None, payload: None, + event_context: None, } } fn schema() -> Vec { + let mut schema = Self::base_schema(); + schema.push(EVENT_CONTEXT_COLUMN.column_schema()); + schema + } + + fn base_schema() -> Vec { column_schemas([ &CATALOG_NAME_COLUMN, &SCHEMA_NAME_COLUMN, @@ -132,12 +141,19 @@ impl Event for RepartitionEvent { } fn extra_rows(&self) -> Result> { + let event_context = self + .event_context + .as_ref() + .map(serde_json::to_value) + .transpose() + .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.schema_name.as_deref()), nullable_string(self.table_name.as_deref()), nullable_value(self.table_id.map(ValueData::U32Value)), + nullable_json(event_context.as_ref()), ], }]) } @@ -307,7 +323,7 @@ impl RepartitionGroupEvent { } fn schema() -> Vec { - let mut schema = RepartitionEvent::schema(); + let mut schema = RepartitionEvent::base_schema(); schema.extend(column_schemas([ &PARENT_PROCEDURE_ID_COLUMN, &REPARTITION_GROUP_ID_COLUMN, @@ -368,6 +384,7 @@ mod tests { PROCEDURE_TRIGGER_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; + use common_meta::rpc::ddl::EventContext; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; use table::table_name::TableName; use uuid::Uuid; @@ -385,6 +402,7 @@ mod tests { TableName::new("greptime", "public", "repartition_events"), 1024, Some(Duration::from_secs(30)), + EventContext::default(), ) } @@ -420,6 +438,9 @@ mod tests { ValueData::StringValue("public".to_string()).into(), ValueData::StringValue("repartition_events".to_string()).into(), ValueData::U32Value(1024).into(), + jsonb_value(&serde_json::json!({ + "reason": "unknown", + })), ], }], ); @@ -674,6 +695,9 @@ mod tests { ValueData::StringValue("public".to_string()).into(), ValueData::StringValue("repartition_events".to_string()).into(), ValueData::U32Value(1024).into(), + jsonb_value(&serde_json::json!({ + "reason": "unknown", + })), ], }], ); @@ -685,11 +709,12 @@ mod tests { &SCHEMA_NAME_COLUMN, &TABLE_NAME_COLUMN, &TABLE_ID_COLUMN, + &EVENT_CONTEXT_COLUMN, ]) } fn group_schema() -> Vec { - let mut schema = parent_schema(); + let mut schema = RepartitionEvent::base_schema(); schema.extend(column_schemas([ &PARENT_PROCEDURE_ID_COLUMN, &REPARTITION_GROUP_ID_COLUMN, diff --git a/src/meta-srv/src/gc/ctx.rs b/src/meta-srv/src/gc/ctx.rs index 79eea6cca33..41027a65390 100644 --- a/src/meta-srv/src/gc/ctx.rs +++ b/src/meta-srv/src/gc/ctx.rs @@ -29,6 +29,7 @@ use common_meta::key::DroppedTableName; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; +use common_meta::rpc::ddl::EventContext; #[cfg(feature = "enterprise")] use common_meta::rpc::ddl::PurgeDroppedTableTask; use common_procedure::{ProcedureManagerRef, ProcedureWithId, watcher}; @@ -67,6 +68,7 @@ pub(crate) trait SchedulerCtx: Send + Sync { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, + event_context: EventContext, ) -> Result; #[cfg(feature = "enterprise")] @@ -254,12 +256,14 @@ impl SchedulerCtx for DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, + event_context: EventContext, ) -> Result { self.gc_regions_inner( region_ids, full_file_listing, timeout, region_routes_override, + event_context, ) .await } @@ -306,6 +310,7 @@ impl DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, + event_context: EventContext, ) -> Result { debug!( "Sending GC instruction for {} regions (full_file_listing: {})", @@ -321,6 +326,7 @@ impl DefaultGcSchedulerCtx { full_file_listing, timeout, region_routes_override, + event_context, ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); diff --git a/src/meta-srv/src/gc/handler.rs b/src/meta-srv/src/gc/handler.rs index 37a71580fa7..4532e9a8ca3 100644 --- a/src/meta-srv/src/gc/handler.rs +++ b/src/meta-srv/src/gc/handler.rs @@ -18,6 +18,7 @@ use std::time::Instant; use common_catalog::consts::MITO_ENGINE; use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::peer::Peer; +use common_meta::rpc::ddl::EventContext; use common_telemetry::tracing::Instrument as _; use common_telemetry::{debug, error, info, warn}; use futures::StreamExt; @@ -36,7 +37,7 @@ use crate::gc::tracker::RegionGcInfo; use crate::metrics::METRIC_META_GC_CANDIDATE_REGIONS; impl GcScheduler { - pub(crate) async fn trigger_gc(&self) -> Result { + pub(crate) async fn trigger_gc(&self, event_context: EventContext) -> Result { let start_time = Instant::now(); info!("Starting GC cycle"); @@ -107,6 +108,7 @@ impl GcScheduler { datanode_to_candidates, dropped_assignment.force_full_listing, dropped_assignment.region_routes_override, + event_context, ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_dispatch_to_datanodes" @@ -210,6 +212,7 @@ impl GcScheduler { datanode_to_candidates: HashMap>, force_full_listing_by_peer: HashMap>, region_routes_override_by_peer: HashMap, + event_context: EventContext, ) -> GcJobReport { let mut per_datanode_reports = HashMap::new(); let mut failed_datanodes: HashMap<_, Vec<_>> = HashMap::new(); @@ -231,6 +234,7 @@ impl GcScheduler { .get(&peer) .cloned() .unwrap_or_default(); + let event_context = event_context.clone(); async move { ( peer, @@ -240,6 +244,7 @@ impl GcScheduler { candidates, force_full_listing, region_routes_override, + event_context, ) .await, ) @@ -278,6 +283,7 @@ impl GcScheduler { candidates: Vec<(TableId, GcCandidate)>, force_full_listing: HashSet, region_routes_override: Region2Peers, + event_context: EventContext, ) -> Result { info!( "Starting GC for datanode {} with {} candidate regions", @@ -327,6 +333,7 @@ impl GcScheduler { false, self.config.mailbox_timeout, region_routes_override.clone(), + event_context.clone(), ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_call_datanode", @@ -359,6 +366,7 @@ impl GcScheduler { true, self.config.mailbox_timeout, region_routes_override, + event_context, ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_call_datanode", diff --git a/src/meta-srv/src/gc/mock.rs b/src/meta-srv/src/gc/mock.rs index 1d007dc30bc..a941476ab44 100644 --- a/src/meta-srv/src/gc/mock.rs +++ b/src/meta-srv/src/gc/mock.rs @@ -30,6 +30,7 @@ use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; +use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::{Region, RegionRoute}; use common_telemetry::debug; use ordered_float::OrderedFloat; @@ -74,6 +75,7 @@ pub struct MockSchedulerCtx { pub candidates: Arc>>>>, pub get_table_to_region_stats_calls: Arc>, pub gc_regions_calls: Arc>, + pub gc_event_contexts: Arc>>, // Error injection fields for testing pub get_table_to_region_stats_error: Arc>>, pub get_table_route_error: Arc>>, @@ -212,8 +214,10 @@ impl SchedulerCtx for MockSchedulerCtx { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, + event_context: EventContext, ) -> Result { *self.gc_regions_calls.lock().unwrap() += 1; + self.gc_event_contexts.lock().unwrap().push(event_context); // Check per-region error injection first (for any region) for ®ion_id in region_ids { diff --git a/src/meta-srv/src/gc/mock/basic.rs b/src/meta-srv/src/gc/mock/basic.rs index 87278abcbf9..45d64910561 100644 --- a/src/meta-srv/src/gc/mock/basic.rs +++ b/src/meta-srv/src/gc/mock/basic.rs @@ -17,6 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::Instant; use common_meta::peer::Peer; +use common_meta::rpc::ddl::{EventContext, TriggerReason}; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -31,7 +32,12 @@ async fn test_parallel_process_datanodes_empty() { let env = TestEnv::new(); let report = env .scheduler - .parallel_process_datanodes(HashMap::new(), HashMap::new(), HashMap::new()) + .parallel_process_datanodes( + HashMap::new(), + HashMap::new(), + HashMap::new(), + EventContext::default(), + ) .await; match report { @@ -95,7 +101,12 @@ async fn test_parallel_process_datanodes_with_candidates() { )]); let report = scheduler - .parallel_process_datanodes(datanode_to_candidates, HashMap::new(), HashMap::new()) + .parallel_process_datanodes( + datanode_to_candidates, + HashMap::new(), + HashMap::new(), + EventContext::default(), + ) .await; match report { @@ -176,6 +187,10 @@ async fn test_handle_tick() { assert_eq!(*ctx.get_table_to_region_stats_calls.lock().unwrap(), 1); assert_eq!(*ctx.gc_regions_calls.lock().unwrap(), 1); + assert_eq!( + ctx.gc_event_contexts.lock().unwrap().as_slice(), + &[EventContext::new(TriggerReason::ScheduledGc)] + ); let tracker = scheduler.region_gc_tracker.lock().await; assert!( @@ -184,3 +199,56 @@ async fn test_handle_tick() { tracker ); } + +#[tokio::test] +async fn test_handle_manual_gc_without_regions_records_manual_event_context() { + init_default_ut_logging(); + + let table_id = 1; + let region_id = RegionId::new(table_id, 1); + let peer = Peer::new(1, ""); + let candidates = HashMap::from([(table_id, vec![new_candidate(region_id, 1.0)])]); + let file_refs = FileRefsManifest { + manifest_version: HashMap::from([(region_id, 1)]), + ..Default::default() + }; + let ctx = Arc::new( + MockSchedulerCtx { + table_to_region_stats: Arc::new(Mutex::new(Some(HashMap::from([( + table_id, + vec![mock_region_stat( + region_id, + RegionRole::Leader, + TEST_REGION_SIZE_200MB, + 10, + )], + )])))), + gc_reports: Arc::new(Mutex::new(HashMap::from([( + region_id, + GcReport::default(), + )]))), + candidates: Arc::new(Mutex::new(Some(candidates))), + file_refs: Arc::new(Mutex::new(Some(file_refs))), + ..Default::default() + } + .with_table_routes(HashMap::from([( + table_id, + (table_id, vec![(region_id, peer)]), + )])), + ); + let scheduler = GcScheduler { + ctx: ctx.clone(), + runtime_switch_manager: crate::gc::scheduler::new_test_runtime_switch_manager(), + receiver: GcScheduler::channel().1, + config: GcSchedulerOptions::default(), + region_gc_tracker: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + last_tracker_cleanup: Arc::new(tokio::sync::Mutex::new(Instant::now())), + }; + + scheduler.handle_manual_gc(None, None, None).await.unwrap(); + + assert_eq!( + ctx.gc_event_contexts.lock().unwrap().as_slice(), + &[EventContext::new(TriggerReason::Manual)] + ); +} diff --git a/src/meta-srv/src/gc/mock/concurrent.rs b/src/meta-srv/src/gc/mock/concurrent.rs index ddaf8128431..7dca378b71d 100644 --- a/src/meta-srv/src/gc/mock/concurrent.rs +++ b/src/meta-srv/src/gc/mock/concurrent.rs @@ -18,6 +18,7 @@ use std::time::{Duration, Instant}; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; +use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::{Region, RegionRoute}; use common_telemetry::{info, init_default_ut_logging}; use store_api::region_engine::RegionRole; @@ -97,7 +98,12 @@ async fn test_concurrent_table_processing_limits() { )]); let report = scheduler - .parallel_process_datanodes(datanode_to_candidates, HashMap::new(), HashMap::new()) + .parallel_process_datanodes( + datanode_to_candidates, + HashMap::new(), + HashMap::new(), + EventContext::default(), + ) .await; // Should process all datanodes @@ -180,7 +186,12 @@ async fn test_datanode_processes_tables_with_partial_gc_failures() { )]); let report = scheduler - .parallel_process_datanodes(datanode_to_candidates, HashMap::new(), HashMap::new()) + .parallel_process_datanodes( + datanode_to_candidates, + HashMap::new(), + HashMap::new(), + EventContext::default(), + ) .await; // Should have one datanode with mixed results @@ -285,6 +296,7 @@ async fn test_region_gc_concurrency_limit() { candidates.into_iter().map(|c| (table_id, c)).collect(), HashSet::new(), HashMap::new(), + EventContext::default(), ) .await .unwrap(); @@ -397,7 +409,12 @@ async fn test_region_gc_concurrency_with_partial_failures() { )]); let report = scheduler - .parallel_process_datanodes(datanode_to_candidates, HashMap::new(), HashMap::new()) + .parallel_process_datanodes( + datanode_to_candidates, + HashMap::new(), + HashMap::new(), + EventContext::default(), + ) .await; let report = match &report { @@ -536,7 +553,12 @@ async fn test_region_gc_concurrency_with_retryable_errors() { candidates.into_iter().map(|c| (table_id, c)).collect(), )]); let report = scheduler - .parallel_process_datanodes(datanode_to_candidates, HashMap::new(), HashMap::new()) + .parallel_process_datanodes( + datanode_to_candidates, + HashMap::new(), + HashMap::new(), + EventContext::default(), + ) .await; let report = match &report { diff --git a/src/meta-srv/src/gc/mock/full_list.rs b/src/meta-srv/src/gc/mock/full_list.rs index 847c0e73e40..a5d5012c6d3 100644 --- a/src/meta-srv/src/gc/mock/full_list.rs +++ b/src/meta-srv/src/gc/mock/full_list.rs @@ -17,6 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use common_meta::peer::Peer; +use common_meta::rpc::ddl::EventContext; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -82,6 +83,7 @@ async fn test_full_file_listing_first_time_gc() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), + EventContext::default(), ) .await .unwrap(); @@ -155,6 +157,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), + EventContext::default(), ) .await .unwrap(); @@ -181,6 +184,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), + EventContext::default(), ) .await .unwrap(); @@ -256,6 +260,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), + EventContext::default(), ) .await .unwrap(); @@ -279,6 +284,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), + EventContext::default(), ) .await .unwrap(); diff --git a/src/meta-srv/src/gc/mock/misc.rs b/src/meta-srv/src/gc/mock/misc.rs index c8b2b43155e..9b3c35cbe9a 100644 --- a/src/meta-srv/src/gc/mock/misc.rs +++ b/src/meta-srv/src/gc/mock/misc.rs @@ -17,6 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::Instant; use common_meta::peer::Peer; +use common_meta::rpc::ddl::EventContext; use common_telemetry::init_default_ut_logging; use store_api::storage::{FileRefsManifest, GcReport, RegionId}; @@ -71,7 +72,12 @@ async fn test_empty_file_refs_manifest() { )]); let report = scheduler - .parallel_process_datanodes(datanode_to_candidates, HashMap::new(), HashMap::new()) + .parallel_process_datanodes( + datanode_to_candidates, + HashMap::new(), + HashMap::new(), + EventContext::default(), + ) .await; match report { @@ -159,7 +165,12 @@ async fn test_multiple_regions_per_table() { )]); let report = scheduler - .parallel_process_datanodes(datanode_to_candidates, HashMap::new(), HashMap::new()) + .parallel_process_datanodes( + datanode_to_candidates, + HashMap::new(), + HashMap::new(), + EventContext::default(), + ) .await; match report { diff --git a/src/meta-srv/src/gc/procedure.rs b/src/meta-srv/src/gc/procedure.rs index 311a14e4a98..d745e88e2c5 100644 --- a/src/meta-srv/src/gc/procedure.rs +++ b/src/meta-srv/src/gc/procedure.rs @@ -24,10 +24,11 @@ use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::lock_key::{RegionLock, TableLock}; use common_meta::peer::Peer; +use common_meta::rpc::ddl::{EventContext, TriggerReason}; use common_procedure::error::ToJsonSnafu; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, - Procedure, ProcedureState, Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, + LockKey, Procedure, ProcedureState, Result as ProcedureResult, Status, }; use common_telemetry::tracing::Instrument as _; use common_telemetry::tracing_context::TracingContext; @@ -199,6 +200,8 @@ pub struct BatchGcData { /// mailbox timeout duration timeout: Duration, gc_report: Option, + #[serde(default)] + event_context: EventContext, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -216,6 +219,7 @@ pub enum State { impl BatchGcProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::BatchGcProcedure"; + #[allow(clippy::too_many_arguments)] pub fn new( mailbox: MailboxRef, table_metadata_manager: TableMetadataManagerRef, @@ -224,6 +228,7 @@ impl BatchGcProcedure { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, + event_context: EventContext, ) -> Self { Self { mailbox, @@ -239,6 +244,7 @@ impl BatchGcProcedure { related_regions: HashMap::new(), file_refs: FileRefsManifest::default(), gc_report: None, + event_context, }, } } @@ -269,6 +275,7 @@ impl BatchGcProcedure { related_regions: HashMap::new(), file_refs, gc_report: Some(GcReport::default()), + event_context: EventContext::default(), }, } } @@ -1065,16 +1072,26 @@ impl Procedure for BatchGcProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx.event_type_filter.allows(BATCH_GC_EVENT_TYPE) { return None; } let event = match &ctx.trigger { - // Keep normal GC low-noise; only terminal events with meaningful reports are recorded. - EventTrigger::Submitted - | EventTrigger::Recovered - | EventTrigger::ChildSubmitted { .. } => return None, + // Keep scheduled GC low-noise; record submitted manual requests for auditability. + EventTrigger::Submitted => (self.data.event_context.reason == TriggerReason::Manual) + .then(|| { + BatchGcEvent::with_config( + &self.data.regions, + self.data.full_file_listing, + self.data.timeout, + Some(self.data.event_context.clone()), + ) + })?, + EventTrigger::Recovered | EventTrigger::ChildSubmitted { .. } => return None, EventTrigger::Succeeded => { let ProcedureState::Done { output: Some(output), @@ -1089,6 +1106,7 @@ impl Procedure for BatchGcProcedure { &self.data.regions, self.data.full_file_listing, self.data.timeout, + None, ), EventTrigger::Failed | EventTrigger::Poisoned => self .data @@ -1100,6 +1118,7 @@ impl Procedure for BatchGcProcedure { &self.data.regions, self.data.full_file_listing, self.data.timeout, + None, ) }), }; @@ -1192,6 +1211,16 @@ mod tests { assert!(report.need_retry_regions.contains(®ion_id)); } + #[test] + fn test_batch_gc_data_defaults_missing_event_context() { + let mut data = serde_json::to_value(&batch_gc_procedure().data).unwrap(); + data.as_object_mut().unwrap().remove("event_context"); + + let recovered: BatchGcData = serde_json::from_value(data).unwrap(); + + assert_eq!(recovered.event_context, EventContext::default()); + } + #[tokio::test] async fn test_send_gc_instructions_preserves_partial_report() { let first_region = RegionId::new(1024, 1); @@ -1226,6 +1255,7 @@ mod tests { true, Duration::from_secs(10), HashMap::new(), + EventContext::default(), ); procedure.data.region_routes = HashMap::from([ (first_region, (first_peer, vec![])), @@ -1266,6 +1296,7 @@ mod tests { true, Duration::from_secs(10), HashMap::new(), + EventContext::default(), ) } } diff --git a/src/meta-srv/src/gc/scheduler.rs b/src/meta-srv/src/gc/scheduler.rs index c637d152526..92b362ad947 100644 --- a/src/meta-srv/src/gc/scheduler.rs +++ b/src/meta-srv/src/gc/scheduler.rs @@ -18,6 +18,7 @@ use std::time::{Duration, Instant}; use common_meta::DatanodeId; use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; +use common_meta::rpc::ddl::{EventContext, TriggerReason}; use common_telemetry::tracing::Instrument as _; use common_telemetry::{error, info}; use snafu::ResultExt; @@ -216,7 +217,10 @@ impl GcScheduler { .await; } let span = common_telemetry::tracing::info_span!("meta_gc_handle_tick"); - let report = self.trigger_gc().instrument(span).await?; + let report = self + .trigger_gc(EventContext::new(TriggerReason::ScheduledGc)) + .instrument(span) + .await?; // Periodically clean up stale tracker entries self.cleanup_tracker_if_needed().await?; @@ -300,7 +304,9 @@ impl GcScheduler { // No specific regions, use default tick behavior let Some(regions) = region_ids else { - let report = self.trigger_gc().await?; + let report = self + .trigger_gc(EventContext::new(TriggerReason::Manual)) + .await?; info!("Finished manual gc request"); return Ok(report); }; @@ -347,6 +353,7 @@ impl GcScheduler { full_listing, gc_timeout, Region2Peers::new(), + EventContext::new(TriggerReason::Manual), ) .await?; combined_report.merge(report); @@ -355,7 +362,13 @@ impl GcScheduler { if !dropped_regions.is_empty() { let report = self .ctx - .gc_regions(&dropped_regions, true, gc_timeout, dropped_routes_override) + .gc_regions( + &dropped_regions, + true, + gc_timeout, + dropped_routes_override, + EventContext::new(TriggerReason::Manual), + ) .await?; combined_report.merge(report); } @@ -477,6 +490,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, + _event_context: EventContext, ) -> Result { self.gc_regions_calls.fetch_add(1, Ordering::Relaxed); panic!("gc_regions should not be called in maintenance mode") @@ -550,6 +564,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, + _event_context: EventContext, ) -> Result { Ok(GcReport::default()) } @@ -883,6 +898,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, + _event_context: EventContext, ) -> Result { crate::error::UnexpectedSnafu { violated: "mock gc failure".to_string(), diff --git a/src/meta-srv/src/procedure/region_migration.rs b/src/meta-srv/src/procedure/region_migration.rs index 23a48381ebd..25fc7d6a309 100644 --- a/src/meta-srv/src/procedure/region_migration.rs +++ b/src/meta-srv/src/procedure/region_migration.rs @@ -45,11 +45,13 @@ use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; use common_meta::lock_key::{CatalogLock, RegionLock, SchemaLock, TableLock}; use common_meta::peer::Peer; use common_meta::region_keeper::{MemoryRegionKeeperRef, OperatingRegionGuard}; +use common_meta::rpc::ddl::{EventContext, TriggerReason}; use common_procedure::error::{ Error as ProcedureError, FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu, }; use common_procedure::{ - Context as ProcedureContext, EventContext, LockKey, Procedure, Status, StringKey, + Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Status, + StringKey, }; use common_telemetry::{debug, error, info}; use manager::RegionMigrationProcedureGuard; @@ -124,6 +126,8 @@ pub struct PersistentContext { /// The trigger reason of region migration. #[serde(default)] pub(crate) trigger_reason: RegionMigrationTriggerReason, + #[serde(default)] + pub(crate) event_context: EventContext, } impl PersistentContext { @@ -133,7 +137,7 @@ impl PersistentContext { to_peer: Peer, region_ids: Vec, timeout: Duration, - trigger_reason: RegionMigrationTriggerReason, + event_context: EventContext, ) -> Self { #[allow(deprecated)] Self { @@ -144,7 +148,37 @@ impl PersistentContext { to_peer, region_ids, timeout, - trigger_reason, + trigger_reason: RegionMigrationTriggerReason::from_trigger_reason(event_context.reason), + event_context, + } + } + + fn normalize_trigger_reasons(&mut self) { + if self.trigger_reason != RegionMigrationTriggerReason::Unknown { + self.event_context.reason = self.trigger_reason.to_trigger_reason(); + } else if self.event_context.reason != TriggerReason::Unknown { + self.trigger_reason = + RegionMigrationTriggerReason::from_trigger_reason(self.event_context.reason); + } + } +} + +impl RegionMigrationTriggerReason { + fn to_trigger_reason(self) -> TriggerReason { + match self { + Self::Manual => TriggerReason::Manual, + Self::AutoRebalance => TriggerReason::AutoRebalance, + Self::Failover => TriggerReason::RegionFailover, + Self::Unknown => TriggerReason::Unknown, + } + } + + fn from_trigger_reason(reason: TriggerReason) -> Self { + match reason { + TriggerReason::Manual => Self::Manual, + TriggerReason::AutoRebalance => Self::AutoRebalance, + TriggerReason::RegionFailover => Self::Failover, + _ => Self::Unknown, } } } @@ -818,9 +852,10 @@ impl RegionMigrationProcedure { tracker: RegionMigrationProcedureTracker, ) -> ProcedureResult { let RegionMigrationDataOwned { - persistent_ctx, + mut persistent_ctx, state, } = serde_json::from_str(json).context(FromJsonSnafu)?; + persistent_ctx.normalize_trigger_reasons(); let guards = persistent_ctx .region_ids .iter() @@ -964,13 +999,14 @@ impl Procedure for RegionMigrationProcedure { LockKey::new(self.context.persistent_ctx.lock_key()) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { if !ctx.event_type_filter.allows(REGION_MIGRATION_EVENT_TYPE) { return None; } Some(Box::new(RegionMigrationEvent::from_persistent_ctx( &self.context.persistent_ctx, + matches!(ctx.trigger, EventTrigger::Submitted), ))) } } @@ -1027,7 +1063,7 @@ mod tests { let procedure = RegionMigrationProcedure::new(persistent_context, context, vec![]); let serialized = procedure.dump().unwrap(); - let expected = r#"{"persistent_ctx":{"catalog_and_schema":[["greptime","public"]],"from_peer":{"id":1,"addr":""},"to_peer":{"id":2,"addr":""},"region_ids":[4398046511105],"timeout":"10s","trigger_reason":"Unknown"},"state":{"region_migration_state":"RegionMigrationStart"}}"#; + let expected = r#"{"persistent_ctx":{"catalog_and_schema":[["greptime","public"]],"from_peer":{"id":1,"addr":""},"to_peer":{"id":2,"addr":""},"region_ids":[4398046511105],"timeout":"10s","trigger_reason":"Unknown","event_context":{"reason":"unknown"}},"state":{"region_migration_state":"RegionMigrationStart"}}"#; assert_eq!(expected, serialized); } @@ -1051,7 +1087,7 @@ mod tests { for trigger in triggers { let event = procedure - .event(&EventContext { + .event(&EventRuntimeContext { procedure_id: common_procedure::ProcedureId::random(), lifecycle_state: &state, trigger, @@ -1076,12 +1112,60 @@ mod tests { region_ids: vec![RegionId::new(1024, 1)], timeout: Duration::from_secs(10), trigger_reason: RegionMigrationTriggerReason::default(), + event_context: EventContext::default(), }; // NOTES: Changes it will break backward compatibility. let serialized = r#"{"catalog":"greptime","schema":"public","from_peer":{"id":1,"addr":""},"to_peer":{"id":2,"addr":""},"region_id":4398046511105}"#; let deserialized: PersistentContext = serde_json::from_str(serialized).unwrap(); assert_eq!(persistent_ctx, deserialized); + assert_eq!(deserialized.event_context, EventContext::default()); + } + + #[test] + fn test_new_persistent_context_keeps_event_context() { + let mut event_context = + EventContext::new(TriggerReason::AutoRebalance).with_protocol("internal"); + event_context + .extensions + .insert("scheduler".into(), serde_json::json!("load_based")); + + let context = PersistentContext::new( + vec![("greptime".into(), "public".into())], + Peer::empty(1), + Peer::empty(2), + vec![RegionId::new(1024, 1)], + Duration::from_secs(10), + event_context.clone(), + ); + + assert_eq!( + context.trigger_reason, + RegionMigrationTriggerReason::AutoRebalance + ); + assert_eq!(context.event_context, event_context); + } + + #[test] + fn test_normalize_trigger_reasons() { + let mut context = new_persistent_context(); + context.trigger_reason = RegionMigrationTriggerReason::Manual; + context.event_context = EventContext::default(); + context.normalize_trigger_reasons(); + assert_eq!(context.event_context.reason, TriggerReason::Manual); + + context.trigger_reason = RegionMigrationTriggerReason::Unknown; + context.event_context = EventContext::new(TriggerReason::RegionFailover); + context.normalize_trigger_reasons(); + assert_eq!( + context.trigger_reason, + RegionMigrationTriggerReason::Failover + ); + + context.trigger_reason = RegionMigrationTriggerReason::Manual; + context.event_context = EventContext::new(TriggerReason::RegionFailover); + context.normalize_trigger_reasons(); + assert_eq!(context.event_context.reason, TriggerReason::Manual); } #[derive(Debug, Serialize, Deserialize, Default)] diff --git a/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs b/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs index 1aee76b89f0..c58c1f8cd1f 100644 --- a/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs +++ b/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs @@ -382,6 +382,7 @@ mod tests { use common_meta::key::table_route::TableRouteValue; use common_meta::key::test_utils::new_test_table_info; use common_meta::peer::Peer; + use common_meta::rpc::ddl::{EventContext, TriggerReason}; use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::wal_provider::RegionWalOptions; use store_api::storage::RegionId; @@ -389,7 +390,6 @@ mod tests { use super::*; use crate::error::Error; - use crate::procedure::region_migration::manager::RegionMigrationTriggerReason; use crate::procedure::region_migration::test_util::{TestingEnv, new_procedure_context}; use crate::procedure::region_migration::{ContextFactory, PersistentContext}; use crate::procedure::test_util::{ @@ -403,7 +403,7 @@ mod tests { Peer::empty(2), vec![RegionId::new(1024, 1)], Duration::from_millis(1000), - RegionMigrationTriggerReason::Manual, + EventContext::new(TriggerReason::Manual), ) } diff --git a/src/meta-srv/src/procedure/region_migration/manager.rs b/src/meta-srv/src/procedure/region_migration/manager.rs index d18458b08f4..3de79735e3b 100644 --- a/src/meta-srv/src/procedure/region_migration/manager.rs +++ b/src/meta-srv/src/procedure/region_migration/manager.rs @@ -21,6 +21,7 @@ use std::time::Duration; use common_meta::key::table_info::TableInfoValue; use common_meta::key::table_route::TableRouteValue; use common_meta::peer::Peer; +use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::RegionRoute; use common_procedure::{ProcedureId, ProcedureManagerRef, ProcedureWithId, watcher}; use common_telemetry::{error, info, warn}; @@ -485,7 +486,7 @@ impl RegionMigrationManager { task.to_peer.clone(), task.region_ids.clone(), task.timeout, - task.trigger_reason, + EventContext::new(task.trigger_reason.to_trigger_reason()), ), self.context_factory.clone(), procedure_guards, @@ -581,7 +582,7 @@ impl RegionMigrationManager { to_peer, vec![region_id], timeout, - trigger_reason, + EventContext::new(trigger_reason.to_trigger_reason()), ), self.context_factory.clone(), vec![guard], diff --git a/src/meta-srv/src/procedure/region_migration/test_util.rs b/src/meta-srv/src/procedure/region_migration/test_util.rs index 269357bc0ac..70342608646 100644 --- a/src/meta-srv/src/procedure/region_migration/test_util.rs +++ b/src/meta-srv/src/procedure/region_migration/test_util.rs @@ -27,6 +27,7 @@ use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::peer::Peer; use common_meta::region_keeper::{MemoryRegionKeeper, MemoryRegionKeeperRef}; +use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::RegionRoute; use common_meta::sequence::SequenceBuilder; use common_meta::state_store::KvStateStore; @@ -45,9 +46,7 @@ use crate::metasrv::MetasrvInfo; use crate::procedure::region_migration::close_downgraded_region::CloseDowngradedRegion; use crate::procedure::region_migration::downgrade_leader_region::DowngradeLeaderRegion; use crate::procedure::region_migration::flush_leader_region::PreFlushRegion; -use crate::procedure::region_migration::manager::{ - RegionMigrationProcedureTracker, RegionMigrationTriggerReason, -}; +use crate::procedure::region_migration::manager::RegionMigrationProcedureTracker; use crate::procedure::region_migration::migration_abort::RegionMigrationAbort; use crate::procedure::region_migration::migration_end::RegionMigrationEnd; use crate::procedure::region_migration::open_candidate_region::OpenCandidateRegion; @@ -191,7 +190,7 @@ pub fn new_persistent_context(from: u64, to: u64, region_id: RegionId) -> Persis Peer::empty(to), vec![region_id], Duration::from_secs(10), - RegionMigrationTriggerReason::default(), + EventContext::default(), ) } diff --git a/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs index debe72f4632..16d0270354b 100644 --- a/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs +++ b/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs @@ -356,6 +356,7 @@ mod tests { use common_meta::key::topic_name::TopicNameKey; use common_meta::key::topic_region::{ReplayCheckpoint, TopicRegionKey, TopicRegionValue}; use common_meta::peer::Peer; + use common_meta::rpc::ddl::{EventContext, TriggerReason}; use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::wal_provider::RegionWalOptions; use common_wal::options::KafkaWalOptions; @@ -363,7 +364,6 @@ mod tests { use super::*; use crate::error::Error; - use crate::procedure::region_migration::manager::RegionMigrationTriggerReason; use crate::procedure::region_migration::test_util::{TestingEnv, new_procedure_context}; use crate::procedure::region_migration::{ContextFactory, PersistentContext}; use crate::procedure::test_util::{ @@ -377,7 +377,7 @@ mod tests { Peer::empty(2), vec![RegionId::new(1024, 1)], Duration::from_millis(1000), - RegionMigrationTriggerReason::Manual, + EventContext::new(TriggerReason::Manual), ) } diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index f8452d196fb..62f9f26aabd 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -44,11 +44,12 @@ use common_meta::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use common_meta::node_manager::NodeManagerRef; use common_meta::region_keeper::{MemoryRegionKeeperRef, OperatingRegionGuard}; use common_meta::region_registry::LeaderRegionRegistryRef; +use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::{RegionRoute, operating_leader_region_roles}; use common_meta::wal_provider::RegionWalOptions; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - BoxedProcedure, Context as ProcedureContext, Error as ProcedureError, EventContext, + BoxedProcedure, Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, LockKey, Procedure, ProcedureManagerRef, Result as ProcedureResult, Status, StringKey, }; @@ -103,6 +104,8 @@ pub struct PersistentContext { #[serde(default)] /// Records table-level partition metadata updated by this repartition. pub partition_metadata_update: Option, + #[serde(default)] + pub event_context: EventContext, } fn default_timeout() -> Duration { @@ -110,7 +113,7 @@ fn default_timeout() -> Duration { } impl PersistentContext { - /// Creates a new [PersistentContext] with the given table name, table id and timeout. + /// Creates a new [PersistentContext] with the table, timeout and event context. /// /// If the timeout is not provided, the default timeout will be used. pub fn new( @@ -121,6 +124,7 @@ impl PersistentContext { }: TableName, table_id: TableId, timeout: Option, + event_context: EventContext, ) -> Self { Self { catalog_name, @@ -132,6 +136,7 @@ impl PersistentContext { unknown_procedures: vec![], timeout: timeout.unwrap_or_else(default_timeout), partition_metadata_update: None, + event_context, } } @@ -792,7 +797,10 @@ impl Procedure for RepartitionProcedure { LockKey::new(self.context.persistent_ctx.lock_key()) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx.event_type_filter.allows(REPARTITION_EVENT_TYPE) { return None; } @@ -861,6 +869,7 @@ impl RepartitionProcedureFactory for GcDisabledRepartitionProcedureFactory { _source: RepartitionSource, _to_exprs: Vec, _timeout: Option, + _event_context: EventContext, ) -> std::result::Result { Err(BoxedError::new( error::InvalidArgumentsSnafu { @@ -899,8 +908,9 @@ impl RepartitionProcedureFactory for DefaultRepartitionProcedureFactory { source: RepartitionSource, to_exprs: Vec, timeout: Option, + event_context: EventContext, ) -> std::result::Result { - let persistent_ctx = PersistentContext::new(table_name, table_id, timeout); + let persistent_ctx = PersistentContext::new(table_name, table_id, timeout, event_context); let from = match source { RepartitionSource::Partitioned { exprs, @@ -1133,6 +1143,7 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, + EventContext::default(), ); Context::new( @@ -1164,6 +1175,7 @@ mod tests { }, vec![], None, + EventContext::default(), ) .err() .expect("GC-disabled factory must reject repartition"); @@ -1265,7 +1277,7 @@ mod tests { let all = Arc::new(EventTypeFilter::All); let submitted = procedure - .event(&EventContext { + .event(&EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, @@ -1276,7 +1288,7 @@ mod tests { assert_ne!(submitted.json_payload().unwrap(), serde_json::Value::Null); let allowed = procedure - .event(&EventContext { + .event(&EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, @@ -1288,7 +1300,7 @@ mod tests { assert_eq!(allowed.event_type(), REPARTITION_EVENT_TYPE); let succeeded = procedure - .event(&EventContext { + .event(&EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Succeeded, @@ -1297,7 +1309,7 @@ mod tests { .unwrap(); assert_eq!(succeeded.json_payload().unwrap(), serde_json::Value::Null); - let filtered = procedure.event(&EventContext { + let filtered = procedure.event(&EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, @@ -1307,7 +1319,7 @@ mod tests { }); assert!(filtered.is_none()); - let empty = procedure.event(&EventContext { + let empty = procedure.event(&EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, @@ -1370,7 +1382,7 @@ mod tests { } #[test] - fn test_persistent_context_partition_metadata_update_serde_default() { + fn test_persistent_context_serde_defaults_for_new_fields() { let json = r#"{ "catalog_name":"test_catalog", "schema_name":"test_schema", @@ -1382,7 +1394,10 @@ mod tests { let persistent_ctx: PersistentContext = serde_json::from_str(json).unwrap(); + assert!(persistent_ctx.failed_procedures.is_empty()); + assert!(persistent_ctx.unknown_procedures.is_empty()); assert!(persistent_ctx.partition_metadata_update.is_none()); + assert_eq!(persistent_ctx.event_context, EventContext::default()); } #[tokio::test] @@ -1459,6 +1474,7 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, + EventContext::default(), ); persistent_ctx.plans = vec![with_rollback_metadata( test_plan(table_id), @@ -1525,6 +1541,7 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, + EventContext::default(), ); persistent_ctx.plans = vec![test_plan(table_id)]; let context = Context::new( @@ -1578,6 +1595,7 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, + EventContext::default(), ); let failed_plan = test_plan(table_id); let failed_plan = with_rollback_metadata( @@ -1711,6 +1729,7 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, + EventContext::default(), ); persistent_ctx.plans = vec![failed_plan, succeeded_plan.clone()]; persistent_ctx.failed_procedures = vec![ProcedureMeta { @@ -1820,6 +1839,7 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, + EventContext::default(), ); persistent_ctx.plans = vec![plan.clone()]; persistent_ctx.unknown_procedures = vec![ProcedureMeta { @@ -1881,6 +1901,7 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, + EventContext::default(), ); persistent_ctx.plans = vec![with_rollback_metadata( test_plan(table_id), @@ -2009,6 +2030,7 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, + EventContext::default(), ); persistent_ctx.plans = vec![failed_merge_plan, succeeded_split_plan.clone()]; persistent_ctx.failed_procedures = vec![ProcedureMeta { diff --git a/src/meta-srv/src/procedure/repartition/collect.rs b/src/meta-srv/src/procedure/repartition/collect.rs index 1a6d0c62574..56d9a007a16 100644 --- a/src/meta-srv/src/procedure/repartition/collect.rs +++ b/src/meta-srv/src/procedure/repartition/collect.rs @@ -136,6 +136,7 @@ mod tests { use common_error::mock::MockError; use common_error::status_code::StatusCode; + use common_meta::rpc::ddl::EventContext; use common_meta::test_util::MockDatanodeManager; use common_procedure::{ Context as ProcedureContext, ContextProvider, Error as ProcedureError, ProcedureId, @@ -193,6 +194,7 @@ mod tests { table::table_name::TableName::new("test_catalog", "test_schema", "test_table"), 1024, None, + EventContext::default(), ); let mut ctx = crate::procedure::repartition::Context::new( &ddl_ctx, @@ -227,6 +229,7 @@ mod tests { table::table_name::TableName::new("test_catalog", "test_schema", "test_table"), 1024, None, + EventContext::default(), ); let mut ctx = crate::procedure::repartition::Context::new( &ddl_ctx, diff --git a/src/meta-srv/src/procedure/repartition/group.rs b/src/meta-srv/src/procedure/repartition/group.rs index e8005a21471..700c2fdcd1d 100644 --- a/src/meta-srv/src/procedure/repartition/group.rs +++ b/src/meta-srv/src/procedure/repartition/group.rs @@ -38,8 +38,8 @@ use common_meta::peer::Peer; use common_meta::rpc::router::RegionRoute; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, - Procedure, ProcedureId, Result as ProcedureResult, Status, StringKey, + Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, + LockKey, Procedure, ProcedureId, Result as ProcedureResult, Status, StringKey, }; use common_telemetry::{error, info}; use serde::{Deserialize, Serialize}; @@ -262,7 +262,10 @@ impl Procedure for RepartitionGroupProcedure { LockKey::new(self.context.persistent_ctx.lock_key()) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx.event_type_filter.allows(REPARTITION_GROUP_EVENT_TYPE) { return None; } @@ -631,7 +634,9 @@ mod tests { use common_event_recorder::EventTypeFilter; use common_meta::key::TableMetadataManager; use common_meta::kv_backend::test_util::MockKvBackendBuilder; - use common_procedure::{EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState}; + use common_procedure::{ + EventRuntimeContext, EventTrigger, Procedure, ProcedureId, ProcedureState, + }; use crate::error::Error; use crate::event::repartition::REPARTITION_GROUP_EVENT_TYPE; @@ -708,7 +713,7 @@ mod tests { context: env.create_context(new_persistent_context(1024, vec![], vec![])), }; let state = ProcedureState::Running; - let event_context = |event_type_filter| EventContext { + let runtime_context = |event_type_filter| EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, @@ -716,7 +721,7 @@ mod tests { }; let allowed = procedure - .event(&event_context(EventTypeFilter::Only(HashSet::from([ + .event(&runtime_context(EventTypeFilter::Only(HashSet::from([ REPARTITION_GROUP_EVENT_TYPE.to_string(), ])))) .unwrap(); @@ -724,14 +729,14 @@ mod tests { assert!( procedure - .event(&event_context(EventTypeFilter::Only(HashSet::from([ + .event(&runtime_context(EventTypeFilter::Only(HashSet::from([ "another_event".to_string(), ])))) .is_none() ); assert!( procedure - .event(&event_context(EventTypeFilter::Only(HashSet::new()))) + .event(&runtime_context(EventTypeFilter::Only(HashSet::new()))) .is_none() ); } diff --git a/src/meta-srv/src/procedure/repartition/test_util.rs b/src/meta-srv/src/procedure/repartition/test_util.rs index 74289cdff9d..a6f2330b83a 100644 --- a/src/meta-srv/src/procedure/repartition/test_util.rs +++ b/src/meta-srv/src/procedure/repartition/test_util.rs @@ -24,6 +24,7 @@ use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::node_manager::NodeManagerRef; use common_meta::peer::Peer; +use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::sequence::SequenceBuilder; use common_meta::test_util::new_ddl_context_with_kv_backend; @@ -224,6 +225,7 @@ pub fn new_parent_context( TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, + EventContext::default(), ); ParentContext::new( diff --git a/src/meta-srv/src/procedure/tests.rs b/src/meta-srv/src/procedure/tests.rs index 4f70dcca95c..dd855c66aa0 100644 --- a/src/meta-srv/src/procedure/tests.rs +++ b/src/meta-srv/src/procedure/tests.rs @@ -29,7 +29,7 @@ use common_meta::ddl::test_util::create_table::{ }; use common_meta::key::table_route::{PhysicalTableRouteValue, TableRouteValue}; use common_meta::node_manager::NodeManagerRef; -use common_meta::rpc::ddl::CreateTableTask; +use common_meta::rpc::ddl::{CreateTableTask, EventContext, QueryContext}; use common_meta::rpc::router::{RegionRoute, find_leaders}; use common_procedure::Status; use store_api::storage::RegionId; @@ -91,6 +91,8 @@ fn create_table_task(table_name: Option<&str>) -> CreateTableTask { fn test_region_request_builder() { let mut procedure = CreateTableProcedure::new( create_table_task(None), + QueryContext::default(), + EventContext::default(), test_data::new_ddl_context(Arc::new(NodeClients::default())), ) .unwrap(); @@ -188,6 +190,8 @@ async fn test_on_datanode_create_regions() { let mut procedure = CreateTableProcedure::new( create_table_task(None), + QueryContext::default(), + EventContext::default(), test_data::new_ddl_context(node_manager), ) .unwrap(); @@ -260,8 +264,12 @@ async fn test_on_datanode_create_logical_regions() { .unwrap() .0; let _ = kv_backend.txn(physical_route_txn).await.unwrap(); - let mut procedure = - CreateLogicalTablesProcedure::new(vec![task1, task2, task3], physical_table_id, ctx); + let mut procedure = CreateLogicalTablesProcedure::new( + vec![task1, task2, task3], + physical_table_id, + EventContext::default(), + ctx, + ); let expected_created_regions = Arc::new(Mutex::new(HashMap::from([(1, 3), (2, 3), (3, 3)]))); diff --git a/src/meta-srv/src/procedure/wal_prune.rs b/src/meta-srv/src/procedure/wal_prune.rs index b5b556a778d..128fc58c44d 100644 --- a/src/meta-srv/src/procedure/wal_prune.rs +++ b/src/meta-srv/src/procedure/wal_prune.rs @@ -25,8 +25,8 @@ use common_meta::lock_key::RemoteWalLock; use common_meta::region_registry::LeaderRegionRegistryRef; use common_procedure::error::ToJsonSnafu; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, - Procedure, ProcedureState, Result as ProcedureResult, Status, StringKey, + Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, + LockKey, Procedure, ProcedureState, Result as ProcedureResult, Status, StringKey, }; use common_telemetry::{info, warn}; use manager::{WalPruneProcedureGuard, WalPruneProcedureTracker}; @@ -213,7 +213,10 @@ impl Procedure for WalPruneProcedure { LockKey::new(vec![lock_key]) } - fn event(&self, ctx: &EventContext<'_>) -> Option> { + fn event( + &self, + ctx: &EventRuntimeContext<'_>, + ) -> Option> { if !ctx.event_type_filter.allows(WAL_PRUNE_EVENT_TYPE) { return None; } @@ -503,7 +506,7 @@ mod tests { let mut procedure = WalPruneProcedure::new(context, None, "test_topic".to_string(), 42, false); let running = ProcedureState::Running; - let event_context = |trigger, lifecycle_state, event_type_filter| EventContext { + let runtime_context = |trigger, lifecycle_state, event_type_filter| EventRuntimeContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, @@ -513,13 +516,13 @@ mod tests { for trigger in [EventTrigger::Submitted, EventTrigger::Recovered] { assert!( procedure - .event(&event_context(trigger, &running, EventTypeFilter::All)) + .event(&runtime_context(trigger, &running, EventTypeFilter::All)) .is_none() ); } let retrying_event = procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Retrying { phase: common_procedure::RetryPhase::Execute, attempt: 1, @@ -537,7 +540,7 @@ mod tests { procedure.observed_latest_offset = Some(100); for trigger in [EventTrigger::Failed, EventTrigger::Poisoned] { let event = procedure - .event(&event_context(trigger, &running, EventTypeFilter::All)) + .event(&runtime_context(trigger, &running, EventTypeFilter::All)) .unwrap(); assert_eq!( event.extra_rows().unwrap()[0].values[2].value_data, @@ -548,7 +551,7 @@ mod tests { let done_without_output = ProcedureState::Done { output: None }; assert!( procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Succeeded, &done_without_output, EventTypeFilter::All, @@ -560,7 +563,7 @@ mod tests { }; assert!( procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Succeeded, &done_with_wrong_output, EventTypeFilter::All, @@ -572,7 +575,7 @@ mod tests { output: Some(Arc::new(WalPruneOutcome) as Output), }; let event = procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Succeeded, &done, EventTypeFilter::All, @@ -589,7 +592,7 @@ mod tests { assert!( procedure - .event(&event_context( + .event(&runtime_context( EventTrigger::Succeeded, &done, EventTypeFilter::Only(HashSet::new()), diff --git a/src/meta-srv/src/service/procedure.rs b/src/meta-srv/src/service/procedure.rs index 87ec36def3e..b20d7f64775 100644 --- a/src/meta-srv/src/service/procedure.rs +++ b/src/meta-srv/src/service/procedure.rs @@ -102,7 +102,6 @@ impl procedure_service_server::ProcedureService for Metasrv { .try_into() .context(error::ConvertProtoDataSnafu)?; restore_create_database_creator(&metadata, header.role, &mut task, &mut query_context)?; - let resp = self .ddl_manager() .submit_ddl_task( diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 2bfcc617cb1..8205932708e 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -36,6 +36,7 @@ use common_grpc_expr::util::ColumnExpr; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::node_manager::{AffectedRows, NodeManagerRef}; use common_meta::peer::Peer; +use common_meta::rpc::ddl::TriggerReason; use common_query::Output; use common_query::native_histogram::{is_native_histogram_value_type, native_histogram_value_type}; use common_query::prelude::{greptime_timestamp, greptime_value}; @@ -656,7 +657,7 @@ impl Inserter { if !alter_tables.is_empty() { // Alter logical tables in batch. statement_executor - .alter_logical_tables(alter_tables, ctx.clone()) + .alter_logical_tables(alter_tables, ctx.clone(), TriggerReason::AutoAlter) .await?; } } @@ -677,7 +678,7 @@ impl Inserter { } for alter_expr in alter_tables.into_iter() { statement_executor - .alter_table_inner(alter_expr, ctx.clone()) + .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter) .await?; } } @@ -786,7 +787,7 @@ impl Inserter { } for alter_expr in alter_tables.into_iter() { statement_executor - .alter_table_inner(alter_expr, ctx.clone()) + .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter) .await?; } } @@ -872,7 +873,12 @@ impl Inserter { // create physical table let res = statement_executor - .create_table_inner(create_table_expr, None, ctx.clone()) + .create_table_inner( + create_table_expr, + None, + ctx.clone(), + TriggerReason::AutoCreate, + ) .await; match res { @@ -1060,7 +1066,12 @@ impl Inserter { info!("Table `{table_ref}` does not exist, try creating table"); } let res = statement_executor - .create_table_inner(&mut create_table_expr, partitions, ctx.clone()) + .create_table_inner( + &mut create_table_expr, + partitions, + ctx.clone(), + TriggerReason::AutoCreate, + ) .await; let table_ref = TableReference::full( @@ -1091,7 +1102,7 @@ impl Inserter { statement_executor: &StatementExecutor, ) -> Result> { let res = statement_executor - .create_logical_tables(&create_table_exprs, ctx.clone()) + .create_logical_tables(&create_table_exprs, ctx.clone(), TriggerReason::AutoCreate) .await; match res { diff --git a/src/operator/src/procedure.rs b/src/operator/src/procedure.rs index cd2a0842f1a..1a849afaf7a 100644 --- a/src/operator/src/procedure.rs +++ b/src/operator/src/procedure.rs @@ -19,7 +19,7 @@ use common_error::ext::BoxedError; use common_function::handlers::ProcedureServiceHandler; use common_meta::key::TableMetadataManagerRef; use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutorRef}; -use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest}; +use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, TriggerReason}; use common_meta::rpc::procedure::{ GcRegionsRequest as MetaGcRegionsRequest, GcResponse as MetaGcResponse, GcTableRequest as MetaGcTableRequest, ManageRegionFollowerRequest, MigrateRegionRequest, @@ -31,7 +31,7 @@ use snafu::ResultExt; use table::table_name::TableName; use crate::error; -use crate::utils::to_meta_query_context; +use crate::utils::to_meta_query_context_with_trigger_reason; /// The operator for procedures which implements [`ProcedureServiceHandler`]. #[derive(Clone)] @@ -76,8 +76,10 @@ impl ProcedureServiceHandler for ProcedureServiceOperator { }) .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu)?; + let meta_query_context = + to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_ctx), + meta_query_context, DdlTask::new_purge_dropped_table(dropped.table_id), ); self.procedure_executor @@ -161,7 +163,9 @@ mod tests { use common_meta::key::test_utils::new_test_table_info_with_name; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::procedure_executor::{ProcedureExecutor, ProcedureExecutorRef}; - use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; + use common_meta::rpc::ddl::{ + DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TRIGGER_REASON_EXTENSION_KEY, + }; use common_meta::rpc::procedure::{MigrateRegionResponse, ProcedureStateResponse}; use session::context::QueryContextBuilder; use table::table_name::TableName; @@ -277,6 +281,13 @@ mod tests { ); assert_eq!(requests[0].query_context.current_catalog, "catalog"); assert_eq!(requests[0].query_context.current_schema, "schema"); + assert_eq!( + Some(&TriggerReason::Manual.as_ref().to_string()), + requests[0] + .query_context + .extensions + .get(TRIGGER_REASON_EXTENSION_KEY) + ); } #[tokio::test] diff --git a/src/operator/src/statement/comment.rs b/src/operator/src/statement/comment.rs index dc3eff0b8d8..824287fc80b 100644 --- a/src/operator/src/statement/comment.rs +++ b/src/operator/src/statement/comment.rs @@ -16,7 +16,9 @@ use api::v1::CommentOnExpr; use common_error::ext::BoxedError; use common_meta::cache_invalidator::Context; use common_meta::procedure_executor::ExecutorContext; -use common_meta::rpc::ddl::{CommentObjectType, CommentOnTask, DdlTask, SubmitDdlTaskRequest}; +use common_meta::rpc::ddl::{ + CommentObjectType, CommentOnTask, DdlTask, SubmitDdlTaskRequest, TriggerReason, +}; use common_query::Output; use session::context::QueryContextRef; use session::table_name::table_idents_to_full_name; @@ -28,7 +30,7 @@ use crate::error::{ self, ExecuteDdlSnafu, ExternalSnafu, InvalidSqlSnafu, Result, TableMetadataManagerSnafu, }; use crate::statement::StatementExecutor; -use crate::utils::to_meta_query_context; +use crate::utils::to_meta_query_context_with_trigger_reason; impl StatementExecutor { /// Adds a comment to a database object (table, column, or flow). @@ -52,10 +54,10 @@ impl StatementExecutor { .context(TableMetadataManagerSnafu)?; let cache_idents = comment_on_task.cache_idents(); - let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_ctx), - DdlTask::new_comment_on(comment_on_task), - ); + let meta_query_context = + to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); + let request = + SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_comment_on(comment_on_task)); self.procedure_executor .submit_ddl_task(&ExecutorContext::default(), request) @@ -86,10 +88,10 @@ impl StatementExecutor { .context(TableMetadataManagerSnafu)?; let cache_idents = comment_on_task.cache_idents(); - let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_ctx), - DdlTask::new_comment_on(comment_on_task), - ); + let meta_query_context = + to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); + let request = + SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_comment_on(comment_on_task)); self.procedure_executor .submit_ddl_task(&ExecutorContext::default(), request) diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 96321b13a3f..ae55acd2fb8 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -55,7 +55,7 @@ use common_meta::rpc::ddl::trigger::CreateTriggerTask; use common_meta::rpc::ddl::trigger::DropTriggerTask; use common_meta::rpc::ddl::{ CreateFlowTask, CreatorGrantIntent, DdlTask, DropFlowTask, DropViewTask, SubmitDdlTaskRequest, - SubmitDdlTaskResponse, + SubmitDdlTaskResponse, TriggerReason, }; use common_query::Output; use common_recordbatch::{RecordBatch, RecordBatches}; @@ -118,7 +118,10 @@ use crate::error::{ use crate::expr_helper::{self, RepartitionRequest, RepartitionSource}; use crate::statement::StatementExecutor; use crate::statement::show::create_partitions_stmt; -use crate::utils::{to_meta_query_context, to_meta_query_context_with_origin_frontend}; +use crate::utils::{ + to_meta_query_context_with_origin_frontend_and_trigger_reason, + to_meta_query_context_with_trigger_reason, +}; #[derive(Debug, Clone, Copy)] struct DdlSubmitOptions { @@ -302,7 +305,7 @@ impl StatementExecutor { } } - self.create_table_inner(create_expr, stmt.partitions, ctx) + self.create_table_inner(create_expr, stmt.partitions, ctx, TriggerReason::Manual) .await } @@ -359,7 +362,8 @@ impl StatementExecutor { ); let create_expr = &mut expr_helper::create_to_expr(&create_stmt, &ctx)?; - self.create_table_inner(create_expr, partitions, ctx).await + self.create_table_inner(create_expr, partitions, ctx, TriggerReason::Manual) + .await } #[tracing::instrument(skip_all)] @@ -371,7 +375,8 @@ impl StatementExecutor { let create_expr = &mut expr_helper::create_external_expr(create_expr, &ctx, &self.local_file_access) .await?; - self.create_table_inner(create_expr, None, ctx).await + self.create_table_inner(create_expr, None, ctx, TriggerReason::Manual) + .await } #[tracing::instrument(skip_all)] @@ -380,6 +385,7 @@ impl StatementExecutor { create_table: &mut CreateTableExpr, partitions: Option, query_ctx: QueryContextRef, + trigger_reason: TriggerReason, ) -> Result { ensure_table_writable(&create_table.schema_name, &create_table.table_name)?; @@ -395,16 +401,20 @@ impl StatementExecutor { .await?; } // Create logical tables - self.create_logical_tables(std::slice::from_ref(create_table), query_ctx) - .await? - .into_iter() - .next() - .context(error::UnexpectedSnafu { - violated: "expected to create logical tables", - }) + self.create_logical_tables( + std::slice::from_ref(create_table), + query_ctx, + trigger_reason, + ) + .await? + .into_iter() + .next() + .context(error::UnexpectedSnafu { + violated: "expected to create logical tables", + }) } else { // Create other normal table - self.create_non_logic_table(create_table, partitions, query_ctx) + self.create_non_logic_table(create_table, partitions, query_ctx, trigger_reason) .await } } @@ -415,6 +425,7 @@ impl StatementExecutor { create_table: &mut CreateTableExpr, partitions: Option, query_ctx: QueryContextRef, + trigger_reason: TriggerReason, ) -> Result { let _timer = crate::metrics::DIST_CREATE_TABLE.start_timer(); @@ -483,6 +494,7 @@ impl StatementExecutor { partitions, table_info.clone(), query_ctx, + trigger_reason, ) .await?; @@ -510,6 +522,7 @@ impl StatementExecutor { &self, create_table_exprs: &[CreateTableExpr], query_context: QueryContextRef, + trigger_reason: TriggerReason, ) -> Result> { let _timer = crate::metrics::DIST_CREATE_TABLES.start_timer(); ensure!( @@ -540,7 +553,7 @@ impl StatementExecutor { .collect::>(); let resp = self - .create_logical_tables_procedure(tables_data, query_context.clone()) + .create_logical_tables_procedure(tables_data, query_context.clone(), trigger_reason) .await?; let table_ids = resp.table_ids; @@ -703,7 +716,7 @@ impl StatementExecutor { .context(error::InvalidExprSnafu)?; let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), DdlTask::new_create_trigger(task), ); @@ -767,7 +780,7 @@ impl StatementExecutor { }) .context(error::InvalidExprSnafu)?; let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), DdlTask::new_create_flow(task), ); @@ -1080,7 +1093,7 @@ impl StatementExecutor { }; let request = SubmitDdlTaskRequest::new( - to_meta_query_context(ctx), + to_meta_query_context_with_trigger_reason(ctx, TriggerReason::Manual), DdlTask::new_create_view(expr, view_info.clone()), ); @@ -1166,7 +1179,7 @@ impl StatementExecutor { query_context: QueryContextRef, ) -> Result { let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), DdlTask::new_drop_flow(expr), ); @@ -1201,7 +1214,7 @@ impl StatementExecutor { query_context: QueryContextRef, ) -> Result { let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), DdlTask::new_drop_trigger(expr), ); @@ -1282,7 +1295,7 @@ impl StatementExecutor { query_context: QueryContextRef, ) -> Result { let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), DdlTask::new_drop_view(expr), ); @@ -1297,6 +1310,7 @@ impl StatementExecutor { &self, alter_table_exprs: Vec, query_context: QueryContextRef, + trigger_reason: TriggerReason, ) -> Result { let _timer = crate::metrics::DIST_ALTER_TABLES.start_timer(); ensure!( @@ -1342,7 +1356,8 @@ impl StatementExecutor { // Submit procedure for each physical table let mut handles = Vec::with_capacity(groups.len()); for (_physical_table_id, exprs) in groups { - let fut = self.alter_logical_tables_procedure(exprs, query_context.clone()); + let fut = + self.alter_logical_tables_procedure(exprs, query_context.clone(), trigger_reason); handles.push(fut); } let _results = futures::future::try_join_all(handles).await?; @@ -1514,7 +1529,8 @@ impl StatementExecutor { } let expr = expr_helper::to_alter_table_expr(alter_table, &query_context)?; - self.alter_table_inner(expr, query_context).await + self.alter_table_inner(expr, query_context, TriggerReason::Manual) + .await } #[tracing::instrument(skip_all)] @@ -1747,8 +1763,10 @@ impl StatementExecutor { source: Some(source), ..Default::default() }; + let meta_query_context = + to_meta_query_context_with_trigger_reason(query_context.clone(), TriggerReason::Manual); let mut req = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context.clone()), + meta_query_context, DdlTask::new_alter_table(AlterTableExpr { catalog_name: request.catalog_name.clone(), schema_name: request.schema_name.clone(), @@ -1803,6 +1821,7 @@ impl StatementExecutor { &self, expr: AlterTableExpr, query_context: QueryContextRef, + trigger_reason: TriggerReason, ) -> Result { ensure_table_writable(&expr.schema_name, &expr.table_name)?; @@ -1852,12 +1871,12 @@ impl StatementExecutor { .await .context(TableMetadataManagerSnafu)?; + let meta_query_context = + to_meta_query_context_with_trigger_reason(query_context, trigger_reason); + let (req, invalidate_keys) = if physical_table_id == table_id { // This is physical table - let req = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), - DdlTask::new_alter_table(expr), - ); + let req = SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_alter_table(expr)); let invalidate_keys = vec![ CacheIdent::TableId(table_id), @@ -1868,7 +1887,7 @@ impl StatementExecutor { } else { // This is logical table let req = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + meta_query_context, DdlTask::new_alter_logical_tables(vec![expr]), ); @@ -1982,6 +2001,7 @@ impl StatementExecutor { partitions: Vec, table_info: TableInfo, query_context: QueryContextRef, + trigger_reason: TriggerReason, ) -> Result { let partitions = partitions .into_iter() @@ -1989,7 +2009,11 @@ impl StatementExecutor { .collect::>>()?; let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_origin_frontend(query_context, &self.origin_frontend_addr), + to_meta_query_context_with_origin_frontend_and_trigger_reason( + query_context, + &self.origin_frontend_addr, + trigger_reason, + ), DdlTask::new_create_table(create_table, partitions, table_info), ); @@ -2003,9 +2027,14 @@ impl StatementExecutor { &self, tables_data: Vec<(CreateTableExpr, TableInfo)>, query_context: QueryContextRef, + trigger_reason: TriggerReason, ) -> Result { let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_origin_frontend(query_context, &self.origin_frontend_addr), + to_meta_query_context_with_origin_frontend_and_trigger_reason( + query_context, + &self.origin_frontend_addr, + trigger_reason, + ), DdlTask::new_create_logical_tables(tables_data), ); @@ -2019,9 +2048,12 @@ impl StatementExecutor { &self, tables_data: Vec, query_context: QueryContextRef, + trigger_reason: TriggerReason, ) -> Result { + let meta_query_context = + to_meta_query_context_with_trigger_reason(query_context, trigger_reason); let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + meta_query_context, DdlTask::new_alter_logical_tables(tables_data), ); @@ -2039,7 +2071,7 @@ impl StatementExecutor { query_context: QueryContextRef, ) -> Result { let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), DdlTask::new_drop_table( table_name.catalog_name.clone(), table_name.schema_name.clone(), @@ -2063,7 +2095,7 @@ impl StatementExecutor { query_context: QueryContextRef, ) -> Result { let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), DdlTask::new_drop_database(catalog, schema, drop_if_exists), ); @@ -2079,7 +2111,7 @@ impl StatementExecutor { query_context: QueryContextRef, ) -> Result { let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), DdlTask::new_alter_database(alter_expr), ); @@ -2096,8 +2128,10 @@ impl StatementExecutor { time_ranges: Vec<(Timestamp, Timestamp)>, query_context: QueryContextRef, ) -> Result { + let meta_query_context = + to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual); let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + meta_query_context, DdlTask::new_truncate_table( table_name.catalog_name.clone(), table_name.schema_name.clone(), @@ -2189,7 +2223,7 @@ impl StatementExecutor { creator: Option, ) -> Result { let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), DdlTask::new_create_database(catalog, database, create_if_not_exists, options, creator), ); @@ -2776,8 +2810,10 @@ async fn execute_undrop_table( table_name: table_name.to_string(), })?; + let meta_query_context = + to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual); let request = SubmitDdlTaskRequest::new( - to_meta_query_context(query_context), + meta_query_context, DdlTask::new_undrop_table(dropped.table_id), ); procedure_executor @@ -2829,7 +2865,9 @@ mod test { ExecutorContext, ProcedureExecutor, ProcedureExecutorRef, }; #[cfg(feature = "enterprise")] - use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; + use common_meta::rpc::ddl::{ + DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TRIGGER_REASON_EXTENSION_KEY, + }; #[cfg(feature = "enterprise")] use common_meta::rpc::procedure::{ MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, @@ -2974,6 +3012,13 @@ mod test { let requests = procedure.requests.lock().unwrap(); assert!(matches!(&requests[0].task, DdlTask::UndropTable(task) if task.table_id == 42)); + assert_eq!( + requests[0] + .query_context + .extensions + .get(TRIGGER_REASON_EXTENSION_KEY), + Some(&TriggerReason::Manual.as_ref().to_string()) + ); assert_eq!( cache.invalidations.lock().unwrap()[0], vec![CacheIdent::TableId(42), CacheIdent::TableName(name)] diff --git a/src/operator/src/utils.rs b/src/operator/src/utils.rs index fbb760c4fca..03649713b0b 100644 --- a/src/operator/src/utils.rs +++ b/src/operator/src/utils.rs @@ -14,6 +14,9 @@ use std::sync::{Arc, RwLock}; +use common_meta::rpc::ddl::{ + ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TRIGGER_REASON_EXTENSION_KEY, TriggerReason, +}; use common_time::Timezone; use session::context::{QueryContextBuilder, QueryContextRef}; use snafu::ResultExt; @@ -34,13 +37,28 @@ pub fn to_meta_query_context( } } -pub fn to_meta_query_context_with_origin_frontend( +/// Converts a session query context and records the DDL trigger reason. +pub fn to_meta_query_context_with_trigger_reason( query_context: QueryContextRef, - origin_frontend_addr: &str, + trigger_reason: TriggerReason, ) -> common_meta::rpc::ddl::QueryContext { let mut meta_query_context = to_meta_query_context(query_context); meta_query_context.extensions.insert( - common_meta::rpc::ddl::ORIGIN_FRONTEND_ADDR_EXTENSION_KEY.to_string(), + TRIGGER_REASON_EXTENSION_KEY.to_string(), + trigger_reason.as_ref().to_string(), + ); + meta_query_context +} + +pub fn to_meta_query_context_with_origin_frontend_and_trigger_reason( + query_context: QueryContextRef, + origin_frontend_addr: &str, + trigger_reason: TriggerReason, +) -> common_meta::rpc::ddl::QueryContext { + let mut meta_query_context = + to_meta_query_context_with_trigger_reason(query_context, trigger_reason); + meta_query_context.extensions.insert( + ORIGIN_FRONTEND_ADDR_EXTENSION_KEY.to_string(), origin_frontend_addr.to_string(), ); meta_query_context @@ -69,12 +87,14 @@ mod tests { use std::collections::HashMap; use std::sync::{Arc, RwLock}; - use common_meta::rpc::ddl::ORIGIN_FRONTEND_ADDR_EXTENSION_KEY; + use common_meta::rpc::ddl::{ + ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TRIGGER_REASON_EXTENSION_KEY, TriggerReason, + }; use common_time::Timezone; use session::context::QueryContextBuilder; use super::{ - to_meta_query_context, to_meta_query_context_with_origin_frontend, + to_meta_query_context, to_meta_query_context_with_origin_frontend_and_trigger_reason, try_to_session_query_context, }; @@ -112,7 +132,11 @@ mod tests { .build(), ); - let meta_ctx = to_meta_query_context_with_origin_frontend(session_ctx, "127.0.0.1:4000"); + let meta_ctx = to_meta_query_context_with_origin_frontend_and_trigger_reason( + session_ctx, + "127.0.0.1:4000", + TriggerReason::Manual, + ); assert_eq!( meta_ctx @@ -121,5 +145,12 @@ mod tests { .map(String::as_str), Some("127.0.0.1:4000") ); + assert_eq!( + meta_ctx + .extensions + .get(TRIGGER_REASON_EXTENSION_KEY) + .map(String::as_str), + Some("manual") + ); } } diff --git a/src/pipeline/src/manager/pipeline_operator.rs b/src/pipeline/src/manager/pipeline_operator.rs index 6c4256db69b..4f60ce4e40f 100644 --- a/src/pipeline/src/manager/pipeline_operator.rs +++ b/src/pipeline/src/manager/pipeline_operator.rs @@ -19,6 +19,7 @@ use std::time::Instant; use api::v1::CreateTableExpr; use catalog::{CatalogManagerRef, RegisterSystemTableRequest}; use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, default_engine}; +use common_meta::rpc::ddl::TriggerReason; use common_telemetry::info; use common_time::FOREVER; use datatypes::timestamp::TimestampNanosecond; @@ -131,7 +132,7 @@ impl PipelineOperator { // create table self.statement_executor - .create_table_inner(&mut expr, None, ctx.clone()) + .create_table_inner(&mut expr, None, ctx.clone(), TriggerReason::AutoCreate) .await .context(CreateTableSnafu)?; diff --git a/src/session/src/context.rs b/src/session/src/context.rs index c554c0955f0..4780523e4b7 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -618,6 +618,7 @@ impl ConnInfo { } } +/// The protocol mapping for these wire values lives in [`common_session::channel_protocol`]. #[derive(Debug, PartialEq, Default, Clone, Copy)] #[repr(u8)] pub enum Channel { @@ -680,23 +681,7 @@ impl Display for Channel { impl AsRef for Channel { fn as_ref(&self) -> &str { - match self { - Channel::Mysql => "mysql", - Channel::Postgres => "postgres", - Channel::HttpSql => "httpsql", - Channel::Prometheus => "prometheus", - Channel::Otlp => "otlp", - Channel::Grpc => "grpc", - Channel::Influx => "influx", - Channel::Opentsdb => "opentsdb", - Channel::Loki => "loki", - Channel::Elasticsearch => "elasticsearch", - Channel::Jaeger => "jaeger", - Channel::Log => "log", - Channel::Promql => "promql", - Channel::Splunk => "splunk", - Channel::Unknown => "unknown", - } + common_session::channel_protocol(*self as u8).unwrap_or("unknown") } } diff --git a/src/standalone/src/procedure.rs b/src/standalone/src/procedure.rs index ec63b11703b..85973e9f643 100644 --- a/src/standalone/src/procedure.rs +++ b/src/standalone/src/procedure.rs @@ -20,6 +20,7 @@ use common_meta::ddl::DdlContext; use common_meta::ddl_manager::{RepartitionProcedureFactory, RepartitionSource}; use common_meta::key::runtime_switch::RuntimeSwitchManager; use common_meta::kv_backend::KvBackendRef; +use common_meta::rpc::ddl::EventContext; use common_meta::state_store::KvStateStore; use common_procedure::local::{EventRecorderHandle, LocalManager, ManagerConfig}; use common_procedure::options::ProcedureConfig; @@ -73,6 +74,7 @@ impl RepartitionProcedureFactory for StandaloneRepartitionProcedureFactory { _source: RepartitionSource, _to_exprs: Vec, _timeout: Option, + _event_context: EventContext, ) -> std::result::Result { Err(BoxedError::new(NoSupportRepartitionProcedureSnafu.build())) } diff --git a/tests-integration/src/tests/gc.rs b/tests-integration/src/tests/gc.rs index badaef5f7ac..5bc20e2e2cd 100644 --- a/tests-integration/src/tests/gc.rs +++ b/tests-integration/src/tests/gc.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use std::time::Duration; use common_meta::key::TableMetadataManagerRef; +use common_meta::rpc::ddl::EventContext; use common_procedure::ProcedureWithId; use common_telemetry::info; use common_test_util::recordbatch::check_output_stream; @@ -213,6 +214,7 @@ async fn test_gc_basic(store_type: &StorageType) { false, // full_file_listing Duration::from_secs(10), // timeout Default::default(), + EventContext::default(), ); // Submit the procedure to the procedure manager diff --git a/tests-integration/tests/database_ddl_event.rs b/tests-integration/tests/database_ddl_event.rs index ab8d3ee0251..79b741b977f 100644 --- a/tests-integration/tests/database_ddl_event.rs +++ b/tests-integration/tests/database_ddl_event.rs @@ -22,7 +22,7 @@ use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; use tests_integration::test_util::{StorageType, get_test_store_config}; -use crate::event_recorder_test_util::assert_single_event; +use crate::event_recorder_test_util::{assert_single_event, find_eventually_string}; const DATABASE_NAME: &str = "database_ddl_events"; @@ -122,6 +122,15 @@ WHERE type = '{event_type}' AND {submitted_payload_predicate}"# ); assert_single_event(instance, &submitted).await; + let actual = find_eventually_string( + instance, + &format!( + "SELECT json_to_string(event_context) AS event_context FROM greptime_private.events WHERE type = '{event_type}' AND procedure_state = 'Running' AND json_path_match(procedure_trigger, '$.type == \"Submitted\"') ORDER BY timestamp DESC LIMIT 1" + ), + "event_context", + ) + .await; + assert_eq!(r#"{"reason":"manual"}"#, actual); let lifecycle = format!( r#"SELECT count(*) AS event_count diff --git a/tests-integration/tests/flow_ddl_event.rs b/tests-integration/tests/flow_ddl_event.rs index 13ca36898aa..004d8382ed2 100644 --- a/tests-integration/tests/flow_ddl_event.rs +++ b/tests-integration/tests/flow_ddl_event.rs @@ -88,6 +88,7 @@ WHERE type = '{CREATE_FLOW_EVENT_TYPE}' ), ) .await; + assert_event_context(instance, CREATE_FLOW_EVENT_TYPE, &procedure_id).await; assert_single_event( instance, &format!( @@ -106,6 +107,22 @@ WHERE type = '{CREATE_FLOW_EVENT_TYPE}' .await; } +async fn assert_event_context( + instance: &Arc, + event_type: &str, + procedure_id: &str, +) { + let actual = find_eventually_string( + instance, + &format!( + "SELECT json_to_string(event_context) AS event_context FROM greptime_private.events WHERE type = '{event_type}' AND procedure_id = '{procedure_id}' AND json_path_match(procedure_trigger, '$.type == \"Submitted\"')" + ), + "event_context", + ) + .await; + assert_eq!(r#"{"reason":"manual"}"#, actual); +} + async fn assert_drop_events(instance: &Arc, flow: &str) { let procedure_id = find_submitted_procedure_id(instance, DROP_FLOW_EVENT_TYPE, flow).await; assert_single_event( @@ -125,6 +142,7 @@ WHERE type = '{DROP_FLOW_EVENT_TYPE}' ), ) .await; + assert_event_context(instance, DROP_FLOW_EVENT_TYPE, &procedure_id).await; assert_single_event( instance, &format!( diff --git a/tests-integration/tests/gc_event.rs b/tests-integration/tests/gc_event.rs index 3cea55bca78..b62ba1d3649 100644 --- a/tests-integration/tests/gc_event.rs +++ b/tests-integration/tests/gc_event.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use std::time::Duration; +use common_meta::rpc::ddl::{EventContext, TriggerReason}; use common_procedure::{ProcedureId, ProcedureWithId, watcher}; use common_test_util::temp_dir::create_temp_dir; use meta_srv::gc::{BatchGcProcedure, GcSchedulerOptions}; @@ -128,6 +129,7 @@ async fn test_batch_gc_event() { false, Duration::from_secs(10), Default::default(), + EventContext::new(TriggerReason::Manual), ); let procedure_id = ProcedureId::parse_str("00000000-0000-0000-0000-00000000bac0").unwrap(); let mut watcher = cluster @@ -142,6 +144,19 @@ async fn test_batch_gc_event() { watcher::wait(&mut watcher).await.unwrap(); assert_sst_count(&cluster, 1).await; + let submitted = format!( + r#"SELECT json_to_string(event_context) AS event_context +FROM greptime_private.events +WHERE type = 'batch_gc' + AND procedure_id = '{procedure_id}' + AND procedure_state = 'Running' + AND json_get_string(procedure_trigger, 'type') = 'Submitted'"#, + ); + assert_eq!( + r#"{"reason":"manual"}"#, + find_eventually_string(instance, &submitted, "event_context").await + ); + let succeeded = format!( r#"SELECT json_to_string(gc_report) AS gc_report FROM greptime_private.events @@ -149,7 +164,8 @@ WHERE type = 'batch_gc' AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') - AND json_is_null(payload)"#, + AND json_is_null(payload) + AND event_context IS NULL"#, ); let actual_report: serde_json::Value = serde_json::from_str(&find_eventually_string(instance, &succeeded, "gc_report").await) diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 49ae28dec0d..aa60109a616 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2567,6 +2567,15 @@ pub async fn test_prometheus_remote_write(store_type: StorageType) { .await; assert_eq!(res.status(), StatusCode::NO_CONTENT); + // The direct Prom remote-write handler creates tables without going through + // the generic inserter. Its submitted events must retain this origin. + wait_for_event_data( + &client, + "SELECT DISTINCT json_to_string(event_context) FROM greptime_private.events WHERE type IN ('create_table', 'create_logical_tables') AND json_get_string(procedure_trigger, 'type') = 'Submitted' AND json_get_string(event_context, 'protocol') = 'prometheus' AND json_get_string(event_context, 'reason') = 'auto_create'", + r#"[["{\"protocol\":\"prometheus\",\"reason\":\"auto_create\"}"]]"#, + ) + .await; + let expected = "[[\"demo\"],[\"demo_metrics\"],[\"demo_metrics_with_nanos\"],[\"greptime_physical_table\"],[\"metric1\"],[\"metric2\"],[\"metric3\"],[\"mito\"],[\"multi_labels\"],[\"numbers\"],[\"phy\"],[\"phy2\"],[\"phy_ns\"]]"; validate_data("prometheus_remote_write", &client, "show tables;", expected).await; @@ -2609,6 +2618,15 @@ pub async fn test_prometheus_remote_write(store_type: StorageType) { assert_eq!(res.status(), StatusCode::NO_CONTENT); + // Adding a label is handled by PendingRowsSchemaAlterer, which is another + // direct DDL path outside the generic inserter. + wait_for_event_data( + &client, + "SELECT DISTINCT json_to_string(event_context) FROM greptime_private.events WHERE type IN ('alter_table', 'alter_logical_tables') AND json_get_string(procedure_trigger, 'type') = 'Submitted' AND json_get_string(event_context, 'protocol') = 'prometheus' AND json_get_string(event_context, 'reason') = 'auto_alter'", + r#"[["{\"protocol\":\"prometheus\",\"reason\":\"auto_alter\"}"]]"#, + ) + .await; + guard.remove_all().await; } @@ -7104,6 +7122,15 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { let res = send_trace_v1_req(&client, existing_int_table_name, existing_int_req, false).await; assert_eq!(StatusCode::OK, res.status()); + // Trace schema reconciliation widens the existing column through a direct + // alter-table call, so it must identify itself as an OTLP auto alteration. + wait_for_event_data( + &client, + "SELECT DISTINCT json_to_string(event_context) FROM greptime_private.events WHERE type = 'alter_table' AND json_get_string(procedure_trigger, 'type') = 'Submitted' AND json_get_string(event_context, 'protocol') = 'otlp' AND json_get_string(event_context, 'reason') = 'auto_alter'", + r#"[["{\"protocol\":\"otlp\",\"reason\":\"auto_alter\"}"]]"#, + ) + .await; + validate_data( "otlp_traces_v1_existing_int_widens_rows", &client, @@ -9854,7 +9881,20 @@ async fn validate_data(test_name: &str, client: &TestClient, sql: &str, expected } async fn wait_for_data(client: &TestClient, sql: &str, expected: &str) { - tokio::time::timeout(Duration::from_secs(10), async { + wait_for_data_with_timeout(client, sql, expected, Duration::from_secs(10)).await; +} + +async fn wait_for_event_data(client: &TestClient, sql: &str, expected: &str) { + wait_for_data_with_timeout(client, sql, expected, Duration::from_secs(30)).await; +} + +async fn wait_for_data_with_timeout( + client: &TestClient, + sql: &str, + expected: &str, + timeout: Duration, +) { + tokio::time::timeout(timeout, async { let encoded_sql = encode(sql); loop { let res = client diff --git a/tests-integration/tests/region_migration.rs b/tests-integration/tests/region_migration.rs index 2d11966dc73..3d27dabc3f7 100644 --- a/tests-integration/tests/region_migration.rs +++ b/tests-integration/tests/region_migration.rs @@ -1347,6 +1347,10 @@ async fn check_region_migration_events_system_table( Expr::cust("json_get_string(procedure_trigger, 'type')"), Alias::new("procedure_trigger"), ) + .expr_as( + Expr::cust("json_to_string(event_context)"), + Alias::new("event_context"), + ) .from((RegionMigrationEvents::Schema, RegionMigrationEvents::Table)) .and_where(Expr::col(RegionMigrationEvents::EventType).eq(REGION_MIGRATION_EVENT_TYPE)) .and_where(Expr::col(RegionMigrationEvents::ProcedureId).eq(procedure_id)) @@ -1362,11 +1366,11 @@ async fn check_region_migration_events_system_table( .remove(0); let expected = "\ -+---------------------------------+-----------------+-------------------+ -| region_migration_trigger_reason | procedure_state | procedure_trigger | -+---------------------------------+-----------------+-------------------+ -| Manual | Running | Submitted | -| Manual | Done | Succeeded | -+---------------------------------+-----------------+-------------------+"; ++---------------------------------+-----------------+-------------------+---------------------+ +| region_migration_trigger_reason | procedure_state | procedure_trigger | event_context | ++---------------------------------+-----------------+-------------------+---------------------+ +| Manual | Running | Submitted | {\"reason\":\"manual\"} | +| Manual | Done | Succeeded | | ++---------------------------------+-----------------+-------------------+---------------------+"; check_output_stream(result.unwrap().data, expected).await; } diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index 1990cc2c43d..3dac860066b 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -19,6 +19,7 @@ use std::time::Duration; use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_error::root_source; use common_meta::key::table_name::TableNameKey; +use common_meta::rpc::ddl::EventContext; use common_procedure::{ProcedureWithId, watcher}; use common_query::Output; use common_telemetry::info; @@ -816,6 +817,7 @@ async fn trigger_table_gc(metasrv: &Arc, table_name: &str) { false, // full_file_listing Duration::from_secs(10), // timeout Default::default(), + EventContext::default(), ); // Submit the procedure to the procedure manager diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index 73ce21e828a..45d87f30971 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -16,7 +16,15 @@ use std::sync::Arc; #[cfg(feature = "enterprise")] use std::time::Duration; +use api::v1::ddl_request::Expr as DdlExpr; +use api::v1::greptime_request::Request; +use api::v1::value::ValueData; +use api::v1::{ + AddColumn, AddColumns, AlterTableExpr, ColumnDataType, ColumnDef, CreateTableExpr, DdlRequest, + Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, alter_table_expr, +}; use client::OutputData; +use common_catalog::consts::MITO_ENGINE; use common_event_recorder::event_table::{ CATALOG_NAME_COLUMN, PAYLOAD_COLUMN, PHYSICAL_TABLE_ID_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_TRIGGER_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN, TABLE_NAME_COLUMN, TYPE_COLUMN, @@ -25,8 +33,9 @@ use common_test_util::temp_dir::create_temp_dir; use frontend::instance::Instance; use meta_srv::gc::GcSchedulerOptions; use mito2::gc::GcConfig; +use servers::query_handler::grpc::GrpcQueryHandler; use servers::query_handler::sql::SqlQueryHandler; -use session::context::QueryContext; +use session::context::{Channel, QueryContext}; use tests_integration::cluster::GreptimeDbClusterBuilder; use crate::event_recorder_test_util::{ @@ -37,6 +46,9 @@ const EVENTS_TABLE: &str = "greptime_private.events"; const TABLE: &str = "table_ddl_events"; const PHYSICAL_TABLE: &str = "table_ddl_events_phy"; const LOGICAL_TABLE: &str = "table_ddl_events_logical"; +const AUTO_TABLE: &str = "table_ddl_events_auto"; +const AUTO_INFLUX_TABLE: &str = "table_ddl_events_auto_influx"; +const GRPC_TABLE: &str = "table_ddl_events_grpc"; #[tokio::test(flavor = "multi_thread")] async fn test_table_ddl_procedure_events() { @@ -70,17 +82,142 @@ async fn test_table_ddl_procedure_events() { .await; let frontend = cluster.fe_instance().clone(); + auto_create_table(&frontend, AUTO_TABLE, Channel::Prometheus).await; + let auto_create_procedure_id = + submitted_procedure_id(&frontend, "create_table", AUTO_TABLE).await; + assert_event_context( + &frontend, + "create_table", + &auto_create_procedure_id, + "auto_create", + "prometheus", + ) + .await; + + auto_create_table(&frontend, AUTO_INFLUX_TABLE, Channel::Influx).await; + let auto_influx_create_procedure_id = + submitted_procedure_id(&frontend, "create_table", AUTO_INFLUX_TABLE).await; + assert_event_context( + &frontend, + "create_table", + &auto_influx_create_procedure_id, + "auto_create", + "influx", + ) + .await; + + // gRPC DDL bypasses the SQL statement executor, so it must set its own + // manual trigger reason before submitting the table procedure. + GrpcQueryHandler::do_query( + frontend.as_ref(), + Request::Ddl(DdlRequest { + expr: Some(DdlExpr::CreateTable(CreateTableExpr { + catalog_name: "greptime".to_string(), + schema_name: "public".to_string(), + table_name: GRPC_TABLE.to_string(), + column_defs: vec![ + ColumnDef { + name: "host".to_string(), + data_type: ColumnDataType::String as i32, + is_nullable: true, + semantic_type: SemanticType::Tag as i32, + ..Default::default() + }, + ColumnDef { + name: "ts".to_string(), + data_type: ColumnDataType::TimestampMillisecond as i32, + is_nullable: false, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }, + ], + time_index: "ts".to_string(), + engine: MITO_ENGINE.to_string(), + ..Default::default() + })), + }), + Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::Grpc, + )), + ) + .await + .unwrap(); + let grpc_create_procedure_id = + submitted_procedure_id(&frontend, "create_table", GRPC_TABLE).await; + assert_event_context( + &frontend, + "create_table", + &grpc_create_procedure_id, + "manual", + "grpc", + ) + .await; + GrpcQueryHandler::do_query( + frontend.as_ref(), + Request::Ddl(DdlRequest { + expr: Some(DdlExpr::AlterTable(AlterTableExpr { + catalog_name: "greptime".to_string(), + schema_name: "public".to_string(), + table_name: GRPC_TABLE.to_string(), + kind: Some(alter_table_expr::Kind::AddColumns(AddColumns { + add_columns: vec![AddColumn { + column_def: Some(ColumnDef { + name: "value".to_string(), + data_type: ColumnDataType::Float64 as i32, + is_nullable: true, + semantic_type: SemanticType::Field as i32, + ..Default::default() + }), + ..Default::default() + }], + })), + })), + }), + Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::Grpc, + )), + ) + .await + .unwrap(); + let grpc_alter_procedure_id = + submitted_procedure_id(&frontend, "alter_table", GRPC_TABLE).await; + assert_event_context( + &frontend, + "alter_table", + &grpc_alter_procedure_id, + "manual", + "grpc", + ) + .await; + // Act / Assert: Create Table retains its rich submitted row and records the // created table ID when it completes. - run_sql( + run_sql_with_context( &frontend, &format!( "CREATE TABLE {TABLE} (host STRING PRIMARY KEY, ts TIMESTAMP TIME INDEX, val DOUBLE)" ), + Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::HttpSql, + )), ) .await; let table_id = find_table_id(&frontend, TABLE).await; let create_table_procedure_id = submitted_procedure_id(&frontend, "create_table", TABLE).await; + assert_event_context( + &frontend, + "create_table", + &create_table_procedure_id, + "manual", + "httpsql", + ) + .await; assert_named_submitted_event( &frontend, "create_table", @@ -110,12 +247,25 @@ async fn test_table_ddl_procedure_events() { // Act / Assert: Alter and Truncate have a rich submitted row and lightweight // terminal lifecycle row. - run_sql( + run_sql_with_context( &frontend, &format!("ALTER TABLE {TABLE} ADD COLUMN extra STRING"), + Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::HttpSql, + )), ) .await; let alter_table_procedure_id = submitted_procedure_id(&frontend, "alter_table", TABLE).await; + assert_event_context( + &frontend, + "alter_table", + &alter_table_procedure_id, + "manual", + "httpsql", + ) + .await; assert_named_submitted_event( &frontend, "alter_table", @@ -140,9 +290,26 @@ async fn test_table_ddl_procedure_events() { &format!("INSERT INTO {TABLE} VALUES ('a', 0, 1, 'b')"), ) .await; - run_sql(&frontend, &format!("TRUNCATE TABLE {TABLE}")).await; + run_sql_with_context( + &frontend, + &format!("TRUNCATE TABLE {TABLE}"), + Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::HttpSql, + )), + ) + .await; let truncate_table_procedure_id = submitted_procedure_id(&frontend, "truncate_table", TABLE).await; + assert_event_context( + &frontend, + "truncate_table", + &truncate_table_procedure_id, + "manual", + "httpsql", + ) + .await; assert_named_submitted_event( &frontend, "truncate_table", @@ -165,8 +332,25 @@ async fn test_table_ddl_procedure_events() { // Act / Assert: Drop records the table name, while Undrop and Purge use the // table ID as their submitted locator. - run_sql(&frontend, &format!("DROP TABLE {TABLE}")).await; + run_sql_with_context( + &frontend, + &format!("DROP TABLE {TABLE}"), + Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::HttpSql, + )), + ) + .await; let drop_table_procedure_id = submitted_procedure_id(&frontend, "drop_table", TABLE).await; + assert_event_context( + &frontend, + "drop_table", + &drop_table_procedure_id, + "manual", + "httpsql", + ) + .await; assert_named_submitted_event( &frontend, "drop_table", @@ -191,10 +375,23 @@ async fn test_table_ddl_procedure_events() { let (undrop_table_procedure_id, _) = cluster .metasrv .ddl_manager() - .submit_undrop_table_task(common_meta::rpc::ddl::UndropTableTask { table_id }) + .submit_undrop_table_task( + common_meta::rpc::ddl::UndropTableTask { table_id }, + common_meta::rpc::ddl::EventContext::new( + common_meta::rpc::ddl::TriggerReason::Manual, + ), + ) .await .unwrap(); let undrop_table_procedure_id = undrop_table_procedure_id.to_string(); + assert_event_context( + &frontend, + "undrop_table", + &undrop_table_procedure_id, + "manual", + "unknown", + ) + .await; assert_id_submitted_event( &frontend, "undrop_table", @@ -218,12 +415,23 @@ async fn test_table_ddl_procedure_events() { let (purge_table_procedure_id, _) = cluster .metasrv .ddl_manager() - .submit_purge_dropped_table_task(common_meta::rpc::ddl::PurgeDroppedTableTask { - table_id, - }) + .submit_purge_dropped_table_task( + common_meta::rpc::ddl::PurgeDroppedTableTask { table_id }, + common_meta::rpc::ddl::EventContext::new( + common_meta::rpc::ddl::TriggerReason::Manual, + ), + ) .await .unwrap(); let purge_table_procedure_id = purge_table_procedure_id.to_string(); + assert_event_context( + &frontend, + "purge_dropped_table", + &purge_table_procedure_id, + "manual", + "unknown", + ) + .await; assert_id_submitted_event( &frontend, "purge_dropped_table", @@ -255,16 +463,29 @@ async fn test_table_ddl_procedure_events() { ) .await; let physical_table_id = find_table_id(&frontend, PHYSICAL_TABLE).await; - run_sql( + run_sql_with_context( &frontend, &format!( "CREATE TABLE {LOGICAL_TABLE} (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) ENGINE=metric WITH (\"on_physical_table\" = \"{PHYSICAL_TABLE}\")" ), + Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::HttpSql, + )), ) .await; let logical_table_id = find_table_id(&frontend, LOGICAL_TABLE).await; let create_logical_tables_procedure_id = submitted_procedure_id(&frontend, "create_logical_tables", LOGICAL_TABLE).await; + assert_event_context( + &frontend, + "create_logical_tables", + &create_logical_tables_procedure_id, + "manual", + "httpsql", + ) + .await; // Assert: logical Create emits one submitted row per logical table and // preserves the logical and physical table IDs on its terminal row. @@ -298,13 +519,26 @@ async fn test_table_ddl_procedure_events() { // Act / Assert: logical Alter preserves the one-row-per-logical-table submitted // contract and emits a lightweight terminal row. - run_sql( + run_sql_with_context( &frontend, &format!("ALTER TABLE {LOGICAL_TABLE} ADD COLUMN rack STRING PRIMARY KEY"), + Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::HttpSql, + )), ) .await; let alter_logical_tables_procedure_id = submitted_procedure_id(&frontend, "alter_logical_tables", LOGICAL_TABLE).await; + assert_event_context( + &frontend, + "alter_logical_tables", + &alter_logical_tables_procedure_id, + "manual", + "httpsql", + ) + .await; assert_named_submitted_event( &frontend, "alter_logical_tables", @@ -327,14 +561,98 @@ async fn test_table_ddl_procedure_events() { } async fn run_sql(instance: &Arc, sql: &str) { - let output = instance - .do_query(sql, QueryContext::arc()) + run_sql_with_context(instance, sql, QueryContext::arc()).await; +} + +async fn auto_create_table(instance: &Arc, table_name: &str, channel: Channel) { + instance + .handle_row_inserts( + RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: table_name.to_string(), + rows: Some(Rows { + schema: vec![ + api::v1::ColumnSchema { + column_name: "host".to_string(), + datatype: ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + datatype_extension: None, + options: None, + }, + api::v1::ColumnSchema { + column_name: "ts".to_string(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + datatype_extension: None, + options: None, + }, + api::v1::ColumnSchema { + column_name: "val".to_string(), + datatype: ColumnDataType::Float64 as i32, + semantic_type: SemanticType::Field as i32, + datatype_extension: None, + options: None, + }, + ], + rows: vec![Row { + values: vec![ + ValueData::StringValue("host".to_string()).into(), + ValueData::TimestampMillisecondValue(0).into(), + ValueData::F64Value(1.0).into(), + ], + }], + }), + }], + }, + Arc::new(QueryContext::with_channel("greptime", "public", channel)), + false, + false, + ) + .await + .unwrap(); +} + +async fn run_sql_with_context( + instance: &Arc, + sql: &str, + query_context: Arc, +) { + let output = SqlQueryHandler::do_query(instance.as_ref(), sql, query_context) .await .remove(0) .unwrap(); assert!(matches!(output.data, OutputData::AffectedRows(_)), "{sql}"); } +async fn assert_event_context( + instance: &Arc, + event_type: &str, + procedure_id: &str, + reason: &str, + protocol: &str, +) { + let event_filter = format!( + "FROM {EVENTS_TABLE} WHERE {} = '{event_type}' AND {} = '{procedure_id}' AND json_get_string({}, 'type') = 'Submitted'", + TYPE_COLUMN.name(), + PROCEDURE_ID_COLUMN.name(), + PROCEDURE_TRIGGER_COLUMN.name(), + ); + assert_single_event( + instance, + &format!("SELECT count(*) AS event_count {event_filter}"), + ) + .await; + + let actual = find_eventually_string( + instance, + &format!("SELECT json_to_string(event_context) AS event_context {event_filter}"), + "event_context", + ) + .await; + let expected = format!(r#"{{"protocol":"{protocol}","reason":"{reason}"}}"#); + assert_eq!(expected, actual); +} + async fn find_table_id(instance: &Arc, table_name: &str) -> u32 { find_eventually_u32( instance, diff --git a/tests-integration/tests/view_ddl_event.rs b/tests-integration/tests/view_ddl_event.rs index 2f362c256fe..a25b58b684a 100644 --- a/tests-integration/tests/view_ddl_event.rs +++ b/tests-integration/tests/view_ddl_event.rs @@ -124,6 +124,7 @@ WHERE type = '{CREATE_VIEW_EVENT_TYPE}' ), ) .await; + assert_event_context(instance, CREATE_VIEW_EVENT_TYPE, &procedure_id).await; assert_single_event( instance, &format!( @@ -143,6 +144,22 @@ WHERE type = '{CREATE_VIEW_EVENT_TYPE}' .await; } +async fn assert_event_context( + instance: &Arc, + event_type: &str, + procedure_id: &str, +) { + let actual = find_eventually_string( + instance, + &format!( + "SELECT json_to_string(event_context) AS event_context FROM greptime_private.events WHERE type = '{event_type}' AND procedure_id = '{procedure_id}' AND json_path_match(procedure_trigger, '$.type == \"Submitted\"')" + ), + "event_context", + ) + .await; + assert_eq!(r#"{"reason":"manual"}"#, actual); +} + async fn assert_drop_events(instance: &Arc, view: &str) { let procedure_id = find_submitted_procedure_id(instance, DROP_VIEW_EVENT_TYPE, view).await; assert_single_event( @@ -163,6 +180,7 @@ WHERE type = '{DROP_VIEW_EVENT_TYPE}' ), ) .await; + assert_event_context(instance, DROP_VIEW_EVENT_TYPE, &procedure_id).await; assert_single_event( instance, &format!( From d90cca4b752b34d7ccdb54fe8445ce3bf832e5e0 Mon Sep 17 00:00:00 2001 From: Gaurav Shokeen Date: Fri, 7 Aug 2026 19:07:00 +0530 Subject: [PATCH 023/126] fix(prometheus): custom column remote reads (#8659) * fix(prometheus): custom column remote reads Resolve timestamp and value column names from the table schema and carry them through query planning and result conversion. Add a remote-read regression test covering custom_ts and custom_value. Signed-off-by: grezzko * fix: resolve remote-read value columns safely Prefer the sole field for custom schemas and greptime_value for multi-field tables. Reject ambiguous schemas and add regression tests. Signed-off-by: grezzko --------- Signed-off-by: grezzko Signed-off-by: Lei, HUANG Co-authored-by: Lei, HUANG --- src/frontend/src/error.rs | 11 ++ src/frontend/src/instance.rs | 213 +++++++++++++++++++++++- src/frontend/src/instance/prom_store.rs | 107 ++++++++++-- src/servers/src/prom_store.rs | 90 +++++++--- 4 files changed, 378 insertions(+), 43 deletions(-) diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs index 144fb96ab94..9a70b716619 100644 --- a/src/frontend/src/error.rs +++ b/src/frontend/src/error.rs @@ -164,6 +164,16 @@ pub enum Error { location: Location, }, + #[snafu(display( + "Ambiguous value column in table '{table_name}', candidates: {field_columns:?}" + ))] + AmbiguousValueColumn { + table_name: String, + field_columns: Vec, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to collect recordbatch"))] CollectRecordbatch { #[snafu(implicit)] @@ -377,6 +387,7 @@ impl ErrorExt for Error { | Error::IllegalPrimaryKeysDef { .. } | Error::SchemaExists { .. } | Error::ColumnNotFound { .. } + | Error::AmbiguousValueColumn { .. } | Error::UnsupportedFormat { .. } | Error::IllegalAuthConfig { .. } | Error::ColumnNoneDefaultValue { .. } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 14acf87787e..23090ef932b 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1675,7 +1675,9 @@ mod tests { use std::time::Duration; use api::prom_store::remote::label_matcher::Type as PromMatcherType; - use api::prom_store::remote::{LabelMatcher, Query as RemoteQuery, ReadRequest}; + use api::prom_store::remote::{ + Label, LabelMatcher, Query as RemoteQuery, ReadRequest, ReadResponse, Sample, + }; use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse}; use auth::{ DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, JAEGER_QUERY, PIPELINE_DELETE, @@ -1684,7 +1686,7 @@ mod tests { use catalog::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer}; use common_base::Plugins; use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME; - use common_error::ext::{BoxedError, PlainError}; + use common_error::ext::{BoxedError, ErrorExt, PlainError}; use common_error::status_code::StatusCode; use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef}; use common_frontend::slow_query_event::SlowQueryEvent; @@ -1695,6 +1697,7 @@ mod tests { use common_meta::rpc::procedure::{ MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, }; + use common_query::prelude::greptime_value; use common_query::{Output, OutputMeta}; use common_recordbatch::{ OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream, @@ -1706,8 +1709,12 @@ mod tests { use datafusion_expr::{LogicalPlanBuilder, LogicalTableSource}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema as GtSchema, SchemaRef as GtSchemaRef}; - use datatypes::vectors::{StringVector, TimestampNanosecondVector, VectorRef}; + use datatypes::vectors::{ + Float64Vector, StringVector, TimestampMillisecondVector, TimestampNanosecondVector, + VectorRef, + }; use log_query::LogQuery; + use prost::Message; use query::query_engine::options::QueryOptions; use servers::query_handler::{ DashboardHandler, JaegerQueryHandler, LogQueryHandler, PipelineHandler, PipelineHandlerRef, @@ -2726,6 +2733,206 @@ mod tests { assert_eq!(CheckedAction { action, targets }, checker.take_check()); } + #[tokio::test] + async fn test_prom_remote_read_with_custom_timestamp_and_value_columns() -> TestResult<()> { + let schema = Arc::new(GtSchema::new(vec![ + ColumnSchema::new( + "custom_ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("custom_value", ConcreteDataType::float64_datatype(), false), + ])); + let recordbatch = RecordBatch::new( + schema, + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef, + Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef, + ], + ) + .unwrap(); + let instance = test_instance_with_tables( + MemTable::table("custom_metric", recordbatch), + test_table(1025, "target")?, + ) + .await?; + + let response = PromStoreProtocolHandler::read( + &instance, + ReadRequest { + queries: vec![RemoteQuery { + start_timestamp_ms: 1500, + end_timestamp_ms: 2500, + matchers: vec![LabelMatcher { + r#type: PromMatcherType::Eq as i32, + name: servers::prom_store::METRIC_NAME_LABEL.to_string(), + value: "custom_metric".to_string(), + }], + ..Default::default() + }], + ..Default::default() + }, + test_query_ctx(1), + ) + .await + .unwrap(); + let body = servers::prom_store::snappy_decompress(&response.body).unwrap(); + let response = ReadResponse::decode(body.as_slice()).unwrap(); + + assert_eq!(1, response.results.len()); + assert_eq!(1, response.results[0].timeseries.len()); + let timeseries = &response.results[0].timeseries[0]; + assert_eq!( + vec![Label { + name: servers::prom_store::METRIC_NAME_LABEL.to_string(), + value: "custom_metric".to_string(), + }], + timeseries.labels + ); + assert_eq!( + vec![Sample { + value: 2.0, + timestamp: 2000, + }], + timeseries.samples + ); + + Ok(()) + } + + #[tokio::test] + async fn test_prom_remote_read_prefers_default_value_column() -> TestResult<()> { + let schema = Arc::new(GtSchema::new(vec![ + ColumnSchema::new( + "custom_ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("extra_field", ConcreteDataType::float64_datatype(), false), + ColumnSchema::new( + greptime_value(), + ConcreteDataType::float64_datatype(), + false, + ), + ])); + let recordbatch = RecordBatch::new( + schema, + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef, + Arc::new(Float64Vector::from_vec(vec![99.0, 99.0, 99.0])) as VectorRef, + Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef, + ], + ) + .unwrap(); + let instance = test_instance_with_tables( + MemTable::table("multi_field_metric", recordbatch), + test_table(1025, "target")?, + ) + .await?; + + let response = PromStoreProtocolHandler::read( + &instance, + ReadRequest { + queries: vec![RemoteQuery { + start_timestamp_ms: 1500, + end_timestamp_ms: 2500, + matchers: vec![LabelMatcher { + r#type: PromMatcherType::Eq as i32, + name: servers::prom_store::METRIC_NAME_LABEL.to_string(), + value: "multi_field_metric".to_string(), + }], + ..Default::default() + }], + ..Default::default() + }, + test_query_ctx(1), + ) + .await + .unwrap(); + let body = servers::prom_store::snappy_decompress(&response.body).unwrap(); + let response = ReadResponse::decode(body.as_slice()).unwrap(); + + assert_eq!(1, response.results.len()); + assert_eq!(1, response.results[0].timeseries.len()); + let timeseries = &response.results[0].timeseries[0]; + assert_eq!( + vec![ + Label { + name: servers::prom_store::METRIC_NAME_LABEL.to_string(), + value: "multi_field_metric".to_string(), + }, + Label { + name: "extra_field".to_string(), + value: "99".to_string(), + }, + ], + timeseries.labels + ); + assert_eq!( + vec![Sample { + value: 2.0, + timestamp: 2000, + }], + timeseries.samples + ); + + Ok(()) + } + + #[tokio::test] + async fn test_prom_remote_read_rejects_ambiguous_value_columns() -> TestResult<()> { + let schema = Arc::new(GtSchema::new(vec![ + ColumnSchema::new( + "custom_ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("field_a", ConcreteDataType::float64_datatype(), false), + ColumnSchema::new("field_b", ConcreteDataType::float64_datatype(), false), + ])); + let recordbatch = RecordBatch::new( + schema, + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as VectorRef, + Arc::new(Float64Vector::from_vec(vec![1.0])) as VectorRef, + Arc::new(Float64Vector::from_vec(vec![2.0])) as VectorRef, + ], + ) + .unwrap(); + let instance = test_instance_with_tables( + MemTable::table("ambiguous_metric", recordbatch), + test_table(1025, "target")?, + ) + .await?; + + let err = PromStoreProtocolHandler::read( + &instance, + ReadRequest { + queries: vec![RemoteQuery { + matchers: vec![LabelMatcher { + r#type: PromMatcherType::Eq as i32, + name: servers::prom_store::METRIC_NAME_LABEL.to_string(), + value: "ambiguous_metric".to_string(), + }], + ..Default::default() + }], + ..Default::default() + }, + test_query_ctx(1), + ) + .await + .err() + .expect("ambiguous value columns should fail remote read"); + + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert!(format!("{err:?}").contains("Ambiguous value column")); + + Ok(()) + } + #[tokio::test] async fn test_event_recorder_is_exposed() -> TestResult<()> { let instance = diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index 7c86f3316a5..1bfb5f56437 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -32,7 +32,7 @@ use common_catalog::{format_full_table_name, parse_optional_catalog_and_schema_f use common_error::ext::BoxedError; use common_meta::rpc::ddl::TriggerReason; use common_query::Output; -use common_query::prelude::GREPTIME_PHYSICAL_TABLE; +use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_value}; use common_recordbatch::RecordBatches; use common_telemetry::{debug, tracing}; use operator::insert::{ @@ -54,17 +54,25 @@ use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt}; use store_api::metric_engine_consts::{METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY}; use store_api::mito_engine_options::SST_FORMAT_KEY; +use table::TableRef; use table::table_reference::TableReference; use tracing::instrument; use crate::error::{ - CatalogSnafu, ExecLogicalPlanSnafu, PromStoreRemoteQueryPlanSnafu, ReadTableSnafu, Result, - TableNotFoundSnafu, + AmbiguousValueColumnSnafu, CatalogSnafu, ColumnNotFoundSnafu, ExecLogicalPlanSnafu, + PromStoreRemoteQueryPlanSnafu, ReadTableSnafu, Result, TableNotFoundSnafu, }; use crate::instance::Instance; const SAMPLES_RESPONSE_TYPE: i32 = ResponseType::Samples as i32; +struct RemoteQueryOutput { + table_name: String, + timestamp_column_name: String, + value_column_name: String, + output: Output, +} + fn auto_create_table_type_for_prom_remote_write( ctx: &QueryContextRef, with_metric_engine: bool, @@ -141,7 +149,12 @@ fn resolve_remote_query_target( } #[instrument(skip_all, fields(table_name))] -async fn to_query_result(table_name: &str, output: Output) -> ServerResult { +async fn to_query_result( + table_name: &str, + timestamp_column_name: &str, + value_column_name: &str, + output: Output, +) -> ServerResult { let OutputData::Stream(stream) = output.data else { unreachable!() }; @@ -149,10 +162,41 @@ async fn to_query_result(table_name: &str, output: Output) -> ServerResult Result { + let columns = table + .field_columns() + .map(|column| column.name) + .collect::>(); + + match columns.as_slice() { + [] => ColumnNotFoundSnafu { + msg: format!("value field in table '{table_name}'"), + } + .fail(), + + [only] => Ok(only.clone()), + + columns if columns.iter().any(|name| name == greptime_value()) => { + Ok(greptime_value().to_string()) + } + + columns => AmbiguousValueColumnSnafu { + table_name: table_name.to_string(), + field_columns: columns.to_vec(), + } + .fail(), + } +} + impl Instance { #[tracing::instrument(skip_all)] async fn handle_remote_query( @@ -162,7 +206,7 @@ impl Instance { schema_name: &str, table_name: &str, query: &Query, - ) -> Result { + ) -> Result { let table = self .catalog_manager .table(catalog_name, schema_name, table_name, Some(ctx)) @@ -172,6 +216,17 @@ impl Instance { table_name: format_full_table_name(catalog_name, schema_name, table_name), })?; + let timestamp_column_name = table + .schema() + .timestamp_column() + .with_context(|| ColumnNotFoundSnafu { + msg: format!("time index in table '{table_name}'"), + })? + .name + .clone(); + + let value_column_name = resolve_column_names(table_name, &table)?; + let dataframe = self .query_engine .read_table(table) @@ -179,8 +234,8 @@ impl Instance { table_name: format_full_table_name(catalog_name, schema_name, table_name), })?; - let logical_plan = - prom_store::query_to_plan(dataframe, query).context(PromStoreRemoteQueryPlanSnafu)?; + let logical_plan = prom_store::query_to_plan(dataframe, query, ×tamp_column_name) + .context(PromStoreRemoteQueryPlanSnafu)?; debug!( "Prometheus remote read, table: {}, logical plan: {}", @@ -188,10 +243,18 @@ impl Instance { logical_plan.display_indent(), ); - self.query_engine + let output = self + .query_engine .execute(logical_plan, ctx.clone()) .await - .context(ExecLogicalPlanSnafu) + .context(ExecLogicalPlanSnafu)?; + + Ok(RemoteQueryOutput { + table_name: table_name.to_string(), + timestamp_column_name, + value_column_name, + output, + }) } #[tracing::instrument(skip_all)] @@ -200,17 +263,17 @@ impl Instance { ctx: QueryContextRef, queries: &[Query], query_targets: &[PermissionTableTarget], - ) -> ServerResult> { + ) -> ServerResult> { let mut results = Vec::with_capacity(queries.len()); for (query, target) in queries.iter().zip(query_targets) { - let output = self + let result = self .handle_remote_query(&ctx, &target.catalog, &target.schema, &target.table, query) .await .map_err(BoxedError::new) .context(error::ExecuteQuerySnafu)?; - results.push((target.table.clone(), output)); + results.push(result); } Ok(results) } @@ -527,9 +590,23 @@ impl PromStoreProtocolHandler for Instance { ResponseType::Samples => { let mut query_results = Vec::with_capacity(results.len()); let mut map = HashMap::new(); - for (table_name, output) in results { + for result in results { + let RemoteQueryOutput { + table_name, + timestamp_column_name, + value_column_name, + output, + } = result; let plan = output.meta.plan.clone(); - query_results.push(to_query_result(&table_name, output).await?); + query_results.push( + to_query_result( + &table_name, + ×tamp_column_name, + &value_column_name, + output, + ) + .await?, + ); if let Some(ref plan) = plan { collect_plan_metrics(plan, &mut [&mut map]); } diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 772689f0e07..33707b29358 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -137,7 +137,11 @@ pub fn extract_schema_from_query(query: &Query) -> Option { /// Create a DataFrame from a remote Query #[tracing::instrument(skip_all)] -pub fn query_to_plan(dataframe: DataFrame, q: &Query) -> Result { +pub fn query_to_plan( + dataframe: DataFrame, + q: &Query, + timestamp_column_name: &str, +) -> Result { let start_timestamp_ms = q.start_timestamp_ms; let end_timestamp_ms = q.end_timestamp_ms; @@ -145,8 +149,9 @@ pub fn query_to_plan(dataframe: DataFrame, q: &Query) -> Result { let mut conditions = Vec::with_capacity(label_matches.len() + 1); - conditions.push(col(greptime_timestamp()).gt_eq(lit_timestamp_millisecond(start_timestamp_ms))); - conditions.push(col(greptime_timestamp()).lt_eq(lit_timestamp_millisecond(end_timestamp_ms))); + conditions + .push(col(timestamp_column_name).gt_eq(lit_timestamp_millisecond(start_timestamp_ms))); + conditions.push(col(timestamp_column_name).lt_eq(lit_timestamp_millisecond(end_timestamp_ms))); for m in label_matches { let name = &m.name; @@ -261,14 +266,18 @@ struct LabelColumn<'a> { values: LabelValues<'a>, } -fn label_columns(recordbatch: &RecordBatch) -> Result>> { +fn label_columns<'a>( + recordbatch: &'a RecordBatch, + timestamp_column_name: &str, + value_column_name: &str, +) -> Result>> { recordbatch .schema .column_schemas() .iter() .enumerate() .filter(|(_, column_schema)| { - column_schema.name != greptime_timestamp() && column_schema.name != greptime_value() + column_schema.name != timestamp_column_name && column_schema.name != value_column_name }) .map(|(index, column_schema)| { let array = recordbatch.column(index); @@ -354,24 +363,31 @@ fn new_timeseries(table: &str, columns: &[LabelColumn<'_>], row: usize) -> TimeS pub fn recordbatches_to_timeseries( table_name: &str, + timestamp_column_name: &str, + value_column_name: &str, recordbatches: RecordBatches, ) -> Result> { Ok(recordbatches .take() .into_iter() - .map(|x| recordbatch_to_timeseries(table_name, x)) + .map(|x| recordbatch_to_timeseries(table_name, timestamp_column_name, value_column_name, x)) .collect::>>()? .into_iter() .flatten() .collect()) } -fn recordbatch_to_timeseries(table: &str, recordbatch: RecordBatch) -> Result> { - let ts_column = recordbatch.column_by_name(greptime_timestamp()).context( - error::InvalidPromRemoteReadQueryResultSnafu { - msg: "missing greptime_timestamp column in query result", - }, - )?; +fn recordbatch_to_timeseries( + table: &str, + timestamp_column_name: &str, + value_column_name: &str, + recordbatch: RecordBatch, +) -> Result> { + let ts_column = recordbatch + .column_by_name(timestamp_column_name) + .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { + msg: format!("missing timestamp column '{timestamp_column_name}' in query result"), + })?; let ts_column = ts_column .as_primitive_opt::() .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { @@ -383,11 +399,11 @@ fn recordbatch_to_timeseries(table: &str, recordbatch: RecordBatch) -> Result() .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { @@ -397,7 +413,7 @@ fn recordbatch_to_timeseries(table: &str, recordbatch: RecordBatch) -> Result = Vec::new(); let mut timeseries_by_hash: HashMap> = HashMap::new(); let mut previous_timeseries: Option = None; @@ -834,7 +850,7 @@ mod tests { let table_provider = Arc::new(DfTableProviderAdapter::new(table)); let dataframe = ctx.read_table(table_provider.clone()).unwrap(); - let plan = query_to_plan(dataframe, &q).unwrap(); + let plan = query_to_plan(dataframe, &q, greptime_timestamp()).unwrap(); let display_string = format!("{}", plan.display_indent()); let ts_col = greptime_timestamp(); @@ -868,7 +884,7 @@ mod tests { }; let dataframe = ctx.read_table(table_provider).unwrap(); - let plan = query_to_plan(dataframe, &q).unwrap(); + let plan = query_to_plan(dataframe, &q, greptime_timestamp()).unwrap(); let display_string = format!("{}", plan.display_indent()); let ts_col = greptime_timestamp(); @@ -1056,7 +1072,13 @@ mod tests { ) .unwrap(); - let timeseries = recordbatches_to_timeseries("metric1", recordbatches).unwrap(); + let timeseries = recordbatches_to_timeseries( + "metric1", + greptime_timestamp(), + greptime_value(), + recordbatches, + ) + .unwrap(); assert_eq!(2, timeseries.len()); assert_eq!( @@ -1131,7 +1153,7 @@ mod tests { ) .unwrap(); let recordbatch = RecordBatch::from_df_record_batch(schema.clone(), batch); - let columns = label_columns(&recordbatch).unwrap(); + let columns = label_columns(&recordbatch, greptime_timestamp(), greptime_value()).unwrap(); assert!(matches!( columns[0].values, LabelValues::DictionaryUtf8 { .. } @@ -1139,7 +1161,13 @@ mod tests { drop(columns); let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap(); - let timeseries = recordbatches_to_timeseries("metric1", recordbatches).unwrap(); + let timeseries = recordbatches_to_timeseries( + "metric1", + greptime_timestamp(), + greptime_value(), + recordbatches, + ) + .unwrap(); assert_eq!(3, timeseries.len()); assert_eq!( @@ -1207,7 +1235,13 @@ mod tests { ) .unwrap(); - let timeseries = recordbatch_to_timeseries("metric1", recordbatch).unwrap(); + let timeseries = recordbatch_to_timeseries( + "metric1", + greptime_timestamp(), + greptime_value(), + recordbatch, + ) + .unwrap(); // The result stays sorted by labels as it was with the previous BTreeMap. assert_eq!("host1", timeseries[0].labels[1].value); @@ -1258,7 +1292,13 @@ mod tests { ) .unwrap(); - let timeseries = recordbatch_to_timeseries("metric1", recordbatch).unwrap(); + let timeseries = recordbatch_to_timeseries( + "metric1", + greptime_timestamp(), + greptime_value(), + recordbatch, + ) + .unwrap(); assert_eq!(2, timeseries.len()); assert_eq!( From 5d4699db1cb3a01dec3b3ec77cef67887051d65a Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 10 Aug 2026 10:45:37 +0800 Subject: [PATCH 024/126] docs: refine coding agent maps (#8790) * docs: refine coding agent maps Signed-off-by: Dennis Zhuang * docs: trim license header guidance Signed-off-by: Dennis Zhuang * docs: update README links and project status Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .agents/README.md | 13 +++-- .agents/architecture-invariants.md | 47 +++++++++--------- .agents/generated-files.md | 14 ++++-- .../skills/greptimedb-release-note/SKILL.md | 7 +-- AGENTS.md | 49 +++++++++++++++++-- README.md | 4 +- src/common/meta/AGENTS.md | 38 ++++++++++++++ src/flow/AGENTS.md | 21 +++++--- src/frontend/AGENTS.md | 9 ++-- src/metric-engine/AGENTS.md | 8 +-- src/operator/AGENTS.md | 40 +++++++++++++++ src/query/AGENTS.md | 39 +++++++++++++++ src/servers/AGENTS.md | 37 ++++++++++++++ tests/perf/AGENTS.md | 5 +- 14 files changed, 273 insertions(+), 58 deletions(-) create mode 100644 src/common/meta/AGENTS.md create mode 100644 src/operator/AGENTS.md create mode 100644 src/query/AGENTS.md create mode 100644 src/servers/AGENTS.md diff --git a/.agents/README.md b/.agents/README.md index d55aac50f06..8d21de0f1a2 100644 --- a/.agents/README.md +++ b/.agents/README.md @@ -11,16 +11,23 @@ Shared agent skills live in `.agents/skills`. - Add or update shared skills under `.agents/skills//SKILL.md`. - If a new skill does not appear, restart the agent or start a new thread. -## Per-crate guides +## Per-directory guides -Hot crates carry an `AGENTS.md` next to their code as a navigation aid (module -map, read/write paths, change-coupling points, test commands, gotchas): +Read every `AGENTS.md` from the repository root down to the path being changed. +The root guide applies everywhere; nested guides add rules and navigation for +their subtree. +- [`src/common/meta/AGENTS.md`](../src/common/meta/AGENTS.md) — metadata keys, KV backends, DDL procedures, and caches +- [`src/query/AGENTS.md`](../src/query/AGENTS.md) — query planning, optimization, and distributed execution +- [`src/servers/AGENTS.md`](../src/servers/AGENTS.md) — wire protocols and network servers +- [`src/operator/AGENTS.md`](../src/operator/AGENTS.md) — statement, DDL/DML, and write orchestration - [`src/mito2/AGENTS.md`](../src/mito2/AGENTS.md) — primary time-series storage engine - [`src/metric-engine/AGENTS.md`](../src/metric-engine/AGENTS.md) — metrics engine (logical/physical regions) - [`src/flow/AGENTS.md`](../src/flow/AGENTS.md) — stream processing / continuous aggregation - [`src/frontend/AGENTS.md`](../src/frontend/AGENTS.md) — request entry point and orchestration - [`src/meta-srv/AGENTS.md`](../src/meta-srv/AGENTS.md) — metadata and cluster coordination +- [`tests/compatibility/AGENTS.md`](../tests/compatibility/AGENTS.md) — persisted/wire compatibility cases +- [`tests/perf/AGENTS.md`](../tests/perf/AGENTS.md) — query regression harness and case DSL ## Architecture invariants diff --git a/.agents/architecture-invariants.md b/.agents/architecture-invariants.md index b7ce954c972..78c78f3a976 100644 --- a/.agents/architecture-invariants.md +++ b/.agents/architecture-invariants.md @@ -20,15 +20,18 @@ metadata KV values (`common-meta` keys, metric-engine metadata), and gRPC messages. A node running an old version may read data written by a new one and vice versa. -- Add fields, don't repurpose or reorder them. For serde types use - `#[serde(default)]` / `#[serde(alias)]`; never change the meaning of an existing - field or the discriminant of an existing enum variant. -- Monotonic version counters (e.g. the mito2 manifest version) only ever move - forward — never reset or skip. +- Preserve the contract of the actual encoding: Protobuf field numbers/types, + serde field and variant names, positional field order, and explicitly encoded + enum discriminants. Do not treat Rust declaration order as a wire contract + for name-based formats such as JSON. +- Use `#[serde(default)]` for newly optional data and `#[serde(alias)]` for a + compatible rename. Do not reset or reuse versions within one persisted + history. - When a change touches a persisted or wire format, add a case to the - compatibility test suite. See `docs/rfcs/2025-07-04-compatibility-test-framework.md` - — compatibility has been broken on releases before (v0.14.1, v0.15.1) precisely - because this step was skipped. + compatibility test suite. Follow [`tests/compatibility/README.md`](../tests/compatibility/README.md) + and [`tests/compatibility/AGENTS.md`](../tests/compatibility/AGENTS.md). The + original compatibility RFC describes a superseded case format and is not the + current runbook. - Wire types are generated from the external `greptime-proto` crate; change the format there first, then bump the dependency (see invariant 6's pattern). @@ -54,9 +57,9 @@ The workspace is layered; dependencies point downward only. Runtimes are partitioned by workload so one workload can't starve another. They live in `common-runtime` (`src/common/runtime/`). -- Use the categorized spawns — `spawn_global`, `spawn_query`, `spawn_ingest`, - `spawn_compact`, `spawn_hb` — instead of constructing your own tokio runtime, - and pick the category that matches the work. +- Runtime construction belongs in process bootstrap or test/benchmark harnesses. + Product components use `spawn_global`, `spawn_query`, `spawn_ingest`, + `spawn_compact`, or `spawn_hb` instead of a component-local Tokio runtime. - Run CPU-bound or synchronous-blocking work via `spawn_blocking_*`; never do heavy CPU or blocking syscalls directly inside an async task. - Never call `block_on*` from inside an async context or an engine worker — it @@ -85,23 +88,19 @@ unfinished work merge without freezing it into the stable config surface. When you stabilize such a feature, drop the prefix and document the migration. -## 6. DataFusion is a pinned fork — two sections, two forms +## 6. DataFusion is a pinned fork — workspace dependencies plus patches GreptimeDB uses a fork at `GreptimeTeam/datafusion`, wired up in the root -`Cargo.toml` through **two sections that hold different things**: +`Cargo.toml` through two sections: -- `[workspace.dependencies]` pins each DataFusion sub-crate to an **exact - crates.io version** (e.g. `datafusion = "=53.1.0"`). -- `[patch.crates-io]` redirects those same crates to the **fork at a git rev** - (e.g. `datafusion = { git = ".../GreptimeTeam/datafusion.git", rev = "..." }`). +- `[workspace.dependencies]` pins directly referenced sub-crates to an exact + crates.io version. +- `[patch.crates-io]` redirects resolution to the fork. It may also contain + transitive-only sub-crates with no workspace dependency. -So: - -- Adding a new DataFusion sub-crate dependency means adding it to **both** - sections — the `=` entry under `[workspace.dependencies]` and the - matching git-rev patch under `[patch.crates-io]`. -- Upgrading DataFusion means bumping the version in `[workspace.dependencies]` - **and** the rev in `[patch.crates-io]` together, for all of them. +Add a direct dependency to `[workspace.dependencies]` and patch it if needed; a +transitive-only override needs only the patch. Upgrade all exact pins and fork +revisions together. ## 7. Enterprise-gated code is licensed differently from the rest diff --git a/.agents/generated-files.md b/.agents/generated-files.md index 1a16fcb28cf..23f7e81daf3 100644 --- a/.agents/generated-files.md +++ b/.agents/generated-files.md @@ -35,7 +35,14 @@ Edit the example TOMLs and/or the template, then regenerate. Do not edit ## Grafana dashboards -Generated artifacts; regenerate rather than editing by hand. +For the metrics dashboards, edit +`grafana/dashboards/metrics/cluster/dashboard.json`. The generator derives the +standalone `dashboard.json` by removing instance filters, then produces +`dashboard.yaml` and `dashboard.md` for both cluster and standalone variants. +Do not hand-edit those five derived files. + +`grafana/dashboards/events/dashboard.json` and +`grafana/dashboards/logs/dashboard.json` are not outputs of this generator. ```bash make dashboards @@ -56,5 +63,6 @@ dependency. ## License headers -Managed by the license-header tooling and the pre-commit hooks; let the tooling -add/fix them rather than editing headers by hand. +License headers are checked by `korandoru/hawkeye@v5` in +`.github/workflows/checks.yml`. For enterprise-gated files, also update both +license lists and run `make check-enterprise-license`. diff --git a/.agents/skills/greptimedb-release-note/SKILL.md b/.agents/skills/greptimedb-release-note/SKILL.md index a5cb47e9b02..893b99445ce 100644 --- a/.agents/skills/greptimedb-release-note/SKILL.md +++ b/.agents/skills/greptimedb-release-note/SKILL.md @@ -167,9 +167,10 @@ The release note is also published as a blog post in **`GreptimeTeam/docs`**. git -C fetch origin main git -C worktree add -b chore/X.Y.Z-release-note /tmp/docs-release-note origin/main ``` -- **The PR body must follow the docs repo's template** (`.github/pull_request_template.md` — - "What's Changed in this PR" + a Checklist). Fill in the description; leave checklist boxes - for the reviewer. +- **Read and follow the current docs repo PR template** + (`.github/pull_request_template.md`) from the checkout used for the release. + Fill in every required section and leave reviewer-owned checklist boxes unchecked; + do not rely on hard-coded section names from an older template. - Commit with sign-off, push, open a **draft** PR, then remove the worktree: ``` git -C /tmp/docs-release-note add blog/release-X-Y-Z.md diff --git a/AGENTS.md b/AGENTS.md index a00889b855f..c99856c2941 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -18,7 +18,8 @@ provides sub-second querying at PB scale with high cost efficiency. | Build (debug) | `make build` | | Build (release) | `make build RELEASE=true` | | Run standalone | `cargo run -- standalone start` | -| Test | `cargo nextest run` (preferred over `cargo test`) | +| Targeted Rust tests | `cargo nextest run -p ` (preferred over `cargo test`) | +| Full local Rust suite | `make test` | | SQL tests | `cargo sqlness bare` (single case: `cargo sqlness bare -t `) | | Format | `make fmt` (and `make fmt-toml` for TOML) | | Lint | `make clippy` (= `cargo clippy --workspace --all-targets --all-features -- -D warnings`) | @@ -43,14 +44,30 @@ GreptimeDB is a Cargo workspace rooted at the repository; most crates live under - **Shared**: `src/common/`, `src/datatypes/`, `src/store-api/` (engine contract), `src/catalog/`, `src/table/` -Hot crates carry their own `AGENTS.md` with a module map, read/write paths, -change-coupling points, and gotchas: +Before editing a path, read every `AGENTS.md` from the repository root down to +that path. The root guide always applies; a nested guide adds or overrides rules +for its subtree. Changes spanning multiple subtrees must follow every applicable +guide. +`AGENTS.md` files are maps, not manuals. Keep them short and stable: record +where code lives, important boundaries or change-coupling, high-cost gotchas, +and the validation entry point. Implementation details belong in code. Do not +copy exhaustive test lists or workflow logic; link to the source of truth. + +High-change areas and specialized test suites carry their own `AGENTS.md` with +a module map, read/write paths, change-coupling points, and gotchas: + +- [`src/common/meta/AGENTS.md`](src/common/meta/AGENTS.md) +- [`src/query/AGENTS.md`](src/query/AGENTS.md) +- [`src/servers/AGENTS.md`](src/servers/AGENTS.md) +- [`src/operator/AGENTS.md`](src/operator/AGENTS.md) - [`src/mito2/AGENTS.md`](src/mito2/AGENTS.md) - [`src/metric-engine/AGENTS.md`](src/metric-engine/AGENTS.md) - [`src/flow/AGENTS.md`](src/flow/AGENTS.md) - [`src/frontend/AGENTS.md`](src/frontend/AGENTS.md) - [`src/meta-srv/AGENTS.md`](src/meta-srv/AGENTS.md) +- [`tests/compatibility/AGENTS.md`](tests/compatibility/AGENTS.md) +- [`tests/perf/AGENTS.md`](tests/perf/AGENTS.md) ## Read before changing code @@ -71,11 +88,35 @@ change-coupling points, and gotchas: - Error handling: `src/common/error/` (`ErrorExt`, `StatusCode`) - Protocol implementations: `src/servers/src/` +## Worktree safety + +- Check `git status --short` before editing. Preserve unrelated tracked changes + and untracked files; re-read files changed by another process. +- Do not rewrite history, force-push, or remove files in bulk unless the user + explicitly requests it. +- Update generated artifacts only through the generators documented in + [`.agents/generated-files.md`](.agents/generated-files.md). + +## Validation by change type + +Use the narrowest command that covers the change, then expand only when its +blast radius requires it. + +| Change | Minimum validation | +| --- | --- | +| One Rust crate | `cargo nextest run -p ` | +| Cross-workspace Rust behavior | `make test`; inspect `.github/workflows/rust.yml` when CI parity matters | +| SQL parsing, planning, execution, or output | `cargo sqlness bare -t `; inspect regenerated `.result` files | +| Persisted metadata or wire format | Add/run a case under `tests/compatibility/`; follow its `README.md` and `AGENTS.md` | +| Public configuration | Update example TOMLs, loading/serialization snapshots, and docs; run `make config-docs` | +| Query regression harness or DSL | Follow `tests/perf/AGENTS.md` | +| Enterprise-gated code | Build/test with `--features enterprise` where applicable and run `make check-enterprise-license` | + ## Before opening a PR 1. `make fmt` 2. `make clippy` -3. `make test` (or `cargo nextest run`) +3. `make test` 4. `make check-udeps` (run `make fix-udeps` if it reports unused dependencies). 5. If you added or changed a public configuration option, update the applicable example TOMLs, configuration-loading and serialized-config snapshot tests, diff --git a/README.md b/README.md index 019a7c848ba..21980b1a22b 100644 --- a/README.md +++ b/README.md @@ -203,13 +203,13 @@ See the [Contribution Guidelines](CONTRIBUTING.md) for the full developer workfl - **Kubernetes**: [GreptimeDB Operator](https://github.com/GreptimeTeam/greptimedb-operator) - **Helm Charts**: [Greptime Helm Charts](https://github.com/GreptimeTeam/helm-charts) - **Dashboard**: [Web UI](https://github.com/GreptimeTeam/dashboard) -- **gRPC Ingester**: [Go](https://github.com/GreptimeTeam/greptimedb-ingester-go), [Java](https://github.com/GreptimeTeam/greptimedb-ingester-java), [C++](https://github.com/GreptimeTeam/greptimedb-ingester-cpp), [Erlang](https://github.com/GreptimeTeam/greptimedb-ingester-erl), [Rust](https://github.com/GreptimeTeam/greptimedb-ingester-rust), [.NET](https://github.com/GreptimeTeam/greptimedb-ingester-dotnet) +- **gRPC Ingester**: [Go](https://github.com/GreptimeTeam/greptimedb-ingester-go), [Java](https://github.com/GreptimeTeam/greptimedb-ingester-java), [C++](https://github.com/GreptimeTeam/greptimedb-ingester-cpp), [Erlang](https://github.com/GreptimeTeam/greptimedb-ingester-erl), [Rust](https://github.com/GreptimeTeam/greptimedb-ingester-rust), [.NET](https://github.com/GreptimeTeam/greptimedb-ingester-dotnet), [TypeScript](https://github.com/GreptimeTeam/greptimedb-ingester-ts) - **Grafana Data Source**: [GreptimeDB Grafana data source plugin](https://github.com/GreptimeTeam/greptimedb-grafana-datasource) - **Grafana Dashboard**: [Official Dashboard for monitoring](https://github.com/GreptimeTeam/greptimedb/blob/main/grafana/README.md) ## Project Status -GreptimeDB is at [v1.0 GA](https://github.com/GreptimeTeam/greptimedb/releases/tag/v1.0.0) with stable APIs and regular releases. It runs in production at scale — [OceanBase Cloud](https://greptime.com/blogs/2025-07-22-user-case-obcloud-log-management-greptimedb) operates 80+ GreptimeDB clusters managing 300 TB of logs, cutting log storage cost by 60% after migrating from Grafana Loki. See more in [case studies](https://greptime.com/blogs/?category=Use%20Case). +GreptimeDB is generally available, with stable APIs and regular releases. It runs in production at scale — [OceanBase Cloud](https://greptime.com/blogs/2025-07-22-user-case-obcloud-log-management-greptimedb) operates 80+ GreptimeDB clusters managing 300 TB of logs, cutting log storage cost by 60% after migrating from Grafana Loki. See more in [case studies](https://greptime.com/blogs/?category=Use%20Case). Read the [v1.0 highlights](https://greptime.com/blogs/2025-11-05-greptimedb-v1-highlights) and [2026 roadmap](https://greptime.com/blogs/2026-02-11-greptimedb-roadmap-2026), or browse the [version reference](https://docs.greptime.com/nightly/reference/about-greptimedb-version). diff --git a/src/common/meta/AGENTS.md b/src/common/meta/AGENTS.md new file mode 100644 index 00000000000..fd5da08a933 --- /dev/null +++ b/src/common/meta/AGENTS.md @@ -0,0 +1,38 @@ +# common-meta — Agent & Contributor Guide + +Navigation map for `src/common/meta`. Repo-wide invariants: +[`.agents/architecture-invariants.md`](../../../.agents/architecture-invariants.md). + +`common-meta` contains shared metadata keys and values, KV backends, caches, +and durable DDL procedures. Metasrv process and service wiring live in +`src/meta-srv`. + +## Module map + +| Area | Path | Entry point | +| --- | --- | --- | +| Metadata model | `src/common/meta/src/key.rs`, `src/common/meta/src/key/` | Typed keys, values, and `TableMetadataManager` | +| KV storage | `src/common/meta/src/kv_backend.rs`, `src/common/meta/src/kv_backend/` | `KvBackend`, transactions, memory/etcd/RDS backends | +| DDL procedures | `src/common/meta/src/ddl.rs`, `src/common/meta/src/ddl/`, `src/common/meta/src/ddl_manager.rs` | Durable DDL state machines and task dispatch | +| Caches | `src/common/meta/src/cache.rs`, `src/common/meta/src/cache/`, `src/common/meta/src/cache_invalidator.rs` | Metadata caches and invalidation | +| RPC types | `src/common/meta/src/rpc.rs`, `src/common/meta/src/rpc/` | Shared metasrv request/response types | +| Recovery | `src/common/meta/src/reconciliation.rs`, `src/common/meta/src/reconciliation/` | Catalog/table/region reconciliation | + +## Change coupling + +- Key/value encoding changes require backward-compatible decoding and a case in + `tests/compatibility/`. +- Multi-key metadata changes must preserve transaction boundaries and cache + invalidation. +- Procedure changes must preserve persisted state and `TYPE_NAME`; new + procedures need loader registration in `src/common/meta/src/ddl_manager.rs`. +- `KvBackend` behavior changes should extend the shared tests in + `src/common/meta/src/kv_backend/test.rs`. + +## Testing + +```bash +cargo nextest run -p common-meta +``` + +Backend-specific coverage may also need `pg_kvbackend` or `mysql_kvbackend`. diff --git a/src/flow/AGENTS.md b/src/flow/AGENTS.md index aba019e02a0..0ac0aa0e243 100644 --- a/src/flow/AGENTS.md +++ b/src/flow/AGENTS.md @@ -8,15 +8,19 @@ Repo-wide rules that apply here: [`.agents/architecture-invariants.md`](../../.a ## What this crate does Flownode is the stream-processing engine behind continuous aggregation / -materialized views. It runs in two modes: +materialized views. It has two execution paths: -- **Batching mode** (the default and the actively developed path): splits data +- **Batching mode** (the actively developed path): splits data into time windows and periodically runs aggregation SQL through the frontend, writing results back to a sink table. -- **Streaming mode** (the older dataflow engine): an incremental DFIR/dataflow +- **Streaming mode** (the legacy dataflow path): an incremental DFIR/dataflow compute graph that processes row-level diffs. -A flow without an explicit `flow_type` is created as **batching**. +Users cannot select a mode directly: `flow_type` is a reserved internal option. +`StatementExecutor::determine_flow_type` in `src/operator/src/statement/ddl.rs` +owns current mode selection. `FlowDualEngine` defaults missing internal +`flow_type` metadata to batching for compatibility. Read those paths before +changing routing rules. ## Module map @@ -64,7 +68,9 @@ Flow metadata lives in `common-meta`, not here: `compute/render.rs` (streaming) and ensure batching SQL handles it. - **Persisted flow metadata**: keep `FlowInfoValue` backward compatible (`serde(default)` / `serde(alias)`). -- A change to one engine often needs the mirror change in the other. +- **Shared routing, metadata, or sink contracts**: check both engines. An + engine-specific implementation change does not automatically need a mirror + change in the other path. ## Testing @@ -76,9 +82,8 @@ Helpers in `src/flow/src/test_utils.rs` (test context, test query engine). ## Gotchas -- Batching vs streaming differ a lot in latency, debuggability, and code path — - confirm which mode a flow uses before reasoning about it. Batching is the - default and the primary target. +- Batching vs streaming differ in latency, state, and execution. Confirm the + selected mode before reasoning about a flow. - Streaming workers are `!Send`; cross-thread interaction goes through `WorkerHandle`, not the worker directly. - Internal flow timestamps (`repr::Timestamp`, ms) are not necessarily the diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index 2321b657327..3afd75634d0 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -36,12 +36,9 @@ distributed mode it talks to remote datanodes via `operator`/`client`. ## Request lifecycles -- **SQL query** (`instance.rs`): `do_query` → `pre_parsing` interceptor → parse → - `post_parsing` interceptor → then per statement: `check_permission` → - `statement_executor.plan` (logical plan) → `query_engine.execute` → for - distributed reads, `region_query.rs` fetches from datanodes → cancellable - `RecordBatch` stream. Interceptors run around parsing, before the per-statement - permission check — preserve that ordering. +- **SQL query** (`instance.rs`): `do_query_inner` handles parsing, interceptors, + permission checks, timeout/cancellation, and delegates planning/execution to + `StatementExecutor`. Distributed scans enter through `region_query.rs`. - **Insert** (`instance/grpc.rs`): `handle_inserts` / `handle_row_inserts` → `check_permission` → `operator`'s `Inserter` (schema validation, optional auto-create, partition routing) → local `RegionServer` (standalone) or RPC to diff --git a/src/metric-engine/AGENTS.md b/src/metric-engine/AGENTS.md index 4471de3faee..d1049478978 100644 --- a/src/metric-engine/AGENTS.md +++ b/src/metric-engine/AGENTS.md @@ -30,7 +30,8 @@ The architecture is documented at the top of `src/metric-engine/src/lib.rs`. | `state` | `src/metric-engine/src/engine/state.rs` | In-memory cache of physical columns and logical column metadata | | `repeated_task` | `src/metric-engine/src/repeated_task.rs` | Periodic metadata-region flush task | | `utils` | `src/metric-engine/src/utils.rs` | `RegionId` conversions (data vs metadata group), manifest encoding | -| `config` | `src/metric-engine/src/config.rs` | `EngineConfig` (metadata flush interval, sparse PK) | +| `config` | `src/metric-engine/src/config.rs` | `EngineConfig` (metadata flush interval) | +| `engine/options` | `src/metric-engine/src/engine/options.rs` | Physical data-region options, including the default sparse primary-key encoding | | `test_util` | `src/metric-engine/src/test_util.rs` | `TestEnv` building the Mito2 + Metric stack | ## Write path @@ -77,8 +78,9 @@ cargo nextest run -p metric-engine writes; operate on logical region ids. - TSID must be stable for the same tag set — it is a hash over sorted tag names + values and may be stored in `__tsid` or encoded into `__primary_key`. -- Metadata is cached (LRU with a TTL); after an alter, stale reads are possible - until invalidation/expiry. +- Metadata reads use an LRU cache. Metadata writes take the per-region cache + write lock and invalidate the cache after the Mito2 write succeeds; preserve + that ordering when adding a metadata mutation path. - Always convert ids via `utils::to_data_region_id` / `to_metadata_region_id`. ## Maintenance contract diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md new file mode 100644 index 00000000000..761d8e5080a --- /dev/null +++ b/src/operator/AGENTS.md @@ -0,0 +1,40 @@ +# operator — Agent & Contributor Guide + +Navigation map for `src/operator`. Repo-wide invariants: +[`.agents/architecture-invariants.md`](../../.agents/architecture-invariants.md). + +`operator` connects frontend statements and writes to query, catalog, metasrv, +and datanodes. It owns statement dispatch, DDL task conversion, and region-level +insert/delete routing. + +## Module map + +| Area | Path | Entry point | +| --- | --- | --- | +| Statements | `src/operator/src/statement.rs`, `src/operator/src/statement/` | `StatementExecutor` and per-statement handlers | +| Inserts | `src/operator/src/insert.rs`, `src/operator/src/bulk_insert.rs` | Insert, auto-create/alter, routing, dispatch | +| Deletes | `src/operator/src/delete.rs` | Delete conversion and dispatch | +| Request conversion | `src/operator/src/req_convert.rs`, `src/operator/src/req_convert/` | Table/row/column requests to region requests | +| Region requests | `src/operator/src/region_req_factory.rs` | Region-level DDL/DML request construction | +| Procedures | `src/operator/src/procedure.rs` | Metasrv procedure administration | +| Flow | `src/operator/src/flow.rs`, `src/operator/src/statement/ddl.rs` | Flow requests and internal mode selection | + +## Change coupling + +- New statement kinds usually touch parser/AST handling, frontend permission + and interception, executor dispatch, and sqlness cases. +- DDL task changes require matching `common-meta`, metasrv, recovery, and + compatibility updates. +- Insert/delete conversion changes must cover both request formats, partition + routing, and schema-on-write behavior. +- Flow creation changes must stay aligned with `src/flow/AGENTS.md` and + `common-meta` flow metadata. + +## Testing + +```bash +cargo nextest run -p operator +cargo sqlness bare -t +``` + +Check standalone and distributed routing when changing region dispatch. diff --git a/src/query/AGENTS.md b/src/query/AGENTS.md new file mode 100644 index 00000000000..8dc03a22877 --- /dev/null +++ b/src/query/AGENTS.md @@ -0,0 +1,39 @@ +# query — Agent & Contributor Guide + +Navigation map for `src/query`. Repo-wide invariants: +[`.agents/architecture-invariants.md`](../../.agents/architecture-invariants.md). + +`query` owns SQL/PromQL/log planning, GreptimeDB optimizer rules, DataFusion +execution, and distributed query plans. Protocols live in `servers`; statement +side effects live in `operator`. + +## Module map + +| Area | Path | Entry point | +| --- | --- | --- | +| Logical planning | `src/query/src/planner.rs`, `src/query/src/parser.rs` | `LogicalPlanner`, `QueryStatement` | +| Query engine | `src/query/src/query_engine.rs`, `src/query/src/query_engine/` | Engine traits, factory, state, serialization | +| DataFusion execution | `src/query/src/datafusion.rs`, `src/query/src/datafusion/` | Physical planning and execution | +| Optimizers | `src/query/src/optimizer.rs`, `src/query/src/optimizer/` | GreptimeDB logical/physical rules | +| Distributed plans | `src/query/src/dist_plan.rs`, `src/query/src/dist_plan/` | `MergeScan`, pruning, merge/sort, remote filters | +| Remote reads | `src/query/src/region_query.rs` | `RegionQueryHandler` boundary | +| Query languages | `src/query/src/promql/`, `src/query/src/log_query/` | PromQL and log planners | + +## Change coupling + +- Optimizer or `MergeScan` changes usually touch rule ordering, physical + planning, serialization, and distributed-plan tests together. +- User-visible SQL/PromQL/log behavior needs planner coverage and a sqlness case. +- `RegionQueryHandler` changes require the frontend implementation to move with + it. +- DataFusion dependency changes follow the pinned-fork rule in the repo-wide + invariants. + +## Testing + +```bash +cargo nextest run -p query +cargo sqlness bare -t +``` + +Check both standalone and distributed plans when changing `src/query/src/dist_plan/`. diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md new file mode 100644 index 00000000000..2a446033f62 --- /dev/null +++ b/src/servers/AGENTS.md @@ -0,0 +1,37 @@ +# servers — Agent & Contributor Guide + +Navigation map for `src/servers`. Repo-wide invariants: +[`.agents/architecture-invariants.md`](../../.agents/architecture-invariants.md). + +`servers` owns network listeners and protocol translation. Database behavior is +provided through handler traits, mostly implemented by `frontend`. + +## Module map + +| Area | Path | Entry point | +| --- | --- | --- | +| Lifecycle | `src/servers/src/server.rs` | Shared `Server` trait | +| Handler boundary | `src/servers/src/query_handler.rs`, `src/servers/src/query_handler/` | Protocol-independent handler traits | +| gRPC / Flight | `src/servers/src/grpc.rs`, `src/servers/src/grpc/` | gRPC services, auth, cancellation, builders | +| HTTP | `src/servers/src/http.rs`, `src/servers/src/http/` | Routes, middleware, protocol endpoints, output formats | +| SQL protocols | `src/servers/src/mysql.rs`, `src/servers/src/mysql/`, `src/servers/src/postgres.rs`, `src/servers/src/postgres/` | MySQL/PostgreSQL servers and type conversion | +| Observability protocols | `src/servers/src/prom_remote_write/`, `src/servers/src/prometheus.rs`, `src/servers/src/otlp/`, `src/servers/src/otel_arrow.rs` | Prometheus and OpenTelemetry conversion | +| Resources / TLS | `src/servers/src/request_memory_limiter.rs`, `src/servers/src/tls.rs`, `src/servers/src/addrs.rs` | Memory admission, TLS, listener addresses | + +## Change coupling + +- Handler trait changes require matching frontend and mock implementations. +- gRPC wire changes start in `greptime-proto`; update builders and handlers here + after bumping it. +- Auth or `QueryContext` changes must cover HTTP, gRPC, MySQL, and PostgreSQL + entry points. +- New protocols or externally visible routes also require frontend service and + configuration wiring. + +## Testing + +```bash +cargo nextest run -p servers +``` + +Keep protocol translation here and permissions/database behavior in frontend. diff --git a/tests/perf/AGENTS.md b/tests/perf/AGENTS.md index e7c24f16e62..1f13dea58dc 100644 --- a/tests/perf/AGENTS.md +++ b/tests/perf/AGENTS.md @@ -19,9 +19,10 @@ - Keep the direct-SST generator generic. Issue-specific behavior belongs in case files and thresholds, not in Rust generator logic. - Before pushing perf harness changes, run at least: - - `uv run --no-project python -m py_compile .github/scripts/query-regression-run.py .github/scripts/query-regression-summary.py .github/scripts/query-regression-pr-metadata.py tests/perf/test_query_regression_runner_compaction_toctou.py tests/perf/test_query_regression_runner_otlp_trace_load.py` - - `uv run --no-project python tests/perf/test_query_regression_runner_compaction_toctou.py && uv run --no-project python tests/perf/test_query_regression_runner_otlp_trace_load.py` + - the Python tests in the `Test query regression tooling` step of + `.github/workflows/query-regression.yml` - `cargo fmt --all -- --check` - `cargo build -p cmd --bin query_perf_fixture --features dev-tools` + - `cargo build -p cmd --bin query_regression_runner --features dev-tools` - exercise the outer lifecycle script and Rust fixture generator against all built-in cases when the DSL or workflow case selection changes. From 606e5fd6a0c2389321eab6e06ef6abd98aeff683 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 10 Aug 2026 11:42:53 +0800 Subject: [PATCH 025/126] feat: update opentelemetry family to 0.32 series (#8776) * feat: update openetelemtry family to 0.32 series * chore: resolve warning * fix: update tests --- Cargo.lock | 175 ++++++------------ Cargo.toml | 4 +- src/common/telemetry/Cargo.toml | 6 +- src/common/telemetry/src/tracing_context.rs | 2 +- src/common/telemetry/src/tracing_sampler.rs | 6 +- src/frontend/src/instance/otlp.rs | 1 + .../src/instance/otlp/trace_ingest/tests.rs | 1 + src/mito2/src/gc.rs | 2 +- src/servers/src/otlp/logs.rs | 12 ++ src/servers/src/otlp/trace/attributes.rs | 12 ++ src/servers/src/otlp/trace/span.rs | 1 + src/servers/src/otlp/trace/v1.rs | 7 + src/servers/src/otlp/utils.rs | 11 ++ 13 files changed, 112 insertions(+), 128 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 476133db9e7..841cc035f77 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2262,7 +2262,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -2897,10 +2897,10 @@ dependencies = [ "humantime-serde", "lazy_static", "once_cell", - "opentelemetry 0.30.0", + "opentelemetry", "opentelemetry-otlp", "opentelemetry-semantic-conventions 0.30.0", - "opentelemetry_sdk 0.30.0", + "opentelemetry_sdk", "parking_lot 0.12.4", "prometheus 0.14.0", "serde", @@ -5599,7 +5599,7 @@ dependencies = [ "meta-client", "meta-srv", "num_cpus", - "opentelemetry-proto 0.31.0", + "opentelemetry-proto", "opentelemetry-semantic-conventions 0.31.0", "operator", "otel-arrow-rust", @@ -6662,7 +6662,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.57.0", + "windows-core 0.61.2", ] [[package]] @@ -7831,7 +7831,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.48.5", + "windows-targets 0.53.5", ] [[package]] @@ -9795,23 +9795,9 @@ checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e" [[package]] name = "opentelemetry" -version = "0.30.0" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aaf416e4cb72756655126f7dd7bb0af49c674f4c1b9903e80c009e0c37e552e6" -dependencies = [ - "futures-core", - "futures-sink", - "js-sys", - "pin-project-lite", - "thiserror 2.0.17", - "tracing", -] - -[[package]] -name = "opentelemetry" -version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b84bcd6ae87133e903af7ef497404dda70c60d0ea14895fc8a5e6722754fc2a0" +checksum = "b0142c63252a9e054e68a4c61a5778f7b14f576274d593f8ce883d191a099682" dependencies = [ "futures-core", "futures-sink", @@ -9823,61 +9809,48 @@ dependencies = [ [[package]] name = "opentelemetry-http" -version = "0.30.0" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50f6639e842a97dbea8886e3439710ae463120091e2e064518ba8e716e6ac36d" +checksum = "5683015d09e2df236ef005b17f6f196f0d5f6313c4fa43a7b6a53b52776e4331" dependencies = [ "async-trait", "bytes", "http 1.3.1", - "opentelemetry 0.30.0", - "reqwest 0.12.24", + "opentelemetry", + "reqwest 0.13.4", ] [[package]] name = "opentelemetry-otlp" -version = "0.30.0" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dbee664a43e07615731afc539ca60c6d9f1a9425e25ca09c57bc36c87c55852b" +checksum = "9966929966d17620d7c316c643ba62631826e10021409357772d5eea84f62c35" dependencies = [ "http 1.3.1", - "opentelemetry 0.30.0", + "opentelemetry", "opentelemetry-http", - "opentelemetry-proto 0.30.0", - "opentelemetry_sdk 0.30.0", - "prost 0.13.5", - "reqwest 0.12.24", + "opentelemetry-proto", + "opentelemetry_sdk", + "prost 0.14.1", + "reqwest 0.13.4", "thiserror 2.0.17", "tokio", - "tonic 0.13.1", - "tracing", + "tonic 0.14.2", + "tonic-types", ] [[package]] name = "opentelemetry-proto" -version = "0.30.0" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e046fd7660710fe5a05e8748e70d9058dc15c94ba914e7c4faa7c728f0e8ddc" -dependencies = [ - "opentelemetry 0.30.0", - "opentelemetry_sdk 0.30.0", - "prost 0.13.5", - "tonic 0.13.1", -] - -[[package]] -name = "opentelemetry-proto" -version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7175df06de5eaee9909d4805a3d07e28bb752c34cab57fa9cff549da596b30f" +checksum = "56d658ba1faf63f7b9c492cfbe6e0ec365440a16132d3270c1065f7b33f1b638" dependencies = [ "base64 0.22.1", "const-hex", - "opentelemetry 0.31.0", - "opentelemetry_sdk 0.31.0", + "opentelemetry", + "opentelemetry_sdk", "prost 0.14.1", "serde", - "serde_json", "tonic 0.14.2", "tonic-prost", ] @@ -9896,37 +9869,22 @@ checksum = "e62e29dfe041afb8ed2a6c9737ab57db4907285d999ef8ad3a59092a36bdc846" [[package]] name = "opentelemetry_sdk" -version = "0.30.0" +version = "0.32.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11f644aa9e5e31d11896e024305d7e3c98a88884d9f8919dbf37a9991bc47a4b" +checksum = "9b59f80e1ac4d5ff7a2db8fb6c80badb7f0f3f858211fba08dd9aaec750894f9" dependencies = [ "futures-channel", "futures-executor", "futures-util", - "opentelemetry 0.30.0", + "opentelemetry", "percent-encoding", + "portable-atomic", "rand 0.9.4", - "serde_json", "thiserror 2.0.17", "tokio", "tokio-stream", ] -[[package]] -name = "opentelemetry_sdk" -version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e14ae4f5991976fd48df6d843de219ca6d31b01daaab2dad5af2badeded372bd" -dependencies = [ - "futures-channel", - "futures-executor", - "futures-util", - "opentelemetry 0.31.0", - "percent-encoding", - "rand 0.9.4", - "thiserror 2.0.17", -] - [[package]] name = "operator" version = "1.3.0" @@ -11287,7 +11245,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ "heck 0.5.0", - "itertools 0.11.0", + "itertools 0.14.0", "log", "multimap", "once_cell", @@ -11335,7 +11293,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -11348,7 +11306,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -12247,7 +12205,6 @@ checksum = "9d0946410b9f7b082a427e4ef5c8ff541a88b357bc6c637c40db3a68ac70a36f" dependencies = [ "base64 0.22.1", "bytes", - "futures-channel", "futures-core", "futures-util", "h2 0.4.11", @@ -12291,6 +12248,7 @@ checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" dependencies = [ "base64 0.22.1", "bytes", + "futures-channel", "futures-core", "futures-util", "http 1.3.1", @@ -13349,7 +13307,7 @@ dependencies = [ "once_cell", "openmetrics-parser", "opensrv-mysql", - "opentelemetry-proto 0.31.0", + "opentelemetry-proto", "operator", "otel-arrow-rust", "parking_lot 0.12.4", @@ -14857,7 +14815,7 @@ dependencies = [ "mito2", "mysql_async", "object-store", - "opentelemetry-proto 0.31.0", + "opentelemetry-proto", "operator", "otel-arrow-rust", "partition", @@ -15318,32 +15276,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "tonic" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e581ba15a835f4d9ea06c55ab1bd4dce26fc53752c69a04aac00703bfb49ba9" -dependencies = [ - "async-trait", - "base64 0.22.1", - "bytes", - "http 1.3.1", - "http-body 1.0.1", - "http-body-util", - "hyper 1.6.0", - "hyper-timeout 0.5.2", - "hyper-util", - "percent-encoding", - "pin-project", - "prost 0.13.5", - "tokio", - "tokio-stream", - "tower 0.5.2", - "tower-layer", - "tower-service", - "tracing", -] - [[package]] name = "tonic" version = "0.14.2" @@ -15430,6 +15362,17 @@ dependencies = [ "tonic-prost", ] +[[package]] +name = "tonic-types" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a875a902255423d34c1f20838ab374126db8eb41625b7947a1d54113b0b7399" +dependencies = [ + "prost 0.14.1", + "prost-types 0.14.1", + "tonic 0.14.2", +] + [[package]] name = "tonic-web" version = "0.14.5" @@ -15552,9 +15495,9 @@ checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" [[package]] name = "tracing" -version = "0.1.41" +version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "784e0ac535deb450455cbfa28a6f0df145ea1bb7ae51b821cf5e7927fdcfbdd0" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "log", "pin-project-lite", @@ -15576,9 +15519,9 @@ dependencies = [ [[package]] name = "tracing-attributes" -version = "0.1.28" +version = "0.1.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "395ae124c09f9e6918a2310af6038fba074bcf474ac352496d5910dd59a2226d" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", @@ -15587,9 +15530,9 @@ dependencies = [ [[package]] name = "tracing-core" -version = "0.1.34" +version = "0.1.36" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d12581f227e93f094d3af2ae690a574abb8a2b9b7a96e7cfe9647b2b617678" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", "valuable", @@ -15608,14 +15551,12 @@ dependencies = [ [[package]] name = "tracing-opentelemetry" -version = "0.31.0" +version = "0.33.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddcf5959f39507d0d04d6413119c04f33b623f4f951ebcbdddddfad2d0623a9c" +checksum = "adbc64cba7137545b8044cb1fe9814f7aacf3c6b5f9b45be8bb5db538befdb26" dependencies = [ "js-sys", - "once_cell", - "opentelemetry 0.30.0", - "opentelemetry_sdk 0.30.0", + "opentelemetry", "smallvec", "tracing", "tracing-core", @@ -15636,9 +15577,9 @@ dependencies = [ [[package]] name = "tracing-subscriber" -version = "0.3.20" +version = "0.3.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2054a14f5307d601f88daf0553e1cbf472acc4f2c51afab632431cdcd72124d5" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" dependencies = [ "matchers", "nu-ansi-term", @@ -16462,7 +16403,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 9582b12bd85..c9b199b3673 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -182,7 +182,7 @@ notify = "8.0" num_cpus = "1.16" object_store_opendal = "0.57" once_cell = "1.18" -opentelemetry-proto = { version = "0.31", features = [ +opentelemetry-proto = { version = "0.32", features = [ "gen-tonic", "metrics", "trace", @@ -260,7 +260,7 @@ tower = "0.5" tower-http = "0.6" tracing = "0.1" tracing-appender = "0.2" -tracing-opentelemetry = "0.31.0" +tracing-opentelemetry = "0.33.0" tracing-subscriber = { version = "0.3", features = ["env-filter", "json", "fmt"] } typetag = "0.2" uuid = { version = "1.17", features = ["serde", "v4", "v7", "fast-rng"] } diff --git a/src/common/telemetry/Cargo.toml b/src/common/telemetry/Cargo.toml index 92c3304d53f..ae25a43c5dc 100644 --- a/src/common/telemetry/Cargo.toml +++ b/src/common/telemetry/Cargo.toml @@ -21,12 +21,12 @@ greptime-proto.workspace = true humantime-serde.workspace = true lazy_static.workspace = true once_cell.workspace = true -opentelemetry = { version = "0.30.0", default-features = false, features = [ +opentelemetry = { version = "0.32.0", default-features = false, features = [ "trace", ] } -opentelemetry-otlp = { version = "0.30.0", features = ["trace", "grpc-tonic", "http-proto"] } +opentelemetry-otlp = { version = "0.32.0", features = ["trace", "grpc-tonic", "http-proto"] } opentelemetry-semantic-conventions = { version = "0.30.0", features = ["semconv_experimental"] } -opentelemetry_sdk = { version = "0.30.0", features = ["rt-tokio", "trace"] } +opentelemetry_sdk = { version = "0.32.1", features = ["rt-tokio", "trace"] } parking_lot.workspace = true prometheus.workspace = true serde.workspace = true diff --git a/src/common/telemetry/src/tracing_context.rs b/src/common/telemetry/src/tracing_context.rs index fd389bccf59..e427b40effd 100644 --- a/src/common/telemetry/src/tracing_context.rs +++ b/src/common/telemetry/src/tracing_context.rs @@ -73,7 +73,7 @@ impl TracingContext { /// Attach the given span as a child of the context. Returns the attached span. pub fn attach(&self, span: tracing::Span) -> tracing::Span { - span.set_parent(self.0.clone()); + let _ = span.set_parent(self.0.clone()); span } diff --git a/src/common/telemetry/src/tracing_sampler.rs b/src/common/telemetry/src/tracing_sampler.rs index ccd029b781b..facc0e5d40f 100644 --- a/src/common/telemetry/src/tracing_sampler.rs +++ b/src/common/telemetry/src/tracing_sampler.rs @@ -15,10 +15,8 @@ use std::collections::HashSet; use opentelemetry::KeyValue; -use opentelemetry::trace::{ - Link, SamplingDecision, SamplingResult, SpanKind, TraceContextExt, TraceId, TraceState, -}; -use opentelemetry_sdk::trace::{Sampler, ShouldSample}; +use opentelemetry::trace::{Link, SpanKind, TraceContextExt, TraceId, TraceState}; +use opentelemetry_sdk::trace::{Sampler, SamplingDecision, SamplingResult, ShouldSample}; use serde::{Deserialize, Serialize}; #[derive(Clone, Debug, Serialize, Deserialize)] diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index dc31416dfe4..c4616ae578d 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -291,6 +291,7 @@ mod tests { value: Some(AnyValue { value: Some(any_value::Value::StringValue("frontend".to_string())), }), + ..Default::default() }], ..Default::default() }), diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index 949ab1535ac..0f18a7194be 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -234,6 +234,7 @@ fn trace_attribute(key: &str, value: OtlpValue) -> KeyValue { KeyValue { key: key.to_string(), value: Some(AnyValue { value: Some(value) }), + ..Default::default() } } diff --git a/src/mito2/src/gc.rs b/src/mito2/src/gc.rs index babebabb741..a5474099cbf 100644 --- a/src/mito2/src/gc.rs +++ b/src/mito2/src/gc.rs @@ -410,7 +410,7 @@ impl LocalGcWorker { region_present = region.is_some() ) )] - pub async fn do_region_gc( + pub(crate) async fn do_region_gc( &self, region_id: RegionId, region: Option, diff --git a/src/servers/src/otlp/logs.rs b/src/servers/src/otlp/logs.rs index a956c7977c8..33d4f4eb5e5 100644 --- a/src/servers/src/otlp/logs.rs +++ b/src/servers/src/otlp/logs.rs @@ -24,6 +24,7 @@ use api::v1::{ RowInsertRequest, Rows, SemanticType, Value as GreptimeValue, }; use bytes::Bytes; +use common_telemetry::warn; use common_time::Timestamp; use common_time::timestamp::TimeUnit; use jsonb::{Number as JsonbNumber, Value as JsonbValue}; @@ -1031,6 +1032,16 @@ fn any_value_to_vrl_value(value: any_value::Value) -> VrlValue { VrlValue::Object(key_value_to_map(key_value_list.values)) } any_value::Value::BytesValue(items) => VrlValue::Bytes(Bytes::from(items)), + // `StringValueStrindex` references the Profiling signal's + // `ProfilesDictionary.string_table`, which is unavailable to logs/traces. + // Per the OTLP spec, non-Profiling receivers must treat it as a non-fatal + // issue and process the value as if it were absent. + any_value::Value::StringValueStrindex(_) => { + warn!( + "encountered a profiling-only `StringValueStrindex` value in a non-profiling signal; ignoring" + ); + VrlValue::Null + } } } @@ -1082,6 +1093,7 @@ mod tests { KeyValue { key: key.to_string(), value: Some(AnyValue { value: Some(value) }), + ..Default::default() } } diff --git a/src/servers/src/otlp/trace/attributes.rs b/src/servers/src/otlp/trace/attributes.rs index 08554863432..471832b8920 100644 --- a/src/servers/src/otlp/trace/attributes.rs +++ b/src/servers/src/otlp/trace/attributes.rs @@ -16,6 +16,7 @@ use std::fmt::Display; use opentelemetry_proto::tonic::common::v1::any_value::Value::{ ArrayValue, BoolValue, BytesValue, DoubleValue, IntValue, KvlistValue, StringValue, + StringValueStrindex, }; use opentelemetry_proto::tonic::common::v1::{AnyValue, KeyValue}; use serde::Serialize; @@ -78,6 +79,11 @@ impl Serialize for OtlpAnyValue<'_> { map.end() } BytesValue(v) => zer.serialize_bytes(v), + // `StringValueStrindex` is profiling-signal-only and references the + // Profiling `ProfilesDictionary.string_table`, which is unavailable + // here. Per the OTLP spec, non-Profiling receivers must treat it as + // a non-fatal issue and process the value as if it were absent. + StringValueStrindex(_) => zer.serialize_none(), }, None => zer.serialize_none(), } @@ -293,6 +299,7 @@ mod tests { vec![KeyValue { key: "key1".into(), value: None, + ..Default::default() }], ), ( @@ -300,6 +307,7 @@ mod tests { vec![KeyValue { key: "key1".into(), value: Some(AnyValue { value: None }), + ..Default::default() }], ), ( @@ -309,6 +317,7 @@ mod tests { value: Some(AnyValue { value: Some(Value::StringValue(String::from("val1"))), }), + ..Default::default() }], ), ]; @@ -331,6 +340,7 @@ mod tests { vec![KeyValue { key: "key1".into(), value: None, + ..Default::default() }], ), ( @@ -338,6 +348,7 @@ mod tests { vec![KeyValue { key: "key1".into(), value: Some(AnyValue { value: None }), + ..Default::default() }], ), ( @@ -347,6 +358,7 @@ mod tests { value: Some(AnyValue { value: Some(Value::StringValue(String::from("val1"))), }), + ..Default::default() }], ), ]; diff --git a/src/servers/src/otlp/trace/span.rs b/src/servers/src/otlp/trace/span.rs index 19103240f62..fe8b37fed0a 100644 --- a/src/servers/src/otlp/trace/span.rs +++ b/src/servers/src/otlp/trace/span.rs @@ -319,6 +319,7 @@ mod tests { value: Some(AnyValue { value: Some(any_value::Value::StringValue(value.to_string())), }), + ..Default::default() } } diff --git a/src/servers/src/otlp/trace/v1.rs b/src/servers/src/otlp/trace/v1.rs index eb1612eada6..8ddeac24f33 100644 --- a/src/servers/src/otlp/trace/v1.rs +++ b/src/servers/src/otlp/trace/v1.rs @@ -633,6 +633,12 @@ fn write_attributes_with_schema( row, ); } + // `StringValueStrindex` is profiling-signal-only and references the + // Profiling `ProfilesDictionary.string_table`, which is unavailable to + // traces. Per the OTLP spec, non-Profiling receivers must treat it as a + // non-fatal issue and process the value as if it were absent. Like the + // `None` arm, no field is written for the attribute. + Some(OtlpValue::StringValueStrindex(_)) => {} None => {} } } @@ -656,6 +662,7 @@ mod tests { KeyValue { key: key.to_string(), value: Some(AnyValue { value: Some(value) }), + ..Default::default() } } diff --git a/src/servers/src/otlp/utils.rs b/src/servers/src/otlp/utils.rs index d69f8896d38..5ae5fe6ef50 100644 --- a/src/servers/src/otlp/utils.rs +++ b/src/servers/src/otlp/utils.rs @@ -14,6 +14,7 @@ use api::v1::ColumnDataType; use api::v1::value::ValueData; +use common_telemetry::warn; use jsonb::{Number as JsonbNumber, Value as JsonbValue}; use opentelemetry_proto::tonic::common::v1::{KeyValue, any_value}; @@ -40,6 +41,16 @@ pub fn any_value_to_jsonb(value: any_value::Value) -> JsonbValue<'static> { } any_value::Value::KvlistValue(kv) => key_value_to_jsonb(kv.values), any_value::Value::BytesValue(b) => JsonbValue::String(bytes_to_hex_string(&b).into()), + // `StringValueStrindex` references the Profiling signal's + // `ProfilesDictionary.string_table`, which is unavailable to logs/traces. + // Per the OTLP spec, non-Profiling receivers must treat it as a non-fatal + // issue and process the value as if it were absent. + any_value::Value::StringValueStrindex(_) => { + warn!( + "encountered a profiling-only `StringValueStrindex` value in a non-profiling signal; ignoring" + ); + JsonbValue::Null + } } } From 42c94211a2878c402d98bafd2304596a9d848ac7 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 10 Aug 2026 12:35:20 +0800 Subject: [PATCH 026/126] chore(ci): warn when a member has too many open pull requests (#8813) * chore(ci): warn when a member has too many open pull requests Review capacity is the bottleneck. Add a `pull_request_target` workflow that counts an org member's open pull requests (drafts included) on open/reopen and posts a warning comment when the count exceeds the limit. Advisory only for now: nothing is closed and no check fails. Signed-off-by: Dennis Zhuang * chore(ci): address review on the pr-open-limit workflow - Only match marker comments authored by the Actions bot; a marker pasted by anyone else would otherwise be picked up and fail the edit with 403. - Validate MAX_OPEN_PRS and fall back to 5 on a non-numeric variable. - Drop pull-requests write permission; commenting goes through the issues API. - continue-on-error so a script failure never marks the pull request red. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .github/scripts/pr-open-limit.js | 117 ++++++++++++++++++++++++++++ .github/workflows/pr-open-limit.yml | 48 ++++++++++++ 2 files changed, 165 insertions(+) create mode 100644 .github/scripts/pr-open-limit.js create mode 100644 .github/workflows/pr-open-limit.yml diff --git a/.github/scripts/pr-open-limit.js b/.github/scripts/pr-open-limit.js new file mode 100644 index 00000000000..1ba908b7450 --- /dev/null +++ b/.github/scripts/pr-open-limit.js @@ -0,0 +1,117 @@ +// Warns when a repository member has too many open pull requests. +// +// Review capacity is the bottleneck, so authors are expected to land or close +// existing work before opening more. Drafts count: they still occupy attention. +// This is advisory only — it never closes a pull request or fails the job. + +import { appendFileSync } from "node:fs"; + +const MARKER = ""; +// GITHUB_TOKEN can only edit comments authored by the Actions bot itself. +const COMMENT_AUTHOR = "github-actions[bot]"; +// Associations that identify an organization member. Outside collaborators and +// community contributors are out of scope. +const MEMBER_ASSOCIATIONS = ["OWNER", "MEMBER"]; +// Keep the comment readable for authors who are far over the limit. +const MAX_LISTED_PRS = 10; + +function summary(text) { + if (process.env.GITHUB_STEP_SUMMARY) { + appendFileSync(process.env.GITHUB_STEP_SUMMARY, `${text}\n`); + } + console.log(text); +} + +function buildComment(author, total, limit, openPrs) { + const listed = openPrs + .slice(0, MAX_LISTED_PRS) + .map((pr) => `- #${pr.number} ${pr.title}${pr.draft ? " _(draft)_" : ""}`); + const hidden = openPrs.length - listed.length; + const list = hidden > 0 ? `${listed.join("\n")}\n- ...and ${hidden} more` : listed.join("\n"); + + return `${MARKER} +> [!WARNING] +> @${author} has **${total}** open pull requests in this repository, over the limit of **${limit}**. + +Review is the scarcest resource here. Please land or close some of these before +pushing this one forward: + +${list} + +This check is advisory for now and blocks nothing.`; +} + +async function upsertComment(octokit, params, body) { + const comments = await octokit.paginate(octokit.issues.listComments, { + ...params, + per_page: 100, + }); + const existing = comments.find( + (comment) => + comment.user?.login === COMMENT_AUTHOR && comment.body?.includes(MARKER) + ); + + if (existing) { + await octokit.issues.updateComment({ + owner: params.owner, + repo: params.repo, + comment_id: existing.id, + body, + }); + } else { + await octokit.issues.createComment({ ...params, body }); + } +} + +(async () => { + const { Octokit } = await import("@octokit/rest"); + + const [owner, repo] = process.env.GITHUB_REPOSITORY.split("/"); + const prNumber = Number(process.env.PR_NUMBER); + const author = process.env.PR_AUTHOR; + const association = process.env.PR_AUTHOR_ASSOCIATION; + const rawLimit = Number(process.env.MAX_OPEN_PRS); + const limit = Number.isInteger(rawLimit) && rawLimit >= 0 ? rawLimit : 5; + + if (author.endsWith("[bot]")) { + summary(`Skipping bot author \`${author}\`.`); + return; + } + + if (!MEMBER_ASSOCIATIONS.includes(association)) { + summary(`Skipping \`${author}\`: association is \`${association}\`, not an org member.`); + return; + } + + const octokit = new Octokit({ auth: process.env.GITHUB_TOKEN }); + + const allOpen = await octokit.paginate(octokit.pulls.list, { + owner, + repo, + state: "open", + sort: "created", + direction: "asc", + per_page: 100, + }); + // The pull request that triggered this run is already open, so exclude it and + // report the total separately. + const others = allOpen.filter( + (pr) => pr.user?.login === author && pr.number !== prNumber + ); + const total = others.length + 1; + + if (others.length < limit) { + summary(`\`${author}\` has ${total} open pull requests (limit ${limit}). OK.`); + return; + } + + summary(`\`${author}\` has ${total} open pull requests, over the limit of ${limit}.`); + await upsertComment( + octokit, + { owner, repo, issue_number: prNumber }, + buildComment(author, total, limit, others) + ); +})().catch((error) => { + console.error(error); + process.exit(1); +}); diff --git a/.github/workflows/pr-open-limit.yml b/.github/workflows/pr-open-limit.yml new file mode 100644 index 00000000000..d22929911a8 --- /dev/null +++ b/.github/workflows/pr-open-limit.yml @@ -0,0 +1,48 @@ +name: PR Open Limit + +on: + pull_request_target: + types: + - opened + - reopened + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number }} + cancel-in-progress: true + +# The warning comment goes through the issues API, hence `issues: write`. +permissions: + contents: read + pull-requests: read + issues: write + +jobs: + check-open-limit: + name: Check Open Pull Request Limit + runs-on: ubuntu-latest + timeout-minutes: 10 + # Advisory: a script failure must not put a red cross on the pull request. + continue-on-error: true + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Setup Node.js + uses: actions/setup-node@v4 + with: + node-version: '20' + + - name: Install dependencies + working-directory: .github/scripts + run: npm ci + + - name: Check open pull request limit + working-directory: .github/scripts + run: node pr-open-limit.js + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + PR_NUMBER: ${{ github.event.pull_request.number }} + PR_AUTHOR: ${{ github.event.pull_request.user.login }} + PR_AUTHOR_ASSOCIATION: ${{ github.event.pull_request.author_association }} + MAX_OPEN_PRS: ${{ vars.MAX_OPEN_PRS_PER_AUTHOR || 5 }} From 0a2b10e476635ea9ffd2ac7fd21e45d48a1ecb24 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 10 Aug 2026 13:30:13 +0800 Subject: [PATCH 027/126] fix(mito2): avoid region worker panic when building a WAL entry fails (#8810) When `add_wal_entry` fails for a region, the worker only sets the error on the write context, which stays in `region_ctxs`. The region's entry is not in the batch, so a successful `write_to_wal` returns no last entry id for it and the success branch panics on `response.last_entry_ids.get(region_id).unwrap()`, killing the region worker. When the failed region is the only one in the batch, the batch is empty and `append_batch` always returns an empty response, so the panic is guaranteed. No in-tree log store can fail to build an entry at runtime today (the provider/log store combination is validated when the region opens), so this is a latent panic rather than a reachable crash. Skip contexts already marked as failed when updating next entry ids; their waiters are already notified with the error. Extract the WAL phase of `handle_write_requests` into `write_wal` and cover the failure paths with unit tests. Signed-off-by: jeremyhi --- src/mito2/src/region_write_ctx.rs | 11 + src/mito2/src/worker/handle_write.rs | 287 ++++++++++++++++++++++++--- 2 files changed, 269 insertions(+), 29 deletions(-) diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index c6abdb0872c..b71e3e5ab17 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -218,11 +218,22 @@ impl RegionWriteCtx { self.failed = true; } + /// Returns whether the write operation is already marked as failed. + pub(crate) fn is_failed(&self) -> bool { + self.failed + } + /// Updates next entry id. pub(crate) fn set_next_entry_id(&mut self, next_entry_id: EntryId) { self.next_entry_id = next_entry_id } + /// Returns the next entry id to write. + #[cfg(test)] + pub(crate) fn next_entry_id(&self) -> EntryId { + self.next_entry_id + } + /// Consumes mutations and writes them into mutable memtable. pub(crate) async fn write_memtable(&mut self) { debug_assert_eq!(self.notifiers.len(), self.wal_entry.mutations.len()); diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index c52572df114..3be40553145 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -35,6 +35,7 @@ use crate::metrics::{ use crate::region::{RegionLeaderState, RegionRoleState}; use crate::region_write_ctx::RegionWriteCtx; use crate::request::{SenderBulkRequest, SenderWriteRequest, WriteRequest}; +use crate::wal::Wal; use crate::worker::RegionWorkerLoop; impl RegionWorkerLoop { @@ -110,35 +111,9 @@ impl RegionWorkerLoop { let _timer = WRITE_STAGE_ELAPSED .with_label_values(&["write_wal"]) .start_timer(); - let mut wal_writer = self.wal.writer(); - for region_ctx in region_ctxs.values_mut() { - if region_ctx.skip_wal() { - continue; - } - if let Err(e) = region_ctx.add_wal_entry(&mut wal_writer).map_err(Arc::new) { - region_ctx.set_error(e); - } - } - match wal_writer.write_to_wal().await.map_err(Arc::new) { - Ok(response) => { - for (region_id, region_ctx) in region_ctxs.iter_mut() { - if region_ctx.skip_wal() { - continue; - } - - // Safety: the log store implementation ensures that either the `write_to_wal` fails and no - // response is returned or the last entry ids for each region do exist. - let last_entry_id = response.last_entry_ids.get(region_id).unwrap(); - region_ctx.set_next_entry_id(last_entry_id + 1); - } - } - Err(e) => { - // Failed to write wal. - for mut region_ctx in region_ctxs.into_values() { - region_ctx.set_error(e.clone()); - } - return; - } + if !write_wal(&self.wal, &mut region_ctxs).await { + // Failed to write to the WAL, all waiters are notified with the error. + return; } } @@ -625,6 +600,55 @@ impl RegionWorkerLoop { } } +/// Writes WAL entries of all region contexts to the WAL in one batch and updates +/// the next entry id of each region on success. +/// +/// Returns `false` if the batch fails to be written to the WAL. In this case all +/// contexts are consumed and their waiters are notified with the error, so the +/// caller should skip the memtable phase. +async fn write_wal( + wal: &Wal, + region_ctxs: &mut HashMap, +) -> bool { + let mut wal_writer = wal.writer(); + for region_ctx in region_ctxs.values_mut() { + if region_ctx.skip_wal() { + continue; + } + if let Err(e) = region_ctx.add_wal_entry(&mut wal_writer).map_err(Arc::new) { + region_ctx.set_error(e); + } + } + match wal_writer.write_to_wal().await.map_err(Arc::new) { + Ok(response) => { + for (region_id, region_ctx) in region_ctxs.iter_mut() { + if region_ctx.skip_wal() { + continue; + } + // The entry of a failed region (e.g. failed to build its WAL entry) is + // not in the batch so the response has no last entry id for it. Its + // waiters are already notified with the error. + if region_ctx.is_failed() { + continue; + } + + // Safety: the log store implementation ensures that either the `write_to_wal` fails and no + // response is returned or the last entry ids for each region in the batch do exist. + let last_entry_id = response.last_entry_ids.get(region_id).unwrap(); + region_ctx.set_next_entry_id(last_entry_id + 1); + } + true + } + Err(e) => { + // Failed to write wal. + for (_, mut region_ctx) in region_ctxs.drain() { + region_ctx.set_error(e.clone()); + } + false + } + } +} + /// Send rejected error to all `write_requests`. fn reject_write_requests( write_requests: &mut Vec, @@ -707,3 +731,208 @@ fn check_partition_expr_version( } Ok(()) } + +#[cfg(test)] +mod tests { + use api::v1::{Row, Rows}; + use futures::stream; + use log_store::error::{ + Error as LogStoreError, IllegalStateSnafu, InvalidProviderSnafu, Result as LogStoreResult, + }; + use store_api::logstore::entry::{Entry, NaiveEntry}; + use store_api::logstore::provider::Provider; + use store_api::logstore::{AppendBatchResponse, EntryId, SendableEntryStream, WalIndex}; + use store_api::region_request::AffectedRows; + use tokio::sync::oneshot; + + use super::*; + use crate::request::OptionOutputTx; + use crate::test_util::version_util::VersionControlBuilder; + + /// A log store that fails to build entries for `failing_region` and fails the + /// whole batch when `fail_append` is true. + #[derive(Debug, Default)] + struct MockLogStore { + failing_region: Option, + fail_append: bool, + } + + #[async_trait::async_trait] + impl LogStore for MockLogStore { + type Error = LogStoreError; + + async fn stop(&self) -> LogStoreResult<()> { + Ok(()) + } + + async fn append_batch(&self, entries: Vec) -> LogStoreResult { + if self.fail_append { + return IllegalStateSnafu {}.fail(); + } + let mut last_entry_ids = HashMap::new(); + for entry in &entries { + let last_entry_id = last_entry_ids.entry(entry.region_id()).or_insert(0); + *last_entry_id = entry.entry_id().max(*last_entry_id); + } + Ok(AppendBatchResponse { last_entry_ids }) + } + + async fn read( + &self, + _provider: &Provider, + _entry_id: EntryId, + _index: Option, + ) -> LogStoreResult> { + Ok(Box::pin(stream::empty())) + } + + async fn create_namespace(&self, _ns: &Provider) -> LogStoreResult<()> { + Ok(()) + } + + async fn delete_namespace(&self, _ns: &Provider) -> LogStoreResult<()> { + Ok(()) + } + + async fn list_namespaces(&self) -> LogStoreResult> { + Ok(vec![]) + } + + async fn obsolete( + &self, + _provider: &Provider, + _region_id: RegionId, + _entry_id: EntryId, + ) -> LogStoreResult<()> { + Ok(()) + } + + async fn obsolete_all( + &self, + _provider: &Provider, + _region_id: RegionId, + ) -> LogStoreResult<()> { + Ok(()) + } + + fn entry( + &self, + data: Vec, + entry_id: EntryId, + region_id: RegionId, + provider: &Provider, + ) -> LogStoreResult { + if self.failing_region == Some(region_id) { + return InvalidProviderSnafu { + expected: "raft_engine", + actual: "mock", + } + .fail(); + } + Ok(Entry::Naive(NaiveEntry { + provider: provider.clone(), + region_id, + entry_id, + data, + })) + } + + fn latest_entry_id(&self, _provider: &Provider) -> LogStoreResult { + Ok(0) + } + } + + /// Creates a write context for `region_id` with one pending mutation of one row. + fn new_region_ctx( + region_id: RegionId, + ) -> (RegionWriteCtx, oneshot::Receiver>) { + let version_control = Arc::new(VersionControlBuilder::new().build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + let (tx, rx) = oneshot::channel(); + ctx.push_mutation( + OpType::Put as i32, + Some(Rows { + schema: vec![], + rows: vec![Row { values: vec![] }], + }), + None, + OptionOutputTx::from(tx), + None, + ); + (ctx, rx) + } + + #[tokio::test] + async fn test_write_wal_skips_region_failed_to_build_entry() { + let failing_region = RegionId::new(1, 1); + let ok_region = RegionId::new(1, 2); + let wal = Wal::new(Arc::new(MockLogStore { + failing_region: Some(failing_region), + ..Default::default() + })); + + let mut region_ctxs = HashMap::new(); + let (ctx, failing_rx) = new_region_ctx(failing_region); + region_ctxs.insert(failing_region, ctx); + let (ctx, ok_rx) = new_region_ctx(ok_region); + region_ctxs.insert(ok_region, ctx); + let entry_id = region_ctxs[&ok_region].next_entry_id(); + + // The failed region must not fail the batch or panic the worker. + assert!(write_wal(&wal, &mut region_ctxs).await); + + assert!(region_ctxs[&failing_region].is_failed()); + assert!(!region_ctxs[&ok_region].is_failed()); + assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); + + // Waiters of the failed region get the error while others get the result. + drop(region_ctxs); + assert!(failing_rx.await.unwrap().is_err()); + assert_eq!(1, ok_rx.await.unwrap().unwrap()); + } + + #[tokio::test] + async fn test_write_wal_all_regions_failed_to_build_entries() { + let failing_region = RegionId::new(1, 1); + let wal = Wal::new(Arc::new(MockLogStore { + failing_region: Some(failing_region), + ..Default::default() + })); + + let mut region_ctxs = HashMap::new(); + let (ctx, rx) = new_region_ctx(failing_region); + region_ctxs.insert(failing_region, ctx); + + // Writing an empty batch to the WAL succeeds, the failed region must not panic + // the worker. + assert!(write_wal(&wal, &mut region_ctxs).await); + + assert!(region_ctxs[&failing_region].is_failed()); + drop(region_ctxs); + assert!(rx.await.unwrap().is_err()); + } + + #[tokio::test] + async fn test_write_wal_append_batch_failure() { + let region_id = RegionId::new(1, 1); + let wal = Wal::new(Arc::new(MockLogStore { + fail_append: true, + ..Default::default() + })); + + let mut region_ctxs = HashMap::new(); + let (ctx, rx) = new_region_ctx(region_id); + region_ctxs.insert(region_id, ctx); + + assert!(!write_wal(&wal, &mut region_ctxs).await); + + // All contexts are consumed and waiters are notified with the error. + assert!(region_ctxs.is_empty()); + assert!(rx.await.unwrap().is_err()); + } +} From d4af650ec01c335b6ff392d6c78018ae73c8883b Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:18:16 +0800 Subject: [PATCH 028/126] perf: reduce cold workspace compile time (#8801) * refactor: remove datanode and meta-srv dep from frontend Signed-off-by: shuiyisong * refactor: use on-device protoc if possible Signed-off-by: shuiyisong * refactor: remove unused dep Signed-off-by: shuiyisong * fix: CR issue Signed-off-by: shuiyisong * fix: version and docs Signed-off-by: shuiyisong * chore: update logs Co-authored-by: fys <40801205+fengys1996@users.noreply.github.com> --------- Signed-off-by: shuiyisong Co-authored-by: fys <40801205+fengys1996@users.noreply.github.com> --- .agents/architecture-invariants.md | 4 +- Cargo.lock | 138 +++++++++++++----- Cargo.toml | 2 + src/cmd/src/standalone.rs | 6 +- src/common/stat/Cargo.toml | 4 +- src/common/test-util/Cargo.toml | 1 - src/common/test-util/src/recordbatch.rs | 28 ---- src/file-engine/Cargo.toml | 1 - src/frontend/AGENTS.md | 12 +- src/frontend/Cargo.toml | 3 - src/frontend/src/error.rs | 14 -- src/frontend/src/frontend.rs | 6 +- src/frontend/src/instance.rs | 2 - src/index/Cargo.toml | 4 +- src/log-store/Cargo.toml | 5 +- src/log-store/build.rs | 95 +++++++++++- src/standalone/Cargo.toml | 2 + .../src/datanode_manager.rs} | 4 +- src/standalone/src/error.rs | 21 ++- src/standalone/src/lib.rs | 2 + tests-integration/src/standalone.rs | 4 +- 21 files changed, 244 insertions(+), 114 deletions(-) rename src/{frontend/src/instance/standalone.rs => standalone/src/datanode_manager.rs} (97%) diff --git a/.agents/architecture-invariants.md b/.agents/architecture-invariants.md index 78c78f3a976..f290e1ce386 100644 --- a/.agents/architecture-invariants.md +++ b/.agents/architecture-invariants.md @@ -46,8 +46,8 @@ The workspace is layered; dependencies point downward only. `file-engine`) implement it; `datanode` drives engines **through the trait**, not through engine internals. - `frontend` reaches storage through `operator` / `query` / `catalog`, not by - depending on `datanode` internals. Standalone mode is the one bridge, via a - `RegionServer` wrapper (`src/frontend/src/instance/standalone.rs`). + depending on `datanode` internals. Standalone mode is the one bridge, via the + `RegionServer` adapter in `src/standalone/src/datanode_manager.rs`. - Do not introduce circular dependencies. New deps go through `[workspace.dependencies]` in the root `Cargo.toml`, not per-crate version literals. diff --git a/Cargo.lock b/Cargo.lock index 841cc035f77..4ba0a7f7e27 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -989,15 +989,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" -[[package]] -name = "autotools" -version = "0.2.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef941527c41b0fc0dd48511a8154cd5fc7e29200a0ff8b7203c5d777dbc795cf" -dependencies = [ - "cc", -] - [[package]] name = "aws-lc-rs" version = "1.17.0" @@ -2917,7 +2908,6 @@ dependencies = [ name = "common-test-util" version = "1.3.0" dependencies = [ - "client", "common-grpc", "common-query", "common-recordbatch", @@ -5274,7 +5264,6 @@ dependencies = [ "common-datasource", "common-error", "common-macro", - "common-procedure", "common-query", "common-recordbatch", "common-telemetry", @@ -5587,7 +5576,6 @@ dependencies = [ "dashmap", "datafusion", "datafusion-expr", - "datanode", "datatypes", "futures", "hostname 0.4.1", @@ -5597,7 +5585,6 @@ dependencies = [ "lazy_static", "log-query", "meta-client", - "meta-srv", "num_cpus", "opentelemetry-proto", "opentelemetry-semantic-conventions 0.31.0", @@ -8013,7 +8000,8 @@ dependencies = [ "pin-project", "prometheus 0.14.0", "protobuf 2.28.0", - "protobuf-build", + "protoc-bin-vendored", + "protoc-rust", "raft-engine", "rand 0.9.4", "rskafka", @@ -11379,19 +11367,6 @@ dependencies = [ "thiserror 1.0.69", ] -[[package]] -name = "protobuf-build" -version = "0.15.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c852d9625b912c3e50480cdc701f60f49890b5d7ad46198dd583600f15e7c6ec" -dependencies = [ - "bitflags 1.3.2", - "protobuf 2.28.0", - "protobuf-codegen", - "protobuf-src", - "regex", -] - [[package]] name = "protobuf-codegen" version = "2.28.0" @@ -11401,15 +11376,6 @@ dependencies = [ "protobuf 2.28.0", ] -[[package]] -name = "protobuf-src" -version = "1.1.0+21.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7ac8852baeb3cc6fb83b93646fb93c0ffe5d14bf138c945ceb4b9948ee0e3c1" -dependencies = [ - "autotools", -] - [[package]] name = "protobuf-support" version = "3.7.2" @@ -11419,6 +11385,92 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "protoc" +version = "2.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0218039c514f9e14a5060742ecd50427f8ac4f85a6dc58f2ddb806e318c55ee" +dependencies = [ + "log", + "which", +] + +[[package]] +name = "protoc-bin-vendored" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1c381df33c98266b5f08186583660090a4ffa0889e76c7e9a5e175f645a67fa" +dependencies = [ + "protoc-bin-vendored-linux-aarch_64", + "protoc-bin-vendored-linux-ppcle_64", + "protoc-bin-vendored-linux-s390_64", + "protoc-bin-vendored-linux-x86_32", + "protoc-bin-vendored-linux-x86_64", + "protoc-bin-vendored-macos-aarch_64", + "protoc-bin-vendored-macos-x86_64", + "protoc-bin-vendored-win32", +] + +[[package]] +name = "protoc-bin-vendored-linux-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c350df4d49b5b9e3ca79f7e646fde2377b199e13cfa87320308397e1f37e1a4c" + +[[package]] +name = "protoc-bin-vendored-linux-ppcle_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a55a63e6c7244f19b5c6393f025017eb5d793fd5467823a099740a7a4222440c" + +[[package]] +name = "protoc-bin-vendored-linux-s390_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1dba5565db4288e935d5330a07c264a4ee8e4a5b4a4e6f4e83fad824cc32f3b0" + +[[package]] +name = "protoc-bin-vendored-linux-x86_32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8854774b24ee28b7868cd71dccaae8e02a2365e67a4a87a6cd11ee6cdbdf9cf5" + +[[package]] +name = "protoc-bin-vendored-linux-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b38b07546580df720fa464ce124c4b03630a6fb83e05c336fea2a241df7e5d78" + +[[package]] +name = "protoc-bin-vendored-macos-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89278a9926ce312e51f1d999fee8825d324d603213344a9a706daa009f1d8092" + +[[package]] +name = "protoc-bin-vendored-macos-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81745feda7ccfb9471d7a4de888f0652e806d5795b61480605d4943176299756" + +[[package]] +name = "protoc-bin-vendored-win32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95067976aca6421a523e491fce939a3e65249bac4b977adee0ee9771568e8aa3" + +[[package]] +name = "protoc-rust" +version = "2.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22f8a182bb17c485f20bdc4274a8c39000a61024cfe461c799b50fec77267838" +dependencies = [ + "protobuf 2.28.0", + "protobuf-codegen", + "protoc", + "tempfile", +] + [[package]] name = "psl" version = "2.1.123" @@ -14066,6 +14118,7 @@ dependencies = [ name = "standalone" version = "1.3.0" dependencies = [ + "api", "async-trait", "catalog", "client", @@ -14079,6 +14132,7 @@ dependencies = [ "common-options", "common-procedure", "common-query", + "common-recordbatch", "common-stat", "common-telemetry", "common-time", @@ -16360,6 +16414,18 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix 0.38.44", +] + [[package]] name = "whoami" version = "1.6.0" diff --git a/Cargo.toml b/Cargo.toml index c9b199b3673..502e652ffc8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -202,6 +202,8 @@ prometheus = { version = "0.14", features = ["process"] } promql-parser = { version = "0.10.0", features = ["ser"] } prost = { version = "0.14", features = ["no-recursion-limit"] } prost-types = "0.14" +protoc-bin-vendored = "3.2" +protoc-rust = "2.28" raft-engine = { version = "0.4.1", default-features = false } rand = "0.9" ratelimit = "0.10" diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index aec373612a4..9f8155b9cc2 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -60,7 +60,6 @@ use flow::{ FrontendInvoker, GrpcQueryHandlerWithBoxedError, }; use frontend::frontend::Frontend; -use frontend::instance::StandaloneDatanodeManager; use frontend::instance::builder::FrontendBuilder; use frontend::server::Services; use meta_srv::metasrv::{FLOW_ID_SEQ, TABLE_ID_SEQ}; @@ -72,7 +71,10 @@ use plugins::standalone::context::DdlManagerConfigureContext; use servers::tls::{TlsMode, TlsOption, merge_tls_option}; use snafu::{OptionExt, ResultExt}; use standalone::options::StandaloneOptions; -use standalone::{StandaloneInformationExtension, StandaloneRepartitionProcedureFactory}; +use standalone::{ + StandaloneDatanodeManager, StandaloneInformationExtension, + StandaloneRepartitionProcedureFactory, +}; use tracing_appender::non_blocking::WorkerGuard; use crate::error::{OtherSnafu, Result, StartFlownodeSnafu}; diff --git a/src/common/stat/Cargo.toml b/src/common/stat/Cargo.toml index d0e8b5448ff..2a99107c909 100644 --- a/src/common/stat/Cargo.toml +++ b/src/common/stat/Cargo.toml @@ -9,11 +9,13 @@ common-base.workspace = true common-runtime.workspace = true common-telemetry.workspace = true lazy_static.workspace = true -nix.workspace = true num_cpus.workspace = true prometheus.workspace = true sysinfo.workspace = true tokio.workspace = true +[target.'cfg(target_os = "linux")'.dependencies] +nix.workspace = true + [lints] workspace = true diff --git a/src/common/test-util/Cargo.toml b/src/common/test-util/Cargo.toml index 18dc70e100b..46b374746ab 100644 --- a/src/common/test-util/Cargo.toml +++ b/src/common/test-util/Cargo.toml @@ -8,7 +8,6 @@ license.workspace = true workspace = true [dependencies] -client = { workspace = true, features = ["testing"] } common-grpc.workspace = true common-query.workspace = true common-recordbatch.workspace = true diff --git a/src/common/test-util/src/recordbatch.rs b/src/common/test-util/src/recordbatch.rs index aa68f793566..15e754ca7a6 100644 --- a/src/common/test-util/src/recordbatch.rs +++ b/src/common/test-util/src/recordbatch.rs @@ -12,15 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use client::Database; use common_query::OutputData; use common_recordbatch::util; -pub enum ExpectedOutput<'a> { - AffectedRows(usize), - QueryResult(&'a str), -} - pub async fn check_output_stream(output: OutputData, expected: &str) { let recordbatches = match output { OutputData::Stream(stream) => util::collect_batches(stream).await.unwrap(), @@ -30,25 +24,3 @@ pub async fn check_output_stream(output: OutputData, expected: &str) { let pretty_print = recordbatches.pretty_print().unwrap(); assert_eq!(pretty_print, expected.trim(), "actual: \n{}", pretty_print); } - -pub async fn execute_and_check_output(db: &Database, sql: &str, expected: ExpectedOutput<'_>) { - let output = db.sql(sql).await.unwrap(); - let output = output.data; - - match (&output, expected) { - (OutputData::AffectedRows(x), ExpectedOutput::AffectedRows(y)) => { - assert_eq!( - *x, y, - r#" -expected: {y} -actual: {x} -"# - ) - } - (OutputData::RecordBatches(_), ExpectedOutput::QueryResult(x)) - | (OutputData::Stream(_), ExpectedOutput::QueryResult(x)) => { - check_output_stream(output, x).await - } - _ => panic!(), - } -} diff --git a/src/file-engine/Cargo.toml b/src/file-engine/Cargo.toml index 9d031cb279f..ec904addf1e 100644 --- a/src/file-engine/Cargo.toml +++ b/src/file-engine/Cargo.toml @@ -18,7 +18,6 @@ common-catalog.workspace = true common-datasource.workspace = true common-error.workspace = true common-macro.workspace = true -common-procedure.workspace = true common-query.workspace = true common-recordbatch.workspace = true common-telemetry.workspace = true diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index 3afd75634d0..b5c0871fc11 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -16,8 +16,9 @@ reads go to the query engine (`query` crate), writes go to the inserter/deleter Boundary with `servers`: the `servers` crate implements the wire protocols and network I/O; `frontend` provides the business logic by implementing handler traits (`SqlQueryHandler`, `GrpcQueryHandler`, `InfluxdbLineProtocolHandler`, -...). In standalone mode the frontend embeds a datanode `RegionServer`; in -distributed mode it talks to remote datanodes via `operator`/`client`. +...). In standalone mode the `standalone` crate bridges the frontend to an +embedded datanode `RegionServer`; in distributed mode the frontend talks to +remote datanodes via `operator`/`client`. ## Module map @@ -26,7 +27,6 @@ distributed mode it talks to remote datanodes via `operator`/`client`. | `instance` | `src/frontend/src/instance.rs` | `Instance`: the core handler; implements `SqlQueryHandler`, `PrometheusHandler`, etc. | | `instance/builder` | `src/frontend/src/instance/builder.rs` | `FrontendBuilder` assembles `Instance` from its dependencies | | `instance/grpc` | `src/frontend/src/instance/grpc.rs` | `GrpcQueryHandler`: insert/delete/query/promql over gRPC | -| `instance/standalone` | `src/frontend/src/instance/standalone.rs` | Calls the local `RegionServer` instead of RPC | | `instance/region_query` | `src/frontend/src/instance/region_query.rs` | Routes distributed region reads to datanodes | | `instance/*` | `src/frontend/src/instance/` | Per-protocol handlers (`influxdb.rs`, `promql.rs`, `otlp/`, `jaeger.rs`, `logs.rs`, `prom_store.rs`, ...) | | `frontend` | `src/frontend/src/frontend.rs` | `Frontend` lifecycle wrapper (`FrontendOptions`, start/shutdown) | @@ -70,9 +70,9 @@ cargo nextest run -p frontend - Keep the frontend/servers split straight: wire format and network live in `servers`; permissions, planning, and routing live here. -- Standalone vs distributed diverge in datanode access (local `RegionServer` vs - `NodeClients` RPC), MetaClient usage, and whether heartbeat matters. In - standalone, the cache invalidator is a no-op. +- Standalone vs distributed diverge in datanode access (the `standalone` crate's + local `RegionServer` adapter vs `NodeClients` RPC), MetaClient usage, and + whether heartbeat matters. In standalone, the cache invalidator is a no-op. ## Maintenance contract diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index a5a589e2ac7..29ca224a55a 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -52,7 +52,6 @@ common-version.workspace = true dashmap.workspace = true datafusion.workspace = true datafusion-expr.workspace = true -datanode.workspace = true datatypes.workspace = true futures.workspace = true hostname.workspace = true @@ -91,11 +90,9 @@ tracing.workspace = true [dev-dependencies] catalog = { workspace = true, features = ["testing"] } common-test-util.workspace = true -datanode.workspace = true datatypes.workspace = true futures.workspace = true hyper-util = { workspace = true, features = ["tokio"] } -meta-srv.workspace = true reqwest.workspace = true serde_json.workspace = true strfmt = "0.2" diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs index 9a70b716619..c0f9e2fe6bd 100644 --- a/src/frontend/src/error.rs +++ b/src/frontend/src/error.rs @@ -203,13 +203,6 @@ pub enum Error { source: query::error::Error, }, - #[snafu(display("Operation to region server failed"))] - InvokeRegionServer { - #[snafu(implicit)] - location: Location, - source: servers::error::Error, - }, - #[snafu(display("Not supported: {}", feat))] NotSupported { feat: String }, @@ -296,9 +289,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Invalid region request, reason: {}", reason))] - InvalidRegionRequest { reason: String }, - #[snafu(display("Table operation error"))] TableOperation { source: operator::error::Error, @@ -426,8 +416,6 @@ impl ErrorExt for Error { Error::CacheRequired { .. } => StatusCode::Internal, - Error::InvalidRegionRequest { .. } => StatusCode::IllegalState, - Error::TableNotFound { .. } => StatusCode::TableNotFound, Error::Catalog { source, .. } => source.status_code(), @@ -438,7 +426,6 @@ impl ErrorExt for Error { | Error::ReadTable { source, .. } | Error::ExecLogicalPlan { source, .. } => source.status_code(), - Error::InvokeRegionServer { source, .. } => source.status_code(), Error::External { source, .. } | Error::InitPlugin { source, .. } => { source.status_code() } @@ -472,7 +459,6 @@ impl ErrorExt for Error { Error::StartServer { source, .. } | Error::ShutdownServer { source, .. } - | Error::InvokeRegionServer { source, .. } | Error::ExecutePromql { source, .. } | Error::PromStoreRemoteQueryPlan { source, .. } | Error::PrometheusMetricNamesQueryPlan { source, .. } => source.retry_hint(), diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 30fcd394922..7a09b40e736 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -165,6 +165,7 @@ impl Frontend { mod tests { use std::any::Any; use std::net::SocketAddr; + use std::pin::Pin; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::time::Duration; @@ -186,9 +187,9 @@ mod tests { use common_meta::heartbeat::handler::suspend::SuspendHandler; use common_meta::instruction::Instruction; use common_stat::ResourceStatImpl; + use futures::Stream; use meta_client::MetaClientRef; use meta_client::client::MetaClientBuilder; - use meta_srv::service::GrpcStream; use servers::grpc::{FlightCompression, GRPC_SERVER}; use servers::http::HTTP_SERVER; use servers::http::result::greptime_result_v1::GreptimedbV1Response; @@ -206,6 +207,9 @@ mod tests { use crate::instance::builder::FrontendBuilder; use crate::server::Services; + type GrpcStream = + Pin> + Send + Sync + 'static>>; + #[test] fn test_toml() { let opts = FrontendOptions::default(); diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 23090ef932b..bdf4d66c7e9 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -25,7 +25,6 @@ mod otlp; pub mod prom_store; mod promql; mod region_query; -pub mod standalone; use std::collections::HashSet; use std::pin::Pin; @@ -100,7 +99,6 @@ use sql::statements::statement::Statement; use sql::statements::tql::Tql; use sql::util::{extract_tables_from_prom_expr_checked, extract_tables_from_statement_checked}; use sqlparser::ast::{AnalyzeFormat, ObjectName}; -pub use standalone::StandaloneDatanodeManager; use table::requests::{OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM}; use tracing::Span; diff --git a/src/index/Cargo.toml b/src/index/Cargo.toml index 167c1c0df14..b8518112610 100644 --- a/src/index/Cargo.toml +++ b/src/index/Cargo.toml @@ -8,7 +8,7 @@ license.workspace = true workspace = true [features] -vector_index = ["dep:usearch"] +vector_index = ["dep:nalgebra", "dep:usearch"] [dependencies] async-trait.workspace = true @@ -29,7 +29,7 @@ itertools.workspace = true jieba-rs = "0.10" lazy_static.workspace = true mockall.workspace = true -nalgebra.workspace = true +nalgebra = { workspace = true, optional = true } pin-project.workspace = true prost.workspace = true puffin.workspace = true diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index b98657f5609..c8372c578bd 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -5,9 +5,8 @@ edition.workspace = true license.workspace = true [build-dependencies] -protobuf-build = { version = "0.15", default-features = false, features = [ - "protobuf-codec", -] } +protoc-bin-vendored.workspace = true +protoc-rust.workspace = true [lints] workspace = true diff --git a/src/log-store/build.rs b/src/log-store/build.rs index c6c15cdfe2c..a8b6f791487 100644 --- a/src/log-store/build.rs +++ b/src/log-store/build.rs @@ -12,13 +12,94 @@ // See the License for the specific language governing permissions and // limitations under the License. -use protobuf_build::Builder; +//! Generates Rust protobuf bindings for `proto/logstore.proto`. +//! +//! The protoc compiler is selected in the following order: +//! +//! 1. Use the compiler specified by the `PROTOC` environment variable. +//! If the configured compiler is unusable or does not meet the minimum +//! supported version requirement, the build will panic instead of falling +//! back to another compiler. +//! +//! 2. Use a compatible `protoc` compiler available from `PATH`. +//! If it is unavailable or does not meet the minimum supported version +//! requirement, continue with the vendored compiler fallback. +//! +//! 3. Use a prebuilt vendored `protoc` compiler. +//! This provides a portable fallback when no supported system compiler is +//! available. + +use std::env; +use std::path::{Path, PathBuf}; +use std::process::Command; + +use protoc_rust::{Codegen, Customize}; + +/// Returns whether `protoc` runs and meets GreptimeDB's minimum supported version. +fn usable_protoc(protoc: &Path) -> bool { + let Ok(output) = Command::new(protoc).arg("--version").output() else { + return false; + }; + if !output.status.success() { + return false; + } + + let Some(version) = std::str::from_utf8(&output.stdout) + .ok() + .and_then(|output| output.trim().strip_prefix("libprotoc ")) + else { + return false; + }; + let mut components = version.split('.'); + matches!( + ( + components.next().and_then(|part| part.parse::().ok()), + components.next().and_then(|part| part.parse::().ok()), + ), + (Some(major), Some(minor)) if (major, minor) >= (3, 15) + ) +} + +/// Selects a compiler from `PROTOC`, `PATH`, or the vendored fallback, in that order. +/// +/// An invalid explicit `PROTOC` is rejected so configuration errors are not +/// silently masked. +fn protoc_path() -> PathBuf { + if let Some(protoc) = env::var_os("PROTOC") { + let protoc = PathBuf::from(protoc); + assert!(usable_protoc(&protoc), "PROTOC version not usable"); + return protoc; + } + + let protoc = PathBuf::from("protoc"); + if usable_protoc(&protoc) { + return protoc; + } + + protoc_bin_vendored::protoc_bin_path().expect("no bundled protoc for this platform") +} fn main() { - let base = std::env::var("CARGO_MANIFEST_DIR").unwrap_or_else(|_| ".".to_string()); - Builder::new() - .search_dir_for_protos(&format!("{base}/proto")) - .includes(&[format!("{base}/include"), format!("{base}/proto")]) - .include_google_protos() - .generate() + let base = + PathBuf::from(env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not set")); + let proto_dir = base.join("proto"); + let proto = proto_dir.join("logstore.proto"); + let out_dir = PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR not set")).join("protos"); + + println!("cargo:rerun-if-env-changed=PROTOC"); + println!("cargo:rerun-if-env-changed=PATH"); + println!("cargo:rerun-if-changed={}", proto.display()); + std::fs::create_dir_all(&out_dir).expect("failed to create protobuf output directory"); + + Codegen::new() + .protoc_path(protoc_path()) + .out_dir(out_dir) + .input(proto) + .include(proto_dir) + .customize(Customize { + gen_mod_rs: Some(true), + ..Default::default() + }) + .run() + .expect("failed to generate log-store protobuf bindings") } diff --git a/src/standalone/Cargo.toml b/src/standalone/Cargo.toml index 16479991f41..836f891fa49 100644 --- a/src/standalone/Cargo.toml +++ b/src/standalone/Cargo.toml @@ -8,6 +8,7 @@ license.workspace = true workspace = true [dependencies] +api.workspace = true async-trait.workspace = true catalog.workspace = true client.workspace = true @@ -21,6 +22,7 @@ common-meta.workspace = true common-options.workspace = true common-procedure.workspace = true common-query.workspace = true +common-recordbatch.workspace = true common-stat.workspace = true common-telemetry.workspace = true common-time.workspace = true diff --git a/src/frontend/src/instance/standalone.rs b/src/standalone/src/datanode_manager.rs similarity index 97% rename from src/frontend/src/instance/standalone.rs rename to src/standalone/src/datanode_manager.rs index d34922956b3..13a17ea2b43 100644 --- a/src/frontend/src/instance/standalone.rs +++ b/src/standalone/src/datanode_manager.rs @@ -54,12 +54,12 @@ impl FlownodeManager for StandaloneDatanodeManager { } /// Relative to [client::region::RegionRequester] -pub struct RegionInvoker { +struct RegionInvoker { region_server: RegionServer, } impl RegionInvoker { - pub fn arc(region_server: RegionServer) -> Arc { + fn arc(region_server: RegionServer) -> Arc { Arc::new(Self { region_server }) } diff --git a/src/standalone/src/error.rs b/src/standalone/src/error.rs index 37859490392..8301c104b59 100644 --- a/src/standalone/src/error.rs +++ b/src/standalone/src/error.rs @@ -14,7 +14,7 @@ use std::any::Any; -use common_error::ext::{BoxedError, ErrorExt}; +use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; use snafu::{Location, Snafu}; @@ -30,6 +30,13 @@ pub enum Error { source: BoxedError, }, + #[snafu(display("Operation to region server failed"))] + InvokeRegionServer { + #[snafu(implicit)] + location: Location, + source: servers::error::Error, + }, + #[snafu(display("External error"))] External { source: BoxedError, @@ -49,6 +56,9 @@ pub enum Error { #[snafu(display("Invalid url scheme: {}", scheme))] InvalidUrlScheme { scheme: String }, + #[snafu(display("Invalid region request, reason: {}", reason))] + InvalidRegionRequest { reason: String }, + #[snafu(display("Repartition procedure is not supported in standalone mode"))] NoSupportRepartitionProcedure { #[snafu(implicit)] @@ -62,9 +72,11 @@ impl ErrorExt for Error { fn status_code(&self) -> StatusCode { match self { Error::OpenMetadataKvBackend { source, .. } => source.status_code(), + Error::InvokeRegionServer { source, .. } => source.status_code(), Error::External { source, .. } => source.status_code(), Error::ParseUrl { .. } => StatusCode::InvalidArguments, Error::InvalidUrlScheme { .. } => StatusCode::InvalidArguments, + Error::InvalidRegionRequest { .. } => StatusCode::IllegalState, Error::NoSupportRepartitionProcedure { .. } => StatusCode::Unsupported, } } @@ -72,4 +84,11 @@ impl ErrorExt for Error { fn as_any(&self) -> &dyn Any { self } + + fn retry_hint(&self) -> RetryHint { + match self { + Error::InvokeRegionServer { source, .. } => source.retry_hint(), + _ => RetryHint::NonRetryable, + } + } } diff --git a/src/standalone/src/lib.rs b/src/standalone/src/lib.rs index 5f3d2a2de4d..57c4efee249 100644 --- a/src/standalone/src/lib.rs +++ b/src/standalone/src/lib.rs @@ -12,12 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod datanode_manager; pub mod error; pub mod information_extension; pub mod metadata; pub mod options; pub mod procedure; +pub use datanode_manager::StandaloneDatanodeManager; pub use information_extension::StandaloneInformationExtension; pub use metadata::{build_metadata_kv_from_url, build_metadata_kvbackend}; pub use procedure::{StandaloneRepartitionProcedureFactory, build_procedure_manager}; diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index f93d7da804e..1085c21983e 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -49,14 +49,14 @@ use common_wal::config::{DatanodeWalConfig, MetasrvWalConfig}; use datanode::datanode::DatanodeBuilder; use flow::{FlownodeBuilder, FrontendClient, GrpcQueryHandlerWithBoxedError}; use frontend::frontend::Frontend; +use frontend::instance::Instance; use frontend::instance::builder::FrontendBuilder; -use frontend::instance::{Instance, StandaloneDatanodeManager}; use frontend::server::Services; use meta_srv::metasrv::{FLOW_ID_SEQ, TABLE_ID_SEQ}; use servers::grpc::GrpcOptions; use snafu::ResultExt; -use standalone::StandaloneRepartitionProcedureFactory; use standalone::options::StandaloneOptions; +use standalone::{StandaloneDatanodeManager, StandaloneRepartitionProcedureFactory}; use crate::test_util::{self, StorageType, TestGuard, create_tmp_dir_and_datanode_opts}; From 7fa6f5f98e0811f2f16bba686ed5ab2b2f9286ee Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:11:21 +0800 Subject: [PATCH 029/126] feat: add riscv64 cross-build support (#8820) * feat: add riscv64 cross-build support Add the missing build infrastructure for riscv64gc-unknown-linux-gnu. The codebase itself already compiles cleanly for riscv64 (verified with `cargo check --workspace --target riscv64gc-unknown-linux-gnu`): all architecture-sensitive dependencies support it (tikv-jemalloc-sys, aws-lc-sys, ring, pprof, simd-json). - .cargo/config.toml: set riscv64-linux-gnu-gcc as the linker for the riscv64gc-unknown-linux-gnu target - rust.yml: add a check-riscv64 CI job that cross-checks the whole workspace to prevent regressions from future dependency changes - docker/dev-builder/riscv64/Dockerfile: new cross dev-builder image with gcc/g++-riscv64-linux-gnu and the riscv64 rust target - Makefile: add dev-builder-riscv64 and build-riscv64-bin targets Verified end-to-end: the produced riscv64 binary starts in standalone mode under qemu and serves SQL (create/insert/select) over HTTP. Signed-off-by: Lei, HUANG * ci: build riscv64 artifacts in the release workflow - release.yml: add build-linux-riscv64-artifacts job that cross-compiles greptime for riscv64gc-unknown-linux-gnu on the amd64 runner with the dev-builder-riscv64 image, and uploads greptime-linux-riscv64-* artifacts. The job is wired into the needs of publish-github-release, release-cn-artifacts and stop-linux-amd64-runner. Integration tests are skipped since the cross-compiled binary cannot run on the host. - release-dev-builder-images.yaml + build-dev-builder-images action: build and push the dev-builder-riscv64 image to DockerHub, and sync it to ECR and ACR via skopeo like the other dev-builder images. - Makefile: DEV_BUILDER_RISCV64_IMAGE_TAG now defaults to DEV_BUILDER_IMAGE_TAG so the existing tag-bump automation keeps the riscv64 image tag in sync. Signed-off-by: Lei, HUANG * fix: forward cargo extension in riscv64 build Pass CARGO_EXTENSION through build-riscv64-bin just like the existing build-by-dev-builder target, so wrappers such as sccache are preserved inside the cross-build container. Signed-off-by: Lei, HUANG * ci: check riscv64 release feature graph Check all workspace targets with the servers/dashboard feature enabled so the riscv64 CI job covers the same optional dependency graph used by the release artifact build. Signed-off-by: Lei, HUANG * fix: gate riscv64 latest tags to main pushes Manual dev-builder workflow dispatches now publish only their immutable version tag. Update DockerHub and ECR latest tags only for the workflow's main-branch push event, preventing feature-branch builds from replacing the shared latest image. Signed-off-by: Lei, HUANG * docs: include riscv64 in release input description Update the build_linux_artifacts workflow input description to reflect that it now triggers amd64, arm64, and riscv64 artifact builds. Signed-off-by: Lei, HUANG * fix: fall back when riscv64 builder is unpublished Before building a release artifact, pull the pinned RISC-V dev-builder from ECR. If the image has not been published yet, build the same tag locally from the current Dockerfile so releases remain unblocked during the builder-tag update window. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- .cargo/config.toml | 3 + .../build-dev-builder-images/action.yml | 19 +++++++ .../workflows/release-dev-builder-images.yaml | 50 ++++++++++++++++- .github/workflows/release.yml | 52 +++++++++++++++++- .github/workflows/rust.yml | 26 +++++++++ Makefile | 26 +++++++++ docker/dev-builder/riscv64/Dockerfile | 55 +++++++++++++++++++ 7 files changed, 229 insertions(+), 2 deletions(-) create mode 100644 docker/dev-builder/riscv64/Dockerfile diff --git a/.cargo/config.toml b/.cargo/config.toml index 53dcb05f335..9ffc1961b09 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,6 +1,9 @@ [target.aarch64-unknown-linux-gnu] linker = "aarch64-linux-gnu-gcc" +[target.riscv64gc-unknown-linux-gnu] +linker = "riscv64-linux-gnu-gcc" + [alias] sqlness = "run --bin sqlness-runner --target-dir target/sqlness --" diff --git a/.github/actions/build-dev-builder-images/action.yml b/.github/actions/build-dev-builder-images/action.yml index 41f69193969..1ca37101302 100644 --- a/.github/actions/build-dev-builder-images/action.yml +++ b/.github/actions/build-dev-builder-images/action.yml @@ -31,6 +31,10 @@ inputs: description: Build dev-builder-android image required: false default: "true" + build-dev-builder-riscv64: + description: Build dev-builder-riscv64 image + required: false + default: "true" runs: using: composite steps: @@ -70,6 +74,21 @@ runs: IMAGE_NAMESPACE=${{ inputs.dockerhub-image-namespace }} \ DEV_BUILDER_IMAGE_TAG=${{ inputs.version }} + - name: Build and push dev-builder-riscv64 image # Only build image for amd64 platform (it cross-compiles riscv64 binaries). + shell: bash + if: ${{ inputs.build-dev-builder-riscv64 == 'true' }} + run: | + make dev-builder-riscv64 \ + IMAGE_REGISTRY=${{ inputs.dockerhub-image-registry }} \ + IMAGE_NAMESPACE=${{ inputs.dockerhub-image-namespace }} \ + DEV_BUILDER_RISCV64_IMAGE_TAG=${{ inputs.version }} + docker push ${{ inputs.dockerhub-image-registry }}/${{ inputs.dockerhub-image-namespace }}/dev-builder-riscv64:${{ inputs.version }} + if [[ "${{ github.event_name }}" == "push" ]]; then + docker tag ${{ inputs.dockerhub-image-registry }}/${{ inputs.dockerhub-image-namespace }}/dev-builder-riscv64:${{ inputs.version }} \ + ${{ inputs.dockerhub-image-registry }}/${{ inputs.dockerhub-image-namespace }}/dev-builder-riscv64:latest + docker push ${{ inputs.dockerhub-image-registry }}/${{ inputs.dockerhub-image-namespace }}/dev-builder-riscv64:latest + fi + - name: Build and push dev-builder-android image # Only build image for amd64 platform. shell: bash if: ${{ inputs.build-dev-builder-android == 'true' }} diff --git a/.github/workflows/release-dev-builder-images.yaml b/.github/workflows/release-dev-builder-images.yaml index e6a33dcceb0..bebd5ed61e7 100644 --- a/.github/workflows/release-dev-builder-images.yaml +++ b/.github/workflows/release-dev-builder-images.yaml @@ -24,6 +24,11 @@ on: description: Release dev-builder-android image required: false default: false + release_dev_builder_riscv64_image: + type: boolean + description: Release dev-builder-riscv64 image + required: false + default: false update_dev_builder_image_tag: type: boolean description: Update the DEV_BUILDER_IMAGE_TAG in Makefile and create a PR @@ -36,7 +41,7 @@ jobs: # The jobs are triggered by the following events: # 1. Manually triggered workflow_dispatch event # 2. Push event when the PR that modifies the `rust-toolchain.toml` or `docker/dev-builder/**` is merged to main - if: ${{ github.event_name == 'push' || inputs.release_dev_builder_ubuntu_image || inputs.release_dev_builder_centos_image || inputs.release_dev_builder_android_image }} + if: ${{ github.event_name == 'push' || inputs.release_dev_builder_ubuntu_image || inputs.release_dev_builder_centos_image || inputs.release_dev_builder_android_image || inputs.release_dev_builder_riscv64_image }} runs-on: ubuntu-latest outputs: version: ${{ steps.set-version.outputs.version }} @@ -68,6 +73,7 @@ jobs: build-dev-builder-ubuntu: ${{ inputs.release_dev_builder_ubuntu_image || github.event_name == 'push' }} build-dev-builder-centos: ${{ inputs.release_dev_builder_centos_image || github.event_name == 'push' }} build-dev-builder-android: ${{ inputs.release_dev_builder_android_image || github.event_name == 'push' }} + build-dev-builder-riscv64: ${{ inputs.release_dev_builder_riscv64_image || github.event_name == 'push' }} release-dev-builder-images-ecr: name: Release dev builder images to AWS ECR @@ -154,6 +160,35 @@ jobs: copy -a docker://docker.io/$IMAGE_NAMESPACE/dev-builder-android:latest \ docker://$ECR_IMAGE_REGISTRY/$ECR_IMAGE_NAMESPACE/dev-builder-android:latest + - name: Push dev-builder-riscv64 image + shell: bash + if: ${{ inputs.release_dev_builder_riscv64_image || github.event_name == 'push' }} + env: + IMAGE_VERSION: ${{ needs.release-dev-builder-images.outputs.version }} + IMAGE_NAMESPACE: ${{ vars.IMAGE_NAMESPACE }} + ECR_IMAGE_REGISTRY: ${{ vars.ECR_IMAGE_REGISTRY }} + ECR_IMAGE_NAMESPACE: ${{ vars.ECR_IMAGE_NAMESPACE }} + run: | + docker run -v "${DOCKER_CONFIG:-$HOME/.docker}:/root/.docker:ro" \ + -e "REGISTRY_AUTH_FILE=/root/.docker/config.json" \ + quay.io/skopeo/stable:latest \ + copy -a docker://docker.io/$IMAGE_NAMESPACE/dev-builder-riscv64:$IMAGE_VERSION \ + docker://$ECR_IMAGE_REGISTRY/$ECR_IMAGE_NAMESPACE/dev-builder-riscv64:$IMAGE_VERSION + + - name: Push latest dev-builder-riscv64 image + shell: bash + if: ${{ github.event_name == 'push' }} + env: + IMAGE_NAMESPACE: ${{ vars.IMAGE_NAMESPACE }} + ECR_IMAGE_REGISTRY: ${{ vars.ECR_IMAGE_REGISTRY }} + ECR_IMAGE_NAMESPACE: ${{ vars.ECR_IMAGE_NAMESPACE }} + run: | + docker run -v "${DOCKER_CONFIG:-$HOME/.docker}:/root/.docker:ro" \ + -e "REGISTRY_AUTH_FILE=/root/.docker/config.json" \ + quay.io/skopeo/stable:latest \ + copy -a docker://docker.io/$IMAGE_NAMESPACE/dev-builder-riscv64:latest \ + docker://$ECR_IMAGE_REGISTRY/$ECR_IMAGE_NAMESPACE/dev-builder-riscv64:latest + release-dev-builder-images-cn: # Note: Be careful issue: https://github.com/containers/skopeo/issues/1874 and we decide to use the latest stable skopeo container. name: Release dev builder images to CN region runs-on: ubuntu-latest @@ -209,6 +244,19 @@ jobs: quay.io/skopeo/stable:latest \ copy -a docker://docker.io/$IMAGE_NAMESPACE/dev-builder-android:$IMAGE_VERSION \ docker://$ACR_IMAGE_REGISTRY/$IMAGE_NAMESPACE/dev-builder-android:$IMAGE_VERSION + - name: Push dev-builder-riscv64 image + shell: bash + if: ${{ inputs.release_dev_builder_riscv64_image || github.event_name == 'push' }} + env: + IMAGE_VERSION: ${{ needs.release-dev-builder-images.outputs.version }} + IMAGE_NAMESPACE: ${{ vars.IMAGE_NAMESPACE }} + ACR_IMAGE_REGISTRY: ${{ vars.ACR_IMAGE_REGISTRY }} + run: | + docker run -v "${DOCKER_CONFIG:-$HOME/.docker}:/root/.docker:ro" \ + -e "REGISTRY_AUTH_FILE=/root/.docker/config.json" \ + quay.io/skopeo/stable:latest \ + copy -a docker://docker.io/$IMAGE_NAMESPACE/dev-builder-riscv64:$IMAGE_VERSION \ + docker://$ACR_IMAGE_REGISTRY/$IMAGE_NAMESPACE/dev-builder-riscv64:$IMAGE_VERSION update-dev-builder-image-tag: name: Update dev-builder image tag diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index dbb0f3c3146..52d62299a07 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -52,7 +52,7 @@ on: default: true build_linux_artifacts: type: boolean - description: Build linux artifacts (both amd64 and arm64) + description: Build linux artifacts (amd64, arm64, and riscv64) required: false default: false build_macos_artifacts: @@ -225,6 +225,52 @@ jobs: image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} large-page-size: true + build-linux-riscv64-artifacts: + name: Build linux-riscv64 artifacts + if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + needs: [ + allocate-runners, + ] + runs-on: ${{ needs.allocate-runners.outputs.linux-amd64-runner }} + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Prepare riscv64 dev-builder image + shell: bash + env: + IMAGE_REGISTRY: ${{ vars.ECR_IMAGE_REGISTRY }} + IMAGE_NAMESPACE: ${{ vars.ECR_IMAGE_NAMESPACE }} + run: | + builder_tag=$(grep '^DEV_BUILDER_IMAGE_TAG ?=' Makefile | awk '{print $3}') + builder_image="${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-riscv64:${builder_tag}" + if ! docker pull "${builder_image}"; then + make dev-builder-riscv64 \ + IMAGE_REGISTRY="${IMAGE_REGISTRY}" \ + IMAGE_NAMESPACE="${IMAGE_NAMESPACE}" + fi + + # Cross-compile greptime for riscv64gc-unknown-linux-gnu on the amd64 runner + # with the riscv64 cross dev-builder image. Integration tests are skipped + # because the cross-compiled binary cannot run on the amd64 host. + - name: Build greptime riscv64 binary + shell: bash + run: | + make build-riscv64-bin \ + CARGO_PROFILE=${{ env.CARGO_PROFILE }} \ + FEATURES=servers/dashboard \ + IMAGE_REGISTRY=${{ vars.ECR_IMAGE_REGISTRY }} \ + IMAGE_NAMESPACE=${{ vars.ECR_IMAGE_NAMESPACE }} + + - name: Upload artifacts + uses: ./.github/actions/upload-artifacts + with: + artifacts-dir: greptime-linux-riscv64-${{ needs.allocate-runners.outputs.version }} + target-files: ./target/riscv64gc-unknown-linux-gnu/${{ env.CARGO_PROFILE }}/greptime + version: ${{ needs.allocate-runners.outputs.version }} + build-macos-artifacts: name: Build macOS artifacts strategy: @@ -520,6 +566,7 @@ jobs: build-linux-amd64-artifacts, build-linux-arm64-artifacts, build-macos-artifacts, + build-linux-riscv64-artifacts, build-windows-artifacts, release-images-to-dockerhub, ] @@ -568,6 +615,7 @@ jobs: (needs.prepare-release-validation.outputs.run-query-regression == 'false' && needs.query-regression-release.result == 'skipped')) && needs.build-linux-amd64-artifacts.result == 'success' && needs.build-linux-arm64-artifacts.result == 'success' && + needs.build-linux-riscv64-artifacts.result == 'success' && (needs.build-macos-artifacts.result == 'success' || needs.build-macos-artifacts.result == 'skipped') && (needs.build-windows-artifacts.result == 'success' || needs.build-windows-artifacts.result == 'skipped') && (needs.release-images-to-dockerhub.result == 'success' || @@ -581,6 +629,7 @@ jobs: build-linux-amd64-artifacts, build-linux-arm64-artifacts, build-macos-artifacts, + build-linux-riscv64-artifacts, build-windows-artifacts, release-images-to-dockerhub, prepare-release-validation, @@ -619,6 +668,7 @@ jobs: needs: [ allocate-runners, build-linux-amd64-artifacts, + build-linux-riscv64-artifacts, ] steps: - name: Checkout diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index b5bfa1a2968..6d12566b684 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -96,6 +96,32 @@ jobs: - name: Check unused dependencies run: make check-udeps + check-riscv64: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} + name: Check riscv64 build + runs-on: ubuntu-latest + timeout-minutes: 90 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: arduino/setup-protoc@v3 + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + - name: Install riscv64 cross toolchain + run: sudo apt-get update && sudo apt-get install -y gcc-riscv64-linux-gnu g++-riscv64-linux-gnu + - uses: actions-rust-lang/setup-rust-toolchain@v1 + - name: Add riscv64 target + run: rustup target add riscv64gc-unknown-linux-gnu + - name: Rust Cache + uses: Swatinem/rust-cache@v2 + with: + shared-key: "check-riscv64" + cache-all-crates: "true" + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Run cargo check for riscv64 + run: cargo check --locked --workspace --all-targets --features servers/dashboard --target riscv64gc-unknown-linux-gnu + conflict-check: if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} name: Check for conflict diff --git a/Makefile b/Makefile index de93a2f113a..5bdb63d3520 100644 --- a/Makefile +++ b/Makefile @@ -9,6 +9,9 @@ IMAGE_REGISTRY ?= docker.io IMAGE_NAMESPACE ?= greptime IMAGE_TAG ?= latest DEV_BUILDER_IMAGE_TAG ?= 2026-03-21-9c9d9e9e-20260331090344 +DEV_BUILDER_RISCV64_IMAGE_TAG ?= $(DEV_BUILDER_IMAGE_TAG) +DEV_BUILDER_RISCV64_BASE_IMAGE ?= ubuntu:22.04 +RISCV64_TARGET ?= riscv64gc-unknown-linux-gnu BUILDX_MULTI_PLATFORM_BUILD ?= false BUILDX_BUILDER_NAME ?= gtbuilder BASE_IMAGE ?= ubuntu @@ -98,6 +101,21 @@ build-by-dev-builder: ## Build greptime by dev-builder. RELEASE=${RELEASE} \ CARGO_BUILD_EXTRA_OPTS="${CARGO_BUILD_EXTRA_OPTS}" +.PHONY: build-riscv64-bin +build-riscv64-bin: ## Build greptime binary for riscv64 (linux-gnu) by the riscv64 cross dev-builder image. + docker run --network=host \ + ${ASSEMBLED_EXTRA_BUILD_ENV} \ + -v ${PWD}:/greptimedb -v ${CARGO_REGISTRY_CACHE}:/root/.cargo/registry -v ${CARGO_GIT_CACHE}:/root/.cargo/git \ + -w /greptimedb ${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-riscv64:${DEV_BUILDER_RISCV64_IMAGE_TAG} \ + make build \ + CARGO_EXTENSION="${CARGO_EXTENSION}" \ + CARGO_PROFILE=${CARGO_PROFILE} \ + FEATURES=${FEATURES} \ + TARGET_DIR=${TARGET_DIR} \ + TARGET=${RISCV64_TARGET} \ + RELEASE=${RELEASE} \ + CARGO_BUILD_EXTRA_OPTS="${CARGO_BUILD_EXTRA_OPTS}" + .PHONY: build-android-bin build-android-bin: ## Build greptime binary for android. docker run --network=host \ @@ -158,6 +176,14 @@ dev-builder: multi-platform-buildx ## Build dev-builder image. -f docker/dev-builder/${BASE_IMAGE}/Dockerfile \ -t ${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-${BASE_IMAGE}:${DEV_BUILDER_IMAGE_TAG} ${BUILDX_MULTI_PLATFORM_BUILD_OPTS} . +.PHONY: dev-builder-riscv64 +dev-builder-riscv64: ## Build the riscv64 cross dev-builder image. + docker build --network=host \ + --build-arg="BASE_IMAGE=${DEV_BUILDER_RISCV64_BASE_IMAGE}" \ + --build-arg="RUST_TOOLCHAIN=${RUST_TOOLCHAIN}" \ + -f docker/dev-builder/riscv64/Dockerfile \ + -t ${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-riscv64:${DEV_BUILDER_RISCV64_IMAGE_TAG} . + .PHONY: multi-platform-buildx multi-platform-buildx: ## Create buildx multi-platform builder. docker buildx inspect ${BUILDX_BUILDER_NAME} || docker buildx create --name ${BUILDX_BUILDER_NAME} --driver docker-container --bootstrap --use diff --git a/docker/dev-builder/riscv64/Dockerfile b/docker/dev-builder/riscv64/Dockerfile new file mode 100644 index 00000000000..8505fa3b49a --- /dev/null +++ b/docker/dev-builder/riscv64/Dockerfile @@ -0,0 +1,55 @@ +# Cross-build environment for producing riscv64 (riscv64gc-unknown-linux-gnu) +# greptime binaries on x86_64/aarch64 hosts. +ARG BASE_IMAGE=ubuntu:22.04 +FROM ${BASE_IMAGE} + +ENV LANG=en_US.utf8 +WORKDIR /greptimedb + +RUN apt-get update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y software-properties-common +# Install dependencies, including the riscv64 cross toolchain. +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \ + libssl-dev \ + tzdata \ + curl \ + unzip \ + ca-certificates \ + git \ + build-essential \ + pkg-config \ + gcc-riscv64-linux-gnu \ + g++-riscv64-linux-gnu + +ARG TARGETPLATFORM +RUN echo "host platform: $TARGETPLATFORM" + +ARG PROTOBUF_VERSION=29.3 + +# Install protobuf for the host arch (build scripts and protoc run on the host), +# because the one in the apt is too old (v3.12). +RUN if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \ + curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOBUF_VERSION}/protoc-${PROTOBUF_VERSION}-linux-aarch_64.zip && \ + unzip -o protoc-${PROTOBUF_VERSION}-linux-aarch_64.zip -d protoc3; \ +else \ + curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOBUF_VERSION}/protoc-${PROTOBUF_VERSION}-linux-x86_64.zip && \ + unzip -o protoc-${PROTOBUF_VERSION}-linux-x86_64.zip -d protoc3; \ +fi +RUN mv -f protoc3/bin/* /usr/local/bin/ +# Use cp -r instead of mv so that it also works when the base image already +# contains protobuf headers under /usr/local/include. +RUN cp -r protoc3/include/* /usr/local/include/ + +# Silence all `safe.directory` warnings, to avoid the "detect dubious repository" +# error when building with the repository mounted from a different user. +RUN git config --global --add safe.directory '*' + +# Install Rust. +SHELL ["/bin/bash", "-c"] +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- --no-modify-path --default-toolchain none -y +ENV PATH=/root/.cargo/bin/:$PATH + +# Install Rust toolchains and the riscv64 target std. +ARG RUST_TOOLCHAIN +RUN rustup toolchain install ${RUST_TOOLCHAIN} +RUN rustup target add --toolchain ${RUST_TOOLCHAIN} riscv64gc-unknown-linux-gnu From 9fcfeff9ce876522fe2c47d2c47d7ea5fffa3c89 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 10 Aug 2026 16:28:11 +0800 Subject: [PATCH 030/126] feat(auth): support HTTP bearer-token authentication (#8719) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(auth): support HTTP bearer-token authentication (#8718) Adds an opt-in bearer-token (JWT / OAuth2) authentication path to the HTTP layer, so clients can authenticate with `Authorization: Bearer ` against any `/v1/` interface. Today such requests are rejected with `UnsupportedAuthScheme("bearer")` -> 401 before any handler runs. The token is treated as opaque by the server; validation and identity derivation stay in the UserProvider, so JWT/JWKS/OIDC policy remains pluggable and out of core. Changes: - `auth::UserProvider` gains `auth_token(token, catalog, schema) -> Result` with a default that rejects (`Error::UnsupportedAuthMethod`), so password-only providers keep today's behavior. A provider that supports token auth overrides it to validate the token, resolve it to a user, and authorize the connection. - `auth::Error::UnsupportedAuthMethod` for the default-reject case. - `servers::http::authorize::inner_auth` extracts a bearer token (`extract_bearer_token`) and, when present, authenticates via `UserProvider::auth_token`; otherwise it falls through unchanged to the username/password path (Basic / influxdb / splunk). Basic and bearer coexist on the same server. Backward compatible: the default impl preserves existing behavior, and non-bearer requests take the exact same path as before. Tests: - `extract_bearer_token` recognizes `Bearer` (either header) and ignores Basic/Token/Splunk/empty. - `inner_auth` dispatches a bearer token to `auth_token` and populates the QueryContext user on success; rejects on failure. - A password-only provider (default `auth_token`) rejects bearer tokens. Refs: #8718 * chore: fmt * refactor(auth): address bearer-auth review feedback (#8719) Address the review comments on the HTTP bearer-token authentication PR: - Match the `Bearer` scheme case-insensitively (RFC 9110 §11.1) via `eq_ignore_ascii_case`. `extract_bearer_token` previously only accepted `Bearer`/`bearer`, so a valid `BEARER ` fell through to `UnsupportedAuthScheme` and never reached the provider. The opaque token itself is deliberately not lowercased. - Return `Option<&str>` (borrowing the request headers) instead of `Option`, avoiding an allocation per bearer request. - Rename `UserProvider::auth_token` -> `auth_bearer_token` for clarity. - Route bearer-auth failures on Splunk HEC requests through `splunk_hec_err` (FORBIDDEN, code 4) instead of the generic 401 `ErrorResponse`, so HEC clients retain their `{"text":"Invalid token","code":4}` endpoint contract. Adds test coverage for case-insensitive scheme parsing (token preserved verbatim) and a regression test for the bearer/splunk routing path. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- src/auth/src/error.rs | 8 ++ src/auth/src/user_provider.rs | 29 +++- src/servers/src/http/authorize.rs | 226 +++++++++++++++++++++++++++++- 3 files changed, 259 insertions(+), 4 deletions(-) diff --git a/src/auth/src/error.rs b/src/auth/src/error.rs index 2b339b86554..f99a485db4e 100644 --- a/src/auth/src/error.rs +++ b/src/auth/src/error.rs @@ -88,6 +88,13 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + + #[snafu(display("Unsupported authentication method: {}", method))] + UnsupportedAuthMethod { + method: String, + #[snafu(implicit)] + location: Location, + }, } impl ErrorExt for Error { @@ -105,6 +112,7 @@ impl ErrorExt for Error { Error::UserPasswordMismatch { .. } => StatusCode::UserPasswordMismatch, Error::AccessDenied { .. } => StatusCode::AccessDenied, Error::PermissionDenied { .. } => StatusCode::PermissionDenied, + Error::UnsupportedAuthMethod { .. } => StatusCode::UserPasswordMismatch, } } diff --git a/src/auth/src/user_provider.rs b/src/auth/src/user_provider.rs index 60c178fb598..985e25978e8 100644 --- a/src/auth/src/user_provider.rs +++ b/src/auth/src/user_provider.rs @@ -34,8 +34,8 @@ use crate::common::{ PgScramSha256Verifier, auth_mysql_with_hash_stage_2, }; use crate::error::{ - IllegalParamSnafu, InvalidConfigSnafu, IoSnafu, Result, UnsupportedPasswordTypeSnafu, - UserNotFoundSnafu, UserPasswordMismatchSnafu, + IllegalParamSnafu, InvalidConfigSnafu, IoSnafu, Result, UnsupportedAuthMethodSnafu, + UnsupportedPasswordTypeSnafu, UserNotFoundSnafu, UserPasswordMismatchSnafu, }; use crate::user_info::{DefaultUserInfo, PermissionMode}; use crate::{UserInfoRef, auth_mysql}; @@ -66,6 +66,31 @@ pub trait UserProvider: Send + Sync { Ok(user_info) } + /// Authenticates and authorizes an opaque bearer token (e.g. a JWT or an + /// OAuth2 access token). + /// + /// Unlike [auth()](Self::auth), the caller has no `Identity`/`Password` — + /// the provider validates the token and *derives* the identity from it. + /// The token is opaque to the server, so JWT/JWKS/OIDC validation policy + /// stays pluggable and out of core. + /// + /// The default rejects token auth with + /// [`Error::UnsupportedAuthMethod`], so password-only providers keep + /// today's behavior. Providers that support token auth override this to + /// validate the token, resolve it to a user, and + /// [`authorize`](Self::authorize) the connection. + async fn auth_bearer_token( + &self, + _token: &str, + _catalog: &str, + _schema: &str, + ) -> Result { + UnsupportedAuthMethodSnafu { + method: "bearer token", + } + .fail() + } + async fn postgres_auth_info(&self, _id: Identity<'_>) -> Result { Ok(PgAuthInfo::Cleartext) } diff --git a/src/servers/src/http/authorize.rs b/src/servers/src/http/authorize.rs index a6d8ce2438e..755bd4d4c24 100644 --- a/src/servers/src/http/authorize.rs +++ b/src/servers/src/http/authorize.rs @@ -80,7 +80,36 @@ pub async fn inner_auth( return Ok(req); }; - // 3. get username and pwd + // 3. bearer token auth (JWT / OAuth2). When an `Authorization: Bearer + // ` header is present, authenticate via + // [`UserProvider::auth_bearer_token`]; otherwise fall through to the + // username/password path (Basic / influxdb / splunk) below. + if let Some(token) = extract_bearer_token(&req) { + match user_provider + .auth_bearer_token(token, &catalog, &schema) + .await + { + Ok(userinfo) => { + query_ctx.set_current_user(userinfo); + let _ = req.extensions_mut().insert(query_ctx); + return Ok(req); + } + Err(e) => { + warn!(e; "bearer token authentication failed"); + crate::metrics::METRIC_AUTH_FAILURE + .with_label_values(&[e.status_code().as_ref()]) + .inc(); + // Splunk HEC clients expect `{"text":"Invalid token","code":4}` + // (FORBIDDEN), not the generic 401 `ErrorResponse`. + if is_splunk_request(&req) { + return Err(splunk_hec_err(StatusCode::FORBIDDEN, 4)); + } + return Err(err_response(e)); + } + } + } + + // 4. get username and pwd let (username, password) = match extract_username_and_password(&req) { Ok((username, password)) => (username, password), Err(e) => { @@ -100,7 +129,7 @@ pub async fn inner_auth( } }; - // 4. auth + // 5. auth match user_provider .auth( auth::Identity::UserId(&username, None), @@ -307,6 +336,29 @@ impl From for api::v1::auth_header::AuthScheme { type Credential<'a> = &'a str; +/// Extracts an opaque bearer token from an `Authorization: Bearer ` +/// header (the standard or `x-greptime-auth` header). +/// +/// Returns `None` for any other scheme (`Basic`, influxdb `Token`, splunk +/// `Splunk`, …) so the caller can fall through to the username/password path. +fn extract_bearer_token(req: &Request) -> Option<&str> { + let header = req + .headers() + .get(AUTHORIZATION_HEADER) + .or_else(|| req.headers().get(http::header::AUTHORIZATION))?; + let value = header.to_str().ok()?; + // HTTP authentication schemes are case-insensitive (RFC 9110 §11.1), so + // match the scheme with `eq_ignore_ascii_case` — but never lowercase the + // *token* itself, which is opaque and case-sensitive. Returns a borrow + // into the request's headers, so there is no allocation. + let (scheme, token) = value.split_once(' ')?; + if !scheme.eq_ignore_ascii_case("bearer") { + return None; + } + let token = token.trim_start(); + (!token.is_empty()).then_some(token) +} + fn auth_header(req: &Request) -> Result { let auth_header = req .headers() @@ -610,4 +662,174 @@ mod tests { (Some("123"), Some("4")) ); } + + #[test] + fn test_extract_bearer_token() { + let bearer = |scheme: &str, val: &str| { + mock_http_request(Some(&format!("{scheme} {val}")), None).unwrap() + }; + + // Standard bearer scheme, on either header. + assert_eq!( + extract_bearer_token(&bearer("Bearer", "abc.def.ghi")), + Some("abc.def.ghi") + ); + assert_eq!(extract_bearer_token(&bearer("bearer", "tok")), Some("tok")); + // HTTP schemes are case-insensitive (RFC 9110 §11.1); the token + // itself is opaque and must NOT be lowercased. + assert_eq!( + extract_bearer_token(&bearer("BEARER", "ABC.DEF.GHI")), + Some("ABC.DEF.GHI") + ); + assert_eq!(extract_bearer_token(&bearer("BeArEr", "tok")), Some("tok")); + let mut req = mock_http_request(Some("Bearer xyz"), None).unwrap(); + req.headers_mut().insert( + AUTHORIZATION_HEADER, + "Bearer from-x-greptime".parse().unwrap(), + ); + assert_eq!(extract_bearer_token(&req), Some("from-x-greptime")); + + // Non-bearer schemes are ignored so the caller falls through to Basic. + assert_eq!(extract_bearer_token(&bearer("Basic", "dXNlcjpwYXNz")), None); + assert_eq!(extract_bearer_token(&bearer("Token", "u:p")), None); + assert_eq!(extract_bearer_token(&bearer("Splunk", "u:p")), None); + + // No header, empty token. + assert_eq!( + extract_bearer_token(&mock_http_request(None, None).unwrap()), + None + ); + assert_eq!(extract_bearer_token(&bearer("Bearer", "")), None); + } + + /// A `UserProvider` that resolves exactly one bearer token to a known user + /// and rejects everything else (including password auth). + struct TokenUserProvider { + token: String, + user: auth::UserInfoRef, + } + + #[async_trait::async_trait] + impl auth::UserProvider for TokenUserProvider { + fn name(&self) -> &str { + "token-test" + } + + async fn authenticate( + &self, + _: auth::Identity<'_>, + _: auth::Password<'_>, + ) -> auth::error::Result { + unreachable!("password auth should not be reached for a bearer request") + } + + async fn authorize( + &self, + _: &str, + _: &str, + _: &auth::UserInfoRef, + ) -> auth::error::Result<()> { + Ok(()) + } + + async fn auth_bearer_token( + &self, + token: &str, + _: &str, + _: &str, + ) -> auth::error::Result { + if token == self.token { + Ok(self.user.clone()) + } else { + auth::error::UnsupportedAuthMethodSnafu { + method: "bearer token", + } + .fail() + } + } + } + + #[tokio::test] + async fn test_bearer_token_dispatches_to_auth_bearer_token() { + let provider = TokenUserProvider { + token: "good-token".to_string(), + user: auth::userinfo_by_name(Some("alice".into())), + }; + let req = mock_http_request(Some("Bearer good-token"), None).unwrap(); + + let req = inner_auth::<()>( + Some(std::sync::Arc::new(provider) as auth::UserProviderRef), + req, + ) + .await + .expect("valid bearer token authenticates"); + + let user = req + .extensions() + .get::() + .expect("query context is populated") + .current_user(); + assert_eq!(user.username(), "alice"); + } + + #[tokio::test] + async fn test_bearer_token_failure_rejects() { + let provider = TokenUserProvider { + token: "good-token".to_string(), + user: auth::userinfo_by_name(Some("alice".into())), + }; + let req = mock_http_request(Some("Bearer bad-token"), None).unwrap(); + + let result = inner_auth::<()>( + Some(std::sync::Arc::new(provider) as auth::UserProviderRef), + req, + ) + .await; + assert!(result.is_err(), "an invalid bearer token is rejected"); + } + + #[tokio::test] + async fn test_default_provider_rejects_bearer() { + // A password-only provider uses the default `auth_bearer_token`, which rejects. + let provider = + auth::user_provider_from_option("static_user_provider:cmd:alice=s3cret").unwrap(); + let req = mock_http_request(Some("Bearer some-jwt"), None).unwrap(); + let result = inner_auth::<()>(Some(provider), req).await; + assert!( + result.is_err(), + "password-only providers reject bearer tokens" + ); + } + + /// A bearer-token failure on a Splunk HEC request must keep the HEC + /// contract — `{"text":"Invalid token","code":4}` with FORBIDDEN — + /// instead of falling through to the generic 401 `ErrorResponse`. + /// Regression for the bearer/splunk routing gap. + #[tokio::test] + async fn test_bearer_failure_on_splunk_keeps_hec_contract() { + let provider = TokenUserProvider { + token: "good-token".to_string(), + user: auth::userinfo_by_name(Some("alice".into())), + }; + let req = mock_http_request( + Some("Bearer bad-token"), + Some("http://127.0.0.1/v1/splunk/services/collector/event"), + ) + .unwrap(); + + let resp = inner_auth::<()>( + Some(std::sync::Arc::new(provider) as auth::UserProviderRef), + req, + ) + .await + .expect_err("an invalid bearer token is rejected"); + + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let payload: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(payload["code"], 4); + assert_eq!(payload["text"], "Invalid token"); + } } From 4ce49ec33eb0b416cbd86f737631a0c37cf12774 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:06:20 +0800 Subject: [PATCH 031/126] feat(promql): support mixed sample ranges (#8784) * feat(promql): support mixed sample ranges Signed-off-by: shuiyisong * chore: update test Signed-off-by: shuiyisong * chore: add comments Signed-off-by: shuiyisong * fix: CR issue Signed-off-by: shuiyisong * fix: return error than panic Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/promql/src/functions.rs | 2 +- src/promql/src/functions/native_histogram.rs | 997 +++++++++++++++++- src/query/src/promql/planner.rs | 343 +++++- .../src/query_engine/default_serializer.rs | 109 +- 4 files changed, 1419 insertions(+), 32 deletions(-) diff --git a/src/promql/src/functions.rs b/src/promql/src/functions.rs index b534bb06928..b34c21c5569 100644 --- a/src/promql/src/functions.rs +++ b/src/promql/src/functions.rs @@ -45,7 +45,7 @@ pub use double_exponential_smoothing::DoubleExponentialSmoothing; pub use extrapolate_rate::{Delta, Increase, Rate}; pub use idelta::IDelta; pub use native_histogram::{ - NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, + MixedRange, NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index 4a4e3adad5c..1863e663b82 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -25,7 +25,8 @@ use datafusion::arrow::array::{ Array, ArrayRef, BooleanArray, Float64Array, Float64Builder, Int64Array, StringBuilder, StructArray, TimestampMillisecondArray, UInt64Array, }; -use datafusion::arrow::datatypes::{DataType, TimeUnit}; +use datafusion::arrow::compute::filter; +use datafusion::arrow::datatypes::{DataType, Field, TimeUnit}; use datafusion::common::{DataFusionError, Result as DfResult}; use datafusion::logical_expr::{Accumulator as DfAccumulator, AggregateUDF, ScalarUDF, Volatility}; use datafusion::physical_plan::ColumnarValue; @@ -33,7 +34,11 @@ use datafusion_common::ScalarValue; use datafusion_expr::function::AccumulatorArgs; use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, create_udaf, create_udf}; -use crate::functions::{extract_array, extract_range_dict}; +use crate::functions::{ + AvgOverTime, Deriv, DoubleExponentialSmoothing, IDelta, Increase, LastOverTime, MaxOverTime, + MinOverTime, PredictLinear, QuantileOverTime, Rate, StddevOverTime, StdvarOverTime, + SumOverTime, extract_array, extract_range_dict, +}; use crate::range_array::{RangeArray, unpack}; fn extract_histogram_array(value: &ColumnarValue, func_name: &str) -> DfResult { @@ -1264,6 +1269,783 @@ impl NativeHistogramResets { } } +/// Coordinated float/native-histogram range evaluation. +/// +/// The function name is passed as the first scalar argument so these two UDF names are enough for +/// distributed plan decoding. The remaining leading arguments are timestamp, float, and histogram +/// ranges, followed by the ordinary function arguments. +pub struct MixedRange; + +impl MixedRange { + const fn float_name() -> &'static str { + "prom_mixed_range_float" + } + + const fn histogram_name() -> &'static str { + "prom_mixed_range_histogram" + } + + pub fn float_udf(collector: Option) -> ScalarUDF { + ScalarUDF::new_from_impl(MixedRangeUdf::new(MixedRangeOutput::Float, collector)) + } + + pub fn histogram_udf(collector: Option) -> ScalarUDF { + ScalarUDF::new_from_impl(MixedRangeUdf::new(MixedRangeOutput::Histogram, collector)) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum MixedRangeOutput { + Float, + Histogram, +} + +impl MixedRangeOutput { + fn name(self) -> &'static str { + match self { + Self::Float => MixedRange::float_name(), + Self::Histogram => MixedRange::histogram_name(), + } + } + + fn data_type(self) -> DataType { + match self { + Self::Float => DataType::Float64, + Self::Histogram => native_histogram_arrow_type(), + } + } +} + +#[derive(Debug, Clone)] +struct MixedRangeUdf { + output: MixedRangeOutput, + signature: Signature, + collector: Option, +} + +impl MixedRangeUdf { + fn new(output: MixedRangeOutput, collector: Option) -> Self { + Self { + output, + signature: Signature::variadic_any(Volatility::Volatile), + collector, + } + } +} + +impl PartialEq for MixedRangeUdf { + fn eq(&self, other: &Self) -> bool { + self.output == other.output + } +} + +impl Eq for MixedRangeUdf {} + +impl Hash for MixedRangeUdf { + fn hash(&self, state: &mut H) { + self.output.hash(state); + } +} + +impl ScalarUDFImpl for MixedRangeUdf { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + self.output.name() + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DfResult { + Ok(self.output.data_type()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult { + let collector = args + .config_options + .extensions + .get::() + .cloned() + .or_else(|| self.collector.clone()); + mixed_range(&args, self.output, collector) + } +} + +#[derive(Debug, Clone, Copy)] +enum MixedRangeFunction { + Rate, + Increase, + Delta, + IDelta, + IRate, + Changes, + Resets, + AvgOverTime, + MinOverTime, + MaxOverTime, + SumOverTime, + CountOverTime, + LastOverTime, + AbsentOverTime, + PresentOverTime, + StddevOverTime, + StdvarOverTime, + QuantileOverTime, + Deriv, + PredictLinear, + DoubleExponentialSmoothing, + HoltWinters, +} + +impl MixedRangeFunction { + fn parse(value: &ColumnarValue) -> DfResult { + let name = match value { + ColumnarValue::Scalar(ScalarValue::Utf8(Some(name))) + | ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(name))) + | ColumnarValue::Scalar(ScalarValue::Utf8View(Some(name))) => name.as_str(), + other => { + return Err(DataFusionError::Execution(format!( + "mixed range function name must be a non-null string scalar, found {}", + other.data_type() + ))); + } + }; + + match name { + "rate" => Ok(Self::Rate), + "increase" => Ok(Self::Increase), + "delta" => Ok(Self::Delta), + "idelta" => Ok(Self::IDelta), + "irate" => Ok(Self::IRate), + "changes" => Ok(Self::Changes), + "resets" => Ok(Self::Resets), + "avg_over_time" => Ok(Self::AvgOverTime), + "min_over_time" => Ok(Self::MinOverTime), + "max_over_time" => Ok(Self::MaxOverTime), + "sum_over_time" => Ok(Self::SumOverTime), + "count_over_time" => Ok(Self::CountOverTime), + "last_over_time" => Ok(Self::LastOverTime), + "absent_over_time" => Ok(Self::AbsentOverTime), + "present_over_time" => Ok(Self::PresentOverTime), + "stddev_over_time" => Ok(Self::StddevOverTime), + "stdvar_over_time" => Ok(Self::StdvarOverTime), + "quantile_over_time" => Ok(Self::QuantileOverTime), + "deriv" => Ok(Self::Deriv), + "predict_linear" => Ok(Self::PredictLinear), + "double_exponential_smoothing" => Ok(Self::DoubleExponentialSmoothing), + "holt_winters" => Ok(Self::HoltWinters), + _ => Err(DataFusionError::Execution(format!( + "unsupported mixed range function: {name}" + ))), + } + } + + fn name(self) -> &'static str { + match self { + Self::Rate => "rate", + Self::Increase => "increase", + Self::Delta => "delta", + Self::IDelta => "idelta", + Self::IRate => "irate", + Self::Changes => "changes", + Self::Resets => "resets", + Self::AvgOverTime => "avg_over_time", + Self::MinOverTime => "min_over_time", + Self::MaxOverTime => "max_over_time", + Self::SumOverTime => "sum_over_time", + Self::CountOverTime => "count_over_time", + Self::LastOverTime => "last_over_time", + Self::AbsentOverTime => "absent_over_time", + Self::PresentOverTime => "present_over_time", + Self::StddevOverTime => "stddev_over_time", + Self::StdvarOverTime => "stdvar_over_time", + Self::QuantileOverTime => "quantile_over_time", + Self::Deriv => "deriv", + Self::PredictLinear => "predict_linear", + Self::DoubleExponentialSmoothing => "double_exponential_smoothing", + Self::HoltWinters => "holt_winters", + } + } + + fn policy(self) -> MixedRangePolicy { + match self { + Self::Rate | Self::Increase | Self::Delta | Self::AvgOverTime | Self::SumOverTime => { + MixedRangePolicy::DropMixed + } + Self::IDelta | Self::IRate => MixedRangePolicy::LastTwo, + Self::LastOverTime => MixedRangePolicy::Last, + Self::Changes + | Self::Resets + | Self::CountOverTime + | Self::AbsentOverTime + | Self::PresentOverTime => MixedRangePolicy::Combined, + Self::MinOverTime + | Self::MaxOverTime + | Self::StddevOverTime + | Self::StdvarOverTime + | Self::QuantileOverTime + | Self::Deriv + | Self::PredictLinear + | Self::DoubleExponentialSmoothing + | Self::HoltWinters => MixedRangePolicy::FloatOnly, + } + } + + fn float_udf(self) -> Option { + match self { + Self::Rate => Some(Rate::scalar_udf()), + Self::Increase => Some(Increase::scalar_udf()), + Self::Delta => Some(crate::functions::Delta::scalar_udf()), + Self::IDelta => Some(IDelta::::scalar_udf()), + Self::IRate => Some(IDelta::::scalar_udf()), + Self::AvgOverTime => Some(AvgOverTime::scalar_udf()), + Self::MinOverTime => Some(MinOverTime::scalar_udf()), + Self::MaxOverTime => Some(MaxOverTime::scalar_udf()), + Self::SumOverTime => Some(SumOverTime::scalar_udf()), + Self::LastOverTime => Some(LastOverTime::scalar_udf()), + Self::StddevOverTime => Some(StddevOverTime::scalar_udf()), + Self::StdvarOverTime => Some(StdvarOverTime::scalar_udf()), + Self::QuantileOverTime => Some(QuantileOverTime::scalar_udf()), + Self::Deriv => Some(Deriv::scalar_udf()), + Self::PredictLinear => Some(PredictLinear::scalar_udf()), + Self::DoubleExponentialSmoothing | Self::HoltWinters => { + Some(DoubleExponentialSmoothing::scalar_udf()) + } + Self::Changes + | Self::Resets + | Self::CountOverTime + | Self::AbsentOverTime + | Self::PresentOverTime => None, + } + } + + fn histogram_udf(self, collector: Option) -> Option { + match self { + Self::Rate => Some(NativeHistogramRate::scalar_udf_with_collector(collector)), + Self::Increase => Some(NativeHistogramIncrease::scalar_udf_with_collector( + collector, + )), + Self::Delta => Some(NativeHistogramDelta::scalar_udf_with_collector(collector)), + Self::IDelta => Some(NativeHistogramIDelta::scalar_udf_with_collector(collector)), + Self::IRate => Some(NativeHistogramIRate::scalar_udf_with_collector(collector)), + Self::AvgOverTime => Some(NativeHistogramAvgOverTime::scalar_udf_with_collector( + collector, + )), + Self::SumOverTime => Some(NativeHistogramSumOverTime::scalar_udf_with_collector( + collector, + )), + Self::LastOverTime => Some(NativeHistogramLastOverTime::scalar_udf()), + _ => None, + } + } +} + +#[derive(Debug, Clone, Copy)] +enum MixedRangePolicy { + DropMixed, + LastTwo, + Last, + Combined, + FloatOnly, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SampleLane { + Float, + Histogram, +} + +fn mixed_range( + args: &ScalarFunctionArgs, + output: MixedRangeOutput, + collector: Option, +) -> DfResult { + if args.args.len() < 4 { + return Err(DataFusionError::Plan(format!( + "{} function should have at least 4 inputs", + output.name() + ))); + } + let function = MixedRangeFunction::parse(&args.args[0])?; + let name = function.name(); + let ts_range = extract_range_dict( + &args.args[1], + name, + "timestamp range vector", + &DataType::Timestamp(TimeUnit::Millisecond, None), + )?; + let float_range = extract_range_dict( + &args.args[2], + name, + "float range vector", + &DataType::Float64, + )?; + let histogram_range = extract_range_dict( + &args.args[3], + name, + "native histogram range vector", + &native_histogram_arrow_type(), + )?; + let keys = ts_range.keys().values(); + if float_range.keys().values() != keys || histogram_range.keys().values() != keys { + return Err(DataFusionError::Execution(format!( + "{name}: timestamp, float, and native histogram ranges should have the same layout" + ))); + } + if args.number_rows != keys.len() { + return Err(DataFusionError::Execution(format!( + "{name}: range inputs have {} windows but the batch has {} rows", + keys.len(), + args.number_rows + ))); + } + + let timestamps = ts_range + .values() + .as_any() + .downcast_ref::() + .expect("validated timestamp range"); + let floats = float_range + .values() + .as_any() + .downcast_ref::() + .expect("validated float range"); + let histograms = histogram_range + .values() + .as_any() + .downcast_ref::() + .expect("validated native histogram range"); + if floats.len() != timestamps.len() || histograms.len() != timestamps.len() { + return Err(DataFusionError::Execution(format!( + "{name}: timestamp, float, and native histogram values should be row-aligned" + ))); + } + + let bounds = checked_window_bounds(keys, timestamps.len(), name)?; + let float_valid = (0..floats.len()) + .map(|row| floats.is_valid(row)) + .collect::>(); + let histogram_valid = (0..histograms.len()) + .map(|row| histograms.is_valid(row)) + .collect::>(); + let float_prefix = validity_prefix(&float_valid, name)?; + let histogram_prefix = validity_prefix(&histogram_valid, name)?; + + if matches!(function.policy(), MixedRangePolicy::Combined) { + if output != MixedRangeOutput::Float { + return Err(DataFusionError::Execution(format!( + "{name} does not return native histograms" + ))); + } + return combined_range_float( + function, + &bounds, + timestamps, + floats, + histograms, + &float_valid, + &histogram_valid, + &float_prefix, + &histogram_prefix, + ); + } + + let selections = select_mixed_range_lanes( + function, + &bounds, + timestamps, + &float_valid, + &histogram_valid, + &float_prefix, + &histogram_prefix, + &collector, + ); + let (lane, values, valid, prefix) = match output { + MixedRangeOutput::Float => ( + SampleLane::Float, + floats as &dyn Array, + float_valid.as_slice(), + float_prefix.as_slice(), + ), + MixedRangeOutput::Histogram => ( + SampleLane::Histogram, + histograms as &dyn Array, + histogram_valid.as_slice(), + histogram_prefix.as_slice(), + ), + }; + let input = compact_lane_input( + timestamps, + values, + valid, + prefix, + &bounds, + &selections, + lane, + name, + )?; + + let udf = match output { + MixedRangeOutput::Float => function.float_udf(), + MixedRangeOutput::Histogram => function.histogram_udf(collector), + } + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{name} does not return {} values", + match output { + MixedRangeOutput::Float => "float", + MixedRangeOutput::Histogram => "native histogram", + } + )) + })?; + let mut input = input; + input.extend_from_slice(&args.args[4..]); + invoke_range_udf(udf, input, args, bounds.len()) +} + +/// Decodes packed range keys into validated half-open `(offset, end)` bounds. +fn checked_window_bounds( + keys: &[i64], + value_len: usize, + name: &str, +) -> DfResult> { + keys.iter() + .map(|key| { + let (offset, length) = unpack(*key); + let offset = offset as usize; + let end = offset + .checked_add(length as usize) + .filter(|end| *end <= value_len) + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{name}: invalid range ({offset}, {length}) for {value_len} values" + )) + })?; + Ok((offset, end)) + }) + .collect() +} + +/// Builds prefix counts where `prefix[i]` is the number of valid samples in `valid[..i]`. +fn validity_prefix(valid: &[bool], name: &str) -> DfResult> { + let capacity = valid.len().checked_add(1).ok_or_else(|| { + DataFusionError::Execution(format!("{name}: sample validity length overflow")) + })?; + let mut prefix = Vec::with_capacity(capacity); + prefix.push(0usize); + for is_valid in valid { + prefix.push( + prefix + .last() + .copied() + .unwrap() + .checked_add(usize::from(*is_valid)) + .ok_or_else(|| { + DataFusionError::Execution(format!("{name}: sample count overflow")) + })?, + ); + } + Ok(prefix) +} + +/// Returns the number of valid samples in the half-open window `[offset, end)`. +fn valid_count(prefix: &[usize], offset: usize, end: usize) -> usize { + prefix[end] + .checked_sub(prefix[offset]) + .expect("validity prefix is monotonic") +} + +/// Selects the sample lane for each window and records policy-required annotations. +#[allow(clippy::too_many_arguments)] +fn select_mixed_range_lanes( + function: MixedRangeFunction, + bounds: &[(usize, usize)], + timestamps: &TimestampMillisecondArray, + float_valid: &[bool], + histogram_valid: &[bool], + float_prefix: &[usize], + histogram_prefix: &[usize], + collector: &Option, +) -> Vec> { + bounds + .iter() + .map(|(offset, end)| { + let float_count = valid_count(float_prefix, *offset, *end); + let histogram_count = valid_count(histogram_prefix, *offset, *end); + match function.policy() { + MixedRangePolicy::DropMixed => match (float_count > 0, histogram_count > 0) { + (true, true) => { + record_warning( + collector, + format!( + "{}: encountered a mix of float and native histogram samples", + function.name() + ), + ); + None + } + (true, false) => Some(SampleLane::Float), + (false, true) => Some(SampleLane::Histogram), + (false, false) => None, + }, + MixedRangePolicy::FloatOnly => { + if float_count > 0 && histogram_count > 0 { + record_info( + collector, + format!( + "{}: ignored native histogram samples", + function.name() + ), + ); + } + (float_count > 0).then_some(SampleLane::Float) + } + MixedRangePolicy::Last => (*offset..*end).rev().find_map(|row| { + if histogram_valid[row] { + Some(SampleLane::Histogram) + } else if float_valid[row] { + Some(SampleLane::Float) + } else { + None + } + }), + MixedRangePolicy::LastTwo => { + let mut last_two = Vec::with_capacity(2); + for row in (*offset..*end).rev() { + // Prometheus keeps a float as the newest sample if both lanes have the + // same timestamp, while the histogram becomes the preceding sample. + if float_valid[row] { + last_two.push((SampleLane::Float, timestamps.value(row))); + } + if last_two.len() < 2 && histogram_valid[row] { + last_two.push((SampleLane::Histogram, timestamps.value(row))); + } + if last_two.len() == 2 { + break; + } + } + if last_two.len() < 2 || last_two[0].1 == last_two[1].1 { + None + } else if last_two[0].0 == last_two[1].0 { + Some(last_two[0].0) + } else { + record_warning( + collector, + format!( + "{}: encountered a mix of float and native histogram samples in the last two points", + function.name() + ), + ); + None + } + } + MixedRangePolicy::Combined => unreachable!(), + } + }) + .collect() +} + +/// Filters null placeholders from one sample lane and remaps its selected windows. +/// Windows assigned to the other lane become empty. +#[allow(clippy::too_many_arguments)] +fn compact_lane_input( + timestamps: &TimestampMillisecondArray, + values: &dyn Array, + valid: &[bool], + prefix: &[usize], + bounds: &[(usize, usize)], + selections: &[Option], + lane: SampleLane, + name: &str, +) -> DfResult> { + let mask = BooleanArray::from(valid.to_vec()); + let filtered_timestamps = filter(timestamps, &mask)?; + let filtered_values = filter(values, &mask)?; + let ranges = bounds + .iter() + .zip(selections) + .map(|((offset, end), selected)| { + let compact_offset = prefix[*offset]; + let compact_length = if *selected == Some(lane) { + valid_count(prefix, *offset, *end) + } else { + 0 + }; + Ok(( + u32::try_from(compact_offset).map_err(|_| { + DataFusionError::Execution(format!( + "{name}: compacted range offset exceeds u32" + )) + })?, + u32::try_from(compact_length).map_err(|_| { + DataFusionError::Execution(format!( + "{name}: compacted range length exceeds u32" + )) + })?, + )) + }) + .collect::>>()?; + let timestamp_range = RangeArray::from_ranges(filtered_timestamps, ranges.clone()) + .map_err(DataFusionError::from)?; + let value_range = + RangeArray::from_ranges(filtered_values, ranges).map_err(DataFusionError::from)?; + Ok(vec![ + ColumnarValue::Array(Arc::new(timestamp_range.into_dict())), + ColumnarValue::Array(Arc::new(value_range.into_dict())), + ]) +} + +fn invoke_range_udf( + udf: ScalarUDF, + input: Vec, + outer_args: &ScalarFunctionArgs, + number_rows: usize, +) -> DfResult { + let arg_fields = input + .iter() + .enumerate() + .map(|(index, value)| Arc::new(Field::new(format!("arg_{index}"), value.data_type(), true))) + .collect(); + udf.invoke_with_args(ScalarFunctionArgs { + args: input, + arg_fields, + number_rows, + return_field: outer_args.return_field.clone(), + config_options: outer_args.config_options.clone(), + }) +} + +enum MixedSample { + Float(f64), + Histogram(NativeHistogram), +} + +#[allow(clippy::too_many_arguments)] +fn combined_range_float( + function: MixedRangeFunction, + bounds: &[(usize, usize)], + timestamps: &TimestampMillisecondArray, + floats: &Float64Array, + histograms: &StructArray, + float_valid: &[bool], + histogram_valid: &[bool], + float_prefix: &[usize], + histogram_prefix: &[usize], +) -> DfResult { + let mut result = Float64Builder::with_capacity(bounds.len()); + for (offset, end) in bounds { + let sample_count = valid_count(float_prefix, *offset, *end) + .checked_add(valid_count(histogram_prefix, *offset, *end)) + .ok_or_else(|| { + DataFusionError::Execution(format!("{}: sample count overflow", function.name())) + })?; + match function { + MixedRangeFunction::CountOverTime => { + if sample_count == 0 { + result.append_null(); + } else { + result.append_value(sample_count as f64); + } + } + MixedRangeFunction::AbsentOverTime => { + if sample_count == 0 { + result.append_value(1.0); + } else { + result.append_null(); + } + } + MixedRangeFunction::PresentOverTime => { + if sample_count == 0 { + result.append_null(); + } else { + result.append_value(1.0); + } + } + MixedRangeFunction::Changes | MixedRangeFunction::Resets => { + if sample_count == 0 { + result.append_null(); + continue; + } + let mut count = 0usize; + let mut previous = None; + for row in *offset..*end { + if float_valid[row] { + count += mixed_transition( + function, + &mut previous, + timestamps.value(row), + MixedSample::Float(floats.value(row)), + ); + } + if histogram_valid[row] { + let histogram = read_histogram(histograms, row)? + .expect("validated native histogram sample"); + count += mixed_transition( + function, + &mut previous, + timestamps.value(row), + MixedSample::Histogram(histogram), + ); + } + } + result.append_value(count as f64); + } + _ => { + return Err(DataFusionError::Internal(format!( + "{} does not support combined range evaluation", + function.name() + ))); + } + } + } + Ok(ColumnarValue::Array(Arc::new(result.finish()))) +} + +fn mixed_transition( + function: MixedRangeFunction, + previous: &mut Option<(i64, MixedSample)>, + timestamp: i64, + current: MixedSample, +) -> usize { + let changed = previous.as_ref().is_some_and(|(previous_ts, previous)| { + match (function, previous, ¤t) { + ( + MixedRangeFunction::Changes, + MixedSample::Float(previous), + MixedSample::Float(current), + ) => current != previous && !(current.is_nan() && previous.is_nan()), + ( + MixedRangeFunction::Changes, + MixedSample::Histogram(previous), + MixedSample::Histogram(current), + ) => !current.promql_eq(previous), + (MixedRangeFunction::Changes, _, _) => true, + ( + MixedRangeFunction::Resets, + MixedSample::Float(previous), + MixedSample::Float(current), + ) => current < previous, + ( + MixedRangeFunction::Resets, + MixedSample::Histogram(previous), + MixedSample::Histogram(current), + ) => { + (previous.reset_hint == GAUGE_RESET_HINT) + != (current.reset_hint == GAUGE_RESET_HINT) + || current.detect_counter_reset(previous, *previous_ts, timestamp) + } + (MixedRangeFunction::Resets, _, _) => true, + _ => unreachable!(), + } + }); + *previous = Some((timestamp, current)); + usize::from(changed) +} + fn native_histogram_scalar(histogram: Option) -> ScalarValue { let array = build_histogram_array(&[histogram]); let histogram = array @@ -2024,6 +2806,53 @@ mod tests { ] } + fn mixed_range_input( + name: &str, + floats: Vec>, + histograms: Vec>, + ) -> Vec { + assert_eq!(floats.len(), histograms.len()); + let timestamps = Arc::new(TimestampMillisecondArray::from_iter( + (0..floats.len()).map(|idx| Some((idx as i64 + 1) * 1000)), + )); + let floats = Arc::new(Float64Array::from(floats)); + let histograms = build_histogram_array(&histograms); + let range = [(0, u32::try_from(floats.len()).unwrap())]; + vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(Some(name.to_string()))), + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(timestamps, range) + .unwrap() + .into_dict(), + )), + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(floats, range).unwrap().into_dict(), + )), + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(histograms, range) + .unwrap() + .into_dict(), + )), + ] + } + + fn mixed_float_result(udf: ScalarUDF, input: Vec) -> Option { + let result = run_udf(udf, input, DataType::Float64); + let result = extract_array(&result).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + result.is_valid(0).then(|| result.value(0)) + } + + fn mixed_histogram_result( + udf: ScalarUDF, + input: Vec, + ) -> Option { + let result = run_udf(udf, input, native_histogram_arrow_type()); + let result = extract_array(&result).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + read_histogram(result, 0).unwrap() + } + fn run_histogram_range_udf( udf: ScalarUDF, histograms: Vec, @@ -2098,6 +2927,170 @@ mod tests { warnings } + fn collected_infos(collector: &PromqlAnnotationCollector) -> Vec { + let mut infos = Vec::new(); + collector.append_to(&mut Vec::new(), &mut infos); + infos + } + + #[test] + fn mixed_ranges_follow_prometheus_sample_type_semantics() { + let first = sample_histogram(1.0, 1.0, vec![1.0]); + let second = sample_histogram(3.0, 3.0, vec![3.0]); + let collector = PromqlAnnotationCollector::default(); + + let mut rate = mixed_range_input( + "rate", + vec![Some(1.0), None, Some(3.0)], + vec![None, Some(first.clone()), None], + ); + rate.push(ColumnarValue::Array(Arc::new( + TimestampMillisecondArray::from(vec![3000]), + ))); + rate.push(ColumnarValue::Array(Arc::new(Int64Array::from(vec![3000])))); + assert_eq!( + mixed_float_result(MixedRange::float_udf(Some(collector.clone())), rate.clone()), + None + ); + assert_eq!( + mixed_histogram_result(MixedRange::histogram_udf(Some(collector.clone())), rate), + None + ); + assert!( + collected_warnings(&collector) + .iter() + .any(|warning| warning.contains("mix of float and native histogram")) + ); + + let pure_rate = |floats, histograms| { + let mut input = mixed_range_input("rate", floats, histograms); + input.push(ColumnarValue::Array(Arc::new( + TimestampMillisecondArray::from(vec![3000]), + ))); + input.push(ColumnarValue::Array(Arc::new(Int64Array::from(vec![3000])))); + input + }; + let pure_float = pure_rate( + vec![Some(1.0), Some(2.0), Some(3.0)], + vec![None, None, None], + ); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), pure_float.clone()), + Some(1.0) + ); + assert_eq!( + mixed_histogram_result(MixedRange::histogram_udf(None), pure_float), + None + ); + let pure_histogram = pure_rate( + vec![None, None, None], + vec![ + Some(sample_histogram(1.0, 1.0, vec![1.0])), + Some(sample_histogram(2.0, 2.0, vec![2.0])), + Some(sample_histogram(3.0, 3.0, vec![3.0])), + ], + ); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), pure_histogram.clone()), + None + ); + assert_eq!( + mixed_histogram_result(MixedRange::histogram_udf(None), pure_histogram) + .unwrap() + .count, + 1.0 + ); + + let idelta = mixed_range_input( + "idelta", + vec![Some(10.0), None, None], + vec![None, Some(first.clone()), Some(second.clone())], + ); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), idelta.clone()), + None + ); + assert_eq!( + mixed_histogram_result(MixedRange::histogram_udf(None), idelta) + .unwrap() + .count, + 2.0 + ); + + let alternating = || { + mixed_range_input( + "changes", + vec![Some(1.0), None, None, Some(1.0)], + vec![None, Some(first.clone()), Some(first.clone()), None], + ) + }; + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), alternating()), + Some(2.0) + ); + let mut resets = alternating(); + resets[0] = ColumnarValue::Scalar(ScalarValue::Utf8(Some("resets".to_string()))); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), resets), + Some(2.0) + ); + + for (name, expected) in [ + ("count_over_time", Some(4.0)), + ("present_over_time", Some(1.0)), + ("absent_over_time", None), + ] { + let mut input = alternating(); + input[0] = ColumnarValue::Scalar(ScalarValue::Utf8(Some(name.to_string()))); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), input), + expected, + "{name}" + ); + } + + let last = mixed_range_input( + "last_over_time", + vec![None, Some(4.0)], + vec![Some(first.clone()), None], + ); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), last.clone()), + Some(4.0) + ); + assert_eq!( + mixed_histogram_result(MixedRange::histogram_udf(None), last), + None + ); + + let collector = PromqlAnnotationCollector::default(); + let histogram_only_min = + mixed_range_input("min_over_time", vec![None], vec![Some(first.clone())]); + assert_eq!( + mixed_float_result( + MixedRange::float_udf(Some(collector.clone())), + histogram_only_min, + ), + None + ); + assert!(collected_infos(&collector).is_empty()); + + let min = mixed_range_input( + "min_over_time", + vec![Some(3.0), None, Some(1.0)], + vec![None, Some(first), None], + ); + assert_eq!( + mixed_float_result(MixedRange::float_udf(Some(collector.clone())), min), + Some(1.0) + ); + assert!( + collected_infos(&collector) + .iter() + .any(|info| info.contains("ignored native histogram")) + ); + } + #[test] fn count_sum_and_avg_read_struct() { let histograms = vec![Some(sample_histogram(6.0, 10.0, vec![2.0, 4.0]))]; diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index a463f176c71..d21a90bf123 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -63,15 +63,15 @@ use promql::extension_plan::{ }; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, - IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, NativeHistogramAbsentOverTime, - NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, - NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDrop, - NativeHistogramFraction, NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, - NativeHistogramLastOverTime, NativeHistogramPresentOverTime, NativeHistogramQuantile, - NativeHistogramRate, NativeHistogramResets, NativeHistogramStddev, NativeHistogramStdvar, - NativeHistogramSum, NativeHistogramSumOverTime, PredictLinear, PresentOverTime, - QuantileOverTime, Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, - quantile_udaf, + IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, + NativeHistogramAbsentOverTime, NativeHistogramAvg, NativeHistogramAvgOverTime, + NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, + NativeHistogramDelta, NativeHistogramDrop, NativeHistogramFraction, NativeHistogramIDelta, + NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, + NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, + NativeHistogramResets, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSum, + NativeHistogramSumOverTime, PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, + Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; @@ -2833,6 +2833,90 @@ impl PromPlanner { Ok(result) } + fn create_mixed_range_function_exprs( + &mut self, + func: &Function, + mut other_input_exprs: VecDeque, + float_field: &str, + histogram_field: &str, + ) -> Result>> { + let returns_histogram = matches!( + func.name, + "rate" + | "increase" + | "delta" + | "idelta" + | "irate" + | "avg_over_time" + | "sum_over_time" + | "last_over_time" + ); + if !returns_histogram + && !matches!( + func.name, + "changes" + | "resets" + | "deriv" + | "min_over_time" + | "max_over_time" + | "count_over_time" + | "absent_over_time" + | "present_over_time" + | "stddev_over_time" + | "stdvar_over_time" + | "quantile_over_time" + | "predict_linear" + | "double_exponential_smoothing" + | "holt_winters" + ) + { + return Ok(None); + } + + if func.name == "predict_linear" { + other_input_exprs[0] = DfExpr::Cast(Cast { + expr: Box::new(other_input_exprs[0].clone()), + data_type: ArrowDataType::Int64, + }); + } + + let mut args = Vec::with_capacity(other_input_exprs.len() + 6); + args.push(lit(func.name)); + args.push(DfExpr::Column(Column::from_name( + RangeManipulate::build_timestamp_range_name( + self.ctx.time_index_column.as_ref().unwrap(), + ), + ))); + args.push(DfExpr::Column(Column::from_name(float_field))); + args.push(DfExpr::Column(Column::from_name(histogram_field))); + args.extend(other_input_exprs); + if matches!(func.name, "rate" | "increase" | "delta") { + args.push(self.create_time_index_column_expr()?); + args.push(lit(self.ctx.range.context(ExpectRangeSelectorSnafu)?)); + } + + let float_expr = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())), + args: args.clone(), + }); + let exprs = if returns_histogram { + self.ctx.field_columns = vec![float_field.to_string(), histogram_field.to_string()]; + vec![ + float_expr.alias(float_field), + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(MixedRange::histogram_udf(self.promql_annotations.clone())), + args, + }) + .alias(histogram_field), + ] + } else { + let display_name = float_expr.schema_name().to_string(); + self.ctx.field_columns = vec![display_name.clone()]; + vec![float_expr.alias(display_name)] + }; + Ok(Some(exprs)) + } + /// Creates function expressions for projection and returns the expressions and new tags. /// /// # Side Effects @@ -2847,6 +2931,18 @@ impl PromPlanner { ) -> Result<(Vec, Vec)> { // TODO(ruihang): check function args list let mut other_input_exprs: VecDeque = other_input_exprs.into(); + if let Some((float_field, histogram_field)) = + Self::alternative_sample_range_columns(input_schema, &self.ctx.field_columns) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())) + && let Some(exprs) = self.create_mixed_range_function_exprs( + func, + other_input_exprs.clone(), + &float_field, + &histogram_field, + )? + { + return Ok((exprs, vec![])); + } let alternative_samples = Self::field_columns_are_alternative_samples(input_schema, &self.ctx.field_columns); let all_field_columns_are_native_histogram_ranges = @@ -4444,6 +4540,17 @@ impl PromPlanner { .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type()) } + fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool { + Self::field_column_type(schema, field_column).is_some_and(|data_type| { + matches!( + data_type, + ArrowDataType::Dictionary(key_type, value_type) + if key_type.as_ref() == &ArrowDataType::Int64 + && value_type.as_ref() == &ArrowDataType::Float64 + ) + }) + } + fn field_columns_are_alternative_samples( schema: &DFSchemaRef, field_columns: &[String], @@ -4461,11 +4568,13 @@ impl PromPlanner { let canonical_float = field_columns.iter().find(|field| { field.as_str() == greptime_value() - && Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64) + && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64) + || Self::field_column_is_float_range(schema, field)) }); let canonical_histogram = field_columns.iter().find(|field| { field.as_str() == greptime_native_histogram() - && Self::field_column_is_native_histogram(schema, field) + && (Self::field_column_is_native_histogram(schema, field) + || Self::field_column_is_native_histogram_range(schema, field)) }); if let (Some(float), Some(histogram)) = (canonical_float, canonical_histogram) { return Some((float, histogram)); @@ -4473,21 +4582,34 @@ impl PromPlanner { let float = field_columns.iter().find(|field| { field.starts_with(OR_FLOAT_FIELD_PREFIX) - && Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64) + && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64) + || Self::field_column_is_float_range(schema, field)) })?; let histogram = field_columns.iter().find(|field| { field.starts_with(OR_HISTOGRAM_FIELD_PREFIX) - && Self::field_column_is_native_histogram(schema, field) + && (Self::field_column_is_native_histogram(schema, field) + || Self::field_column_is_native_histogram_range(schema, field)) })?; Some((float, histogram)) } + fn alternative_sample_range_columns<'a>( + schema: &DFSchemaRef, + field_columns: &'a [String], + ) -> Option<(&'a str, &'a str)> { + Self::alternative_sample_columns(schema, field_columns).filter(|(float, histogram)| { + Self::field_column_is_float_range(schema, float) + && Self::field_column_is_native_histogram_range(schema, histogram) + }) + } + fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool { Self::field_column_type(schema, field_column).is_some_and(|data_type| { matches!( data_type, - ArrowDataType::Dictionary(_, value_type) - if value_type.as_ref() == &Self::native_histogram_arrow_type() + ArrowDataType::Dictionary(key_type, value_type) + if key_type.as_ref() == &ArrowDataType::Int64 + && value_type.as_ref() == &Self::native_histogram_arrow_type() ) }) } @@ -5617,7 +5739,10 @@ mod test { use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; use common_base::Plugins; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; - use common_query::prelude::{greptime_native_histogram, greptime_timestamp}; + use common_query::native_histogram::{ + CounterResetHint, NativeHistogram, build_histogram_array, + }; + use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use common_query::test_util::DummyDecoder; use datafusion::arrow::array::{ Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, @@ -5791,6 +5916,23 @@ mod test { } } + fn direct_or_histogram() -> NativeHistogram { + NativeHistogram { + schema: 0, + zero_threshold: 0.0, + sum: 1.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: vec![], + positive_spans: vec![], + negative_spans: vec![], + count: 1.0, + zero_count: 1.0, + positive_buckets: vec![], + negative_buckets: vec![], + } + } + struct DirectOrSource { name: &'static str, empty: bool, @@ -8921,6 +9063,175 @@ mod test { } } + #[tokio::test] + async fn mixed_native_histogram_ranges_use_coordinated_udfs() { + let dual_output = [ + "increase(some_metric[5m])", + "rate(some_metric[5m])", + "delta(some_metric[5m])", + "idelta(some_metric[5m])", + "irate(some_metric[5m])", + "avg_over_time(some_metric[5m])", + "sum_over_time(some_metric[5m])", + "last_over_time(some_metric[5m])", + ]; + let float_output = [ + "resets(some_metric[5m])", + "changes(some_metric[5m])", + "deriv(some_metric[5m])", + "min_over_time(some_metric[5m])", + "max_over_time(some_metric[5m])", + "count_over_time(some_metric[5m])", + "absent_over_time(some_metric[5m])", + "present_over_time(some_metric[5m])", + "stddev_over_time(some_metric[5m])", + "stdvar_over_time(some_metric[5m])", + "quantile_over_time(0.9, some_metric[5m])", + "predict_linear(some_metric[5m], 60)", + "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", + ]; + + for query in dual_output.iter().chain(float_output.iter()) { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}"); + assert_eq!( + plan.contains("prom_mixed_range_histogram"), + dual_output.contains(query), + "{query}\n{plan}" + ); + } + + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; + assert_eq!(plan, expected); + } + + #[tokio::test] + async fn mixed_native_histogram_rate_executes_real_ranges() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new(greptime_value(), ArrowDataType::Float64, true), + Field::new( + greptime_native_histogram(), + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])), + Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])), + build_histogram_array(&[None, Some(direct_or_histogram()), None]), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let collector = PromqlAnnotationCollector::default(); + let mut planner = PromPlanner { + table_provider: build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await, + ctx: PromPlannerContext { + start: 3000, + end: 3000, + interval: 1000, + range: Some(3000), + time_index_column: Some("timestamp".to_string()), + field_columns: vec![ + greptime_native_histogram().to_string(), + greptime_value().to_string(), + ], + ..Default::default() + }, + promql_annotations: Some(collector.clone()), + }; + let input = LogicalPlan::Extension(Extension { + node: Arc::new( + RangeManipulate::new( + 3000, + 3000, + 1000, + 3000, + "timestamp".to_string(), + planner.ctx.field_columns.clone(), + input, + ) + .unwrap(), + ), + }); + let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else { + unreachable!() + }; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, vec![], input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = Vec::new(); + collector.append_to(&mut warnings, &mut Vec::new()); + assert!( + warnings + .iter() + .any(|warning| warning.contains("mix of float and native histogram")) + ); + } + #[tokio::test] async fn native_histogram_mixed_field_table_behaves() { // Metric Engine exposes float and histogram samples as alternative nullable fields. diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index ca4bdd81044..030ed27b4e3 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -30,18 +30,18 @@ use datafusion_expr::UserDefinedLogicalNode; use greptime_proto::substrait_extension::MergeScan as PbMergeScan; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, - IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, NativeHistogramAbsentOverTime, - NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, - NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, - NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDivScalar, - NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, NativeHistogramIDelta, - NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, - NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq, - NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, - NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, - NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, NativeHistogramToString, - PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, Round, StddevOverTime, - StdvarOverTime, SumOverTime, quantile_udaf, + IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, + NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, + NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, + NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, + NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, + NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, + NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg, + NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, + NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, + NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, + NativeHistogramToString, PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, Round, + StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, }; use prost::Message; use session::context::QueryContextRef; @@ -195,6 +195,8 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { NativeHistogramIRate::scalar_udf(), NativeHistogramIncrease::scalar_udf(), NativeHistogramLastOverTime::scalar_udf(), + MixedRange::float_udf(None), + MixedRange::histogram_udf(None), NativeHistogramMulScalar::scalar_udf(), NativeHistogramNeg::scalar_udf(), NativeHistogramNotEq::scalar_udf(), @@ -240,10 +242,14 @@ mod tests { use common_query::native_histogram::native_histogram_value_type; use datafusion::catalog::TableProvider; use datafusion::datasource::MemTable; + use datafusion::logical_expr::Extension; use datafusion_expr::expr::ScalarFunction; use datafusion_expr::{Expr, LogicalPlanBuilder, LogicalTableSource, col, lit}; - use datatypes::arrow::datatypes::{Field, Schema, SchemaRef}; + use datatypes::arrow::datatypes::{ + DataType as ArrowDataType, Field, Schema, SchemaRef, TimeUnit, + }; use datatypes::data_type::DataType; + use promql::extension_plan::RangeManipulate; use session::context::QueryContext; use super::*; @@ -374,5 +380,82 @@ mod tests { assert!(decoded.contains("prom_native_histogram_count")); assert!(decoded.contains("prom_native_histogram_drop_bool")); assert!(decoded.contains("prom_native_histogram_drop_float")); + + let schema = Arc::new(Schema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("float", ArrowDataType::Float64, true), + Field::new( + "histogram", + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let input = LogicalPlanBuilder::scan( + "devices", + Arc::new(LogicalTableSource::new(schema.clone())), + None, + ) + .unwrap() + .build() + .unwrap(); + let input = LogicalPlan::Extension(Extension { + node: Arc::new( + RangeManipulate::new( + 0, + 1000, + 1000, + 1000, + "timestamp".to_string(), + vec!["float".to_string(), "histogram".to_string()], + input, + ) + .unwrap(), + ), + }); + let plan = LogicalPlanBuilder::from(input) + .project(vec![ + Expr::ScalarFunction(ScalarFunction { + func: Arc::new(MixedRange::float_udf(None)), + args: vec![ + lit("last_over_time"), + col("timestamp_range"), + col("float"), + col("histogram"), + ], + }) + .alias("mixed_float"), + Expr::ScalarFunction(ScalarFunction { + func: Arc::new(MixedRange::histogram_udf(None)), + args: vec![ + lit("last_over_time"), + col("timestamp_range"), + col("float"), + col("histogram"), + ], + }) + .alias("mixed_histogram"), + ]) + .unwrap() + .build() + .unwrap(); + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let table_provider = Arc::new(MemTable::try_new(schema, vec![vec![]]).unwrap()); + let decoded = plan_decoder + .decode( + bytes, + Arc::new(DummyCatalogList::with_table_provider(table_provider)), + false, + ) + .await + .unwrap() + .to_string(); + assert!(decoded.contains("prom_mixed_range_float")); + assert!(decoded.contains("prom_mixed_range_histogram")); } } From 584d66e3c3c2ae878b6e7bc71910fda6db12b5ca Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 10 Aug 2026 17:49:35 +0800 Subject: [PATCH 032/126] fix(meta): validate leaders before enabling skip_wal (#8823) Signed-off-by: jeremyhi --- src/common/meta/src/ddl/alter_table.rs | 11 +++- src/common/meta/src/ddl/tests/alter_table.rs | 60 ++++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index 9504e8672ca..acc02c4d1a8 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -198,6 +198,12 @@ impl AlterTableProcedure { // DDL manager can submit another procedure with locks for the current route. return Ok(Status::done_with_output(RegionRouteChanged)); } + ensure!( + !find_leaders(&physical_table_route.region_routes).is_empty(), + NoLeaderSnafu { + table_id: physical_table_id + } + ); self.data.region_distribution = Some(region_distribution(&physical_table_route.region_routes)); } @@ -264,9 +270,8 @@ impl AlterTableProcedure { MultipleResults::PartialRetryable(error) => Err(error), MultipleResults::PartialNonRetryable(error) | MultipleResults::AllNonRetryable(error) => { - // The metadata already enables skip-WAL. Keep retrying until every - // replica applies the runtime option instead of leaving the table in - // a permanently inconsistent state. + // The metadata already enables skip-WAL. Retry the idempotent request + // so later attempts can update the remaining replicas. Err(BoxedError::new(error)).context(RetryLaterSnafu { clean_poisons: true, }) diff --git a/src/common/meta/src/ddl/tests/alter_table.rs b/src/common/meta/src/ddl/tests/alter_table.rs index af72aaa65df..d68d1c96dfa 100644 --- a/src/common/meta/src/ddl/tests/alter_table.rs +++ b/src/common/meta/src/ddl/tests/alter_table.rs @@ -850,6 +850,66 @@ async fn test_skip_wal_detects_region_route_change() { assert!(!table_info.meta.options.skip_wal); } +#[tokio::test] +async fn test_skip_wal_rejects_no_leader_before_updating_metadata() { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let table_name = "foo"; + let table_id = 1024; + let task = test_create_table_task(table_name, table_id); + let mut region_routes = prepare_table_route(table_id) + .region_routes() + .unwrap() + .clone(); + for route in &mut region_routes { + route.leader_peer = None; + } + let region_locks = region_routes.iter().map(|route| route.region.id).collect(); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(region_routes), + HashMap::new(), + ) + .await + .unwrap(); + + let alter_task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![api::v1::Option { + key: SKIP_WAL_KEY.to_string(), + value: "true".to_string(), + }], + })), + }, + }; + let mut procedure = AlterTableProcedure::new_with_region_locks( + table_id, + alter_task, + EventContext::default(), + region_locks, + ddl_context.clone(), + ) + .unwrap(); + + let error = procedure.on_prepare().await.unwrap_err(); + assert_matches!(error, Error::NoLeader { .. }); + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert!(!table_info.meta.options.skip_wal); +} + #[tokio::test] async fn test_skip_wal_updates_metadata_before_all_replicas() { let (tx, mut rx) = mpsc::channel(8); From 60c7bbcaf30a0e65ad7fa2b2b8ad0c1199a662b1 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 10 Aug 2026 18:47:50 +0800 Subject: [PATCH 033/126] refactor: share query channel definition (#8825) refactor: share query channel protocol mapping Signed-off-by: WenyXu --- Cargo.lock | 3 +- src/common/base/Cargo.toml | 1 + src/common/base/src/lib.rs | 1 + src/common/base/src/protocol.rs | 106 ++++++++++++++++++++++++++++++++ src/common/meta/Cargo.toml | 1 - src/common/meta/src/rpc/ddl.rs | 16 ++--- src/common/session/src/lib.rs | 49 +-------------- src/session/Cargo.toml | 1 + src/session/src/context.rs | 71 +++------------------ src/session/src/lib.rs | 4 +- 10 files changed, 126 insertions(+), 127 deletions(-) create mode 100644 src/common/base/src/protocol.rs diff --git a/Cargo.lock b/Cargo.lock index 4ba0a7f7e27..e2210b80bfe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2292,6 +2292,7 @@ dependencies = [ "pin-project", "regex", "serde", + "strum 0.27.1", "tokio", "toml 0.8.23", "zeroize", @@ -2633,7 +2634,6 @@ dependencies = [ "common-query", "common-recordbatch", "common-runtime", - "common-session", "common-telemetry", "common-test-util", "common-time", @@ -13426,6 +13426,7 @@ dependencies = [ "api", "arc-swap", "auth", + "common-base", "common-catalog", "common-error", "common-macro", diff --git a/src/common/base/Cargo.toml b/src/common/base/Cargo.toml index 44c30cd548d..760241a3e6a 100644 --- a/src/common/base/Cargo.toml +++ b/src/common/base/Cargo.toml @@ -21,6 +21,7 @@ lazy_static.workspace = true pin-project.workspace = true regex.workspace = true serde = { version = "1.0", features = ["derive"] } +strum.workspace = true tokio.workspace = true zeroize = { version = "1.6", default-features = false, features = ["alloc"] } diff --git a/src/common/base/src/lib.rs b/src/common/base/src/lib.rs index 18e6ffb3e0f..04fc7038018 100644 --- a/src/common/base/src/lib.rs +++ b/src/common/base/src/lib.rs @@ -18,6 +18,7 @@ pub mod cancellation; pub mod hash; pub mod memory_limit; pub mod plugins; +pub mod protocol; pub mod range_read; #[allow(clippy::all)] pub mod readable_size; diff --git a/src/common/base/src/protocol.rs b/src/common/base/src/protocol.rs new file mode 100644 index 00000000000..66704dbc18b --- /dev/null +++ b/src/common/base/src/protocol.rs @@ -0,0 +1,106 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt::{Display, Formatter}; + +/// The protocol through which a query is received. +#[derive(Debug, PartialEq, Default, Clone, Copy, strum::FromRepr)] +#[repr(u8)] +pub enum Channel { + #[default] + Unknown = 0, + + Mysql = 1, + Postgres = 2, + HttpSql = 3, + Prometheus = 4, + Otlp = 5, + Grpc = 6, + Influx = 7, + Opentsdb = 8, + Loki = 9, + Elasticsearch = 10, + Jaeger = 11, + Log = 12, + Promql = 13, + Splunk = 14, +} + +impl From for Channel { + fn from(value: u32) -> Self { + u8::try_from(value) + .ok() + .and_then(Self::from_repr) + .unwrap_or_default() + } +} + +impl Display for Channel { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_ref()) + } +} + +impl AsRef for Channel { + fn as_ref(&self) -> &str { + match self { + Self::Unknown => "unknown", + Self::Mysql => "mysql", + Self::Postgres => "postgres", + Self::HttpSql => "httpsql", + Self::Prometheus => "prometheus", + Self::Otlp => "otlp", + Self::Grpc => "grpc", + Self::Influx => "influx", + Self::Opentsdb => "opentsdb", + Self::Loki => "loki", + Self::Elasticsearch => "elasticsearch", + Self::Jaeger => "jaeger", + Self::Log => "log", + Self::Promql => "promql", + Self::Splunk => "splunk", + } + } +} + +#[cfg(test)] +mod tests { + use super::Channel; + + #[test] + fn test_channel_name() { + let expected = [ + (1, "mysql"), + (2, "postgres"), + (3, "httpsql"), + (4, "prometheus"), + (5, "otlp"), + (6, "grpc"), + (7, "influx"), + (8, "opentsdb"), + (9, "loki"), + (10, "elasticsearch"), + (11, "jaeger"), + (12, "log"), + (13, "promql"), + (14, "splunk"), + ]; + + for (value, name) in expected { + assert_eq!(name, Channel::from(value).as_ref()); + } + assert_eq!("unknown", Channel::from(0).as_ref()); + assert_eq!("unknown", Channel::from(15).as_ref()); + } +} diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 43714b7d599..8e0cbbdf354 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -41,7 +41,6 @@ common-procedure.workspace = true common-query.workspace = true common-recordbatch.workspace = true common-runtime.workspace = true -common-session.workspace = true common-telemetry.workspace = true common-time.workspace = true common-wal.workspace = true diff --git a/src/common/meta/src/rpc/ddl.rs b/src/common/meta/src/rpc/ddl.rs index 9cb75aa29c0..8aef3325e64 100644 --- a/src/common/meta/src/rpc/ddl.rs +++ b/src/common/meta/src/rpc/ddl.rs @@ -42,9 +42,9 @@ use api::v1::{ }; use base64::Engine as _; use base64::engine::general_purpose; +use common_base::protocol::Channel; use common_catalog::{format_full_flow_name, format_full_table_name}; use common_error::ext::BoxedError; -use common_session::channel_protocol; use common_time::{DatabaseTimeToLive, Timestamp}; use prost::Message; use serde::{Deserialize, Serialize}; @@ -1668,13 +1668,6 @@ pub struct QueryContext { pub sst_min_sequences: HashMap, } -impl QueryContext { - /// Returns the protocol name represented by the wire channel value. - pub fn channel_protocol(&self) -> Option<&'static str> { - channel_protocol(self.channel) - } -} - /// The stable context recorded for a procedure event. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct EventContext { @@ -1708,11 +1701,12 @@ impl EventContext { .get(TRIGGER_REASON_EXTENSION_KEY) .map(|reason| TriggerReason::from_extension(reason)) .unwrap_or_default(); + let channel = Channel::from(u32::from(query_context.channel)); let context = Self::new(reason); - if let Some(protocol) = query_context.channel_protocol() { - context.with_protocol(protocol) - } else { + if channel == Channel::Unknown { context + } else { + context.with_protocol(channel.as_ref()) } } } diff --git a/src/common/session/src/lib.rs b/src/common/session/src/lib.rs index 418425df15e..f09baf92866 100644 --- a/src/common/session/src/lib.rs +++ b/src/common/session/src/lib.rs @@ -15,27 +15,6 @@ use serde::{Deserialize, Serialize}; use strum::{AsRefStr, Display, EnumString}; -/// Returns the protocol name represented by a wire query channel value. -pub fn channel_protocol(channel: u8) -> Option<&'static str> { - match channel { - 1 => Some("mysql"), - 2 => Some("postgres"), - 3 => Some("httpsql"), - 4 => Some("prometheus"), - 5 => Some("otlp"), - 6 => Some("grpc"), - 7 => Some("influx"), - 8 => Some("opentsdb"), - 9 => Some("loki"), - 10 => Some("elasticsearch"), - 11 => Some("jaeger"), - 12 => Some("log"), - 13 => Some("promql"), - 14 => Some("splunk"), - _ => None, - } -} - /// Defines the read preference for frontend route operations, /// determining whether to read from the region leader or follower. #[derive( @@ -52,33 +31,7 @@ pub enum ReadPreference { mod tests { use std::str::FromStr; - use crate::{ReadPreference, channel_protocol}; - - #[test] - fn test_channel_protocol() { - let expected = [ - (1, "mysql"), - (2, "postgres"), - (3, "httpsql"), - (4, "prometheus"), - (5, "otlp"), - (6, "grpc"), - (7, "influx"), - (8, "opentsdb"), - (9, "loki"), - (10, "elasticsearch"), - (11, "jaeger"), - (12, "log"), - (13, "promql"), - (14, "splunk"), - ]; - - for (channel, protocol) in expected { - assert_eq!(Some(protocol), channel_protocol(channel)); - } - assert_eq!(None, channel_protocol(0)); - assert_eq!(None, channel_protocol(15)); - } + use crate::ReadPreference; #[test] fn test_read_preference() { diff --git a/src/session/Cargo.toml b/src/session/Cargo.toml index c7be8c2f7e8..889b911c3cc 100644 --- a/src/session/Cargo.toml +++ b/src/session/Cargo.toml @@ -15,6 +15,7 @@ ahash.workspace = true api.workspace = true arc-swap = "1.5" auth.workspace = true +common-base.workspace = true common-catalog.workspace = true common-error.workspace = true common-macro.workspace = true diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 4780523e4b7..d383aef477e 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -22,6 +22,7 @@ use api::v1::ExplainOptions; use api::v1::region::RegionRequestHeader; use arc_swap::ArcSwap; use auth::UserInfoRef; +pub use common_base::protocol::Channel; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_catalog::{build_db_string, parse_catalog_and_schema_from_db_string}; use common_recordbatch::cursor::RecordBatchStreamCursor; @@ -618,70 +619,12 @@ impl ConnInfo { } } -/// The protocol mapping for these wire values lives in [`common_session::channel_protocol`]. -#[derive(Debug, PartialEq, Default, Clone, Copy)] -#[repr(u8)] -pub enum Channel { - #[default] - Unknown = 0, - - Mysql = 1, - Postgres = 2, - HttpSql = 3, - Prometheus = 4, - Otlp = 5, - Grpc = 6, - Influx = 7, - Opentsdb = 8, - Loki = 9, - Elasticsearch = 10, - Jaeger = 11, - Log = 12, - Promql = 13, - Splunk = 14, -} - -impl From for Channel { - fn from(value: u32) -> Self { - match value { - 1 => Self::Mysql, - 2 => Self::Postgres, - 3 => Self::HttpSql, - 4 => Self::Prometheus, - 5 => Self::Otlp, - 6 => Self::Grpc, - 7 => Self::Influx, - 8 => Self::Opentsdb, - 9 => Self::Loki, - 10 => Self::Elasticsearch, - 11 => Self::Jaeger, - 12 => Self::Log, - 13 => Self::Promql, - 14 => Self::Splunk, - _ => Self::Unknown, - } - } -} - -impl Channel { - pub fn dialect(&self) -> Arc { - match self { - Channel::Mysql => Arc::new(MySqlDialect {}), - Channel::Postgres => Arc::new(PostgreSqlDialect {}), - _ => Arc::new(GenericDialect {}), - } - } -} - -impl Display for Channel { - fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { - write!(f, "{}", self.as_ref()) - } -} - -impl AsRef for Channel { - fn as_ref(&self) -> &str { - common_session::channel_protocol(*self as u8).unwrap_or("unknown") +/// Returns the SQL dialect for the given query channel. +pub fn dialect_for_channel(channel: Channel) -> Arc { + match channel { + Channel::Mysql => Arc::new(MySqlDialect {}), + Channel::Postgres => Arc::new(PostgreSqlDialect {}), + _ => Arc::new(GenericDialect {}), } } diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index 81ee7b0cdc2..110560b3243 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -34,7 +34,7 @@ use common_time::timezone::get_timezone; use context::{ConfigurationVariables, QueryContextBuilder}; use derive_more::Debug; -use crate::context::{Channel, ConnInfo, QueryContextRef}; +use crate::context::{Channel, ConnInfo, QueryContextRef, dialect_for_channel}; /// Maximum number of warnings to store per session (similar to MySQL's max_error_count) const MAX_WARNINGS: usize = 64; @@ -102,7 +102,7 @@ impl Session { // string here .current_catalog(self.catalog.read().unwrap().clone()) .mutable_session_data(self.mutable_inner.clone()) - .sql_dialect(self.conn_info.channel.dialect()) + .sql_dialect(dialect_for_channel(self.conn_info.channel)) .configuration_parameter(self.configuration_variables.clone()) .channel(self.conn_info.channel) .process_id(self.process_id) From 335a95a369c728c355f1c8452d8688525beb24b9 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 10 Aug 2026 20:18:28 +0800 Subject: [PATCH 034/126] feat: declared edges and the derivation contract for the entity graph (#8794) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(frontend): run entity-graph derivation as the caller The derivation contract requires the computed graph tables to run under the outer query's identity. Capture the caller's QueryContext when the computed table is resolved, thread it through EntityGraphProvider, and: - authorize every contributing source table against the caller via the new semantic_graph.query permission action, silently excluding denied sources (entities, edges and source_tables never appear); - execute the derivation plan under the caller's context so it inherits permissions, cancellation and deadline instead of a fresh default. Signed-off-by: Dennis Zhuang * feat(operator): derive the entity-graph window from the scan's time predicate Implements the RFC window contract for the computed graph tables: - table: add extract_time_range_strict, a strict variant of the lenient time-range extraction that distinguishes an absent observed_at filter from one that cannot be safely turned into a range; - operator: replace GraphWindow with GraphQueryWindow, splitting the queried observed_at range from the source-scan range widened to whole 60s buckets, so boundary buckets aggregate over their full extent; - frontend: resolve the window from ScanRequest filters — no predicate keeps the last-hour default, a missing upper bound means now, and a missing lower bound or unextractable shape is an explicit error, never a silent fallback. Signed-off-by: Dennis Zhuang * feat(operator): system-defined declared-edge table for the entity graph Reintroduces greptime_private.semantic_relationships_declared with a canonical, system-owned definition: - the CREATE TABLE expr (8-tag primary key, business validity columns, RED fields, 30d TTL); attributes is now a json column so the future union branch matches the computed table without a per-scan parse; - created on first use on every write path: SQL INSERT creates it before executing, and the gRPC row-insert auto-create substitutes the canonical expr instead of deriving a schema from the request; - user DDL (CREATE/ALTER/DROP/RENAME/TRUNCATE) and write-path auto-ALTER are rejected via the new is_ddl_reserved_table guard, while INSERT/DELETE stay allowed. Signed-off-by: Dennis Zhuang * feat(operator): union declared edges into semantic_relationships Adds the declared-edge branch to the relationship derivation (build_relationships_plan replaces build_calls_plan): - latest revision per edge key first (mito dedups on primary key plus observed_at, so a re-asserted edge stores a new revision), then the business-validity overlap against the queried window; valid_from defaults to the declaration time and a NULL valid_until means the edge holds while its row exists; - the projected observed_at is synthesized inside the queried range (Inexact pushdown re-applies the scan's filters above the computed table, which would drop rows keyed by the physical revision time); window_end/fresh_until of open-ended edges take the window's upper bound so 'fresh_until >= now() - ...' queries see them; - tag columns are cast out of dictionary encoding, and the union is re-projected to the 16-column contract; - the frontend feeds the branch only when the physical table exists, the caller may read it, and its schema still matches the canonical definition (mismatch is an explicit error, not a silent drop). Signed-off-by: Dennis Zhuang * test: cover declared edges, window contract and caller authorization - sqlness: system auto-create on first INSERT, latest-revision reads, open-ended vs retired validity, explicit/lower-only/upper-only window behavior, user-DDL rejection, rename-into rejection, DELETE cleanup; - integration: a permission checker denying one trace table excludes it from both semantic_relationships and semantic_entities. Signed-off-by: Dennis Zhuang * fix: allow DROP/TRUNCATE on the declared-edge table and fix CI lints The definition guard rejected every DDL, which left sqlness (and any shared deployment) no way to remove the table the semantic_graph case creates — its extra region then broke unrelated region/partition case expectations. Narrow the guard to what actually protects the canonical definition: user CREATE, ALTER, RENAME-into and repartition stay rejected, while DROP and TRUNCATE are allowed — dropping loses nothing structural, the next INSERT recreates the table canonically, and DROP doubles as the recovery path if the canonical definition ever changes. The sqlness case now verifies drop-then-recreate and cleans up after itself. Also: rustfmt for the catalog crate and two typo fixes. Signed-off-by: Dennis Zhuang * chore: adapt canonical declared-table create to TriggerReason Signed-off-by: Dennis Zhuang * fix: address review on the declared-edge table lifecycle and revision reads - gRPC first writes actually work now: the reserved table's creation went through the generic create_table_inner, which the definition guard itself rejects; both branches of create_or_alter_tables_on_demand route it to create_declared_relationships_table instead, and being a system action it also bypasses the auto_create_table config/hint; - revision selection is as-of the queried window: revisions recorded after the window's end, or whose validity starts after it, no longer outrank (and hide) the revision that was in effect inside it; - the canonical-schema check validates the whole definition the union semantics lean on — time index, primary key, engine, append/merge mode — not just column names and types; - UNDROP TABLE of the reserved name is rejected like CREATE: it could resurrect a pre-canonical shape, and the next INSERT recreates the table anyway. Signed-off-by: Dennis Zhuang * chore: trim over-commenting in the entity-graph code Comments that restated adjacent code or narrated justification are cut; the ones stating non-obvious contracts and gotchas stay. Signed-off-by: Dennis Zhuang * fix: reject CREATE VIEW against DDL-reserved table names A view named greptime_private.semantic_relationships_declared would squat the reserved name: the first INSERT then skips the canonical create (an object already exists) and graph reads fail on the schema mismatch. CREATE VIEW now passes the same definition guard as CREATE TABLE. Signed-off-by: Dennis Zhuang * feat: debug-log authorization exclusions; declared-edge TTL to 90d Sources the derivation contract silently excludes (per-table denial, whole-scan denial, the declared-edge table) are invisible from outside; a debug log at each names what was excluded and why. The declared-edge table's default TTL becomes 90d, overridable at creation time via GREPTIMEDB_DECLARED_RELATIONSHIPS_TTL (a proper configuration option is a TODO). Signed-off-by: Dennis Zhuang * fix: rank declared-edge revisions by the visible edge identity Ranking partitioned by the full primary key, but the projection drops scope and generation_id: two assertions of the same visible edge under different generations both ranked first and came out as duplicate, indistinguishable rows. Rank by the exposed identity (endpoints, rel_type, provenance) instead, with generation_id/scope as deterministic tie-breakers for same-timestamp assertions. Signed-off-by: Dennis Zhuang * test: drop redundant declared-edge tests The generations regression is already asserted by the revision and as-of tests; the DDL shape test restated the declarative builder against itself. Its one non-tautological check (attributes maps to the json type) moves into the schema-matcher test. Signed-off-by: Dennis Zhuang * fix: reject disjunctive graph windows and unmatchable future windows - OR/IN over observed_at collapse disjoint ranges into their convex hull; a declared edge's synthesized timestamp can land in a gap and be dropped by the re-applied filter even though the edge is valid at a requested instant. The strict extractor now rejects those shapes. - A lower bound in the future inverts against the implicit up-to-now upper bound; the declared branch then fabricated an edge observed at the future bound. Such windows now derive nothing. - The reserved-table gRPC create path classifies an instant-TTL table like every sibling path. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/auth/src/lib.rs | 2 +- src/auth/src/permission.rs | 5 + src/catalog/src/kvbackend/manager.rs | 8 +- .../src/system_schema/semantic_graph.rs | 57 +- src/common/catalog/src/consts.rs | 11 + src/frontend/src/instance/builder.rs | 2 + src/frontend/src/instance/entity_graph.rs | 269 ++++- src/operator/src/error.rs | 12 +- src/operator/src/insert.rs | 53 +- src/operator/src/statement/ddl.rs | 51 +- src/operator/src/statement/dml.rs | 35 +- src/operator/src/statement/semantic_graph.rs | 1004 +++++++++++++++-- src/table/src/predicate.rs | 171 +++ tests-integration/src/grpc.rs | 172 ++- tests-integration/src/instance.rs | 104 ++ .../common/system/semantic_graph.result | 124 ++ .../common/system/semantic_graph.sql | 69 ++ 17 files changed, 2004 insertions(+), 145 deletions(-) diff --git a/src/auth/src/lib.rs b/src/auth/src/lib.rs index 44382737ee8..d2a1c05b1cd 100644 --- a/src/auth/src/lib.rs +++ b/src/auth/src/lib.rs @@ -34,7 +34,7 @@ pub use permission::{ DefaultPermissionChecker, INFLUXDB_WRITE, JAEGER_QUERY, LOG_QUERY, LOG_WRITE, OPENTSDB_WRITE, OTLP_WRITE, PIPELINE_DELETE, PIPELINE_INSERT, PIPELINE_QUERY, PROM_STORE_READ, PROM_STORE_WRITE, PROMQL_QUERY, PermissionAction, PermissionChecker, PermissionReq, - PermissionResp, PermissionTableTarget, PermissionTableTargets, + PermissionResp, PermissionTableTarget, PermissionTableTargets, SEMANTIC_GRAPH_QUERY, }; pub use user_info::UserInfo; pub use user_provider::static_user_provider::StaticUserProvider; diff --git a/src/auth/src/permission.rs b/src/auth/src/permission.rs index 60a89225738..901b935d6a2 100644 --- a/src/auth/src/permission.rs +++ b/src/auth/src/permission.rs @@ -150,6 +150,10 @@ pub const PIPELINE_DELETE: PermissionAction = PermissionAction::write("pipeline. pub const DASHBOARD_QUERY: PermissionAction = PermissionAction::read("dashboard.query"); pub const DASHBOARD_SAVE: PermissionAction = PermissionAction::write("dashboard.save"); pub const DASHBOARD_DELETE: PermissionAction = PermissionAction::write("dashboard.delete"); +/// Reading a telemetry table through the entity-graph derivation +/// (`greptime_private.semantic_entities` / `semantic_relationships`); checked +/// once per candidate source table with that table as the target. +pub const SEMANTIC_GRAPH_QUERY: PermissionAction = PermissionAction::read("semantic_graph.query"); /// All permission actions built into GreptimeDB. /// @@ -171,6 +175,7 @@ pub const ALL_ACTIONS: &[PermissionAction] = &[ DASHBOARD_QUERY, DASHBOARD_SAVE, DASHBOARD_DELETE, + SEMANTIC_GRAPH_QUERY, ]; #[derive(Debug, Clone)] diff --git a/src/catalog/src/kvbackend/manager.rs b/src/catalog/src/kvbackend/manager.rs index db7361615d0..2afc77c0d23 100644 --- a/src/catalog/src/kvbackend/manager.rs +++ b/src/catalog/src/kvbackend/manager.rs @@ -686,8 +686,12 @@ impl SystemCatalog { // Constructed on demand (the provider is a name + a weak ref); the // system catalog is consulted before physical resolution, so the // computed tables shadow same-named physical tables by design. - SemanticGraphTableProvider::new(catalog.to_string(), self.catalog_manager.clone()) - .table(table_name) + SemanticGraphTableProvider::new( + catalog.to_string(), + self.catalog_manager.clone(), + query_ctx.cloned().map(Arc::new), + ) + .table(table_name) } else { None } diff --git a/src/catalog/src/system_schema/semantic_graph.rs b/src/catalog/src/system_schema/semantic_graph.rs index 049377f5a38..028ce79b906 100644 --- a/src/catalog/src/system_schema/semantic_graph.rs +++ b/src/catalog/src/system_schema/semantic_graph.rs @@ -18,10 +18,9 @@ //! They live in `greptime_private`, not `information_schema`: scanning them //! triggers read-time derivation over telemetry tables (trace self-joins, ...), //! which breaks the "cheap, metadata-only" expectation users have of -//! `information_schema`. `greptime_private` already signals "system-managed, -//! computed data objects", and it is also where the physical declared-edge -//! table will live (a follow-up), so derived and declared edges will share one -//! schema. +//! `information_schema`. `greptime_private` also hosts the physical +//! declared-edge table (`semantic_relationships_declared`), whose rows are +//! unioned into the computed `semantic_relationships`. //! //! These are thin forwarders: their rows are derived at read time by the injected //! [`EntityGraphProvider`], which enumerates the `table_semantics` declarations, @@ -49,6 +48,7 @@ use common_recordbatch::{ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use futures::StreamExt; +use session::context::QueryContextRef; use snafu::ResultExt; use store_api::storage::{ScanRequest, TableId}; use table::TableRef; @@ -59,17 +59,18 @@ use crate::system_schema::{SystemSchemaProviderInner, SystemTable, SystemTableRe pub type EntityGraphProviderRef = Arc; -/// Produces the rows of the computed entity-graph tables -/// (`semantic_entities` / `semantic_relationships`) at read time. +/// Produces the rows of the computed entity-graph tables at read time. /// -/// The computed tables are thin forwarders to this provider, which is -/// implemented above the query engine (in the frontend): it enumerates the entity -/// declarations from `table_semantics`, builds the typed derivation plans, and -/// executes them. It is injected into the catalog manager *after* construction — -/// the provider needs the engine, which needs the catalog manager — so this late -/// binding breaks the `catalog -> query` dependency cycle. Keeping derivation out -/// of `catalog` also respects the crate layering (the plan builders live in -/// `operator`). See `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. +/// Implemented above the query engine (in the frontend) and injected into the +/// catalog manager *after* construction — the provider needs the engine, which +/// needs the catalog manager — so this late binding breaks the +/// `catalog -> query` dependency cycle. +/// +/// `query_ctx` is the outer query's context, captured when the computed table +/// is resolved: the derivation must read only sources that context may read +/// and execute under it, inheriting the caller's permissions, cancellation and +/// deadline (the RFC's derivation contract). `None` (context-less internal +/// resolution) keeps the provider's default behaviour. #[async_trait::async_trait] pub trait EntityGraphProvider: Send + Sync { /// Produces the entity registry (`semantic_entities`) rows for `catalog`. @@ -78,6 +79,7 @@ pub trait EntityGraphProvider: Send + Sync { &self, catalog: &str, request: ScanRequest, + query_ctx: Option, ) -> std::result::Result, BoxedError>; /// Produces the relationship set (`semantic_relationships`) rows for `catalog`. @@ -85,6 +87,7 @@ pub trait EntityGraphProvider: Send + Sync { &self, catalog: &str, request: ScanRequest, + query_ctx: Option, ) -> std::result::Result, BoxedError>; } @@ -94,13 +97,21 @@ pub trait EntityGraphProvider: Send + Sync { pub(crate) struct SemanticGraphTableProvider { catalog_name: String, catalog_manager: Weak, + /// The resolving query's context; captured per resolution (the provider is + /// built on demand, never cached), so it cannot leak across sessions. + query_ctx: Option, } impl SemanticGraphTableProvider { - pub(crate) fn new(catalog_name: String, catalog_manager: Weak) -> Self { + pub(crate) fn new( + catalog_name: String, + catalog_manager: Weak, + query_ctx: Option, + ) -> Self { Self { catalog_name, catalog_manager, + query_ctx, } } @@ -139,6 +150,7 @@ impl SystemSchemaProviderInner for SemanticGraphTableProvider { kind, self.catalog_name.clone(), self.catalog_manager.clone(), + self.query_ctx.clone(), )) as _) } } @@ -252,6 +264,7 @@ struct SemanticGraphTable { schema: SchemaRef, catalog_name: String, catalog_manager: Weak, + query_ctx: Option, } impl SemanticGraphTable { @@ -259,6 +272,7 @@ impl SemanticGraphTable { kind: GraphTableKind, catalog_name: String, catalog_manager: Weak, + query_ctx: Option, ) -> Self { let schema = match kind { GraphTableKind::Entities => ENTITIES_SCHEMA.clone(), @@ -269,6 +283,7 @@ impl SemanticGraphTable { schema, catalog_name, catalog_manager, + query_ctx, } } @@ -277,6 +292,7 @@ impl SemanticGraphTable { catalog: String, catalog_manager: Weak, request: ScanRequest, + query_ctx: Option, ) -> Result> { let provider = utils::entity_graph_provider(&catalog_manager)?; // No provider (engine not up / non-frontend node): stream empty. @@ -284,8 +300,12 @@ impl SemanticGraphTable { return Ok(None); }; match kind { - GraphTableKind::Entities => provider.scan_entities(&catalog, request).await, - GraphTableKind::Relationships => provider.scan_relationships(&catalog, request).await, + GraphTableKind::Entities => provider.scan_entities(&catalog, request, query_ctx).await, + GraphTableKind::Relationships => { + provider + .scan_relationships(&catalog, request, query_ctx) + .await + } } .context(InternalSnafu) } @@ -338,10 +358,11 @@ impl SystemTable for SemanticGraphTable { let kind = self.kind; let catalog = self.catalog_name.clone(); let catalog_manager = self.catalog_manager.clone(); + let query_ctx = self.query_ctx.clone(); let stream_schema = schema.clone(); let stream = async move { - let stream = Self::derive(kind, catalog, catalog_manager, request) + let stream = Self::derive(kind, catalog, catalog_manager, request, query_ctx) .await .map_err(BoxedError::new) .context(common_recordbatch::error::ExternalSnafu)?; diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index 64bbbcbde86..6053d5a67c8 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -167,6 +167,14 @@ pub fn is_readonly_table(schema: &str, table: &str) -> bool { )) } +/// Returns true for tables whose *definition* is system-owned while their rows +/// are user-written: created with a canonical schema on first write; user +/// CREATE/ALTER/RENAME-into rejected; DML and DROP/TRUNCATE allowed (the next +/// write recreates the table). Currently only the declared-edge table. +pub fn is_ddl_reserved_table(schema: &str, table: &str) -> bool { + schema == DEFAULT_PRIVATE_SCHEMA_NAME && table == SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME +} + // ---- special table and fields ---- pub const TRACE_ID_COLUMN: &str = "trace_id"; pub const SPAN_ID_COLUMN: &str = "span_id"; @@ -208,4 +216,7 @@ pub const SEMANTIC_ENTITIES_TABLE_NAME: &str = "semantic_entities"; /// Computed relationship set: the edge set of the observability graph, derived at /// read time (`calls`/`runs_on`/... ) and unioned with the declared-edge table. pub const SEMANTIC_RELATIONSHIPS_TABLE_NAME: &str = "semantic_relationships"; +/// Physical table of hand-declared edges, unioned into the computed +/// `semantic_relationships`; see [`is_ddl_reserved_table`] for its lifecycle. +pub const SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME: &str = "semantic_relationships_declared"; // ---- End of entity relationship graph tables ---- diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 8e69c5951b8..85a2b8ff183 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicBool; +use auth::PermissionCheckerRef; use cache::{PARTITION_INFO_CACHE_NAME, TABLE_FLOWNODE_SET_CACHE_NAME, TABLE_ROUTE_CACHE_NAME}; use catalog::CatalogManagerRef; use catalog::kvbackend::KvBackendCatalogManager; @@ -283,6 +284,7 @@ impl FrontendBuilder { let provider: EntityGraphProviderRef = Arc::new(EntityGraphProviderImpl::new( query_engine.clone(), Arc::downgrade(&self.catalog_manager), + plugins.get::(), )); kv_catalog.set_entity_graph_provider(provider); } diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs index 1d012b7fc60..e55e5f2b8d7 100644 --- a/src/frontend/src/instance/entity_graph.rs +++ b/src/frontend/src/instance/entity_graph.rs @@ -27,29 +27,37 @@ use std::collections::HashMap; use std::sync::Weak; use async_trait::async_trait; +use auth::{ + PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget, + PermissionTableTargets, SEMANTIC_GRAPH_QUERY, +}; use catalog::CatalogManager; use catalog::system_schema::semantic_graph::EntityGraphProvider; use common_catalog::consts::{ DEFAULT_PRIVATE_SCHEMA_NAME, DEFAULT_SCHEMA_NAME, INFORMATION_SCHEMA_NAME, PG_CATALOG_NAME, - SERVICE_NAME_COLUMN, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SERVICE_NAME_COLUMN, }; -use common_error::ext::BoxedError; +use common_error::ext::{BoxedError, ErrorExt}; +use common_error::status_code::StatusCode; use common_query::OutputData; use common_recordbatch::SendableRecordBatchStream; -use common_telemetry::warn; +use common_telemetry::{debug, warn}; +use common_time::timestamp::TimeUnit; use datafusion::dataframe::DataFrame; use datafusion_expr::LogicalPlan; use futures::TryStreamExt; use operator::statement::semantic_graph::{ - CallsSource, EntityDeclaration, GraphWindow, RegistrySource, build_calls_plan, - build_registry_plan, + CallsSource, DeclaredSource, EntityDeclaration, GraphQueryWindow, OBSERVED_AT_COLUMN, + RegistrySource, build_registry_plan, build_relationships_plan, + declared_relationships_schema_matches, }; use query::QueryEngineRef; -use session::context::{QueryContextBuilder, QueryContextRef}; +use session::context::{QueryContext, QueryContextBuilder, QueryContextRef}; use snafu::ResultExt; use store_api::storage::ScanRequest; use table::TableRef; use table::metadata::TableInfo; +use table::predicate::{TimeRangeExtraction, extract_time_range_strict}; use table::requests::{ EntityRole, is_trace_v1_table, parse_entity_columns, parse_entity_option_key, }; @@ -60,6 +68,7 @@ use crate::error; pub struct EntityGraphProviderImpl { query_engine: QueryEngineRef, catalog_manager: Weak, + permission_checker: Option, } struct EntitySource { @@ -68,10 +77,40 @@ struct EntitySource { } impl EntityGraphProviderImpl { - pub fn new(query_engine: QueryEngineRef, catalog_manager: Weak) -> Self { + pub fn new( + query_engine: QueryEngineRef, + catalog_manager: Weak, + permission_checker: Option, + ) -> Self { Self { query_engine, catalog_manager, + permission_checker, + } + } + + /// Whether the caller may read the derivation sources named by `targets`. + /// `Ok(false)` = denied: the source is silently excluded, per the + /// derivation contract. Errors other than a denial abort the scan. + fn authorize_sources( + &self, + query_ctx: Option<&QueryContext>, + targets: PermissionTableTargets, + ) -> Result { + let Some(ctx) = query_ctx else { + return Ok(true); + }; + match self + .permission_checker + .as_ref() + .check_permission_with_table_targets( + ctx.current_user(), + PermissionReq::Action(SEMANTIC_GRAPH_QUERY), + targets, + ) { + Ok(_) => Ok(true), + Err(err) if err.status_code() == StatusCode::PermissionDenied => Ok(false), + Err(err) => Err(BoxedError::new(err)), } } @@ -182,15 +221,37 @@ impl EntityGraphProviderImpl { async fn enumerate( &self, catalog: &str, + query_ctx: Option<&QueryContext>, ) -> Result<(Vec, Vec<(EntityDeclaration, TableRef)>), BoxedError> { let Some(catalog_manager) = self.catalog_manager.upgrade() else { return Ok((vec![], vec![])); }; + // A target-blind checker (e.g. the default mode-based one) answers the + // same for every table: ask once up front instead of per table. + let per_table_auth = match self.permission_checker.as_ref() { + Some(checker) if query_ctx.is_some() => { + if checker.uses_table_targets() { + true + } else if self + .authorize_sources(query_ctx, PermissionTableTargets::resolved(vec![]))? + { + false + } else { + debug!( + "Caller lacks the entity-graph read permission; deriving an empty graph \ + (catalog: {catalog})" + ); + return Ok((vec![], vec![])); + } + } + _ => false, + }; + let mut declarations = vec![]; let mut traces = vec![]; let schemas = catalog_manager - .schema_names(catalog, None) + .schema_names(catalog, query_ctx) .await .map_err(BoxedError::new)?; for schema in schemas { @@ -202,11 +263,31 @@ impl EntityGraphProviderImpl { { continue; } - let mut tables = catalog_manager.tables(catalog, &schema, None); + let mut tables = catalog_manager.tables(catalog, &schema, query_ctx); while let Some(table) = tables.try_next().await.map_err(BoxedError::new)? { let table_info = table.table_info(); let table_declarations = Self::declarations_for(&table_info); - if is_trace_v1_table(&table_info) { + let is_trace = is_trace_v1_table(&table_info); + // Authorize only tables that would contribute rows. + if per_table_auth + && (is_trace || !table_declarations.is_empty()) + && !self.authorize_sources( + query_ctx, + PermissionTableTargets::resolved(vec![PermissionTableTarget::new( + catalog, + &schema, + &table_info.name, + )]), + )? + { + debug!( + "Excluding `{schema}.{}` from the entity-graph derivation: caller lacks \ + read permission", + table_info.name + ); + continue; + } + if is_trace { // TODO(entity-graph): validate the complete fixed trace-v1 // schema before adding the table to calls derivation, and // skip malformed/stale tables instead of letting one poison @@ -235,37 +316,122 @@ impl EntityGraphProviderImpl { Ok((declarations, traces)) } - /// The time window to derive over, taken from the scan's time predicate. + /// The time window to derive over, taken from the scan's `observed_at` + /// predicate. /// - /// TODO(entity-graph): read the `observed_at` bounds out of `request.filters` - /// and build the window from them. Until then this always returns the default - /// (last hour), so a query whose time filter falls entirely outside the last - /// hour derives nothing and returns empty — a known limitation tracked for the - /// scan-time-predicate-pushdown follow-up. - fn query_window(_request: &ScanRequest) -> GraphWindow { - GraphWindow::default_last_hour() + /// The contract (RFC "The contract"): no `observed_at` predicate → the + /// product default (last hour); a missing upper bound means "up to now"; a + /// predicate without a lower bound or in a shape that cannot be safely + /// extracted is an error asking for an explicit range — never a silent + /// fallback into incomplete results. + /// `Ok(None)` = the window cannot match anything (e.g. a lower bound in + /// the future with the implicit "up to now" upper bound): the scan streams + /// empty instead of deriving. + fn query_window(request: &ScanRequest) -> Result, BoxedError> { + let invalid = + |err_msg: String| Err(BoxedError::new(error::InvalidSqlSnafu { err_msg }.build())); + match extract_time_range_strict(OBSERVED_AT_COLUMN, TimeUnit::Millisecond, &request.filters) + { + TimeRangeExtraction::Absent => Ok(Some(GraphQueryWindow::default_last_hour())), + TimeRangeExtraction::Extracted(range) => { + let Some(start) = range.start() else { + return invalid(format!( + "the {OBSERVED_AT_COLUMN} filter has no lower bound; the graph cannot \ + derive over unbounded history — add e.g. {OBSERVED_AT_COLUMN} >= \ + '2026-01-01 00:00:00'" + )); + }; + let end_ms = range + .end() + .map(|ts| ts.value()) + .unwrap_or_else(common_time::util::current_time_millis); + if start.value() >= end_ms { + return Ok(None); + } + Ok(Some(GraphQueryWindow::from_observed(start.value(), end_ms))) + } + TimeRangeExtraction::Unsupported => invalid(format!( + "cannot derive the graph window from the {OBSERVED_AT_COLUMN} filter; use plain \ + range predicates (>=, <, BETWEEN) with literal bounds" + )), + } } fn read_table(&self, table: TableRef) -> Result { self.query_engine.read_table(table).map_err(BoxedError::new) } - /// Executes a derivation plan and returns its live result stream. - /// - /// TODO(entity-graph): the QueryContext must come from the outer query - /// instead of being built here, so the derivation inherits the caller's - /// permissions, cancellation and deadline. Requires threading the context - /// through the computed-table scan path (planned next PR). + /// The declared-edge branch source, when the physical table exists, the + /// caller may read it, and it still matches the canonical definition. A + /// mismatch (upgrade skew) is an explicit error, not a silent skip: + /// dropping declared edges would misrepresent the graph. + async fn declared_source( + &self, + catalog: &str, + query_ctx: Option<&QueryContext>, + ) -> Result, BoxedError> { + let Some(catalog_manager) = self.catalog_manager.upgrade() else { + return Ok(None); + }; + let Some(table) = catalog_manager + .table( + catalog, + DEFAULT_PRIVATE_SCHEMA_NAME, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + query_ctx, + ) + .await + .map_err(BoxedError::new)? + else { + return Ok(None); + }; + if !self.authorize_sources( + query_ctx, + PermissionTableTargets::resolved(vec![PermissionTableTarget::new( + catalog, + DEFAULT_PRIVATE_SCHEMA_NAME, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + )]), + )? { + debug!( + "Excluding declared edges: caller lacks read permission on \ + `{DEFAULT_PRIVATE_SCHEMA_NAME}.{SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME}`" + ); + return Ok(None); + } + let table_info = table.table_info(); + if !declared_relationships_schema_matches(&table_info) { + return Err(BoxedError::new( + error::InvalidSqlSnafu { + err_msg: format!( + "{DEFAULT_PRIVATE_SCHEMA_NAME}.{SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME} \ + does not match its canonical schema; declared edges cannot be derived" + ), + } + .build(), + )); + } + Ok(Some(DeclaredSource { + scan: self.read_table(table)?, + })) + } + + /// Executes a derivation plan under the caller's context (inheriting its + /// permissions, cancellation and deadline); context-less internal scans + /// get a minimal default. async fn execute_plan( &self, catalog: &str, plan: LogicalPlan, + query_ctx: Option, ) -> Result, BoxedError> { - let query_ctx: QueryContextRef = QueryContextBuilder::default() - .current_catalog(catalog.to_string()) - .current_schema(DEFAULT_SCHEMA_NAME.to_string()) - .build() - .into(); + let query_ctx = query_ctx.unwrap_or_else(|| { + QueryContextBuilder::default() + .current_catalog(catalog.to_string()) + .current_schema(DEFAULT_SCHEMA_NAME.to_string()) + .build() + .into() + }); let output = self .query_engine .execute(plan, query_ctx) @@ -286,8 +452,9 @@ impl EntityGraphProvider for EntityGraphProviderImpl { &self, catalog: &str, request: ScanRequest, + query_ctx: Option, ) -> Result, BoxedError> { - let (sources, _) = self.enumerate(catalog).await?; + let (sources, _) = self.enumerate(catalog, query_ctx.as_deref()).await?; let mut plans = Vec::with_capacity(sources.len()); for source in sources { plans.push(RegistrySource { @@ -295,22 +462,25 @@ impl EntityGraphProvider for EntityGraphProviderImpl { scan: self.read_table(source.table)?, }); } - let window = Self::query_window(&request); + let Some(window) = Self::query_window(&request)? else { + return Ok(None); + }; let Some(plan) = build_registry_plan(plans, &window) .context(error::DataFusionSnafu) .map_err(BoxedError::new)? else { return Ok(None); }; - self.execute_plan(catalog, plan).await + self.execute_plan(catalog, plan, query_ctx).await } async fn scan_relationships( &self, catalog: &str, request: ScanRequest, + query_ctx: Option, ) -> Result, BoxedError> { - let (_, traces) = self.enumerate(catalog).await?; + let (_, traces) = self.enumerate(catalog, query_ctx.as_deref()).await?; let mut scans = Vec::with_capacity(traces.len()); for (service, trace) in traces { scans.push(CallsSource { @@ -318,14 +488,17 @@ impl EntityGraphProvider for EntityGraphProviderImpl { scan: self.read_table(trace)?, }); } - let window = Self::query_window(&request); - let Some(plan) = build_calls_plan(scans, &window) + let declared = self.declared_source(catalog, query_ctx.as_deref()).await?; + let Some(window) = Self::query_window(&request)? else { + return Ok(None); + }; + let Some(plan) = build_relationships_plan(scans, declared, &window) .context(error::DataFusionSnafu) .map_err(BoxedError::new)? else { return Ok(None); }; - self.execute_plan(catalog, plan).await + self.execute_plan(catalog, plan, query_ctx).await } } @@ -393,6 +566,30 @@ mod tests { .unwrap() } + #[test] + fn future_lower_only_window_matches_nothing() { + let future_ms = common_time::util::current_time_millis() + 86_400_000; + let request = ScanRequest { + filters: vec![ + datafusion_expr::col(OBSERVED_AT_COLUMN).gt_eq(datafusion_expr::lit( + datafusion::common::ScalarValue::TimestampMillisecond(Some(future_ms), None), + )), + ], + ..Default::default() + }; + assert!( + EntityGraphProviderImpl::query_window(&request) + .unwrap() + .is_none() + ); + + assert!( + EntityGraphProviderImpl::query_window(&ScanRequest::default()) + .unwrap() + .is_some() + ); + } + #[test] fn declaration_referencing_missing_column_is_skipped() { let info = table_info( diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index 324825d4589..b5e76f0f2cc 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -356,6 +356,15 @@ pub enum Error { location: Location, }, + #[snafu(display( + "The definition of table `{name}` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write)" + ))] + TableDdlReserved { + name: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Table occurs error"))] Table { #[snafu(implicit)] @@ -986,7 +995,8 @@ impl ErrorExt for Error { Error::NotSupported { .. } | Error::ShowCreateTableBaseOnly { .. } | Error::SchemaReadOnly { .. } - | Error::TableReadOnly { .. } => StatusCode::Unsupported, + | Error::TableReadOnly { .. } + | Error::TableDdlReserved { .. } => StatusCode::Unsupported, Error::TableMetadataManager { source, .. } => source.status_code(), Error::ParseSql { source, .. } => source.status_code(), Error::InvalidateTableCache { source, .. } => source.status_code(), diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 8205932708e..b5b6be093d9 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -29,8 +29,8 @@ use catalog::CatalogManagerRef; use client::{OutputData, OutputMeta}; use common_catalog::consts::{ PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN, TRACE_TABLE_NAME, - TRACE_TABLE_NAME_SESSION_KEY, default_engine, trace_operations_table_name, - trace_services_table_name, + TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table, + trace_operations_table_name, trace_services_table_name, }; use common_grpc_expr::util::ColumnExpr; use common_meta::cache::TableFlownodeSetCacheRef; @@ -578,15 +578,25 @@ impl Inserter { if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? { let mut instant_table_ids = HashSet::new(); for req in &requests.inserts { - let table = self - .get_table(catalog, &schema, &req.table_name) - .await? - .context(InvalidInsertRequestSnafu { - reason: format!( - "Table `{}` does not exist, and {}", - req.table_name, disabled_reason - ), - })?; + let table = match self.get_table(catalog, &schema, &req.table_name).await? { + Some(table) => table, + // System-defined table: created canonically by the system, + // so the auto-create config/hint does not apply. + None if is_ddl_reserved_table(&schema, &req.table_name) => { + statement_executor + .create_declared_relationships_table(catalog, ctx.clone()) + .await? + } + None => { + return InvalidInsertRequestSnafu { + reason: format!( + "Table `{}` does not exist, and {}", + req.table_name, disabled_reason + ), + } + .fail(); + } + }; let table_info = table.table_info(); if table_info.is_ttl_instant_table() { instant_table_ids.insert(table_info.table_id()); @@ -630,6 +640,19 @@ impl Inserter { table_infos.insert(table_info.table_id(), table.table_info()); } } + // A DDL-reserved table's definition never derives from the + // write request; the system creates it canonically, below the + // user-DDL guard that rejects the generic create path. + None if is_ddl_reserved_table(&schema, &req.table_name) => { + let table = statement_executor + .create_declared_relationships_table(catalog, ctx.clone()) + .await?; + let table_info = table.table_info(); + if table_info.is_ttl_instant_table() { + instant_table_ids.insert(table_info.table_id()); + } + table_infos.insert(table_info.table_id(), table_info); + } None => { let create_expr = self.get_create_table_expr_on_demand(req, &auto_create_table_type, ctx)?; @@ -911,6 +934,7 @@ impl Inserter { create_type: &AutoCreateTableType, ctx: &QueryContextRef, ) -> Result { + let schema = ctx.current_schema(); let mut table_options = std::collections::HashMap::with_capacity(4); fill_table_options_for_create(&mut table_options, create_type, ctx); apply_per_table_semantic_options(&mut table_options, ctx, &req.table_name); @@ -922,7 +946,6 @@ impl Inserter { default_engine() }; - let schema = ctx.current_schema(); let table_ref = TableReference::full(ctx.current_catalog(), &schema, &req.table_name); // SAFETY: `req.rows` is guaranteed to be `Some` by `handle_row_inserts_with_create_type()`. let request_schema = req.rows.as_ref().unwrap().schema.as_slice(); @@ -963,6 +986,12 @@ impl Inserter { let schema_name = ctx.current_schema(); let table_name = table.table_info().name.clone(); + // Never auto-alter a system-defined table to fit a write; a request + // with unknown columns fails instead. + if is_ddl_reserved_table(&schema_name, &table_name) { + return Ok(None); + } + let request_schema = req.rows.as_ref().unwrap().schema.as_slice(); let request_field_count = request_schema .iter() diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index ae55acd2fb8..3f637ebebcf 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -32,7 +32,8 @@ use catalog::CatalogManagerRef; use chrono::Utc; use common_base::regex_pattern::NAME_PATTERN_REG; use common_catalog::consts::{ - DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, is_readonly_schema, is_readonly_table, + DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, is_ddl_reserved_table, is_readonly_schema, + is_readonly_table, }; use common_catalog::{format_full_flow_name, format_full_table_name}; use common_error::ext::BoxedError; @@ -112,8 +113,8 @@ use crate::error::{ InvalidPartitionSnafu, InvalidSqlSnafu, InvalidTableNameSnafu, InvalidViewNameSnafu, InvalidViewStmtSnafu, NotSupportedSnafu, PartitionExprToPbSnafu, Result, SchemaInUseSnafu, SchemaNotFoundSnafu, SchemaReadOnlySnafu, SerializePartitionExprSnafu, SubstraitCodecSnafu, - TableAlreadyExistsSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu, TableReadOnlySnafu, - UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu, + TableAlreadyExistsSnafu, TableDdlReservedSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu, + TableReadOnlySnafu, UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu, }; use crate::expr_helper::{self, RepartitionRequest, RepartitionSource}; use crate::statement::StatementExecutor; @@ -379,6 +380,19 @@ impl StatementExecutor { .await } + /// Creates the declared-edge table with its canonical schema, entering + /// below the user-DDL guard ([`ensure_table_definition_writable`]). + /// `create_if_not_exists` makes concurrent first inserts race safely. + pub async fn create_declared_relationships_table( + &self, + catalog: &str, + query_ctx: QueryContextRef, + ) -> Result { + let mut expr = super::semantic_graph::build_declared_relationships_expr(catalog); + self.create_non_logic_table(&mut expr, None, query_ctx, TriggerReason::AutoCreate) + .await + } + #[tracing::instrument(skip_all)] pub async fn create_table_inner( &self, @@ -387,7 +401,7 @@ impl StatementExecutor { query_ctx: QueryContextRef, trigger_reason: TriggerReason, ) -> Result { - ensure_table_writable(&create_table.schema_name, &create_table.table_name)?; + ensure_table_definition_writable(&create_table.schema_name, &create_table.table_name)?; if create_table.engine == METRIC_ENGINE_NAME && create_table @@ -1002,6 +1016,9 @@ impl StatementExecutor { expr: CreateViewExpr, ctx: QueryContextRef, ) -> Result { + // A view could otherwise squat a reserved name and block the canonical + // table's first-write creation. + ensure_table_definition_writable(&expr.schema_name, &expr.view_name)?; ensure! { !(expr.create_if_not_exists & expr.or_replace), InvalidSqlSnafu { @@ -1540,7 +1557,7 @@ impl StatementExecutor { query_context: &QueryContextRef, ) -> Result { // Check if the schema is read-only. - ensure_table_writable(&request.schema_name, &request.table_name)?; + ensure_table_definition_writable(&request.schema_name, &request.table_name)?; let table_ref = TableReference::full( &request.catalog_name, @@ -1823,7 +1840,7 @@ impl StatementExecutor { query_context: QueryContextRef, trigger_reason: TriggerReason, ) -> Result { - ensure_table_writable(&expr.schema_name, &expr.table_name)?; + ensure_table_definition_writable(&expr.schema_name, &expr.table_name)?; let catalog_name = if expr.catalog_name.is_empty() { DEFAULT_CATALOG_NAME.to_string() @@ -2328,7 +2345,7 @@ pub fn verify_alter( ); // Renaming INTO a computed graph table's name would let the overlay // shadow the renamed physical table, orphaning its data. - ensure_table_writable(&table_info.schema_name, new_table_name)?; + ensure_table_definition_writable(&table_info.schema_name, new_table_name)?; } else if let AlterKind::AddColumns { columns } = alter_kind { // If all the columns are marked as add_if_not_exists and they already exist in the table, // there is no need to perform the alter. @@ -2548,6 +2565,22 @@ fn ensure_table_writable(schema: &str, table: &str) -> Result<()> { Ok(()) } +/// [`ensure_table_writable`] plus the definition guard for system-defined +/// tables: user CREATE, ALTER and RENAME-into are rejected so the canonical +/// schema cannot be squatted or mutated. DROP and TRUNCATE stay allowed — the +/// next INSERT recreates the table canonically, which is also the recovery +/// path if the canonical definition changes across an upgrade. +fn ensure_table_definition_writable(schema: &str, table: &str) -> Result<()> { + ensure_table_writable(schema, table)?; + ensure!( + !is_ddl_reserved_table(schema, table), + TableDdlReservedSnafu { + name: table.to_string() + } + ); + Ok(()) +} + /// Validates `greptime.semantic.entity..{id|descriptive|scope}` options /// against the table schema: every named column must exist, and `id` columns must /// additionally be tag (primary-key) columns so entity identity stays indexable @@ -2800,7 +2833,9 @@ async fn execute_undrop_table( table_name: TableName, query_context: QueryContextRef, ) -> Result { - ensure_table_writable(&table_name.schema_name, &table_name.table_name)?; + // Undropping restores a table definition and could resurrect a + // pre-canonical shape of a DDL-reserved table; rejected like CREATE. + ensure_table_definition_writable(&table_name.schema_name, &table_name.table_name)?; let dropped = table_metadata_manager .get_dropped_table(&table_name) diff --git a/src/operator/src/statement/dml.rs b/src/operator/src/statement/dml.rs index 827bfd8b663..d72daaaf1aa 100644 --- a/src/operator/src/statement/dml.rs +++ b/src/operator/src/statement/dml.rs @@ -12,19 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_catalog::consts::is_ddl_reserved_table; +use common_error::ext::BoxedError; use common_query::Output; use common_telemetry::tracing; use query::parser::QueryStatement; use session::context::QueryContextRef; +use session::table_name::table_idents_to_full_name; +use snafu::ResultExt; use sql::statements::insert::Insert; use sql::statements::statement::Statement; -use crate::error::Result; +use crate::error::{ExternalSnafu, ParseSqlSnafu, Result}; use crate::statement::StatementExecutor; impl StatementExecutor { #[tracing::instrument(skip_all)] pub async fn insert(&self, insert: Box, query_ctx: QueryContextRef) -> Result { + self.create_ddl_reserved_target_on_demand(&insert, &query_ctx) + .await?; if insert.can_extract_values() { // Fast path: plain insert ("insert with literal values") is executed directly self.inserter @@ -36,4 +42,31 @@ impl StatementExecutor { self.plan_exec(statement, query_ctx).await } } + + /// A DDL-reserved table is defined by the system, not the user: its first + /// INSERT creates it here with the canonical schema. + async fn create_ddl_reserved_target_on_demand( + &self, + insert: &Insert, + query_ctx: &QueryContextRef, + ) -> Result<()> { + let table_name = insert.table_name().context(ParseSqlSnafu)?; + let (catalog, schema, table) = table_idents_to_full_name(table_name, query_ctx) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + if !is_ddl_reserved_table(&schema, &table) { + return Ok(()); + } + let exists = self + .catalog_manager + .table_exists(&catalog, &schema, &table, Some(query_ctx)) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + if !exists { + self.create_declared_relationships_table(&catalog, query_ctx.clone()) + .await?; + } + Ok(()) + } } diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index b845dce3b75..15f76905555 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -12,9 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Typed DataFusion plan builders for the read-time derivation behind the +//! The entity-relationship graph: the physical declared-edge table DDL and the +//! typed DataFusion plan builders for the read-time derivation behind the //! computed `semantic_entities` / `semantic_relationships` tables. //! +//! In OSS the graph is derived at read time, so the *only* stored part is the +//! declared-edge table: edges a user asserts by hand (`provenance = 'declared'`). +//! //! The derivation is built as typed [`Expr`]s over [`DataFrame`]s (never as SQL //! text), so user-controlled identifiers are plain values — no quoting or SQL //! injection surface — and the plans compose with DataFusion's optimizer, @@ -23,10 +27,16 @@ use std::sync::{Arc, LazyLock}; +use api::v1::column_data_type_extension::TypeExt; +use api::v1::{ + ColumnDataType, ColumnDataTypeExtension, ColumnDef, CreateTableExpr, JsonTypeExtension, + SemanticType, +}; use common_catalog::consts::{ - DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, - SPAN_KIND_COLUMN, SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, - TRACE_ID_COLUMN, TRACE_TIMESTAMP_COLUMN, + DEFAULT_PRIVATE_SCHEMA_NAME, DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, SPAN_KIND_COLUMN, + SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, TRACE_ID_COLUMN, + TRACE_TIMESTAMP_COLUMN, }; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; @@ -35,13 +45,221 @@ use datafusion::dataframe::DataFrame; use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns}; use datafusion::functions_aggregate::expr_fn::{count, sum}; use datafusion::functions_nested::expr_fn::make_array; +use datafusion::functions_window::expr_fn::row_number; use datafusion_common::{Column, Result as DfResult, ScalarValue}; use datafusion_expr::{Expr, ExprFunctionExt, JoinType, LogicalPlan, ScalarUDF, cast, ident, lit}; +use store_api::mito_engine_options::{APPEND_MODE_KEY, MERGE_MODE_KEY, TTL_KEY}; + +/// Whether an existing declared-edge table still matches the canonical +/// definition ([`build_declared_relationships_expr`]): columns, time index, +/// primary key, engine, and merge behaviour — the union branch's revision and +/// dedup semantics lean on all of these. A mismatch (upgrade skew) must be +/// surfaced, not silently derived wrong; DROP resets the table. +pub fn declared_relationships_schema_matches(table_info: &table::metadata::TableInfo) -> bool { + let meta = &table_info.meta; + if meta.engine != common_catalog::consts::MITO_ENGINE { + return false; + } + let schema = &meta.schema; + if schema + .timestamp_column() + .is_none_or(|column| column.name != OBSERVED_AT_COLUMN) + { + return false; + } + let primary_keys = meta + .primary_key_indices + .iter() + .filter_map(|idx| schema.column_schemas().get(*idx)) + .map(|column| column.name.as_str()); + if !primary_keys.eq(DECLARED_PRIMARY_KEY_COLUMNS) { + return false; + } + // LastRow merge is what makes a re-asserted edge a *revision*; append mode + // or another merge mode would change the read-side dedup semantics. + let options = &meta.options.extra_options; + if options.get(APPEND_MODE_KEY).map(String::as_str) == Some("true") + || options + .get(MERGE_MODE_KEY) + .is_some_and(|mode| mode != "last_row") + { + return false; + } + + let canonical = build_declared_relationships_expr("greptime"); + if schema.column_schemas().len() != canonical.column_defs.len() { + return false; + } + canonical + .column_defs + .iter() + .zip(schema.column_schemas()) + .all(|(def, column)| { + let Ok(wrapper) = api::helper::ColumnDataTypeWrapper::try_new( + def.data_type, + def.datatype_extension.clone(), + ) else { + return false; + }; + def.name == column.name + && datatypes::prelude::ConcreteDataType::from(wrapper) == column.data_type + }) +} /// Bin width for the temporal window of derived rows: 60s buckets, matching the /// service-graph convention. const BIN_NANOS: i64 = 60 * 1_000_000_000; +/// Time index of the graph tables: when an observation was recorded. +pub const OBSERVED_AT_COLUMN: &str = "observed_at"; + +/// Default retention for the declared-edge table; expiry slides the topology window. +const DEFAULT_DECLARED_RELATIONSHIPS_TTL: &str = "90d"; +/// Environment variable overriding the declared-edge table's TTL at creation +/// time (e.g. `180d`, `forever`). +// TODO(entity-graph): promote this to a real configuration option. +const DECLARED_RELATIONSHIPS_TTL_ENV: &str = "GREPTIMEDB_DECLARED_RELATIONSHIPS_TTL"; + +fn declared_relationships_ttl() -> String { + std::env::var(DECLARED_RELATIONSHIPS_TTL_ENV) + .ok() + .map(|ttl| ttl.trim().to_string()) + .filter(|ttl| !ttl.is_empty()) + .unwrap_or_else(|| DEFAULT_DECLARED_RELATIONSHIPS_TTL.to_string()) +} + +/// The primary-key (tag) columns, in key order. Starting with the source endpoint +/// makes out-edge lookup (`WHERE src_type=? AND src_id=?`) a key-prefix scan; +/// `provenance` and `generation_id` are in the key so a declared edge and a +/// (future) derived edge for the same pair coexist without clobbering. +pub const DECLARED_PRIMARY_KEY_COLUMNS: [&str; 8] = [ + "src_type", + "src_id", + "rel_type", + "dst_type", + "dst_id", + "provenance", + "scope", + "generation_id", +]; + +/// The externally visible edge identity: the primary key minus `scope` and +/// `generation_id`, which the computed table does not expose. Revision ranking +/// uses this identity, or assertions differing only in those two columns would +/// surface as indistinguishable duplicate rows. +const DECLARED_EDGE_IDENTITY_COLUMNS: [&str; 6] = [ + "src_type", + "src_id", + "rel_type", + "dst_type", + "dst_id", + "provenance", +]; + +fn column( + name: &str, + data_type: ColumnDataType, + semantic_type: SemanticType, + nullable: bool, +) -> ColumnDef { + ColumnDef { + name: name.to_string(), + data_type: data_type as i32, + is_nullable: nullable, + default_constraint: vec![], + semantic_type: semantic_type as i32, + comment: String::new(), + datatype_extension: None, + options: None, + } +} + +fn tag(name: &str) -> ColumnDef { + column(name, ColumnDataType::String, SemanticType::Tag, false) +} + +fn field(name: &str, data_type: ColumnDataType) -> ColumnDef { + column(name, data_type, SemanticType::Field, true) +} + +fn json_field(name: &str) -> ColumnDef { + let mut def = field(name, ColumnDataType::Binary); + def.datatype_extension = Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), + }); + def +} + +/// Builds the `CREATE TABLE` request for the declared-edge table. Columns mirror +/// the computed `semantic_relationships` shape (temporal window + endpoints + +/// provenance/confidence + RED metrics) plus the declared-only business validity +/// window (`valid_from` / `valid_until`), which — unlike TTL (physical retention) +/// — expresses whether a hand-declared edge is still in effect. +/// +/// The default `LastRow` merge dedups on primary key **plus** `observed_at` +/// (the time index is part of mito's dedup key): re-asserting an edge at a new +/// `observed_at` stores a new revision, and the read-side union keeps only the +/// latest revision per edge key. +pub fn build_declared_relationships_expr(catalog: &str) -> CreateTableExpr { + let column_defs = vec![ + // Temporal: when this revision of the edge was declared (time index, + // also the TTL clock). + column( + OBSERVED_AT_COLUMN, + ColumnDataType::TimestampMillisecond, + SemanticType::Timestamp, + false, + ), + field("window_start", ColumnDataType::TimestampMillisecond), + field("window_end", ColumnDataType::TimestampMillisecond), + field("fresh_until", ColumnDataType::TimestampMillisecond), + // Declared-only business validity. NULL valid_from = valid since the + // declaration; NULL valid_until = valid for as long as the row exists + // (TTL expiry retires the edge with the row). + field("valid_from", ColumnDataType::TimestampMillisecond), + field("valid_until", ColumnDataType::TimestampMillisecond), + // Endpoints + edge identity (all tags, in primary-key order). + tag("src_type"), + tag("src_id"), + tag("rel_type"), + tag("dst_type"), + tag("dst_id"), + tag("provenance"), + tag("scope"), + tag("generation_id"), + // Confidence + RED metrics (populated for derived edges; usually NULL here). + field("confidence", ColumnDataType::Float64), + field("request_count", ColumnDataType::Int64), + field("error_count", ColumnDataType::Int64), + field("duration_sum", ColumnDataType::Float64), + field("duration_count", ColumnDataType::Int64), + // JSONB, so the union matches the computed table's json column + // without a per-scan parse. + json_field("attributes"), + ]; + + let table_options = [(TTL_KEY.to_string(), declared_relationships_ttl())] + .into_iter() + .collect(); + + CreateTableExpr { + catalog_name: catalog.to_string(), + schema_name: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(), + table_name: SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME.to_string(), + desc: "Hand-declared edges of the entity-relationship graph".to_string(), + column_defs, + time_index: OBSERVED_AT_COLUMN.to_string(), + primary_keys: DECLARED_PRIMARY_KEY_COLUMNS + .iter() + .map(|c| c.to_string()) + .collect(), + create_if_not_exists: true, + table_options, + table_id: None, + engine: common_catalog::consts::MITO_ENGINE.to_string(), + } +} + /// A single table's entity-identity declaration, projected from /// `information_schema.table_semantics` (`greptime.semantic.entity..*`). #[derive(Debug, Clone)] @@ -62,33 +280,85 @@ pub struct EntityDeclaration { pub scope_columns: Vec, } -/// The read-time query window as scalar [`Expr`]s. Kept as expressions so the -/// provider can fill them from a scan's time predicate or the product default. +/// The read-time query window, resolved from the scan's `observed_at` predicate +/// (or the product default). Two half-open `[start, end)` millisecond ranges: +/// +/// - *observed*: the queried `observed_at` range. The scan's filters are +/// re-applied above the computed table (`FilterPushDownType::Inexact`), so +/// every emitted row's `observed_at` must fall inside it. The declared-edge +/// branch also runs its validity overlap against this range. +/// - *source scan*: `observed` widened to whole 60s buckets. Derived rows are +/// `date_bin('60s')` aggregates keyed by the bucket start, so a bucket the +/// observed range selects must be scanned over its full extent — filtering +/// the source rows to the observed range instead would silently truncate the +/// RED numbers of the boundary buckets. +/// +/// The bounds are plain millisecond values turned into timestamp literals on +/// demand — wall-clock snapshot semantics like `now()`, but as constants they +/// prune the source table scans without depending on constant-folding. #[derive(Debug, Clone)] -pub struct GraphWindow { - pub start: Expr, - pub end: Expr, +pub struct GraphQueryWindow { + observed_start_ms: i64, + observed_end_ms: i64, + source_start_ms: i64, + source_end_ms: i64, } -impl GraphWindow { +impl GraphQueryWindow { + /// Builds the window for the queried `observed_at` range `[start_ms, end_ms)`. + pub fn from_observed(start_ms: i64, end_ms: i64) -> Self { + const BIN_MS: i64 = BIN_NANOS / 1_000_000; + // Ceiling division that stays correct for pre-epoch (negative) values. + let ceil_to_bin = |ms: i64| { + ms.div_euclid(BIN_MS) * BIN_MS + + if ms.rem_euclid(BIN_MS) == 0 { + 0 + } else { + BIN_MS + } + }; + // A bucket `b` (a 60s multiple, `date_bin`'s floor of its rows) is + // selected when `start <= b < end`: the first is ceil(start) and the + // last one's rows extend to ceil(end). + Self { + observed_start_ms: start_ms, + observed_end_ms: end_ms, + source_start_ms: ceil_to_bin(start_ms), + source_end_ms: ceil_to_bin(end_ms), + } + } + /// Conservative default when a query carries no explicit time predicate: the /// last hour, so a bare `SELECT * FROM semantic_entities` never scans every /// declaring table's full history. This is a product default, not a cap. - /// - /// The bounds are timestamp literals taken from the wall clock at plan-build - /// time — the same snapshot semantics `now()` would give, but as plain - /// constants they prune the source table scans without depending on - /// constant-folding. pub fn default_last_hour() -> Self { let end_ms = common_time::util::current_time_millis(); - Self { - start: lit(ScalarValue::TimestampMillisecond( - Some(end_ms - 60 * 60 * 1000), - None, - )), - end: lit(ScalarValue::TimestampMillisecond(Some(end_ms), None)), - } + Self::from_observed(end_ms - 60 * 60 * 1000, end_ms) } + + /// Start of the queried `observed_at` range (inclusive). + pub fn observed_start(&self) -> Expr { + ts_ms_lit(self.observed_start_ms) + } + + /// End of the queried `observed_at` range (exclusive). + pub fn observed_end(&self) -> Expr { + ts_ms_lit(self.observed_end_ms) + } + + /// Start of the source-table scan range (inclusive). + pub fn source_start(&self) -> Expr { + ts_ms_lit(self.source_start_ms) + } + + /// End of the source-table scan range (exclusive). + pub fn source_end(&self) -> Expr { + ts_ms_lit(self.source_end_ms) + } +} + +fn ts_ms_lit(ms: i64) -> Expr { + lit(ScalarValue::TimestampMillisecond(Some(ms), None)) } /// An `INTERVAL` literal of `nanos` nanoseconds. @@ -254,14 +524,14 @@ fn registry_source( first: &EntityDeclaration, rest: &[EntityDeclaration], df: DataFrame, - window: &GraphWindow, + window: &GraphQueryWindow, ) -> DfResult { let ts = ident(&first.time_index); let bin = bin_ms(ts.clone()); let window_predicate = ts .clone() - .gt_eq(window.start.clone()) - .and(ts.lt(window.end.clone())); + .gt_eq(window.source_start()) + .and(ts.lt(window.source_end())); let mut arrays = vec![Vec::new(); REGISTRY_COLUMNS.len() + 1]; for decl in std::iter::once(first).chain(rest) { @@ -350,19 +620,25 @@ pub struct RegistrySource { } /// A trace table's scan paired with the `service` declaration it derives -/// `calls` edges for — the unit `build_calls_plan` consumes. +/// `calls` edges for — a unit of `build_relationships_plan`. pub struct CallsSource { pub service: EntityDeclaration, pub scan: DataFrame, } +/// The declared-edge table's scan (`semantic_relationships_declared`), whose +/// rows `build_relationships_plan` unions into the edge set. +pub struct DeclaredSource { + pub scan: DataFrame, +} + /// Builds the `semantic_entities` registry plan: one branch and source scan per /// declaring table, filtered to `window`, then `UNION ALL` across source tables. /// Returns `None` when nothing declared an entity, so the computed table streams /// empty. pub fn build_registry_plan( sources: Vec, - window: &GraphWindow, + window: &GraphQueryWindow, ) -> DfResult> { let mut union_df: Option = None; for source in sources { @@ -375,11 +651,11 @@ pub fn build_registry_plan( } /// The projected columns of `semantic_relationships`, in order. Every derived -/// branch and the declared-edge scan must project exactly these so the top-level -/// `UNION ALL` type-aligns. (The physical declared table additionally stores -/// `valid_from`/`valid_until`, which are applied as a filter, not projected.) -/// Test-only until the declared-edge union branch lands and enforces it in code. -#[cfg(test)] +/// branch and the declared-edge branch must project exactly these so the +/// top-level `UNION ALL` type-aligns; `build_relationships_plan` re-selects +/// them over the union to enforce the contract. (The physical declared table +/// additionally stores `valid_from`/`valid_until`, which feed the validity +/// filter and the projected window columns.) const RELATIONSHIP_COLUMNS: [&str; 16] = [ "observed_at", "window_start", @@ -406,21 +682,117 @@ const RELATIONSHIP_COLUMNS: [&str; 16] = [ const CHILD_SPAN_EARLY_NANOS: i64 = 5 * 60 * 1_000_000_000; const CHILD_SPAN_LATE_NANOS: i64 = 60 * 60 * 1_000_000_000; -/// Builds the `calls` derivation (RFC §3a) over `traces`: pair each client span -/// with its child server span on `trace_id` + `parent_span_id`, project to -/// `service`, union the pairs of all trace tables, and aggregate to RED metrics -/// per 60s window in one pass — so an edge observed across several trace tables -/// yields one row, not per-table fragments. This is the plan form of the Tempo -/// servicegraph connector. Virtual-node edges (uninstrumented peers) are a -/// separate branch, added on top of this. Column names are the fixed -/// `greptime_trace_v1` schema (the reason `table_data_model = greptime_trace_v1` -/// is required); `span_status_code` is a string column (`STATUS_CODE_ERROR`), -/// verified against the trace ingest path. Returns `None` when there is no -/// trace table. -pub fn build_calls_plan( +/// Builds the `semantic_relationships` plan: the trace-derived `calls` branch +/// unioned with the declared-edge branch, re-projected to the 16-column +/// contract. Returns `None` when there is neither a trace table nor a +/// declared-edge table, so the computed table streams empty. +pub fn build_relationships_plan( traces: Vec, - window: &GraphWindow, + declared: Option, + window: &GraphQueryWindow, ) -> DfResult> { + let mut union_df = calls_branch(traces, window)?; + if let Some(declared) = declared { + union_df = union_all(union_df, declared_edges(declared.scan, window)?)?; + } + union_df + .map(|df| { + Ok(df + .select_columns(&RELATIONSHIP_COLUMNS)? + .into_unoptimized_plan()) + }) + .transpose() +} + +/// Ranking column used to pick the latest declared-edge revision per edge key. +const DECLARED_REVISION_COLUMN: &str = "__declared_revision"; + +/// The declared-edge branch: the latest revision per edge key *as of the +/// queried window* (mito dedups on primary key + `observed_at`, so re-asserting +/// an edge stores a new revision), filtered by business-validity overlap. +/// +/// `valid_from` defaults to the declaration time; a NULL `valid_until` means +/// the edge holds while its row exists (TTL retires it), so `window_end` / +/// `fresh_until` take the window's upper bound. The output `observed_at` is +/// synthesized inside the queried range: the scan's filters are re-applied +/// above the computed table, and the physical revision time would fail them. +fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult { + let eff_valid_from = + core_fns::coalesce().call(vec![ident("valid_from"), ident(OBSERVED_AT_COLUMN)]); + let eff_valid_until = + core_fns::coalesce().call(vec![ident("valid_until"), window.observed_end()]); + + // As-of the queried window: a revision recorded after its end, or whose + // validity starts after it, was not the edge's state inside the window and + // must not outrank (and thereby hide) the revision that was. + let as_of = ident(OBSERVED_AT_COLUMN) + .lt(window.observed_end()) + .and(eff_valid_from.clone().lt(window.observed_end())); + + let revision = row_number() + .partition_by( + DECLARED_EDGE_IDENTITY_COLUMNS + .iter() + .map(|c| ident(*c)) + .collect(), + ) + // generation_id/scope break observed_at ties deterministically (rows + // differing only in them are not merged by the storage dedup). + .order_by(vec![ + ident(OBSERVED_AT_COLUMN).sort(false, false), + ident("generation_id").sort(false, false), + ident("scope").sort(false, false), + ]) + .build()? + .alias(DECLARED_REVISION_COLUMN); + + // The as-of filter already bounds eff_valid_from below the window's end; + // only the retirement side of the overlap remains to check. + let still_valid = ident("valid_until") + .is_null() + .or(ident("valid_until").gt(window.observed_start())); + // Tags come out of the storage engine dictionary-encoded; cast to plain + // strings so the union with the derived branches type-aligns. + let tag_utf8 = |name: &str| cast(ident(name), DataType::Utf8).alias(name); + + scan.filter(as_of)? + .window(vec![revision])? + .filter(ident(DECLARED_REVISION_COLUMN).eq(lit(1_u64)))? + .filter(still_valid)? + .select(vec![ + core_fns::greatest() + .call(vec![eff_valid_from.clone(), window.observed_start()]) + .alias(OBSERVED_AT_COLUMN), + eff_valid_from.alias("window_start"), + eff_valid_until.clone().alias("window_end"), + eff_valid_until.alias("fresh_until"), + tag_utf8("src_type"), + tag_utf8("src_id"), + tag_utf8("dst_type"), + tag_utf8("dst_id"), + tag_utf8("rel_type"), + tag_utf8("provenance"), + ident("confidence"), + ident("request_count"), + ident("error_count"), + ident("duration_sum"), + ident("duration_count"), + ident("attributes"), + ]) +} + +/// The `calls` derivation: pair each client span with its child server span on +/// `trace_id` + `parent_span_id`, project to `service`, union the pairs of all +/// trace tables, and aggregate RED metrics per 60s window in one pass — an edge +/// observed across several trace tables yields one row, not per-table +/// fragments (the plan form of the Tempo servicegraph connector). Column names +/// are the fixed `greptime_trace_v1` schema, the reason +/// `table_data_model = greptime_trace_v1` is required. Returns `None` when +/// there is no trace table. +fn calls_branch( + traces: Vec, + window: &GraphQueryWindow, +) -> DfResult> { let mut union_df: Option = None; for trace in traces { union_df = union_all(union_df, calls_pairs(&trace.service, trace.scan, window)?)?; @@ -429,7 +801,7 @@ pub fn build_calls_plan( return Ok(None); }; - let plan = pairs + let df = pairs .aggregate( vec![ident("observed_at"), ident("src_id"), ident("dst_id")], vec![ @@ -460,9 +832,8 @@ pub fn build_calls_plan( .alias("duration_sum"), ident("request_count").alias("duration_count"), null_json().alias("attributes"), - ])? - .into_unoptimized_plan(); - Ok(Some(plan)) + ])?; + Ok(Some(df)) } /// One trace table's client/server span pairs, projected to the aggregation @@ -473,12 +844,12 @@ pub fn build_calls_plan( fn calls_pairs( service: &EntityDeclaration, trace: DataFrame, - window: &GraphWindow, + window: &GraphQueryWindow, ) -> DfResult { let mut client_pred = ident(SPAN_KIND_COLUMN) .eq(lit(SPAN_KIND_CLIENT)) - .and(ident(TRACE_TIMESTAMP_COLUMN).gt_eq(window.start.clone())) - .and(ident(TRACE_TIMESTAMP_COLUMN).lt(window.end.clone())); + .and(ident(TRACE_TIMESTAMP_COLUMN).gt_eq(window.source_start())) + .and(ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end())); let mut server_pred = ident(SPAN_KIND_COLUMN) .eq(lit(SPAN_KIND_SERVER)) // Static bounds implied by the window and the join's time-proximity @@ -486,10 +857,10 @@ fn calls_pairs( // cannot prune the server-side scan. .and( ident(TRACE_TIMESTAMP_COLUMN) - .gt_eq(window.start.clone() - interval(CHILD_SPAN_EARLY_NANOS)), + .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)), ) .and( - ident(TRACE_TIMESTAMP_COLUMN).lt(window.end.clone() + interval(CHILD_SPAN_LATE_NANOS)), + ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS)), ); // A NULL identity component identifies nothing, on either endpoint. for id in &service.id_columns { @@ -541,14 +912,34 @@ mod tests { use super::*; - fn test_window() -> GraphWindow { - GraphWindow { - start: lit(ScalarValue::TimestampMillisecond(Some(0), None)), - end: lit(ScalarValue::TimestampMillisecond( - Some(10 * 60 * 1000), - None, - )), - } + fn test_window() -> GraphQueryWindow { + GraphQueryWindow::from_observed(0, 10 * 60 * 1000) + } + + #[test] + fn window_source_scan_covers_whole_buckets() { + // Buckets selected by [10:30.5', 11:30') is exactly {11:00'}; its source + // rows span [11:00', 12:00'). + let window = GraphQueryWindow::from_observed(630_000, 690_000); + assert_eq!( + (window.source_start_ms, window.source_end_ms), + (660_000, 720_000) + ); + + // Aligned bounds stay put. + let window = GraphQueryWindow::from_observed(600_000, 720_000); + assert_eq!( + (window.source_start_ms, window.source_end_ms), + (600_000, 720_000) + ); + + // A sub-bucket range that selects no bucket start scans nothing. + let window = GraphQueryWindow::from_observed(610_000, 650_000); + assert!(window.source_start_ms >= window.source_end_ms); + + // Pre-epoch bounds round toward the correct buckets. + let window = GraphQueryWindow::from_observed(-90_000, -30_000); + assert_eq!((window.source_start_ms, window.source_end_ms), (-60_000, 0)); } async fn collect(ctx: &SessionContext, plan: LogicalPlan) -> Vec { @@ -954,7 +1345,7 @@ mod tests { async fn calls_plan_aggregates_red_metrics() { let ctx = trace_table_ctx(); let trace = ctx.table("opentelemetry_traces").await.unwrap(); - let plan = build_calls_plan( + let plan = build_relationships_plan_for_test( vec![CallsSource { service: trace_service_decl(&[SERVICE_NAME_COLUMN]), scan: trace, @@ -1016,7 +1407,7 @@ mod tests { let ctx = trace_table_ctx(); let t1 = ctx.table("opentelemetry_traces").await.unwrap(); let t2 = ctx.table("opentelemetry_traces_2").await.unwrap(); - let plan = build_calls_plan( + let plan = build_relationships_plan_for_test( vec![ CallsSource { service: trace_service_decl(&[SERVICE_NAME_COLUMN]), @@ -1062,7 +1453,7 @@ mod tests { async fn calls_endpoints_follow_service_declaration() { let ctx = trace_table_ctx(); let trace = ctx.table("opentelemetry_traces").await.unwrap(); - let plan = build_calls_plan( + let plan = build_relationships_plan_for_test( vec![CallsSource { service: trace_service_decl(&[SERVICE_NAME_COLUMN, "service_namespace"]), scan: trace, @@ -1087,4 +1478,487 @@ mod tests { vec!["service_name=cart,service_namespace=ns1"] ); } + + fn build_relationships_plan_for_test( + traces: Vec, + window: &GraphQueryWindow, + ) -> DfResult> { + build_relationships_plan(traces, None, window) + } + + /// A `TableInfo` derived from the canonical declared-edge definition, with + /// injection points for each aspect the schema matcher must check. + fn declared_table_info( + mutate_columns: impl FnOnce(&mut Vec), + engine: &str, + primary_key_indices: Vec, + extra_options: &[(&str, &str)], + ) -> table::metadata::TableInfo { + let canonical = build_declared_relationships_expr("greptime"); + let mut columns: Vec = canonical + .column_defs + .iter() + .map(|def| { + let wrapper = api::helper::ColumnDataTypeWrapper::try_new( + def.data_type, + def.datatype_extension.clone(), + ) + .unwrap(); + datatypes::schema::ColumnSchema::new( + &def.name, + datatypes::prelude::ConcreteDataType::from(wrapper), + def.is_nullable, + ) + .with_time_index(def.name == OBSERVED_AT_COLUMN) + }) + .collect(); + mutate_columns(&mut columns); + let schema = Arc::new( + datatypes::schema::SchemaBuilder::try_from_columns(columns) + .unwrap() + .build() + .unwrap(), + ); + let options = table::requests::TableOptions { + extra_options: extra_options + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(), + ..Default::default() + }; + let meta = table::metadata::TableMeta { + schema, + primary_key_indices, + value_indices: vec![], + engine: engine.to_string(), + next_column_id: 1, + options, + created_on: Default::default(), + updated_on: Default::default(), + partition_key_indices: vec![], + column_ids: vec![], + }; + table::metadata::TableInfoBuilder::default() + .table_id(1) + .name(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .catalog_name("greptime") + .schema_name(DEFAULT_PRIVATE_SCHEMA_NAME) + .table_version(0) + .table_type(table::metadata::TableType::Base) + .meta(meta) + .build() + .unwrap() + } + + #[test] + fn declared_schema_match_checks_more_than_columns() { + let mito = common_catalog::consts::MITO_ENGINE; + let canonical_pk = || (6..14).collect::>(); + + let ok = declared_table_info(|_| {}, mito, canonical_pk(), &[(TTL_KEY, "180d")]); + assert!(declared_relationships_schema_matches(&ok)); + // The canonical expr must declare attributes as json, matching the + // computed table's column. + assert_eq!( + ok.meta + .schema + .column_schema_by_name("attributes") + .map(|c| c.data_type.clone()), + Some(datatypes::prelude::ConcreteDataType::json_datatype()) + ); + let ok = declared_table_info( + |_| {}, + mito, + canonical_pk(), + &[(MERGE_MODE_KEY, "last_row")], + ); + assert!(declared_relationships_schema_matches(&ok)); + + let wrong_engine = declared_table_info(|_| {}, "metric", canonical_pk(), &[]); + assert!(!declared_relationships_schema_matches(&wrong_engine)); + + let pk_missing_generation_id = declared_table_info(|_| {}, mito, (6..13).collect(), &[]); + assert!(!declared_relationships_schema_matches( + &pk_missing_generation_id + )); + + let wrong_time_index = declared_table_info( + |columns| { + columns[0] = columns[0].clone().with_time_index(false); + columns[4] = datatypes::schema::ColumnSchema::new( + "valid_from", + datatypes::prelude::ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true); + }, + mito, + canonical_pk(), + &[], + ); + assert!(!declared_relationships_schema_matches(&wrong_time_index)); + + let append_mode = + declared_table_info(|_| {}, mito, canonical_pk(), &[(APPEND_MODE_KEY, "true")]); + assert!(!declared_relationships_schema_matches(&append_mode)); + + let wrong_column_type = declared_table_info( + |columns| { + columns[14] = datatypes::schema::ColumnSchema::new( + "confidence", + datatypes::prelude::ConcreteDataType::int64_datatype(), + true, + ); + }, + mito, + canonical_pk(), + &[], + ); + assert!(!declared_relationships_schema_matches(&wrong_column_type)); + } + + fn ts_values(batch: &RecordBatch, column: usize) -> Vec { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()).map(|i| array.value(i)).collect() + } + + /// Registers a declared-edge table holding: + /// - frontend->db: two revisions (0.5 then 1.0), the newer retiring the edge at 300s; + /// - api->cache: expired at 100s; + /// - agent-1->search: open-ended `provenance = 'agent'` with JSON attributes; + /// - webapp->auth: two revisions, the newer declared at 900s; + /// - batch->queue: declared at 100s but only valid from 700s; + /// - gateway->redis: asserted twice at the same instant as generations g1/g2. + fn declared_table(ctx: &SessionContext) { + let ts_field = |name: &str| { + Field::new( + name, + DataType::Timestamp(TimeUnit::Millisecond, None), + name != OBSERVED_AT_COLUMN, + ) + }; + let schema = Arc::new(Schema::new(vec![ + ts_field(OBSERVED_AT_COLUMN), + ts_field("window_start"), + ts_field("window_end"), + ts_field("fresh_until"), + ts_field("valid_from"), + ts_field("valid_until"), + Field::new("src_type", DataType::Utf8, false), + Field::new("src_id", DataType::Utf8, false), + Field::new("rel_type", DataType::Utf8, false), + Field::new("dst_type", DataType::Utf8, false), + Field::new("dst_id", DataType::Utf8, false), + Field::new("provenance", DataType::Utf8, false), + Field::new("scope", DataType::Utf8, false), + Field::new("generation_id", DataType::Utf8, false), + Field::new("confidence", DataType::Float64, true), + Field::new("request_count", DataType::Int64, true), + Field::new("error_count", DataType::Int64, true), + Field::new("duration_sum", DataType::Float64, true), + Field::new("duration_count", DataType::Int64, true), + Field::new("attributes", DataType::Binary, true), + ])); + let attrs = jsonb::parse_value(br#"{"connection_type":"virtual_node"}"#) + .unwrap() + .to_vec(); + let no_ts = || Arc::new(TimestampMillisecondArray::from(vec![None::; 9])) as ArrayRef; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 100_000, 200_000, 50_000, 90_000, 100_000, 900_000, 100_000, 100_000, 100_000, + ])) as ArrayRef, + no_ts(), // window_start + no_ts(), // window_end + no_ts(), // fresh_until + Arc::new(TimestampMillisecondArray::from(vec![ + None, + None, + None, + None, + None, + None, + Some(700_000), + None, + None, + ])), + Arc::new(TimestampMillisecondArray::from(vec![ + None, + Some(300_000), + Some(100_000), + None, + None, + None, + None, + None, + None, + ])), + Arc::new(StringArray::from(vec![ + "service", "service", "service", "agent", "service", "service", "service", + "service", "service", + ])), + Arc::new(StringArray::from(vec![ + "frontend", "frontend", "api", "agent-1", "webapp", "webapp", "batch", + "gateway", "gateway", + ])), + Arc::new(StringArray::from(vec![ + "depends_on", + "depends_on", + "depends_on", + "uses", + "depends_on", + "depends_on", + "depends_on", + "depends_on", + "depends_on", + ])), + Arc::new(StringArray::from(vec![ + "service", "service", "service", "tool", "service", "service", "service", + "service", "service", + ])), + Arc::new(StringArray::from(vec![ + "db", "db", "cache", "search", "auth", "auth", "queue", "redis", "redis", + ])), + Arc::new(StringArray::from(vec![ + "declared", "declared", "declared", "agent", "declared", "declared", + "declared", "declared", "declared", + ])), + Arc::new(StringArray::from(vec![""; 9])), + Arc::new(StringArray::from(vec![ + "", "", "", "", "", "", "", "g1", "g2", + ])), + Arc::new(Float64Array::from(vec![ + Some(0.5), + Some(1.0), + Some(1.0), + Some(0.8), + Some(0.5), + Some(1.0), + Some(1.0), + Some(0.5), + Some(1.0), + ])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(Float64Array::from(vec![None::; 9])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(BinaryArray::from(vec![ + None, + None, + None, + Some(attrs.as_slice()), + None, + None, + None, + None, + None, + ])), + ], + ) + .unwrap(); + ctx.register_table( + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + } + + #[tokio::test] + async fn declared_edges_latest_revision_validity_and_synthetic_time() { + let ctx = SessionContext::new(); + declared_table(&ctx); + let scan = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + // Queried window [120s, 600s). + let window = GraphQueryWindow::from_observed(120_000, 600_000); + let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) + .unwrap() + .unwrap(); + let names = plan + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(); + assert_eq!(names, RELATIONSHIP_COLUMNS); + + // (src_id, observed_at, window_start, window_end, fresh_until, + // provenance, attributes) + type DeclaredRow = (String, i64, i64, i64, i64, String, Option); + + let batches = collect(&ctx, plan).await; + let mut rows: Vec = vec![]; + for batch in &batches { + let observed = ts_values(batch, 0); + let window_start = ts_values(batch, 1); + let window_end = ts_values(batch, 2); + let fresh_until = ts_values(batch, 3); + let src = strings(batch, 5); + let provenance = strings(batch, 9); + let attributes = json_texts(batch, 15); + for i in 0..batch.num_rows() { + rows.push(( + src[i].clone(), + observed[i], + window_start[i], + window_end[i], + fresh_until[i], + provenance[i].clone(), + attributes[i].clone(), + )); + } + } + rows.sort(); + + // api->cache expired before the window; batch->queue is not yet valid; + // frontend->db keeps only the latest revision; webapp->auth keeps the + // first revision (the second postdates the window); gateway->redis + // collapses to one row across generations. + assert_eq!(rows.len(), 4); + + // agent->tool: open-ended validity declared at 90s. The synthesized + // observed_at is clamped into the queried window; window_end and + // fresh_until take the window's upper bound. + assert_eq!( + rows[0], + ( + "agent-1".to_string(), + 120_000, + 90_000, + 600_000, + 600_000, + "agent".to_string(), + Some(r#"{"connection_type":"virtual_node"}"#.to_string()), + ) + ); + + // frontend->db: the latest revision (declared at 200s, retired at 300s) + // supersedes the open-ended first revision. + assert_eq!( + rows[1], + ( + "frontend".to_string(), + 200_000, + 200_000, + 300_000, + 300_000, + "declared".to_string(), + None, + ) + ); + + assert_eq!( + rows[2], + ( + "gateway".to_string(), + 120_000, + 100_000, + 600_000, + 600_000, + "declared".to_string(), + None, + ) + ); + + assert_eq!( + rows[3], + ( + "webapp".to_string(), + 120_000, + 100_000, + 600_000, + 600_000, + "declared".to_string(), + None, + ) + ); + } + + #[tokio::test] + async fn declared_edges_pick_the_revision_as_of_the_window() { + let ctx = SessionContext::new(); + declared_table(&ctx); + let scan = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + // A historical window that predates frontend->db's second revision: + // the first revision was the edge's state then and must not be hidden + // by the newer one. + let window = GraphQueryWindow::from_observed(0, 150_000); + let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let mut rows: Vec<(String, f64)> = vec![]; + for batch in &batches { + let src = strings(batch, 5); + let confidence = batch + .column(10) + .as_any() + .downcast_ref::() + .unwrap(); + for (i, src) in src.into_iter().enumerate() { + rows.push((src, confidence.value(i))); + } + } + rows.sort_by(|a, b| a.0.cmp(&b.0)); + + // api->cache is valid inside this window; batch->queue is not yet. + assert_eq!( + rows, + vec![ + ("agent-1".to_string(), 0.8), + ("api".to_string(), 1.0), + ("frontend".to_string(), 0.5), + ("gateway".to_string(), 1.0), + ("webapp".to_string(), 0.5), + ] + ); + } + + #[tokio::test] + async fn declared_edges_union_calls_branch() { + let ctx = trace_table_ctx(); + declared_table(&ctx); + let trace = ctx.table("opentelemetry_traces").await.unwrap(); + let declared = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + let plan = build_relationships_plan( + vec![CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: trace, + }], + Some(DeclaredSource { scan: declared }), + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let mut provenance: Vec = batches.iter().flat_map(|b| strings(b, 9)).collect(); + provenance.sort(); + // One trace-derived calls edge plus the declared edges valid somewhere + // inside [0s, 600s): frontend, api, webapp, gateway, and the + // agent-asserted one. + assert_eq!( + provenance, + vec![ + "agent", "declared", "declared", "declared", "declared", "trace" + ] + ); + } } diff --git a/src/table/src/predicate.rs b/src/table/src/predicate.rs index 29b10430492..5fb341bc63b 100644 --- a/src/table/src/predicate.rs +++ b/src/table/src/predicate.rs @@ -23,6 +23,7 @@ use datafusion::common::ScalarValue; use datafusion::physical_optimizer::pruning::PruningPredicate; use datafusion_common::ToDFSchema; use datafusion_common::pruning::PruningStatistics; +use datafusion_common::tree_node::TreeNode; use datafusion_expr::expr::{Expr, InList}; use datafusion_expr::{Between, BinaryExpr, Operator}; use datafusion_physical_expr::execution_props::ExecutionProps; @@ -217,6 +218,83 @@ pub fn build_time_range_predicate( res } +/// The outcome of strictly extracting the time range of `ts_col_name` from a +/// scan's filters. Unlike [`build_time_range_predicate`], which quietly widens +/// to `min_to_max` on anything it cannot parse (fine for pruning), this +/// distinguishes "the column is not filtered at all" from "it is filtered in a +/// way that cannot be safely turned into a range" — for callers whose contract +/// forbids silently falling back to a default window. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TimeRangeExtraction { + /// No filter references the column. + Absent, + /// Every filter referencing the column was folded into this range, which + /// over-approximates the filters' satisfying set. + Extracted(TimestampRange), + /// At least one filter references the column in a shape that cannot be + /// safely extracted (e.g. under `OR`/`NOT`, or compared to a non-literal). + Unsupported, +} + +/// Strictly extracts the time range of `ts_col_name` from the (implicitly +/// AND-ed) scan filters. See [`TimeRangeExtraction`]. +pub fn extract_time_range_strict( + ts_col_name: &str, + ts_col_unit: TimeUnit, + filters: &[Expr], +) -> TimeRangeExtraction { + let mut range: Option = None; + for expr in filters { + if !expr + .column_refs() + .iter() + .any(|column| column.name == ts_col_name) + { + continue; + } + // Disjunctive shapes (`OR`, `IN`) collapse disjoint ranges into their + // convex hull, which is not exactly representable as one contiguous + // range; callers deriving synthetic timestamps from the bounds would + // emit points inside the gaps. + if contains_disjunction_over_column(expr, ts_col_name) { + return TimeRangeExtraction::Unsupported; + } + // `Some` from the lenient extractor over-approximates the expression's + // satisfying set (an `AND` side it cannot parse is dropped, which only + // widens), so intersecting extracted conjuncts stays an + // over-approximation. + match extract_time_range_from_expr(ts_col_name, ts_col_unit, expr) { + Some(extracted) => { + range = Some(match range { + Some(acc) => acc.and(&extracted), + None => extracted, + }); + } + None => return TimeRangeExtraction::Unsupported, + } + } + match range { + Some(range) => TimeRangeExtraction::Extracted(range), + None => TimeRangeExtraction::Absent, + } +} + +fn contains_disjunction_over_column(expr: &Expr, ts_col_name: &str) -> bool { + let is_disjunction = |expr: &Expr| { + matches!( + expr, + Expr::BinaryExpr(BinaryExpr { + op: Operator::Or, + .. + }) | Expr::InList(_) + ) && expr + .column_refs() + .iter() + .any(|column| column.name == ts_col_name) + }; + expr.exists(|expr| Ok(is_disjunction(expr))).unwrap_or(true) +} + /// Extract time range filter from `WHERE`/`IN (...)`/`BETWEEN` clauses. /// Return None if no time range can be found in expr. pub fn extract_time_range_from_expr( @@ -593,6 +671,99 @@ mod tests { ); } + #[test] + fn test_extract_time_range_strict() { + fn ts_lit(ms: i64) -> Expr { + lit(ScalarValue::TimestampMillisecond(Some(ms), None)) + } + let extract = + |filters: &[Expr]| extract_time_range_strict("ts", TimeUnit::Millisecond, filters); + let range = |start: i64, end: i64| { + TimestampRange::new( + Timestamp::new_millisecond(start), + Timestamp::new_millisecond(end), + ) + .unwrap() + }; + + // No filter references the column. + assert_eq!(extract(&[]), TimeRangeExtraction::Absent); + assert_eq!( + extract(&[col("host").eq(lit("a"))]), + TimeRangeExtraction::Absent + ); + + // Both bounds across conjuncts; unrelated filters are ignored. + assert_eq!( + extract(&[ + col("ts").gt_eq(ts_lit(1000)), + col("ts").lt(ts_lit(2000)), + col("host").eq(lit("a")), + ]), + TimeRangeExtraction::Extracted(range(1000, 2000)) + ); + + // Lower bound only / upper bound only. + assert_eq!( + extract(&[col("ts").gt_eq(ts_lit(1000))]), + TimeRangeExtraction::Extracted(TimestampRange::from_start(Timestamp::new_millisecond( + 1000 + ))) + ); + assert_eq!( + extract(&[col("ts").lt(ts_lit(2000))]), + TimeRangeExtraction::Extracted(TimestampRange::until_end( + Timestamp::new_millisecond(2000), + false + )) + ); + + // BETWEEN is inclusive on both ends. + assert_eq!( + extract(&[col("ts").between(ts_lit(1000), ts_lit(2000))]), + TimeRangeExtraction::Extracted(range(1000, 2001)) + ); + + // Equality pins a single point. + assert_eq!( + extract(&[col("ts").eq(ts_lit(1500))]), + TimeRangeExtraction::Extracted(TimestampRange::single(Timestamp::new_millisecond( + 1500 + ))) + ); + + // An unparsable side under AND only widens; the extraction stays safe. + assert_eq!( + extract(&[col("ts").gt_eq(ts_lit(1000)).and(col("ts").lt(col("t2")))]), + TimeRangeExtraction::Extracted(TimestampRange::from_start(Timestamp::new_millisecond( + 1000 + ))) + ); + + // Contradictory bounds collapse to the empty range, not an error. + let TimeRangeExtraction::Extracted(empty) = + extract(&[col("ts").gt_eq(ts_lit(2000)), col("ts").lt(ts_lit(1000))]) + else { + panic!("expected an extraction"); + }; + assert!(empty.is_empty()); + + // Shapes that could widen the satisfying set beyond what is extractable + // must be refused: OR with an unparsable side, NOT, non-literal bounds. + assert_eq!( + extract(&[col("ts").gt(ts_lit(1000)).or(col("host").eq(lit("a")))]), + TimeRangeExtraction::Unsupported + ); + assert_eq!( + extract(&[!col("ts").gt(ts_lit(1000))]), + TimeRangeExtraction::Unsupported + ); + assert_eq!( + extract(&[col("ts").gt_eq(col("t2"))]), + TimeRangeExtraction::Unsupported + ); + } + async fn gen_test_parquet_file(dir: &TempDir, cnt: usize) -> (String, Arc) { let path = dir .path() diff --git a/tests-integration/src/grpc.rs b/tests-integration/src/grpc.rs index 7c2148e130e..904cd0f1116 100644 --- a/tests-integration/src/grpc.rs +++ b/tests-integration/src/grpc.rs @@ -63,7 +63,10 @@ mod test { }; use client::OutputData; use common_base::Plugins; - use common_catalog::consts::MITO_ENGINE; + use common_catalog::consts::{ + DEFAULT_CATALOG_NAME, DEFAULT_PRIVATE_SCHEMA_NAME, MITO_ENGINE, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + }; use common_meta::rpc::router::region_distribution; use common_query::Output; use common_query::logical_plan::breakup_insert_plan; @@ -1342,4 +1345,171 @@ CREATE TABLE {table_name} ( +--------------------+---------------------------------------------------+"#; execute_sql_and_expect(&frontend, sql, expected).await; } + + fn declared_relationships_row_insert() -> Request { + use api::v1::value::ValueData; + use api::v1::{ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, Value}; + + let string_column = |name: &str, semantic: SemanticType| ColumnSchema { + column_name: name.to_string(), + datatype: ColumnDataType::String as i32, + semantic_type: semantic as i32, + ..Default::default() + }; + let string_value = |value: &str| Value { + value_data: Some(ValueData::StringValue(value.to_string())), + }; + let mut schema = vec![ColumnSchema { + column_name: "observed_at".to_string(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }]; + let mut values = vec![Value { + value_data: Some(ValueData::TimestampMillisecondValue(1000)), + }]; + for (tag, value) in [ + ("src_type", "service"), + ("src_id", "frontend"), + ("rel_type", "depends_on"), + ("dst_type", "service"), + ("dst_id", "users-db"), + ("provenance", "declared"), + ("scope", ""), + ("generation_id", ""), + ] { + schema.push(string_column(tag, SemanticType::Tag)); + values.push(string_value(value)); + } + + Request::RowInserts(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME.to_string(), + rows: Some(Rows { + schema, + rows: vec![Row { values }], + }), + }], + }) + } + + async fn assert_declared_relationships_table_is_canonical(instance: &Instance) { + let table = instance + .catalog_manager() + .table( + DEFAULT_CATALOG_NAME, + DEFAULT_PRIVATE_SCHEMA_NAME, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + None, + ) + .await + .unwrap() + .expect("declared-edge table must exist"); + let table_info = table.table_info(); + assert_eq!(table_info.meta.schema.column_schemas().len(), 20); + assert_eq!( + table_info + .meta + .schema + .timestamp_column() + .map(|c| c.name.as_str()), + Some("observed_at") + ); + let primary_keys: Vec<&str> = table_info + .meta + .primary_key_indices + .iter() + .map(|idx| table_info.meta.schema.column_schemas()[*idx].name.as_str()) + .collect(); + assert_eq!( + primary_keys, + [ + "src_type", + "src_id", + "rel_type", + "dst_type", + "dst_id", + "provenance", + "scope", + "generation_id", + ] + ); + } + + /// The declared-edge table's first gRPC write must create it with the + /// canonical definition (never the request's shape), entering below the + /// user-DDL guard that rejects generic CREATE against the reserved name. + #[tokio::test(flavor = "multi_thread")] + async fn test_grpc_first_write_creates_declared_relationships_table() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_grpc_declared_first_write") + .build() + .await; + let instance = standalone.fe_instance(); + + let ctx = Arc::new(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_PRIVATE_SCHEMA_NAME, + )); + let output = + GrpcQueryHandler::do_query(instance.as_ref(), declared_relationships_row_insert(), ctx) + .await + .unwrap(); + assert!(matches!(output.data, OutputData::AffectedRows(1))); + + assert_declared_relationships_table_is_canonical(instance).await; + } + + /// The system-defined table is created regardless of the auto-create hint: + /// its creation is a system action, not user auto-create. + #[tokio::test(flavor = "multi_thread")] + async fn test_grpc_declared_relationships_bypasses_auto_create_hint() { + use api::v1::value::ValueData; + use api::v1::{ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, Value}; + + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_grpc_declared_no_auto_create") + .build() + .await; + let instance = standalone.fe_instance(); + + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_PRIVATE_SCHEMA_NAME); + ctx.set_extension(table::requests::AUTO_CREATE_TABLE_KEY, "false"); + let ctx = Arc::new(ctx); + + let output = GrpcQueryHandler::do_query( + instance.as_ref(), + declared_relationships_row_insert(), + ctx.clone(), + ) + .await + .unwrap(); + assert!(matches!(output.data, OutputData::AffectedRows(1))); + assert_declared_relationships_table_is_canonical(instance).await; + + // An ordinary table stays subject to the hint. + let request = Request::RowInserts(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "ordinary_table".to_string(), + rows: Some(Rows { + schema: vec![ColumnSchema { + column_name: "ts".to_string(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }], + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::TimestampMillisecondValue(1000)), + }], + }], + }), + }], + }); + let err = GrpcQueryHandler::do_query(instance.as_ref(), request, ctx) + .await + .unwrap_err(); + let msg = common_error::ext::ErrorExt::output_msg(&err); + assert!(msg.contains("auto_create_table"), "unexpected error: {msg}"); + } } diff --git a/tests-integration/src/instance.rs b/tests-integration/src/instance.rs index 17435a3b7a0..547f64436f6 100644 --- a/tests-integration/src/instance.rs +++ b/tests-integration/src/instance.rs @@ -536,4 +536,108 @@ mod tests { unreachable!(); } } + + /// The entity-graph derivation must run as the caller: a source table the + /// caller cannot read contributes neither entities nor edges. + #[tokio::test(flavor = "multi_thread")] + async fn test_entity_graph_derivation_runs_as_caller() { + common_telemetry::init_default_ut_logging(); + + /// Denies the entity-graph derivation on `secret_traces` only. + struct DenySecretTraces; + + impl auth::PermissionChecker for DenySecretTraces { + fn check_permission( + &self, + _user_info: auth::UserInfoRef, + _req: auth::PermissionReq, + ) -> auth::error::Result { + Ok(auth::PermissionResp::Allow) + } + + fn check_permission_with_table_targets( + &self, + _user_info: auth::UserInfoRef, + req: auth::PermissionReq, + targets: auth::PermissionTableTargets, + ) -> auth::error::Result { + if matches!(req, auth::PermissionReq::Action(auth::SEMANTIC_GRAPH_QUERY)) + && let auth::PermissionTableTargets::Resolved(targets) = targets + && targets.iter().any(|target| target.table == "secret_traces") + { + return Ok(auth::PermissionResp::Reject); + } + Ok(auth::PermissionResp::Allow) + } + } + + let plugins = Plugins::new(); + plugins.insert::(Arc::new(DenySecretTraces)); + + let standalone = GreptimeDbStandaloneBuilder::new("test_entity_graph_runs_as_caller") + .with_plugin(plugins) + .build() + .await; + let instance = standalone.fe_instance().clone(); + + for table in ["open_traces", "secret_traces"] { + let create = format!( + r#"create table {table} ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) + ) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true')"# + ); + create_table(&instance, &create).await; + let insert = format!( + "insert into {table} values \ + (now(), 't0', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'client-{table}', 0), \ + (now(), 't0', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'server-{table}', 100)" + ); + let output = query(&instance, &insert).await; + let OutputData::AffectedRows(x) = output.data else { + unreachable!() + }; + assert_eq!(x, 2); + } + + let sql = "select src_id, dst_id from greptime_private.semantic_relationships \ + order by src_id"; + let output = query(&instance, sql).await; + let OutputData::Stream(s) = output.data else { + unreachable!() + }; + let batches = common_recordbatch::util::collect_batches(s).await.unwrap(); + let pretty_print = batches.pretty_print().unwrap(); + assert!( + pretty_print.contains("client-open_traces"), + "allowed edge missing:\n{pretty_print}" + ); + assert!( + !pretty_print.contains("secret_traces"), + "denied source leaked into edges:\n{pretty_print}" + ); + + let sql = "select entity_id from greptime_private.semantic_entities order by entity_id"; + let output = query(&instance, sql).await; + let OutputData::Stream(s) = output.data else { + unreachable!() + }; + let batches = common_recordbatch::util::collect_batches(s).await.unwrap(); + let pretty_print = batches.pretty_print().unwrap(); + assert!( + pretty_print.contains("client-open_traces"), + "allowed entity missing:\n{pretty_print}" + ); + assert!( + !pretty_print.contains("secret_traces"), + "denied source leaked into entities:\n{pretty_print}" + ); + } } diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index bda28755f1e..73f00f57c6f 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -153,3 +153,127 @@ drop table graph_traces; Affected Rows: 0 +-- Declared edges: the physical table is created by the system on the first +-- INSERT with its canonical schema. Re-asserting an edge stores a new revision; +-- reads keep only the latest one per edge key. +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, confidence) +values (now() - interval '10 minute', 'service', 'frontend', 'depends_on', 'service', 'users-db', 'declared', '', '', 0.5); + +Affected Rows: 1 + +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, confidence) +values (now() - interval '5 minute', 'service', 'frontend', 'depends_on', 'service', 'users-db', 'declared', '', '', 1.0); + +Affected Rows: 1 + +-- An old open-ended declaration stays valid until its row expires; an edge +-- retired in the past must not appear. +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, valid_until) +values + ('1970-01-01 00:00:01', 'service', 'legacy', 'depends_on', 'service', 'mainframe', 'declared', '', '', NULL), + ('1970-01-01 00:00:01', 'service', 'retired', 'depends_on', 'service', 'oldsys', 'declared', '', '', '2000-01-01 00:00:00'); + +Affected Rows: 2 + +select src_id, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by src_id; + ++----------+-----------+------------+------------+------------+ +| src_id | dst_id | rel_type | provenance | confidence | ++----------+-----------+------------+------------+------------+ +| frontend | users-db | depends_on | declared | 1.0 | +| legacy | mainframe | depends_on | declared | | ++----------+-----------+------------+------------+------------+ + +-- An explicit observed_at window replaces the default last hour; the emitted +-- timestamps of declared edges are synthesized inside the queried window. +select observed_at, window_start, fresh_until, src_id, dst_id +from greptime_private.semantic_relationships +where observed_at >= '2001-01-01 00:00:00' and observed_at < '2001-01-02 00:00:00' +order by src_id; + ++---------------------+---------------------+---------------------+--------+-----------+ +| observed_at | window_start | fresh_until | src_id | dst_id | ++---------------------+---------------------+---------------------+--------+-----------+ +| 2001-01-01T00:00:00 | 1970-01-01T00:00:01 | 2001-01-02T00:00:00 | legacy | mainframe | ++---------------------+---------------------+---------------------+--------+-----------+ + +-- A lower bound alone is fine (the upper bound defaults to now)... +select src_id, dst_id, provenance +from greptime_private.semantic_relationships +where observed_at >= now() - interval '30 minute' +order by src_id; + ++----------+-----------+------------+ +| src_id | dst_id | provenance | ++----------+-----------+------------+ +| frontend | users-db | declared | +| legacy | mainframe | declared | ++----------+-----------+------------+ + +-- ...but an upper bound alone would scan unbounded history: explicit error. +select src_id from greptime_private.semantic_relationships +where observed_at < '2001-01-02 00:00:00'; + +Error: 3001(EngineExecuteQuery), Invalid SQL, error: the observed_at filter has no lower bound; the graph cannot derive over unbounded history — add e.g. observed_at >= '2026-01-01 00:00:00' + +-- The declared-edge table's definition is system-owned: user CREATE/ALTER are +-- rejected and a physical table cannot be renamed into the name, while plain +-- DML stays allowed. +create table greptime_private.semantic_relationships_declared (ts timestamp time index); + +Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) + +alter table greptime_private.semantic_relationships_declared add column extra string; + +Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) + +create table greptime_private.declared_rename_probe (ts timestamp time index); + +Affected Rows: 0 + +alter table greptime_private.declared_rename_probe rename semantic_relationships_declared; + +Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) + +drop table greptime_private.declared_rename_probe; + +Affected Rows: 0 + +delete from greptime_private.semantic_relationships_declared; + +Affected Rows: 4 + +select src_id from greptime_private.semantic_relationships order by src_id; + +++ +++ + +-- DROP is allowed (nothing structural is lost: the next INSERT recreates the +-- canonical table) and cleans up after this test. +drop table greptime_private.semantic_relationships_declared; + +Affected Rows: 0 + +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id) +values (now(), 'service', 'reborn', 'depends_on', 'service', 'db', 'declared', '', ''); + +Affected Rows: 1 + +select src_id from greptime_private.semantic_relationships order by src_id; + ++--------+ +| src_id | ++--------+ +| reborn | ++--------+ + +drop table greptime_private.semantic_relationships_declared; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index 000401e1666..c5a93d6c6b7 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -92,3 +92,72 @@ from greptime_private.semantic_entities order by entity_id; drop table graph_traces; + +-- Declared edges: the physical table is created by the system on the first +-- INSERT with its canonical schema. Re-asserting an edge stores a new revision; +-- reads keep only the latest one per edge key. +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, confidence) +values (now() - interval '10 minute', 'service', 'frontend', 'depends_on', 'service', 'users-db', 'declared', '', '', 0.5); + +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, confidence) +values (now() - interval '5 minute', 'service', 'frontend', 'depends_on', 'service', 'users-db', 'declared', '', '', 1.0); + +-- An old open-ended declaration stays valid until its row expires; an edge +-- retired in the past must not appear. +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, valid_until) +values + ('1970-01-01 00:00:01', 'service', 'legacy', 'depends_on', 'service', 'mainframe', 'declared', '', '', NULL), + ('1970-01-01 00:00:01', 'service', 'retired', 'depends_on', 'service', 'oldsys', 'declared', '', '', '2000-01-01 00:00:00'); + +select src_id, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by src_id; + +-- An explicit observed_at window replaces the default last hour; the emitted +-- timestamps of declared edges are synthesized inside the queried window. +select observed_at, window_start, fresh_until, src_id, dst_id +from greptime_private.semantic_relationships +where observed_at >= '2001-01-01 00:00:00' and observed_at < '2001-01-02 00:00:00' +order by src_id; + +-- A lower bound alone is fine (the upper bound defaults to now)... +select src_id, dst_id, provenance +from greptime_private.semantic_relationships +where observed_at >= now() - interval '30 minute' +order by src_id; + +-- ...but an upper bound alone would scan unbounded history: explicit error. +select src_id from greptime_private.semantic_relationships +where observed_at < '2001-01-02 00:00:00'; + +-- The declared-edge table's definition is system-owned: user CREATE/ALTER are +-- rejected and a physical table cannot be renamed into the name, while plain +-- DML stays allowed. +create table greptime_private.semantic_relationships_declared (ts timestamp time index); + +alter table greptime_private.semantic_relationships_declared add column extra string; + +create table greptime_private.declared_rename_probe (ts timestamp time index); + +alter table greptime_private.declared_rename_probe rename semantic_relationships_declared; + +drop table greptime_private.declared_rename_probe; + +delete from greptime_private.semantic_relationships_declared; + +select src_id from greptime_private.semantic_relationships order by src_id; + +-- DROP is allowed (nothing structural is lost: the next INSERT recreates the +-- canonical table) and cleans up after this test. +drop table greptime_private.semantic_relationships_declared; + +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id) +values (now(), 'service', 'reborn', 'depends_on', 'service', 'db', 'declared', '', ''); + +select src_id from greptime_private.semantic_relationships order by src_id; + +drop table greptime_private.semantic_relationships_declared; From 78084a9d4485ee61ff03e071f6428fc57c906332 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 10 Aug 2026 20:19:27 +0800 Subject: [PATCH 035/126] feat: add admin function to discard unflushed data (#8768) * feat: add admin function to discard unflushed data Signed-off-by: evenyag * test: cover discarding unflushed data by table Signed-off-by: evenyag * chore: fix license header Signed-off-by: evenyag * fix: reject discarding logical metric table data Signed-off-by: evenyag * refactor: defer table name formatting in error paths Signed-off-by: evenyag * chore(deps): update greptime-proto revision Signed-off-by: evenyag * refactor: rename discard unflushed admin function Signed-off-by: evenyag --------- Signed-off-by: evenyag --- AGENTS.md | 21 +- Cargo.lock | 2 +- Cargo.toml | 2 +- src/common/function/src/admin.rs | 4 +- .../src/admin/discard_unflushed_data.rs | 232 ++++++++++++++++++ src/common/function/src/handlers.rs | 14 ++ src/common/function/src/state.rs | 17 ++ src/operator/src/request.rs | 96 +++++++- src/operator/src/table.rs | 25 ++ src/store-api/src/region_request.rs | 39 +++ tests-integration/src/tests/instance_test.rs | 75 ++++++ .../admin/discard_unflushed_data.result | 133 ++++++++++ .../function/admin/discard_unflushed_data.sql | 62 +++++ 13 files changed, 709 insertions(+), 13 deletions(-) create mode 100644 src/common/function/src/admin/discard_unflushed_data.rs create mode 100644 tests/cases/standalone/common/function/admin/discard_unflushed_data.result create mode 100644 tests/cases/standalone/common/function/admin/discard_unflushed_data.sql diff --git a/AGENTS.md b/AGENTS.md index c99856c2941..70f2ae20b41 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -114,22 +114,25 @@ blast radius requires it. ## Before opening a PR -1. `make fmt` -2. `make clippy` -3. `make test` -4. `make check-udeps` (run `make fix-udeps` if it reports unused dependencies). -5. If you added or changed a public configuration option, update the applicable +1. If you added a `.rs`, `.py`, or `.ts` file, apply and verify its license + header with `hawkeye format` followed by `hawkeye check`. Use the inception + year from `licenserc.toml`, not the current year. +2. `make fmt` +3. `make clippy` +4. `make test` +5. `make check-udeps` (run `make fix-udeps` if it reports unused dependencies). +6. If you added or changed a public configuration option, update the applicable example TOMLs, configuration-loading and serialized-config snapshot tests, and related user-facing documentation. Run `make config-docs` (needs Docker) and commit the regenerated `config/config.md`. -6. If you changed a persisted or wire format, add a compatibility test case (see +7. If you changed a persisted or wire format, add a compatibility test case (see `.agents/architecture-invariants.md`). -7. If you added or gated an enterprise-only file, give it the enterprise license +8. If you added or gated an enterprise-only file, give it the enterprise license header, list it in `licenserc-enterprise.toml` (`includes`) and `licenserc.toml` (`excludes`), and run `make check-enterprise-license`. -8. Use a conventional-commit title, sign off commits (`git commit -s`), and sign +9. Use a conventional-commit title, sign off commits (`git commit -s`), and sign the CLA. -9. When creating or updating a pull request, follow +10. When creating or updating a pull request, follow [`.github/pull_request_template.md`](.github/pull_request_template.md): include the CLA statement, fill the change-intention section with enough detail, and update checklist items accurately. diff --git a/Cargo.lock b/Cargo.lock index e2210b80bfe..d6385dd7657 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6032,7 +6032,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=032510ded061277b7d1bbb29d4046abb5b1cbf4b#032510ded061277b7d1bbb29d4046abb5b1cbf4b" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=e7be20ff855b7522efe235383d82238d56e676c3#e7be20ff855b7522efe235383d82238d56e676c3" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", diff --git a/Cargo.toml b/Cargo.toml index 502e652ffc8..64e6dba2f7f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -159,7 +159,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "032510ded061277b7d1bbb29d4046abb5b1cbf4b" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "e7be20ff855b7522efe235383d82238d56e676c3" } hex = "0.4" http = "1" humantime = "2.1" diff --git a/src/common/function/src/admin.rs b/src/common/function/src/admin.rs index 73bed86f7ad..ff8172ebb5a 100644 --- a/src/common/function/src/admin.rs +++ b/src/common/function/src/admin.rs @@ -13,6 +13,7 @@ // limitations under the License. mod build_index_table; +mod discard_unflushed_data; mod flush_compact_region; mod flush_compact_table; mod gc; @@ -34,6 +35,7 @@ use reconcile_database::ReconcileDatabaseFunction; use reconcile_table::ReconcileTableFunction; use crate::admin::build_index_table::BuildIndexFunction; +use crate::admin::discard_unflushed_data::DiscardUnflushedDataFunction; use crate::flush_flow::FlushFlowFunction; use crate::function_registry::FunctionRegistry; @@ -58,8 +60,8 @@ impl AdminFunction { } /// Register functions that must only be resolved by an ADMIN statement. - #[cfg_attr(not(feature = "enterprise"), allow(unused_variables))] pub fn register_admin_only(registry: &FunctionRegistry) { + registry.register(DiscardUnflushedDataFunction::factory()); #[cfg(feature = "enterprise")] registry.register(PurgeTableFunction::factory()); } diff --git a/src/common/function/src/admin/discard_unflushed_data.rs b/src/common/function/src/admin/discard_unflushed_data.rs new file mode 100644 index 00000000000..b9f46f62f96 --- /dev/null +++ b/src/common/function/src/admin/discard_unflushed_data.rs @@ -0,0 +1,232 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use arrow::datatypes::DataType as ArrowDataType; +use common_error::ext::BoxedError; +use common_macro::admin_fn; +use common_query::error::{ + InvalidFuncArgsSnafu, MissingTableMutationHandlerSnafu, Result, TableMutationSnafu, + UnsupportedInputDataTypeSnafu, +}; +use datafusion_expr::{Signature, TypeSignature, Volatility}; +use datatypes::data_type::DataType; +use datatypes::prelude::*; +use session::context::QueryContextRef; +use session::table_name::table_name_to_full_name; +use snafu::{ResultExt, ensure}; +use store_api::storage::RegionId; +use table::table_name::TableName; + +use crate::handlers::TableMutationHandlerRef; +use crate::helper::cast_u64; + +/// Discards all unflushed data from a region. +#[admin_fn( + name = DiscardUnflushedDataFunction, + display_name = discard_unflushed, + sig_fn = signature, + ret = uint64, + single_row +)] +pub(crate) async fn discard_unflushed_data( + table_mutation_handler: &TableMutationHandlerRef, + query_ctx: &QueryContextRef, + params: &[ValueRef<'_>], +) -> Result { + ensure!( + params.len() == 1, + InvalidFuncArgsSnafu { + err_msg: format!( + "The length of the args is not correct, expect 1, have: {}", + params.len() + ), + } + ); + + let affected_rows = match params[0] { + ValueRef::String(table_name) => { + let (catalog_name, schema_name, table_name) = + table_name_to_full_name(table_name, query_ctx) + .map_err(BoxedError::new) + .context(TableMutationSnafu)?; + table_mutation_handler + .discard_unflushed_data_by_table( + TableName::new(catalog_name, schema_name, table_name), + query_ctx.clone(), + ) + .await? + } + _ => { + let Some(region_id) = cast_u64(¶ms[0])? else { + return UnsupportedInputDataTypeSnafu { + function: "discard_unflushed", + datatypes: params + .iter() + .map(|value| value.data_type()) + .collect::>(), + } + .fail(); + }; + table_mutation_handler + .discard_unflushed_data(RegionId::from_u64(region_id), query_ctx.clone()) + .await? + } + }; + + Ok(Value::from(affected_rows as u64)) +} + +fn signature() -> Signature { + Signature::one_of( + vec![ + TypeSignature::Uniform( + 1, + ConcreteDataType::numerics() + .into_iter() + .map(|data_type| data_type.as_arrow_type()) + .collect(), + ), + TypeSignature::Exact(vec![ArrowDataType::Utf8]), + ], + Volatility::Immutable, + ) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{StringArray, UInt64Array}; + use arrow::datatypes::{DataType, Field}; + use datafusion_expr::{ColumnarValue, TypeSignature}; + + use super::*; + use crate::function::FunctionContext; + use crate::function_factory::ScalarFunctionFactory; + use crate::function_registry::{FUNCTION_REGISTRY, get_admin_function}; + + #[test] + fn test_discard_unflushed_data_is_admin_only() { + assert!(get_admin_function("discard_unflushed").is_some()); + assert!( + FUNCTION_REGISTRY + .get_function("discard_unflushed") + .is_none() + ); + } + + #[test] + fn test_discard_unflushed_data_signature() { + let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into(); + let function = factory.provide(FunctionContext::mock()); + + assert_eq!("discard_unflushed", function.name()); + assert_eq!(DataType::UInt64, function.return_type(&[]).unwrap()); + assert!(matches!( + function.signature(), + Signature { + type_signature: TypeSignature::OneOf(valid_types), + volatility: Volatility::Immutable, + .. + } if valid_types == &vec![ + TypeSignature::Uniform( + 1, + ConcreteDataType::numerics() + .into_iter() + .map(|data_type| { + use datatypes::data_type::DataType; + data_type.as_arrow_type() + }) + .collect::>(), + ), + TypeSignature::Exact(vec![DataType::Utf8]), + ] + )); + } + + #[tokio::test] + async fn test_discard_unflushed_data() { + let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into(); + let function = factory.provide(FunctionContext::mock()); + let args = datafusion::logical_expr::ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new(UInt64Array::from(vec![99])))], + arg_fields: vec![Arc::new(Field::new("arg_0", DataType::UInt64, false))], + return_field: Arc::new(Field::new("result", DataType::UInt64, false)), + number_rows: 1, + config_options: Arc::new(datafusion_common::config::ConfigOptions::default()), + }; + + let result = function + .as_async() + .unwrap() + .invoke_async_with_args(args) + .await + .unwrap(); + let ColumnarValue::Array(array) = result else { + panic!("expected array output"); + }; + let array = array.as_any().downcast_ref::().unwrap(); + assert_eq!(42, array.value(0)); + } + + #[tokio::test] + async fn test_discard_unflushed_data_by_table() { + let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into(); + let function = factory.provide(FunctionContext::mock()); + let args = datafusion::logical_expr::ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![ + "my_table", + ])))], + arg_fields: vec![Arc::new(Field::new("arg_0", DataType::Utf8, false))], + return_field: Arc::new(Field::new("result", DataType::UInt64, false)), + number_rows: 1, + config_options: Arc::new(datafusion_common::config::ConfigOptions::default()), + }; + + let result = function + .as_async() + .unwrap() + .invoke_async_with_args(args) + .await + .unwrap(); + let ColumnarValue::Array(array) = result else { + panic!("expected array output"); + }; + let array = array.as_any().downcast_ref::().unwrap(); + assert_eq!(42, array.value(0)); + } + + #[tokio::test] + async fn test_discard_unflushed_data_rejects_multiple_rows() { + let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into(); + let function = factory.provide(FunctionContext::mock()); + let args = datafusion::logical_expr::ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new(UInt64Array::from(vec![ + 1, 2, + ])))], + arg_fields: vec![Arc::new(Field::new("arg_0", DataType::UInt64, false))], + return_field: Arc::new(Field::new("result", DataType::UInt64, false)), + number_rows: 2, + config_options: Arc::new(datafusion_common::config::ConfigOptions::default()), + }; + + let error = function + .as_async() + .unwrap() + .invoke_async_with_args(args) + .await + .unwrap_err(); + assert!(error.to_string().contains("received 2")); + } +} diff --git a/src/common/function/src/handlers.rs b/src/common/function/src/handlers.rs index 47de2824b8f..24ce115a0d6 100644 --- a/src/common/function/src/handlers.rs +++ b/src/common/function/src/handlers.rs @@ -69,6 +69,20 @@ pub trait TableMutationHandler: Send + Sync { region_id: RegionId, ctx: QueryContextRef, ) -> Result; + + /// Discard all unflushed data from a table region. + async fn discard_unflushed_data( + &self, + region_id: RegionId, + ctx: QueryContextRef, + ) -> Result; + + /// Discard all unflushed data from all regions of a table. + async fn discard_unflushed_data_by_table( + &self, + table_name: TableName, + ctx: QueryContextRef, + ) -> Result; } /// A trait for handling procedure service requests in `QueryEngine`. diff --git a/src/common/function/src/state.rs b/src/common/function/src/state.rs index 726785f0741..b1802b62578 100644 --- a/src/common/function/src/state.rs +++ b/src/common/function/src/state.rs @@ -48,6 +48,7 @@ impl FunctionState { BuildIndexTableRequest, CompactTableRequest, DeleteRequest, FlushTableRequest, InsertRequest, }; + use table::table_name::TableName; use crate::handlers::{FlowServiceHandler, ProcedureServiceHandler, TableMutationHandler}; struct MockProcedureServiceHandler; @@ -171,6 +172,22 @@ impl FunctionState { ) -> Result { Ok(ROWS) } + + async fn discard_unflushed_data( + &self, + _region_id: RegionId, + _ctx: QueryContextRef, + ) -> Result { + Ok(ROWS) + } + + async fn discard_unflushed_data_by_table( + &self, + _table_name: TableName, + _ctx: QueryContextRef, + ) -> Result { + Ok(ROWS) + } } #[async_trait] diff --git a/src/operator/src/request.rs b/src/operator/src/request.rs index aea395d8eed..b46a6b657a6 100644 --- a/src/operator/src/request.rs +++ b/src/operator/src/request.rs @@ -18,9 +18,11 @@ use api::helper::to_pb_time_unit; use api::v1::region::region_request::Body as RegionRequestBody; use api::v1::region::{ BuildIndexRequest, CompactRequest, CompactionTimeRange, FlushRequest, RegionRequestHeader, + TruncateRequest, Unflushed, truncate_request, }; use catalog::CatalogManagerRef; use common_catalog::build_db_string; +use common_catalog::consts::METRIC_ENGINE; use common_meta::node_manager::{AffectedRows, NodeManagerRef}; use common_meta::peer::Peer; use common_telemetry::tracing_context::TracingContext; @@ -34,10 +36,12 @@ use session::context::QueryContextRef; use snafu::prelude::*; use store_api::storage::RegionId; use table::requests::{BuildIndexTableRequest, CompactTableRequest, FlushTableRequest}; +use table::table_name::TableName; use crate::error::{ CatalogSnafu, FindRegionLeaderSnafu, FindTablePartitionRuleSnafu, JoinTaskSnafu, - RequestRegionSnafu, Result, TableNotFoundSnafu, UnsupportedRegionRequestSnafu, + NotSupportedSnafu, RequestRegionSnafu, Result, TableNotFoundSnafu, + UnsupportedRegionRequestSnafu, }; use crate::region_req_factory::RegionRequestFactory; @@ -206,6 +210,76 @@ impl Requester { info!("Handle region manual compaction request: {region_id}"); self.do_request(vec![request], None, &ctx).await } + + /// Discard all unflushed data from the region. + pub async fn handle_discard_unflushed_data( + &self, + region_id: RegionId, + ctx: QueryContextRef, + ) -> Result { + let request = RegionRequestBody::Truncate(TruncateRequest { + region_id: region_id.into(), + kind: Some(truncate_request::Kind::Unflushed(Unflushed {})), + }); + + info!("Handle region discard unflushed data request: {region_id}"); + self.do_request(vec![request], None, &ctx).await + } + + /// Discard all unflushed data from all regions of the table. + pub async fn handle_discard_unflushed_data_by_table( + &self, + table_name: TableName, + ctx: QueryContextRef, + ) -> Result { + let table = self + .catalog_manager + .table( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + None, + ) + .await + .context(CatalogSnafu)?; + let table = table.with_context(|| TableNotFoundSnafu { + table_name: table_name.to_string(), + })?; + let table_info = table.table_info(); + ensure_discard_unflushed_supported( + &table_info.meta.engine, + table_info.is_physical_table(), + )?; + + let partitions = &self + .partition_manager + .find_physical_partition_info(table_info.ident.table_id) + .await + .with_context(|_| FindTablePartitionRuleSnafu { + table_name: table_name.to_string(), + })? + .partitions; + let requests = partitions + .iter() + .map(|partition| { + RegionRequestBody::Truncate(TruncateRequest { + region_id: partition.id.into(), + kind: Some(truncate_request::Kind::Unflushed(Unflushed {})), + }) + }) + .collect(); + + info!("Handle table discard unflushed data request: {table_name}"); + self.do_request( + requests, + Some(build_db_string( + &table_name.catalog_name, + &table_name.schema_name, + )), + &ctx, + ) + .await + } } fn to_pb_compaction_time_range(range: TimestampRange) -> Result { @@ -291,6 +365,7 @@ impl Requester { RegionRequestBody::Flush(req) => req.region_id, RegionRequestBody::Compact(req) => req.region_id, RegionRequestBody::BuildIndex(req) => req.region_id, + RegionRequestBody::Truncate(req) => req.region_id, _ => { error!("Unsupported region request: {:?}", req); return UnsupportedRegionRequestSnafu {}.fail(); @@ -329,6 +404,16 @@ impl Requester { } } +fn ensure_discard_unflushed_supported(engine: &str, is_physical_table: bool) -> Result<()> { + ensure!( + engine != METRIC_ENGINE || is_physical_table, + NotSupportedSnafu { + feat: "discarding unflushed data from a Metric Engine logical table" + } + ); + Ok(()) +} + #[cfg(test)] mod tests { use api::v1::TimeUnit; @@ -351,6 +436,15 @@ mod tests { assert_eq!(TimeUnit::Second as i32, pb_range.time_unit); } + #[test] + fn test_discard_unflushed_rejects_metric_logical_table() { + let error = ensure_discard_unflushed_supported(METRIC_ENGINE, false).unwrap_err(); + assert!(matches!(error, crate::error::Error::NotSupported { .. })); + + ensure_discard_unflushed_supported(METRIC_ENGINE, true).unwrap(); + ensure_discard_unflushed_supported("mito", false).unwrap(); + } + #[test] fn test_to_pb_compaction_time_range() { let range = TimestampRange::new( diff --git a/src/operator/src/table.rs b/src/operator/src/table.rs index 13ed57200c2..117351280b9 100644 --- a/src/operator/src/table.rs +++ b/src/operator/src/table.rs @@ -26,6 +26,7 @@ use table::requests::{ BuildIndexTableRequest, CompactTableRequest, DeleteRequest as TableDeleteRequest, FlushTableRequest, InsertRequest as TableInsertRequest, }; +use table::table_name::TableName; use crate::delete::DeleterRef; use crate::insert::InserterRef; @@ -132,4 +133,28 @@ impl TableMutationHandler for TableMutationOperator { .map_err(BoxedError::new) .context(query_error::TableMutationSnafu) } + + async fn discard_unflushed_data( + &self, + region_id: RegionId, + ctx: QueryContextRef, + ) -> QueryResult { + self.requester + .handle_discard_unflushed_data(region_id, ctx) + .await + .map_err(BoxedError::new) + .context(query_error::TableMutationSnafu) + } + + async fn discard_unflushed_data_by_table( + &self, + table_name: TableName, + ctx: QueryContextRef, + ) -> QueryResult { + self.requester + .handle_discard_unflushed_data_by_table(table_name, ctx) + .await + .map_err(BoxedError::new) + .context(query_error::TableMutationSnafu) + } } diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 490e6c4d6e6..3bb2bb916c1 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -458,6 +458,10 @@ fn make_region_truncate(truncate: TruncateRequest) -> Result Ok(vec![( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::Unflushed), + )]), } } @@ -1811,6 +1815,41 @@ mod tests { ); } + #[test] + fn test_make_region_truncate_unflushed() { + let region_id = RegionId::new(42, 3); + let requests = + RegionRequest::try_from_request_body(region_request::Body::Truncate(TruncateRequest { + region_id: region_id.as_u64(), + kind: Some(truncate_request::Kind::Unflushed( + api::v1::region::Unflushed {}, + )), + })) + .unwrap(); + + assert_eq!(region_id, requests[0].0); + assert!(matches!( + requests[0].1, + RegionRequest::Truncate(RegionTruncateRequest::Unflushed) + )); + } + + #[test] + fn test_make_region_truncate_requires_kind() { + let error = + RegionRequest::try_from_request_body(region_request::Body::Truncate(TruncateRequest { + region_id: RegionId::new(42, 3).as_u64(), + kind: None, + })) + .unwrap_err(); + + assert!( + error + .to_string() + .contains("missing kind in TruncateRequest") + ); + } + #[test] fn test_from_proto_location() { let proto_location = v1::AddColumnLocation { diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 9c64e841254..f21b459c2a7 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -90,6 +90,81 @@ async fn test_create_database_and_insert_query(instance: Arc) } } +#[apply(both_instances_cases)] +async fn test_admin_discard_unflushed_data(instance: Arc) { + let instance = instance.frontend(); + + execute_sql( + &instance, + r#"CREATE TABLE discard_unflushed_data_test ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX + ) ENGINE = mito"#, + ) + .await; + execute_sql( + &instance, + "INSERT INTO discard_unflushed_data_test VALUES ('persisted', 1, 1)", + ) + .await; + execute_sql( + &instance, + "ADMIN FLUSH_TABLE('discard_unflushed_data_test')", + ) + .await; + execute_sql( + &instance, + "INSERT INTO discard_unflushed_data_test VALUES ('unflushed', 2, 2)", + ) + .await; + + let region_output = execute_sql( + &instance, + "SELECT greptime_partition_id FROM information_schema.partitions \ + WHERE table_name = 'discard_unflushed_data_test' LIMIT 1", + ) + .await; + let OutputData::Stream(stream) = region_output.data else { + panic!("expected region id stream"); + }; + let batches = util::collect(stream).await.unwrap(); + let region_ids = batches[0].column(0); + let region_ids = region_ids.as_any().downcast_ref::().unwrap(); + let region_id = region_ids.value(0); + + let admin_sql = format!("ADMIN discard_unflushed({region_id})"); + for _ in 0..2 { + let output = execute_sql(&instance, &admin_sql).await; + let OutputData::RecordBatches(batches) = output.data else { + panic!("expected ADMIN result batches"); + }; + let result = batches.iter().next().unwrap().column(0); + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(0, result.value(0)); + } + + let output = execute_sql( + &instance, + "SELECT host, val FROM discard_unflushed_data_test ORDER BY host", + ) + .await; + assert_eq!( + "+-----------+-----+\n\ + | host | val |\n\ + +-----------+-----+\n\ + | persisted | 1.0 |\n\ + +-----------+-----+", + output.data.pretty_print().await + ); + + assert!( + try_execute_sql(&instance, &format!("SELECT discard_unflushed({region_id})"),) + .await + .is_err() + ); +} + #[tokio::test(flavor = "multi_thread")] async fn test_distributed_scalar_latest_with_query_parallelism_below_regions() { common_telemetry::init_default_ut_logging(); diff --git a/tests/cases/standalone/common/function/admin/discard_unflushed_data.result b/tests/cases/standalone/common/function/admin/discard_unflushed_data.result new file mode 100644 index 00000000000..9723baa8638 --- /dev/null +++ b/tests/cases/standalone/common/function/admin/discard_unflushed_data.result @@ -0,0 +1,133 @@ +CREATE TABLE discard_unflushed_data_test ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = mito; + +Affected Rows: 0 + +INSERT INTO discard_unflushed_data_test VALUES ('persisted', 1, 1); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('discard_unflushed_data_test'); + ++--------------------------------------------------+ +| ADMIN FLUSH_TABLE('discard_unflushed_data_test') | ++--------------------------------------------------+ +| 0 | ++--------------------------------------------------+ + +INSERT INTO discard_unflushed_data_test VALUES ('unflushed', 2, 2); + +Affected Rows: 1 + +ADMIN discard_unflushed('discard_unflushed_data_test'); + ++--------------------------------------------------------+ +| ADMIN discard_unflushed('discard_unflushed_data_test') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +-- Repeating the operation is idempotent. +ADMIN discard_unflushed('discard_unflushed_data_test'); + ++--------------------------------------------------------+ +| ADMIN discard_unflushed('discard_unflushed_data_test') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +SELECT host, val FROM discard_unflushed_data_test ORDER BY host; + ++-----------+-----+ +| host | val | ++-----------+-----+ +| persisted | 1.0 | ++-----------+-----+ + +-- The function must only be available through ADMIN. +-- SQLNESS REPLACE \nDid\syou\smean.* +SELECT discard_unflushed(0); + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Invalid function 'discard_unflushed'. + +DROP TABLE discard_unflushed_data_test; + +Affected Rows: 0 + +CREATE TABLE discard_unflushed_data_test ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = mito; + +Affected Rows: 0 + +INSERT INTO discard_unflushed_data_test VALUES ('unflushed', 1, 1); + +Affected Rows: 1 + +-- The table has no persisted data, so discarding leaves it empty. +ADMIN discard_unflushed('discard_unflushed_data_test'); + ++--------------------------------------------------------+ +| ADMIN discard_unflushed('discard_unflushed_data_test') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +SELECT COUNT(*) FROM discard_unflushed_data_test; + ++----------+ +| count(*) | ++----------+ +| 0 | ++----------+ + +DROP TABLE discard_unflushed_data_test; + +Affected Rows: 0 + +CREATE TABLE discard_unflushed_data_physical ( + ts TIMESTAMP TIME INDEX, + val DOUBLE +) ENGINE = metric WITH ("physical_metric_table" = ""); + +Affected Rows: 0 + +CREATE TABLE discard_unflushed_data_logical ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = metric WITH ("on_physical_table" = "discard_unflushed_data_physical"); + +Affected Rows: 0 + +INSERT INTO discard_unflushed_data_logical VALUES ('unflushed', 1, 1); + +Affected Rows: 1 + +-- Logical Metric Engine tables share their physical regions with other logical tables. +-- Discarding by logical table name must not truncate those shared regions. +ADMIN discard_unflushed('discard_unflushed_data_logical'); + +Error: 1002(Unexpected), Failed to execute admin function discard_unflushed: Execution error: Not supported: discarding unflushed data from a Metric Engine logical table + +SELECT host, val FROM discard_unflushed_data_logical; + ++-----------+-----+ +| host | val | ++-----------+-----+ +| unflushed | 1.0 | ++-----------+-----+ + +DROP TABLE discard_unflushed_data_logical; + +Affected Rows: 0 + +DROP TABLE discard_unflushed_data_physical; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/admin/discard_unflushed_data.sql b/tests/cases/standalone/common/function/admin/discard_unflushed_data.sql new file mode 100644 index 00000000000..9183489495f --- /dev/null +++ b/tests/cases/standalone/common/function/admin/discard_unflushed_data.sql @@ -0,0 +1,62 @@ +CREATE TABLE discard_unflushed_data_test ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = mito; + +INSERT INTO discard_unflushed_data_test VALUES ('persisted', 1, 1); + +ADMIN FLUSH_TABLE('discard_unflushed_data_test'); + +INSERT INTO discard_unflushed_data_test VALUES ('unflushed', 2, 2); + +ADMIN discard_unflushed('discard_unflushed_data_test'); + +-- Repeating the operation is idempotent. +ADMIN discard_unflushed('discard_unflushed_data_test'); + +SELECT host, val FROM discard_unflushed_data_test ORDER BY host; + +-- The function must only be available through ADMIN. +-- SQLNESS REPLACE \nDid\syou\smean.* +SELECT discard_unflushed(0); + +DROP TABLE discard_unflushed_data_test; + +CREATE TABLE discard_unflushed_data_test ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = mito; + +INSERT INTO discard_unflushed_data_test VALUES ('unflushed', 1, 1); + +-- The table has no persisted data, so discarding leaves it empty. +ADMIN discard_unflushed('discard_unflushed_data_test'); + +SELECT COUNT(*) FROM discard_unflushed_data_test; + +DROP TABLE discard_unflushed_data_test; + +CREATE TABLE discard_unflushed_data_physical ( + ts TIMESTAMP TIME INDEX, + val DOUBLE +) ENGINE = metric WITH ("physical_metric_table" = ""); + +CREATE TABLE discard_unflushed_data_logical ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = metric WITH ("on_physical_table" = "discard_unflushed_data_physical"); + +INSERT INTO discard_unflushed_data_logical VALUES ('unflushed', 1, 1); + +-- Logical Metric Engine tables share their physical regions with other logical tables. +-- Discarding by logical table name must not truncate those shared regions. +ADMIN discard_unflushed('discard_unflushed_data_logical'); + +SELECT host, val FROM discard_unflushed_data_logical; + +DROP TABLE discard_unflushed_data_logical; + +DROP TABLE discard_unflushed_data_physical; From f88226c3ff64f3cf410950b483f43f170fb07d0e Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:27:47 +0800 Subject: [PATCH 036/126] ci: update dev-builder image tag (#8831) Signed-off-by: greptimedb-ci Co-authored-by: greptimedb-ci --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 5bdb63d3520..0f280b7b9be 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,7 @@ CARGO_BUILD_OPTS := --locked IMAGE_REGISTRY ?= docker.io IMAGE_NAMESPACE ?= greptime IMAGE_TAG ?= latest -DEV_BUILDER_IMAGE_TAG ?= 2026-03-21-9c9d9e9e-20260331090344 +DEV_BUILDER_IMAGE_TAG ?= 2026-03-21-7fa6f5f9-20260810091725 DEV_BUILDER_RISCV64_IMAGE_TAG ?= $(DEV_BUILDER_IMAGE_TAG) DEV_BUILDER_RISCV64_BASE_IMAGE ?= ubuntu:22.04 RISCV64_TARGET ?= riscv64gc-unknown-linux-gnu From e2ab11763a945504988d4e246181b68a1db07197 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 10 Aug 2026 21:14:56 +0800 Subject: [PATCH 037/126] fix(ci): identify team members by repository permission (#8822) * fix(ci): identify team members by repository permission `author_association` is computed from what the caller can see, so GITHUB_TOKEN reports a private organization member as CONTRIBUTOR. Only 5 of GreptimeTeam's members have public membership, so the open-pull-request check skipped almost everyone it was written for. Use the repository permission of the author instead, which is viewer-independent. On error, apply the limit rather than skipping, so a token that cannot read permissions cannot silently disable the check again. Signed-off-by: Dennis Zhuang * fix(ci): do not log repository permission levels Job logs are public. Resolving the author's permission is fine; printing the level is not. Signed-off-by: Dennis Zhuang * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --------- Signed-off-by: Dennis Zhuang Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .github/scripts/pr-open-limit.js | 37 ++++++++++++++++++++++------- .github/workflows/pr-open-limit.yml | 1 - 2 files changed, 29 insertions(+), 9 deletions(-) diff --git a/.github/scripts/pr-open-limit.js b/.github/scripts/pr-open-limit.js index 1ba908b7450..dd1650b0d64 100644 --- a/.github/scripts/pr-open-limit.js +++ b/.github/scripts/pr-open-limit.js @@ -9,9 +9,8 @@ import { appendFileSync } from "node:fs"; const MARKER = ""; // GITHUB_TOKEN can only edit comments authored by the Actions bot itself. const COMMENT_AUTHOR = "github-actions[bot]"; -// Associations that identify an organization member. Outside collaborators and -// community contributors are out of scope. -const MEMBER_ASSOCIATIONS = ["OWNER", "MEMBER"]; +// Repository permissions that mark someone as part of the team. +const TEAM_PERMISSIONS = ["admin", "maintain", "write"]; // Keep the comment readable for authors who are far over the limit. const MAX_LISTED_PRS = 10; @@ -22,6 +21,29 @@ function summary(text) { console.log(text); } +// `author_association` cannot be used here: it is computed from what the caller +// is allowed to see, so GITHUB_TOKEN reports a private organization member as +// CONTRIBUTOR. Repository permission is viewer-independent. +// +// On error this returns true rather than false. A token that cannot read +// permissions would otherwise make the whole check silently pass everyone. +async function isTeamMember(octokit, owner, repo, username) { + try { + const { data } = await octokit.repos.getCollaboratorPermissionLevel({ + owner, + repo, + username, + }); + // Never log the level itself: job logs are public. + return TEAM_PERMISSIONS.includes(data.permission); + } catch (error) { + console.log( + `Cannot read repository permission for \`${username}\` (HTTP ${error.status}); applying the limit anyway.` + ); + return true; + } +} + function buildComment(author, total, limit, openPrs) { const listed = openPrs .slice(0, MAX_LISTED_PRS) @@ -69,7 +91,6 @@ async function upsertComment(octokit, params, body) { const [owner, repo] = process.env.GITHUB_REPOSITORY.split("/"); const prNumber = Number(process.env.PR_NUMBER); const author = process.env.PR_AUTHOR; - const association = process.env.PR_AUTHOR_ASSOCIATION; const rawLimit = Number(process.env.MAX_OPEN_PRS); const limit = Number.isInteger(rawLimit) && rawLimit >= 0 ? rawLimit : 5; @@ -78,13 +99,13 @@ async function upsertComment(octokit, params, body) { return; } - if (!MEMBER_ASSOCIATIONS.includes(association)) { - summary(`Skipping \`${author}\`: association is \`${association}\`, not an org member.`); + const octokit = new Octokit({ auth: process.env.GITHUB_TOKEN }); + + if (!(await isTeamMember(octokit, owner, repo, author))) { + summary(`Skipping \`${author}\`: not a team member.`); return; } - const octokit = new Octokit({ auth: process.env.GITHUB_TOKEN }); - const allOpen = await octokit.paginate(octokit.pulls.list, { owner, repo, diff --git a/.github/workflows/pr-open-limit.yml b/.github/workflows/pr-open-limit.yml index d22929911a8..de5b723cb25 100644 --- a/.github/workflows/pr-open-limit.yml +++ b/.github/workflows/pr-open-limit.yml @@ -44,5 +44,4 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} PR_NUMBER: ${{ github.event.pull_request.number }} PR_AUTHOR: ${{ github.event.pull_request.user.login }} - PR_AUTHOR_ASSOCIATION: ${{ github.event.pull_request.author_association }} MAX_OPEN_PRS: ${{ vars.MAX_OPEN_PRS_PER_AUTHOR || 5 }} From e4175475bb7e630cefefa189b17282b063868a49 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 10 Aug 2026 21:30:34 +0800 Subject: [PATCH 038/126] fix(mito): re-encode bulk WAL entry after filling missing columns (#8808) When a bulk insert request carries a stale schema, the worker fills the missing columns via BulkPart::fill_missing_columns before writing. The method replaced the batch but kept raw_data (the original Arrow Flight IPC bytes), while BulkWalEntry::try_from(&BulkPart) prefers raw_data, so the memtable received the filled batch but the WAL recorded the pre-fill bytes. Replaying such an entry restores a batch that misses the filled columns: - Bulk memtable (flat format): convert_bulk_part fails with ColumnNotFound; the error is swallowed by the no-op write notifier and the rows are silently lost after restart. - Time series memtable: BulkPart::to_mutation builds rows shorter than the declared schema and the region worker panics with index out of bounds during replay, hanging the region open. Fixes: - fill_missing_columns clears raw_data so the WAL entry is re-encoded from the filled batch. - replay_memtable fills missing columns for replayed bulk parts of dense regions, so entries already written by affected versions replay correctly. Signed-off-by: jeremyhi --- src/mito2/src/engine/alter_test.rs | 214 +++++++++++++++++++++++++++- src/mito2/src/memtable/bulk/part.rs | 64 +++++++++ src/mito2/src/region/opener.rs | 13 +- 3 files changed, 284 insertions(+), 7 deletions(-) diff --git a/src/mito2/src/engine/alter_test.rs b/src/mito2/src/engine/alter_test.rs index be4daa1d88e..2461f2ca0cd 100644 --- a/src/mito2/src/engine/alter_test.rs +++ b/src/mito2/src/engine/alter_test.rs @@ -17,21 +17,26 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; +use api::v1::bulk_wal_entry::Body; use api::v1::helper::{row, tag_column_schema}; use api::v1::value::ValueData; -use api::v1::{ColumnDataType, Row, Rows, SemanticType, Value}; +use api::v1::{ArrowIpc, BulkWalEntry, ColumnDataType, Row, Rows, SemanticType, Value, WalEntry}; use common_error::ext::ErrorExt; use common_meta::ddl::utils::{parse_column_metadatas, parse_manifest_infos_from_extensions}; -use common_recordbatch::RecordBatches; +use common_recordbatch::{DfRecordBatch, RecordBatches}; +use common_test_util::flight::encode_to_flight_data; use datafusion_expr::col; +use datatypes::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; +use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextBackend, FulltextOptions}; +use store_api::logstore::provider::Provider; use store_api::metadata::ColumnMetadata; use store_api::metric_engine_consts::TABLE_COLUMN_METADATA_EXTENSION_KEY; use store_api::region_engine::{RegionEngine, RegionManifestInfo, RegionRole}; use store_api::region_request::{ - AddColumn, AddColumnLocation, AlterKind, PathType, RegionAlterRequest, RegionOpenRequest, - RegionRequest, SetIndexOption, SetRegionOption, + AddColumn, AddColumnLocation, AlterKind, PathType, RegionAlterRequest, + RegionBulkInsertsRequest, RegionOpenRequest, RegionRequest, SetIndexOption, SetRegionOption, }; use store_api::storage::{ColumnId, RegionId, ScanRequest}; @@ -42,9 +47,10 @@ use crate::error; use crate::sst::FormatType; use crate::test_util::batch_util::sort_batches_and_print; use crate::test_util::{ - CreateRequestBuilder, TestEnv, build_rows, build_rows_for_key, - column_metadata_to_column_schema, flush_region, put_rows, rows_schema, + CreateRequestBuilder, LogStoreImpl, TestEnv, build_rows, build_rows_for_key, + column_metadata_to_column_schema, flush_region, put_rows, reopen_region, rows_schema, }; +use crate::wal::Wal; async fn scan_check_after_alter(engine: &MitoEngine, region_id: RegionId, expected: &str) { let request = ScanRequest::default(); @@ -1978,3 +1984,199 @@ async fn test_alter_region_append_mode_invalid() { // append_mode should still be true check_append_mode(&engine, true); } + +/// Builds a batch against the schema before [add_nullable_field1] +/// (schema version 0): `(tag_0, field_0, ts)`. +fn build_schema_v0_batch() -> DfRecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("tag_0", DataType::Utf8, true), + Field::new("field_0", DataType::Float64, true), + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])); + let tag = Arc::new(StringArray::from_iter_values( + (0..3).map(|v| format!("a{v}")), + )) as ArrayRef; + let field = Arc::new(Float64Array::from_iter_values((0..3).map(|v| v as f64))) as ArrayRef; + let ts = Arc::new(TimestampMillisecondArray::from_iter_values( + (0..3).map(|v| v as i64 * 1000), + )) as ArrayRef; + DfRecordBatch::try_new(schema, vec![tag, field, ts]).unwrap() +} + +fn encode_arrow_ipc(batch: &DfRecordBatch) -> ArrowIpc { + let (schema, record_batch) = encode_to_flight_data(batch.clone()); + ArrowIpc { + schema: schema.data_header, + data_header: record_batch.data_header, + payload: record_batch.data_body, + } +} + +/// Builds a bulk insert request whose batch was built against the schema +/// before [add_nullable_field1] (schema version 0): `(tag_0, field_0, ts)`. +fn build_schema_v0_bulk_request(region_id: RegionId) -> RegionBulkInsertsRequest { + let payload = build_schema_v0_batch(); + let raw_data = encode_arrow_ipc(&payload); + + RegionBulkInsertsRequest { + region_id, + payload, + raw_data, + partition_expr_version: None, + // The writer only knows schema version 0 while the region is already + // at version 1, so the worker has to fill the missing columns. + aligned_schema_version: Some(0), + } +} + +async fn scan_all_sorted(engine: &MitoEngine, region_id: RegionId) -> String { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + sort_batches_and_print(&batches, &["tag_0", "ts"]) +} + +const SCHEMA_V0_BULK_ROWS: &str = "\ ++-------+---------+---------------------+---------+ +| tag_0 | field_0 | ts | field_1 | ++-------+---------+---------------------+---------+ +| a0 | 0.0 | 1970-01-01T00:00:00 | | +| a1 | 1.0 | 1970-01-01T00:00:01 | | +| a2 | 2.0 | 1970-01-01T00:00:02 | | ++-------+---------+---------------------+---------+"; + +#[tokio::test] +async fn test_bulk_insert_stale_schema_wal_replay() { + test_bulk_insert_stale_schema_wal_replay_with_format(false).await; + test_bulk_insert_stale_schema_wal_replay_with_format(true).await; +} + +async fn test_bulk_insert_stale_schema_wal_replay_with_format(flat_format: bool) { + common_telemetry::init_default_ut_logging(); + + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Bumps the region schema version to 1. + engine + .handle_request(region_id, RegionRequest::Alter(add_nullable_field1())) + .await + .unwrap(); + + // Bulk insert whose batch misses `field_1`. The worker fills the missing + // column before writing to the memtable and the WAL. + let response = engine + .handle_request( + region_id, + RegionRequest::BulkInserts(build_schema_v0_bulk_request(region_id)), + ) + .await + .unwrap(); + assert_eq!(3, response.affected_rows); + + assert_eq!( + SCHEMA_V0_BULK_ROWS, + scan_all_sorted(&engine, region_id).await, + "flat_format: {flat_format}" + ); + + // Reopens the region to replay the WAL. The rows must survive the replay. + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + + assert_eq!( + SCHEMA_V0_BULK_ROWS, + scan_all_sorted(&engine, region_id).await, + "flat_format: {flat_format}" + ); +} + +/// Replays a bulk WAL entry that misses columns added by an alter. This +/// simulates entries written by old versions that keep the stale raw data +/// when filling missing columns. +#[tokio::test] +async fn test_bulk_insert_stale_wal_entry_replay() { + test_bulk_insert_stale_wal_entry_replay_with_format(false).await; + test_bulk_insert_stale_wal_entry_replay_with_format(true).await; +} + +async fn test_bulk_insert_stale_wal_entry_replay_with_format(flat_format: bool) { + common_telemetry::init_default_ut_logging(); + + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Bumps the region schema version to 1. + engine + .handle_request(region_id, RegionRequest::Alter(add_nullable_field1())) + .await + .unwrap(); + + // Writes a bulk entry whose batch misses `field_1` to the WAL directly. + let batch = build_schema_v0_batch(); + let entry = WalEntry { + mutations: vec![], + bulk_entries: vec![BulkWalEntry { + sequence: 1, + max_ts: 2000, + min_ts: 0, + timestamp_index: 2, + body: Some(Body::ArrowIpc(encode_arrow_ipc(&batch))), + }], + }; + let LogStoreImpl::RaftEngine(log_store) = env.get_log_store().unwrap() else { + unreachable!() + }; + let wal = Wal::new(log_store); + let mut writer = wal.writer(); + writer + .add_entry( + region_id, + 1, + &entry, + &Provider::raft_engine_provider(region_id.as_u64()), + ) + .unwrap(); + writer.write_to_wal().await.unwrap(); + + // Reopens the region to replay the WAL. The replay must fill the missing + // column instead of losing the rows. + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + + assert_eq!( + SCHEMA_V0_BULK_ROWS, + scan_all_sorted(&engine, region_id).await, + "flat_format: {flat_format}" + ); +} diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index fcba4519478..20482794079 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -262,6 +262,12 @@ impl BulkPart { // Update the batch self.batch = new_batch; + // The raw data no longer matches the batch after filling missing + // columns. Clears it so the WAL entry is re-encoded from the filled + // batch, otherwise replaying the entry restores a batch without the + // filled columns and the memtable rejects it. + self.raw_data = None; + Ok(()) } @@ -2896,4 +2902,62 @@ mod tests { let total_rows: usize = reader.map(|r| r.unwrap().num_rows()).sum(); assert_eq!(total_rows, 6); } + + #[test] + fn test_fill_missing_columns_resets_raw_data() { + let metadata = metadata_for_test(); + + // A batch missing the `v1` field column, e.g. built against an older + // schema before `v1` was added. + let k0_array = Arc::new(arrow::array::StringArray::from(vec!["key1", "key2"])); + let k1_array = Arc::new(arrow::array::UInt32Array::from(vec![1u32, 2])); + let v0_array = Arc::new(arrow::array::Int64Array::from(vec![100i64, 200])); + let ts_array = Arc::new(TimestampMillisecondArray::from(vec![1000i64, 2000])); + let input_schema = Arc::new(Schema::new(vec![ + Field::new("k0", ArrowDataType::Utf8, false), + Field::new("k1", ArrowDataType::UInt32, false), + Field::new("v0", ArrowDataType::Int64, true), + Field::new( + "ts", + ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + ])); + let batch = + RecordBatch::try_new(input_schema, vec![k0_array, k1_array, v0_array, ts_array]) + .unwrap(); + + // Encodes the batch as raw data like the bulk insert request does. + let mut encoder = FlightEncoder::default(); + let schema_bytes = encoder.encode_schema(batch.schema().as_ref()).data_header; + let [flight_data] = encoder + .encode(FlightMessage::RecordBatch(batch.clone())) + .try_into() + .unwrap(); + let mut part = BulkPart { + batch, + max_timestamp: 2000, + min_timestamp: 1000, + sequence: 5, + timestamp_index: 3, + raw_data: Some(ArrowIpc { + schema: schema_bytes, + data_header: flight_data.data_header, + payload: flight_data.data_body, + }), + }; + + part.fill_missing_columns(&metadata).unwrap(); + assert!(part.batch.column_by_name("v1").is_some()); + // The raw data no longer matches the batch after filling missing + // columns and must be cleared, otherwise the WAL entry built from + // this part loses the filled columns. + assert!(part.raw_data.is_none()); + + // The WAL entry round trip keeps the filled column. + let entry = BulkWalEntry::try_from(&part).unwrap(); + let replayed = BulkPart::try_from(entry).unwrap(); + assert_eq!(2, replayed.num_rows()); + assert!(replayed.batch.column_by_name("v1").is_some()); + } } diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 9388f943ceb..5d2d36df107 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -32,6 +32,7 @@ use object_store::manager::ObjectStoreManagerRef; use object_store::util::{is_object_storage, normalize_dir}; use parquet::file::metadata::PageIndexPolicy; use snafu::{OptionExt, ResultExt, ensure}; +use store_api::codec::PrimaryKeyEncoding; use store_api::logstore::LogStore; use store_api::logstore::provider::Provider; use store_api::metadata::{ @@ -832,6 +833,7 @@ where // data in the WAL. let mut last_entry_id = flushed_entry_id; let replay_from_entry_id = flushed_entry_id + 1; + let region_metadata = version_control.current().version.metadata.clone(); let mut wal_stream = wal_entry_reader.read(provider, replay_from_entry_id)?; while let Some(res) = wal_stream.next().await { @@ -876,7 +878,16 @@ where } for bulk_entry in entry.bulk_entries { - let part = BulkPart::try_from(bulk_entry)?; + let mut part = BulkPart::try_from(bulk_entry)?; + // The entry may miss columns added by a concurrent alter if it was + // written by a writer with a stale schema (older versions kept the + // stale raw data when filling missing columns). Fills missing + // columns like the write path, otherwise the memtable rejects the + // batch. Skips sparse batches as they don't carry tag columns by + // design. + if region_metadata.primary_key_encoding != PrimaryKeyEncoding::Sparse { + part.fill_missing_columns(®ion_metadata)?; + } rows_replayed += part.num_rows(); // During replay, we should adopt the sequence from WAL. let bulk_sequence_from_wal = part.sequence; From 00e45c860222195367943942b8ec21fef46b7cf5 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 10 Aug 2026 22:57:08 +0800 Subject: [PATCH 039/126] test: fix unknown channel event context expectation (#8832) Signed-off-by: WenyXu --- tests-integration/tests/table_ddl_event.rs | 31 ++++++++++++---------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index 45d87f30971..c53842d66d6 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -90,7 +90,7 @@ async fn test_table_ddl_procedure_events() { "create_table", &auto_create_procedure_id, "auto_create", - "prometheus", + Some("prometheus"), ) .await; @@ -102,7 +102,7 @@ async fn test_table_ddl_procedure_events() { "create_table", &auto_influx_create_procedure_id, "auto_create", - "influx", + Some("influx"), ) .await; @@ -151,7 +151,7 @@ async fn test_table_ddl_procedure_events() { "create_table", &grpc_create_procedure_id, "manual", - "grpc", + Some("grpc"), ) .await; GrpcQueryHandler::do_query( @@ -190,7 +190,7 @@ async fn test_table_ddl_procedure_events() { "alter_table", &grpc_alter_procedure_id, "manual", - "grpc", + Some("grpc"), ) .await; @@ -215,7 +215,7 @@ async fn test_table_ddl_procedure_events() { "create_table", &create_table_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; assert_named_submitted_event( @@ -263,7 +263,7 @@ async fn test_table_ddl_procedure_events() { "alter_table", &alter_table_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; assert_named_submitted_event( @@ -307,7 +307,7 @@ async fn test_table_ddl_procedure_events() { "truncate_table", &truncate_table_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; assert_named_submitted_event( @@ -348,7 +348,7 @@ async fn test_table_ddl_procedure_events() { "drop_table", &drop_table_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; assert_named_submitted_event( @@ -389,7 +389,7 @@ async fn test_table_ddl_procedure_events() { "undrop_table", &undrop_table_procedure_id, "manual", - "unknown", + None, ) .await; assert_id_submitted_event( @@ -429,7 +429,7 @@ async fn test_table_ddl_procedure_events() { "purge_dropped_table", &purge_table_procedure_id, "manual", - "unknown", + None, ) .await; assert_id_submitted_event( @@ -483,7 +483,7 @@ async fn test_table_ddl_procedure_events() { "create_logical_tables", &create_logical_tables_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; @@ -536,7 +536,7 @@ async fn test_table_ddl_procedure_events() { "alter_logical_tables", &alter_logical_tables_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; assert_named_submitted_event( @@ -629,7 +629,7 @@ async fn assert_event_context( event_type: &str, procedure_id: &str, reason: &str, - protocol: &str, + protocol: Option<&str>, ) { let event_filter = format!( "FROM {EVENTS_TABLE} WHERE {} = '{event_type}' AND {} = '{procedure_id}' AND json_get_string({}, 'type') = 'Submitted'", @@ -649,7 +649,10 @@ async fn assert_event_context( "event_context", ) .await; - let expected = format!(r#"{{"protocol":"{protocol}","reason":"{reason}"}}"#); + let expected = match protocol { + Some(protocol) => format!(r#"{{"protocol":"{protocol}","reason":"{reason}"}}"#), + None => format!(r#"{{"reason":"{reason}"}}"#), + }; assert_eq!(expected, actual); } From cd6efa0abf6ec99737ad8bce36376c4946023edc Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 11 Aug 2026 11:43:17 +0800 Subject: [PATCH 040/126] fix(query): ignore field metadata in MergeScan remote schema validation (#8818) * fix(query): ignore field metadata in MergeScan remote schema validation validate_remote_schema compared Arrow Fields with the default Field equality, which includes field metadata. A column whose SST region metadata carries greptime:skipping_index (or greptime:inverted_index) while the frontend table schema does not declare it was misreported as a schema mismatch (HTTP 500 'advertised remote stream schema field mismatch'), even though name, data_type, and nullability were identical. Field metadata is auxiliary (index/encoding info) and is not part of field semantics. Compare name + data_type + nullability only; JSON fields keep their existing semantic comparison (wire Binary vs decoded Struct) and timestamp timezone-only differences remain accepted. * refactor(query): drop all field metadata comparison from MergeScan schema validation fields_semantically_equal now compares name + data_type + nullability only. The previous version still compared JSON identity metadata keys (TYPE_KEY, EXTENSION_TYPE_METADATA_KEY) and is_json_field; those are auxiliary and must not participate in field semantics. JSON fields keep their separate physical-type exemption (wire Binary vs decoded Struct) via json_fields_compatible, which never applies to non-JSON fields. --- src/query/src/dist_plan/merge_scan.rs | 127 ++++++++++++++++++++------ 1 file changed, 101 insertions(+), 26 deletions(-) diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 0245f98ddc4..aafc17c469f 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -113,6 +113,33 @@ fn merge_scan_schema_error_count_for_test() -> u64 { TEST_MERGE_SCAN_SCHEMA_ERRORS.with(Cell::get) } +/// Returns true when the field is a JSON column, identified either by its +/// Arrow extension type (`greptime.json` / `greptime.json2`) or by its +/// `greptime:type=Json` marker. +fn is_json_field(field: &Field) -> bool { + is_any_json_extension_type(field) + || field + .metadata() + .get(datatypes::schema::TYPE_KEY) + .map(String::as_ref) + == Some("Json") +} + +/// Returns true when two fields describe the same column semantics: same +/// name, same data type, and same nullability. +/// +/// Arrow `Field` metadata is auxiliary information (indexing, encoding, +/// compression hints, ...) and does not participate in the comparison. This +/// applies to every field, JSON columns included: JSON identity keys +/// (`greptime:type`, `ARROW:extension:metadata`) are never compared here. +/// The wire/decoded physical-type difference of JSON columns is exempted +/// separately by [`json_fields_compatible`]. +fn fields_semantically_equal(expected_field: &Field, actual_field: &Field) -> bool { + expected_field.name() == actual_field.name() + && expected_field.data_type() == actual_field.data_type() + && expected_field.is_nullable() == actual_field.is_nullable() +} + /// Returns true when two fields are semantically equivalent JSON columns. /// /// A JSON column is carried on the wire in its binary-encoded form (e.g. @@ -124,17 +151,8 @@ fn merge_scan_schema_error_count_for_test() -> u64 { /// settings (`ARROW:extension:metadata`, e.g. type hints) are the semantic /// parts of the field and must match. fn json_fields_compatible(expected_field: &Field, actual_field: &Field) -> bool { - let is_json = |field: &Field| { - is_any_json_extension_type(field) - || field - .metadata() - .get(datatypes::schema::TYPE_KEY) - .map(String::as_ref) - == Some("Json") - }; - - is_json(expected_field) - && is_json(actual_field) + is_json_field(expected_field) + && is_json_field(actual_field) && expected_field.name() == actual_field.name() && expected_field.is_nullable() == actual_field.is_nullable() // Both must carry the same JSON marker. @@ -151,12 +169,15 @@ fn json_fields_compatible(expected_field: &Field, actual_field: &Field) -> bool /// Validates the remote schema before positional column handling. /// -/// A timestamp timezone difference is the only intentional exception. It is -/// accepted by directly comparing the same timestamp unit, distinct timezones, -/// and equal name, nullability, and field metadata. Top-level Arrow schema -/// metadata is non-semantic at this boundary. JSON columns are additionally -/// compared semantically (see [`json_fields_compatible`]) because their wire -/// and decoded representations use different physical arrow types. +/// Field metadata (indexing/encoding/compression hints) is non-semantic at +/// this boundary and is ignored: two fields are equal when their name, data +/// type, and nullability match (see [`fields_semantically_equal`]). A +/// timestamp timezone difference is the only intentional exception, accepted +/// when the fields share the same timestamp unit, distinct timezones, and +/// equal name and nullability. Top-level Arrow schema metadata is +/// non-semantic at this boundary as well. JSON columns are compared +/// semantically (see [`json_fields_compatible`]) because their wire and +/// decoded representations use different physical arrow types. fn validate_remote_schema( expected: &ArrowSchema, actual: &ArrowSchema, @@ -176,7 +197,8 @@ fn validate_remote_schema( .zip(actual.fields().iter()) .enumerate() { - if expected_field == actual_field { + // Field metadata does not participate in the semantic comparison. + if fields_semantically_equal(expected_field, actual_field) { continue; } @@ -197,7 +219,6 @@ fn validate_remote_schema( && expected_timezone != actual_timezone && expected_field.name() == actual_field.name() && expected_field.is_nullable() == actual_field.is_nullable() - && expected_field.metadata() == actual_field.metadata() ); if !timezone_only_difference { return Err(remote_schema_mismatch(format!( @@ -2693,9 +2714,12 @@ mod tests { } #[test] - fn merge_scan_remote_schema_identity_rejects_json_vs_plain_binary_field() { - // A JSON column must not be accepted as compatible with a plain Binary - // column lacking the JSON extension metadata. + fn merge_scan_remote_schema_identity_ignores_json_extension_metadata_when_physical_type_matches() + { + // Field metadata (including the JSON extension identity) is not part + // of the semantic field identity: a JSON column in its binary wire + // form and a plain Binary column share the same name, physical type, + // and nullability, so they validate as equal. let expected = ArrowSchema::new(vec![ Field::new("j", TestArrowDataType::Binary, true).with_metadata(json_field_metadata( true, @@ -2703,11 +2727,26 @@ mod tests { )), ]); let actual = ArrowSchema::new(vec![Field::new("j", TestArrowDataType::Binary, true)]); - assert!(validate_remote_schema(&expected, &actual, "test").is_err()); + assert!(validate_remote_schema(&expected, &actual, "test").is_ok()); + + // The JSON compatibility exemption is scoped to JSON columns only: a + // pair of non-JSON fields with different physical types is still + // rejected even though one side is Binary. + let plain_binary = ArrowSchema::new(vec![Field::new("v", TestArrowDataType::Binary, true)]); + let plain_struct = ArrowSchema::new(vec![Field::new( + "v", + TestArrowDataType::Struct(arrow_schema::Fields::from(vec![Arc::new(Field::new( + "a", + TestArrowDataType::Utf8View, + true, + ))])), + true, + )]); + assert!(validate_remote_schema(&plain_binary, &plain_struct, "test").is_err()); } #[test] - fn merge_scan_remote_schema_identity_rejects_field_metadata_and_nullability_mismatch() { + fn merge_scan_remote_schema_identity_ignores_field_metadata_but_rejects_nullability_mismatch() { let expected = expected_int64_schema(); let metadata_mismatch = ArrowSchema::new_with_metadata( vec![ @@ -2730,10 +2769,42 @@ mod tests { ], expected.metadata().clone(), ); - assert!(validate_remote_schema(&expected, &metadata_mismatch, "test").is_err()); + // Field metadata is auxiliary and does not participate in the + // semantic field comparison. + assert!(validate_remote_schema(&expected, &metadata_mismatch, "test").is_ok()); assert!(validate_remote_schema(&expected, &nullability_mismatch, "test").is_err()); } + #[test] + fn merge_scan_remote_schema_identity_ignores_skipping_index_field_metadata() { + // Regression: a remote stream may advertise auxiliary field metadata + // (e.g. `greptime:skipping_index`) that the expected schema does not + // carry. Field metadata is not part of the semantic field identity + // (name + data type + nullability) and must not fail validation. + let skipping_index_metadata = StdHashMap::from([( + "greptime:skipping_index".to_string(), + r#"{"granularity":1,"false-positive-rate-in-10000":100,"index-type":"BloomFilter"}"# + .to_string(), + )]); + let expected = + ArrowSchema::new(vec![Field::new("value", TestArrowDataType::Float64, true)]); + let actual_with_metadata = ArrowSchema::new(vec![ + Field::new("value", TestArrowDataType::Float64, true) + .with_metadata(skipping_index_metadata), + ]); + assert!(validate_remote_schema(&expected, &actual_with_metadata, "test").is_ok()); + + // A real data type mismatch is still rejected. + let wrong_type = + ArrowSchema::new(vec![Field::new("value", TestArrowDataType::Int64, true)]); + assert!(validate_remote_schema(&expected, &wrong_type, "test").is_err()); + + // A name mismatch is still rejected. + let wrong_name = + ArrowSchema::new(vec![Field::new("other", TestArrowDataType::Float64, true)]); + assert!(validate_remote_schema(&expected, &wrong_name, "test").is_err()); + } + #[test] fn merge_scan_remote_schema_identity_rejects_timestamp_timezone_plus_field_mismatches() { let expected = ArrowSchema::new(vec![Field::new( @@ -2767,9 +2838,13 @@ mod tests { "different".to_string(), )])), ]); - for actual in [&different_unit, &different_name, &nullability, &metadata] { + for actual in [&different_unit, &different_name, &nullability] { assert!(validate_remote_schema(&expected, actual, "test").is_err()); } + // Field metadata is non-semantic and is ignored even for timezone + // pairs: a metadata difference on an otherwise timezone-only + // difference is accepted. + assert!(validate_remote_schema(&expected, &metadata, "test").is_ok()); } #[test] From b3faf22290e75650f1bc356ccf6b38355f1cbb29 Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 11 Aug 2026 13:52:31 +0800 Subject: [PATCH 041/126] fix(tests): make two Windows CI failures deterministic (Nightly CI #8837) (#8840) * fix(tests): reject overlay directories before opening on all platforms DatanodeOverlay::load() opened the target before checking is_file(). On Unix, File::open on a directory succeeds and the loader rejects it with "must be a regular file". On Windows, File::open on a directory fails up front with "Access is denied", so the type check was never reached and the rejects_directories_and_parse_errors test failed 4/4 in Nightly CI (issue #8837). Check std::fs::metadata before File::open: metadata succeeds on directories on both platforms, so the error message is now identical everywhere and the test assertion holds on Windows too. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(datanode): make test_region_error deterministic across platforms The second phase raced a 100ms mock handle delay against a 200ms replay_timeout; on busy Windows CI runners the error could land after the timeout fired, flaking reply.error.is_some() (Nightly CI, issue #8837). Use a mock handle that returns the error on its first poll with no delay: the catchup future completes before replay_timeout can ever fire, so the test no longer depends on wall-clock scheduling. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/heartbeat/handler/upgrade_region.rs | 31 ++++++++++++++++--- tests/runner/src/cmd/datanode_overlay.rs | 24 +++++++++----- 2 files changed, 43 insertions(+), 12 deletions(-) diff --git a/src/datanode/src/heartbeat/handler/upgrade_region.rs b/src/datanode/src/heartbeat/handler/upgrade_region.rs index 05292415be8..e2a5677fcad 100644 --- a/src/datanode/src/heartbeat/handler/upgrade_region.rs +++ b/src/datanode/src/heartbeat/handler/upgrade_region.rs @@ -408,9 +408,12 @@ mod tests { #[tokio::test] async fn test_region_error() { common_telemetry::init_default_ut_logging(); - let mock_region_server = mock_region_server(); let region_id = RegionId::new(1024, 1); + // With the default (zero) `replay_timeout`, the handler times out before the + // mock handle's 100ms delay elapses: it didn't wait for handle returns; it + // had no idea about the error. + let region_server = mock_region_server(); let (mock_engine, _) = MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| { // Region is not ready. @@ -424,9 +427,9 @@ mod tests { // Note: Don't change. region_engine.handle_request_delay = Some(Duration::from_millis(100)); }); - mock_region_server.register_test_region(region_id, mock_engine); + region_server.register_test_region(region_id, mock_engine); let kv_backend = Arc::new(MemoryKvBackend::new()); - let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend); + let handler_context = HandlerContext::new_for_test(region_server, kv_backend); let reply = UpgradeRegionsHandler::new_test() .handle( &handler_context, @@ -437,12 +440,32 @@ mod tests { ) .await; - // It didn't wait for handle returns; it had no idea about the error. let reply = &reply.unwrap().expect_upgrade_regions_reply()[0]; assert!(!reply.ready); assert!(reply.exists); assert!(reply.error.is_none()); + // With a 200ms `replay_timeout` and a mock handle that returns the error + // immediately (no delay), the handler waits for the handle and propagates + // its error deterministically: the catchup future completes on its first + // poll, so the `replay_timeout` can never fire first. This avoids racing + // real wall-clock time (a 100ms delay vs 200ms timeout is flaky on + // Windows CI). + let mock_region_server = mock_region_server(); + let (mock_engine, _) = + MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| { + // Region is not ready. + region_engine.mock_role = Some(Some(RegionRole::Follower)); + region_engine.handle_request_mock_fn = Some(Box::new(|_, _| { + error::UnexpectedSnafu { + violated: "mock_error".to_string(), + } + .fail() + })); + }); + mock_region_server.register_test_region(region_id, mock_engine); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend); let reply = UpgradeRegionsHandler::new_test() .handle( &handler_context, diff --git a/tests/runner/src/cmd/datanode_overlay.rs b/tests/runner/src/cmd/datanode_overlay.rs index 21892b9a091..5feefd99632 100644 --- a/tests/runner/src/cmd/datanode_overlay.rs +++ b/tests/runner/src/cmd/datanode_overlay.rs @@ -106,13 +106,11 @@ impl DatanodeOverlay { )); } - let mut file = File::open(&canonical_target).map_err(|error| { - format!( - "Failed to open datanode overlay {}: {error}", - canonical_target.display() - ) - })?; - let metadata = file.metadata().map_err(|error| { + // Check the type before opening so the error is identical on every + // platform: `std::fs::metadata` succeeds on directories on both Unix + // and Windows, whereas `File::open` on a directory fails up front on + // Windows (with a platform-specific message). + let metadata = std::fs::metadata(&canonical_target).map_err(|error| { format!( "Failed to inspect datanode overlay {}: {error}", canonical_target.display() @@ -125,6 +123,13 @@ impl DatanodeOverlay { )); } + let mut file = File::open(&canonical_target).map_err(|error| { + format!( + "Failed to open datanode overlay {}: {error}", + canonical_target.display() + ) + })?; + let mut content = String::new(); file.read_to_string(&mut content).map_err(|error| { format!( @@ -613,7 +618,10 @@ mod tests { let directory_error = DatanodeOverlay::load(temp_dir.path(), Path::new("directory.toml")).unwrap_err(); - assert!(directory_error.contains("regular file")); + assert!( + directory_error.contains("regular file"), + "unexpected error for directory.toml: {directory_error}" + ); let parse_error = DatanodeOverlay::load(temp_dir.path(), Path::new("broken.toml")).unwrap_err(); assert!(parse_error.contains("Failed to parse datanode overlay")); From 4eeb4052d6e4ed486266f6c599e60280671a4aac Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 11 Aug 2026 14:20:08 +0800 Subject: [PATCH 042/126] feat(servers): stamp prometheus remote write v2 metadata as semantic table options (#8797) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(servers): stamp prometheus remote write v2 metadata as semantic table options Remote write 2.0 carries per-series metadata (type, unit, help) that the v2 ingest decoded and dropped; tables kept the name-based 'inferred' quality. Wire it into the semantic layer: - generalize the OTLP per-table semantic index into a shared, schema- aware servers::semantic module: v2 lets each series override its target schema, so the index is keyed {schema -> table -> options} and the same metric name in two schemas no longer collapses; - into_write_requests records metric type and unit per written table; an explicit type upgrades the table's metadata quality to declared, UNSPECIFIED series keep the request-level inferred stamp, and units are canonicalised from OpenMetrics words to the UCUM codes the vocabulary is defined in (unknown units are dropped, help text is not persisted); - the consumer folds the index in on both auto-create paths: the operator row-insert path and the pending-rows batched create, which bypasses the former. Table options are stamped at auto-create only; updating existing tables from later metadata is future work (a metadata registry, see the native histograms RFC). Signed-off-by: Dennis Zhuang * chore: trim over-commenting in the remote write metadata path Signed-off-by: Dennis Zhuang * fix: parse the per-table semantic index once per create round The index was re-parsed from JSON for every table being created — a first write creating N tables (a fleet's first scrape) paid O(N x index size). Parsing now happens lazily once per create-planning round, on both consumers: the operator row-insert auto-create (also serving OTLP metrics) and the pending-rows batched create. Also validate every non-zero metadata symbol reference up front, as the remote write 2.0 spec requires: help_ref was never checked, and unit_ref escaped checking when the metric type was UNSPECIFIED. Signed-off-by: Dennis Zhuang * fix(servers): stamp remote-write v2 unit independently of metric type OpenMetrics models TYPE and UNIT as independent MetricFamily metadata, and the Prometheus v2 sender emits UNSPECIFIED-type series that still carry a unit. The early return on UNSPECIFIED dropped that unit, which is unrecoverable after table auto-create (units are not stored in rows). Stamp the mapped unit whenever present; the type and the metadata_quality=declared upgrade still require an explicit type. Signed-off-by: Dennis Zhuang * chore: trim restating comments in the v2 metadata path Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/instance/otlp.rs | 5 +- src/frontend/src/instance/prom_store.rs | 15 +- src/operator/src/insert.rs | 142 +++++++--- src/servers/src/http/prom_store.rs | 23 +- src/servers/src/lib.rs | 1 + src/servers/src/otlp/metrics.rs | 20 +- src/servers/src/otlp/metrics/semantic.rs | 177 ------------ src/servers/src/prom_remote_write/v2.rs | 300 ++++++++++++++++++++- src/servers/src/semantic.rs | 325 +++++++++++++++++++++++ tests-integration/tests/http.rs | 97 ++++++- 10 files changed, 865 insertions(+), 240 deletions(-) delete mode 100644 src/servers/src/otlp/metrics/semantic.rs create mode 100644 src/servers/src/semantic.rs diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index c4616ae578d..459a05321cb 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -132,8 +132,9 @@ impl OpenTelemetryProtocolHandler for Instance { c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC); c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY); // Per-table metric specifics + resource/scope lineage ride this - // internal channel; the auto-create path folds them per table name. - if let Some(index) = semantic_index.encode() { + // internal channel; the auto-create path folds them per schema and + // table name. + if let Some(index) = semantic_index.encode(&c.current_schema()) { c.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index); } if !is_legacy { diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index 1bfb5f56437..428628df123 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -36,7 +36,8 @@ use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_value}; use common_recordbatch::RecordBatches; use common_telemetry::{debug, tracing}; use operator::insert::{ - AutoCreateTableType, InserterRef, build_create_table_expr, fill_table_options_for_create, + AutoCreateTableType, InserterRef, PerTableSemanticIndex, apply_per_table_semantic_options, + build_create_table_expr, fill_table_options_for_create, parse_per_table_semantic_index, }; use operator::statement::StatementExecutor; use prost::Message; @@ -313,6 +314,7 @@ impl PendingRowsSchemaAlterer for Instance { // Check which tables actually still need to be created (may have been // concurrently created by another request). let mut create_exprs: Vec = Vec::with_capacity(tables.len()); + let mut per_table_semantics: Option> = None; for &(table_name, request_schema) in tables { let existing = self .catalog_manager() @@ -331,6 +333,17 @@ impl PendingRowsSchemaAlterer for Instance { let mut table_options = std::collections::HashMap::with_capacity(4); fill_table_options_for_create(&mut table_options, &create_type, &ctx); + // The batched create path bypasses the operator's auto-create, so + // fold the per-table semantic index in here too. + let semantic_index = per_table_semantics + .get_or_insert_with(|| parse_per_table_semantic_index(&ctx)) + .as_ref(); + apply_per_table_semantic_options( + &mut table_options, + semantic_index, + schema, + table_name, + ); create_table_expr.table_options.extend(table_options); create_exprs.push(create_table_expr); } diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index b5b6be093d9..fabc1b3778b 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::BTreeMap; use std::sync::Arc; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; @@ -614,6 +615,7 @@ impl Inserter { let mut alter_tables = vec![]; let mut need_refresh_table_infos = HashSet::new(); let mut instant_table_ids = HashSet::new(); + let mut per_table_semantics: Option> = None; for req in &mut requests.inserts { match self.get_table(catalog, &schema, &req.table_name).await? { @@ -654,8 +656,15 @@ impl Inserter { table_infos.insert(table_info.table_id(), table_info); } None => { - let create_expr = - self.get_create_table_expr_on_demand(req, &auto_create_table_type, ctx)?; + let semantic_index = per_table_semantics + .get_or_insert_with(|| parse_per_table_semantic_index(ctx)) + .as_ref(); + let create_expr = self.get_create_table_expr_on_demand( + req, + &auto_create_table_type, + ctx, + semantic_index, + )?; create_tables.push(create_expr); } } @@ -933,11 +942,17 @@ impl Inserter { req: &RowInsertRequest, create_type: &AutoCreateTableType, ctx: &QueryContextRef, + semantic_index: Option<&PerTableSemanticIndex>, ) -> Result { let schema = ctx.current_schema(); let mut table_options = std::collections::HashMap::with_capacity(4); fill_table_options_for_create(&mut table_options, create_type, ctx); - apply_per_table_semantic_options(&mut table_options, ctx, &req.table_name); + apply_per_table_semantic_options( + &mut table_options, + semantic_index, + ctx.current_schema().as_str(), + &req.table_name, + ); let engine_name = if let AutoCreateTableType::Logical(_) = create_type { // engine should be metric engine when creating logical tables. @@ -1284,32 +1299,45 @@ pub fn fill_table_options_for_create( } } -/// Folds the semantic keys for `table_name` carried on the internal per-table -/// index extension into `table_options`. +/// The parsed per-table semantic index: `{schema -> {table -> {key -> value}}}`, +/// produced by the OTLP metrics encode path (where one metric can fan out into +/// several tables with distinct keys) and the Prometheus remote write v2 path +/// (where per-series metadata declares type/unit, and a series may override its +/// target schema). +pub type PerTableSemanticIndex = BTreeMap>>; + +/// Parses the per-table semantic index off the context extension. Call once per +/// create-planning round: a first write creating N tables would otherwise +/// re-parse the whole index N times. `None` when the request carries no index +/// (logs, traces, Prom RW v1) or it fails to parse. +pub fn parse_per_table_semantic_index(ctx: &QueryContextRef) -> Option { + let raw = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY)?; + match serde_json::from_str(raw) { + Ok(index) => Some(index), + Err(_) => { + warn!("failed to parse semantic per-table index, skipping per-table options"); + None + } + } +} + +/// Folds the semantic keys of the table being created into `table_options`. /// -/// The index is a `{table_name -> {semantic_key: value}}` JSON blob produced by -/// the OTLP metrics encode path (where one metric can fan out into several -/// tables with distinct keys). Common keys shared by every table in a request -/// travel as plain semantic extensions and are handled by -/// [`fill_table_options_for_create`]; this carries only the per-table tail. -/// Keys are re-checked against the vocabulary defensively. Ingestion paths -/// without a per-table index (logs, traces, Prom RW) carry no extension, so this -/// is a no-op for them. -fn apply_per_table_semantic_options( +/// Common keys shared by every table in a request travel as plain semantic +/// extensions and are handled by [`fill_table_options_for_create`]; this +/// carries only the per-table tail and is applied after it, so a per-table +/// value (e.g. `declared` quality) wins. Keys are re-checked against the +/// vocabulary defensively. +pub fn apply_per_table_semantic_options( table_options: &mut std::collections::HashMap, - ctx: &QueryContextRef, + index: Option<&PerTableSemanticIndex>, + schema: &str, table_name: &str, ) { - let Some(raw) = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY) else { - return; - }; - let Ok(index) = serde_json::from_str::< - std::collections::BTreeMap>, - >(raw) else { - warn!("failed to parse semantic per-table index, skipping per-table options"); - return; - }; - let Some(entry) = index.get(table_name) else { + let Some(entry) = index + .and_then(|index| index.get(schema)) + .and_then(|tables| tables.get(table_name)) + else { return; }; for (key, value) in entry { @@ -1734,22 +1762,41 @@ mod tests { SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, SEMANTIC_PER_TABLE_INDEX_KEY, }; - let index = r#"{ - "http_requests_total": { - "greptime.semantic.metric.type": "counter", - "greptime.semantic.metric.unit": "By", - "greptime.semantic.metric.type_BOGUS": "x" - }, - "other_table": { - "greptime.semantic.metric.type": "gauge" - } - }"#; + let index = format!( + r#"{{ + "{DEFAULT_SCHEMA_NAME}": {{ + "http_requests_total": {{ + "greptime.semantic.metric.type": "counter", + "greptime.semantic.metric.unit": "By", + "greptime.semantic.metric.type_BOGUS": "x" + }}, + "other_table": {{ + "greptime.semantic.metric.type": "gauge" + }} + }}, + "other_schema": {{ + "http_requests_total": {{ + "greptime.semantic.metric.type": "gauge" + }} + }} + }}"# + ); let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index); let ctx = Arc::new(ctx); + let index = parse_per_table_semantic_index(&ctx); + assert!(index.is_some()); + let index = index.as_ref(); + let mut table_options = std::collections::HashMap::new(); - apply_per_table_semantic_options(&mut table_options, &ctx, "http_requests_total"); + apply_per_table_semantic_options( + &mut table_options, + index, + DEFAULT_SCHEMA_NAME, + "http_requests_total", + ); + // The write schema's entry applies — not other_schema's `gauge`. assert_eq!( table_options.get(SEMANTIC_METRIC_TYPE).map(String::as_str), Some("counter") @@ -1764,16 +1811,33 @@ mod tests { assert_eq!(table_options.len(), 2); let mut empty = std::collections::HashMap::new(); - apply_per_table_semantic_options(&mut empty, &ctx, "not_in_index"); + apply_per_table_semantic_options(&mut empty, index, DEFAULT_SCHEMA_NAME, "not_in_index"); assert!(empty.is_empty()); - // No extension at all is a no-op (e.g. logs / Prom RW). + // A schema with no entry is a no-op even when the table name matches + // elsewhere. + let mut opts = std::collections::HashMap::new(); + apply_per_table_semantic_options( + &mut opts, + index, + "schema_without_entry", + "http_requests_total", + ); + assert!(opts.is_empty()); + + // No extension at all parses to no index (e.g. logs / Prom RW v1). let bare = Arc::new(QueryContext::with( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, )); + assert!(parse_per_table_semantic_index(&bare).is_none()); let mut opts = std::collections::HashMap::new(); - apply_per_table_semantic_options(&mut opts, &bare, "http_requests_total"); + apply_per_table_semantic_options( + &mut opts, + None, + DEFAULT_SCHEMA_NAME, + "http_requests_total", + ); assert!(opts.is_empty()); } diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index f659656e03e..ac4f39628d1 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -35,8 +35,9 @@ use serde::{Deserialize, Serialize}; use session::context::{Channel, QueryContext, QueryContextRef}; use snafu::prelude::*; use table::requests::{ - METADATA_QUALITY_INFERRED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_SIGNAL_TYPE, - SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, SOURCE_PROMETHEUS, + METADATA_QUALITY_INFERRED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_PER_TABLE_INDEX_KEY, + SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, + SOURCE_PROMETHEUS, }; use crate::error::{self, InternalSnafu, PipelineSnafu, Result}; @@ -155,6 +156,7 @@ async fn remote_write_v1( let (db, query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V1_VERSION); + let query_ctx = Arc::new(query_ctx); let mut processor = PromSeriesProcessor::default_processor(); @@ -230,7 +232,7 @@ async fn remote_write_v2( // optional pipeline parameter and ingest samples directly. let _ = pipeline_info; - let (db, query_ctx, _timer) = + let (db, mut query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V2_VERSION); let request = match decode_remote_write_v2_request(is_zstd, body) { @@ -253,6 +255,12 @@ async fn remote_write_v2( Ok(req) => req, Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), }; + // The v2 per-series metadata upgrades the written tables' semantic options + // (metric type/unit, declared quality) at auto-create time. + if let Some(index) = req.semantic_index.encode(&query_ctx.current_schema()) { + query_ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index); + } + let query_ctx = Arc::new(query_ctx); let sample_count = req.sample_count; let histogram_count = req.histogram_count; let sample_batches = into_prom_write_batches(req.samples, query_ctx.clone()); @@ -310,11 +318,14 @@ fn vm_proto_version_response(params: &RemoteWriteQuery) -> Option (String, Arc, HistogramTimer) { +) -> (String, QueryContext, HistogramTimer) { let db = params.db.clone().unwrap_or_default(); query_ctx.set_channel(Channel::Prometheus); let physical_table = params @@ -325,12 +336,12 @@ fn prepare_remote_write_context( // Stamp the Prometheus metric identity here, before `as_req_iter` splits into the // batched and direct write paths, so both inherit it (the batched path bypasses // `PromStoreProtocolHandler::write`). Prometheus remote-write metadata is weak - // here, so the type is inferred from naming. + // here, so the type is inferred from naming; v2 upgrades tables whose series + // carry inline metadata via the per-table index. query_ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC); query_ctx.set_extension(SEMANTIC_SOURCE, SOURCE_PROMETHEUS); query_ctx.set_extension(SEMANTIC_SOURCE_VERSION, remote_write_version); query_ctx.set_extension(SEMANTIC_METRIC_METADATA_QUALITY, METADATA_QUALITY_INFERRED); - let query_ctx = Arc::new(query_ctx); let timer = crate::metrics::METRIC_HTTP_PROM_STORE_WRITE_ELAPSED .with_label_values(&[db.as_str()]) .start_timer(); diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index 55d3b65fed5..2a3796aab21 100644 --- a/src/servers/src/lib.rs +++ b/src/servers/src/lib.rs @@ -50,6 +50,7 @@ pub mod repeated_field; pub mod request_memory_limiter; pub mod request_memory_metrics; mod row_writer; +pub mod semantic; pub mod server; pub mod tls; diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 56ab0530338..b43cc73338f 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -29,11 +29,14 @@ use table::requests::{ use crate::error::Result; use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME}; use crate::row_writer::{self, MultiTableData, TableData}; +pub use crate::semantic::SemanticIndex; +use crate::semantic::{ + METRIC_TYPE_COUNTER, METRIC_TYPE_GAUGE, METRIC_TYPE_HISTOGRAM, METRIC_TYPE_SUMMARY, + METRIC_TYPE_UPDOWN_COUNTER, +}; -mod semantic; mod translator; -pub use semantic::SemanticIndex; pub use translator::legacy_normalize_otlp_name; use translator::{translate_label_name, translate_metric_name}; @@ -44,15 +47,6 @@ const COUNT_TABLE_SUFFIX: &str = "_count"; const SUM_TABLE_SUFFIX: &str = "_sum"; const BUCKET_TABLE_SUFFIX: &str = "_bucket"; -// `greptime.semantic.metric.type` values stamped per emitted table. Must stay -// within the domain accepted by `validate_semantic_option`; the drift-guard test -// asserts this. -const METRIC_TYPE_COUNTER: &str = "counter"; -const METRIC_TYPE_UPDOWN_COUNTER: &str = "updown_counter"; -const METRIC_TYPE_GAUGE: &str = "gauge"; -const METRIC_TYPE_HISTOGRAM: &str = "histogram"; -const METRIC_TYPE_SUMMARY: &str = "summary"; - const JOB_KEY: &str = "job"; const INSTANCE_KEY: &str = "instance"; @@ -1183,7 +1177,9 @@ mod tests { use table::requests::validate_semantic_option; fn decode(index: &SemanticIndex) -> BTreeMap> { - serde_json::from_str(&index.encode().expect("non-empty index")).unwrap() + let nested: BTreeMap>> = + serde_json::from_str(&index.encode("public").expect("non-empty index")).unwrap(); + nested.into_values().next().unwrap() } fn record(metric: &Metric, metric_type: MetricType, name: &str) -> SemanticIndex { diff --git a/src/servers/src/otlp/metrics/semantic.rs b/src/servers/src/otlp/metrics/semantic.rs deleted file mode 100644 index 0726dc15c12..00000000000 --- a/src/servers/src/otlp/metrics/semantic.rs +++ /dev/null @@ -1,177 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Per-table semantic metadata accumulated during one OTLP metrics encode pass. -//! -//! A metric emits one or more tables (histogram and summary fan out into -//! `_bucket`/`_sum`/`_count` companions). Each emitted table collects the -//! metric's scalar semantic keys. The resulting index is serialized onto the -//! `greptime.internal.semantic.per_table_index` context extension and folded -//! into each table's options at auto-create time. -//! -//! Conflict handling follows the RFC: when two sources disagree on a -//! single-valued key the value collapses to `mixed` (or `unknown` for keys whose -//! domain has no `mixed`). - -use std::collections::{BTreeMap, HashMap}; - -use table::requests::{SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, validate_semantic_option}; - -/// Index of `{table_name -> {semantic_key -> value}}` built while encoding. -#[derive(Debug, Default)] -pub struct SemanticIndex { - /// Per-table scalar keys; conflicting values collapse to `mixed`/`unknown`. - tables: HashMap>, -} - -impl SemanticIndex { - pub fn is_empty(&self) -> bool { - self.tables.is_empty() - } - - /// Records a scalar semantic key for `table`. A value conflicting with one - /// already recorded collapses the key to `mixed`/`unknown`; once collapsed - /// it stays collapsed. - pub fn record_scalar(&mut self, table: &str, key: &'static str, value: &str) { - // Avoid allocating the table name (and an empty map) on the common path - // where the table is already present. - if let Some(scalars) = self.tables.get_mut(table) { - match scalars.get(key).map(String::as_str) { - Some(existing) if existing == value => {} - Some(SEMANTIC_VALUE_MIXED) | Some(SEMANTIC_VALUE_UNKNOWN) => {} - Some(_) => { - scalars.insert(key, collapse_value(key)); - } - None => { - scalars.insert(key, value.to_string()); - } - } - } else { - self.tables.insert( - table.to_string(), - BTreeMap::from([(key, value.to_string())]), - ); - } - } - - /// Serializes to the JSON `{table -> {key -> value}}` carried on the context - /// extension, or `None` when nothing was recorded. - pub fn encode(&self) -> Option { - if self.tables.is_empty() { - return None; - } - serde_json::to_string(&self.tables).ok() - } - - #[cfg(test)] - fn options_of(&self, table: &str) -> Option<&BTreeMap<&'static str, String>> { - self.tables.get(table) - } -} - -/// The collapsed value for a conflicting scalar key: `mixed` when the key's -/// domain accepts it, else `unknown`. Uses the vocabulary validator as the -/// single source of truth for which keys allow `mixed`. -fn collapse_value(key: &str) -> String { - if validate_semantic_option(key, SEMANTIC_VALUE_MIXED) { - SEMANTIC_VALUE_MIXED.to_string() - } else { - SEMANTIC_VALUE_UNKNOWN.to_string() - } -} - -#[cfg(test)] -mod tests { - use table::requests::{ - SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, - }; - - use super::*; - - #[test] - fn test_scalar_recording_keeps_first_then_collapses_on_conflict() { - let mut index = SemanticIndex::default(); - index.record_scalar("t", SEMANTIC_METRIC_TYPE, "counter"); - index.record_scalar("t", SEMANTIC_METRIC_TYPE, "counter"); - assert_eq!( - index - .options_of("t") - .unwrap() - .get(SEMANTIC_METRIC_TYPE) - .map(String::as_str), - Some("counter") - ); - - // Conflict on a key whose domain has `mixed` collapses to `mixed`. - index.record_scalar("t", SEMANTIC_METRIC_TYPE, "gauge"); - assert_eq!( - index - .options_of("t") - .unwrap() - .get(SEMANTIC_METRIC_TYPE) - .map(String::as_str), - Some("mixed") - ); - // Further writes stay collapsed. - index.record_scalar("t", SEMANTIC_METRIC_TYPE, "histogram"); - assert_eq!( - index - .options_of("t") - .unwrap() - .get(SEMANTIC_METRIC_TYPE) - .map(String::as_str), - Some("mixed") - ); - } - - #[test] - fn test_scalar_conflict_without_mixed_domain_collapses_to_unknown() { - let mut index = SemanticIndex::default(); - index.record_scalar("t", SEMANTIC_METRIC_METADATA_QUALITY, "declared"); - index.record_scalar("t", SEMANTIC_METRIC_METADATA_QUALITY, "inferred"); - // metadata_quality accepts only declared/inferred/unknown, so a conflict - // is `unknown`. - assert_eq!( - index - .options_of("t") - .unwrap() - .get(SEMANTIC_METRIC_METADATA_QUALITY) - .map(String::as_str), - Some("unknown") - ); - } - - #[test] - fn test_encode_is_none_when_empty_and_round_trips() { - let index = SemanticIndex::default(); - assert!(index.is_empty()); - assert_eq!(index.encode(), None); - - let mut index = SemanticIndex::default(); - index.record_scalar("metric_a", SEMANTIC_METRIC_TYPE, "counter"); - index.record_scalar("metric_a", SEMANTIC_METRIC_UNIT, "By"); - let json = index.encode().unwrap(); - let parsed: BTreeMap> = - serde_json::from_str(&json).unwrap(); - let table = parsed.get("metric_a").unwrap(); - assert_eq!( - table.get(SEMANTIC_METRIC_TYPE).map(String::as_str), - Some("counter") - ); - assert_eq!( - table.get(SEMANTIC_METRIC_UNIT).map(String::as_str), - Some("By") - ); - } -} diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index cfa772f6919..9940ad2e7c2 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -17,10 +17,10 @@ use std::collections::hash_map::Entry; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; use api::greptime_proto::io::prometheus::write::v2::{ - BucketSpan, Histogram, Request, Sample, TimeSeries, + BucketSpan, Histogram, Request, Sample, TimeSeries, metadata, }; #[cfg(test)] -use api::greptime_proto::io::prometheus::write::v2::{Exemplar, Metadata, metadata}; +use api::greptime_proto::io::prometheus::write::v2::{Exemplar, Metadata}; use api::helper::ColumnDataTypeWrapper; use api::v1::value::ValueData; use api::v1::{ColumnSchema, ListValue, RowInsertRequest, Rows, SemanticType, Value}; @@ -31,6 +31,10 @@ use common_query::prelude::{greptime_native_histogram, greptime_timestamp, grept use pipeline::{ContextOpt, ContextReq}; use prost::Message; use snafu::{OptionExt, ResultExt, ensure}; +use table::requests::{ + METADATA_QUALITY_DECLARED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_TYPE, + SEMANTIC_METRIC_UNIT, +}; use crate::error::{self, Result}; use crate::prom_remote_write::row_builder::PromCtx; @@ -42,6 +46,11 @@ use crate::prom_store::{ PHYSICAL_TABLE_LABEL_ALT, SCHEMA_LABEL, }; use crate::row_writer::{self, TableData}; +use crate::semantic::{ + METRIC_TYPE_COUNTER, METRIC_TYPE_GAUGE, METRIC_TYPE_GAUGE_HISTOGRAM, METRIC_TYPE_HISTOGRAM, + METRIC_TYPE_INFO, METRIC_TYPE_STATESET, METRIC_TYPE_SUMMARY, SemanticIndexes, + openmetrics_unit_to_ucum, +}; type PromTags = Vec<(String, String)>; type ResolvedSeriesLabels = (PromCtx, String, PromTags); @@ -67,6 +76,9 @@ pub(crate) struct RemoteWriteV2WriteRequests { pub histograms: ContextReq, pub sample_count: u64, pub histogram_count: u64, + /// Per-table semantic metadata from the series' inline `Metadata`, folded + /// into table options at auto-create time. + pub semantic_index: SemanticIndexes, } /// Converts a PRW v2 request into normal sample writes and native histogram writes. @@ -92,6 +104,7 @@ pub(crate) fn into_write_requests(request: Request) -> Result Result 0 { histogram_count > 0 || histogram_tables @@ -166,9 +186,78 @@ pub(crate) fn into_write_requests(request: Request) -> Result Result<()> { + let Some(series_metadata) = &series.metadata else { + return Ok(()); + }; + // Symbol 0 is the mandatory empty string: no help / no unit. + if series_metadata.help_ref != 0 { + symbol_ref(symbols, series_metadata.help_ref, "metadata help")?; + } + let unit = if series_metadata.unit_ref != 0 { + Some(symbol_ref( + symbols, + series_metadata.unit_ref, + "metadata unit", + )?) + } else { + None + }; + let metric_type = metric_type_value(series_metadata.r#type); + let ucum = unit.and_then(|unit| openmetrics_unit_to_ucum(unit.trim())); + if metric_type.is_none() && ucum.is_none() { + return Ok(()); + } + + let index = index.index_for(prom_ctx.schema.as_deref()); + if let Some(metric_type) = metric_type { + index.record_scalar(table_name, SEMANTIC_METRIC_TYPE, metric_type); + index.record_scalar( + table_name, + SEMANTIC_METRIC_METADATA_QUALITY, + METADATA_QUALITY_DECLARED, + ); + } + if let Some(ucum) = ucum { + index.record_scalar(table_name, SEMANTIC_METRIC_UNIT, ucum); + } + Ok(()) +} + +/// The `greptime.semantic.metric.type` value for a wire metric type; `None` +/// for `UNSPECIFIED` (nothing was declared) and out-of-range values. +fn metric_type_value(wire_type: i32) -> Option<&'static str> { + match metadata::MetricType::try_from(wire_type).ok()? { + metadata::MetricType::Unspecified => None, + metadata::MetricType::Counter => Some(METRIC_TYPE_COUNTER), + metadata::MetricType::Gauge => Some(METRIC_TYPE_GAUGE), + metadata::MetricType::Histogram => Some(METRIC_TYPE_HISTOGRAM), + metadata::MetricType::Gaugehistogram => Some(METRIC_TYPE_GAUGE_HISTOGRAM), + metadata::MetricType::Summary => Some(METRIC_TYPE_SUMMARY), + metadata::MetricType::Info => Some(METRIC_TYPE_INFO), + metadata::MetricType::Stateset => Some(METRIC_TYPE_STATESET), + } +} + fn get_or_create_table_data( tables: &mut HashMap>, prom_ctx: PromCtx, @@ -1823,4 +1912,211 @@ mod tests { fn is_empty_list(value: Option) -> bool { matches!(value, Some(ValueData::ListValue(list)) if list.items.is_empty()) } + + fn push_test_symbol(symbols: &mut Vec, symbol: &str) -> u32 { + if let Some(idx) = symbols.iter().position(|s| s == symbol) { + return idx as u32; + } + let idx = symbols.len(); + symbols.push(symbol.to_string()); + idx as u32 + } + + /// One sample-carrying test series: `(labels, metric_type, unit)`. + type MetadataSeries<'a> = (Vec<(&'a str, &'a str)>, i32, Option<&'a str>); + + fn metadata_request(series: Vec>) -> Request { + let mut symbols = vec!["".to_string()]; + let timeseries = series + .into_iter() + .map(|(labels, metric_type, unit)| { + let mut labels_refs = Vec::with_capacity(labels.len() * 2); + for (name, value) in labels { + labels_refs.push(push_test_symbol(&mut symbols, name)); + labels_refs.push(push_test_symbol(&mut symbols, value)); + } + let unit_ref = unit.map_or(0, |unit| push_test_symbol(&mut symbols, unit)); + TimeSeries { + labels_refs, + samples: vec![Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + histograms: Vec::new(), + exemplars: Vec::new(), + metadata: Some(Metadata { + r#type: metric_type, + help_ref: 0, + unit_ref, + }), + } + }) + .collect(); + Request { + symbols, + timeseries, + } + } + + type DecodedIndex = std::collections::BTreeMap< + String, + std::collections::BTreeMap>, + >; + + fn decoded_index(request: Request) -> DecodedIndex { + let encoded = into_write_requests(request) + .unwrap() + .semantic_index + .encode("public") + .expect("non-empty semantic index"); + serde_json::from_str(&encoded).unwrap() + } + + #[test] + fn test_metadata_stamps_semantic_index() { + use table::requests::METADATA_QUALITY_DECLARED; + + let index = decoded_index(metadata_request(vec![ + ( + vec![(METRIC_NAME_LABEL, "http_requests_total")], + metadata::MetricType::Counter as i32, + Some("seconds"), + ), + ( + vec![(METRIC_NAME_LABEL, "queue_depth")], + metadata::MetricType::Gauge as i32, + // Outside the OpenMetrics base set: dropped, not passed through. + Some("requests"), + ), + ])); + + let tables = &index["public"]; + let typed = &tables["http_requests_total"]; + assert_eq!(typed[SEMANTIC_METRIC_TYPE], "counter"); + assert_eq!( + typed[SEMANTIC_METRIC_METADATA_QUALITY], + METADATA_QUALITY_DECLARED + ); + assert_eq!(typed[SEMANTIC_METRIC_UNIT], "s"); + + let unitless = &tables["queue_depth"]; + assert_eq!(unitless[SEMANTIC_METRIC_TYPE], "gauge"); + assert!(!unitless.contains_key(SEMANTIC_METRIC_UNIT)); + } + + #[test] + fn test_metadata_unspecified_stamps_unit_but_not_type() { + let index = decoded_index(metadata_request(vec![( + vec![(METRIC_NAME_LABEL, "untyped_total")], + metadata::MetricType::Unspecified as i32, + Some("seconds"), + )])); + let untyped = &index["public"]["untyped_total"]; + assert_eq!(untyped[SEMANTIC_METRIC_UNIT], "s"); + assert!(!untyped.contains_key(SEMANTIC_METRIC_TYPE)); + assert!(!untyped.contains_key(SEMANTIC_METRIC_METADATA_QUALITY)); + + let requests = into_write_requests(metadata_request(vec![( + vec![(METRIC_NAME_LABEL, "untyped_unitless_total")], + metadata::MetricType::Unspecified as i32, + None, + )])) + .unwrap(); + assert!(requests.semantic_index.is_empty()); + + // No metadata at all behaves the same. + let requests = into_write_requests(test_util::request_with_labels_and_samples( + vec![(METRIC_NAME_LABEL, "bare_total")], + vec![Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + )) + .unwrap(); + assert!(requests.semantic_index.is_empty()); + } + + #[test] + fn test_metadata_type_conflict_collapses_to_mixed() { + let index = decoded_index(metadata_request(vec![ + ( + vec![(METRIC_NAME_LABEL, "flappy_metric"), ("job", "a")], + metadata::MetricType::Counter as i32, + None, + ), + ( + vec![(METRIC_NAME_LABEL, "flappy_metric"), ("job", "b")], + metadata::MetricType::Gauge as i32, + None, + ), + ])); + assert_eq!( + index["public"]["flappy_metric"][SEMANTIC_METRIC_TYPE], + "mixed" + ); + } + + #[test] + fn test_metadata_schema_overrides_stay_apart() { + // The same metric name written into two schemas by one request must not + // collapse each other's metadata. + let index = decoded_index(metadata_request(vec![ + ( + vec![(METRIC_NAME_LABEL, "cpu_usage")], + metadata::MetricType::Counter as i32, + None, + ), + ( + vec![ + (METRIC_NAME_LABEL, "cpu_usage"), + (DATABASE_LABEL, "tenant_b"), + ], + metadata::MetricType::Gauge as i32, + None, + ), + ])); + assert_eq!( + index["public"]["cpu_usage"][SEMANTIC_METRIC_TYPE], + "counter" + ); + assert_eq!( + index["tenant_b"]["cpu_usage"][SEMANTIC_METRIC_TYPE], + "gauge" + ); + } + + #[test] + fn test_metadata_out_of_range_refs_are_rejected() { + let mut request = metadata_request(vec![( + vec![(METRIC_NAME_LABEL, "broken_total")], + metadata::MetricType::Counter as i32, + None, + )]); + request.timeseries[0].metadata.as_mut().unwrap().unit_ref = 999; + let err = into_write_requests(request).err().unwrap(); + assert!(err.to_string().contains("out of range"), "{err}"); + + // unit_ref must be validated even when the type is UNSPECIFIED + // (which persists nothing). + let mut request = metadata_request(vec![( + vec![(METRIC_NAME_LABEL, "broken_total")], + metadata::MetricType::Unspecified as i32, + None, + )]); + request.timeseries[0].metadata.as_mut().unwrap().unit_ref = 999; + let err = into_write_requests(request).err().unwrap(); + assert!(err.to_string().contains("out of range"), "{err}"); + + // help_ref is validated although help is never persisted. + let mut request = metadata_request(vec![( + vec![(METRIC_NAME_LABEL, "broken_total")], + metadata::MetricType::Counter as i32, + None, + )]); + request.timeseries[0].metadata.as_mut().unwrap().help_ref = 999; + let err = into_write_requests(request).err().unwrap(); + assert!(err.to_string().contains("out of range"), "{err}"); + } } diff --git a/src/servers/src/semantic.rs b/src/servers/src/semantic.rs new file mode 100644 index 00000000000..63501dffc31 --- /dev/null +++ b/src/servers/src/semantic.rs @@ -0,0 +1,325 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Per-table semantic metadata accumulated during one ingest pass (OTLP +//! metrics, Prometheus remote write v2). +//! +//! Each written table collects the scalar semantic keys its wire metadata +//! declares (an OTLP histogram fans out into `_bucket`/`_sum`/`_count` +//! companions; a remote write series names one metric family). The resulting +//! index is serialized onto the `greptime.internal.semantic.per_table_index` +//! context extension as `{schema -> {table -> {key -> value}}}` and folded into +//! each table's options at auto-create time. +//! +//! Conflict handling follows the RFC: when two sources disagree on a +//! single-valued key the value collapses to `mixed` (or `unknown` for keys whose +//! domain has no `mixed`). + +use std::collections::{BTreeMap, HashMap}; + +use table::requests::{SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, validate_semantic_option}; + +// `greptime.semantic.metric.type` values stamped per emitted table. Must stay +// within the domain accepted by `validate_semantic_option`; the drift-guard test +// asserts this. +pub const METRIC_TYPE_COUNTER: &str = "counter"; +pub const METRIC_TYPE_UPDOWN_COUNTER: &str = "updown_counter"; +pub const METRIC_TYPE_GAUGE: &str = "gauge"; +pub const METRIC_TYPE_HISTOGRAM: &str = "histogram"; +pub const METRIC_TYPE_GAUGE_HISTOGRAM: &str = "gauge_histogram"; +pub const METRIC_TYPE_SUMMARY: &str = "summary"; +pub const METRIC_TYPE_INFO: &str = "info"; +pub const METRIC_TYPE_STATESET: &str = "stateset"; + +/// Maps an OpenMetrics unit name (open vocabulary full words: `seconds`, +/// `bytes`) to the UCUM code the `greptime.semantic.metric.unit` option is +/// defined in. Only the OpenMetrics base units are mapped; anything else is +/// dropped — a missing unit beats a corrupted cross-protocol vocabulary. +pub fn openmetrics_unit_to_ucum(unit: &str) -> Option<&'static str> { + Some(match unit { + "seconds" => "s", + "celsius" => "Cel", + "meters" => "m", + "bytes" => "By", + "ratios" => "1", + "volts" => "V", + "amperes" => "A", + "joules" => "J", + "grams" => "g", + _ => return None, + }) +} + +/// Index of `{table_name -> {semantic_key -> value}}` for one target schema. +#[derive(Debug, Default)] +pub struct SemanticIndex { + /// Per-table scalar keys; conflicting values collapse to `mixed`/`unknown`. + tables: HashMap>, +} + +impl SemanticIndex { + pub fn is_empty(&self) -> bool { + self.tables.is_empty() + } + + /// Records a scalar semantic key for `table`. A value conflicting with one + /// already recorded collapses the key to `mixed`/`unknown`; once collapsed + /// it stays collapsed. + pub fn record_scalar(&mut self, table: &str, key: &'static str, value: &str) { + // Avoid allocating the table name (and an empty map) on the common path + // where the table is already present. + if let Some(scalars) = self.tables.get_mut(table) { + match scalars.get(key).map(String::as_str) { + Some(existing) if existing == value => {} + Some(SEMANTIC_VALUE_MIXED) | Some(SEMANTIC_VALUE_UNKNOWN) => {} + Some(_) => { + scalars.insert(key, collapse_value(key)); + } + None => { + scalars.insert(key, value.to_string()); + } + } + } else { + self.tables.insert( + table.to_string(), + BTreeMap::from([(key, value.to_string())]), + ); + } + } + + /// Serializes to the JSON `{schema -> {table -> {key -> value}}}` carried + /// on the context extension, with every table under `schema`. `None` when + /// nothing was recorded. + pub fn encode(&self, schema: &str) -> Option { + if self.tables.is_empty() { + return None; + } + serde_json::to_string(&BTreeMap::from([(schema, &self.tables)])).ok() + } + + fn merge_from(&mut self, other: &SemanticIndex) { + for (table, scalars) in &other.tables { + for (key, value) in scalars { + self.record_scalar(table, key, value); + } + } + } + + #[cfg(test)] + fn options_of(&self, table: &str) -> Option<&BTreeMap<&'static str, String>> { + self.tables.get(table) + } +} + +/// Schema-aware collection of [`SemanticIndex`]es: Prometheus remote write lets +/// each series override the target schema with a special label, so one request +/// may write the same metric name into several schemas — their metadata must +/// not collapse into each other. +#[derive(Debug, Default)] +pub struct SemanticIndexes { + /// Tables written into the request's default schema (no override). + default: SemanticIndex, + /// Tables written under a per-series schema override. + overrides: HashMap, +} + +impl SemanticIndexes { + pub fn is_empty(&self) -> bool { + self.default.is_empty() && self.overrides.values().all(SemanticIndex::is_empty) + } + + /// The index for `schema` (`None` = the request's default schema). + pub fn index_for(&mut self, schema: Option<&str>) -> &mut SemanticIndex { + match schema { + None => &mut self.default, + Some(schema) => { + if !self.overrides.contains_key(schema) { + self.overrides + .insert(schema.to_string(), SemanticIndex::default()); + } + self.overrides.get_mut(schema).expect("just inserted") + } + } + } + + /// Serializes to the JSON `{schema -> {table -> {key -> value}}}` carried on + /// the context extension, resolving the default index to `default_schema`. + /// An override explicitly naming `default_schema` merges into the default + /// with the usual conflict collapse. + pub fn encode(&self, default_schema: &str) -> Option { + if self.is_empty() { + return None; + } + let mut by_schema: BTreeMap<&str, &HashMap>> = + BTreeMap::new(); + let mut merged_default; + if let Some(aliased) = self.overrides.get(default_schema) { + merged_default = SemanticIndex::default(); + merged_default.merge_from(&self.default); + merged_default.merge_from(aliased); + by_schema.insert(default_schema, &merged_default.tables); + } else if !self.default.is_empty() { + by_schema.insert(default_schema, &self.default.tables); + } + for (schema, index) in &self.overrides { + if schema != default_schema && !index.is_empty() { + by_schema.insert(schema, &index.tables); + } + } + serde_json::to_string(&by_schema).ok() + } +} + +/// The collapsed value for a conflicting scalar key: `mixed` when the key's +/// domain accepts it, else `unknown`. Uses the vocabulary validator as the +/// single source of truth for which keys allow `mixed`. +fn collapse_value(key: &str) -> String { + if validate_semantic_option(key, SEMANTIC_VALUE_MIXED) { + SEMANTIC_VALUE_MIXED.to_string() + } else { + SEMANTIC_VALUE_UNKNOWN.to_string() + } +} + +#[cfg(test)] +mod tests { + use table::requests::{ + SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, + }; + + use super::*; + + type Decoded = BTreeMap>>; + + #[test] + fn test_scalar_recording_keeps_first_then_collapses_on_conflict() { + let mut index = SemanticIndex::default(); + index.record_scalar("t", SEMANTIC_METRIC_TYPE, "counter"); + index.record_scalar("t", SEMANTIC_METRIC_TYPE, "counter"); + assert_eq!( + index + .options_of("t") + .unwrap() + .get(SEMANTIC_METRIC_TYPE) + .map(String::as_str), + Some("counter") + ); + + // Conflict on a key whose domain has `mixed` collapses to `mixed`. + index.record_scalar("t", SEMANTIC_METRIC_TYPE, "gauge"); + assert_eq!( + index + .options_of("t") + .unwrap() + .get(SEMANTIC_METRIC_TYPE) + .map(String::as_str), + Some("mixed") + ); + // Further writes stay collapsed. + index.record_scalar("t", SEMANTIC_METRIC_TYPE, "histogram"); + assert_eq!( + index + .options_of("t") + .unwrap() + .get(SEMANTIC_METRIC_TYPE) + .map(String::as_str), + Some("mixed") + ); + } + + #[test] + fn test_scalar_conflict_without_mixed_domain_collapses_to_unknown() { + let mut index = SemanticIndex::default(); + index.record_scalar("t", SEMANTIC_METRIC_METADATA_QUALITY, "declared"); + index.record_scalar("t", SEMANTIC_METRIC_METADATA_QUALITY, "inferred"); + // metadata_quality accepts only declared/inferred/unknown, so a conflict + // is `unknown`. + assert_eq!( + index + .options_of("t") + .unwrap() + .get(SEMANTIC_METRIC_METADATA_QUALITY) + .map(String::as_str), + Some("unknown") + ); + } + + #[test] + fn test_encode_is_none_when_empty_and_round_trips() { + let index = SemanticIndex::default(); + assert!(index.is_empty()); + assert_eq!(index.encode("public"), None); + + let mut index = SemanticIndex::default(); + index.record_scalar("metric_a", SEMANTIC_METRIC_TYPE, "counter"); + index.record_scalar("metric_a", SEMANTIC_METRIC_UNIT, "By"); + let json = index.encode("public").unwrap(); + let parsed: Decoded = serde_json::from_str(&json).unwrap(); + let table = parsed.get("public").unwrap().get("metric_a").unwrap(); + assert_eq!( + table.get(SEMANTIC_METRIC_TYPE).map(String::as_str), + Some("counter") + ); + assert_eq!( + table.get(SEMANTIC_METRIC_UNIT).map(String::as_str), + Some("By") + ); + } + + #[test] + fn test_indexes_keep_schemas_apart_and_merge_default_alias() { + let mut indexes = SemanticIndexes::default(); + assert!(indexes.is_empty()); + assert_eq!(indexes.encode("public"), None); + + // The same metric name in two schemas must not collapse to `mixed`. + indexes + .index_for(None) + .record_scalar("cpu_usage", SEMANTIC_METRIC_TYPE, "counter"); + indexes.index_for(Some("tenant_b")).record_scalar( + "cpu_usage", + SEMANTIC_METRIC_TYPE, + "gauge", + ); + let parsed: Decoded = serde_json::from_str(&indexes.encode("public").unwrap()).unwrap(); + assert_eq!( + parsed["public"]["cpu_usage"][SEMANTIC_METRIC_TYPE], + "counter" + ); + assert_eq!( + parsed["tenant_b"]["cpu_usage"][SEMANTIC_METRIC_TYPE], + "gauge" + ); + + // An override naming the default schema merges into it — and a real + // conflict then collapses. + indexes + .index_for(Some("public")) + .record_scalar("cpu_usage", SEMANTIC_METRIC_TYPE, "gauge"); + let parsed: Decoded = serde_json::from_str(&indexes.encode("public").unwrap()).unwrap(); + assert_eq!(parsed["public"]["cpu_usage"][SEMANTIC_METRIC_TYPE], "mixed"); + } + + #[test] + fn test_openmetrics_unit_mapping() { + assert_eq!(openmetrics_unit_to_ucum("seconds"), Some("s")); + assert_eq!(openmetrics_unit_to_ucum("bytes"), Some("By")); + assert_eq!(openmetrics_unit_to_ucum("ratios"), Some("1")); + // Outside the OpenMetrics base set: dropped, not passed through. + assert_eq!(openmetrics_unit_to_ucum("requests"), None); + assert_eq!(openmetrics_unit_to_ucum(""), None); + // No fuzzy matching: UCUM codes are not OpenMetrics names. + assert_eq!(openmetrics_unit_to_ucum("By"), None); + } +} diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index aa60109a616..9d2f1eb8cb5 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -18,8 +18,10 @@ use std::str::FromStr; use std::time::Duration; use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; +use api::greptime_proto::io::prometheus::write::v2::metadata::MetricType as RemoteWriteV2MetricType; use api::greptime_proto::io::prometheus::write::v2::{ - BucketSpan, Histogram, Sample as RemoteWriteV2Sample, TimeSeries as RemoteWriteV2TimeSeries, + BucketSpan, Histogram, Metadata as RemoteWriteV2Metadata, Sample as RemoteWriteV2Sample, + TimeSeries as RemoteWriteV2TimeSeries, }; use api::prom_store::remote::label_matcher::Type as MatcherType; use api::prom_store::remote::{ @@ -2693,6 +2695,52 @@ pub async fn test_prometheus_remote_write_v2(store_type: StorageType) { ) .await; + // A series carrying inline metadata (v2 senders SHOULD) upgrades its + // table's quality to declared, with the metric type and the unit + // canonicalised from the OpenMetrics word to UCUM. + let mut write_request = remote_write_v2::request_with_labels_and_samples( + vec![ + (prom_store::METRIC_NAME_LABEL, "remote_write_v2_typed_total"), + ("job", "api"), + ], + vec![RemoteWriteV2Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + ); + let unit_ref = write_request.symbols.len() as u32; + write_request.symbols.push("seconds".to_string()); + write_request.timeseries[0].metadata = Some(RemoteWriteV2Metadata { + r#type: RemoteWriteV2MetricType::Counter as i32, + help_ref: 0, + unit_ref, + }); + let compressed_request = prom_store::snappy_compress(&write_request.encode_to_vec()) + .expect("failed to encode snappy"); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .body(compressed_request) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + validate_data( + "prometheus_remote_write_v2_declared_metadata", + &client, + "select count(*) from information_schema.tables where table_name = 'remote_write_v2_typed_total' \ + and create_options like '%greptime.semantic.metric.type=counter%' \ + and create_options like '%greptime.semantic.metric.unit=s%' \ + and create_options like '%greptime.semantic.metric.metadata_quality=declared%';", + "[[1]]", + ) + .await; + guard.remove_all().await; } @@ -2892,6 +2940,53 @@ pub async fn test_prometheus_remote_write_batched(store_type: StorageType) { ) .await; + // The batched path bypasses the operator's auto-create: v2 per-series + // metadata must reach its table options through the batch create too. + let mut write_request = remote_write_v2::request_with_labels_and_samples( + vec![ + ( + prom_store::METRIC_NAME_LABEL, + "remote_write_batched_typed_total", + ), + ("job", "api"), + ], + vec![RemoteWriteV2Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + ); + let unit_ref = write_request.symbols.len() as u32; + write_request.symbols.push("bytes".to_string()); + write_request.timeseries[0].metadata = Some(RemoteWriteV2Metadata { + r#type: RemoteWriteV2MetricType::Gauge as i32, + help_ref: 0, + unit_ref, + }); + let compressed_request = prom_store::snappy_compress(&write_request.encode_to_vec()) + .expect("failed to encode snappy"); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .body(compressed_request) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + wait_for_data( + &client, + "select count(*) from information_schema.tables where table_name = 'remote_write_batched_typed_total' \ + and create_options like '%greptime.semantic.metric.type=gauge%' \ + and create_options like '%greptime.semantic.metric.unit=By%' \ + and create_options like '%greptime.semantic.metric.metadata_quality=declared%'", + "[[1]]", + ) + .await; + guard.remove_all().await; } From 32e215cad0aa0803991d751bc7b4f44ac1865e62 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 11 Aug 2026 15:12:45 +0800 Subject: [PATCH 043/126] fix(event): preserve procedure lifecycle locators (#8787) * fix(event): preserve procedure lifecycle locators Signed-off-by: WenyXu * fix(event): preserve dropped table lifecycle locators Signed-off-by: WenyXu * test(event): cover lifecycle locators Signed-off-by: WenyXu * test(event): fix lifecycle context expectations Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/common/meta/src/ddl/alter_database.rs | 2 +- .../meta/src/ddl/alter_logical_tables.rs | 13 +- src/common/meta/src/ddl/alter_table.rs | 9 +- src/common/meta/src/ddl/create_database.rs | 2 +- src/common/meta/src/ddl/create_flow.rs | 11 +- .../meta/src/ddl/create_logical_tables.rs | 66 +++++----- src/common/meta/src/ddl/create_table.rs | 17 +-- src/common/meta/src/ddl/create_view.rs | 29 ++++- src/common/meta/src/ddl/drop_database.rs | 2 +- src/common/meta/src/ddl/drop_flow.rs | 6 +- src/common/meta/src/ddl/drop_table.rs | 20 ++- src/common/meta/src/ddl/drop_view.rs | 10 +- src/common/meta/src/ddl/event/database.rs | 18 +-- src/common/meta/src/ddl/event/flow.rs | 31 +++-- src/common/meta/src/ddl/event/table.rs | 31 +++-- src/common/meta/src/ddl/event/view.rs | 64 +++++++--- .../meta/src/ddl/purge_dropped_table.rs | 29 +++-- .../meta/src/ddl/tests/event/database.rs | 29 +++-- src/common/meta/src/ddl/tests/event/flow.rs | 43 +++++-- src/common/meta/src/ddl/tests/event/table.rs | 44 +++++-- src/common/meta/src/ddl/tests/event/view.rs | 114 +++++++++++++----- src/common/meta/src/ddl/truncate_table.rs | 20 ++- src/common/meta/src/ddl/undrop_table.rs | 33 ++--- src/meta-srv/src/event/repartition.rs | 54 +++++---- src/meta-srv/src/procedure/repartition.rs | 2 +- .../src/procedure/repartition/group.rs | 2 +- tests-integration/tests/database_ddl_event.rs | 4 +- tests-integration/tests/flow_ddl_event.rs | 10 +- tests-integration/tests/repartition_event.rs | 49 ++++++++ tests-integration/tests/table_ddl_event.rs | 99 +++++++-------- tests-integration/tests/view_ddl_event.rs | 14 +-- 31 files changed, 562 insertions(+), 315 deletions(-) diff --git a/src/common/meta/src/ddl/alter_database.rs b/src/common/meta/src/ddl/alter_database.rs index 7e16275da3c..f9344213025 100644 --- a/src/common/meta/src/ddl/alter_database.rs +++ b/src/common/meta/src/ddl/alter_database.rs @@ -196,7 +196,7 @@ impl Procedure for AlterDatabaseProcedure { self.data.event_context.clone(), ) } else { - DatabaseDdlEvent::alter_lifecycle() + DatabaseDdlEvent::alter_lifecycle(self.data.catalog(), self.data.schema()) }; Some(Box::new(event)) } diff --git a/src/common/meta/src/ddl/alter_logical_tables.rs b/src/common/meta/src/ddl/alter_logical_tables.rs index 816db0a6afb..f2f15430933 100644 --- a/src/common/meta/src/ddl/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/alter_logical_tables.rs @@ -335,6 +335,11 @@ impl Procedure for AlterLogicalTablesProcedure { if ctx.trigger != EventTrigger::Submitted { return Some(Box::new(TableDdlEvent::lifecycle( TableDdlEventType::AlterLogicalTables, + self.data.tasks.iter().map(|task| { + let table_ref = task.table_ref(); + TableDdlLocator::new(table_ref.catalog, table_ref.schema, table_ref.table) + .with_physical_table_id(self.data.physical_table_id) + }), ))); } @@ -375,12 +380,12 @@ pub struct AlterTablesData { } impl AlterTablesData { - /// Clears all data fields except `state` and `table_cache_keys_to_invalidate` after metadata update. - /// This is done to avoid persisting unnecessary data after the update metadata step. + /// Clears metadata snapshots after the update metadata step. + /// + /// Keep the tasks and physical table ID until the procedure finishes: lifecycle + /// events use them to retain the logical table locator. fn clear_metadata_fields(&mut self) { - self.tasks.clear(); self.table_info_values.clear(); - self.physical_table_id = 0; self.physical_table_info = None; self.physical_columns.clear(); } diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index acc02c4d1a8..21975fab714 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -579,12 +579,11 @@ impl Procedure for AlterTableProcedure { { return None; } + let table_ref = self.data.table_ref(); + let locator = TableDdlLocator::new(table_ref.catalog, table_ref.schema, table_ref.table) + .with_table_id(self.data.table_id()); let event = match &ctx.trigger { EventTrigger::Submitted => { - let table_ref = self.data.table_ref(); - let locator = - TableDdlLocator::new(table_ref.catalog, table_ref.schema, table_ref.table) - .with_table_id(self.data.table_id()); let kind = self .data .task @@ -594,7 +593,7 @@ impl Procedure for AlterTableProcedure { .and_then(alter_table_kind_name); TableDdlEvent::alter_table_submitted(locator, kind, self.data.event_context.clone()) } - _ => TableDdlEvent::lifecycle(TableDdlEventType::AlterTable), + _ => TableDdlEvent::lifecycle(TableDdlEventType::AlterTable, [locator]), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/create_database.rs b/src/common/meta/src/ddl/create_database.rs index 8f9c5bee1a5..b48e700a519 100644 --- a/src/common/meta/src/ddl/create_database.rs +++ b/src/common/meta/src/ddl/create_database.rs @@ -297,7 +297,7 @@ impl Procedure for CreateDatabaseProcedure { self.data.event_context.clone(), ) } else { - DatabaseDdlEvent::create_lifecycle() + DatabaseDdlEvent::create_lifecycle(&self.data.catalog, &self.data.schema) }; Some(Box::new(event)) } diff --git a/src/common/meta/src/ddl/create_flow.rs b/src/common/meta/src/ddl/create_flow.rs index 66083f5949c..3ef4dadbc63 100644 --- a/src/common/meta/src/ddl/create_flow.rs +++ b/src/common/meta/src/ddl/create_flow.rs @@ -440,9 +440,16 @@ impl Procedure for CreateFlowProcedure { .or(self.data.flow_id), _ => self.data.flow_id, }; - FlowDdlEvent::create_succeeded(flow_id) + FlowDdlEvent::create_succeeded( + &self.data.task.catalog_name, + &self.data.task.flow_name, + flow_id, + ) } - _ => FlowDdlEvent::create_lifecycle(), + _ => FlowDdlEvent::create_lifecycle( + &self.data.task.catalog_name, + &self.data.task.flow_name, + ), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/create_logical_tables.rs b/src/common/meta/src/ddl/create_logical_tables.rs index 8b93ecfcc0f..1c9efb865e4 100644 --- a/src/common/meta/src/ddl/create_logical_tables.rs +++ b/src/common/meta/src/ddl/create_logical_tables.rs @@ -212,6 +212,19 @@ impl CreateLogicalTablesProcedure { } } +impl CreateLogicalTablesProcedure { + fn event_locators(&self) -> impl Iterator + '_ { + self.data.tasks.iter().map(|task| { + TableDdlLocator::new( + &task.create_table.catalog_name, + &task.create_table.schema_name, + &task.create_table.table_name, + ) + .with_physical_table_id(self.data.physical_table_id) + }) + } +} + #[async_trait] impl Procedure for CreateLogicalTablesProcedure { fn type_name(&self) -> &str { @@ -268,21 +281,11 @@ impl Procedure for CreateLogicalTablesProcedure { return None; } let event = match &ctx.trigger { - EventTrigger::Submitted => { - let locators = self.data.tasks.iter().map(|task| { - TableDdlLocator::new( - &task.create_table.catalog_name, - &task.create_table.schema_name, - &task.create_table.table_name, - ) - .with_physical_table_id(self.data.physical_table_id) - }); - TableDdlEvent::create_logical_tables_submitted( - locators, - self.data.tasks.len(), - self.data.event_context.clone(), - ) - } + EventTrigger::Submitted => TableDdlEvent::create_logical_tables_submitted( + self.event_locators(), + self.data.tasks.len(), + self.data.event_context.clone(), + ), EventTrigger::Succeeded => match ctx.lifecycle_state { ProcedureState::Done { output: Some(output), @@ -290,28 +293,27 @@ impl Procedure for CreateLogicalTablesProcedure { .downcast_ref::>() .map(|table_ids| { debug_assert_eq!(self.data.tasks.len(), table_ids.len()); - let locators = - self.data - .tasks - .iter() - .zip(table_ids) - .map(|(task, table_id)| { - TableDdlLocator::new( - &task.create_table.catalog_name, - &task.create_table.schema_name, - &task.create_table.table_name, - ) - .with_table_id(*table_id) - .with_physical_table_id(self.data.physical_table_id) - }); + let locators = self + .event_locators() + .zip(table_ids) + .map(|(locator, table_id)| locator.with_table_id(*table_id)); TableDdlEvent::create_logical_tables_succeeded(locators) }) .unwrap_or_else(|| { - TableDdlEvent::lifecycle(TableDdlEventType::CreateLogicalTables) + TableDdlEvent::lifecycle( + TableDdlEventType::CreateLogicalTables, + self.event_locators(), + ) }), - _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateLogicalTables), + _ => TableDdlEvent::lifecycle( + TableDdlEventType::CreateLogicalTables, + self.event_locators(), + ), }, - _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateLogicalTables), + _ => TableDdlEvent::lifecycle( + TableDdlEventType::CreateLogicalTables, + self.event_locators(), + ), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/create_table.rs b/src/common/meta/src/ddl/create_table.rs index c1690416a2e..c63d7f7a34d 100644 --- a/src/common/meta/src/ddl/create_table.rs +++ b/src/common/meta/src/ddl/create_table.rs @@ -406,11 +406,10 @@ impl Procedure for CreateTableProcedure { { return None; } + let table_ref = self.data.table_ref(); + let locator = TableDdlLocator::new(table_ref.catalog, table_ref.schema, table_ref.table); let event = match &ctx.trigger { EventTrigger::Submitted => { - let table_ref = self.data.table_ref(); - let locator = - TableDdlLocator::new(table_ref.catalog, table_ref.schema, table_ref.table); let create_table = &self.data.task.create_table; TableDdlEvent::create_table_submitted( locator, @@ -425,11 +424,15 @@ impl Procedure for CreateTableProcedure { } => output .downcast_ref::() .copied() - .map(TableDdlEvent::create_table_succeeded) - .unwrap_or_else(|| TableDdlEvent::lifecycle(TableDdlEventType::CreateTable)), - _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateTable), + .map(|table_id| { + TableDdlEvent::create_table_succeeded(locator.clone(), table_id) + }) + .unwrap_or_else(|| { + TableDdlEvent::lifecycle(TableDdlEventType::CreateTable, [locator.clone()]) + }), + _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateTable, [locator.clone()]), }, - _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateTable), + _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateTable, [locator]), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/create_view.rs b/src/common/meta/src/ddl/create_view.rs index 0cb622774f0..a7175854b76 100644 --- a/src/common/meta/src/ddl/create_view.rs +++ b/src/common/meta/src/ddl/create_view.rs @@ -297,12 +297,33 @@ impl Procedure for CreateViewProcedure { ProcedureState::Done { output: Some(output), } => output.downcast_ref::().copied().map_or_else( - ViewDdlEvent::create_lifecycle, - ViewDdlEvent::create_succeeded, + || { + ViewDdlEvent::create_lifecycle( + &self.data.task.create_view.catalog_name, + &self.data.task.create_view.schema_name, + &self.data.task.create_view.view_name, + ) + }, + |view_id| { + ViewDdlEvent::create_succeeded( + &self.data.task.create_view.catalog_name, + &self.data.task.create_view.schema_name, + &self.data.task.create_view.view_name, + view_id, + ) + }, + ), + _ => ViewDdlEvent::create_lifecycle( + &self.data.task.create_view.catalog_name, + &self.data.task.create_view.schema_name, + &self.data.task.create_view.view_name, ), - _ => ViewDdlEvent::create_lifecycle(), }, - _ => ViewDdlEvent::create_lifecycle(), + _ => ViewDdlEvent::create_lifecycle( + &self.data.task.create_view.catalog_name, + &self.data.task.create_view.schema_name, + &self.data.task.create_view.view_name, + ), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/drop_database.rs b/src/common/meta/src/ddl/drop_database.rs index d84a5e6dc86..7772e966cb2 100644 --- a/src/common/meta/src/ddl/drop_database.rs +++ b/src/common/meta/src/ddl/drop_database.rs @@ -202,7 +202,7 @@ impl Procedure for DropDatabaseProcedure { self.event_context.clone(), ) } else { - DatabaseDdlEvent::drop_lifecycle() + DatabaseDdlEvent::drop_lifecycle(&self.context.catalog, &self.context.schema) }; Some(Box::new(event)) } diff --git a/src/common/meta/src/ddl/drop_flow.rs b/src/common/meta/src/ddl/drop_flow.rs index badded386ad..99c18ebc13b 100644 --- a/src/common/meta/src/ddl/drop_flow.rs +++ b/src/common/meta/src/ddl/drop_flow.rs @@ -239,7 +239,11 @@ impl Procedure for DropFlowProcedure { self.data.task.drop_if_exists, self.data.event_context.clone(), ), - _ => FlowDdlEvent::drop_lifecycle(), + _ => FlowDdlEvent::drop_lifecycle( + &self.data.task.catalog_name, + &self.data.task.flow_name, + self.data.task.flow_id, + ), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs index 2dcfa2f7d82..7c8dec55e49 100644 --- a/src/common/meta/src/ddl/drop_table.rs +++ b/src/common/meta/src/ddl/drop_table.rs @@ -421,18 +421,16 @@ impl Procedure for DropTableProcedure { { return None; } + let task = &self.data.task; + let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) + .with_table_id(task.table_id); let event = match &ctx.trigger { - EventTrigger::Submitted => { - let task = &self.data.task; - let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) - .with_table_id(task.table_id); - TableDdlEvent::drop_table_submitted( - locator, - task.drop_if_exists, - self.data.event_context.clone(), - ) - } - _ => TableDdlEvent::lifecycle(TableDdlEventType::DropTable), + EventTrigger::Submitted => TableDdlEvent::drop_table_submitted( + locator, + task.drop_if_exists, + self.data.event_context.clone(), + ), + _ => TableDdlEvent::lifecycle(TableDdlEventType::DropTable, [locator]), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/drop_view.rs b/src/common/meta/src/ddl/drop_view.rs index 53a62e76159..ece1d2b06b7 100644 --- a/src/common/meta/src/ddl/drop_view.rs +++ b/src/common/meta/src/ddl/drop_view.rs @@ -231,7 +231,15 @@ impl Procedure for DropViewProcedure { self.data.event_context.clone(), ) } - _ => ViewDdlEvent::drop_lifecycle(), + _ => { + let table_ref = self.data.table_ref(); + ViewDdlEvent::drop_lifecycle( + table_ref.catalog, + table_ref.schema, + table_ref.table, + self.data.view_id(), + ) + } }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/event/database.rs b/src/common/meta/src/ddl/event/database.rs index 9d70ea35a5e..b24611799c6 100644 --- a/src/common/meta/src/ddl/event/database.rs +++ b/src/common/meta/src/ddl/event/database.rs @@ -179,16 +179,16 @@ impl DatabaseDdlEvent { ) } - pub(crate) fn create_lifecycle() -> Self { - Self::lifecycle(CREATE_DATABASE_EVENT_TYPE) + pub(crate) fn create_lifecycle(catalog_name: &str, schema_name: &str) -> Self { + Self::lifecycle(CREATE_DATABASE_EVENT_TYPE, catalog_name, schema_name) } - pub(crate) fn alter_lifecycle() -> Self { - Self::lifecycle(ALTER_DATABASE_EVENT_TYPE) + pub(crate) fn alter_lifecycle(catalog_name: &str, schema_name: &str) -> Self { + Self::lifecycle(ALTER_DATABASE_EVENT_TYPE, catalog_name, schema_name) } - pub(crate) fn drop_lifecycle() -> Self { - Self::lifecycle(DROP_DATABASE_EVENT_TYPE) + pub(crate) fn drop_lifecycle(catalog_name: &str, schema_name: &str) -> Self { + Self::lifecycle(DROP_DATABASE_EVENT_TYPE, catalog_name, schema_name) } fn submitted( @@ -207,11 +207,11 @@ impl DatabaseDdlEvent { } } - fn lifecycle(event_type: &'static str) -> Self { + fn lifecycle(event_type: &'static str, catalog_name: &str, schema_name: &str) -> Self { Self { event_type, - catalog_name: None, - schema_name: None, + catalog_name: Some(catalog_name.to_string()), + schema_name: Some(schema_name.to_string()), payload: None, event_context: None, } diff --git a/src/common/meta/src/ddl/event/flow.rs b/src/common/meta/src/ddl/event/flow.rs index 8a52486ddfe..7d0d6cb539b 100644 --- a/src/common/meta/src/ddl/event/flow.rs +++ b/src/common/meta/src/ddl/event/flow.rs @@ -117,29 +117,36 @@ impl FlowDdlEvent { } } - /// Builds a lightweight Create Flow lifecycle event. - pub(crate) fn create_lifecycle() -> Self { - Self::lifecycle(CREATE_FLOW_EVENT_TYPE) + /// Builds a Create Flow lifecycle event with its submitted locator. + pub(crate) fn create_lifecycle(catalog_name: &str, flow_name: &str) -> Self { + Self::lifecycle(CREATE_FLOW_EVENT_TYPE, catalog_name, flow_name) } - /// Builds a successful Create Flow event containing only a resolved ID. - pub(crate) fn create_succeeded(flow_id: Option) -> Self { + /// Builds a successful Create Flow event with its submitted locator and resolved ID. + pub(crate) fn create_succeeded( + catalog_name: &str, + flow_name: &str, + flow_id: Option, + ) -> Self { Self { flow_id, - ..Self::lifecycle(CREATE_FLOW_EVENT_TYPE) + ..Self::lifecycle(CREATE_FLOW_EVENT_TYPE, catalog_name, flow_name) } } - /// Builds a lightweight Drop Flow lifecycle event. - pub(crate) fn drop_lifecycle() -> Self { - Self::lifecycle(DROP_FLOW_EVENT_TYPE) + /// Builds a Drop Flow lifecycle event with its submitted locator. + pub(crate) fn drop_lifecycle(catalog_name: &str, flow_name: &str, flow_id: u32) -> Self { + Self { + flow_id: Some(flow_id), + ..Self::lifecycle(DROP_FLOW_EVENT_TYPE, catalog_name, flow_name) + } } - fn lifecycle(event_type: &'static str) -> Self { + fn lifecycle(event_type: &'static str, catalog_name: &str, flow_name: &str) -> Self { Self { event_type, - catalog_name: None, - flow_name: None, + catalog_name: Some(catalog_name.to_string()), + flow_name: Some(flow_name.to_string()), flow_id: None, payload: None, event_context: None, diff --git a/src/common/meta/src/ddl/event/table.rs b/src/common/meta/src/ddl/event/table.rs index 8608b3151a6..9ae6f530761 100644 --- a/src/common/meta/src/ddl/event/table.rs +++ b/src/common/meta/src/ddl/event/table.rs @@ -103,6 +103,7 @@ impl TableDdlLocator { } /// Creates a locator containing only a table ID. + #[cfg(feature = "enterprise")] pub(crate) fn from_table_id(table_id: TableId) -> Self { Self { table_id: Some(table_id), @@ -361,36 +362,32 @@ impl TableDdlEvent { ) } - /// Builds a lightweight lifecycle event with null domain columns and payload. - pub(crate) fn lifecycle(event_type: TableDdlEventType) -> Self { + /// Builds a lifecycle event with stable object locators and no intent payload. + pub(crate) fn lifecycle( + event_type: TableDdlEventType, + locators: impl IntoIterator, + ) -> Self { Self { event_type, - locators: vec![TableDdlLocator::default()], + locators: locators.into_iter().collect(), payload: None, event_context: None, } } - /// Builds a Create Table success event containing only the allocated table ID. - pub(crate) fn create_table_succeeded(table_id: TableId) -> Self { - Self { - event_type: TableDdlEventType::CreateTable, - locators: vec![TableDdlLocator::from_table_id(table_id)], - payload: None, - event_context: None, - } + /// Builds a Create Table success event containing the submitted locator and allocated ID. + pub(crate) fn create_table_succeeded(locator: TableDdlLocator, table_id: TableId) -> Self { + Self::lifecycle( + TableDdlEventType::CreateTable, + [locator.with_table_id(table_id)], + ) } /// Builds Create Logical Tables success rows from their allocated locators. pub(crate) fn create_logical_tables_succeeded( locators: impl IntoIterator, ) -> Self { - Self { - event_type: TableDdlEventType::CreateLogicalTables, - locators: locators.into_iter().collect(), - payload: None, - event_context: None, - } + Self::lifecycle(TableDdlEventType::CreateLogicalTables, locators) } fn submitted( diff --git a/src/common/meta/src/ddl/event/view.rs b/src/common/meta/src/ddl/event/view.rs index d6dcf2c5f17..c1d5846fcf2 100644 --- a/src/common/meta/src/ddl/event/view.rs +++ b/src/common/meta/src/ddl/event/view.rs @@ -121,19 +121,38 @@ impl ViewDdlEvent { ) } - /// Builds a lightweight create-view lifecycle event with no locator data. - pub(crate) fn create_lifecycle() -> Self { - Self::lifecycle(CREATE_VIEW_EVENT_TYPE) + /// Builds a create-view lifecycle event with its submitted locator. + pub(crate) fn create_lifecycle(catalog_name: &str, schema_name: &str, view_name: &str) -> Self { + Self::lifecycle(CREATE_VIEW_EVENT_TYPE, catalog_name, schema_name, view_name) } - /// Builds the successful create-view row that carries only the allocated id. - pub(crate) fn create_succeeded(view_id: u32) -> Self { - Self::succeeded(CREATE_VIEW_EVENT_TYPE, view_id) + /// Builds the successful create-view row with its submitted locator and allocated ID. + pub(crate) fn create_succeeded( + catalog_name: &str, + schema_name: &str, + view_name: &str, + view_id: u32, + ) -> Self { + Self::succeeded( + CREATE_VIEW_EVENT_TYPE, + catalog_name, + schema_name, + view_name, + view_id, + ) } - /// Builds a lightweight drop-view lifecycle event with no locator data. - pub(crate) fn drop_lifecycle() -> Self { - Self::lifecycle(DROP_VIEW_EVENT_TYPE) + /// Builds a drop-view lifecycle event with its submitted locator and ID. + pub(crate) fn drop_lifecycle( + catalog_name: &str, + schema_name: &str, + view_name: &str, + view_id: u32, + ) -> Self { + Self { + view_id: Some(view_id), + ..Self::lifecycle(DROP_VIEW_EVENT_TYPE, catalog_name, schema_name, view_name) + } } fn submitted( @@ -156,24 +175,35 @@ impl ViewDdlEvent { } } - fn lifecycle(event_type: &'static str) -> Self { + fn lifecycle( + event_type: &'static str, + catalog_name: &str, + schema_name: &str, + view_name: &str, + ) -> Self { Self { event_type, - catalog_name: None, - schema_name: None, - view_name: None, + catalog_name: Some(catalog_name.to_string()), + schema_name: Some(schema_name.to_string()), + view_name: Some(view_name.to_string()), view_id: None, payload: None, event_context: None, } } - fn succeeded(event_type: &'static str, view_id: u32) -> Self { + fn succeeded( + event_type: &'static str, + catalog_name: &str, + schema_name: &str, + view_name: &str, + view_id: u32, + ) -> Self { Self { event_type, - catalog_name: None, - schema_name: None, - view_name: None, + catalog_name: Some(catalog_name.to_string()), + schema_name: Some(schema_name.to_string()), + view_name: Some(view_name.to_string()), view_id: Some(view_id), payload: None, event_context: None, diff --git a/src/common/meta/src/ddl/purge_dropped_table.rs b/src/common/meta/src/ddl/purge_dropped_table.rs index 111b79cabe8..c81aa043dc2 100644 --- a/src/common/meta/src/ddl/purge_dropped_table.rs +++ b/src/common/meta/src/ddl/purge_dropped_table.rs @@ -226,6 +226,20 @@ impl PurgeDroppedTableProcedure { fn executor(&self) -> DropTableExecutor { DropTableExecutor::new(self.data.table_name().clone(), self.data.table_id(), false) } + + fn event_locator(&self) -> TableDdlLocator { + let table_id = self.data.task.table_id; + let Some(table_name) = self.data.table_name.as_ref() else { + return TableDdlLocator::from_table_id(table_id); + }; + + TableDdlLocator::new( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + ) + .with_table_id(table_id) + } } #[async_trait] @@ -274,17 +288,12 @@ impl Procedure for PurgeDroppedTableProcedure { { return None; } - let event = match &ctx.trigger { - EventTrigger::Submitted => { - let locator = TableDdlLocator::from_table_id(self.data.task.table_id); - TableDdlEvent::purge_dropped_table_submitted( - locator, - self.data.event_context.clone(), - ) - } - _ => TableDdlEvent::lifecycle(TableDdlEventType::PurgeDroppedTable), + let locator = self.event_locator(); + let event = if ctx.trigger == EventTrigger::Submitted { + TableDdlEvent::purge_dropped_table_submitted(locator, self.data.event_context.clone()) + } else { + TableDdlEvent::lifecycle(TableDdlEventType::PurgeDroppedTable, [locator]) }; - Some(Box::new(event)) } } diff --git a/src/common/meta/src/ddl/tests/event/database.rs b/src/common/meta/src/ddl/tests/event/database.rs index 8aa85c2fbbd..29cb48f7b80 100644 --- a/src/common/meta/src/ddl/tests/event/database.rs +++ b/src/common/meta/src/ddl/tests/event/database.rs @@ -156,30 +156,33 @@ fn test_drop_database_submitted_event_contract() { } #[test] -fn test_database_lifecycle_events_have_fixed_schema_and_null_intent() { +fn test_database_lifecycle_events_preserve_locator_and_null_intent() { for (event, event_type) in [ ( - DatabaseDdlEvent::create_lifecycle(), + DatabaseDdlEvent::create_lifecycle("greptime", "metrics"), CREATE_DATABASE_EVENT_TYPE, ), ( - DatabaseDdlEvent::alter_lifecycle(), + DatabaseDdlEvent::alter_lifecycle("greptime", "metrics"), ALTER_DATABASE_EVENT_TYPE, ), - (DatabaseDdlEvent::drop_lifecycle(), DROP_DATABASE_EVENT_TYPE), + ( + DatabaseDdlEvent::drop_lifecycle("greptime", "metrics"), + DROP_DATABASE_EVENT_TYPE, + ), ] { - assert_event_locator(&event, event_type, None, None); + assert_event_locator(&event, event_type, Some("greptime"), Some("metrics")); assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); } assert_eq!( - DatabaseDdlEvent::create_lifecycle().extra_schema(), + DatabaseDdlEvent::create_lifecycle("greptime", "metrics").extra_schema(), DatabaseDdlEvent::create_submitted( "c", "s", false, &HashMap::new(), - EventContext::default() + EventContext::default(), ) .extra_schema() ); @@ -202,7 +205,7 @@ fn test_database_events_preserve_procedure_envelope_contract() { ); let lifecycle = ProcedureEvent::new( procedure_id, - Box::new(DatabaseDdlEvent::create_lifecycle()), + Box::new(DatabaseDdlEvent::create_lifecycle("greptime", "metrics")), ProcedureState::Done { output: None }, EventTrigger::Succeeded, ); @@ -220,8 +223,8 @@ fn test_database_events_preserve_procedure_envelope_contract() { CREATE_DATABASE_EVENT_TYPE, "Done", "Succeeded", - None, - None, + Some("greptime"), + Some("metrics"), ); assert_eq!(lifecycle.json_payload().unwrap(), serde_json::Value::Null); } @@ -282,6 +285,7 @@ fn assert_event_locator( catalog_name: Option<&str>, schema_name: Option<&str>, ) { + let has_event_context = !event.json_payload().unwrap().is_null(); assert_event_contract( event, event_type, @@ -298,7 +302,7 @@ fn assert_event_locator( Value { value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, - if catalog_name.is_some() { + if has_event_context { default_event_context_value() } else { Value { value_data: None } @@ -316,6 +320,7 @@ fn assert_procedure_event_contract( catalog_name: Option<&str>, schema_name: Option<&str>, ) { + let has_event_context = procedure_trigger == "Submitted"; assert_event_contract( event, event_type, @@ -348,7 +353,7 @@ fn assert_procedure_event_contract( Value { value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, - if catalog_name.is_some() { + if has_event_context { default_event_context_value() } else { Value { value_data: None } diff --git a/src/common/meta/src/ddl/tests/event/flow.rs b/src/common/meta/src/ddl/tests/event/flow.rs index b2758197917..39f24b74e51 100644 --- a/src/common/meta/src/ddl/tests/event/flow.rs +++ b/src/common/meta/src/ddl/tests/event/flow.rs @@ -99,9 +99,18 @@ fn test_flow_submitted_event_contracts() { #[test] fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { for (event, event_type) in [ - (FlowDdlEvent::create_lifecycle(), CREATE_FLOW_EVENT_TYPE), - (FlowDdlEvent::create_succeeded(None), CREATE_FLOW_EVENT_TYPE), - (FlowDdlEvent::drop_lifecycle(), DROP_FLOW_EVENT_TYPE), + ( + FlowDdlEvent::create_lifecycle("greptime", "metrics"), + CREATE_FLOW_EVENT_TYPE, + ), + ( + FlowDdlEvent::create_succeeded("greptime", "metrics", None), + CREATE_FLOW_EVENT_TYPE, + ), + ( + FlowDdlEvent::drop_lifecycle("greptime", "metrics", 42), + DROP_FLOW_EVENT_TYPE, + ), ] { assert_event_contract( &event, @@ -109,9 +118,13 @@ fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { &flow_schema(), &[Row { values: vec![ - Value { value_data: None }, - Value { value_data: None }, - Value { value_data: None }, + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue("metrics".to_string()).into(), + if event_type == DROP_FLOW_EVENT_TYPE { + ValueData::U32Value(42).into() + } else { + Value { value_data: None } + }, Value { value_data: None }, ], }], @@ -119,15 +132,15 @@ fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); } - let event = FlowDdlEvent::create_succeeded(Some(42)); + let event = FlowDdlEvent::create_succeeded("greptime", "metrics", Some(42)); assert_event_contract( &event, CREATE_FLOW_EVENT_TYPE, &flow_schema(), &[Row { values: vec![ - Value { value_data: None }, - Value { value_data: None }, + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue("metrics".to_string()).into(), ValueData::U32Value(42).into(), Value { value_data: None }, ], @@ -156,7 +169,11 @@ fn test_flow_events_preserve_procedure_envelope_contract() { ); let succeeded = ProcedureEvent::new( procedure_id, - Box::new(FlowDdlEvent::create_succeeded(Some(42))), + Box::new(FlowDdlEvent::create_succeeded( + "greptime", + "metrics", + Some(42), + )), ProcedureState::Done { output: None }, EventTrigger::Succeeded, ); @@ -178,8 +195,8 @@ fn test_flow_events_preserve_procedure_envelope_contract() { "Done", "Succeeded", FlowEventLocator { - catalog_name: None, - flow_name: None, + catalog_name: Some("greptime"), + flow_name: Some("metrics"), flow_id: Some(42), }, ); @@ -257,7 +274,7 @@ fn assert_procedure_event_contract( .map(ValueData::U32Value) .map(Into::into) .unwrap_or(Value { value_data: None }), - if locator.catalog_name.is_some() { + if trigger == "Submitted" { default_event_context_value() } else { Value { value_data: None } diff --git a/src/common/meta/src/ddl/tests/event/table.rs b/src/common/meta/src/ddl/tests/event/table.rs index 90e7d83b3f2..3326cb06808 100644 --- a/src/common/meta/src/ddl/tests/event/table.rs +++ b/src/common/meta/src/ddl/tests/event/table.rs @@ -108,13 +108,26 @@ fn submitted_event_contracts_are_bounded_and_fixed() { .all(|column| column.semantic_type == SemanticType::Field as i32) ); - let lifecycle = TableDdlEvent::lifecycle(case.event_type); + let lifecycle = TableDdlEvent::lifecycle( + case.event_type, + [TableDdlLocator::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + "lifecycle", + )], + ); assert_eq!(lifecycle.extra_schema(), schema); assert_eq!(lifecycle.json_payload().unwrap(), JsonValue::Null); - assert_eq!( - lifecycle.extra_rows().unwrap()[0].values, - vec![Value::default(); schema.len()] - ); + assert_eq!(lifecycle.extra_rows().unwrap()[0].values, { + let mut values = table_locator_values(Some("lifecycle"), None); + if matches!( + case.event_type, + TableDdlEventType::CreateLogicalTables | TableDdlEventType::AlterLogicalTables + ) { + values.push(Value::default()); + } + values + }); } } @@ -158,6 +171,16 @@ fn later_lifecycle_events_are_uniform() { for case in procedure_cases() { let submitted = event_for(case.procedure.as_ref(), EventTrigger::Submitted); let schema = submitted.extra_schema(); + let expected_rows = submitted + .extra_rows() + .unwrap() + .into_iter() + .map(|mut row| { + // Lifecycle events retain locators but not submitted event context. + *row.values.last_mut().unwrap() = Value::default(); + row.values + }) + .collect::>(); for trigger in &triggers { let event = event_for(case.procedure.as_ref(), trigger.clone()); @@ -166,8 +189,13 @@ fn later_lifecycle_events_are_uniform() { assert_eq!(event.extra_schema(), schema); assert_eq!(event.json_payload().unwrap(), JsonValue::Null); assert_eq!( - event.extra_rows().unwrap()[0].values, - vec![Value::default(); schema.len()] + event + .extra_rows() + .unwrap() + .into_iter() + .map(|row| row.values) + .collect::>(), + expected_rows ); } } @@ -193,7 +221,7 @@ fn create_success_events_keep_allocated_ids() { assert_eq!(event.json_payload().unwrap(), JsonValue::Null); assert_eq!( event.extra_rows().unwrap()[0].values, - table_locator_values(None, Some(42)) + table_locator_values(Some("create_success"), Some(42)) ); let logical_tables = CreateLogicalTablesProcedure::new( diff --git a/src/common/meta/src/ddl/tests/event/view.rs b/src/common/meta/src/ddl/tests/event/view.rs index 10947e45ee8..cabe7418ffb 100644 --- a/src/common/meta/src/ddl/tests/event/view.rs +++ b/src/common/meta/src/ddl/tests/event/view.rs @@ -102,26 +102,44 @@ fn test_view_submitted_event_contracts() { #[test] fn test_view_lifecycle_event_contracts() { for (event, event_type) in [ - (ViewDdlEvent::create_lifecycle(), CREATE_VIEW_EVENT_TYPE), - (ViewDdlEvent::drop_lifecycle(), DROP_VIEW_EVENT_TYPE), + ( + ViewDdlEvent::create_lifecycle("greptime", "public", "v_metrics"), + CREATE_VIEW_EVENT_TYPE, + ), + ( + ViewDdlEvent::drop_lifecycle("greptime", "public", "v_metrics", 42), + DROP_VIEW_EVENT_TYPE, + ), ] { - assert_lightweight_event(&event, event_type); + assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); + assert_view_event_contract( + &event, + event_type, + ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("v_metrics"), + view_id: (event_type == DROP_VIEW_EVENT_TYPE).then_some(42), + }, + ); } - let event = ViewDdlEvent::create_succeeded(84); + let event = ViewDdlEvent::create_succeeded("greptime", "public", "v_metrics", 84); assert_view_event_contract( &event, CREATE_VIEW_EVENT_TYPE, ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("v_metrics"), view_id: Some(84), - ..Default::default() }, ); assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); } #[test] -fn test_view_procedures_emit_lightweight_lifecycle_events() { +fn test_view_procedures_preserve_lifecycle_locators() { let create = CreateViewProcedure::new( test_create_view_task("view_name"), EventContext::default(), @@ -147,18 +165,36 @@ fn test_view_procedures_emit_lightweight_lifecycle_events() { EventTrigger::Poisoned, ]; - for (procedure, event_type) in [ - (&create as &dyn Procedure, CREATE_VIEW_EVENT_TYPE), - (&drop as &dyn Procedure, DROP_VIEW_EVENT_TYPE), + for (procedure, event_type, view_id) in [ + (&create as &dyn Procedure, CREATE_VIEW_EVENT_TYPE, None), + (&drop as &dyn Procedure, DROP_VIEW_EVENT_TYPE, Some(42)), ] { for trigger in &triggers { let event = event_for(procedure, trigger.clone()); - assert_lightweight_event(event.as_ref(), event_type); + assert_view_event_contract( + event.as_ref(), + event_type, + ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("view_name"), + view_id, + }, + ); } } let event = event_for(&drop, EventTrigger::Succeeded); - assert_lightweight_event(event.as_ref(), DROP_VIEW_EVENT_TYPE); + assert_view_event_contract( + event.as_ref(), + DROP_VIEW_EVENT_TYPE, + ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("view_name"), + view_id: Some(42), + }, + ); } #[test] @@ -177,8 +213,10 @@ fn test_create_view_succeeded_output_mapping() { event.as_ref(), CREATE_VIEW_EVENT_TYPE, ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("view_name"), view_id: Some(84), - ..Default::default() }, ); assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); @@ -187,7 +225,16 @@ fn test_create_view_succeeded_output_mapping() { for output in invalid_outputs { let state = ProcedureState::Done { output }; let event = event_for_state(&procedure, EventTrigger::Succeeded, &state); - assert_lightweight_event(event.as_ref(), CREATE_VIEW_EVENT_TYPE); + assert_view_event_contract( + event.as_ref(), + CREATE_VIEW_EVENT_TYPE, + ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("view_name"), + view_id: None, + }, + ); } } @@ -233,7 +280,12 @@ fn test_view_event_procedure_envelope_contract() { ); let succeeded = ProcedureEvent::new( procedure_id, - Box::new(ViewDdlEvent::create_succeeded(42)), + Box::new(ViewDdlEvent::create_succeeded( + "greptime", + "public", + "view_name", + 42, + )), ProcedureState::Done { output: None }, EventTrigger::Succeeded, ); @@ -256,8 +308,10 @@ fn test_view_event_procedure_envelope_contract() { "Done", "Succeeded", ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("view_name"), view_id: Some(42), - ..Default::default() }, ); } @@ -280,11 +334,6 @@ impl ViewEventLocator<'_> { .map(ValueData::U32Value) .map(Into::into) .unwrap_or_default(), - if self.catalog_name.is_some() { - default_event_context_value() - } else { - Value { value_data: None } - }, ] } } @@ -300,19 +349,13 @@ fn view_schema() -> Vec { } fn assert_view_event_contract(event: &dyn Event, event_type: &str, locator: ViewEventLocator<'_>) { - assert_event_contract( - event, - event_type, - &view_schema(), - &[Row { - values: locator.values(), - }], - ); -} - -fn assert_lightweight_event(event: &dyn Event, event_type: &str) { - assert_view_event_contract(event, event_type, ViewEventLocator::default()); - assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); + let mut values = locator.values(); + values.push(if event.json_payload().unwrap().is_null() { + Value { value_data: None } + } else { + default_event_context_value() + }); + assert_event_contract(event, event_type, &view_schema(), &[Row { values }]); } fn assert_procedure_event_contract( @@ -337,6 +380,11 @@ fn assert_procedure_event_contract( procedure_trigger_value(trigger), ]; values.extend(locator.values()); + values.push(if trigger == "Submitted" { + default_event_context_value() + } else { + Value { value_data: None } + }); assert_event_contract(event, event_type, &schema, &[Row { values }]); } diff --git a/src/common/meta/src/ddl/truncate_table.rs b/src/common/meta/src/ddl/truncate_table.rs index 9d0ed09e8c8..6a561b62afb 100644 --- a/src/common/meta/src/ddl/truncate_table.rs +++ b/src/common/meta/src/ddl/truncate_table.rs @@ -99,18 +99,16 @@ impl Procedure for TruncateTableProcedure { { return None; } + let task = &self.data.task; + let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) + .with_table_id(task.table_id); let event = match &ctx.trigger { - EventTrigger::Submitted => { - let task = &self.data.task; - let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) - .with_table_id(task.table_id); - TableDdlEvent::truncate_table_submitted( - locator, - task.time_ranges.len(), - self.data.event_context.clone(), - ) - } - _ => TableDdlEvent::lifecycle(TableDdlEventType::TruncateTable), + EventTrigger::Submitted => TableDdlEvent::truncate_table_submitted( + locator, + task.time_ranges.len(), + self.data.event_context.clone(), + ), + _ => TableDdlEvent::lifecycle(TableDdlEventType::TruncateTable, [locator]), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs index 635d15ee6a5..3295b1163e2 100644 --- a/src/common/meta/src/ddl/undrop_table.rs +++ b/src/common/meta/src/ddl/undrop_table.rs @@ -380,6 +380,23 @@ impl UndropTableProcedure { } } +impl UndropTableProcedure { + fn event_locator(&self) -> TableDdlLocator { + self.data + .table_name + .as_ref() + .map(|table_name| { + TableDdlLocator::new( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + ) + }) + .unwrap_or_default() + .with_table_id(self.data.task.table_id) + } +} + #[async_trait] impl Procedure for UndropTableProcedure { fn type_name(&self) -> &str { @@ -437,24 +454,12 @@ impl Procedure for UndropTableProcedure { { return None; } + let locator = self.event_locator(); let event = match &ctx.trigger { EventTrigger::Submitted => { - let locator = self - .data - .table_name - .as_ref() - .map(|table_name| { - TableDdlLocator::new( - &table_name.catalog_name, - &table_name.schema_name, - &table_name.table_name, - ) - }) - .unwrap_or_default() - .with_table_id(self.data.task.table_id); TableDdlEvent::undrop_table_submitted(locator, self.data.event_context.clone()) } - _ => TableDdlEvent::lifecycle(TableDdlEventType::UndropTable), + _ => TableDdlEvent::lifecycle(TableDdlEventType::UndropTable, [locator]), }; Some(Box::new(event)) diff --git a/src/meta-srv/src/event/repartition.rs b/src/meta-srv/src/event/repartition.rs index ad47fbf9d42..958bc4e11d8 100644 --- a/src/meta-srv/src/event/repartition.rs +++ b/src/meta-srv/src/event/repartition.rs @@ -95,12 +95,12 @@ impl RepartitionEvent { } } - pub(crate) fn lifecycle() -> Self { + pub(crate) fn lifecycle(persistent_ctx: &RepartitionPersistentContext) -> Self { Self { - catalog_name: None, - schema_name: None, - table_name: None, - table_id: None, + catalog_name: Some(persistent_ctx.catalog_name.clone()), + schema_name: Some(persistent_ctx.schema_name.clone()), + table_name: Some(persistent_ctx.table_name.clone()), + table_id: Some(persistent_ctx.table_id), payload: None, event_context: None, } @@ -279,14 +279,14 @@ impl RepartitionGroupEvent { } } - pub(crate) fn lifecycle() -> Self { + pub(crate) fn lifecycle(persistent_ctx: &GroupPersistentContext) -> Self { Self { - catalog_name: None, - schema_name: None, - table_name: None, - table_id: None, - parent_procedure_id: None, - group_id: None, + catalog_name: Some(persistent_ctx.catalog_name.clone()), + schema_name: Some(persistent_ctx.schema_name.clone()), + table_name: persistent_ctx.table_name.clone(), + table_id: Some(persistent_ctx.table_id), + parent_procedure_id: persistent_ctx.parent_procedure_id.map(|id| id.to_string()), + group_id: Some(persistent_ctx.group_id.to_string()), topology: None, payload: None, } @@ -377,7 +377,7 @@ mod tests { use std::time::Duration; use api::v1::value::ValueData; - use api::v1::{ColumnSchema, Row, Value}; + use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::event_table::{ PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, @@ -641,22 +641,32 @@ mod tests { } #[test] - fn test_lifecycle_events_are_lightweight() { - let parent = RepartitionEvent::lifecycle(); + fn test_lifecycle_events_preserve_locators_and_null_payloads() { + let parent_ctx = parent_persistent_ctx(); + let parent = RepartitionEvent::lifecycle(&parent_ctx); assert_event_contract( &parent, REPARTITION_EVENT_TYPE, &parent_schema(), - &[null_row(parent_schema().len())], + &[Row { + values: vec![ + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue("public".to_string()).into(), + ValueData::StringValue("repartition_events".to_string()).into(), + ValueData::U32Value(1024).into(), + Default::default(), + ], + }], ); assert_eq!(parent.json_payload().unwrap(), serde_json::Value::Null); - let group = RepartitionGroupEvent::lifecycle(); + let group_ctx = new_persistent_context(1024, vec![], vec![]); + let group = RepartitionGroupEvent::lifecycle(&group_ctx); assert_event_contract( &group, REPARTITION_GROUP_EVENT_TYPE, &group_schema(), - &[null_row(group_schema().len())], + &[group.extra_row(None)], ); assert_eq!(group.json_payload().unwrap(), serde_json::Value::Null); } @@ -762,12 +772,4 @@ mod tests { ], } } - - fn null_row(column_count: usize) -> Row { - Row { - values: (0..column_count) - .map(|_| Value { value_data: None }) - .collect(), - } - } } diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index 62f9f26aabd..d8b585b28ee 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -809,7 +809,7 @@ impl Procedure for RepartitionProcedure { let start = self.state.as_any().downcast_ref::()?; RepartitionEvent::submitted(&self.context.persistent_ctx, start) } else { - RepartitionEvent::lifecycle() + RepartitionEvent::lifecycle(&self.context.persistent_ctx) }; Some(Box::new(event)) } diff --git a/src/meta-srv/src/procedure/repartition/group.rs b/src/meta-srv/src/procedure/repartition/group.rs index 700c2fdcd1d..a3c38b59557 100644 --- a/src/meta-srv/src/procedure/repartition/group.rs +++ b/src/meta-srv/src/procedure/repartition/group.rs @@ -273,7 +273,7 @@ impl Procedure for RepartitionGroupProcedure { let event = if matches!(ctx.trigger, EventTrigger::Submitted) { RepartitionGroupEvent::submitted(&self.context.persistent_ctx) } else { - RepartitionGroupEvent::lifecycle() + RepartitionGroupEvent::lifecycle(&self.context.persistent_ctx) }; Some(Box::new(event)) } diff --git a/tests-integration/tests/database_ddl_event.rs b/tests-integration/tests/database_ddl_event.rs index 79b741b977f..6cbc65a2673 100644 --- a/tests-integration/tests/database_ddl_event.rs +++ b/tests-integration/tests/database_ddl_event.rs @@ -138,8 +138,8 @@ FROM greptime_private.events WHERE type = '{event_type}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') - AND catalog_name IS NULL - AND schema_name IS NULL + AND catalog_name = 'greptime' + AND schema_name = '{DATABASE_NAME}' AND json_is_null(payload)"# ); assert_single_event(instance, &lifecycle).await; diff --git a/tests-integration/tests/flow_ddl_event.rs b/tests-integration/tests/flow_ddl_event.rs index 004d8382ed2..c5fc9be94c9 100644 --- a/tests-integration/tests/flow_ddl_event.rs +++ b/tests-integration/tests/flow_ddl_event.rs @@ -98,8 +98,8 @@ WHERE type = '{CREATE_FLOW_EVENT_TYPE}' AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') - AND catalog_name IS NULL - AND flow_name IS NULL + AND catalog_name = 'greptime' + AND flow_name = '{flow}' AND flow_id IS NOT NULL AND json_is_null(payload)"#, ), @@ -152,9 +152,9 @@ WHERE type = '{DROP_FLOW_EVENT_TYPE}' AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') - AND catalog_name IS NULL - AND flow_name IS NULL - AND flow_id IS NULL + AND catalog_name = 'greptime' + AND flow_name = '{flow}' + AND flow_id IS NOT NULL AND json_is_null(payload)"#, ), ) diff --git a/tests-integration/tests/repartition_event.rs b/tests-integration/tests/repartition_event.rs index 3871406cfa2..65503a4a8ee 100644 --- a/tests-integration/tests/repartition_event.rs +++ b/tests-integration/tests/repartition_event.rs @@ -129,9 +129,32 @@ WHERE type = 'repartition' +-------------+--------------+-------------+--------------------+----------+-----------------+----------------------+---------------+"; assert_eventually_eq(instance, &query, expected).await; + assert_repartition_lifecycle_event(instance, &procedure_id).await; procedure_id } +async fn assert_repartition_lifecycle_event( + instance: &Arc, + procedure_id: &str, +) { + let query = format!( + r#"SELECT type, catalog_name, schema_name, table_name, table_id +FROM greptime_private.events +WHERE type = 'repartition' + AND procedure_id = '{procedure_id}' + AND json_path_match(procedure_trigger, '$.type == "Succeeded"') + AND json_is_null(payload)"# + ); + let expected = "\ ++-------------+--------------+-------------+--------------------+----------+ +| type | catalog_name | schema_name | table_name | table_id | ++-------------+--------------+-------------+--------------------+----------+ +| repartition | greptime | public | repartition_events | 1024 | ++-------------+--------------+-------------+--------------------+----------+"; + + assert_eventually_eq(instance, &query, expected).await; +} + async fn assert_repartition_group_event( instance: &Arc, parent_procedure_id: &str, @@ -170,6 +193,32 @@ ORDER BY target_region_id"# | repartition_group | greptime | public | repartition_events | 1024 | true | true | 4398046511104 | 0 | true | 4398046511105 | 1 | host >= 10 | +-------------------+--------------+-------------+--------------------+----------+-----------------------------+--------------+------------------+----------------------+----------------+------------------+----------------------+-----------------------+"; + assert_eventually_eq(instance, &query, expected).await; + assert_repartition_group_lifecycle_event(instance, &procedure_id, parent_procedure_id).await; +} + +async fn assert_repartition_group_lifecycle_event( + instance: &Arc, + procedure_id: &str, + parent_procedure_id: &str, +) { + let query = format!( + r#"SELECT type, catalog_name, schema_name, table_name, table_id, + parent_procedure_id = '{parent_procedure_id}' AS matches_parent_procedure_id, + repartition_group_id IS NOT NULL AS has_group_id +FROM greptime_private.events +WHERE type = 'repartition_group' + AND procedure_id = '{procedure_id}' + AND json_path_match(procedure_trigger, '$.type == "Succeeded"') + AND json_is_null(payload)"# + ); + let expected = "\ ++-------------------+--------------+-------------+--------------------+----------+-----------------------------+--------------+ +| type | catalog_name | schema_name | table_name | table_id | matches_parent_procedure_id | has_group_id | ++-------------------+--------------+-------------+--------------------+----------+-----------------------------+--------------+ +| repartition_group | greptime | public | repartition_events | 1024 | true | true | ++-------------------+--------------+-------------+--------------------+----------+-----------------------------+--------------+"; + assert_eventually_eq(instance, &query, expected).await; } diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index c53842d66d6..5c23e101de2 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -230,23 +230,17 @@ async fn test_table_ddl_procedure_events() { +--------------+------------------+", ) .await; - assert_id_terminal_event( + assert_terminal_event( &frontend, "create_table", &create_table_procedure_id, - table_id, "Succeeded", - "\ -+--------------+ -| type | -+--------------+ -| create_table | -+--------------+", + Some(TABLE), + Some(table_id), ) .await; - // Act / Assert: Alter and Truncate have a rich submitted row and lightweight - // terminal lifecycle row. + // Act / Assert: Alter and Truncate retain their submitted locators on terminal rows. run_sql_with_context( &frontend, &format!("ALTER TABLE {TABLE} ADD COLUMN extra STRING"), @@ -278,11 +272,13 @@ async fn test_table_ddl_procedure_events() { +-------------+------------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "alter_table", &alter_table_procedure_id, "Succeeded", + Some(TABLE), + Some(table_id), ) .await; run_sql( @@ -322,11 +318,13 @@ async fn test_table_ddl_procedure_events() { +----------------+------------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "truncate_table", &truncate_table_procedure_id, "Succeeded", + Some(TABLE), + Some(table_id), ) .await; @@ -363,11 +361,13 @@ async fn test_table_ddl_procedure_events() { +------------+------------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "drop_table", &drop_table_procedure_id, "Succeeded", + Some(TABLE), + Some(table_id), ) .await; #[cfg(feature = "enterprise")] @@ -404,11 +404,13 @@ async fn test_table_ddl_procedure_events() { +--------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "undrop_table", &undrop_table_procedure_id, "Succeeded", + Some(TABLE), + Some(table_id), ) .await; run_sql(&frontend, &format!("DROP TABLE {TABLE}")).await; @@ -444,11 +446,13 @@ async fn test_table_ddl_procedure_events() { +---------------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "purge_dropped_table", &purge_table_procedure_id, "Succeeded", + Some(TABLE), + Some(table_id), ) .await; } @@ -518,7 +522,7 @@ async fn test_table_ddl_procedure_events() { .await; // Act / Assert: logical Alter preserves the one-row-per-logical-table submitted - // contract and emits a lightweight terminal row. + // contract and retains the logical-table locator on its terminal row. run_sql_with_context( &frontend, &format!("ALTER TABLE {LOGICAL_TABLE} ADD COLUMN rack STRING PRIMARY KEY"), @@ -551,11 +555,13 @@ async fn test_table_ddl_procedure_events() { +----------------------+--------------------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "alter_logical_tables", &alter_logical_tables_procedure_id, "Succeeded", + Some(LOGICAL_TABLE), + None, ) .await; } @@ -719,27 +725,46 @@ async fn assert_id_submitted_event( assert_eventually_eq(instance, &query, expected).await; } -async fn assert_id_terminal_event( +async fn assert_terminal_event( instance: &Arc, event_type: &str, procedure_id: &str, - table_id: u32, terminal_trigger: &str, - expected: &str, + table_name: Option<&str>, + table_id: Option, ) { + let table_name_predicate = table_name.map_or_else( + || format!("{} IS NULL", TABLE_NAME_COLUMN.name()), + |table_name| format!("{} = '{table_name}'", TABLE_NAME_COLUMN.name()), + ); + let catalog_and_schema_predicate = table_name.map_or_else( + || { + format!( + "{} IS NULL AND {} IS NULL", + CATALOG_NAME_COLUMN.name(), + SCHEMA_NAME_COLUMN.name(), + ) + }, + |_| { + format!( + "{} = 'greptime' AND {} = 'public'", + CATALOG_NAME_COLUMN.name(), + SCHEMA_NAME_COLUMN.name(), + ) + }, + ); + let table_id_predicate = table_id.map_or_else( + || format!("{} IS NULL", TABLE_ID_COLUMN.name()), + |table_id| format!("{} = {table_id}", TABLE_ID_COLUMN.name()), + ); let query = format!( - "SELECT {} FROM {EVENTS_TABLE} WHERE {} = '{event_type}' AND {} = '{procedure_id}' AND json_path_match({}, '$.type == \"{terminal_trigger}\"') AND {} = {table_id} AND json_is_null({}) AND {} IS NULL AND {} IS NULL AND {} IS NULL", - TYPE_COLUMN.name(), + "SELECT count(*) AS event_count FROM {EVENTS_TABLE} WHERE {} = '{event_type}' AND {} = '{procedure_id}' AND json_path_match({}, '$.type == \"{terminal_trigger}\"') AND json_is_null({}) AND {catalog_and_schema_predicate} AND {table_name_predicate} AND {table_id_predicate}", TYPE_COLUMN.name(), PROCEDURE_ID_COLUMN.name(), PROCEDURE_TRIGGER_COLUMN.name(), - TABLE_ID_COLUMN.name(), PAYLOAD_COLUMN.name(), - CATALOG_NAME_COLUMN.name(), - SCHEMA_NAME_COLUMN.name(), - TABLE_NAME_COLUMN.name(), ); - assert_eventually_eq(instance, &query, expected).await; + assert_single_event(instance, &query).await; } async fn assert_logical_terminal_event( @@ -764,23 +789,3 @@ async fn assert_logical_terminal_event( ); assert_eventually_eq(instance, &query, expected).await; } - -async fn assert_lightweight_terminal_event( - instance: &Arc, - event_type: &str, - procedure_id: &str, - terminal_trigger: &str, -) { - let query = format!( - "SELECT count(*) AS event_count FROM {EVENTS_TABLE} WHERE {} = '{event_type}' AND {} = '{procedure_id}' AND json_path_match({}, '$.type == \"{terminal_trigger}\"') AND json_is_null({}) AND {} IS NULL AND {} IS NULL AND {} IS NULL AND {} IS NULL", - TYPE_COLUMN.name(), - PROCEDURE_ID_COLUMN.name(), - PROCEDURE_TRIGGER_COLUMN.name(), - PAYLOAD_COLUMN.name(), - CATALOG_NAME_COLUMN.name(), - SCHEMA_NAME_COLUMN.name(), - TABLE_NAME_COLUMN.name(), - TABLE_ID_COLUMN.name(), - ); - assert_single_event(instance, &query).await; -} diff --git a/tests-integration/tests/view_ddl_event.rs b/tests-integration/tests/view_ddl_event.rs index a25b58b684a..fce33d09bc8 100644 --- a/tests-integration/tests/view_ddl_event.rs +++ b/tests-integration/tests/view_ddl_event.rs @@ -134,9 +134,9 @@ WHERE type = '{CREATE_VIEW_EVENT_TYPE}' AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') - AND catalog_name IS NULL - AND schema_name IS NULL - AND view_name IS NULL + AND catalog_name = 'greptime' + AND schema_name = 'public' + AND view_name = '{view}' AND view_id IS NOT NULL AND json_is_null(payload)"#, ), @@ -190,10 +190,10 @@ WHERE type = '{DROP_VIEW_EVENT_TYPE}' AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') - AND catalog_name IS NULL - AND schema_name IS NULL - AND view_name IS NULL - AND view_id IS NULL + AND catalog_name = 'greptime' + AND schema_name = 'public' + AND view_name = '{view}' + AND view_id IS NOT NULL AND json_is_null(payload)"#, ), ) From c851430349faf7be24ed1ed530b4f361d1588b18 Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 11 Aug 2026 15:16:06 +0800 Subject: [PATCH 044/126] chore(deps): bump datafusion to 452cb4b (support Dictionary literals in substrait) (#8839) Bump the GreptimeTeam/datafusion fork rev from 6d6ae9a to 452cb4b, which includes fix(substrait): support Dictionary literals in producer. This fixes flow queries against dictionary-encoded PK string columns (metric tables) failing with: Failed to encode DataFusion plan: NotImplemented("Unsupported literal: Dictionary(UInt32, Utf8(...))") The substrait producer now encodes ScalarValue::Dictionary as its inner value wrapped in a cast to the dictionary type, so the original SQL works without CAST workarounds. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 116 ++++++++++++++++++++++++++--------------------------- Cargo.toml | 30 +++++++------- 2 files changed, 73 insertions(+), 73 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d6385dd7657..2e6fe8ae92c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2226,7 +2226,7 @@ checksum = "fe6d2e5af09e8c8ad56c969f2157a3d4238cebc7c55f0a517728c38f7b200f81" dependencies = [ "serde", "termcolor", - "unicode-width 0.2.1", + "unicode-width 0.1.14", ] [[package]] @@ -2237,7 +2237,7 @@ checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" dependencies = [ "serde", "termcolor", - "unicode-width 0.2.1", + "unicode-width 0.1.14", ] [[package]] @@ -2253,7 +2253,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -3726,7 +3726,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "arrow-schema 58.3.0", @@ -3780,7 +3780,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3804,7 +3804,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3826,7 +3826,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -3850,7 +3850,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "futures", "log", @@ -3860,7 +3860,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-compression", @@ -3894,7 +3894,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "arrow-ipc 58.3.0", @@ -3917,7 +3917,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3939,7 +3939,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3962,7 +3962,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3991,12 +3991,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" [[package]] name = "datafusion-execution" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4018,7 +4018,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4040,7 +4040,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4052,7 +4052,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4083,7 +4083,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4104,7 +4104,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4116,7 +4116,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "arrow-ord 58.3.0", @@ -4140,7 +4140,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4155,7 +4155,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4172,7 +4172,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4181,7 +4181,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "datafusion-doc", "quote", @@ -4191,7 +4191,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "chrono", @@ -4240,7 +4240,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4263,7 +4263,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4277,7 +4277,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4293,7 +4293,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4311,7 +4311,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4342,7 +4342,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "chrono", @@ -4369,7 +4369,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4379,7 +4379,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4395,7 +4395,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "async-trait", "datafusion-common", @@ -4408,7 +4408,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "bigdecimal 0.4.8", @@ -4426,7 +4426,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "async-recursion", "async-trait", @@ -5119,7 +5119,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -5681,7 +5681,7 @@ checksum = "94e7099f6313ecacbe1256e8ff9d617b75d1bcb16a6fddef94866d225a01a14a" dependencies = [ "io-lifetimes 2.0.4", "rustix 1.0.7", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -6578,7 +6578,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.5.10", "tokio", "tower-service", "tracing", @@ -6649,7 +6649,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.61.2", + "windows-core 0.57.0", ] [[package]] @@ -7077,7 +7077,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20fd6de4ccfcc187e38bc21cfa543cb5a302cb86a8b114eb7f0bf0dc9f8ac00f" dependencies = [ "io-lifetimes 3.0.1", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -7115,7 +7115,7 @@ checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -7818,7 +7818,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.53.5", + "windows-targets 0.48.5", ] [[package]] @@ -11212,7 +11212,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.5.0", + "heck 0.4.1", "itertools 0.12.1", "log", "multimap", @@ -11232,8 +11232,8 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.5.0", - "itertools 0.14.0", + "heck 0.4.1", + "itertools 0.10.5", "log", "multimap", "once_cell", @@ -11281,7 +11281,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -11294,7 +11294,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -11760,7 +11760,7 @@ dependencies = [ "once_cell", "socket2 0.5.10", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -12709,7 +12709,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.4.15", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -12722,7 +12722,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.9.4", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -12813,7 +12813,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -13697,7 +13697,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.117", @@ -14112,7 +14112,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -14722,7 +14722,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.0.7", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -16470,7 +16470,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -16914,7 +16914,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f3fd376f71958b862e7afb20cfe5a22830e1963462f3a17f49d82a6c1d1f42d" dependencies = [ "bitflags 2.12.1", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 64e6dba2f7f..a138a880ec6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -345,21 +345,21 @@ git = "https://github.com/GreptimeTeam/greptime-meter.git" rev = "5618e779cf2bb4755b499c630fba4c35e91898cb" [patch.crates-io] -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "2aefa08a8d69c96eec2d6d6703598a009bba6e4c" } # on branch v0.61.x [profile.release] From f27f27da661c69033de7b2e4a5b70b4f058f8dd2 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 11 Aug 2026 15:59:12 +0800 Subject: [PATCH 045/126] fix(ci): grant pull-requests write and stop counting drafts (#8844) Posting to `/issues/{n}/comments` is authorized against the target object, and that object is a pull request, so `issues: write` alone is refused with 403 and the warning comment never lands. Drafts are no longer counted and no longer warned about. `ready_for_review` is added to the trigger types so that opening as a draft and flipping it to ready still goes through the check. Signed-off-by: Dennis Zhuang --- .github/scripts/pr-open-limit.js | 12 +++++++++--- .github/workflows/pr-open-limit.yml | 5 +++-- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/.github/scripts/pr-open-limit.js b/.github/scripts/pr-open-limit.js index dd1650b0d64..46f770074f3 100644 --- a/.github/scripts/pr-open-limit.js +++ b/.github/scripts/pr-open-limit.js @@ -1,7 +1,8 @@ // Warns when a repository member has too many open pull requests. // // Review capacity is the bottleneck, so authors are expected to land or close -// existing work before opening more. Drafts count: they still occupy attention. +// existing work before opening more. Drafts are not counted and never warned +// about: they are not asking for review yet. // This is advisory only — it never closes a pull request or fails the job. import { appendFileSync } from "node:fs"; @@ -47,7 +48,7 @@ async function isTeamMember(octokit, owner, repo, username) { function buildComment(author, total, limit, openPrs) { const listed = openPrs .slice(0, MAX_LISTED_PRS) - .map((pr) => `- #${pr.number} ${pr.title}${pr.draft ? " _(draft)_" : ""}`); + .map((pr) => `- #${pr.number} ${pr.title}`); const hidden = openPrs.length - listed.length; const list = hidden > 0 ? `${listed.join("\n")}\n- ...and ${hidden} more` : listed.join("\n"); @@ -99,6 +100,11 @@ async function upsertComment(octokit, params, body) { return; } + if (process.env.PR_IS_DRAFT === "true") { + summary(`Skipping #${prNumber}: draft.`); + return; + } + const octokit = new Octokit({ auth: process.env.GITHUB_TOKEN }); if (!(await isTeamMember(octokit, owner, repo, author))) { @@ -117,7 +123,7 @@ async function upsertComment(octokit, params, body) { // The pull request that triggered this run is already open, so exclude it and // report the total separately. const others = allOpen.filter( - (pr) => pr.user?.login === author && pr.number !== prNumber + (pr) => pr.user?.login === author && pr.number !== prNumber && !pr.draft ); const total = others.length + 1; diff --git a/.github/workflows/pr-open-limit.yml b/.github/workflows/pr-open-limit.yml index de5b723cb25..081454cb00e 100644 --- a/.github/workflows/pr-open-limit.yml +++ b/.github/workflows/pr-open-limit.yml @@ -5,15 +5,15 @@ on: types: - opened - reopened + - ready_for_review concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number }} cancel-in-progress: true -# The warning comment goes through the issues API, hence `issues: write`. permissions: contents: read - pull-requests: read + pull-requests: write issues: write jobs: @@ -44,4 +44,5 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} PR_NUMBER: ${{ github.event.pull_request.number }} PR_AUTHOR: ${{ github.event.pull_request.user.login }} + PR_IS_DRAFT: ${{ github.event.pull_request.draft }} MAX_OPEN_PRS: ${{ vars.MAX_OPEN_PRS_PER_AUTHOR || 5 }} From 97cbf79fb6dfe9216e75ab2f22bbca32b61f6f67 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Tue, 11 Aug 2026 16:01:06 +0800 Subject: [PATCH 046/126] feat(promql): support native histogram vector operators (#8798) * feat(promql): support native histogram vector operators Signed-off-by: shuiyisong * fix: CR issue Signed-off-by: shuiyisong * chore: rebase main Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- .../src/extension_plan/instant_manipulate.rs | 72 +- src/promql/src/functions/native_histogram.rs | 20 +- src/query/src/promql/planner.rs | 1741 +++++++++++++++-- 3 files changed, 1668 insertions(+), 165 deletions(-) diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index f92a4e4490f..4851916a017 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -18,6 +18,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use common_query::prelude::{greptime_native_histogram, greptime_value}; use datafusion::arrow::array::{Array, TimestampMillisecondArray, UInt64Array}; use datafusion::arrow::datatypes::{DataType, SchemaRef}; use datafusion::arrow::record_batch::RecordBatch; @@ -52,6 +53,15 @@ use crate::metrics::PROMQL_SERIES_COUNT; const MAX_INSTANT_MANIPULATE_OUTPUT_POINTS: usize = 1_000_000; +fn mixed_sample_fields(field: Option<&str>) -> [Option<&str>; 2] { + let companion = match field { + Some(field) if field == greptime_value() => Some(greptime_native_histogram()), + Some(field) if field == greptime_native_histogram() => Some(greptime_value()), + _ => None, + }; + [field, companion] +} + /// Manipulate the input record batch to make it suitable for Instant Operator. /// /// This plan will try to align the input time series, for every timestamp between @@ -66,7 +76,7 @@ pub struct InstantManipulate { time_index_column: String, // Planner-provided tag-column hint for execution fast paths. tag_columns: Vec, - /// A optional column for validating staleness + /// Primary sample column used to derive the columns checked for staleness. field_column: Option, input: LogicalPlan, unfix: Option, @@ -97,9 +107,7 @@ impl UserDefinedLogicalNodeCore for InstantManipulate { } let mut exprs = vec![col(&self.time_index_column)]; - if let Some(field) = &self.field_column { - exprs.push(col(field)); - } + exprs.extend(self.staleness_field_columns().map(col)); exprs } @@ -116,7 +124,7 @@ impl UserDefinedLogicalNodeCore for InstantManipulate { let mut required = output_columns.to_vec(); required.push(input_schema.index_of_column_by_name(None, &self.time_index_column)?); - if let Some(field) = &self.field_column { + for field in self.staleness_field_columns() { required.push(input_schema.index_of_column_by_name(None, field)?); } @@ -223,6 +231,21 @@ impl InstantManipulate { "InstantManipulate" } + fn staleness_field_columns(&self) -> impl Iterator { + let [field, companion] = mixed_sample_fields(self.field_column.as_deref()); + [ + field, + companion.filter(|companion| { + self.input + .schema() + .index_of_column_by_name(None, companion) + .is_some() + }), + ] + .into_iter() + .flatten() + } + fn resolve_tag_columns(input: &LogicalPlan, tag_columns: &[String]) -> Vec { if !tag_columns.is_empty() { return tag_columns.to_vec(); @@ -385,11 +408,9 @@ impl ExecutionPlan for InstantManipulateExec { .column_with_name(&self.time_index_column) .expect("time index column not found") .0; - let field_index = self - .field_column - .as_ref() - .and_then(|name| schema.column_with_name(name)) - .map(|x| x.0); + let field_indices = mixed_sample_fields(self.field_column.as_deref()).map(|field| { + field.and_then(|field| schema.column_with_name(field).map(|(index, _)| index)) + }); let tsid_index = schema .column_with_name("__tsid") .filter(|(_, field)| field.data_type() == &DataType::UInt64) @@ -400,7 +421,7 @@ impl ExecutionPlan for InstantManipulateExec { lookback_delta: self.lookback_delta, interval: self.interval, time_index, - field_index, + field_indices, tsid_index, reuse_tsid_column: self.reuse_tsid_column && tsid_index.is_some(), schema, @@ -467,7 +488,7 @@ pub struct InstantManipulateStream { interval: Millisecond, // Column index of TIME INDEX column's position in schema time_index: usize, - field_index: Option, + field_indices: [Option; 2], tsid_index: Option, reuse_tsid_column: bool, @@ -532,11 +553,16 @@ impl InstantManipulateStream { return Ok(input); } - // Field column for staleness checks, classified once per batch. - let stale_sample_column = self - .field_index - .map(|index| input.column(index).as_ref()) - .and_then(prometheus_stale_sample_column); + // Field columns for staleness checks, classified once per batch. + let stale_sample_columns = self.field_indices.map(|index| { + index.and_then(|index| prometheus_stale_sample_column(input.column(index).as_ref())) + }); + let is_stale = |row| { + stale_sample_columns + .iter() + .flatten() + .any(|column| is_prometheus_stale_sample(*column, row)) + }; // Optimize iteration range based on actual data bounds let first_ts = ts_column.value(0); @@ -581,9 +607,7 @@ impl InstantManipulateStream { let curr = ts_column.value(cursor); match curr.cmp(&expected_ts) { Ordering::Equal => { - if stale_sample_column - .is_some_and(|column| is_prometheus_stale_sample(column, cursor)) - { + if is_stale(cursor) { // Ignore the stale marker. } else { take_indices.push(cursor as u64); @@ -615,9 +639,7 @@ impl InstantManipulateStream { let prev_ts = ts_column.value(prev_cursor); if prev_ts + self.lookback_delta > expected_ts { // only use the point in the time range - if stale_sample_column - .is_some_and(|column| is_prometheus_stale_sample(column, prev_cursor)) - { + if is_stale(prev_cursor) { // Do not use a stale marker as the newest value. continue; } @@ -626,9 +648,7 @@ impl InstantManipulateStream { aligned_ts.push(expected_ts); } } - } else if stale_sample_column - .is_some_and(|column| is_prometheus_stale_sample(column, cursor)) - { + } else if is_stale(cursor) { // Do not use a stale marker as the newest value. } else { // use this point diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index 1863e663b82..e6c9e84d59f 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -173,12 +173,20 @@ impl ScalarUDFImpl for NativeHistogramAnnotationUdf { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult { - if let Some(collector) = args - .config_options - .extensions - .get::() - .cloned() - .or_else(|| self.collector.clone()) + let has_dropped_sample = !args.args.is_empty() + && (0..args.number_rows).any(|row| { + args.args.iter().all(|arg| match arg { + ColumnarValue::Array(array) => array.is_valid(row), + ColumnarValue::Scalar(value) => !value.is_null(), + }) + }); + if has_dropped_sample + && let Some(collector) = args + .config_options + .extensions + .get::() + .cloned() + .or_else(|| self.collector.clone()) { collector.record_info(self.message.clone()); } diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index d21a90bf123..10462ea1828 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -64,14 +64,17 @@ use promql::extension_plan::{ use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, - NativeHistogramAbsentOverTime, NativeHistogramAvg, NativeHistogramAvgOverTime, - NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, - NativeHistogramDelta, NativeHistogramDrop, NativeHistogramFraction, NativeHistogramIDelta, + NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAvg, + NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, + NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDivScalar, + NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, + NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, - NativeHistogramResets, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSum, - NativeHistogramSumOverTime, PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, - Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, + NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, + NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, PredictLinear, + PresentOverTime, QuantileOverTime, Rate, Resets, Round, StddevOverTime, StdvarOverTime, + SumOverTime, quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; @@ -392,6 +395,8 @@ pub struct PromPlanner { promql_annotations: Option, } +type BinaryFieldPair<'a> = (&'a String, &'a String); + impl PromPlanner { pub async fn stmt_to_plan( table_provider: DfTableSourceProvider, @@ -823,8 +828,20 @@ impl PromPlanner { let UnaryExpr { expr } = unary_expr; // Unary Expr in PromQL implys the `-` operator let input = self.prom_expr_to_plan(expr, query_engine_state).await?; + self.negate_field_columns(input) + } + + fn negate_field_columns(&mut self, input: LogicalPlan) -> Result { + let input_schema = input.schema().clone(); self.projection_for_each_field_column(input, |col| { - Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into())))) + if Self::field_column_is_native_histogram(&input_schema, col) { + Ok(DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramNeg::scalar_udf()), + args: vec![DfExpr::Column(col.into())], + })) + } else { + Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into())))) + } }) } @@ -866,6 +883,16 @@ impl PromPlanner { }); } + if planned_leaves.iter().any(|leaf| { + Self::field_columns_contain_native_histogram( + leaf.plan.schema(), + &leaf.ctx.field_columns, + ) + }) { + self.ctx = original_ctx; + return Ok(None); + } + if !Self::binary_island_join_contexts_supported(&planned_leaves) { self.ctx = original_ctx; return Ok(None); @@ -1214,10 +1241,46 @@ impl PromPlanner { if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) { expr = time_expr } + let input_schema = input.schema().clone(); + let preserve_any_value = Self::field_columns_are_alternative_samples( + &input_schema, + &self.ctx.field_columns, + ); + let has_native_histogram = Self::field_columns_contain_native_histogram( + &input_schema, + &self.ctx.field_columns, + ); + let retain_field_columns = self + .ctx + .field_columns + .iter() + .map(|col| { + Self::binary_result_is_histogram( + *op, + false, + Self::field_column_is_native_histogram(&input_schema, col), + ) + .is_some() + }) + .collect(); + let promql_annotations = self.promql_annotations.clone(); let bin_expr_builder = |col: &String| { let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let mut binary_expr = - binary_expr_builder(expr.clone(), DfExpr::Column(col.into()))?; + let rhs_is_histogram = + Self::field_column_is_native_histogram(&input_schema, col); + let rhs = DfExpr::Column(col.into()); + let mut binary_expr = match Self::native_histogram_binary_expr( + *op, + expr.clone(), + false, + rhs.clone(), + rhs_is_histogram, + is_comparison_op && !should_return_bool, + promql_annotations.clone(), + )? { + Some(expr) => expr, + None => binary_expr_builder(expr.clone(), rhs)?, + }; if is_comparison_op && should_return_bool { binary_expr = DfExpr::Cast(Cast { @@ -1230,7 +1293,14 @@ impl PromPlanner { if is_comparison_op && !should_return_bool { self.filter_on_field_column(input, bin_expr_builder) } else { - self.projection_for_each_field_column(input, bin_expr_builder) + let projected = + self.projection_for_each_field_column(input, bin_expr_builder)?; + self.filter_binary_projection( + projected, + has_native_histogram, + preserve_any_value, + retain_field_columns, + ) } } // lhs is a column, rhs is a literal @@ -1240,10 +1310,46 @@ impl PromPlanner { if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) { expr = time_expr } + let input_schema = input.schema().clone(); + let preserve_any_value = Self::field_columns_are_alternative_samples( + &input_schema, + &self.ctx.field_columns, + ); + let has_native_histogram = Self::field_columns_contain_native_histogram( + &input_schema, + &self.ctx.field_columns, + ); + let retain_field_columns = self + .ctx + .field_columns + .iter() + .map(|col| { + Self::binary_result_is_histogram( + *op, + Self::field_column_is_native_histogram(&input_schema, col), + false, + ) + .is_some() + }) + .collect(); + let promql_annotations = self.promql_annotations.clone(); let bin_expr_builder = |col: &String| { let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let mut binary_expr = - binary_expr_builder(DfExpr::Column(col.into()), expr.clone())?; + let lhs_is_histogram = + Self::field_column_is_native_histogram(&input_schema, col); + let lhs = DfExpr::Column(col.into()); + let mut binary_expr = match Self::native_histogram_binary_expr( + *op, + lhs.clone(), + lhs_is_histogram, + expr.clone(), + false, + is_comparison_op && !should_return_bool, + promql_annotations.clone(), + )? { + Some(expr) => expr, + None => binary_expr_builder(lhs, expr.clone())?, + }; if is_comparison_op && should_return_bool { binary_expr = DfExpr::Cast(Cast { @@ -1256,7 +1362,14 @@ impl PromPlanner { if is_comparison_op && !should_return_bool { self.filter_on_field_column(input, bin_expr_builder) } else { - self.projection_for_each_field_column(input, bin_expr_builder) + let projected = + self.projection_for_each_field_column(input, bin_expr_builder)?; + self.filter_binary_projection( + projected, + has_native_histogram, + preserve_any_value, + retain_field_columns, + ) } } // both are columns. join them on time index @@ -1291,6 +1404,14 @@ impl PromPlanner { ); } + let has_native_histogram = Self::field_columns_contain_native_histogram( + left_input.schema(), + &left_field_columns, + ) || Self::field_columns_contain_native_histogram( + right_input.schema(), + &right_field_columns, + ); + // normal join if left_table_ref == right_table_ref { // rename table references to avoid ambiguity @@ -1308,21 +1429,41 @@ impl PromPlanner { self.ctx.table_name = Some("rhs".to_string()); } } - let (output_field_columns, field_columns) = - Self::align_binary_field_columns(&left_field_columns, &right_field_columns); - let left_aligned_field_columns = field_columns + // Computed scalars reach this join path instead of the literal projection paths. + // Broadcast them for arithmetic in the same way as literal scalars. + let broadcast_scalar = !is_comparison_op; + let (field_groups, invalid_field_pairs) = Self::align_binary_field_columns( + left_input.schema(), + right_input.schema(), + &left_field_columns, + &right_field_columns, + *op, + broadcast_scalar && lhs.value_type() == ValueType::Scalar, + broadcast_scalar && rhs.value_type() == ValueType::Scalar, + ); + let left_aligned_field_columns = field_groups .iter() - .map(|(left_col_name, _)| (*left_col_name).clone()) + .flat_map(|(_, pairs)| { + pairs + .iter() + .map(|(left_col_name, _)| (*left_col_name).clone()) + }) .collect::>(); - let right_aligned_field_columns = field_columns + let right_aligned_field_columns = field_groups .iter() - .map(|(_, right_col_name)| (*right_col_name).clone()) + .flat_map(|(_, pairs)| { + pairs + .iter() + .map(|(_, right_col_name)| (*right_col_name).clone()) + }) .collect::>(); - // PromQL binary arithmetic only combines the shared prefix of value columns. - // Keep the output field count aligned with that zipped prefix so planning - // remains stable even when the two sides have uneven multi-field schemas. - self.ctx.field_columns = output_field_columns; - let mut field_columns = field_columns.into_iter(); + // Regular multi-field vectors combine their shared prefix. Alternative + // float/histogram lanes instead align by valid PromQL sample combinations. + self.ctx.field_columns = field_groups + .iter() + .map(|(output, _)| output.clone()) + .collect(); + let mut field_groups = field_groups.into_iter(); let join_plan = self.join_on_non_field_columns( left_input, @@ -1339,28 +1480,103 @@ impl PromPlanner { &right_context, )?; let join_plan_schema = join_plan.schema().clone(); + let promql_annotations = self.promql_annotations.clone(); + // These predicates always pass; they only evaluate otherwise-discarded pairs + // while collecting annotations. + let invalid_pair_predicates = invalid_field_pairs + .into_iter() + .filter(|_| promql_annotations.is_some()) + .map(|(left_col_name, right_col_name)| { + let left_field = join_plan_schema + .qualified_field_with_name(Some(&left_table_ref), left_col_name) + .context(DataFusionPlanningSnafu)?; + let right_field = join_plan_schema + .qualified_field_with_name(Some(&right_table_ref), right_col_name) + .context(DataFusionPlanningSnafu)?; + let left_is_histogram = + left_field.1.data_type() == &Self::native_histogram_arrow_type(); + let right_is_histogram = + right_field.1.data_type() == &Self::native_histogram_arrow_type(); + let drop_expr = Self::native_histogram_binary_expr( + *op, + DfExpr::Column(left_field.into()), + left_is_histogram, + DfExpr::Column(right_field.into()), + right_is_histogram, + true, + promql_annotations.clone(), + )? + .with_context(|| UnexpectedPlanExprSnafu { + desc: "invalid native histogram pair produced no drop expression", + })?; + Ok(DfExpr::Not(Box::new(drop_expr))) + }) + .collect::>>()?; + let join_plan = if let Some(predicate) = conjunction(invalid_pair_predicates) { + LogicalPlanBuilder::from(join_plan) + .filter(predicate) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)? + } else { + join_plan + }; let bin_expr_builder = |_: &String| { - let (left_col_name, right_col_name) = field_columns.next().unwrap(); - let left_col = join_plan_schema - .qualified_field_with_name(Some(&left_table_ref), left_col_name) - .context(DataFusionPlanningSnafu)? - .into(); - let right_col = join_plan_schema - .qualified_field_with_name(Some(&right_table_ref), right_col_name) - .context(DataFusionPlanningSnafu)? - .into(); + let (_, field_pairs) = + field_groups + .next() + .with_context(|| UnexpectedPlanExprSnafu { + desc: "missing binary field group", + })?; + let binary_exprs = field_pairs + .into_iter() + .map(|(left_col_name, right_col_name)| { + let left_field = join_plan_schema + .qualified_field_with_name(Some(&left_table_ref), left_col_name) + .context(DataFusionPlanningSnafu)?; + let right_field = join_plan_schema + .qualified_field_with_name(Some(&right_table_ref), right_col_name) + .context(DataFusionPlanningSnafu)?; + let left_is_histogram = + left_field.1.data_type() == &Self::native_histogram_arrow_type(); + let right_is_histogram = + right_field.1.data_type() == &Self::native_histogram_arrow_type(); + let left_col = left_field.into(); + let right_col = right_field.into(); - let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let mut binary_expr = - binary_expr_builder(DfExpr::Column(left_col), DfExpr::Column(right_col))?; - if is_comparison_op && should_return_bool { - binary_expr = DfExpr::Cast(Cast { - expr: Box::new(binary_expr), - data_type: ArrowDataType::Float64, - }); + let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; + let lhs = DfExpr::Column(left_col); + let rhs = DfExpr::Column(right_col); + let mut binary_expr = match Self::native_histogram_binary_expr( + *op, + lhs.clone(), + left_is_histogram, + rhs.clone(), + right_is_histogram, + is_comparison_op && !should_return_bool, + promql_annotations.clone(), + )? { + Some(expr) => expr, + None => binary_expr_builder(lhs, rhs)?, + }; + if is_comparison_op && should_return_bool { + binary_expr = DfExpr::Cast(Cast { + expr: Box::new(binary_expr), + data_type: ArrowDataType::Float64, + }); + } + Ok(binary_expr) + }) + .collect::>>()?; + if let [binary_expr] = binary_exprs.as_slice() { + Ok(binary_expr.clone()) + } else { + Ok(DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: binary_exprs, + })) } - Ok(binary_expr) }; if is_comparison_op && !should_return_bool { // PromQL comparison operators without `bool` are filters: @@ -1386,12 +1602,90 @@ impl PromPlanner { project_context.field_columns = project_field_columns; self.project_binary_join_side(filtered, project_table_ref, &project_context) } else { - self.projection_for_each_field_column(join_plan, bin_expr_builder) + let projected = + self.projection_for_each_field_column(join_plan, bin_expr_builder)?; + let preserve_any_value = Self::field_columns_are_alternative_samples( + projected.schema(), + &self.ctx.field_columns, + ); + let retain_field_columns = vec![true; self.ctx.field_columns.len()]; + self.filter_binary_projection( + projected, + has_native_histogram, + preserve_any_value, + retain_field_columns, + ) } } } } + fn filter_binary_projection( + &mut self, + input: LogicalPlan, + has_native_histogram: bool, + preserve_any_value: bool, + retain_field_columns: Vec, + ) -> Result { + if !has_native_histogram { + return Ok(input); + } + + ensure!( + retain_field_columns.len() == self.ctx.field_columns.len(), + UnexpectedPlanExprSnafu { + desc: "binary output field count changed unexpectedly", + } + ); + + let filtered = LogicalPlanBuilder::from(input) + .filter(self.create_empty_values_filter_expr(preserve_any_value)?) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + if retain_field_columns.iter().all(|retain| *retain) { + return Ok(filtered); + } + + let retained = self + .ctx + .field_columns + .iter() + .zip(retain_field_columns) + .filter(|(_, retain)| *retain) + .map(|(field, _)| field.clone()) + .collect::>(); + if retained.is_empty() { + return Ok(filtered); + } + self.ctx.field_columns = retained; + + let mut output_columns = self + .ctx + .field_columns + .iter() + .chain(&self.ctx.tag_columns) + .cloned() + .collect::>(); + output_columns.extend(self.ctx.time_index_column.iter().cloned()); + if self.ctx.use_tsid { + output_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string()); + } + let project_exprs = filtered + .schema() + .iter() + .filter(|(_, field)| output_columns.contains(field.name())) + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + LogicalPlanBuilder::from(filtered) + .project(project_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + fn project_binary_join_side( &mut self, input: LogicalPlan, @@ -4416,6 +4710,69 @@ impl PromPlanner { } } + fn native_histogram_binary_expr( + token: TokenType, + lhs: DfExpr, + lhs_is_histogram: bool, + rhs: DfExpr, + rhs_is_histogram: bool, + filter_context: bool, + promql_annotations: Option, + ) -> Result> { + if !lhs_is_histogram && !rhs_is_histogram { + return Ok(None); + } + + let scalar_fn = |func: ScalarUdfDef, args| { + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(func), + args, + }) + }; + let invalid_expr = || { + let message = format!( + "{}: dropped native histogram samples because this binary operation is not supported for native histograms", + token + ); + let func = if filter_context { + NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone()) + } else { + NativeHistogramDrop::float_null_udf(message, promql_annotations.clone()) + }; + let args = vec![lhs.clone(), rhs.clone()]; + scalar_fn(func, args) + }; + + let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) { + (token::T_ADD, true, true) => scalar_fn( + NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()), + vec![lhs, rhs], + ), + (token::T_SUB, true, true) => scalar_fn( + NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()), + vec![lhs, rhs], + ), + (token::T_MUL, true, false) => { + scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs]) + } + (token::T_MUL, false, true) => { + scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs]) + } + (token::T_DIV, true, false) => { + scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs]) + } + (token::T_EQLC, true, true) => { + scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs]) + } + (token::T_NEQ, true, true) => { + scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs]) + } + _ => invalid_expr(), + }; + + Ok(Some(expr)) + } + /// Return a lambda to build binary expression from token. /// Because some binary operator are function in DataFusion like `atan2` or `^`. #[allow(clippy::type_complexity)] @@ -4501,18 +4858,117 @@ impl PromPlanner { } fn align_binary_field_columns<'a>( + left_schema: &DFSchemaRef, + right_schema: &DFSchemaRef, left_field_columns: &'a [String], right_field_columns: &'a [String], - ) -> (Vec, Vec<(&'a String, &'a String)>) { - let field_pairs = left_field_columns - .iter() - .zip(right_field_columns.iter()) - .collect::>(); - let output_field_columns = field_pairs - .iter() - .map(|(left_col_name, _)| (*left_col_name).clone()) - .collect(); - (output_field_columns, field_pairs) + op: TokenType, + left_is_scalar: bool, + right_is_scalar: bool, + ) -> ( + Vec<(String, Vec>)>, + Vec>, + ) { + // Mixed vectors store mutually exclusive float and histogram samples in two columns. + // Retain each valid sample combination and group expressions by their output lane. + let left_alternative = Self::alternative_sample_columns(left_schema, left_field_columns); + let right_alternative = Self::alternative_sample_columns(right_schema, right_field_columns); + let alternative_alignment = match (left_alternative, right_alternative) { + (Some(output_names), Some(_)) => Some(( + output_names, + left_field_columns + .iter() + .flat_map(|left| right_field_columns.iter().map(move |right| (left, right))) + .collect::>(), + )), + (Some(output_names), None) if right_field_columns.len() == 1 => Some(( + output_names, + left_field_columns + .iter() + .map(|left| (left, &right_field_columns[0])) + .collect::>(), + )), + (None, Some(output_names)) if left_field_columns.len() == 1 => Some(( + output_names, + right_field_columns + .iter() + .map(|right| (&left_field_columns[0], right)) + .collect::>(), + )), + _ => None, + }; + let mut invalid_pairs = Vec::new(); + if let Some(((float_output, histogram_output), field_pairs)) = alternative_alignment { + let mut float_pairs = Vec::new(); + let mut histogram_pairs = Vec::new(); + for (left, right) in field_pairs { + let left_is_histogram = Self::field_column_is_native_histogram(left_schema, left); + let right_is_histogram = + Self::field_column_is_native_histogram(right_schema, right); + match Self::binary_result_is_histogram(op, left_is_histogram, right_is_histogram) { + Some(false) => float_pairs.push((left, right)), + Some(true) => histogram_pairs.push((left, right)), + None => invalid_pairs.push((left, right)), + } + } + if !float_pairs.is_empty() || !histogram_pairs.is_empty() { + return ( + [ + (!float_pairs.is_empty()).then(|| (float_output.to_string(), float_pairs)), + (!histogram_pairs.is_empty()) + .then(|| (histogram_output.to_string(), histogram_pairs)), + ] + .into_iter() + .flatten() + .collect(), + invalid_pairs, + ); + } + } + + if left_is_scalar && !right_is_scalar && left_field_columns.len() == 1 { + return ( + right_field_columns + .iter() + .map(|right| (right.clone(), vec![(&left_field_columns[0], right)])) + .collect(), + invalid_pairs, + ); + } + if right_is_scalar && !left_is_scalar && right_field_columns.len() == 1 { + return ( + left_field_columns + .iter() + .map(|left| (left.clone(), vec![(left, &right_field_columns[0])])) + .collect(), + invalid_pairs, + ); + } + + ( + left_field_columns + .iter() + .zip(right_field_columns.iter()) + .map(|(left, right)| (left.clone(), vec![(left, right)])) + .collect(), + invalid_pairs, + ) + } + + fn binary_result_is_histogram( + token: TokenType, + lhs_is_histogram: bool, + rhs_is_histogram: bool, + ) -> Option { + match (token.id(), lhs_is_histogram, rhs_is_histogram) { + (_, false, false) => Some(false), + (token::T_ADD | token::T_SUB, true, true) + | (token::T_MUL, true, false) + | (token::T_MUL, false, true) + | (token::T_DIV, true, false) => Some(true), + (token::T_EQLC | token::T_NEQ, true, true) => Some(false), + _ => None, + } } fn plan_has_tsid_column(plan: &LogicalPlan) -> bool { @@ -4540,6 +4996,15 @@ impl PromPlanner { .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type()) } + fn field_columns_contain_native_histogram( + schema: &DFSchemaRef, + field_columns: &[String], + ) -> bool { + field_columns + .iter() + .any(|field| Self::field_column_is_native_histogram(schema, field)) + } + fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool { Self::field_column_type(schema, field_column).is_some_and(|data_type| { matches!( @@ -4940,9 +5405,13 @@ impl PromPlanner { } ensure!( - left_context.field_columns.len() == 1, + left_context.field_columns.len() == 1 + || Self::field_columns_are_alternative_samples( + left.schema(), + &left_context.field_columns, + ), MultiFieldsNotSupportedSnafu { - operator: "AND operator" + operator: "AND/UNLESS operator" } ); // Generate join plan. @@ -5072,24 +5541,28 @@ impl PromPlanner { (false, false) => {} } - // checks ensure!( - left_context.field_columns.len() == right_context.field_columns.len(), - CombineTableColumnMismatchSnafu { - left: left_context.field_columns.clone(), - right: right_context.field_columns.clone() + !left.schema().fields().is_empty() && !right.schema().fields().is_empty(), + UnexpectedPlanExprSnafu { + desc: "OR operator input has zero columns", } ); + let left_has_alternative_samples = + Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns); + let right_has_alternative_samples = Self::field_columns_are_alternative_samples( + right.schema(), + &right_context.field_columns, + ); ensure!( - left_context.field_columns.len() == 1, + left_context.field_columns.len() == 1 || left_has_alternative_samples, MultiFieldsNotSupportedSnafu { operator: "OR operator" } ); ensure!( - !left.schema().fields().is_empty() && !right.schema().fields().is_empty(), - UnexpectedPlanExprSnafu { - desc: "OR operator input has zero columns", + right_context.field_columns.len() == 1 || right_has_alternative_samples, + MultiFieldsNotSupportedSnafu { + operator: "OR operator" } ); @@ -5122,62 +5595,125 @@ impl PromPlanner { .with_context(|| TimeIndexNotFoundSnafu { table: right_qualifier_string.clone(), })?; - // Take the name of first field column. The length is checked above. - let left_field_col = left_context.field_columns.first().unwrap(); - let right_field_col = right_context.field_columns.first().unwrap(); - let left_field = left - .schema() + let native_histogram_type = Self::native_histogram_arrow_type(); + let is_numeric = |data_type: &ArrowDataType| { + matches!( + data_type, + ArrowDataType::Int8 + | ArrowDataType::Int16 + | ArrowDataType::Int32 + | ArrowDataType::Int64 + | ArrowDataType::UInt8 + | ArrowDataType::UInt16 + | ArrowDataType::UInt32 + | ArrowDataType::UInt64 + | ArrowDataType::Float32 + | ArrowDataType::Float64 + ) + }; + let left_fields = left_context + .field_columns .iter() - .find(|(_, field)| field.name() == left_field_col) - .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { - col: left_field_col.clone(), - })?; - let right_field = right - .schema() + .map(|name| { + left.schema() + .iter() + .find(|(_, field)| field.name() == name) + .map(|(qualifier, field)| { + (name.clone(), qualifier.cloned(), field.data_type().clone()) + }) + .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) + }) + .collect::>>()?; + let right_fields = right_context + .field_columns .iter() - .find(|(_, field)| field.name() == right_field_col) - .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { - col: right_field_col.clone(), - })?; - let target_field_type = if left_field.1 == right_field.1 { - left_field.1.clone() - } else if matches!( - left_field.1, - ArrowDataType::Int8 - | ArrowDataType::Int16 - | ArrowDataType::Int32 - | ArrowDataType::Int64 - | ArrowDataType::UInt8 - | ArrowDataType::UInt16 - | ArrowDataType::UInt32 - | ArrowDataType::UInt64 - | ArrowDataType::Float32 - | ArrowDataType::Float64 - ) && matches!( - right_field.1, - ArrowDataType::Int8 - | ArrowDataType::Int16 - | ArrowDataType::Int32 - | ArrowDataType::Int64 - | ArrowDataType::UInt8 - | ArrowDataType::UInt16 - | ArrowDataType::UInt32 - | ArrowDataType::UInt64 - | ArrowDataType::Float32 - | ArrowDataType::Float64 - ) { + .map(|name| { + right + .schema() + .iter() + .find(|(_, field)| field.name() == name) + .map(|(qualifier, field)| { + (name.clone(), qualifier.cloned(), field.data_type().clone()) + }) + .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) + }) + .collect::>>()?; + let left_field = &left_fields[0]; + let right_field = &right_fields[0]; + let left_field_col = &left_field.0; + let right_field_col = &right_field.0; + let fields_are_samples = |fields: &[(String, Option, ArrowDataType)]| { + fields.iter().all(|(_, _, data_type)| { + is_numeric(data_type) || data_type == &native_histogram_type + }) + }; + let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples { + if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR value fields have incompatible types: {:?} and {:?}", + left_fields + .iter() + .map(|(_, _, data_type)| data_type) + .collect::>(), + right_fields + .iter() + .map(|(_, _, data_type)| data_type) + .collect::>() + ), + } + .fail(); + } + true + } else { + (left_field.2 == native_histogram_type && is_numeric(&right_field.2)) + || (right_field.2 == native_histogram_type && is_numeric(&left_field.2)) + }; + let target_field_type = if mixed_sample_types { + // Mixed vectors use the existing response representation: one nullable float column + // and one nullable native-histogram column. + ArrowDataType::Float64 + } else if left_field.2 == right_field.2 { + left_field.2.clone() + } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) { ArrowDataType::Float64 } else { return UnexpectedPlanExprSnafu { desc: format!( "OR value fields have incompatible types: {:?} and {:?}", - left_field.1, right_field.1 + left_field.2, right_field.2 ), } .fail(); }; + let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types { + let mut reserved_names = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); + for (name, _, _) in left_fields.iter().chain(&right_fields) { + reserved_names.remove(name); + } + reserved_names.extend(all_tags.iter().cloned()); + let unique_name = |prefix: &str, reserved_names: &mut HashSet| { + let mut index = 0; + loop { + let name = format!("{prefix}{index}"); + index += 1; + if reserved_names.insert(name.clone()) { + break name; + } + } + }; + let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names); + let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names); + (float_field, histogram_field) + } else { + (left_field_col.clone(), String::new()) + }; let left_tag_types = left_tag_cols_set .iter() .map(|label| { @@ -5244,8 +5780,15 @@ impl PromPlanner { // remove time index column all_columns_set.remove(&left_time_index_column); all_columns_set.remove(&right_time_index_column); - // remove field column in the right - if left_field_col != right_field_col { + if mixed_sample_types { + for (name, _, _) in left_fields.iter().chain(&right_fields) { + all_columns_set.remove(name); + } + all_columns_set.extend(all_tags.iter().cloned()); + all_columns_set.insert(mixed_float_field_col.clone()); + all_columns_set.insert(mixed_histogram_field_col.clone()); + } else if left_field_col != right_field_col { + // remove field column in the right all_columns_set.remove(right_field_col); } let mut all_columns = all_columns_set.into_iter().collect::>(); @@ -5284,11 +5827,52 @@ impl PromPlanner { .alias(col.clone()) } }; + let null_histogram = + ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?; + let mixed_value_expr = |fields: &[(String, Option, ArrowDataType)], + output_col: &String| { + if output_col == &mixed_float_field_col { + if let Some((name, qualifier, data_type)) = fields + .iter() + .find(|(_, _, data_type)| is_numeric(data_type)) + { + let expr = DfExpr::Column(Column::new(qualifier.clone(), name)); + if data_type == &ArrowDataType::Float64 { + expr.alias(output_col) + } else { + DfExpr::Cast(Cast { + expr: Box::new(expr), + data_type: ArrowDataType::Float64, + }) + .alias(output_col) + } + } else { + DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col) + } + } else { + fields + .iter() + .find(|(_, _, data_type)| data_type == &native_histogram_type) + .map(|(name, qualifier, _)| { + DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col) + }) + .unwrap_or_else(|| { + DfExpr::Literal(null_histogram.clone(), None).alias(output_col) + }) + } + }; let left_proj_exprs = all_columns.iter().map(|col| { - if col == left_field_col && left_field.1 != target_field_type { + if mixed_sample_types + && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) + { + mixed_value_expr(&left_fields, col) + } else if !mixed_sample_types + && col == left_field_col + && left_field.2 != target_field_type + { DfExpr::Cast(Cast { expr: Box::new(DfExpr::Column(Column::new( - left_field.0.clone(), + left_field.1.clone(), left_field_col, ))), data_type: target_field_type.clone(), @@ -5310,9 +5894,13 @@ impl PromPlanner { // `skip（1)` to skip the time index column let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| { // expr - if col == left_field_col { - let expr = DfExpr::Column(Column::new(right_field.0.clone(), right_field_col)); - if right_field.1 != target_field_type { + if mixed_sample_types + && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) + { + mixed_value_expr(&right_fields, col) + } else if !mixed_sample_types && col == left_field_col { + let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col)); + if right_field.2 != target_field_type { DfExpr::Cast(Cast { expr: Box::new(expr), data_type: target_field_type.clone(), @@ -5529,7 +6117,11 @@ impl PromPlanner { visible_tags.sort_unstable(); output_context.time_index_column = Some(left_time_index_column); output_context.tag_columns = visible_tags; - output_context.field_columns = vec![output_field_col]; + output_context.field_columns = if mixed_sample_types { + vec![mixed_float_field_col, mixed_histogram_field_col] + } else { + vec![output_field_col] + }; output_context.use_tsid = left_has_tsid && right_has_tsid; self.ctx = output_context; @@ -5551,6 +6143,10 @@ impl PromPlanner { where F: FnMut(&String) -> Result, { + // Keep the generated float/histogram lane names while an element-wise operation + // preserves both sample types, so downstream operators still recognize the pair. + let preserve_field_names = + Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns); let table_ref = self.ctx.table_name.clone().map(TableReference::bare); let non_field_columns_iter = self .ctx @@ -5572,10 +6168,12 @@ impl PromPlanner { .collect::>>()?; // alias the computation exprs to remove qualifier - self.ctx.field_columns = result_field_columns - .iter() - .map(|expr| expr.schema_name().to_string()) - .collect(); + if !preserve_field_names { + self.ctx.field_columns = result_field_columns + .iter() + .map(|expr| expr.schema_name().to_string()) + .collect(); + } let field_columns_iter = result_field_columns .into_iter() .zip(self.ctx.field_columns.iter()) @@ -5594,24 +6192,32 @@ impl PromPlanner { .context(DataFusionPlanningSnafu) } - /// Build a filter plan that filter on value column. Notice that only one value column - /// is expected. - fn filter_on_field_column( - &self, - input: LogicalPlan, - mut name_to_expr: F, - ) -> Result + /// Build a filter plan on one value column or a float/histogram alternative pair. + fn filter_on_field_column(&self, input: LogicalPlan, name_to_expr: F) -> Result where F: FnMut(&String) -> Result, { ensure!( - self.ctx.field_columns.len() == 1, + self.ctx.field_columns.len() == 1 + || Self::field_columns_are_alternative_samples( + input.schema(), + &self.ctx.field_columns, + ), UnsupportedExprSnafu { name: "filter on multi-value input" } ); - let field_column_filter = name_to_expr(&self.ctx.field_columns[0])?; + let field_column_filters = self + .ctx + .field_columns + .iter() + .map(name_to_expr) + .collect::>>()?; + let field_column_filter = + disjunction(field_column_filters).context(UnsupportedExprSnafu { + name: "filter on empty input", + })?; LogicalPlanBuilder::from(input) .filter(field_column_filter) @@ -5743,7 +6349,9 @@ mod test { CounterResetHint, NativeHistogram, build_histogram_array, }; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; + use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use common_query::test_util::DummyDecoder; + use common_recordbatch::RecordBatch as GreptimeRecordBatch; use datafusion::arrow::array::{ Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, }; @@ -5761,8 +6369,9 @@ mod test { use promql_parser::parser; use session::context::QueryContext; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; - use table::metadata::{TableInfoBuilder, TableMetaBuilder}; - use table::test_util::EmptyTable; + use table::Table; + use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder}; + use table::test_util::{EmptyTable, MemTable as GreptimeMemTable}; use super::*; use crate::QueryEngineContext; @@ -5896,6 +6505,7 @@ mod test { enum DirectOrValue { Float64(f64), Int64(i64), + NativeHistogram(NativeHistogram), Utf8(&'static str), } @@ -5904,6 +6514,7 @@ mod test { match self { Self::Float64(_) => ArrowDataType::Float64, Self::Int64(_) => ArrowDataType::Int64, + Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(), Self::Utf8(_) => ArrowDataType::Utf8, } } @@ -5911,6 +6522,7 @@ mod test { match self { Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])), Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])), + Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]), Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])), } } @@ -5933,6 +6545,121 @@ mod test { } } + fn operator_metric_table( + name: &str, + table_id: u32, + tag: &str, + value: DirectOrValue, + ) -> table::TableRef { + let value_type = match &value { + DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(), + DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(), + DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(), + DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(), + }; + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("v".to_string(), value_type, true), + ])); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from(vec![tag])) as Arc, + Arc::new(TimestampMillisecondArray::from(vec![1_000])), + value.array(), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + name, + GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), + table_id, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(3) + .build() + .unwrap(); + let info = Arc::new( + TableInfoBuilder::default() + .table_id(table_id) + .name(name) + .meta(meta) + .build() + .unwrap(), + ); + Arc::new(Table::new( + info, + FilterPushDownType::Unsupported, + backing.data_source(), + )) + } + + fn operator_table_provider() -> DfTableSourceProvider { + let catalog = MemoryCatalogManager::with_default_setup(); + let tables = [ + operator_metric_table("lf", 2_001, "a", DirectOrValue::Float64(2.0)), + operator_metric_table( + "lh", + 2_002, + "b", + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("rf", 2_003, "b", DirectOrValue::Float64(3.0)), + operator_metric_table( + "rh", + 2_004, + "a", + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("fallback", 2_005, "c", DirectOrValue::Float64(7.0)), + ]; + for table in tables { + let info = table.table_info(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: info.name.clone(), + table_id: info.ident.table_id, + table, + }) + .unwrap(); + } + DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) + } + + fn operator_eval_stmt(expr: &str) -> EvalStmt { + let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap(); + EvalStmt { + expr: parser::parse(expr).unwrap(), + start: time, + end: time, + interval: Duration::from_secs(1), + lookback_delta: Duration::from_secs(5), + } + } + struct DirectOrSource { name: &'static str, empty: bool, @@ -6093,6 +6820,66 @@ mod test { execute(plan, &build_query_engine_state()).await } + async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) { + let sample = |histogram: bool| { + if histogram { + DirectOrValue::NativeHistogram(direct_or_histogram()) + } else { + DirectOrValue::Float64(1.25) + } + }; + let left = tagged_source( + "lhs", + false, + ( + "k", + Some(if histogram_on_left { + "histogram" + } else { + "float" + }), + ), + sample(histogram_on_left), + ); + let right = tagged_source( + "rhs", + false, + ( + "k", + Some(if histogram_on_left { + "float" + } else { + "histogram" + }), + ), + sample(!histogram_on_left), + ); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left_context = direct_or_context("lhs", &["job", "k"], "v"); + let right_context = direct_or_context("rhs", &["job", "k"], "v"); + let plan = planner + .or_operator( + scan(&left), + scan(&right), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on(k) rhs"), + ) + .unwrap(); + (planner, plan) + } + fn assert_no_internal_or_keys(schema: &DFSchema) { assert!( schema @@ -6119,6 +6906,23 @@ mod test { .collect() } + fn histograms(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + let values = batch + .column_by_name(column) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..values.len()).filter_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .collect() + } + fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option)> { let mut rows = batches .iter() @@ -8833,6 +9637,110 @@ mod test { assert!(!plan.contains("PromHistogramFold"), "{plan}"); } + #[tokio::test] + async fn timestamp_filters_native_histogram_stale_marker_before_projection() { + let mut stale = direct_or_histogram(); + stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); + let table = operator_metric_table( + "stale_histogram", + 2_100, + "a", + DirectOrValue::NativeHistogram(stale), + ); + let catalog = MemoryCatalogManager::with_default_setup(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "stale_histogram".to_string(), + table_id: 2_100, + table, + }) + .unwrap(); + let provider = DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + provider, + &operator_eval_stmt("timestamp(stale_histogram)"), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}"); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + } + + #[tokio::test] + async fn timestamp_filters_stale_marker_from_mixed_sample_companion() { + let histograms = build_histogram_array(&[None]); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new( + greptime_native_histogram(), + histograms.data_type().clone(), + true, + ), + Field::new(greptime_value(), ArrowDataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000])), + histograms, + Arc::new(Float64Array::from(vec![f64::from_bits( + PROMETHEUS_STALE_NAN_BITS, + )])), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let input = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesDivide::new( + Vec::new(), + "timestamp".to_string(), + input, + )), + }); + let input = LogicalPlan::Extension(Extension { + node: Arc::new(InstantManipulate::new( + 1_000, + 1_000, + 5_000, + 1_000, + "timestamp".to_string(), + Vec::new(), + Some(greptime_native_histogram().to_string()), + input, + )), + }); + // Match timestamp()'s parent projection, which otherwise prunes the companion lane. + let plan = LogicalPlanBuilder::from(input) + .project([col("timestamp")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + } + #[tokio::test] async fn native_histogram_rate_can_feed_count() { let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await; @@ -10654,7 +11562,6 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom .contains("OR value fields have incompatible types") ); } - #[tokio::test] async fn test_or_with_histogram_quantile_missing_le_column() { let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#; @@ -10871,4 +11778,572 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom "{plan:?}" ); } + + #[tokio::test] + async fn test_direct_or_preserves_float_and_native_histogram_samples() { + for histogram_on_left in [false, true] { + let (planner, plan) = mixed_direct_or(histogram_on_left).await; + + let float_field = &planner.ctx.field_columns[0]; + let histogram_field = &planner.ctx.field_columns[1]; + assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX)); + assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)); + assert_eq!( + plan.schema() + .field_with_name(None, float_field) + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!( + plan.schema() + .field_with_name(None, histogram_field) + .unwrap() + .data_type(), + &native_histogram_value_type().as_arrow_type() + ); + + let (optimized, batches) = execute(plan, &build_query_engine_state()).await; + assert_no_internal_or_keys(optimized.schema()); + let mut sample_kinds = batches + .iter() + .flat_map(|batch| { + let values = batch.column_by_name(float_field).unwrap(); + let histograms = batch.column_by_name(histogram_field).unwrap(); + (0..batch.num_rows()) + .map(|row| (values.is_valid(row), histograms.is_valid(row))) + }) + .collect::>(); + sample_kinds.sort_unstable(); + assert_eq!(sample_kinds, vec![(false, true), (true, false)]); + } + } + + #[tokio::test] + async fn test_mixed_binary_operator_aligns_both_alternative_inputs() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!( + plan_text.contains("prom_native_histogram_mul_scalar"), + "{plan_text}" + ); + assert!( + plan_text.contains("prom_native_histogram_scalar_mul"), + "{plan_text}" + ); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + assert!(values(&batches, &float_field).is_empty()); + let mut sums = histograms(&batches, &histogram_field) + .into_iter() + .map(|histogram| histogram.sum) + .collect::>(); + sums.sort_by(f64::total_cmp); + assert_eq!(sums, vec![2.0, 3.0]); + } + + #[tokio::test] + async fn test_mixed_binary_operator_reports_only_dropped_samples() { + for (query, expected_rows, expected_infos) in [ + ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1), + ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0), + ("(lf or on(tag) lh) % on(tag) lh", 0, 1), + ] { + let state = build_query_engine_state(); + let annotations = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt(query), + &state, + Some(annotations.clone()), + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected_rows, + "{query}" + ); + let mut warnings = vec![]; + let mut infos = vec![]; + annotations.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty(), "{query}: {warnings:?}"); + assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}"); + } + } + + #[tokio::test] + async fn test_mixed_or_can_feed_another_or() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + let mut float_values = values(&batches, &float_field); + float_values.sort_by(f64::total_cmp); + assert_eq!(float_values, vec![2.0, 7.0]); + assert_eq!(histograms(&batches, &histogram_field).len(), 1); + } + + #[tokio::test] + async fn test_mixed_fields_align_with_single_float_vector() { + let (planner, mixed) = mixed_direct_or(false).await; + let scale = tagged_source( + "scale", + false, + ("k", Some("float")), + DirectOrValue::Float64(2.0), + ); + let scale = scan(&scale); + let scale_fields = vec!["v".to_string()]; + let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else { + unreachable!() + }; + + let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( + mixed.schema(), + scale.schema(), + &planner.ctx.field_columns, + &scale_fields, + binary.op, + false, + false, + ); + assert!(invalid_pairs.is_empty()); + assert_eq!( + groups + .iter() + .map(|(output, _)| output.clone()) + .collect::>(), + planner.ctx.field_columns + ); + assert_eq!(groups.len(), 2); + assert!( + groups + .iter() + .flat_map(|(_, pairs)| pairs) + .all(|(_, right)| *right == &scale_fields[0]) + ); + + let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( + scale.schema(), + mixed.schema(), + &scale_fields, + &planner.ctx.field_columns, + binary.op, + false, + false, + ); + assert!(invalid_pairs.is_empty()); + assert_eq!( + groups + .iter() + .map(|(output, _)| output.clone()) + .collect::>(), + planner.ctx.field_columns + ); + assert_eq!(groups.len(), 2); + assert!( + groups + .iter() + .flat_map(|(_, pairs)| pairs) + .all(|(left, _)| *left == &scale_fields[0]) + ); + } + + #[tokio::test] + async fn test_non_bool_comparison_filters_mixed_sample_lanes() { + let (planner, input) = mixed_direct_or(false).await; + let input_schema = input.schema().clone(); + let plan = planner + .filter_on_field_column(input, |field| { + if PromPlanner::field_column_is_native_histogram(&input_schema, field) { + Ok(lit(false)) + } else { + Ok(col(field).gt(lit(0.0))) + } + }) + .unwrap(); + let float_field = planner.ctx.field_columns[0].clone(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_field), vec![1.25]); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + } + + #[tokio::test] + async fn test_mixed_left_and_unless_preserve_sample_lanes() { + for (expression, expected_sample_kind) in [ + ("lhs and on(k) mask", (false, true)), + ("lhs unless on(k) mask", (true, false)), + ] { + let (mut planner, left) = mixed_direct_or(false).await; + let left_context = planner.ctx.clone(); + let float_field = left_context.field_columns[0].clone(); + let histogram_field = left_context.field_columns[1].clone(); + let mask = tagged_source( + "mask", + false, + ("k", Some("histogram")), + DirectOrValue::Float64(1.0), + ); + let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else { + unreachable!() + }; + let plan = planner + .set_op_on_non_field_columns( + left, + scan(&mask), + left_context, + direct_or_context("mask", &["job", "k"], "v"), + binary.op, + &binary.modifier, + ) + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let sample_kinds = batches + .iter() + .flat_map(|batch| { + let floats = batch.column_by_name(&float_field).unwrap(); + let histograms = batch.column_by_name(&histogram_field).unwrap(); + (0..batch.num_rows()) + .map(|row| (floats.is_valid(row), histograms.is_valid(row))) + }) + .collect::>(); + assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}"); + } + } + + #[tokio::test] + async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt("some_metric * scalar(vector(2))"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let schema = plan.schema(); + assert_eq!( + schema + .field_with_unqualified_name(greptime_value()) + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!( + schema + .field_with_unqualified_name(greptime_native_histogram()) + .unwrap() + .data_type(), + &native_histogram_value_type().as_arrow_type() + ); + assert!( + plan.display_indent_schema() + .to_string() + .contains("prom_native_histogram_mul_scalar"), + "{plan:?}" + ); + } + + #[tokio::test] + async fn test_unsupported_histogram_binary_does_not_block_or_fallback() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type()) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(values(&batches, &float_field), vec![0.0]); + assert_eq!(histograms(&batches, &histogram_field).len(), 1); + } + + #[tokio::test] + async fn test_unary_negates_mixed_float_and_native_histogram_samples() { + for histogram_on_left in [false, true] { + let (mut planner, input) = mixed_direct_or(histogram_on_left).await; + let plan = planner.negate_field_columns(input).unwrap(); + assert!(PromPlanner::field_columns_are_alternative_samples( + plan.schema(), + &planner.ctx.field_columns + )); + let float_field = planner + .ctx + .field_columns + .iter() + .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap(); + let histogram_field = planner + .ctx + .field_columns + .iter() + .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, float_field), vec![-1.25]); + let histogram = batches + .iter() + .find_map(|batch| { + let values = batch + .column_by_name(histogram_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..values.len()).find_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .unwrap(); + assert_eq!(histogram.count, -1.0); + assert_eq!(histogram.sum, -1.0); + assert_eq!(histogram.reset_hint, CounterResetHint::Gauge); + } + } + + #[tokio::test] + async fn test_mixed_or_value_aliases_do_not_replace_labels() { + let left = source( + "lhs", + false, + 1, + vec![("job", Some("job")), ("k", Some("float"))], + DirectOrValue::Float64(1.0), + ); + let right = source( + "rhs", + false, + 1, + vec![ + ("job", Some("job")), + ("k", Some("histogram")), + (greptime_value(), Some("value-label")), + ], + DirectOrValue::NativeHistogram(direct_or_histogram()), + ); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left = LogicalPlanBuilder::from(scan(&left)) + .project(vec![ + col("ts"), + col("job"), + col("k"), + col("v").alias(greptime_value()), + ]) + .unwrap() + .build() + .unwrap(); + let left_context = direct_or_context("lhs", &["job", "k"], greptime_value()); + let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v"); + let plan = planner + .or_operator( + left, + scan(&right), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on(k) rhs"), + ) + .unwrap(); + + assert_eq!( + plan.schema() + .field_with_name(None, greptime_value()) + .unwrap() + .data_type(), + &ArrowDataType::Utf8 + ); + assert!( + planner + .ctx + .field_columns + .iter() + .all(|field| { field != greptime_value() && field != greptime_native_histogram() }) + ); + assert!(PromPlanner::field_columns_are_alternative_samples( + plan.schema(), + &planner.ctx.field_columns + )); + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let labels = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(greptime_value()) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .collect::>(); + assert_eq!(labels, vec!["value-label"]); + } + + #[tokio::test] + async fn test_mixed_or_routes_float_histogram_and_label_functions() { + for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] { + let (mut planner, input) = mixed_direct_or(false).await; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else { + unreachable!() + }; + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, vec![], input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + exprs.extend(planner.create_tag_column_exprs().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + let values = batches + .iter() + .flat_map(|batch| { + batch + .schema() + .fields() + .iter() + .position(|field| field.data_type() == &ArrowDataType::Float64) + .map(|index| { + batch + .column(index) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .into_iter() + .flatten() + }) + .collect::>(); + assert_eq!(values, vec![expected], "{function}"); + } + + let (mut planner, input) = mixed_direct_or(false).await; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let PromExpr::Call(call) = + parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap() + else { + unreachable!() + }; + let args = planner.create_function_args(&call.args.args).unwrap(); + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, args.literals, input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + exprs.extend(planner.create_tag_column_exprs().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + let sample_count = batches.iter().map(RecordBatch::num_rows).sum::(); + assert_eq!(sample_count, 2); + } } From 7375be06352146e84d8eab89313eabb341cb92c3 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Tue, 11 Aug 2026 16:27:03 +0800 Subject: [PATCH 047/126] ci: update code owners (#8843) Signed-off-by: evenyag --- .github/CODEOWNERS | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 549d6300ea1..e9e4d6ea561 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -6,8 +6,8 @@ ## [Module] Database Engine /src/index @evenyag @discord9 @WenyXu -/src/mito2 @evenyag @v0y4g3r @waynexia -/src/query @evenyag @waynexia @discord9 +/src/mito2 @evenyag @v0y4g3r +/src/query @evenyag @discord9 ## [Module] Distributed /src/common/meta @MichaelScofield @WenyXu @@ -20,8 +20,8 @@ /src/store-api @v0y4g3r @evenyag ## [Module] Metrics Engine -/src/metric-engine @waynexia @WenyXu -/src/promql @waynexia @evenyag @discord9 +/src/metric-engine @WenyXu @v0y4g3r +/src/promql @evenyag @discord9 ## [Module] Flow -/src/flow @discord9 @waynexia +/src/flow @discord9 @WenyXu From 133eda683670cfc497470ca8a827fad8cc879eb1 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 11 Aug 2026 16:28:23 +0800 Subject: [PATCH 048/126] refactor(operator): split semantic_graph relationship builders into a submodule (#8841) Signed-off-by: Dennis Zhuang --- src/operator/src/statement/semantic_graph.rs | 950 +----------------- .../statement/semantic_graph/relationships.rs | 905 +++++++++++++++++ 2 files changed, 959 insertions(+), 896 deletions(-) create mode 100644 src/operator/src/statement/semantic_graph/relationships.rs diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index 15f76905555..af9108198d4 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -25,6 +25,8 @@ //! including filter pushdown into the source table scans. See //! `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. +mod relationships; + use std::sync::{Arc, LazyLock}; use api::v1::column_data_type_extension::TypeExt; @@ -33,21 +35,17 @@ use api::v1::{ SemanticType, }; use common_catalog::consts::{ - DEFAULT_PRIVATE_SCHEMA_NAME, DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, - SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, SPAN_KIND_COLUMN, - SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, TRACE_ID_COLUMN, - TRACE_TIMESTAMP_COLUMN, + DEFAULT_PRIVATE_SCHEMA_NAME, SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, }; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; use datafusion::arrow::datatypes::{DataType, TimeUnit}; use datafusion::dataframe::DataFrame; use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns}; -use datafusion::functions_aggregate::expr_fn::{count, sum}; use datafusion::functions_nested::expr_fn::make_array; -use datafusion::functions_window::expr_fn::row_number; use datafusion_common::{Column, Result as DfResult, ScalarValue}; -use datafusion_expr::{Expr, ExprFunctionExt, JoinType, LogicalPlan, ScalarUDF, cast, ident, lit}; +use datafusion_expr::{Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; +pub use relationships::{CallsSource, DeclaredSource, build_relationships_plan}; use store_api::mito_engine_options::{APPEND_MODE_KEY, MERGE_MODE_KEY, TTL_KEY}; /// Whether an existing declared-edge table still matches the canonical @@ -619,19 +617,6 @@ pub struct RegistrySource { pub scan: DataFrame, } -/// A trace table's scan paired with the `service` declaration it derives -/// `calls` edges for — a unit of `build_relationships_plan`. -pub struct CallsSource { - pub service: EntityDeclaration, - pub scan: DataFrame, -} - -/// The declared-edge table's scan (`semantic_relationships_declared`), whose -/// rows `build_relationships_plan` unions into the edge set. -pub struct DeclaredSource { - pub scan: DataFrame, -} - /// Builds the `semantic_entities` registry plan: one branch and source scan per /// declaring table, filtered to `window`, then `UNION ALL` across source tables. /// Returns `None` when nothing declared an entity, so the computed table streams @@ -650,266 +635,71 @@ pub fn build_registry_plan( Ok(union_df.map(DataFrame::into_unoptimized_plan)) } -/// The projected columns of `semantic_relationships`, in order. Every derived -/// branch and the declared-edge branch must project exactly these so the -/// top-level `UNION ALL` type-aligns; `build_relationships_plan` re-selects -/// them over the union to enforce the contract. (The physical declared table -/// additionally stores `valid_from`/`valid_until`, which feed the validity -/// filter and the projected window columns.) -const RELATIONSHIP_COLUMNS: [&str; 16] = [ - "observed_at", - "window_start", - "window_end", - "fresh_until", - "src_type", - "src_id", - "dst_type", - "dst_id", - "rel_type", - "provenance", - "confidence", - "request_count", - "error_count", - "duration_sum", - "duration_count", - "attributes", -]; +/// RecordBatch readers shared by this module's tests and the `relationships` +/// tests. +#[cfg(test)] +pub(crate) mod test_util { + use datafusion::arrow::array::{Array, BinaryArray, StringArray, TimestampMillisecondArray}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::prelude::SessionContext; + use datafusion_expr::LogicalPlan; -/// A child server span starts no earlier than 5 minutes before its client span -/// (clock-skew allowance) and no later than 1 hour after it; the bounds keep the -/// join windowed instead of pairing arbitrarily distant spans of a long-lived -/// trace. -const CHILD_SPAN_EARLY_NANOS: i64 = 5 * 60 * 1_000_000_000; -const CHILD_SPAN_LATE_NANOS: i64 = 60 * 60 * 1_000_000_000; - -/// Builds the `semantic_relationships` plan: the trace-derived `calls` branch -/// unioned with the declared-edge branch, re-projected to the 16-column -/// contract. Returns `None` when there is neither a trace table nor a -/// declared-edge table, so the computed table streams empty. -pub fn build_relationships_plan( - traces: Vec, - declared: Option, - window: &GraphQueryWindow, -) -> DfResult> { - let mut union_df = calls_branch(traces, window)?; - if let Some(declared) = declared { - union_df = union_all(union_df, declared_edges(declared.scan, window)?)?; - } - union_df - .map(|df| { - Ok(df - .select_columns(&RELATIONSHIP_COLUMNS)? - .into_unoptimized_plan()) - }) - .transpose() -} - -/// Ranking column used to pick the latest declared-edge revision per edge key. -const DECLARED_REVISION_COLUMN: &str = "__declared_revision"; - -/// The declared-edge branch: the latest revision per edge key *as of the -/// queried window* (mito dedups on primary key + `observed_at`, so re-asserting -/// an edge stores a new revision), filtered by business-validity overlap. -/// -/// `valid_from` defaults to the declaration time; a NULL `valid_until` means -/// the edge holds while its row exists (TTL retires it), so `window_end` / -/// `fresh_until` take the window's upper bound. The output `observed_at` is -/// synthesized inside the queried range: the scan's filters are re-applied -/// above the computed table, and the physical revision time would fail them. -fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult { - let eff_valid_from = - core_fns::coalesce().call(vec![ident("valid_from"), ident(OBSERVED_AT_COLUMN)]); - let eff_valid_until = - core_fns::coalesce().call(vec![ident("valid_until"), window.observed_end()]); - - // As-of the queried window: a revision recorded after its end, or whose - // validity starts after it, was not the edge's state inside the window and - // must not outrank (and thereby hide) the revision that was. - let as_of = ident(OBSERVED_AT_COLUMN) - .lt(window.observed_end()) - .and(eff_valid_from.clone().lt(window.observed_end())); - - let revision = row_number() - .partition_by( - DECLARED_EDGE_IDENTITY_COLUMNS - .iter() - .map(|c| ident(*c)) - .collect(), - ) - // generation_id/scope break observed_at ties deterministically (rows - // differing only in them are not merged by the storage dedup). - .order_by(vec![ - ident(OBSERVED_AT_COLUMN).sort(false, false), - ident("generation_id").sort(false, false), - ident("scope").sort(false, false), - ]) - .build()? - .alias(DECLARED_REVISION_COLUMN); - - // The as-of filter already bounds eff_valid_from below the window's end; - // only the retirement side of the overlap remains to check. - let still_valid = ident("valid_until") - .is_null() - .or(ident("valid_until").gt(window.observed_start())); - // Tags come out of the storage engine dictionary-encoded; cast to plain - // strings so the union with the derived branches type-aligns. - let tag_utf8 = |name: &str| cast(ident(name), DataType::Utf8).alias(name); - - scan.filter(as_of)? - .window(vec![revision])? - .filter(ident(DECLARED_REVISION_COLUMN).eq(lit(1_u64)))? - .filter(still_valid)? - .select(vec![ - core_fns::greatest() - .call(vec![eff_valid_from.clone(), window.observed_start()]) - .alias(OBSERVED_AT_COLUMN), - eff_valid_from.alias("window_start"), - eff_valid_until.clone().alias("window_end"), - eff_valid_until.alias("fresh_until"), - tag_utf8("src_type"), - tag_utf8("src_id"), - tag_utf8("dst_type"), - tag_utf8("dst_id"), - tag_utf8("rel_type"), - tag_utf8("provenance"), - ident("confidence"), - ident("request_count"), - ident("error_count"), - ident("duration_sum"), - ident("duration_count"), - ident("attributes"), - ]) -} - -/// The `calls` derivation: pair each client span with its child server span on -/// `trace_id` + `parent_span_id`, project to `service`, union the pairs of all -/// trace tables, and aggregate RED metrics per 60s window in one pass — an edge -/// observed across several trace tables yields one row, not per-table -/// fragments (the plan form of the Tempo servicegraph connector). Column names -/// are the fixed `greptime_trace_v1` schema, the reason -/// `table_data_model = greptime_trace_v1` is required. Returns `None` when -/// there is no trace table. -fn calls_branch( - traces: Vec, - window: &GraphQueryWindow, -) -> DfResult> { - let mut union_df: Option = None; - for trace in traces { - union_df = union_all(union_df, calls_pairs(&trace.service, trace.scan, window)?)?; - } - let Some(pairs) = union_df else { - return Ok(None); - }; - - let df = pairs - .aggregate( - vec![ident("observed_at"), ident("src_id"), ident("dst_id")], - vec![ - count(lit(1)).alias("request_count"), - count(lit(1)) - .filter(ident("status_code").eq(lit(SPAN_STATUS_ERROR))) - .build()? - .alias("error_count"), - sum(ident("duration_nano")).alias("duration_nano_sum"), - ], - )? - .select(vec![ - ident("observed_at"), - ident("observed_at").alias("window_start"), - (ident("observed_at") + bin_interval()).alias("window_end"), - (ident("observed_at") + bin_interval()).alias("fresh_until"), - lit("service").alias("src_type"), - ident("src_id"), - lit("service").alias("dst_type"), - ident("dst_id"), - lit("calls").alias("rel_type"), - lit("trace").alias("provenance"), - lit(1.0_f64).alias("confidence"), - ident("request_count"), - ident("error_count"), - // duration_nano sums in nanoseconds; the contract column is seconds. - (cast(ident("duration_nano_sum"), DataType::Float64) / lit(1e9_f64)) - .alias("duration_sum"), - ident("request_count").alias("duration_count"), - null_json().alias("attributes"), - ])?; - Ok(Some(df)) -} - -/// One trace table's client/server span pairs, projected to the aggregation -/// inputs `(observed_at, src_id, dst_id, status_code, duration_nano)`. -/// Endpoint ids are built from the table's `service` entity declaration, so -/// edges land on exactly the entity ids the registry emits (a composite -/// service identity renders the same sorted `k=v` form on both sides). -fn calls_pairs( - service: &EntityDeclaration, - trace: DataFrame, - window: &GraphQueryWindow, -) -> DfResult { - let mut client_pred = ident(SPAN_KIND_COLUMN) - .eq(lit(SPAN_KIND_CLIENT)) - .and(ident(TRACE_TIMESTAMP_COLUMN).gt_eq(window.source_start())) - .and(ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end())); - let mut server_pred = ident(SPAN_KIND_COLUMN) - .eq(lit(SPAN_KIND_SERVER)) - // Static bounds implied by the window and the join's time-proximity - // conditions below; the join bounds reference client.timestamp and - // cannot prune the server-side scan. - .and( - ident(TRACE_TIMESTAMP_COLUMN) - .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)), - ) - .and( - ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS)), - ); - // A NULL identity component identifies nothing, on either endpoint. - for id in &service.id_columns { - client_pred = client_pred.and(ident(id).is_not_null()); - server_pred = server_pred.and(ident(id).is_not_null()); + pub(crate) async fn collect(ctx: &SessionContext, plan: LogicalPlan) -> Vec { + ctx.execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap() } - let client = trace.clone().filter(client_pred)?.alias("client")?; - let server = trace.filter(server_pred)?.alias("server")?; + pub(crate) fn json_texts(batch: &RecordBatch, column: usize) -> Vec> { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()) + .map(|i| { + array + .is_valid(i) + .then(|| jsonb::from_slice(array.value(i)).unwrap().to_string()) + }) + .collect() + } - let join_conditions = vec![ - qcol("client", TRACE_ID_COLUMN).eq(qcol("server", TRACE_ID_COLUMN)), - qcol("server", PARENT_SPAN_ID_COLUMN).eq(qcol("client", SPAN_ID_COLUMN)), - qcol("server", TRACE_TIMESTAMP_COLUMN) - .gt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) - interval(CHILD_SPAN_EARLY_NANOS)), - qcol("server", TRACE_TIMESTAMP_COLUMN) - .lt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) + interval(CHILD_SPAN_LATE_NANOS)), - ]; + pub(crate) fn strings(batch: &RecordBatch, column: usize) -> Vec { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()) + .map(|i| array.value(i).to_string()) + .collect() + } - client - .join_on(server, JoinType::Inner, join_conditions)? - .select(vec![ - bin_ms(qcol("client", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), - // The cast inside entity_id_expr also normalizes tag columns, which - // come out of the storage engine dictionary-encoded. - entity_id_expr(&service.id_columns, &|c| qcol("client", c)).alias("src_id"), - entity_id_expr(&service.id_columns, &|c| qcol("server", c)).alias("dst_id"), - qcol("server", SPAN_STATUS_CODE_COLUMN).alias("status_code"), - qcol("server", DURATION_NANO_COLUMN).alias("duration_nano"), - ])? - // Exclude self-calls on the composed identity: an edge needs two - // distinct service entities. - .filter(ident("src_id").not_eq(ident("dst_id"))) + pub(crate) fn ts_values(batch: &RecordBatch, column: usize) -> Vec { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()).map(|i| array.value(i)).collect() + } } #[cfg(test)] mod tests { use std::sync::Arc; - use common_catalog::consts::SERVICE_NAME_COLUMN; - use datafusion::arrow::array::{ - Array, ArrayRef, BinaryArray, Float64Array, Int64Array, StringArray, - TimestampMillisecondArray, TimestampNanosecondArray, UInt64Array, - }; + use datafusion::arrow::array::{ArrayRef, Int64Array, StringArray, TimestampMillisecondArray}; use datafusion::arrow::datatypes::{Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::MemTable; use datafusion::prelude::SessionContext; + use super::test_util::{collect, json_texts, strings}; use super::*; fn test_window() -> GraphQueryWindow { @@ -942,41 +732,6 @@ mod tests { assert_eq!((window.source_start_ms, window.source_end_ms), (-60_000, 0)); } - async fn collect(ctx: &SessionContext, plan: LogicalPlan) -> Vec { - ctx.execute_logical_plan(plan) - .await - .unwrap() - .collect() - .await - .unwrap() - } - - fn json_texts(batch: &RecordBatch, column: usize) -> Vec> { - let array = batch - .column(column) - .as_any() - .downcast_ref::() - .unwrap(); - (0..array.len()) - .map(|i| { - array - .is_valid(i) - .then(|| jsonb::from_slice(array.value(i)).unwrap().to_string()) - }) - .collect() - } - - fn strings(batch: &RecordBatch, column: usize) -> Vec { - let array = batch - .column(column) - .as_any() - .downcast_ref::() - .unwrap(); - (0..array.len()) - .map(|i| array.value(i).to_string()) - .collect() - } - /// A metric-like table: ms timestamps, service/pid identity, nullable /// descriptive column with JSON-hostile characters. fn metric_table_ctx() -> SessionContext { @@ -1234,258 +989,6 @@ mod tests { ); } - /// A trace-like table in the fixed `greptime_trace_v1` shape (ns timestamps). - fn trace_table_ctx() -> SessionContext { - let schema = Arc::new(Schema::new(vec![ - Field::new( - TRACE_TIMESTAMP_COLUMN, - DataType::Timestamp(TimeUnit::Nanosecond, None), - false, - ), - Field::new(TRACE_ID_COLUMN, DataType::Utf8, false), - Field::new(SPAN_ID_COLUMN, DataType::Utf8, false), - Field::new(PARENT_SPAN_ID_COLUMN, DataType::Utf8, true), - Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false), - Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false), - Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false), - Field::new("service_namespace", DataType::Utf8, false), - Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false), - ])); - const MS: i64 = 1_000_000; - // Two client->server pairs frontend->cart (one errored), one pair - // cart->cart (self-call, excluded), one unmatched client span. - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampNanosecondArray::from(vec![ - 1_000 * MS, // client frontend->cart - 1_010 * MS, // server cart - 2_000 * MS, // client frontend->cart (error) - 2_010 * MS, // server cart (error) - 3_000 * MS, // client cart->cart (self-call) - 3_010 * MS, // server cart - 4_000 * MS, // client with no matching server - ])) as ArrayRef, - Arc::new(StringArray::from(vec![ - "t1", "t1", "t2", "t2", "t3", "t3", "t4", - ])), - Arc::new(StringArray::from(vec![ - "c1", "s1", "c2", "s2", "c3", "s3", "c4", - ])), - Arc::new(StringArray::from(vec![ - None, - Some("c1"), - None, - Some("c2"), - None, - Some("c3"), - None, - ])), - Arc::new(StringArray::from(vec![ - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - ])), - Arc::new(StringArray::from(vec![ - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_ERROR", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - ])), - Arc::new(StringArray::from(vec![ - "frontend", "cart", "frontend", "cart", "cart", "cart", "frontend", - ])), - Arc::new(StringArray::from(vec!["ns1"; 7])), - Arc::new(UInt64Array::from(vec![ - 0, - 500_000_000, // 0.5s - 0, - 1_500_000_000, // 1.5s - 0, - 100, - 0, - ])), - ], - ) - .unwrap(); - let ctx = SessionContext::new(); - ctx.register_table( - "opentelemetry_traces", - Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch.clone()]]).unwrap()), - ) - .unwrap(); - ctx.register_table( - "opentelemetry_traces_2", - Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), - ) - .unwrap(); - ctx - } - - fn trace_service_decl(id_columns: &[&str]) -> EntityDeclaration { - EntityDeclaration { - schema: "public".to_string(), - table: "opentelemetry_traces".to_string(), - time_index: TRACE_TIMESTAMP_COLUMN.to_string(), - entity_type: "service".to_string(), - id_columns: id_columns.iter().map(|s| s.to_string()).collect(), - descriptive_columns: vec![], - scope_columns: vec![], - } - } - - #[tokio::test] - async fn calls_plan_aggregates_red_metrics() { - let ctx = trace_table_ctx(); - let trace = ctx.table("opentelemetry_traces").await.unwrap(); - let plan = build_relationships_plan_for_test( - vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: trace, - }], - &test_window(), - ) - .unwrap() - .unwrap(); - - let names = plan - .schema() - .fields() - .iter() - .map(|f| f.name().as_str()) - .collect::>(); - assert_eq!(names, RELATIONSHIP_COLUMNS); - - let batches = collect(&ctx, plan).await; - let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - // frontend->cart only: the self-call and the unmatched client drop out; - // both pairs land in the same 60s bin. - assert_eq!(total, 1); - let batch = &batches[0]; - assert_eq!(strings(batch, 5), vec!["frontend"]); - assert_eq!(strings(batch, 7), vec!["cart"]); - assert_eq!(strings(batch, 8), vec!["calls"]); - assert_eq!(strings(batch, 9), vec!["trace"]); - - let request_count = batch - .column(11) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(request_count.value(0), 2); - let error_count = batch - .column(12) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(error_count.value(0), 1); - let duration_sum = batch - .column(13) - .as_any() - .downcast_ref::() - .unwrap(); - assert!((duration_sum.value(0) - 2.0).abs() < 1e-9); - let duration_count = batch - .column(14) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(duration_count.value(0), 2); - // Derived calls edges carry no attributes: a typed-JSON NULL. - assert!(json_texts(batch, 15)[0].is_none()); - } - - #[tokio::test] - async fn calls_plan_merges_edges_across_trace_tables() { - let ctx = trace_table_ctx(); - let t1 = ctx.table("opentelemetry_traces").await.unwrap(); - let t2 = ctx.table("opentelemetry_traces_2").await.unwrap(); - let plan = build_relationships_plan_for_test( - vec![ - CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: t1, - }, - CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: t2, - }, - ], - &test_window(), - ) - .unwrap() - .unwrap(); - - let batches = collect(&ctx, plan).await; - let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - // The same frontend->cart edge from both tables folds into one row with - // summed RED metrics. - assert_eq!(total, 1); - let batch = &batches[0]; - let request_count = batch - .column(11) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(request_count.value(0), 4); - let error_count = batch - .column(12) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(error_count.value(0), 2); - let duration_sum = batch - .column(13) - .as_any() - .downcast_ref::() - .unwrap(); - assert!((duration_sum.value(0) - 4.0).abs() < 1e-9); - } - - #[tokio::test] - async fn calls_endpoints_follow_service_declaration() { - let ctx = trace_table_ctx(); - let trace = ctx.table("opentelemetry_traces").await.unwrap(); - let plan = build_relationships_plan_for_test( - vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN, "service_namespace"]), - scan: trace, - }], - &test_window(), - ) - .unwrap() - .unwrap(); - - let batches = collect(&ctx, plan).await; - let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total, 1); - let batch = &batches[0]; - // Composite service identity renders the same sorted `k=v` form the - // registry emits, so edges land on registry entity ids. - assert_eq!( - strings(batch, 5), - vec!["service_name=frontend,service_namespace=ns1"] - ); - assert_eq!( - strings(batch, 7), - vec!["service_name=cart,service_namespace=ns1"] - ); - } - - fn build_relationships_plan_for_test( - traces: Vec, - window: &GraphQueryWindow, - ) -> DfResult> { - build_relationships_plan(traces, None, window) - } - /// A `TableInfo` derived from the canonical declared-edge definition, with /// injection points for each aspect the schema matcher must check. fn declared_table_info( @@ -1616,349 +1119,4 @@ mod tests { ); assert!(!declared_relationships_schema_matches(&wrong_column_type)); } - - fn ts_values(batch: &RecordBatch, column: usize) -> Vec { - let array = batch - .column(column) - .as_any() - .downcast_ref::() - .unwrap(); - (0..array.len()).map(|i| array.value(i)).collect() - } - - /// Registers a declared-edge table holding: - /// - frontend->db: two revisions (0.5 then 1.0), the newer retiring the edge at 300s; - /// - api->cache: expired at 100s; - /// - agent-1->search: open-ended `provenance = 'agent'` with JSON attributes; - /// - webapp->auth: two revisions, the newer declared at 900s; - /// - batch->queue: declared at 100s but only valid from 700s; - /// - gateway->redis: asserted twice at the same instant as generations g1/g2. - fn declared_table(ctx: &SessionContext) { - let ts_field = |name: &str| { - Field::new( - name, - DataType::Timestamp(TimeUnit::Millisecond, None), - name != OBSERVED_AT_COLUMN, - ) - }; - let schema = Arc::new(Schema::new(vec![ - ts_field(OBSERVED_AT_COLUMN), - ts_field("window_start"), - ts_field("window_end"), - ts_field("fresh_until"), - ts_field("valid_from"), - ts_field("valid_until"), - Field::new("src_type", DataType::Utf8, false), - Field::new("src_id", DataType::Utf8, false), - Field::new("rel_type", DataType::Utf8, false), - Field::new("dst_type", DataType::Utf8, false), - Field::new("dst_id", DataType::Utf8, false), - Field::new("provenance", DataType::Utf8, false), - Field::new("scope", DataType::Utf8, false), - Field::new("generation_id", DataType::Utf8, false), - Field::new("confidence", DataType::Float64, true), - Field::new("request_count", DataType::Int64, true), - Field::new("error_count", DataType::Int64, true), - Field::new("duration_sum", DataType::Float64, true), - Field::new("duration_count", DataType::Int64, true), - Field::new("attributes", DataType::Binary, true), - ])); - let attrs = jsonb::parse_value(br#"{"connection_type":"virtual_node"}"#) - .unwrap() - .to_vec(); - let no_ts = || Arc::new(TimestampMillisecondArray::from(vec![None::; 9])) as ArrayRef; - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![ - 100_000, 200_000, 50_000, 90_000, 100_000, 900_000, 100_000, 100_000, 100_000, - ])) as ArrayRef, - no_ts(), // window_start - no_ts(), // window_end - no_ts(), // fresh_until - Arc::new(TimestampMillisecondArray::from(vec![ - None, - None, - None, - None, - None, - None, - Some(700_000), - None, - None, - ])), - Arc::new(TimestampMillisecondArray::from(vec![ - None, - Some(300_000), - Some(100_000), - None, - None, - None, - None, - None, - None, - ])), - Arc::new(StringArray::from(vec![ - "service", "service", "service", "agent", "service", "service", "service", - "service", "service", - ])), - Arc::new(StringArray::from(vec![ - "frontend", "frontend", "api", "agent-1", "webapp", "webapp", "batch", - "gateway", "gateway", - ])), - Arc::new(StringArray::from(vec![ - "depends_on", - "depends_on", - "depends_on", - "uses", - "depends_on", - "depends_on", - "depends_on", - "depends_on", - "depends_on", - ])), - Arc::new(StringArray::from(vec![ - "service", "service", "service", "tool", "service", "service", "service", - "service", "service", - ])), - Arc::new(StringArray::from(vec![ - "db", "db", "cache", "search", "auth", "auth", "queue", "redis", "redis", - ])), - Arc::new(StringArray::from(vec![ - "declared", "declared", "declared", "agent", "declared", "declared", - "declared", "declared", "declared", - ])), - Arc::new(StringArray::from(vec![""; 9])), - Arc::new(StringArray::from(vec![ - "", "", "", "", "", "", "", "g1", "g2", - ])), - Arc::new(Float64Array::from(vec![ - Some(0.5), - Some(1.0), - Some(1.0), - Some(0.8), - Some(0.5), - Some(1.0), - Some(1.0), - Some(0.5), - Some(1.0), - ])), - Arc::new(Int64Array::from(vec![None::; 9])), - Arc::new(Int64Array::from(vec![None::; 9])), - Arc::new(Float64Array::from(vec![None::; 9])), - Arc::new(Int64Array::from(vec![None::; 9])), - Arc::new(BinaryArray::from(vec![ - None, - None, - None, - Some(attrs.as_slice()), - None, - None, - None, - None, - None, - ])), - ], - ) - .unwrap(); - ctx.register_table( - SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, - Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), - ) - .unwrap(); - } - - #[tokio::test] - async fn declared_edges_latest_revision_validity_and_synthetic_time() { - let ctx = SessionContext::new(); - declared_table(&ctx); - let scan = ctx - .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) - .await - .unwrap(); - - // Queried window [120s, 600s). - let window = GraphQueryWindow::from_observed(120_000, 600_000); - let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) - .unwrap() - .unwrap(); - let names = plan - .schema() - .fields() - .iter() - .map(|f| f.name().as_str()) - .collect::>(); - assert_eq!(names, RELATIONSHIP_COLUMNS); - - // (src_id, observed_at, window_start, window_end, fresh_until, - // provenance, attributes) - type DeclaredRow = (String, i64, i64, i64, i64, String, Option); - - let batches = collect(&ctx, plan).await; - let mut rows: Vec = vec![]; - for batch in &batches { - let observed = ts_values(batch, 0); - let window_start = ts_values(batch, 1); - let window_end = ts_values(batch, 2); - let fresh_until = ts_values(batch, 3); - let src = strings(batch, 5); - let provenance = strings(batch, 9); - let attributes = json_texts(batch, 15); - for i in 0..batch.num_rows() { - rows.push(( - src[i].clone(), - observed[i], - window_start[i], - window_end[i], - fresh_until[i], - provenance[i].clone(), - attributes[i].clone(), - )); - } - } - rows.sort(); - - // api->cache expired before the window; batch->queue is not yet valid; - // frontend->db keeps only the latest revision; webapp->auth keeps the - // first revision (the second postdates the window); gateway->redis - // collapses to one row across generations. - assert_eq!(rows.len(), 4); - - // agent->tool: open-ended validity declared at 90s. The synthesized - // observed_at is clamped into the queried window; window_end and - // fresh_until take the window's upper bound. - assert_eq!( - rows[0], - ( - "agent-1".to_string(), - 120_000, - 90_000, - 600_000, - 600_000, - "agent".to_string(), - Some(r#"{"connection_type":"virtual_node"}"#.to_string()), - ) - ); - - // frontend->db: the latest revision (declared at 200s, retired at 300s) - // supersedes the open-ended first revision. - assert_eq!( - rows[1], - ( - "frontend".to_string(), - 200_000, - 200_000, - 300_000, - 300_000, - "declared".to_string(), - None, - ) - ); - - assert_eq!( - rows[2], - ( - "gateway".to_string(), - 120_000, - 100_000, - 600_000, - 600_000, - "declared".to_string(), - None, - ) - ); - - assert_eq!( - rows[3], - ( - "webapp".to_string(), - 120_000, - 100_000, - 600_000, - 600_000, - "declared".to_string(), - None, - ) - ); - } - - #[tokio::test] - async fn declared_edges_pick_the_revision_as_of_the_window() { - let ctx = SessionContext::new(); - declared_table(&ctx); - let scan = ctx - .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) - .await - .unwrap(); - - // A historical window that predates frontend->db's second revision: - // the first revision was the edge's state then and must not be hidden - // by the newer one. - let window = GraphQueryWindow::from_observed(0, 150_000); - let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) - .unwrap() - .unwrap(); - - let batches = collect(&ctx, plan).await; - let mut rows: Vec<(String, f64)> = vec![]; - for batch in &batches { - let src = strings(batch, 5); - let confidence = batch - .column(10) - .as_any() - .downcast_ref::() - .unwrap(); - for (i, src) in src.into_iter().enumerate() { - rows.push((src, confidence.value(i))); - } - } - rows.sort_by(|a, b| a.0.cmp(&b.0)); - - // api->cache is valid inside this window; batch->queue is not yet. - assert_eq!( - rows, - vec![ - ("agent-1".to_string(), 0.8), - ("api".to_string(), 1.0), - ("frontend".to_string(), 0.5), - ("gateway".to_string(), 1.0), - ("webapp".to_string(), 0.5), - ] - ); - } - - #[tokio::test] - async fn declared_edges_union_calls_branch() { - let ctx = trace_table_ctx(); - declared_table(&ctx); - let trace = ctx.table("opentelemetry_traces").await.unwrap(); - let declared = ctx - .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) - .await - .unwrap(); - - let plan = build_relationships_plan( - vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: trace, - }], - Some(DeclaredSource { scan: declared }), - &test_window(), - ) - .unwrap() - .unwrap(); - - let batches = collect(&ctx, plan).await; - let mut provenance: Vec = batches.iter().flat_map(|b| strings(b, 9)).collect(); - provenance.sort(); - // One trace-derived calls edge plus the declared edges valid somewhere - // inside [0s, 600s): frontend, api, webapp, gateway, and the - // agent-asserted one. - assert_eq!( - provenance, - vec![ - "agent", "declared", "declared", "declared", "declared", "trace" - ] - ); - } } diff --git a/src/operator/src/statement/semantic_graph/relationships.rs b/src/operator/src/statement/semantic_graph/relationships.rs new file mode 100644 index 00000000000..7ebf5ae86fd --- /dev/null +++ b/src/operator/src/statement/semantic_graph/relationships.rs @@ -0,0 +1,905 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The `semantic_relationships` derivation: the plan builders for every edge +//! branch (trace-derived `calls`, the declared-edge union) behind the computed +//! table. Shared expression helpers and the entity registry live in the parent +//! module. + +use common_catalog::consts::{ + DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, + SPAN_KIND_COLUMN, SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, + TRACE_ID_COLUMN, TRACE_TIMESTAMP_COLUMN, +}; +use datafusion::arrow::datatypes::DataType; +use datafusion::dataframe::DataFrame; +use datafusion::functions::core as core_fns; +use datafusion::functions_aggregate::expr_fn::{count, sum}; +use datafusion::functions_window::expr_fn::row_number; +use datafusion_common::Result as DfResult; +use datafusion_expr::{ExprFunctionExt, JoinType, LogicalPlan, cast, ident, lit}; + +use crate::statement::semantic_graph::{ + DECLARED_EDGE_IDENTITY_COLUMNS, EntityDeclaration, GraphQueryWindow, OBSERVED_AT_COLUMN, + bin_interval, bin_ms, entity_id_expr, interval, null_json, qcol, union_all, +}; + +/// The projected columns of `semantic_relationships`, in order. Every derived +/// branch and the declared-edge branch must project exactly these so the +/// top-level `UNION ALL` type-aligns; `build_relationships_plan` re-selects +/// them over the union to enforce the contract. (The physical declared table +/// additionally stores `valid_from`/`valid_until`, which feed the validity +/// filter and the projected window columns.) +const RELATIONSHIP_COLUMNS: [&str; 16] = [ + "observed_at", + "window_start", + "window_end", + "fresh_until", + "src_type", + "src_id", + "dst_type", + "dst_id", + "rel_type", + "provenance", + "confidence", + "request_count", + "error_count", + "duration_sum", + "duration_count", + "attributes", +]; + +/// A child server span starts no earlier than 5 minutes before its client span +/// (clock-skew allowance) and no later than 1 hour after it; the bounds keep the +/// join windowed instead of pairing arbitrarily distant spans of a long-lived +/// trace. +const CHILD_SPAN_EARLY_NANOS: i64 = 5 * 60 * 1_000_000_000; +const CHILD_SPAN_LATE_NANOS: i64 = 60 * 60 * 1_000_000_000; + +/// A trace table's scan paired with the `service` declaration it derives +/// `calls` edges for — a unit of `build_relationships_plan`. +pub struct CallsSource { + pub service: EntityDeclaration, + pub scan: DataFrame, +} + +/// The declared-edge table's scan (`semantic_relationships_declared`), whose +/// rows `build_relationships_plan` unions into the edge set. +pub struct DeclaredSource { + pub scan: DataFrame, +} + +/// Builds the `semantic_relationships` plan: the trace-derived `calls` branch +/// unioned with the declared-edge branch, re-projected to the 16-column +/// contract. Returns `None` when there is neither a trace table nor a +/// declared-edge table, so the computed table streams empty. +pub fn build_relationships_plan( + traces: Vec, + declared: Option, + window: &GraphQueryWindow, +) -> DfResult> { + let mut union_df = calls_branch(traces, window)?; + if let Some(declared) = declared { + union_df = union_all(union_df, declared_edges(declared.scan, window)?)?; + } + union_df + .map(|df| { + Ok(df + .select_columns(&RELATIONSHIP_COLUMNS)? + .into_unoptimized_plan()) + }) + .transpose() +} + +/// Ranking column used to pick the latest declared-edge revision per edge key. +const DECLARED_REVISION_COLUMN: &str = "__declared_revision"; + +/// The declared-edge branch: the latest revision per edge key *as of the +/// queried window* (mito dedups on primary key + `observed_at`, so re-asserting +/// an edge stores a new revision), filtered by business-validity overlap. +/// +/// `valid_from` defaults to the declaration time; a NULL `valid_until` means +/// the edge holds while its row exists (TTL retires it), so `window_end` / +/// `fresh_until` take the window's upper bound. The output `observed_at` is +/// synthesized inside the queried range: the scan's filters are re-applied +/// above the computed table, and the physical revision time would fail them. +fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult { + let eff_valid_from = + core_fns::coalesce().call(vec![ident("valid_from"), ident(OBSERVED_AT_COLUMN)]); + let eff_valid_until = + core_fns::coalesce().call(vec![ident("valid_until"), window.observed_end()]); + + // As-of the queried window: a revision recorded after its end, or whose + // validity starts after it, was not the edge's state inside the window and + // must not outrank (and thereby hide) the revision that was. + let as_of = ident(OBSERVED_AT_COLUMN) + .lt(window.observed_end()) + .and(eff_valid_from.clone().lt(window.observed_end())); + + let revision = row_number() + .partition_by( + DECLARED_EDGE_IDENTITY_COLUMNS + .iter() + .map(|c| ident(*c)) + .collect(), + ) + // generation_id/scope break observed_at ties deterministically (rows + // differing only in them are not merged by the storage dedup). + .order_by(vec![ + ident(OBSERVED_AT_COLUMN).sort(false, false), + ident("generation_id").sort(false, false), + ident("scope").sort(false, false), + ]) + .build()? + .alias(DECLARED_REVISION_COLUMN); + + // The as-of filter already bounds eff_valid_from below the window's end; + // only the retirement side of the overlap remains to check. + let still_valid = ident("valid_until") + .is_null() + .or(ident("valid_until").gt(window.observed_start())); + // Tags come out of the storage engine dictionary-encoded; cast to plain + // strings so the union with the derived branches type-aligns. + let tag_utf8 = |name: &str| cast(ident(name), DataType::Utf8).alias(name); + + scan.filter(as_of)? + .window(vec![revision])? + .filter(ident(DECLARED_REVISION_COLUMN).eq(lit(1_u64)))? + .filter(still_valid)? + .select(vec![ + core_fns::greatest() + .call(vec![eff_valid_from.clone(), window.observed_start()]) + .alias(OBSERVED_AT_COLUMN), + eff_valid_from.alias("window_start"), + eff_valid_until.clone().alias("window_end"), + eff_valid_until.alias("fresh_until"), + tag_utf8("src_type"), + tag_utf8("src_id"), + tag_utf8("dst_type"), + tag_utf8("dst_id"), + tag_utf8("rel_type"), + tag_utf8("provenance"), + ident("confidence"), + ident("request_count"), + ident("error_count"), + ident("duration_sum"), + ident("duration_count"), + ident("attributes"), + ]) +} + +/// The `calls` derivation: pair each client span with its child server span on +/// `trace_id` + `parent_span_id`, project to `service`, union the pairs of all +/// trace tables, and aggregate RED metrics per 60s window in one pass — an edge +/// observed across several trace tables yields one row, not per-table +/// fragments (the plan form of the Tempo servicegraph connector). Column names +/// are the fixed `greptime_trace_v1` schema, the reason +/// `table_data_model = greptime_trace_v1` is required. Returns `None` when +/// there is no trace table. +fn calls_branch( + traces: Vec, + window: &GraphQueryWindow, +) -> DfResult> { + let mut union_df: Option = None; + for trace in traces { + union_df = union_all(union_df, calls_pairs(&trace.service, trace.scan, window)?)?; + } + let Some(pairs) = union_df else { + return Ok(None); + }; + + let df = pairs + .aggregate( + vec![ident("observed_at"), ident("src_id"), ident("dst_id")], + vec![ + count(lit(1)).alias("request_count"), + count(lit(1)) + .filter(ident("status_code").eq(lit(SPAN_STATUS_ERROR))) + .build()? + .alias("error_count"), + sum(ident("duration_nano")).alias("duration_nano_sum"), + ], + )? + .select(vec![ + ident("observed_at"), + ident("observed_at").alias("window_start"), + (ident("observed_at") + bin_interval()).alias("window_end"), + (ident("observed_at") + bin_interval()).alias("fresh_until"), + lit("service").alias("src_type"), + ident("src_id"), + lit("service").alias("dst_type"), + ident("dst_id"), + lit("calls").alias("rel_type"), + lit("trace").alias("provenance"), + lit(1.0_f64).alias("confidence"), + ident("request_count"), + ident("error_count"), + // duration_nano sums in nanoseconds; the contract column is seconds. + (cast(ident("duration_nano_sum"), DataType::Float64) / lit(1e9_f64)) + .alias("duration_sum"), + ident("request_count").alias("duration_count"), + null_json().alias("attributes"), + ])?; + Ok(Some(df)) +} + +/// One trace table's client/server span pairs, projected to the aggregation +/// inputs `(observed_at, src_id, dst_id, status_code, duration_nano)`. +/// Endpoint ids are built from the table's `service` entity declaration, so +/// edges land on exactly the entity ids the registry emits (a composite +/// service identity renders the same sorted `k=v` form on both sides). +fn calls_pairs( + service: &EntityDeclaration, + trace: DataFrame, + window: &GraphQueryWindow, +) -> DfResult { + let mut client_pred = ident(SPAN_KIND_COLUMN) + .eq(lit(SPAN_KIND_CLIENT)) + .and(ident(TRACE_TIMESTAMP_COLUMN).gt_eq(window.source_start())) + .and(ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end())); + let mut server_pred = ident(SPAN_KIND_COLUMN) + .eq(lit(SPAN_KIND_SERVER)) + // Static bounds implied by the window and the join's time-proximity + // conditions below; the join bounds reference client.timestamp and + // cannot prune the server-side scan. + .and( + ident(TRACE_TIMESTAMP_COLUMN) + .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)), + ) + .and( + ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS)), + ); + // A NULL identity component identifies nothing, on either endpoint. + for id in &service.id_columns { + client_pred = client_pred.and(ident(id).is_not_null()); + server_pred = server_pred.and(ident(id).is_not_null()); + } + + let client = trace.clone().filter(client_pred)?.alias("client")?; + let server = trace.filter(server_pred)?.alias("server")?; + + let join_conditions = vec![ + qcol("client", TRACE_ID_COLUMN).eq(qcol("server", TRACE_ID_COLUMN)), + qcol("server", PARENT_SPAN_ID_COLUMN).eq(qcol("client", SPAN_ID_COLUMN)), + qcol("server", TRACE_TIMESTAMP_COLUMN) + .gt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) - interval(CHILD_SPAN_EARLY_NANOS)), + qcol("server", TRACE_TIMESTAMP_COLUMN) + .lt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) + interval(CHILD_SPAN_LATE_NANOS)), + ]; + + client + .join_on(server, JoinType::Inner, join_conditions)? + .select(vec![ + bin_ms(qcol("client", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), + // The cast inside entity_id_expr also normalizes tag columns, which + // come out of the storage engine dictionary-encoded. + entity_id_expr(&service.id_columns, &|c| qcol("client", c)).alias("src_id"), + entity_id_expr(&service.id_columns, &|c| qcol("server", c)).alias("dst_id"), + qcol("server", SPAN_STATUS_CODE_COLUMN).alias("status_code"), + qcol("server", DURATION_NANO_COLUMN).alias("duration_nano"), + ])? + // Exclude self-calls on the composed identity: an edge needs two + // distinct service entities. + .filter(ident("src_id").not_eq(ident("dst_id"))) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_catalog::consts::{SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SERVICE_NAME_COLUMN}; + use datafusion::arrow::array::{ + ArrayRef, BinaryArray, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, + TimestampNanosecondArray, UInt64Array, + }; + use datafusion::arrow::datatypes::{Field, Schema, TimeUnit}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::MemTable; + use datafusion::prelude::SessionContext; + + use super::super::test_util::{collect, json_texts, strings, ts_values}; + use super::*; + + fn test_window() -> GraphQueryWindow { + GraphQueryWindow::from_observed(0, 10 * 60 * 1000) + } + + /// A trace-like table in the fixed `greptime_trace_v1` shape (ns timestamps). + fn trace_table_ctx() -> SessionContext { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TRACE_TIMESTAMP_COLUMN, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new(TRACE_ID_COLUMN, DataType::Utf8, false), + Field::new(SPAN_ID_COLUMN, DataType::Utf8, false), + Field::new(PARENT_SPAN_ID_COLUMN, DataType::Utf8, true), + Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false), + Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false), + Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false), + Field::new("service_namespace", DataType::Utf8, false), + Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false), + ])); + const MS: i64 = 1_000_000; + // Two client->server pairs frontend->cart (one errored), one pair + // cart->cart (self-call, excluded), one unmatched client span. + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampNanosecondArray::from(vec![ + 1_000 * MS, // client frontend->cart + 1_010 * MS, // server cart + 2_000 * MS, // client frontend->cart (error) + 2_010 * MS, // server cart (error) + 3_000 * MS, // client cart->cart (self-call) + 3_010 * MS, // server cart + 4_000 * MS, // client with no matching server + ])) as ArrayRef, + Arc::new(StringArray::from(vec![ + "t1", "t1", "t2", "t2", "t3", "t3", "t4", + ])), + Arc::new(StringArray::from(vec![ + "c1", "s1", "c2", "s2", "c3", "s3", "c4", + ])), + Arc::new(StringArray::from(vec![ + None, + Some("c1"), + None, + Some("c2"), + None, + Some("c3"), + None, + ])), + Arc::new(StringArray::from(vec![ + "SPAN_KIND_CLIENT", + "SPAN_KIND_SERVER", + "SPAN_KIND_CLIENT", + "SPAN_KIND_SERVER", + "SPAN_KIND_CLIENT", + "SPAN_KIND_SERVER", + "SPAN_KIND_CLIENT", + ])), + Arc::new(StringArray::from(vec![ + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + "STATUS_CODE_ERROR", + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + ])), + Arc::new(StringArray::from(vec![ + "frontend", "cart", "frontend", "cart", "cart", "cart", "frontend", + ])), + Arc::new(StringArray::from(vec!["ns1"; 7])), + Arc::new(UInt64Array::from(vec![ + 0, + 500_000_000, // 0.5s + 0, + 1_500_000_000, // 1.5s + 0, + 100, + 0, + ])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "opentelemetry_traces", + Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch.clone()]]).unwrap()), + ) + .unwrap(); + ctx.register_table( + "opentelemetry_traces_2", + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + ctx + } + + fn trace_service_decl(id_columns: &[&str]) -> EntityDeclaration { + EntityDeclaration { + schema: "public".to_string(), + table: "opentelemetry_traces".to_string(), + time_index: TRACE_TIMESTAMP_COLUMN.to_string(), + entity_type: "service".to_string(), + id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + descriptive_columns: vec![], + scope_columns: vec![], + } + } + + #[tokio::test] + async fn calls_plan_aggregates_red_metrics() { + let ctx = trace_table_ctx(); + let trace = ctx.table("opentelemetry_traces").await.unwrap(); + let plan = build_relationships_plan_for_test( + vec![CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: trace, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let names = plan + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(); + assert_eq!(names, RELATIONSHIP_COLUMNS); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + // frontend->cart only: the self-call and the unmatched client drop out; + // both pairs land in the same 60s bin. + assert_eq!(total, 1); + let batch = &batches[0]; + assert_eq!(strings(batch, 5), vec!["frontend"]); + assert_eq!(strings(batch, 7), vec!["cart"]); + assert_eq!(strings(batch, 8), vec!["calls"]); + assert_eq!(strings(batch, 9), vec!["trace"]); + + let request_count = batch + .column(11) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(request_count.value(0), 2); + let error_count = batch + .column(12) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(error_count.value(0), 1); + let duration_sum = batch + .column(13) + .as_any() + .downcast_ref::() + .unwrap(); + assert!((duration_sum.value(0) - 2.0).abs() < 1e-9); + let duration_count = batch + .column(14) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(duration_count.value(0), 2); + // Derived calls edges carry no attributes: a typed-JSON NULL. + assert!(json_texts(batch, 15)[0].is_none()); + } + + #[tokio::test] + async fn calls_plan_merges_edges_across_trace_tables() { + let ctx = trace_table_ctx(); + let t1 = ctx.table("opentelemetry_traces").await.unwrap(); + let t2 = ctx.table("opentelemetry_traces_2").await.unwrap(); + let plan = build_relationships_plan_for_test( + vec![ + CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: t1, + }, + CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: t2, + }, + ], + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + // The same frontend->cart edge from both tables folds into one row with + // summed RED metrics. + assert_eq!(total, 1); + let batch = &batches[0]; + let request_count = batch + .column(11) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(request_count.value(0), 4); + let error_count = batch + .column(12) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(error_count.value(0), 2); + let duration_sum = batch + .column(13) + .as_any() + .downcast_ref::() + .unwrap(); + assert!((duration_sum.value(0) - 4.0).abs() < 1e-9); + } + + #[tokio::test] + async fn calls_endpoints_follow_service_declaration() { + let ctx = trace_table_ctx(); + let trace = ctx.table("opentelemetry_traces").await.unwrap(); + let plan = build_relationships_plan_for_test( + vec![CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN, "service_namespace"]), + scan: trace, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 1); + let batch = &batches[0]; + // Composite service identity renders the same sorted `k=v` form the + // registry emits, so edges land on registry entity ids. + assert_eq!( + strings(batch, 5), + vec!["service_name=frontend,service_namespace=ns1"] + ); + assert_eq!( + strings(batch, 7), + vec!["service_name=cart,service_namespace=ns1"] + ); + } + + fn build_relationships_plan_for_test( + traces: Vec, + window: &GraphQueryWindow, + ) -> DfResult> { + build_relationships_plan(traces, None, window) + } + + /// Registers a declared-edge table holding: + /// - frontend->db: two revisions (0.5 then 1.0), the newer retiring the edge at 300s; + /// - api->cache: expired at 100s; + /// - agent-1->search: open-ended `provenance = 'agent'` with JSON attributes; + /// - webapp->auth: two revisions, the newer declared at 900s; + /// - batch->queue: declared at 100s but only valid from 700s; + /// - gateway->redis: asserted twice at the same instant as generations g1/g2. + fn declared_table(ctx: &SessionContext) { + let ts_field = |name: &str| { + Field::new( + name, + DataType::Timestamp(TimeUnit::Millisecond, None), + name != OBSERVED_AT_COLUMN, + ) + }; + let schema = Arc::new(Schema::new(vec![ + ts_field(OBSERVED_AT_COLUMN), + ts_field("window_start"), + ts_field("window_end"), + ts_field("fresh_until"), + ts_field("valid_from"), + ts_field("valid_until"), + Field::new("src_type", DataType::Utf8, false), + Field::new("src_id", DataType::Utf8, false), + Field::new("rel_type", DataType::Utf8, false), + Field::new("dst_type", DataType::Utf8, false), + Field::new("dst_id", DataType::Utf8, false), + Field::new("provenance", DataType::Utf8, false), + Field::new("scope", DataType::Utf8, false), + Field::new("generation_id", DataType::Utf8, false), + Field::new("confidence", DataType::Float64, true), + Field::new("request_count", DataType::Int64, true), + Field::new("error_count", DataType::Int64, true), + Field::new("duration_sum", DataType::Float64, true), + Field::new("duration_count", DataType::Int64, true), + Field::new("attributes", DataType::Binary, true), + ])); + let attrs = jsonb::parse_value(br#"{"connection_type":"virtual_node"}"#) + .unwrap() + .to_vec(); + let no_ts = || Arc::new(TimestampMillisecondArray::from(vec![None::; 9])) as ArrayRef; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 100_000, 200_000, 50_000, 90_000, 100_000, 900_000, 100_000, 100_000, 100_000, + ])) as ArrayRef, + no_ts(), // window_start + no_ts(), // window_end + no_ts(), // fresh_until + Arc::new(TimestampMillisecondArray::from(vec![ + None, + None, + None, + None, + None, + None, + Some(700_000), + None, + None, + ])), + Arc::new(TimestampMillisecondArray::from(vec![ + None, + Some(300_000), + Some(100_000), + None, + None, + None, + None, + None, + None, + ])), + Arc::new(StringArray::from(vec![ + "service", "service", "service", "agent", "service", "service", "service", + "service", "service", + ])), + Arc::new(StringArray::from(vec![ + "frontend", "frontend", "api", "agent-1", "webapp", "webapp", "batch", + "gateway", "gateway", + ])), + Arc::new(StringArray::from(vec![ + "depends_on", + "depends_on", + "depends_on", + "uses", + "depends_on", + "depends_on", + "depends_on", + "depends_on", + "depends_on", + ])), + Arc::new(StringArray::from(vec![ + "service", "service", "service", "tool", "service", "service", "service", + "service", "service", + ])), + Arc::new(StringArray::from(vec![ + "db", "db", "cache", "search", "auth", "auth", "queue", "redis", "redis", + ])), + Arc::new(StringArray::from(vec![ + "declared", "declared", "declared", "agent", "declared", "declared", + "declared", "declared", "declared", + ])), + Arc::new(StringArray::from(vec![""; 9])), + Arc::new(StringArray::from(vec![ + "", "", "", "", "", "", "", "g1", "g2", + ])), + Arc::new(Float64Array::from(vec![ + Some(0.5), + Some(1.0), + Some(1.0), + Some(0.8), + Some(0.5), + Some(1.0), + Some(1.0), + Some(0.5), + Some(1.0), + ])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(Float64Array::from(vec![None::; 9])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(BinaryArray::from(vec![ + None, + None, + None, + Some(attrs.as_slice()), + None, + None, + None, + None, + None, + ])), + ], + ) + .unwrap(); + ctx.register_table( + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + } + + #[tokio::test] + async fn declared_edges_latest_revision_validity_and_synthetic_time() { + let ctx = SessionContext::new(); + declared_table(&ctx); + let scan = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + // Queried window [120s, 600s). + let window = GraphQueryWindow::from_observed(120_000, 600_000); + let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) + .unwrap() + .unwrap(); + let names = plan + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(); + assert_eq!(names, RELATIONSHIP_COLUMNS); + + // (src_id, observed_at, window_start, window_end, fresh_until, + // provenance, attributes) + type DeclaredRow = (String, i64, i64, i64, i64, String, Option); + + let batches = collect(&ctx, plan).await; + let mut rows: Vec = vec![]; + for batch in &batches { + let observed = ts_values(batch, 0); + let window_start = ts_values(batch, 1); + let window_end = ts_values(batch, 2); + let fresh_until = ts_values(batch, 3); + let src = strings(batch, 5); + let provenance = strings(batch, 9); + let attributes = json_texts(batch, 15); + for i in 0..batch.num_rows() { + rows.push(( + src[i].clone(), + observed[i], + window_start[i], + window_end[i], + fresh_until[i], + provenance[i].clone(), + attributes[i].clone(), + )); + } + } + rows.sort(); + + // api->cache expired before the window; batch->queue is not yet valid; + // frontend->db keeps only the latest revision; webapp->auth keeps the + // first revision (the second postdates the window); gateway->redis + // collapses to one row across generations. + assert_eq!(rows.len(), 4); + + // agent->tool: open-ended validity declared at 90s. The synthesized + // observed_at is clamped into the queried window; window_end and + // fresh_until take the window's upper bound. + assert_eq!( + rows[0], + ( + "agent-1".to_string(), + 120_000, + 90_000, + 600_000, + 600_000, + "agent".to_string(), + Some(r#"{"connection_type":"virtual_node"}"#.to_string()), + ) + ); + + // frontend->db: the latest revision (declared at 200s, retired at 300s) + // supersedes the open-ended first revision. + assert_eq!( + rows[1], + ( + "frontend".to_string(), + 200_000, + 200_000, + 300_000, + 300_000, + "declared".to_string(), + None, + ) + ); + + assert_eq!( + rows[2], + ( + "gateway".to_string(), + 120_000, + 100_000, + 600_000, + 600_000, + "declared".to_string(), + None, + ) + ); + + assert_eq!( + rows[3], + ( + "webapp".to_string(), + 120_000, + 100_000, + 600_000, + 600_000, + "declared".to_string(), + None, + ) + ); + } + + #[tokio::test] + async fn declared_edges_pick_the_revision_as_of_the_window() { + let ctx = SessionContext::new(); + declared_table(&ctx); + let scan = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + // A historical window that predates frontend->db's second revision: + // the first revision was the edge's state then and must not be hidden + // by the newer one. + let window = GraphQueryWindow::from_observed(0, 150_000); + let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let mut rows: Vec<(String, f64)> = vec![]; + for batch in &batches { + let src = strings(batch, 5); + let confidence = batch + .column(10) + .as_any() + .downcast_ref::() + .unwrap(); + for (i, src) in src.into_iter().enumerate() { + rows.push((src, confidence.value(i))); + } + } + rows.sort_by(|a, b| a.0.cmp(&b.0)); + + // api->cache is valid inside this window; batch->queue is not yet. + assert_eq!( + rows, + vec![ + ("agent-1".to_string(), 0.8), + ("api".to_string(), 1.0), + ("frontend".to_string(), 0.5), + ("gateway".to_string(), 1.0), + ("webapp".to_string(), 0.5), + ] + ); + } + + #[tokio::test] + async fn declared_edges_union_calls_branch() { + let ctx = trace_table_ctx(); + declared_table(&ctx); + let trace = ctx.table("opentelemetry_traces").await.unwrap(); + let declared = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + let plan = build_relationships_plan( + vec![CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: trace, + }], + Some(DeclaredSource { scan: declared }), + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let mut provenance: Vec = batches.iter().flat_map(|b| strings(b, 9)).collect(); + provenance.sort(); + // One trace-derived calls edge plus the declared edges valid somewhere + // inside [0s, 600s): frontend, api, webapp, gateway, and the + // agent-asserted one. + assert_eq!( + provenance, + vec![ + "agent", "declared", "declared", "declared", "declared", "trace" + ] + ); + } +} From 3510ef7d4c13f9afa3ad8fdafa8842f98dfd5cb3 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Tue, 11 Aug 2026 16:55:27 +0800 Subject: [PATCH 049/126] feat!: update native histogram unsigned int types (#8824) * feat(native-histogram): store counts and span lengths as signed integers Native histograms are unreleased, so the on-disk integer payload columns are switched from unsigned to signed types without backward-compat: - count_u64 / zero_count_u64: uint64 -> int64 - positive_span_lengths / negative_span_lengths: list(uint32) -> list(int32) - Span.length (query-time model): u32 -> i32 The Prometheus remote-write v2 source carries these as uint64/uint32, so the unsigned->signed conversion at the ingestion boundary is overflow checked: an integer count >= 2^63 or a span length >= 2^31 is rejected with an explicit error rather than silently wrapping to a negative value. read_spans additionally rejects negative stored lengths to keep the non-negative invariant sound for downstream `as usize` casts. The UDAF accumulator's own observation counter (transient aggregation state, not part of the persisted histogram value) is intentionally left as uint64. Signed-off-by: Ning Sun * refactor(native-histogram): rename count/zero_count fields to _i64 Now that the integer payload columns are stored as int64, rename the field constants and persisted names to match: COUNT_U64_FIELD ("count_u64") -> COUNT_I64_FIELD ("count_i64") ZERO_COUNT_U64_FIELD ("zero_count_u64") -> ZERO_COUNT_I64_FIELD ("zero_count_i64") The local builder variables and the docs/JSON snapshot are updated to match. No backward-compat (unreleased feature). Signed-off-by: Ning Sun * test(native-histogram): refresh planner plan snapshot for signed types The mixed native-histogram range test embeds the full histogram Struct type in its expected plan string, which still carried the pre-rename unsigned fields. Update the snapshot to match the signed schema: positive/negative_span_lengths: List(UInt32) -> List(Int32) count_u64/zero_count_u64: UInt64 -> count_i64/zero_count_i64: Int64 Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- src/common/query/src/native_histogram.rs | 67 ++++++++++--------- src/promql/src/functions/native_histogram.rs | 2 +- src/query/src/promql/planner.rs | 26 +++---- src/servers/src/prom_remote_write/README.md | 2 +- src/servers/src/prom_remote_write/v2.rs | 50 +++++++++++--- .../tests/prom_remote_write_v2_test.rs | 6 +- tests-integration/tests/http.rs | 2 +- 7 files changed, 95 insertions(+), 60 deletions(-) diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index 931fb452ede..2db873487bd 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -23,13 +23,13 @@ use std::collections::BTreeMap; use std::sync::Arc; use datafusion::arrow::array::{ - Array, ArrayRef, Float64Array, Int32Array, ListArray, PrimitiveArray, StructArray, - TimestampMillisecondArray, UInt64Array, + Array, ArrayRef, Float64Array, Int32Array, Int64Array, ListArray, PrimitiveArray, StructArray, + TimestampMillisecondArray, }; use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{ ArrowPrimitiveType, DataType as ArrowDataType, Field, Float64Type, Int32Type, Int64Type, - TimestampMillisecondType, UInt32Type, UInt64Type, + TimestampMillisecondType, }; use datafusion_common::{DataFusionError, Result as DfResult}; use datatypes::data_type::{ConcreteDataType, DataType}; @@ -50,8 +50,8 @@ pub const POSITIVE_SPAN_OFFSETS_FIELD: &str = "positive_span_offsets"; pub const POSITIVE_SPAN_LENGTHS_FIELD: &str = "positive_span_lengths"; pub const NEGATIVE_SPAN_OFFSETS_FIELD: &str = "negative_span_offsets"; pub const NEGATIVE_SPAN_LENGTHS_FIELD: &str = "negative_span_lengths"; -pub const COUNT_U64_FIELD: &str = "count_u64"; -pub const ZERO_COUNT_U64_FIELD: &str = "zero_count_u64"; +pub const COUNT_I64_FIELD: &str = "count_i64"; +pub const ZERO_COUNT_I64_FIELD: &str = "zero_count_i64"; pub const POSITIVE_BUCKETS_I64_FIELD: &str = "positive_buckets_i64"; pub const NEGATIVE_BUCKETS_I64_FIELD: &str = "negative_buckets_i64"; pub const COUNT_F64_FIELD: &str = "count_f64"; @@ -72,8 +72,8 @@ pub const NATIVE_HISTOGRAM_FIELD_NAMES: &[&str] = &[ POSITIVE_SPAN_LENGTHS_FIELD, NEGATIVE_SPAN_OFFSETS_FIELD, NEGATIVE_SPAN_LENGTHS_FIELD, - COUNT_U64_FIELD, - ZERO_COUNT_U64_FIELD, + COUNT_I64_FIELD, + ZERO_COUNT_I64_FIELD, POSITIVE_BUCKETS_I64_FIELD, NEGATIVE_BUCKETS_I64_FIELD, COUNT_F64_FIELD, @@ -108,9 +108,9 @@ pub fn native_histogram_field_type(name: &str) -> Option { ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int32_datatype())), ), POSITIVE_SPAN_LENGTHS_FIELD | NEGATIVE_SPAN_LENGTHS_FIELD => Some( - ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::uint32_datatype())), + ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int32_datatype())), ), - COUNT_U64_FIELD | ZERO_COUNT_U64_FIELD => Some(ConcreteDataType::uint64_datatype()), + COUNT_I64_FIELD | ZERO_COUNT_I64_FIELD => Some(ConcreteDataType::int64_datatype()), POSITIVE_BUCKETS_I64_FIELD | NEGATIVE_BUCKETS_I64_FIELD => Some( ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int64_datatype())), ), @@ -195,7 +195,7 @@ pub struct Span { /// The first bucket index, or the gap after the preceding span. pub offset: i32, /// Number of consecutive buckets in the span. - pub length: u32, + pub length: i32, } /// Inclusion rules for a materialized bucket's lower and upper bounds. @@ -1313,7 +1313,7 @@ where .collect() } -fn read_spans(offsets: Vec, lengths: Vec, name: &str) -> DfResult> { +fn read_spans(offsets: Vec, lengths: Vec, name: &str) -> DfResult> { if offsets.len() != lengths.len() { return Err(DataFusionError::Execution(format!( "native histogram {name} span offsets and lengths mismatch: {} vs {}", @@ -1321,11 +1321,18 @@ fn read_spans(offsets: Vec, lengths: Vec, name: &str) -> DfResult DfResult<()> { @@ -1352,12 +1359,12 @@ pub fn read_histogram(array: &StructArray, row: usize) -> DfResult(array, SCHEMA_FIELD, row)?; let positive_spans = read_spans( list_values::(array, POSITIVE_SPAN_OFFSETS_FIELD, row)?, - list_values::(array, POSITIVE_SPAN_LENGTHS_FIELD, row)?, + list_values::(array, POSITIVE_SPAN_LENGTHS_FIELD, row)?, "positive", )?; let negative_spans = read_spans( list_values::(array, NEGATIVE_SPAN_OFFSETS_FIELD, row)?, - list_values::(array, NEGATIVE_SPAN_LENGTHS_FIELD, row)?, + list_values::(array, NEGATIVE_SPAN_LENGTHS_FIELD, row)?, "negative", )?; @@ -1372,8 +1379,8 @@ pub fn read_histogram(array: &StructArray, row: usize) -> DfResult(array, COUNT_U64_FIELD, row)? as f64, - optional_primitive::(array, ZERO_COUNT_U64_FIELD, row)? + required_primitive::(array, COUNT_I64_FIELD, row)? as f64, + optional_primitive::(array, ZERO_COUNT_I64_FIELD, row)? .unwrap_or_default() as f64, list_values::(array, POSITIVE_BUCKETS_I64_FIELD, row)? .into_iter() @@ -1429,8 +1436,8 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { let mut positive_span_lengths = Vec::with_capacity(values.len()); let mut negative_span_offsets = Vec::with_capacity(values.len()); let mut negative_span_lengths = Vec::with_capacity(values.len()); - let mut count_u64 = Vec::>::with_capacity(values.len()); - let mut zero_count_u64 = Vec::>::with_capacity(values.len()); + let mut count_i64 = Vec::>::with_capacity(values.len()); + let mut zero_count_i64 = Vec::>::with_capacity(values.len()); let mut positive_buckets_i64 = Vec::with_capacity(values.len()); let mut negative_buckets_i64 = Vec::with_capacity(values.len()); let mut count_f64 = Vec::with_capacity(values.len()); @@ -1476,8 +1483,8 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { .map(|span| span.length) .collect(), )); - count_u64.push(None); - zero_count_u64.push(None); + count_i64.push(None); + zero_count_i64.push(None); positive_buckets_i64.push(list_opt(Vec::::new())); negative_buckets_i64.push(list_opt(Vec::::new())); count_f64.push(Some(histogram.count)); @@ -1495,8 +1502,8 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { positive_span_lengths.push(None); negative_span_offsets.push(None); negative_span_lengths.push(None); - count_u64.push(None); - zero_count_u64.push(None); + count_i64.push(None); + zero_count_i64.push(None); positive_buckets_i64.push(None); negative_buckets_i64.push(None); count_f64.push(None); @@ -1532,7 +1539,7 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { ), ( POSITIVE_SPAN_LENGTHS_FIELD, - Arc::new(ListArray::from_iter_primitive::( + Arc::new(ListArray::from_iter_primitive::( positive_span_lengths, )), ), @@ -1544,14 +1551,14 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { ), ( NEGATIVE_SPAN_LENGTHS_FIELD, - Arc::new(ListArray::from_iter_primitive::( + Arc::new(ListArray::from_iter_primitive::( negative_span_lengths, )), ), - (COUNT_U64_FIELD, Arc::new(UInt64Array::from(count_u64))), + (COUNT_I64_FIELD, Arc::new(Int64Array::from(count_i64))), ( - ZERO_COUNT_U64_FIELD, - Arc::new(UInt64Array::from(zero_count_u64)), + ZERO_COUNT_I64_FIELD, + Arc::new(Int64Array::from(zero_count_i64)), ), ( POSITIVE_BUCKETS_I64_FIELD, @@ -1662,7 +1669,7 @@ mod tests { assert_eq!(read_histogram(array, 0).unwrap(), Some(expected)); assert!( - primitive_child::(array, COUNT_U64_FIELD) + primitive_child::(array, COUNT_I64_FIELD) .unwrap() .is_null(0) ); diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index e6c9e84d59f..5b633defa96 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -2751,7 +2751,7 @@ mod tests { custom_values: Vec::new(), positive_spans: vec![Span { offset: 0, - length: positive_buckets.len() as u32, + length: positive_buckets.len() as i32, }], negative_spans: Vec::new(), count, diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 10462ea1828..0933bbcaca4 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -10026,19 +10026,19 @@ mod test { .unwrap() .display_indent_schema() .to_string(); - let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] - PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] - PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; + let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; assert_eq!(plan, expected); } diff --git a/src/servers/src/prom_remote_write/README.md b/src/servers/src/prom_remote_write/README.md index 1f15f8288a7..d21eead1901 100644 --- a/src/servers/src/prom_remote_write/README.md +++ b/src/servers/src/prom_remote_write/README.md @@ -47,7 +47,7 @@ Each histogram row stores one Struct field named - common scalar children: `schema`, `zero_threshold`, `sum`, `reset_hint`, `start_timestamp`; -- count children: `count_u64` / `zero_count_u64` or `count_f64` / `zero_count_f64`; +- count children: `count_i64` / `zero_count_i64` or `count_f64` / `zero_count_f64`; - list children for custom values, spans, and positive/negative buckets; - original Prometheus labels as Greptime tags. diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 9940ad2e7c2..5a25feb566a 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -371,6 +371,8 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { validate_native_histogram(histogram, uses_float_counts)?; let mut items = Vec::with_capacity(NATIVE_HISTOGRAM_FIELD_NAMES.len()); + let positive_span_lengths = i32_span_lengths("positive", &histogram.positive_spans)?; + let negative_span_lengths = i32_span_lengths("negative", &histogram.negative_spans)?; items.extend([ pb_value(ValueData::I32Value(histogram.schema)), pb_value(ValueData::F64Value(histogram.zero_threshold)), @@ -381,9 +383,9 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { )), f64_list_value(histogram.custom_values.iter().copied()), i32_list_value(histogram.positive_spans.iter().map(|span| span.offset)), - u32_list_value(histogram.positive_spans.iter().map(|span| span.length)), + i32_list_value(positive_span_lengths.iter().copied()), i32_list_value(histogram.negative_spans.iter().map(|span| span.offset)), - u32_list_value(histogram.negative_spans.iter().map(|span| span.length)), + i32_list_value(negative_span_lengths.iter().copied()), ]); if uses_float_counts { @@ -418,9 +420,23 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { let positive_buckets = bucket_counts_from_deltas(&histogram.positive_deltas)?; let negative_buckets = bucket_counts_from_deltas(&histogram.negative_deltas)?; validate_integer_native_histogram_counts(histogram, &positive_buckets, &negative_buckets)?; + let count = i64::try_from(count) + .ok() + .context(error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram integer count {count} overflows i64" + ), + })?; + let zero_count = i64::try_from(zero_count).ok().context( + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram integer zero_count {zero_count} overflows i64" + ), + }, + )?; items.extend([ - pb_value(ValueData::U64Value(count)), - pb_value(ValueData::U64Value(zero_count)), + pb_value(ValueData::I64Value(count)), + pb_value(ValueData::I64Value(zero_count)), i64_list_value(positive_buckets.iter().copied()), i64_list_value(negative_buckets.iter().copied()), null_pb_value(), @@ -779,8 +795,20 @@ fn i32_list_value(values: impl IntoIterator) -> Value { list_value(values.into_iter().map(ValueData::I32Value)) } -fn u32_list_value(values: impl IntoIterator) -> Value { - list_value(values.into_iter().map(ValueData::U32Value)) +fn i32_span_lengths(name: &str, spans: &[BucketSpan]) -> Result> { + spans + .iter() + .map(|span| { + i32::try_from(span.length) + .ok() + .context(error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram {name} span length {} overflows i32", + span.length + ), + }) + }) + .collect() } fn i64_list_value(values: impl IntoIterator) -> Value { @@ -1708,8 +1736,8 @@ mod tests { Some(ValueData::I32Value(0)) ); assert_eq!( - histogram_field_value(&rows, 0, COUNT_U64_FIELD), - Some(ValueData::U64Value(0)) + histogram_field_value(&rows, 0, COUNT_I64_FIELD), + Some(ValueData::I64Value(0)) ); assert_eq!(histogram_field_value(&rows, 0, COUNT_F64_FIELD), None); } @@ -1825,8 +1853,8 @@ mod tests { ); assert_eq!( - histogram_field_value(&rows, 0, COUNT_U64_FIELD), - Some(ValueData::U64Value(0)) + histogram_field_value(&rows, 0, COUNT_I64_FIELD), + Some(ValueData::I64Value(0)) ); assert_eq!(histogram_field_value(&rows, 0, COUNT_F64_FIELD), None); assert!(matches!( @@ -1839,7 +1867,7 @@ mod tests { POSITIVE_BUCKETS_F64_FIELD ))); - assert_eq!(histogram_field_value(&rows, 1, COUNT_U64_FIELD), None); + assert_eq!(histogram_field_value(&rows, 1, COUNT_I64_FIELD), None); assert_eq!( histogram_field_value(&rows, 1, COUNT_F64_FIELD), Some(ValueData::F64Value(6.0)) diff --git a/src/servers/tests/prom_remote_write_v2_test.rs b/src/servers/tests/prom_remote_write_v2_test.rs index a44583dccd1..94d260d3fac 100644 --- a/src/servers/tests/prom_remote_write_v2_test.rs +++ b/src/servers/tests/prom_remote_write_v2_test.rs @@ -18,7 +18,7 @@ use api::v1::value::ValueData; use api::v1::{ColumnSchema, Rows}; use bytes::Bytes; use common_query::native_histogram::{ - COUNT_U64_FIELD, NATIVE_HISTOGRAM_FIELD_NAMES, POSITIVE_BUCKETS_F64_FIELD, + COUNT_I64_FIELD, NATIVE_HISTOGRAM_FIELD_NAMES, POSITIVE_BUCKETS_F64_FIELD, POSITIVE_BUCKETS_I64_FIELD, POSITIVE_SPAN_OFFSETS_FIELD, SCHEMA_FIELD, }; use common_query::prelude::greptime_native_histogram; @@ -100,8 +100,8 @@ fn test_decode_remote_write_v2_native_histogram_dump() { Some(ValueData::I32Value(3)) ); assert_eq!( - histogram_field_value(rows, 0, COUNT_U64_FIELD), - Some(ValueData::U64Value(24)) + histogram_field_value(rows, 0, COUNT_I64_FIELD), + Some(ValueData::I64Value(24)) ); assert_eq!( list_i32_values(histogram_field_value(rows, 0, POSITIVE_SPAN_OFFSETS_FIELD)), diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 9d2f1eb8cb5..cd8af8768a4 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2878,7 +2878,7 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag "prometheus_remote_write_v2_native_histogram_rows", &client, "select greptime_timestamp, greptime_native_histogram, job, instance from remote_write_v2_latency_seconds order by greptime_timestamp;", - "[[3000,{\"count_f64\":null,\"count_u64\":8,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_u64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_u64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_u64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]", + "[[3000,{\"count_f64\":null,\"count_i64\":8,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_i64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_i64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_i64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]", ) .await; From 72f6cf09bf7f7d78c84514163554ad2a29fc7f05 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 11 Aug 2026 17:05:26 +0800 Subject: [PATCH 050/126] refactor(procedure): centralize event context handling (#8834) * refactor(procedure): centralize event context handling Signed-off-by: WenyXu * refactor(meta): simplify migration trigger reason handling Signed-off-by: WenyXu * refactor(meta): avoid cloning event context Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 1 + src/common/event-recorder/Cargo.toml | 1 + src/common/event-recorder/src/context.rs | 104 ++++++++ src/common/event-recorder/src/lib.rs | 2 + src/common/meta/src/ddl/alter_database.rs | 23 +- .../meta/src/ddl/alter_logical_tables.rs | 14 +- src/common/meta/src/ddl/alter_table.rs | 33 +-- .../src/ddl/alter_table/region_request.rs | 8 +- src/common/meta/src/ddl/create_database.rs | 19 +- src/common/meta/src/ddl/create_flow.rs | 22 +- .../meta/src/ddl/create_logical_tables.rs | 13 +- src/common/meta/src/ddl/create_table.rs | 34 +-- src/common/meta/src/ddl/create_view.rs | 14 +- src/common/meta/src/ddl/drop_database.rs | 25 +- src/common/meta/src/ddl/drop_flow.rs | 15 +- src/common/meta/src/ddl/drop_table.rs | 23 +- src/common/meta/src/ddl/drop_view.rs | 12 +- src/common/meta/src/ddl/event/database.rs | 25 +- src/common/meta/src/ddl/event/flow.rs | 28 +-- src/common/meta/src/ddl/event/table.rs | 57 +---- src/common/meta/src/ddl/event/view.rs | 25 +- .../meta/src/ddl/purge_dropped_table.rs | 36 +-- src/common/meta/src/ddl/test_util.rs | 9 +- .../src/ddl/tests/alter_logical_tables.rs | 52 +--- src/common/meta/src/ddl/tests/alter_table.rs | 73 ++---- src/common/meta/src/ddl/tests/create_flow.rs | 87 ++----- .../src/ddl/tests/create_logical_tables.rs | 39 +-- src/common/meta/src/ddl/tests/create_table.rs | 78 ++---- src/common/meta/src/ddl/tests/create_view.rs | 34 ++- .../meta/src/ddl/tests/drop_database.rs | 6 +- src/common/meta/src/ddl/tests/drop_flow.rs | 14 +- src/common/meta/src/ddl/tests/drop_table.rs | 225 +++++------------- src/common/meta/src/ddl/tests/drop_view.rs | 14 +- .../meta/src/ddl/tests/event/database.rs | 48 +--- src/common/meta/src/ddl/tests/event/flow.rs | 25 +- src/common/meta/src/ddl/tests/event/table.rs | 166 ++++--------- .../meta/src/ddl/tests/event/test_util.rs | 21 +- src/common/meta/src/ddl/tests/event/view.rs | 67 ++---- src/common/meta/src/ddl/truncate_table.rs | 24 +- src/common/meta/src/ddl/undrop_table.rs | 50 +--- src/common/meta/src/ddl_manager.rs | 210 ++++++++-------- src/common/meta/src/procedure_executor.rs | 5 +- src/common/meta/src/rpc/ddl.rs | 92 ++----- src/common/procedure-test/src/lib.rs | 4 + src/common/procedure/src/event.rs | 82 ++++++- src/common/procedure/src/lib.rs | 4 +- src/common/procedure/src/local.rs | 80 ++++++- src/common/procedure/src/local/runner.rs | 32 ++- src/common/procedure/src/procedure.rs | 44 +++- src/common/procedure/src/store.rs | 46 +++- src/common/procedure/src/watcher.rs | 1 + src/meta-srv/src/event/gc.rs | 86 +++---- src/meta-srv/src/event/region_migration.rs | 22 +- src/meta-srv/src/event/repartition.rs | 46 +--- src/meta-srv/src/gc/ctx.rs | 12 +- src/meta-srv/src/gc/handler.rs | 11 +- src/meta-srv/src/gc/mock.rs | 6 +- src/meta-srv/src/gc/mock/basic.rs | 10 +- src/meta-srv/src/gc/mock/concurrent.rs | 12 +- src/meta-srv/src/gc/mock/full_list.rs | 12 +- src/meta-srv/src/gc/mock/misc.rs | 6 +- src/meta-srv/src/gc/procedure.rs | 36 +-- src/meta-srv/src/gc/scheduler.rs | 16 +- .../src/procedure/region_migration.rs | 152 +++++++----- .../close_downgraded_region.rs | 4 +- .../downgrade_leader_region.rs | 7 +- .../src/procedure/region_migration/manager.rs | 14 +- .../region_migration/migration_abort.rs | 4 +- .../region_migration/open_candidate_region.rs | 30 ++- .../procedure/region_migration/test_util.rs | 8 +- .../upgrade_candidate_region.rs | 7 +- src/meta-srv/src/procedure/repartition.rs | 47 ++-- .../procedure/repartition/allocate_region.rs | 1 + .../src/procedure/repartition/collect.rs | 4 +- .../procedure/repartition/gc_requirement.rs | 1 + .../src/procedure/repartition/group.rs | 22 +- .../src/procedure/repartition/test_util.rs | 4 +- src/meta-srv/src/procedure/tests.rs | 14 +- src/meta-srv/src/procedure/wal_prune.rs | 12 +- src/meta-srv/src/service/procedure.rs | 1 + src/standalone/src/procedure.rs | 2 - tests-integration/src/tests/gc.rs | 2 - tests-integration/tests/gc_event.rs | 6 +- tests-integration/tests/repartition.rs | 2 - tests-integration/tests/table_ddl_event.rs | 4 +- tests/conf/frontend-test.toml.template | 3 + 86 files changed, 1142 insertions(+), 1650 deletions(-) create mode 100644 src/common/event-recorder/src/context.rs diff --git a/Cargo.lock b/Cargo.lock index 2e6fe8ae92c..18248bd442d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2409,6 +2409,7 @@ dependencies = [ "serde_json", "snafu 0.8.6", "store-api", + "strum 0.27.1", "tokio", "tokio-util", "toml 0.8.23", diff --git a/src/common/event-recorder/Cargo.toml b/src/common/event-recorder/Cargo.toml index 7990cd3d183..9a033d7009d 100644 --- a/src/common/event-recorder/Cargo.toml +++ b/src/common/event-recorder/Cargo.toml @@ -23,6 +23,7 @@ serde.workspace = true serde_json.workspace = true snafu.workspace = true store-api.workspace = true +strum.workspace = true tokio.workspace = true tokio-util.workspace = true diff --git a/src/common/event-recorder/src/context.rs b/src/common/event-recorder/src/context.rs new file mode 100644 index 00000000000..601276967e3 --- /dev/null +++ b/src/common/event-recorder/src/context.rs @@ -0,0 +1,104 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use serde::{Deserialize, Serialize}; +use strum::{AsRefStr, EnumString}; + +/// Stable context recorded for a procedure event. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PersistentEventContext { + pub reason: TriggerReason, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub protocol: Option, + #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")] + pub extensions: serde_json::Map, +} + +impl PersistentEventContext { + /// Creates an event context with no additional extensions. + pub fn new(reason: TriggerReason) -> Self { + Self { + reason, + protocol: None, + extensions: Default::default(), + } + } + + /// Adds the protocol that originated the operation. + pub fn with_protocol(mut self, protocol: impl Into) -> Self { + self.protocol = Some(protocol.into()); + self + } +} + +impl Default for PersistentEventContext { + fn default() -> Self { + Self::new(TriggerReason::default()) + } +} + +/// Stable classification of a procedure trigger. +#[derive( + Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, AsRefStr, EnumString, +)] +#[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +pub enum TriggerReason { + Manual, + AutoCreate, + AutoAlter, + AutoRepartition, + AutoRebalance, + RegionFailover, + ScheduledGc, + #[default] + #[serde(other)] + Unknown, +} + +impl TriggerReason { + pub fn from_extension(value: &str) -> Self { + value.parse().unwrap_or_default() + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn test_event_context_serialization() { + let context = PersistentEventContext::new(TriggerReason::Manual).with_protocol("mysql"); + + assert_eq!( + json!({ + "reason": "manual", + "protocol": "mysql", + }), + serde_json::to_value(context).unwrap() + ); + assert_eq!( + json!({ "reason": "manual" }), + serde_json::to_value(PersistentEventContext::new(TriggerReason::Manual)).unwrap() + ); + } + + #[test] + fn test_trigger_reason_deserializes_unknown_value() { + let reason: TriggerReason = serde_json::from_str("\"future_reason\"").unwrap(); + assert_eq!(TriggerReason::Unknown, reason); + } +} diff --git a/src/common/event-recorder/src/lib.rs b/src/common/event-recorder/src/lib.rs index 3e624824980..5b9cb4a7b97 100644 --- a/src/common/event-recorder/src/lib.rs +++ b/src/common/event-recorder/src/lib.rs @@ -14,6 +14,7 @@ #![feature(duration_constructors)] +pub mod context; pub mod error; pub mod event_table; pub mod recorder; @@ -21,4 +22,5 @@ pub mod recorder; #[cfg(any(test, feature = "testing"))] pub mod testing; +pub use context::{PersistentEventContext, TriggerReason}; pub use recorder::*; diff --git a/src/common/meta/src/ddl/alter_database.rs b/src/common/meta/src/ddl/alter_database.rs index f9344213025..ac24ff77e11 100644 --- a/src/common/meta/src/ddl/alter_database.rs +++ b/src/common/meta/src/ddl/alter_database.rs @@ -15,7 +15,7 @@ use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Status, }; use common_telemetry::tracing::info; use serde::{Deserialize, Serialize}; @@ -32,7 +32,7 @@ use crate::key::DeserializedValueWithBytes; use crate::key::schema_name::{SchemaName, SchemaNameKey, SchemaNameValue}; use crate::lock_key::{CatalogLock, SchemaLock}; use crate::rpc::ddl::UnsetDatabaseOption::{self}; -use crate::rpc::ddl::{AlterDatabaseKind, AlterDatabaseTask, EventContext, SetDatabaseOption}; +use crate::rpc::ddl::{AlterDatabaseKind, AlterDatabaseTask, SetDatabaseOption}; pub struct AlterDatabaseProcedure { pub context: DdlContext, @@ -73,14 +73,10 @@ fn build_new_schema_value( impl AlterDatabaseProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::AlterDatabase"; - pub fn new( - task: AlterDatabaseTask, - event_context: EventContext, - context: DdlContext, - ) -> Result { + pub fn new(task: AlterDatabaseTask, context: DdlContext) -> Result { Ok(Self { context, - data: AlterDatabaseData::new(task, event_context)?, + data: AlterDatabaseData::new(task)?, }) } @@ -180,10 +176,7 @@ impl Procedure for AlterDatabaseProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(ALTER_DATABASE_EVENT_TYPE) { return None; } @@ -193,7 +186,6 @@ impl Procedure for AlterDatabaseProcedure { self.data.catalog(), self.data.schema(), &self.data.kind, - self.data.event_context.clone(), ) } else { DatabaseDdlEvent::alter_lifecycle(self.data.catalog(), self.data.schema()) @@ -217,19 +209,16 @@ pub struct AlterDatabaseData { catalog_name: String, schema_name: String, schema_value: Option>, - #[serde(default)] - event_context: EventContext, } impl AlterDatabaseData { - pub fn new(task: AlterDatabaseTask, event_context: EventContext) -> Result { + pub fn new(task: AlterDatabaseTask) -> Result { Ok(Self { state: AlterDatabaseState::Prepare, kind: AlterDatabaseKind::try_from(task.alter_expr.kind.unwrap())?, catalog_name: task.alter_expr.catalog_name, schema_name: task.alter_expr.schema_name, schema_value: None, - event_context, }) } diff --git a/src/common/meta/src/ddl/alter_logical_tables.rs b/src/common/meta/src/ddl/alter_logical_tables.rs index f2f15430933..9207f910577 100644 --- a/src/common/meta/src/ddl/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/alter_logical_tables.rs @@ -20,7 +20,7 @@ use api::region::RegionResponse; use async_trait::async_trait; use common_catalog::format_full_table_name; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; -use common_procedure::{Context, EventRuntimeContext, EventTrigger, LockKey, Procedure, Status}; +use common_procedure::{Context, EventContext, EventTrigger, LockKey, Procedure, Status}; use common_telemetry::{debug, error, info, warn}; pub use executor::make_alter_region_request; use serde::{Deserialize, Serialize}; @@ -47,7 +47,7 @@ use crate::key::table_info::TableInfoValue; use crate::key::table_route::PhysicalTableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; use crate::metrics; -use crate::rpc::ddl::{AlterTableTask, EventContext}; +use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::RegionRoute; pub struct AlterLogicalTablesProcedure { @@ -88,7 +88,6 @@ impl AlterLogicalTablesProcedure { pub fn new( tasks: Vec, physical_table_id: TableId, - event_context: EventContext, context: DdlContext, ) -> Self { Self { @@ -101,7 +100,6 @@ impl AlterLogicalTablesProcedure { physical_table_info: None, physical_columns: vec![], table_cache_keys_to_invalidate: vec![], - event_context, }, physical_table_route: None, } @@ -322,10 +320,7 @@ impl Procedure for AlterLogicalTablesProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::AlterLogicalTables.as_str()) @@ -358,7 +353,6 @@ impl Procedure for AlterLogicalTablesProcedure { locators, self.data.tasks.len(), kinds, - self.data.event_context.clone(), ))) } } @@ -375,8 +369,6 @@ pub struct AlterTablesData { physical_table_info: Option>, physical_columns: Vec, table_cache_keys_to_invalidate: Vec, - #[serde(default)] - event_context: EventContext, } impl AlterTablesData { diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index 21975fab714..63342a7bbb6 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -27,7 +27,7 @@ use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE}; use common_error::ext::BoxedError; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, ContextProvider, Error as ProcedureError, EventRuntimeContext, + Context as ProcedureContext, ContextProvider, Error as ProcedureError, EventContext, EventTrigger, LockKey, PoisonKey, PoisonKeys, Procedure, ProcedureId, Status, StringKey, }; use common_telemetry::{error, info, warn}; @@ -58,7 +58,7 @@ use crate::key::{DeserializedValueWithBytes, RegionDistribution}; use crate::lock_key::{CatalogLock, RegionLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; use crate::poison_key::table_poison_key; -use crate::rpc::ddl::{AlterTableTask, EventContext}; +use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::{RegionRoute, find_leaders, region_distribution}; /// The alter table procedure @@ -97,24 +97,18 @@ fn build_executor_from_alter_expr(alter_data: &AlterTableData) -> AlterTableExec impl AlterTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::AlterTable"; - pub fn new( - table_id: TableId, - task: AlterTableTask, - event_context: EventContext, - context: DdlContext, - ) -> Result { - Self::new_with_region_locks(table_id, task, event_context, vec![], context) + pub fn new(table_id: TableId, task: AlterTableTask, context: DdlContext) -> Result { + Self::new_with_region_locks(table_id, task, vec![], context) } pub(crate) fn new_with_region_locks( table_id: TableId, task: AlterTableTask, - event_context: EventContext, region_locks: Vec, context: DdlContext, ) -> Result { task.validate()?; - let data = AlterTableData::new(task, table_id, event_context, region_locks); + let data = AlterTableData::new(task, table_id, region_locks); let executor = build_executor_from_alter_expr(&data); Ok(Self { context, @@ -569,10 +563,7 @@ impl Procedure for AlterTableProcedure { PoisonKeys::new(vec![self.table_poison_key()]) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::AlterTable.as_str()) @@ -591,7 +582,7 @@ impl Procedure for AlterTableProcedure { .kind .as_ref() .and_then(alter_table_kind_name); - TableDdlEvent::alter_table_submitted(locator, kind, self.data.event_context.clone()) + TableDdlEvent::alter_table_submitted(locator, kind) } _ => TableDdlEvent::lifecycle(TableDdlEventType::AlterTable, [locator]), }; @@ -619,8 +610,6 @@ pub struct AlterTableData { task: AlterTableTask, table_id: TableId, #[serde(default)] - event_context: EventContext, - #[serde(default)] column_metadatas: Vec, /// Table info value before alteration. table_info_value: Option>, @@ -632,17 +621,11 @@ pub struct AlterTableData { } impl AlterTableData { - pub fn new( - task: AlterTableTask, - table_id: TableId, - event_context: EventContext, - region_locks: Vec, - ) -> Self { + pub fn new(task: AlterTableTask, table_id: TableId, region_locks: Vec) -> Self { Self { state: AlterTableState::Prepare, task, table_id, - event_context, column_metadatas: vec![], table_info_value: None, region_distribution: None, diff --git a/src/common/meta/src/ddl/alter_table/region_request.rs b/src/common/meta/src/ddl/alter_table/region_request.rs index ef8d1487c70..73edd4c3615 100644 --- a/src/common/meta/src/ddl/alter_table/region_request.rs +++ b/src/common/meta/src/ddl/alter_table/region_request.rs @@ -144,7 +144,7 @@ mod tests { }; use crate::key::table_route::TableRouteValue; use crate::peer::Peer; - use crate::rpc::ddl::{AlterTableTask, EventContext}; + use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -240,8 +240,7 @@ mod tests { }, }; - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context).unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let alter_kind = procedure.make_region_alter_kind().unwrap(); let Some(Body::Alter(alter_region_request)) = @@ -297,8 +296,7 @@ mod tests { }, }; - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context).unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let alter_kind = procedure.make_region_alter_kind().unwrap(); let Some(Body::Alter(alter_region_request)) = diff --git a/src/common/meta/src/ddl/create_database.rs b/src/common/meta/src/ddl/create_database.rs index b48e700a519..82d467888a5 100644 --- a/src/common/meta/src/ddl/create_database.rs +++ b/src/common/meta/src/ddl/create_database.rs @@ -19,8 +19,8 @@ use async_trait::async_trait; use common_error::ext::{BoxedError, ErrorExt}; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, - ProcedureId, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureId, + Status, }; use serde::{Deserialize, Serialize}; use serde_with::{DefaultOnNull, serde_as}; @@ -34,7 +34,7 @@ use crate::error::{self, Result}; use crate::instruction::{CacheIdent, UserCacheIdent}; use crate::key::schema_name::{SchemaNameKey, SchemaNameValue}; use crate::lock_key::{CatalogLock, SchemaLock}; -use crate::rpc::ddl::{CreatorGrantIntent, EventContext}; +use crate::rpc::ddl::CreatorGrantIntent; /// Describes the creator-access result of an atomic create. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -93,7 +93,6 @@ impl CreateDatabaseProcedure { create_if_not_exists: bool, options: HashMap, creator: Option, - event_context: EventContext, context: DdlContext, ) -> Self { Self { @@ -105,7 +104,6 @@ impl CreateDatabaseProcedure { create_if_not_exists, options, creator, - event_context, }, } } @@ -280,10 +278,7 @@ impl Procedure for CreateDatabaseProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(CREATE_DATABASE_EVENT_TYPE) { return None; } @@ -294,7 +289,6 @@ impl Procedure for CreateDatabaseProcedure { &self.data.schema, self.data.create_if_not_exists, &self.data.options, - self.data.event_context.clone(), ) } else { DatabaseDdlEvent::create_lifecycle(&self.data.catalog, &self.data.schema) @@ -327,8 +321,6 @@ pub struct CreateDatabaseData { /// Authenticated creator whose access is ensured, absent for legacy schema-only requests. #[serde(default)] pub creator: Option, - #[serde(default)] - pub event_context: EventContext, } #[cfg(test)] @@ -423,7 +415,6 @@ mod tests { username: "alice".to_string(), created_at_ns: 1, }), - EventContext::default(), context, ); (procedure, committer) @@ -583,7 +574,6 @@ mod tests { false, HashMap::new(), None, - EventContext::default(), context, ); procedure.data.state = CreateDatabaseState::CreateMetadata; @@ -619,7 +609,6 @@ mod tests { username: "alice".to_string(), created_at_ns: 1, }), - EventContext::default(), context, ); procedure.data.state = CreateDatabaseState::CreateMetadata; diff --git a/src/common/meta/src/ddl/create_flow.rs b/src/common/meta/src/ddl/create_flow.rs index 3ef4dadbc63..a2f306306d8 100644 --- a/src/common/meta/src/ddl/create_flow.rs +++ b/src/common/meta/src/ddl/create_flow.rs @@ -24,8 +24,8 @@ use async_trait::async_trait; use common_catalog::format_full_flow_name; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, - ProcedureState, Result as ProcedureResult, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureState, + Result as ProcedureResult, Status, }; use common_telemetry::info; use common_telemetry::tracing_context::TracingContext; @@ -50,7 +50,7 @@ use crate::key::{DeserializedValueWithBytes, FlowId, FlowPartitionId}; use crate::lock_key::{CatalogLock, FlowNameLock}; use crate::metrics; use crate::peer::Peer; -use crate::rpc::ddl::{CreateFlowTask, EventContext, FlowQueryContext, QueryContext}; +use crate::rpc::ddl::{CreateFlowTask, FlowQueryContext, QueryContext}; /// The procedure of flow creation. pub struct CreateFlowProcedure { @@ -62,12 +62,7 @@ impl CreateFlowProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::CreateFlow"; /// Returns a new [CreateFlowProcedure]. - pub fn new( - task: CreateFlowTask, - query_context: QueryContext, - event_context: EventContext, - context: DdlContext, - ) -> Self { + pub fn new(task: CreateFlowTask, query_context: QueryContext, context: DdlContext) -> Self { Self { context, data: CreateFlowData { @@ -81,7 +76,6 @@ impl CreateFlowProcedure { prev_flow_info_value: None, did_replace: false, flow_type: None, - event_context, }, } } @@ -410,10 +404,7 @@ impl Procedure for CreateFlowProcedure { ]) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(CREATE_FLOW_EVENT_TYPE) { return None; } @@ -428,7 +419,6 @@ impl Procedure for CreateFlowProcedure { expire_after: self.data.task.expire_after, eval_interval_secs: self.data.task.eval_interval_secs, }, - self.data.event_context.clone(), ), EventTrigger::Succeeded => { let flow_id = match ctx.lifecycle_state { @@ -718,8 +708,6 @@ pub struct CreateFlowData { #[serde(default)] pub(crate) did_replace: bool, pub(crate) flow_type: Option, - #[serde(default)] - pub(crate) event_context: EventContext, } impl CreateFlowData { diff --git a/src/common/meta/src/ddl/create_logical_tables.rs b/src/common/meta/src/ddl/create_logical_tables.rs index 1c9efb865e4..9b231f3c538 100644 --- a/src/common/meta/src/ddl/create_logical_tables.rs +++ b/src/common/meta/src/ddl/create_logical_tables.rs @@ -24,8 +24,8 @@ use common_catalog::consts::METRIC_ENGINE; use common_event_recorder::Event; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, - ProcedureState, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureState, + Status, }; use common_telemetry::{debug, error, warn}; use futures::future; @@ -48,7 +48,7 @@ use crate::error::Result; use crate::key::table_route::TableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; -use crate::rpc::ddl::{CreateTableTask, EventContext}; +use crate::rpc::ddl::CreateTableTask; use crate::rpc::router::{RegionRoute, find_leaders}; pub struct CreateLogicalTablesProcedure { @@ -62,7 +62,6 @@ impl CreateLogicalTablesProcedure { pub fn new( tasks: Vec, physical_table_id: TableId, - event_context: EventContext, context: DdlContext, ) -> Self { Self { @@ -75,7 +74,6 @@ impl CreateLogicalTablesProcedure { physical_region_numbers: vec![], physical_columns: vec![], physical_partition_columns: vec![], - event_context, }, } } @@ -273,7 +271,7 @@ impl Procedure for CreateLogicalTablesProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::CreateLogicalTables.as_str()) @@ -284,7 +282,6 @@ impl Procedure for CreateLogicalTablesProcedure { EventTrigger::Submitted => TableDdlEvent::create_logical_tables_submitted( self.event_locators(), self.data.tasks.len(), - self.data.event_context.clone(), ), EventTrigger::Succeeded => match ctx.lifecycle_state { ProcedureState::Done { @@ -329,8 +326,6 @@ pub struct CreateTablesData { physical_region_numbers: Vec, physical_columns: Vec, physical_partition_columns: Vec, - #[serde(default)] - event_context: EventContext, } impl CreateTablesData { diff --git a/src/common/meta/src/ddl/create_table.rs b/src/common/meta/src/ddl/create_table.rs index c63d7f7a34d..ce3e1aa060a 100644 --- a/src/common/meta/src/ddl/create_table.rs +++ b/src/common/meta/src/ddl/create_table.rs @@ -23,8 +23,8 @@ use common_procedure::error::{ }; use common_procedure::local::DynamicKeyLockGuard; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, - ProcedureId, ProcedureState, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureId, + ProcedureState, Status, }; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -47,7 +47,7 @@ use crate::lock_key::{CatalogLock, SchemaLock, TableNameLock}; use crate::metrics; use crate::peer::PeerAllocContext; use crate::region_keeper::OperatingRegionGuard; -use crate::rpc::ddl::{CreateTableTask, EventContext, QueryContext}; +use crate::rpc::ddl::{CreateTableTask, QueryContext}; use crate::rpc::router::{RegionRoute, operating_leader_region_roles}; use crate::wal_provider::{ RegionWalOptions, acquire_remote_wal_read_locks, optional_region_wal_options_serde, @@ -84,17 +84,20 @@ fn build_executor_from_create_table_data( impl CreateTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::CreateTable"; - pub fn new( + pub fn new(task: CreateTableTask, context: DdlContext) -> Result { + Self::new_with_query_context(task, QueryContext::default(), context) + } + + pub fn new_with_query_context( task: CreateTableTask, query_context: QueryContext, - event_context: EventContext, context: DdlContext, ) -> Result { let executor = build_executor_from_create_table_data(&task.create_table)?; Ok(Self { context, - data: CreateTableData::new_with_query_context(task, event_context, query_context), + data: CreateTableData::new(task, query_context), opening_regions: vec![], executor, remote_wal_lock_guards: vec![], @@ -396,10 +399,7 @@ impl Procedure for CreateTableProcedure { ]) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::CreateTable.as_str()) @@ -415,7 +415,6 @@ impl Procedure for CreateTableProcedure { locator, create_table.create_if_not_exists, &create_table.engine, - self.data.event_context.clone(), ) } EventTrigger::Succeeded => match ctx.lifecycle_state { @@ -456,8 +455,6 @@ pub struct CreateTableData { #[serde(default)] pub query_context: QueryContext, #[serde(default)] - pub event_context: EventContext, - #[serde(default)] pub column_metadatas: Vec, /// None stands for not allocated yet. pub(crate) table_route: Option, @@ -468,21 +465,12 @@ pub struct CreateTableData { } impl CreateTableData { - pub fn new(task: CreateTableTask, event_context: EventContext) -> Self { - Self::new_with_query_context(task, event_context, QueryContext::default()) - } - - pub fn new_with_query_context( - task: CreateTableTask, - event_context: EventContext, - query_context: QueryContext, - ) -> Self { + pub fn new(task: CreateTableTask, query_context: QueryContext) -> Self { CreateTableData { state: CreateTableState::Prepare, column_metadatas: vec![], task, query_context, - event_context, table_route: None, region_wal_options: None, } diff --git a/src/common/meta/src/ddl/create_view.rs b/src/common/meta/src/ddl/create_view.rs index a7175854b76..666ef19efdf 100644 --- a/src/common/meta/src/ddl/create_view.rs +++ b/src/common/meta/src/ddl/create_view.rs @@ -16,8 +16,8 @@ use async_trait::async_trait; use common_event_recorder::Event; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, - ProcedureState, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureState, + Status, }; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -35,7 +35,7 @@ use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; use crate::lock_key::{CatalogLock, SchemaLock, TableNameLock}; use crate::metrics; -use crate::rpc::ddl::{CreateViewTask, EventContext}; +use crate::rpc::ddl::CreateViewTask; // The procedure to execute `[CreateViewTask]`. pub struct CreateViewProcedure { @@ -46,14 +46,13 @@ pub struct CreateViewProcedure { impl CreateViewProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::CreateView"; - pub fn new(task: CreateViewTask, event_context: EventContext, context: DdlContext) -> Self { + pub fn new(task: CreateViewTask, context: DdlContext) -> Self { Self { context, data: CreateViewData { state: CreateViewState::Prepare, task, need_update: false, - event_context, }, } } @@ -272,7 +271,7 @@ impl Procedure for CreateViewProcedure { ]) } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(CREATE_VIEW_EVENT_TYPE) { return None; } @@ -290,7 +289,6 @@ impl Procedure for CreateViewProcedure { referenced_table_count: self.data.task.table_names().len(), column_count: self.data.task.columns().len(), }, - self.data.event_context.clone(), ) } EventTrigger::Succeeded => match ctx.lifecycle_state { @@ -344,8 +342,6 @@ pub struct CreateViewData { pub task: CreateViewTask, /// Whether to update the view info. pub need_update: bool, - #[serde(default)] - pub event_context: EventContext, } impl CreateViewData { diff --git a/src/common/meta/src/ddl/drop_database.rs b/src/common/meta/src/ddl/drop_database.rs index 7772e966cb2..5bbd229ec5b 100644 --- a/src/common/meta/src/ddl/drop_database.rs +++ b/src/common/meta/src/ddl/drop_database.rs @@ -23,7 +23,7 @@ use std::fmt::Debug; use common_error::ext::BoxedError; use common_procedure::error::{ExternalSnafu, FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use futures::stream::BoxStream; @@ -38,13 +38,11 @@ use crate::ddl::utils::map_to_procedure_error; use crate::error::Result; use crate::key::table_name::TableNameValue; use crate::lock_key::{CatalogLock, SchemaLock}; -use crate::rpc::ddl::EventContext; pub struct DropDatabaseProcedure { /// The context of procedure runtime. runtime_context: DdlContext, context: DropDatabaseContext, - event_context: EventContext, state: Box, } @@ -87,16 +85,9 @@ pub(crate) trait State: Send + Debug { impl DropDatabaseProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropDatabase"; - pub fn new( - catalog: String, - schema: String, - drop_if_exists: bool, - event_context: EventContext, - context: DdlContext, - ) -> Self { + pub fn new(catalog: String, schema: String, drop_if_exists: bool, context: DdlContext) -> Self { Self { runtime_context: context, - event_context, context: DropDatabaseContext { catalog, schema, @@ -113,13 +104,11 @@ impl DropDatabaseProcedure { catalog, schema, drop_if_exists, - event_context, state, } = serde_json::from_str(json).context(FromJsonSnafu)?; Ok(Self { runtime_context, - event_context, context: DropDatabaseContext { catalog, schema, @@ -170,7 +159,6 @@ impl Procedure for DropDatabaseProcedure { catalog: &self.context.catalog, schema: &self.context.schema, drop_if_exists: self.context.drop_if_exists, - event_context: &self.event_context, state: self.state.as_ref(), }; @@ -186,10 +174,7 @@ impl Procedure for DropDatabaseProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(DROP_DATABASE_EVENT_TYPE) { return None; } @@ -199,7 +184,6 @@ impl Procedure for DropDatabaseProcedure { &self.context.catalog, &self.context.schema, self.context.drop_if_exists, - self.event_context.clone(), ) } else { DatabaseDdlEvent::drop_lifecycle(&self.context.catalog, &self.context.schema) @@ -215,7 +199,6 @@ struct DropDatabaseData<'a> { // The schema name schema: &'a str, drop_if_exists: bool, - event_context: &'a EventContext, state: &'a dyn State, } @@ -226,7 +209,5 @@ struct DropDatabaseOwnedData { // The schema name schema: String, drop_if_exists: bool, - #[serde(default)] - event_context: EventContext, state: Box, } diff --git a/src/common/meta/src/ddl/drop_flow.rs b/src/common/meta/src/ddl/drop_flow.rs index 99c18ebc13b..783fb7796e0 100644 --- a/src/common/meta/src/ddl/drop_flow.rs +++ b/src/common/meta/src/ddl/drop_flow.rs @@ -21,7 +21,7 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; @@ -41,7 +41,7 @@ use crate::key::flow::flow_info::FlowInfoValue; use crate::key::flow::flow_route::FlowRouteValue; use crate::lock_key::{CatalogLock, FlowLock}; use crate::metrics; -use crate::rpc::ddl::{DropFlowTask, EventContext}; +use crate::rpc::ddl::DropFlowTask; /// The procedure for dropping a flow. pub struct DropFlowProcedure { @@ -54,7 +54,7 @@ pub struct DropFlowProcedure { impl DropFlowProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropFlow"; - pub fn new(task: DropFlowTask, event_context: EventContext, context: DdlContext) -> Self { + pub fn new(task: DropFlowTask, context: DdlContext) -> Self { Self { context, data: DropFlowData { @@ -62,7 +62,6 @@ impl DropFlowProcedure { task, flow_info_value: None, flow_route_values: vec![], - event_context, }, } } @@ -223,10 +222,7 @@ impl Procedure for DropFlowProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(DROP_FLOW_EVENT_TYPE) { return None; } @@ -237,7 +233,6 @@ impl Procedure for DropFlowProcedure { &self.data.task.flow_name, self.data.task.flow_id, self.data.task.drop_if_exists, - self.data.event_context.clone(), ), _ => FlowDdlEvent::drop_lifecycle( &self.data.task.catalog_name, @@ -257,8 +252,6 @@ pub(crate) struct DropFlowData { task: DropFlowTask, pub(crate) flow_info_value: Option, pub(crate) flow_route_values: Vec, - #[serde(default)] - pub(crate) event_context: EventContext, } /// The state of drop flow diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs index 7c8dec55e49..b73bd95e1bc 100644 --- a/src/common/meta/src/ddl/drop_table.rs +++ b/src/common/meta/src/ddl/drop_table.rs @@ -22,8 +22,8 @@ use common_error::ext::BoxedError; use common_event_recorder::Event; use common_procedure::error::{ExternalSnafu, FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, - LockKey, Procedure, Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; use common_telemetry::tracing::warn; @@ -48,7 +48,7 @@ use crate::key::table_route::TableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; use crate::region_keeper::OperatingRegionGuard; -use crate::rpc::ddl::{DropTableTask, EventContext}; +use crate::rpc::ddl::DropTableTask; use crate::rpc::router::{RegionRoute, operating_leader_region_roles}; #[cfg(feature = "enterprise")] @@ -74,12 +74,11 @@ pub struct DropTableProcedure { impl DropTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropTable"; - pub fn new(task: DropTableTask, context: DdlContext, event_context: EventContext) -> Self { + pub fn new(task: DropTableTask, context: DdlContext) -> Self { let data = DropTableData::new( task, cfg!(feature = "enterprise") && context.soft_drop_enabled, context.soft_drop_retention, - event_context, ); let executor = data.build_executor(); Self { @@ -414,7 +413,7 @@ impl Procedure for DropTableProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::DropTable.as_str()) @@ -425,11 +424,9 @@ impl Procedure for DropTableProcedure { let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) .with_table_id(task.table_id); let event = match &ctx.trigger { - EventTrigger::Submitted => TableDdlEvent::drop_table_submitted( - locator, - task.drop_if_exists, - self.data.event_context.clone(), - ), + EventTrigger::Submitted => { + TableDdlEvent::drop_table_submitted(locator, task.drop_if_exists) + } _ => TableDdlEvent::lifecycle(TableDdlEventType::DropTable, [locator]), }; @@ -486,8 +483,6 @@ pub struct DropTableData { pub soft_drop_retention_millis: Option, #[serde(default)] pub drop_generation: Option, - #[serde(default)] - pub event_context: EventContext, } impl DropTableData { @@ -495,7 +490,6 @@ impl DropTableData { task: DropTableTask, soft_drop_enabled: bool, soft_drop_retention: Option, - event_context: EventContext, ) -> Self { Self { state: DropTableState::Prepare, @@ -509,7 +503,6 @@ impl DropTableData { retention_expires_at: None, soft_drop_retention_millis: soft_drop_retention_millis(soft_drop_retention), drop_generation: None, - event_context, } } diff --git a/src/common/meta/src/ddl/drop_view.rs b/src/common/meta/src/ddl/drop_view.rs index ece1d2b06b7..32d843fbd48 100644 --- a/src/common/meta/src/ddl/drop_view.rs +++ b/src/common/meta/src/ddl/drop_view.rs @@ -16,7 +16,7 @@ use async_trait::async_trait; use common_event_recorder::Event; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; @@ -35,7 +35,7 @@ use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; use crate::metrics; -use crate::rpc::ddl::{DropViewTask, EventContext}; +use crate::rpc::ddl::DropViewTask; /// The procedure for dropping a view. pub struct DropViewProcedure { @@ -48,13 +48,12 @@ pub struct DropViewProcedure { impl DropViewProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropView"; - pub fn new(task: DropViewTask, event_context: EventContext, context: DdlContext) -> Self { + pub fn new(task: DropViewTask, context: DdlContext) -> Self { Self { context, data: DropViewData { state: DropViewState::Prepare, task, - event_context, }, } } @@ -214,7 +213,7 @@ impl Procedure for DropViewProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(DROP_VIEW_EVENT_TYPE) { return None; } @@ -228,7 +227,6 @@ impl Procedure for DropViewProcedure { table_ref.table, self.data.view_id(), self.data.task.drop_if_exists, - self.data.event_context.clone(), ) } _ => { @@ -251,8 +249,6 @@ impl Procedure for DropViewProcedure { pub(crate) struct DropViewData { state: DropViewState, task: DropViewTask, - #[serde(default)] - event_context: EventContext, } impl DropViewData { diff --git a/src/common/meta/src/ddl/event/database.rs b/src/common/meta/src/ddl/event/database.rs index b24611799c6..cab42020f9f 100644 --- a/src/common/meta/src/ddl/event/database.rs +++ b/src/common/meta/src/ddl/event/database.rs @@ -19,14 +19,13 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, - SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, column_schemas, nullable_json, - nullable_string, + CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, + SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, column_schemas, nullable_string, }; use serde::Serialize; use snafu::ResultExt; -use crate::rpc::ddl::{AlterDatabaseKind, EventContext, SetDatabaseOption, UnsetDatabaseOption}; +use crate::rpc::ddl::{AlterDatabaseKind, SetDatabaseOption, UnsetDatabaseOption}; pub(crate) const CREATE_DATABASE_EVENT_TYPE: &str = "create_database"; pub(crate) const ALTER_DATABASE_EVENT_TYPE: &str = "alter_database"; @@ -40,7 +39,6 @@ pub(crate) struct DatabaseDdlEvent { catalog_name: Option, schema_name: Option, payload: Option, - event_context: Option, } #[derive(Debug, Serialize)] @@ -90,7 +88,6 @@ impl DatabaseDdlEvent { schema_name: &str, create_if_not_exists: bool, options: &HashMap, - event_context: EventContext, ) -> Self { let mut options = options.iter().collect::>(); options.sort_unstable_by_key(|(left, _)| *left); @@ -110,7 +107,6 @@ impl DatabaseDdlEvent { create_if_not_exists, options, }), - event_context, ) } @@ -118,7 +114,6 @@ impl DatabaseDdlEvent { catalog_name: &str, schema_name: &str, kind: &AlterDatabaseKind, - event_context: EventContext, ) -> Self { let intent = match kind { AlterDatabaseKind::SetDatabaseOptions(options) => AlterDatabaseIntent::Set { @@ -157,7 +152,6 @@ impl DatabaseDdlEvent { version: PAYLOAD_VERSION, intent, }), - event_context, ) } @@ -165,7 +159,6 @@ impl DatabaseDdlEvent { catalog_name: &str, schema_name: &str, drop_if_exists: bool, - event_context: EventContext, ) -> Self { Self::submitted( DROP_DATABASE_EVENT_TYPE, @@ -175,7 +168,6 @@ impl DatabaseDdlEvent { version: PAYLOAD_VERSION, drop_if_exists, }), - event_context, ) } @@ -196,14 +188,12 @@ impl DatabaseDdlEvent { catalog_name: &str, schema_name: &str, payload: DatabaseDdlPayload, - event_context: EventContext, ) -> Self { Self { event_type, catalog_name: Some(catalog_name.to_string()), schema_name: Some(schema_name.to_string()), payload: Some(payload), - event_context: Some(event_context), } } @@ -213,7 +203,6 @@ impl DatabaseDdlEvent { catalog_name: Some(catalog_name.to_string()), schema_name: Some(schema_name.to_string()), payload: None, - event_context: None, } } } @@ -234,22 +223,14 @@ impl Event for DatabaseDdlEvent { column_schemas([ &EVENT_TABLE_CATALOG_NAME_COLUMN, &EVENT_TABLE_SCHEMA_NAME_COLUMN, - &EVENT_CONTEXT_COLUMN, ]) } fn extra_rows(&self) -> Result> { - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.schema_name.as_deref()), - nullable_json(event_context.as_ref()), ], }]) } diff --git a/src/common/meta/src/ddl/event/flow.rs b/src/common/meta/src/ddl/event/flow.rs index 7d0d6cb539b..0a9ec2dd244 100644 --- a/src/common/meta/src/ddl/event/flow.rs +++ b/src/common/meta/src/ddl/event/flow.rs @@ -19,14 +19,12 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, column_schemas, - nullable_json, nullable_string, nullable_value, + CATALOG_NAME_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, column_schemas, nullable_string, + nullable_value, }; use serde::Serialize; use snafu::ResultExt; -use crate::rpc::ddl::EventContext; - pub(crate) const CREATE_FLOW_EVENT_TYPE: &str = "create_flow"; pub(crate) const DROP_FLOW_EVENT_TYPE: &str = "drop_flow"; @@ -71,15 +69,14 @@ pub(crate) struct FlowDdlEvent { flow_name: Option, flow_id: Option, payload: Option, - event_context: Option, } impl FlowDdlEvent { + /// Builds the bounded event emitted when creating a Flow is submitted. pub(crate) fn create_submitted( catalog_name: &str, flow_name: &str, intent: CreateFlowEventIntent, - event_context: EventContext, ) -> Self { Self { event_type: CREATE_FLOW_EVENT_TYPE, @@ -93,16 +90,15 @@ impl FlowDdlEvent { expire_after: intent.expire_after, eval_interval_secs: intent.eval_interval_secs, })), - event_context: Some(event_context), } } + /// Builds the bounded event emitted when dropping a Flow is submitted. pub(crate) fn drop_submitted( catalog_name: &str, flow_name: &str, flow_id: u32, drop_if_exists: bool, - event_context: EventContext, ) -> Self { Self { event_type: DROP_FLOW_EVENT_TYPE, @@ -113,7 +109,6 @@ impl FlowDdlEvent { version: PAYLOAD_VERSION, drop_if_exists, })), - event_context: Some(event_context), } } @@ -149,7 +144,6 @@ impl FlowDdlEvent { flow_name: Some(flow_name.to_string()), flow_id: None, payload: None, - event_context: None, } } } @@ -167,27 +161,15 @@ impl Event for FlowDdlEvent { } fn extra_schema(&self) -> Vec { - column_schemas([ - &CATALOG_NAME_COLUMN, - &FLOW_NAME_COLUMN, - &FLOW_ID_COLUMN, - &EVENT_CONTEXT_COLUMN, - ]) + column_schemas([&CATALOG_NAME_COLUMN, &FLOW_NAME_COLUMN, &FLOW_ID_COLUMN]) } fn extra_rows(&self) -> Result> { - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.flow_name.as_deref()), nullable_value(self.flow_id.map(ValueData::U32Value)), - nullable_json(event_context.as_ref()), ], }]) } diff --git a/src/common/meta/src/ddl/event/table.rs b/src/common/meta/src/ddl/event/table.rs index 9ae6f530761..670f3c7c1dc 100644 --- a/src/common/meta/src/ddl/event/table.rs +++ b/src/common/meta/src/ddl/event/table.rs @@ -21,17 +21,14 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, - TABLE_ID_COLUMN, TABLE_NAME_COLUMN, column_schemas, nullable_json, nullable_string, - nullable_value, + CATALOG_NAME_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN, + TABLE_NAME_COLUMN, column_schemas, nullable_string, nullable_value, }; use serde::Serialize; use serde_json::Value as JsonValue; use snafu::ResultExt; use store_api::storage::TableId; -use crate::rpc::ddl::EventContext; - /// Current version of table DDL event payloads. pub(crate) const TABLE_DDL_PAYLOAD_VERSION: u8 = 1; @@ -215,7 +212,6 @@ pub(crate) struct TableDdlEvent { event_type: TableDdlEventType, locators: Vec, payload: Option, - event_context: Option, } impl TableDdlEvent { @@ -224,7 +220,6 @@ impl TableDdlEvent { locator: TableDdlLocator, create_if_not_exists: bool, engine: &str, - event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::CreateTable, @@ -234,7 +229,6 @@ impl TableDdlEvent { create_if_not_exists, engine: engine.to_string(), }), - Some(event_context), ) } @@ -242,7 +236,6 @@ impl TableDdlEvent { pub(crate) fn create_logical_tables_submitted( locators: impl IntoIterator, table_count: usize, - event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::CreateLogicalTables, @@ -251,7 +244,6 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, table_count, }), - Some(event_context), ) } @@ -259,7 +251,6 @@ impl TableDdlEvent { pub(crate) fn alter_table_submitted( locator: TableDdlLocator, kind: Option<&'static str>, - event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::AlterTable, @@ -268,7 +259,6 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, kind, }), - Some(event_context), ) } @@ -277,7 +267,6 @@ impl TableDdlEvent { locators: impl IntoIterator, table_count: usize, kinds: impl IntoIterator, - event_context: EventContext, ) -> Self { let kinds = kinds .into_iter() @@ -292,16 +281,11 @@ impl TableDdlEvent { table_count, kinds, }), - Some(event_context), ) } /// Builds the bounded event emitted when dropping a table is submitted. - pub(crate) fn drop_table_submitted( - locator: TableDdlLocator, - drop_if_exists: bool, - event_context: EventContext, - ) -> Self { + pub(crate) fn drop_table_submitted(locator: TableDdlLocator, drop_if_exists: bool) -> Self { Self::submitted( TableDdlEventType::DropTable, [locator], @@ -309,39 +293,30 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, drop_if_exists, }), - Some(event_context), ) } /// Builds the bounded event emitted when restoring a dropped table is submitted. #[cfg(feature = "enterprise")] - pub(crate) fn undrop_table_submitted( - locator: TableDdlLocator, - event_context: EventContext, - ) -> Self { + pub(crate) fn undrop_table_submitted(locator: TableDdlLocator) -> Self { Self::submitted( TableDdlEventType::UndropTable, [locator], TableDdlPayload::UndropTable(UndropTablePayload { version: TABLE_DDL_PAYLOAD_VERSION, }), - Some(event_context), ) } /// Builds the bounded event emitted when purging a dropped table is submitted. #[cfg(feature = "enterprise")] - pub(crate) fn purge_dropped_table_submitted( - locator: TableDdlLocator, - event_context: EventContext, - ) -> Self { + pub(crate) fn purge_dropped_table_submitted(locator: TableDdlLocator) -> Self { Self::submitted( TableDdlEventType::PurgeDroppedTable, [locator], TableDdlPayload::PurgeDroppedTable(PurgeDroppedTablePayload { version: TABLE_DDL_PAYLOAD_VERSION, }), - Some(event_context), ) } @@ -349,7 +324,6 @@ impl TableDdlEvent { pub(crate) fn truncate_table_submitted( locator: TableDdlLocator, time_range_count: usize, - event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::TruncateTable, @@ -358,7 +332,6 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, time_range_count, }), - Some(event_context), ) } @@ -371,7 +344,6 @@ impl TableDdlEvent { event_type, locators: locators.into_iter().collect(), payload: None, - event_context: None, } } @@ -394,13 +366,11 @@ impl TableDdlEvent { event_type: TableDdlEventType, locators: impl IntoIterator, payload: TableDdlPayload, - event_context: Option, ) -> Self { Self { event_type, locators: locators.into_iter().collect(), payload: Some(payload), - event_context, } } @@ -413,7 +383,7 @@ impl TableDdlEvent { ]) } - fn locator_row(&self, locator: &TableDdlLocator) -> Result { + fn locator_row(&self, locator: &TableDdlLocator) -> Row { let mut values = vec![ nullable_string(locator.catalog_name.as_deref()), nullable_string(locator.schema_name.as_deref()), @@ -423,14 +393,7 @@ impl TableDdlEvent { if self.event_type.has_physical_table_id() { values.push(nullable_table_id(locator.physical_table_id)); } - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; - values.push(nullable_json(event_context.as_ref())); - Ok(Row { values }) + Row { values } } } @@ -451,15 +414,15 @@ impl Event for TableDdlEvent { if self.event_type.has_physical_table_id() { schema.push(PHYSICAL_TABLE_ID_COLUMN.column_schema()); } - schema.push(EVENT_CONTEXT_COLUMN.column_schema()); schema } fn extra_rows(&self) -> Result> { - self.locators + Ok(self + .locators .iter() .map(|locator| self.locator_row(locator)) - .collect() + .collect()) } fn as_any(&self) -> &dyn Any { diff --git a/src/common/meta/src/ddl/event/view.rs b/src/common/meta/src/ddl/event/view.rs index c1d5846fcf2..cdce5954086 100644 --- a/src/common/meta/src/ddl/event/view.rs +++ b/src/common/meta/src/ddl/event/view.rs @@ -19,14 +19,12 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, - VIEW_NAME_COLUMN, column_schemas, nullable_json, nullable_string, nullable_value, + CATALOG_NAME_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, VIEW_NAME_COLUMN, column_schemas, + nullable_string, nullable_value, }; use serde::Serialize; use snafu::ResultExt; -use crate::rpc::ddl::EventContext; - pub(crate) const CREATE_VIEW_EVENT_TYPE: &str = "create_view"; pub(crate) const DROP_VIEW_EVENT_TYPE: &str = "drop_view"; @@ -64,7 +62,6 @@ pub(crate) struct ViewDdlEvent { view_name: Option, view_id: Option, payload: Option, - event_context: Option, } #[derive(Debug, Serialize)] @@ -75,12 +72,12 @@ enum ViewDdlPayload { } impl ViewDdlEvent { + /// Builds the bounded event emitted when creating a View is submitted. pub(crate) fn create_submitted( catalog_name: &str, schema_name: &str, view_name: &str, intent: CreateViewEventIntent, - event_context: EventContext, ) -> Self { Self::submitted( CREATE_VIEW_EVENT_TYPE, @@ -95,17 +92,16 @@ impl ViewDdlEvent { referenced_table_count: intent.referenced_table_count, column_count: intent.column_count, }), - event_context, ) } + /// Builds the bounded event emitted when dropping a View is submitted. pub(crate) fn drop_submitted( catalog_name: &str, schema_name: &str, view_name: &str, view_id: u32, drop_if_exists: bool, - event_context: EventContext, ) -> Self { Self::submitted( DROP_VIEW_EVENT_TYPE, @@ -117,7 +113,6 @@ impl ViewDdlEvent { version: PAYLOAD_VERSION, drop_if_exists, }), - event_context, ) } @@ -162,7 +157,6 @@ impl ViewDdlEvent { view_name: &str, view_id: Option, payload: ViewDdlPayload, - event_context: EventContext, ) -> Self { Self { event_type, @@ -171,7 +165,6 @@ impl ViewDdlEvent { view_name: Some(view_name.to_string()), view_id, payload: Some(payload), - event_context: Some(event_context), } } @@ -188,7 +181,6 @@ impl ViewDdlEvent { view_name: Some(view_name.to_string()), view_id: None, payload: None, - event_context: None, } } @@ -206,7 +198,6 @@ impl ViewDdlEvent { view_name: Some(view_name.to_string()), view_id: Some(view_id), payload: None, - event_context: None, } } } @@ -229,24 +220,16 @@ impl Event for ViewDdlEvent { &SCHEMA_NAME_COLUMN, &VIEW_NAME_COLUMN, &VIEW_ID_COLUMN, - &EVENT_CONTEXT_COLUMN, ]) } fn extra_rows(&self) -> Result> { - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.schema_name.as_deref()), nullable_string(self.view_name.as_deref()), nullable_value(self.view_id.map(ValueData::U32Value)), - nullable_json(event_context.as_ref()), ], }]) } diff --git a/src/common/meta/src/ddl/purge_dropped_table.rs b/src/common/meta/src/ddl/purge_dropped_table.rs index c81aa043dc2..cdafe40de81 100644 --- a/src/common/meta/src/ddl/purge_dropped_table.rs +++ b/src/common/meta/src/ddl/purge_dropped_table.rs @@ -14,7 +14,7 @@ use std::collections::HashMap; use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; @@ -38,7 +38,7 @@ use crate::error::{self, Result}; use crate::key::DroppedTableMetadata; use crate::key::table_route::TableRouteValue; use crate::lock_key::TableLock; -use crate::rpc::ddl::{EventContext, PurgeDroppedTableTask}; +use crate::rpc::ddl::PurgeDroppedTableTask; use crate::rpc::router::RegionRoute; pub struct PurgeDroppedTableProcedure { @@ -50,25 +50,17 @@ impl PurgeDroppedTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::PurgeDroppedTable"; pub const EXPIRED_TYPE_NAME: &'static str = "metasrv-procedure::PurgeExpiredDroppedTable"; - pub fn new( - task: PurgeDroppedTableTask, - context: DdlContext, - event_context: EventContext, - ) -> Self { + pub fn new(task: PurgeDroppedTableTask, context: DdlContext) -> Self { Self { context, - data: PurgeDroppedTableData::new(task, event_context), + data: PurgeDroppedTableData::new(task), } } - pub fn new_if_expired( - task: PurgeDroppedTableTask, - context: DdlContext, - event_context: EventContext, - ) -> Self { + pub fn new_if_expired(task: PurgeDroppedTableTask, context: DdlContext) -> Self { Self { context, - data: PurgeDroppedTableData::new_if_expired(task, event_context), + data: PurgeDroppedTableData::new_if_expired(task), } } @@ -278,10 +270,7 @@ impl Procedure for PurgeDroppedTableProcedure { LockKey::new(vec![TableLock::Write(self.data.task.table_id).into()]) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::PurgeDroppedTable.as_str()) @@ -290,7 +279,7 @@ impl Procedure for PurgeDroppedTableProcedure { } let locator = self.event_locator(); let event = if ctx.trigger == EventTrigger::Submitted { - TableDdlEvent::purge_dropped_table_submitted(locator, self.data.event_context.clone()) + TableDdlEvent::purge_dropped_table_submitted(locator) } else { TableDdlEvent::lifecycle(TableDdlEventType::PurgeDroppedTable, [locator]) }; @@ -316,12 +305,10 @@ pub struct PurgeDroppedTableData { drop_generation_loaded: bool, #[serde(default)] purging_claimed: bool, - #[serde(default)] - event_context: EventContext, } impl PurgeDroppedTableData { - fn new(task: PurgeDroppedTableTask, event_context: EventContext) -> Self { + fn new(task: PurgeDroppedTableTask) -> Self { Self { state: PurgeDroppedTableState::Prepare, task, @@ -334,14 +321,13 @@ impl PurgeDroppedTableData { drop_generation: None, drop_generation_loaded: false, purging_claimed: false, - event_context, } } - fn new_if_expired(task: PurgeDroppedTableTask, event_context: EventContext) -> Self { + fn new_if_expired(task: PurgeDroppedTableTask) -> Self { Self { check_expired: true, - ..Self::new(task, event_context) + ..Self::new(task) } } diff --git a/src/common/meta/src/ddl/test_util.rs b/src/common/meta/src/ddl/test_util.rs index 3686ceb33d2..7d8cdd895b2 100644 --- a/src/common/meta/src/ddl/test_util.rs +++ b/src/common/meta/src/ddl/test_util.rs @@ -45,7 +45,7 @@ use crate::key::node_address::{NodeAddressKey, NodeAddressValue}; use crate::key::table_route::TableRouteValue; use crate::key::{MetadataKey, MetadataValue}; use crate::peer::Peer; -use crate::rpc::ddl::{CreateTableTask, EventContext}; +use crate::rpc::ddl::CreateTableTask; use crate::rpc::store::PutRequest; pub async fn create_physical_table_metadata( @@ -104,12 +104,7 @@ pub async fn create_logical_table( table_name: &str, ) -> TableId { let tasks = vec![test_create_logical_table_task(table_name)]; - let mut procedure = CreateLogicalTablesProcedure::new( - tasks, - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = CreateLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, diff --git a/src/common/meta/src/ddl/tests/alter_logical_tables.rs b/src/common/meta/src/ddl/tests/alter_logical_tables.rs index 09662199389..351c0fd96ec 100644 --- a/src/common/meta/src/ddl/tests/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/tests/alter_logical_tables.rs @@ -43,7 +43,7 @@ use crate::error::Error::{AlterLogicalTablesInvalidArguments, TableNotFound}; use crate::error::Result; use crate::key::table_name::TableNameKey; use crate::key::table_route::{PhysicalTableRouteValue, TableRouteValue}; -use crate::rpc::ddl::{AlterTableTask, EventContext}; +use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -162,12 +162,7 @@ async fn test_on_prepare_check_schema() { ), ]; let physical_table_id = 1024u32; - let mut procedure = AlterLogicalTablesProcedure::new( - tasks, - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -182,12 +177,7 @@ async fn test_on_prepare_check_alter_kind() { "new_table1", )]; let physical_table_id = 1024u32; - let mut procedure = AlterLogicalTablesProcedure::new( - tasks, - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -207,8 +197,7 @@ async fn test_on_prepare_different_physical_table() { make_alter_logical_table_add_column_task(None, "table2", vec!["column2".to_string()]), ]; - let mut procedure = - AlterLogicalTablesProcedure::new(tasks, phy1_id, EventContext::default(), ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy1_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -229,8 +218,7 @@ async fn test_on_prepare_logical_table_not_exists() { make_alter_logical_table_add_column_task(None, "table2", vec!["column2".to_string()]), ]; - let mut procedure = - AlterLogicalTablesProcedure::new(tasks, phy_id, EventContext::default(), ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, TableNotFound { .. }); } @@ -253,8 +241,7 @@ async fn test_on_prepare() { make_alter_logical_table_add_column_task(None, "table3", vec!["column3".to_string()]), ]; - let mut procedure = - AlterLogicalTablesProcedure::new(tasks, phy_id, EventContext::default(), ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); let result = procedure.on_prepare().await; assert_matches!( result, @@ -290,12 +277,7 @@ async fn test_on_update_metadata() { make_alter_logical_table_add_column_task(None, "table3", vec!["new_col".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new( - tasks, - phy_id, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -308,6 +290,7 @@ async fn test_on_update_metadata() { let ctx = common_procedure::Context { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // on_submit_alter_region_requests status = procedure.execute(&ctx).await.unwrap(); @@ -379,12 +362,7 @@ async fn test_on_part_duplicate_alter_request() { make_alter_logical_table_add_column_task(None, "table2", vec!["col_0".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new( - tasks, - phy_id, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -397,6 +375,7 @@ async fn test_on_part_duplicate_alter_request() { let ctx = common_procedure::Context { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // on_submit_alter_region_requests status = procedure.execute(&ctx).await.unwrap(); @@ -469,12 +448,7 @@ async fn test_on_part_duplicate_alter_request() { ), ]; - let mut procedure = AlterLogicalTablesProcedure::new( - tasks, - phy_id, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -487,6 +461,7 @@ async fn test_on_part_duplicate_alter_request() { let ctx = common_procedure::Context { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // on_submit_alter_region_requests status = procedure.execute(&ctx).await.unwrap(); @@ -645,8 +620,7 @@ async fn test_on_submit_alter_region_request() { make_alter_logical_table_add_column_task(None, "table2", vec!["mew_col".to_string()]), ]; - let mut procedure = - AlterLogicalTablesProcedure::new(tasks, phy_id, EventContext::default(), ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); procedure.on_prepare().await.unwrap(); procedure.on_submit_alter_region_requests().await.unwrap(); let mut results = Vec::new(); diff --git a/src/common/meta/src/ddl/tests/alter_table.rs b/src/common/meta/src/ddl/tests/alter_table.rs index d68d1c96dfa..b35666e52e1 100644 --- a/src/common/meta/src/ddl/tests/alter_table.rs +++ b/src/common/meta/src/ddl/tests/alter_table.rs @@ -59,7 +59,7 @@ use crate::lock_key::RegionLock; use crate::node_manager::NodeManagerRef; use crate::peer::Peer; use crate::poison_key::table_poison_key; -use crate::rpc::ddl::{AlterTableTask, EventContext}; +use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -121,8 +121,7 @@ async fn test_on_prepare_table_exists_err() { .unwrap(); let task = test_rename_alter_table_task("non-exists", "foo"); - let mut procedure = - AlterTableProcedure::new(1024, task, EventContext::default(), ddl_context).unwrap(); + let mut procedure = AlterTableProcedure::new(1024, task, ddl_context).unwrap(); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err.status_code(), StatusCode::TableAlreadyExists); } @@ -132,8 +131,7 @@ async fn test_on_prepare_table_not_exists_err() { let node_manager = Arc::new(MockDatanodeManager::new(())); let ddl_context = new_ddl_context(node_manager); let task = test_rename_alter_table_task("non-exists", "foo"); - let mut procedure = - AlterTableProcedure::new(1024, task, EventContext::default(), ddl_context).unwrap(); + let mut procedure = AlterTableProcedure::new(1024, task, ddl_context).unwrap(); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err.status_code(), StatusCode::TableNotFound); } @@ -247,13 +245,8 @@ async fn test_on_submit_alter_request() { let alter_table_task = test_alter_table_task(table_name); let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = AlterTableProcedure::new( - table_id, - alter_table_task, - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -306,13 +299,8 @@ async fn test_on_submit_alter_request_without_sync_request() { let alter_table_task = test_alter_table_task(table_name); let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = AlterTableProcedure::new( - table_id, - alter_table_task, - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -369,13 +357,7 @@ async fn test_on_submit_alter_request_with_outdated_request() { }; let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = AlterTableProcedure::new( - table_id, - alter_table_task, - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, alter_table_task, ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let err = procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -404,9 +386,7 @@ async fn test_on_update_metadata_rename() { .unwrap(); let task = test_rename_alter_table_task(table_name, new_table_name); - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) - .unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); procedure.on_update_metadata().await.unwrap(); @@ -486,9 +466,7 @@ async fn test_on_update_metadata_add_columns() { }; let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) - .unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -585,9 +563,7 @@ async fn test_on_update_table_options() { }; let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) - .unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -658,7 +634,6 @@ async fn test_skip_wal_rejects_mixed_table_options() { })), }, }, - EventContext::default(), ddl_context, ) .unwrap(); @@ -687,13 +662,8 @@ async fn test_skip_wal_rejects_logical_table() { }, }; - let mut procedure = AlterTableProcedure::new( - logical_table_id, - task, - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = + AlterTableProcedure::new(logical_table_id, task, ddl_context.clone()).unwrap(); let error = procedure.on_prepare().await.unwrap_err(); assert_matches!(error, Error::Unsupported { .. }); @@ -739,9 +709,7 @@ async fn test_skip_wal_rejects_file_engine_table() { }, }; - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) - .unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); let error = procedure.on_prepare().await.unwrap_err(); assert_matches!(error, Error::Unsupported { .. }); @@ -778,7 +746,6 @@ fn test_skip_wal_holds_region_locks() { let procedure = AlterTableProcedure::new_with_region_locks( table_id, task, - EventContext::default(), region_ids.clone(), context.clone(), ) @@ -830,7 +797,6 @@ async fn test_skip_wal_detects_region_route_change() { let mut procedure = AlterTableProcedure::new_with_region_locks( table_id, alter_task, - EventContext::default(), stale_region_locks, ddl_context.clone(), ) @@ -890,7 +856,6 @@ async fn test_skip_wal_rejects_no_leader_before_updating_metadata() { let mut procedure = AlterTableProcedure::new_with_region_locks( table_id, alter_task, - EventContext::default(), region_locks, ddl_context.clone(), ) @@ -951,7 +916,6 @@ async fn test_skip_wal_updates_metadata_before_all_replicas() { let mut procedure = AlterTableProcedure::new_with_region_locks( table_id, alter_task, - EventContext::default(), region_locks, ddl_context.clone(), ) @@ -1033,13 +997,8 @@ async fn prepare_alter_table_procedure( }, }; let procedure_id = ProcedureId::random(); - let mut procedure = AlterTableProcedure::new( - table_id, - alter_table_task, - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); (procedure, procedure_id) } diff --git a/src/common/meta/src/ddl/tests/create_flow.rs b/src/common/meta/src/ddl/tests/create_flow.rs index 280a2633e32..b526856bdac 100644 --- a/src/common/meta/src/ddl/tests/create_flow.rs +++ b/src/common/meta/src/ddl/tests/create_flow.rs @@ -37,7 +37,7 @@ use crate::key::flow::flow_info::{ FlowInfoValue, FlowMissedTickPolicy, FlowScheduleConfig, FlowStatus, }; use crate::key::table_route::TableRouteValue; -use crate::rpc::ddl::{CreateFlowTask, EventContext, FlowQueryContext, QueryContext}; +use crate::rpc::ddl::{CreateFlowTask, FlowQueryContext, QueryContext}; use crate::test_util::{MockFlownodeManager, new_ddl_context}; pub(crate) fn test_query_context() -> QueryContext { @@ -120,8 +120,7 @@ async fn test_create_flow_source_table_not_found() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = - CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -144,12 +143,7 @@ async fn test_create_pending_flow_source_table_not_found_with_defer() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context.clone()); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Executing { persist: true, .. }); assert_eq!(procedure.data.unresolved_source_table_names.len(), 1); @@ -189,8 +183,7 @@ async fn test_create_pending_flow_source_table_not_found_with_defer_false() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = - CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -240,12 +233,7 @@ async fn test_create_pending_flow_records_partial_source_resolution() { ); enable_defer_on_missing_source(&mut task); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context.clone()); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Executing { persist: true, .. }); assert_eq!(procedure.data.source_table_ids, vec![existing_table_id]); @@ -440,7 +428,6 @@ fn test_resolved_schedule_defaults_in_create_request() { // Construct CreateFlowData and verify CreateRequest carries internal key. let data = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::CreateFlows, task, flow_id: Some(1024), @@ -454,7 +441,6 @@ fn test_resolved_schedule_defaults_in_create_request() { }; let data2 = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::CreateMetadata, task: data.task.clone(), flow_id: Some(1024), @@ -594,8 +580,7 @@ async fn test_create_flow_rejects_unknown_option_in_meta_task() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = - CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unexpected { .. }); @@ -616,7 +601,6 @@ fn test_create_request_strips_defer_on_missing_source_runtime_option() { enable_defer_on_missing_source(&mut task); let data = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::CreateFlows, task, flow_id: Some(1024), @@ -666,12 +650,7 @@ pub(crate) async fn create_test_flow( false, ); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task.clone(), - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(task.clone(), query_ctx, ddl_context.clone()); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let flow_id = output.downcast_ref::().unwrap(); @@ -692,12 +671,7 @@ pub(crate) async fn create_test_pending_flow( ); enable_defer_on_missing_source(&mut task); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context.clone()); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let flow_id = output.downcast_ref::().unwrap(); @@ -744,12 +718,7 @@ async fn test_create_flow() { true, ); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task.clone(), - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(task.clone(), query_ctx, ddl_context.clone()); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let flow_id = output.downcast_ref::().unwrap(); assert_eq!(*flow_id, 1024); @@ -757,12 +726,7 @@ async fn test_create_flow() { // Creates again let task = test_create_flow_task("my_flow", source_table_names, sink_table_name, false); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task.clone(), - query_ctx, - EventContext::default(), - ddl_context, - ); + let mut procedure = CreateFlowProcedure::new(task.clone(), query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowAlreadyExists { .. }); } @@ -819,12 +783,7 @@ async fn test_replace_pending_flow_with_active_flow_is_unsupported() { ); replace_task.or_replace = true; let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - replace_task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(replace_task, query_ctx, ddl_context.clone()); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -882,12 +841,7 @@ async fn test_replace_active_flow_with_pending_flow_is_unsupported() { enable_defer_on_missing_source(&mut replace_task); replace_task.or_replace = true; let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - replace_task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(replace_task, query_ctx, ddl_context.clone()); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -947,12 +901,7 @@ async fn test_replace_pending_flow_with_pending_flow_updates_metadata() { enable_defer_on_missing_source(&mut replace_task); replace_task.or_replace = true; let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - replace_task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(replace_task, query_ctx, ddl_context.clone()); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let replaced_flow_id = *output.downcast_ref::().unwrap(); assert_eq!(replaced_flow_id, original_flow_id); @@ -1004,8 +953,7 @@ async fn test_create_flow_same_source_and_sink_table() { // Try to create a flow with same source and sink table - should fail let task = test_create_flow_task("my_flow", source_table_names, sink_table_name, false); let query_ctx = test_query_context(); - let mut procedure = - CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); @@ -1038,7 +986,7 @@ fn test_create_flow_lock_key_does_not_lock_sink_table_name() { let task = create_test_flow_task_for_serialization(); let query_ctx = test_query_context(); let ddl_context = new_ddl_context(Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler))); - let procedure = CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let lock_keys = procedure.lock_key().get_keys(); @@ -1113,7 +1061,6 @@ fn test_create_flow_data_new_format_serialization() { }; let data = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::Prepare, task: create_test_flow_task_for_serialization(), flow_id: None, @@ -1182,7 +1129,7 @@ fn test_create_flow_procedure_strips_scheduled_time_extension() { .extensions .insert("flow.other".to_string(), "kept".to_string()); - let procedure = CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); assert!( !procedure @@ -1215,7 +1162,6 @@ fn test_flow_info_conversion_with_flow_context() { }; let data = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::CreateMetadata, task: create_test_flow_task_for_serialization(), flow_id: Some(123), @@ -1258,7 +1204,6 @@ fn test_flow_info_conversion_strips_scheduled_time_extension() { }; let data = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::CreateMetadata, task: create_test_flow_task_for_serialization(), flow_id: Some(123), diff --git a/src/common/meta/src/ddl/tests/create_logical_tables.rs b/src/common/meta/src/ddl/tests/create_logical_tables.rs index d1cc2420e04..89ec486ee03 100644 --- a/src/common/meta/src/ddl/tests/create_logical_tables.rs +++ b/src/common/meta/src/ddl/tests/create_logical_tables.rs @@ -39,7 +39,6 @@ use crate::ddl::test_util::{ }; use crate::error::{Error, Result}; use crate::key::table_route::{PhysicalTableRouteValue, TableRouteValue}; -use crate::rpc::ddl::EventContext; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -97,12 +96,7 @@ async fn test_on_prepare_physical_table_not_found() { let ddl_context = new_ddl_context(node_manager); let tasks = vec![test_create_logical_table_task("foo")]; let physical_table_id = 1024u32; - let mut procedure = CreateLogicalTablesProcedure::new( - tasks, - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = CreateLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableRouteNotFound { .. }); } @@ -132,12 +126,7 @@ async fn test_on_prepare() { // The create logical table procedure. let tasks = vec![test_create_logical_table_task("foo")]; let physical_table_id = table_id; - let mut procedure = CreateLogicalTablesProcedure::new( - tasks, - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = CreateLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -183,12 +172,8 @@ async fn test_on_prepare_logical_table_exists_err() { .unwrap(); // The create logical table procedure. let physical_table_id = table_id; - let mut procedure = CreateLogicalTablesProcedure::new( - vec![task], - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = + CreateLogicalTablesProcedure::new(vec![task], physical_table_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -231,12 +216,8 @@ async fn test_on_prepare_with_create_if_table_exists() { let physical_table_id = table_id; // Sets `create_if_not_exists` task.create_table.create_if_not_exists = true; - let mut procedure = CreateLogicalTablesProcedure::new( - vec![task], - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = + CreateLogicalTablesProcedure::new(vec![task], physical_table_id, ddl_context); let status = procedure.on_prepare().await.unwrap(); let output = status.downcast_output_ref::>().unwrap(); assert_eq!(*output, vec![8192]); @@ -283,7 +264,6 @@ async fn test_on_prepare_part_logical_tables_exist() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, non_exist_task], physical_table_id, - EventContext::default(), ddl_context, ); let status = procedure.on_prepare().await.unwrap(); @@ -331,7 +311,6 @@ async fn test_on_create_metadata() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, yet_another_task], physical_table_id, - EventContext::default(), ddl_context.clone(), ); let status = procedure.on_prepare().await.unwrap(); @@ -345,6 +324,7 @@ async fn test_on_create_metadata() { let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); // Triggers procedure to create table metadata @@ -425,7 +405,6 @@ async fn test_on_create_metadata_part_logical_tables_exist() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, non_exist_task], physical_table_id, - EventContext::default(), ddl_context.clone(), ); let status = procedure.on_prepare().await.unwrap(); @@ -439,6 +418,7 @@ async fn test_on_create_metadata_part_logical_tables_exist() { let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); // Triggers procedure to create table metadata @@ -498,7 +478,6 @@ async fn test_on_create_metadata_err() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task.clone(), yet_another_task], physical_table_id, - EventContext::default(), ddl_context.clone(), ); let status = procedure.on_prepare().await.unwrap(); @@ -512,6 +491,7 @@ async fn test_on_create_metadata_err() { let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); // Creates logical table metadata(different with the task) @@ -559,7 +539,6 @@ async fn test_on_submit_create_request() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, yet_another_task], physical_table_id, - EventContext::default(), ddl_context, ); procedure.on_prepare().await.unwrap(); diff --git a/src/common/meta/src/ddl/tests/create_table.rs b/src/common/meta/src/ddl/tests/create_table.rs index 2eee317dc14..c97e61926c0 100644 --- a/src/common/meta/src/ddl/tests/create_table.rs +++ b/src/common/meta/src/ddl/tests/create_table.rs @@ -48,7 +48,7 @@ use crate::ddl::test_util::{assert_column_name, get_raw_table_info, put_datanode use crate::error::{Error, Result}; use crate::key::table_route::TableRouteValue; use crate::kv_backend::memory::MemoryKvBackend; -use crate::rpc::ddl::{CreateTableTask, EventContext, QueryContext}; +use crate::rpc::ddl::CreateTableTask; use crate::test_util::{MockDatanodeManager, new_ddl_context, new_ddl_context_with_kv_backend}; fn create_request_handler(_peer: Peer, request: RegionRequest) -> Result { @@ -189,13 +189,7 @@ async fn test_on_prepare_table_exists_err() { ) .await .unwrap(); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -218,13 +212,7 @@ async fn test_on_prepare_with_create_if_table_exists() { ) .await .unwrap(); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Done { output: Some(..) }); let table_id = *status.downcast_output_ref::().unwrap(); @@ -237,13 +225,7 @@ async fn test_on_prepare_without_create_if_table_exists() { let ddl_context = new_ddl_context(node_manager); let mut task = test_create_table_task("foo"); task.create_table.create_if_not_exists = true; - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -262,17 +244,12 @@ async fn test_on_datanode_create_regions_should_retry() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; let error = procedure.execute(&ctx).await.unwrap_err(); assert!(error.is_retry_later()); @@ -285,17 +262,12 @@ async fn test_on_datanode_create_regions_should_not_retry() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; let error = procedure.execute(&ctx).await.unwrap_err(); assert!(!error.is_retry_later()); @@ -308,13 +280,7 @@ async fn test_on_datanode_create_regions_remaps_addresses_when_retrying() { let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); let table_route = procedure.data.table_route.as_mut().unwrap(); @@ -335,17 +301,12 @@ async fn test_on_create_metadata_error() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new( - task.clone(), - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task.clone(), ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); let mut task = task; @@ -374,17 +335,12 @@ async fn test_on_create_metadata() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); // Triggers procedure to create table metadata @@ -408,13 +364,7 @@ async fn test_memory_region_keeper_guard_dropped_on_procedure_done() { let ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); let task = test_create_table_task("foo"); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); execute_procedure_until(&mut procedure, |p| { p.data.state == CreateTableState::CreateMetadata diff --git a/src/common/meta/src/ddl/tests/create_view.rs b/src/common/meta/src/ddl/tests/create_view.rs index 618249fd1cd..b73977c79f2 100644 --- a/src/common/meta/src/ddl/tests/create_view.rs +++ b/src/common/meta/src/ddl/tests/create_view.rs @@ -29,7 +29,7 @@ use crate::ddl::create_view::CreateViewProcedure; use crate::ddl::test_util::datanode_handler::NaiveDatanodeHandler; use crate::ddl::tests::create_table::test_create_table_task; use crate::error::Error; -use crate::rpc::ddl::{CreateViewTask, EventContext, QueryContext}; +use crate::rpc::ddl::CreateViewTask; use crate::test_util::{MockDatanodeManager, new_ddl_context}; pub(crate) fn test_table_names() -> HashSet { @@ -112,7 +112,7 @@ async fn test_on_prepare_view_exists_err() { ) .await .unwrap(); - let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = CreateViewProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::ViewAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -138,7 +138,7 @@ async fn test_on_prepare_with_create_if_view_exists() { ) .await .unwrap(); - let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = CreateViewProcedure::new(task, ddl_context); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Done { output: Some(..) }); let table_id = *status.downcast_output_ref::().unwrap(); @@ -151,7 +151,7 @@ async fn test_on_prepare_without_create_if_table_exists() { let ddl_context = new_ddl_context(node_manager); let mut task = test_create_view_task("foo"); task.create_view.create_if_not_exists = true; - let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = CreateViewProcedure::new(task, ddl_context); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -170,11 +170,12 @@ async fn test_on_create_metadata() { let ddl_context = new_ddl_context(node_manager); let task = test_create_view_task("foo"); assert!(!task.create_view.create_if_not_exists); - let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = CreateViewProcedure::new(task, ddl_context); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // Triggers procedure to create view metadata let status = procedure.execute(&ctx).await.unwrap(); @@ -189,12 +190,12 @@ async fn test_replace_view_metadata() { let ddl_context = new_ddl_context(node_manager.clone()); let task = test_create_view_task("foo"); assert!(!task.create_view.create_if_not_exists); - let mut procedure = - CreateViewProcedure::new(task.clone(), EventContext::default(), ddl_context.clone()); + let mut procedure = CreateViewProcedure::new(task.clone(), ddl_context.clone()); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // Triggers procedure to create view metadata let status = procedure.execute(&ctx).await.unwrap(); @@ -215,8 +216,7 @@ async fn test_replace_view_metadata() { let mut task = test_create_view_task("foo"); // The view already exists, prepare should fail { - let mut procedure = - CreateViewProcedure::new(task.clone(), EventContext::default(), ddl_context.clone()); + let mut procedure = CreateViewProcedure::new(task.clone(), ddl_context.clone()); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::ViewAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -227,12 +227,12 @@ async fn test_replace_view_metadata() { task.create_view.logical_plan = vec![4, 5, 6]; task.create_view.definition = "new_definition".to_string(); - let mut procedure = - CreateViewProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = CreateViewProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // Triggers procedure to replace view metadata, but the view_id is unchanged. let status = procedure.execute(&ctx).await.unwrap(); @@ -262,17 +262,12 @@ async fn test_replace_table() { { // Create a `foo` table. let task = test_create_table_task("foo"); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -281,8 +276,7 @@ async fn test_replace_table() { // Try to replace a view named `foo` too. let mut task = test_create_view_task("foo"); task.create_view.or_replace = true; - let mut procedure = - CreateViewProcedure::new(task.clone(), EventContext::default(), ddl_context.clone()); + let mut procedure = CreateViewProcedure::new(task.clone(), ddl_context.clone()); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); diff --git a/src/common/meta/src/ddl/tests/drop_database.rs b/src/common/meta/src/ddl/tests/drop_database.rs index 955625e4bb3..05a70962bf0 100644 --- a/src/common/meta/src/ddl/tests/drop_database.rs +++ b/src/common/meta/src/ddl/tests/drop_database.rs @@ -26,7 +26,6 @@ use crate::ddl::drop_database::executor::DropDatabaseExecutor; use crate::ddl::test_util::datanode_handler::{NaiveDatanodeHandler, RetryErrorDatanodeHandler}; use crate::ddl::test_util::{create_logical_table, create_physical_table}; use crate::key::schema_name::SchemaNameKey; -use crate::rpc::ddl::EventContext; use crate::test_util::{MockDatanodeManager, new_ddl_context}; #[tokio::test] @@ -55,13 +54,13 @@ async fn test_drop_database_with_logical_tables() { DEFAULT_CATALOG_NAME.to_string(), DEFAULT_SCHEMA_NAME.to_string(), false, - EventContext::default(), ddl_context.clone(), ); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; while !procedure.execute(&ctx).await.unwrap().is_done() { @@ -104,13 +103,13 @@ async fn test_drop_database_retryable_error() { DEFAULT_CATALOG_NAME.to_string(), DEFAULT_SCHEMA_NAME.to_string(), false, - EventContext::default(), ddl_context.clone(), ); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; loop { @@ -149,7 +148,6 @@ async fn test_drop_database_recover() { DEFAULT_CATALOG_NAME.to_string(), DEFAULT_SCHEMA_NAME.to_string(), false, - EventContext::default(), ddl_context.clone(), ); let num_operating_regions = 1; diff --git a/src/common/meta/src/ddl/tests/drop_flow.rs b/src/common/meta/src/ddl/tests/drop_flow.rs index 0976aa572f7..6d36fc93ba3 100644 --- a/src/common/meta/src/ddl/tests/drop_flow.rs +++ b/src/common/meta/src/ddl/tests/drop_flow.rs @@ -26,7 +26,7 @@ use crate::ddl::test_util::flownode_handler::NaiveFlownodeHandler; use crate::ddl::tests::create_flow::{create_test_flow, create_test_pending_flow}; use crate::error; use crate::key::table_route::TableRouteValue; -use crate::rpc::ddl::{DropFlowTask, EventContext}; +use crate::rpc::ddl::DropFlowTask; use crate::test_util::{MockFlownodeManager, new_ddl_context}; pub(crate) fn test_drop_flow_task( @@ -48,7 +48,7 @@ async fn test_drop_flow_not_found() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let task = test_drop_flow_task("my_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropFlowProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowNotFound { .. }); } @@ -81,17 +81,17 @@ async fn test_drop_flow() { create_test_flow(&ddl_context, "my_flow", source_table_names, sink_table_name).await; // Drops the flows let task = test_drop_flow_task("my_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = DropFlowProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; // Drops if not exists let task = test_drop_flow_task("my_flow", flow_id, true); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = DropFlowProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; // Drops again let task = test_drop_flow_task("my_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropFlowProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowNotFound { .. }); } @@ -129,11 +129,11 @@ async fn test_drop_pending_flow_without_routes() { assert!(flow_info.flownode_ids().is_empty()); let task = test_drop_flow_task("drop_pending_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = DropFlowProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; let task = test_drop_flow_task("drop_pending_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropFlowProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowNotFound { .. }); } diff --git a/src/common/meta/src/ddl/tests/drop_table.rs b/src/common/meta/src/ddl/tests/drop_table.rs index 176f847c428..a278b951dd5 100644 --- a/src/common/meta/src/ddl/tests/drop_table.rs +++ b/src/common/meta/src/ddl/tests/drop_table.rs @@ -65,9 +65,9 @@ use crate::key::table_route::TableRouteValue; use crate::kv_backend::KvBackend; use crate::kv_backend::memory::MemoryKvBackend; use crate::peer::Peer; -use crate::rpc::ddl::{DropTableTask, EventContext}; +use crate::rpc::ddl::DropTableTask; #[cfg(feature = "enterprise")] -use crate::rpc::ddl::{PurgeDroppedTableTask, QueryContext, UndropTableTask}; +use crate::rpc::ddl::{PurgeDroppedTableTask, UndropTableTask}; use crate::rpc::router::{Region, RegionRoute}; #[cfg(feature = "enterprise")] use crate::rpc::store::{BatchDeleteRequest, PutRequest}; @@ -128,11 +128,7 @@ async fn undrop_at_restore_metadata( table_id: TableId, ) -> UndropTableProcedure { create_dropped_table(context, table_id, Some(1), Some(i64::MAX)).await; - let mut undrop = UndropTableProcedure::new( - new_undrop_table_task(table_id), - context.clone(), - EventContext::default(), - ); + let mut undrop = UndropTableProcedure::new(new_undrop_table_task(table_id), context.clone()); let procedure_context = new_test_procedure_context(); undrop.execute(&procedure_context).await.unwrap(); undrop.execute(&procedure_context).await.unwrap(); @@ -161,7 +157,6 @@ fn test_old_drop_table_json_defaults_to_hard_drop() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", 1024, false), runtime_context.clone(), - EventContext::default(), ); let mut old_data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); old_data @@ -169,7 +164,6 @@ fn test_old_drop_table_json_defaults_to_hard_drop() { .unwrap() .remove("soft_drop_enabled"); old_data.as_object_mut().unwrap().remove("dropped_at"); - old_data.as_object_mut().unwrap().remove("event_context"); let recovered = DropTableProcedure::from_json(&old_data.to_string(), runtime_context).unwrap(); let recovered_data: serde_json::Value = @@ -177,20 +171,12 @@ fn test_old_drop_table_json_defaults_to_hard_drop() { assert_eq!(recovered_data["soft_drop_enabled"], false); assert_eq!(recovered_data["dropped_at"], serde_json::Value::Null); - assert_eq!( - recovered_data["event_context"], - serde_json::json!({"reason": "unknown"}) - ); } #[test] fn test_drop_table_lock_key_includes_table_name() { let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); - let procedure = DropTableProcedure::new( - new_drop_table_task("foo", 1024, false), - context, - EventContext::default(), - ); + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); let keys = procedure .lock_key() @@ -219,11 +205,7 @@ fn test_new_drop_table_procedure_disables_soft_drop_without_enterprise() { let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); context.soft_drop_enabled = true; - let procedure = DropTableProcedure::new( - new_drop_table_task("foo", 1024, false), - context, - EventContext::default(), - ); + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); assert!(!procedure.data.soft_drop_enabled); } @@ -234,11 +216,7 @@ fn test_new_drop_table_procedure_preserves_soft_drop_with_enterprise() { let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); context.soft_drop_enabled = true; - let procedure = DropTableProcedure::new( - new_drop_table_task("foo", 1024, false), - context, - EventContext::default(), - ); + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); assert!(procedure.data.soft_drop_enabled); } @@ -273,7 +251,6 @@ async fn test_recovered_soft_drop_preserves_persisted_mode() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), soft_context.clone(), - EventContext::default(), ); let hard_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); @@ -313,13 +290,9 @@ async fn test_legacy_soft_drop_prepare_preserves_persisted_mode() { .unwrap(); let mut persisted: serde_json::Value = serde_json::from_str( - &DropTableProcedure::new( - new_drop_table_task("foo", table_id, false), - context.clone(), - EventContext::default(), - ) - .dump() - .unwrap(), + &DropTableProcedure::new(new_drop_table_task("foo", table_id, false), context.clone()) + .dump() + .unwrap(), ) .unwrap(); persisted["soft_drop_enabled"] = true.into(); @@ -349,11 +322,7 @@ async fn test_disabled_soft_drop_creates_hard_drop_procedure() { context.soft_drop_enabled = false; context.soft_drop_retention = Some(Duration::from_millis(100)); - let procedure = DropTableProcedure::new( - new_drop_table_task("foo", 1024, false), - context, - EventContext::default(), - ); + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); assert!(!procedure.data.soft_drop_enabled); assert_eq!(None, procedure.data.dropped_at); @@ -387,7 +356,6 @@ async fn test_recovered_hard_drop_ignores_soft_runtime_context() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), hard_context.clone(), - EventContext::default(), ); let mut soft_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); @@ -429,7 +397,7 @@ async fn test_on_prepare_table_not_exists_err() { .unwrap(); let task = new_drop_table_task("bar", table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context, EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context); assert_eq!(procedure.data.dropped_at, None); assert_eq!(procedure.data.retention_expires_at, None); let err = procedure.on_prepare().await.unwrap_err(); @@ -459,7 +427,6 @@ async fn test_soft_drop_prepare_assigns_stable_deadline_and_recovers_it() { let mut procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), ddl_context.clone(), - EventContext::default(), ); procedure.on_prepare().await.unwrap(); @@ -496,7 +463,6 @@ async fn test_soft_drop_prepare_rejects_deadline_overflow_before_metadata_delete let mut procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), ddl_context.clone(), - EventContext::default(), ); let error = procedure.on_prepare().await.unwrap_err(); @@ -536,13 +502,13 @@ async fn test_on_prepare_table() { let task = new_drop_table_task("bar", table_id, true); // Drop if exists - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); assert!(!procedure.rollback_supported()); let task = new_drop_table_task(table_name, table_id, false); // Drop table - let mut procedure = DropTableProcedure::new(task, ddl_context, EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context); procedure.on_prepare().await.unwrap(); } @@ -593,7 +559,7 @@ async fn test_on_datanode_drop_regions() { let task = new_drop_table_task(table_name, table_id, false); // Drop table - let mut procedure = DropTableProcedure::new(task, ddl_context, EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context); procedure.on_prepare().await.unwrap(); procedure.on_datanode_drop_regions(false).await.unwrap(); @@ -662,7 +628,7 @@ async fn test_on_datanode_drop_regions_remaps_addresses_when_retrying() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); put_datanode_address(&ddl_context, 1, "new-leader").await; @@ -721,7 +687,7 @@ async fn test_soft_drop_closes_regions_and_keeps_tombstone() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; @@ -835,7 +801,7 @@ async fn test_soft_drop_keeps_metadata_live_until_regions_close() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); put_datanode_address(&ddl_context, 1, "new-leader").await; @@ -922,7 +888,6 @@ async fn test_soft_drop_timestamp_is_stable_across_retry_and_recovery() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), ddl_context.clone(), - EventContext::default(), ); assert_eq!(procedure.data.dropped_at, None); let mut procedure = procedure; @@ -993,7 +958,7 @@ async fn test_hard_drop_keeps_delete_tombstone_flow() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); execute_procedure_until(&mut procedure, |p| { p.data.state == DropTableState::DeleteTombstone @@ -1036,20 +1001,13 @@ async fn test_create_table_succeeds_while_tombstone_exists() { .unwrap(); let drop_task = new_drop_table_task(table_name, dropped_table_id, false); - let mut drop_procedure = - DropTableProcedure::new(drop_task, ddl_context.clone(), EventContext::default()); + let mut drop_procedure = DropTableProcedure::new(drop_task, ddl_context.clone()); execute_procedure_until_done(&mut drop_procedure).await; let mut create_task = test_create_table_task(table_name, 1025); create_task.create_table.table_id = None; create_task.table_info.ident.table_id = 0; - let mut create_procedure = CreateTableProcedure::new( - create_task, - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut create_procedure = CreateTableProcedure::new(create_task, ddl_context.clone()).unwrap(); execute_procedure_until_done(&mut create_procedure).await; let live_table = ddl_context @@ -1094,8 +1052,7 @@ async fn test_hard_drop_recreated_table_fails_when_soft_tombstone_exists() { .unwrap(); let drop_task = new_drop_table_task(table_name, original_table_id, false); - let mut drop_procedure = - DropTableProcedure::new(drop_task, ddl_context.clone(), EventContext::default()); + let mut drop_procedure = DropTableProcedure::new(drop_task, ddl_context.clone()); execute_procedure_until_done(&mut drop_procedure).await; ddl_context.soft_drop_enabled = false; @@ -1112,7 +1069,6 @@ async fn test_hard_drop_recreated_table_fails_when_soft_tombstone_exists() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context, - EventContext::default(), ); let err = procedure.on_prepare().await.unwrap_err(); @@ -1142,8 +1098,7 @@ async fn test_hard_drop_recreated_table_ignores_previous_orphan_tombstone() { .unwrap(); let drop_task = new_drop_table_task(table_name, original_table_id, false); - let mut drop_procedure = - DropTableProcedure::new(drop_task, ddl_context.clone(), EventContext::default()); + let mut drop_procedure = DropTableProcedure::new(drop_task, ddl_context.clone()); execute_procedure_until(&mut drop_procedure, |p| { p.data.state == DropTableState::DeleteTombstone }) @@ -1162,7 +1117,6 @@ async fn test_hard_drop_recreated_table_ignores_previous_orphan_tombstone() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context, - EventContext::default(), ); procedure.on_prepare().await.unwrap(); @@ -1202,7 +1156,6 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; assert!( @@ -1219,11 +1172,8 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { while rx.try_recv().is_ok() {} - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; let live_table = ddl_context @@ -1314,16 +1264,12 @@ async fn test_undrop_table_opens_regions_before_restoring_live_metadata() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -1387,17 +1333,13 @@ async fn test_undrop_logical_table_skips_datanode_open() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, logical_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(logical_table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(logical_table_id), ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; let live_table = ddl_context @@ -1428,7 +1370,6 @@ async fn test_soft_drop_metric_logical_table_falls_back_to_hard_drop() { let mut procedure = DropTableProcedure::new( new_drop_table_task("foo", logical_table_id, false), ddl_context.clone(), - EventContext::default(), ); procedure.on_prepare().await.unwrap(); assert!(!procedure.data.soft_drop_enabled); @@ -1472,7 +1413,6 @@ async fn test_soft_drop_metric_physical_table_remains_enabled() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut procedure).await; @@ -1517,7 +1457,6 @@ async fn test_soft_drop_file_engine_table_falls_back_to_hard_drop() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); procedure.on_prepare().await.unwrap(); assert!(!procedure.data.soft_drop_enabled); @@ -1554,7 +1493,6 @@ async fn test_file_engine_fallback_is_resolved_before_tombstone_conflict() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context, - EventContext::default(), ); procedure.on_prepare().await.unwrap(); @@ -1570,11 +1508,8 @@ async fn test_undrop_metric_logical_table_fails() { let logical_table_id = create_metric_logical_table_tombstone(&ddl_context, physical_table_id, "foo").await; - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(logical_table_id), - ddl_context, - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(logical_table_id), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::Unsupported); @@ -1592,7 +1527,6 @@ async fn test_purge_metric_logical_table_fails() { let mut procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(logical_table_id), ddl_context, - EventContext::default(), ); let err = procedure .execute(&new_test_procedure_context()) @@ -1624,7 +1558,6 @@ async fn test_undrop_table_fails_when_live_name_exists() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; ddl_context @@ -1637,11 +1570,8 @@ async fn test_undrop_table_fails_when_live_name_exists() { .await .unwrap(); - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(dropped_table_id), - ddl_context, - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(dropped_table_id), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); @@ -1669,15 +1599,11 @@ async fn test_undrop_table_fails_when_live_name_is_created_after_prepare() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(dropped_table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(dropped_table_id), ddl_context.clone()); procedure.on_prepare().await.unwrap(); ddl_context .table_metadata_manager @@ -1741,17 +1667,13 @@ async fn test_undrop_table_closes_opened_regions_when_restore_metadata_races_wit let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} detector_controller.clear().await; - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(dropped_table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(dropped_table_id), ddl_context.clone()); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -1845,7 +1767,6 @@ async fn test_undrop_table_lock_key_includes_original_table_name_before_prepare( let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; @@ -1860,7 +1781,6 @@ async fn test_undrop_table_lock_key_includes_original_table_name_before_prepare( new_undrop_table_task(table_id), ddl_context, Some(original_table_name), - EventContext::default(), ); let keys = procedure @@ -1928,12 +1848,10 @@ async fn test_undrop_rejects_stale_id_after_name_tombstone_is_consumed() { new_undrop_table_task(original_table_id), ddl_context.clone(), Some(original_table_name.clone()), - EventContext::default(), ); let mut hard_drop = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut hard_drop).await; assert!( @@ -1988,15 +1906,11 @@ async fn test_undrop_table_replayed_restore_metadata_is_idempotent() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -2100,7 +2014,6 @@ async fn test_purge_dropped_table_cleans_regions_offline_and_deletes_tombstone() let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} @@ -2109,7 +2022,6 @@ async fn test_purge_dropped_table_cleans_regions_offline_and_deletes_tombstone() let mut procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut procedure).await; @@ -2188,7 +2100,6 @@ async fn test_automatic_purge_rechecks_unexpired_tombstone() { let procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); data["check_expired"] = true.into(); @@ -2245,7 +2156,6 @@ async fn test_recovered_automatic_purge_rechecks_unexpired_tombstone() { let procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); data["check_expired"] = true.into(); @@ -2323,7 +2233,6 @@ async fn test_recovered_automatic_purge_obeys_runtime_disable() { let procedure = PurgeDroppedTableProcedure::new_if_expired( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); let mut recovered = PurgeDroppedTableProcedure::from_json(&procedure.dump().unwrap(), ddl_context.clone()) @@ -2355,7 +2264,6 @@ async fn test_recovered_automatic_purge_finishes_after_cleanup_is_claimed() { let mut procedure = PurgeDroppedTableProcedure::new_if_expired( new_purge_dropped_table_task(table_id), enabled_context.clone(), - EventContext::default(), ); procedure .execute(&new_test_procedure_context()) @@ -2419,7 +2327,6 @@ async fn test_recovered_automatic_purge_ignores_foreign_generation_claim() { let mut procedure = PurgeDroppedTableProcedure::new_if_expired( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); procedure .execute(&new_test_procedure_context()) @@ -2462,11 +2369,8 @@ async fn test_undrop_rejects_table_claimed_for_purge() { .await .unwrap(); - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); let error = procedure.on_prepare().await.unwrap_err(); assert_eq!(StatusCode::TableNotFound, error.status_code()); @@ -2488,11 +2392,8 @@ async fn test_recovered_undrop_rejects_tombstone_deleted_by_purge() { let table_id = 1024; create_dropped_table(&ddl_context, table_id, Some(1), Some(i64::MAX)).await; - let mut undrop = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); undrop.on_prepare().await.unwrap(); let open_regions_state = undrop.dump().unwrap(); undrop.execute(&new_test_procedure_context()).await.unwrap(); @@ -2501,7 +2402,6 @@ async fn test_recovered_undrop_rejects_tombstone_deleted_by_purge() { let mut purge = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut purge).await; @@ -2533,11 +2433,8 @@ async fn test_recovered_undrop_rejects_replacement_generation() { .await .unwrap(); - let mut undrop = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); undrop.on_prepare().await.unwrap(); let persisted = undrop.dump().unwrap(); kv_backend @@ -2603,11 +2500,8 @@ async fn test_legacy_unmarked_tombstone_can_be_undropped_and_purged() { assert_eq!(dropped_tables.len(), 1); assert_eq!(dropped_tables[0].dropped_at, None); - let mut undrop = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); execute_procedure_until_done(&mut undrop).await; ddl_context @@ -2624,7 +2518,6 @@ async fn test_legacy_unmarked_tombstone_can_be_undropped_and_purged() { let mut purge = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut purge).await; assert!( @@ -2667,7 +2560,6 @@ async fn test_purge_dropped_table_by_id_selects_tombstone_when_live_table_exists let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; ddl_context @@ -2684,7 +2576,6 @@ async fn test_purge_dropped_table_by_id_selects_tombstone_when_live_table_exists let mut procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(dropped_table_id), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut procedure).await; @@ -2735,12 +2626,8 @@ async fn test_on_rollback() { let physical_table_id = table_id; // Creates the logical table metadata. let task = test_create_logical_table_task("foo"); - let mut procedure = CreateLogicalTablesProcedure::new( - vec![task], - physical_table_id, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = + CreateLogicalTablesProcedure::new(vec![task], physical_table_id, ddl_context.clone()); procedure.on_prepare().await.unwrap(); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); @@ -2753,8 +2640,7 @@ async fn test_on_rollback() { // Drops the physical table { let task = new_drop_table_task("phy_table", physical_table_id, false); - let mut procedure = - DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); assert!(procedure.rollback_supported()); procedure.on_delete_metadata().await.unwrap(); @@ -2765,8 +2651,7 @@ async fn test_on_rollback() { // The physical table can be dropped again after the metadata rollback. let retry_task = new_drop_table_task("phy_table", physical_table_id, false); - let mut retry = - DropTableProcedure::new(retry_task, ddl_context.clone(), EventContext::default()); + let mut retry = DropTableProcedure::new(retry_task, ddl_context.clone()); retry.on_prepare().await.unwrap(); retry.on_delete_metadata().await.unwrap(); retry.rollback(&ctx).await.unwrap(); @@ -2783,7 +2668,7 @@ async fn test_on_rollback() { // Drops the logical table ddl_context.soft_drop_enabled = false; let task = new_drop_table_task("foo", table_ids[0], false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); assert!(!procedure.rollback_supported()); } @@ -2880,8 +2765,7 @@ async fn test_memory_region_keeper_guard_dropped_on_procedure_done() { let logical_table_id = create_logical_table(ddl_context.clone(), physical_table_id, "s").await; let inner_test = |task: DropTableTask| async { - let mut procedure = - DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); execute_procedure_until(&mut procedure, |p| { p.data.state == DropTableState::InvalidateTableCache }) @@ -2933,8 +2817,7 @@ async fn test_from_json() { let physical_table_id = create_physical_table(&ddl_context, "t").await; let task = new_drop_table_task("t", physical_table_id, false); - let mut procedure = - DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); execute_procedure_until(&mut procedure, |p| p.data.state == state).await; let data = procedure.dump().unwrap(); assert_eq!( @@ -2963,7 +2846,7 @@ async fn test_from_json() { let physical_table_id = create_physical_table(&ddl_context, "t").await; let task = new_drop_table_task("t", physical_table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; let data = procedure.dump().unwrap(); assert_eq!( diff --git a/src/common/meta/src/ddl/tests/drop_view.rs b/src/common/meta/src/ddl/tests/drop_view.rs index 09b2a450c52..412c5a78f0b 100644 --- a/src/common/meta/src/ddl/tests/drop_view.rs +++ b/src/common/meta/src/ddl/tests/drop_view.rs @@ -24,7 +24,7 @@ use crate::ddl::drop_view::{DropViewProcedure, DropViewState}; use crate::ddl::test_util::create_table::test_create_table_task; use crate::ddl::tests::create_view::{test_create_view_task, test_table_names}; use crate::key::table_route::TableRouteValue; -use crate::rpc::ddl::{DropViewTask, EventContext}; +use crate::rpc::ddl::DropViewTask; use crate::test_util::{MockDatanodeManager, new_ddl_context}; pub(crate) fn new_drop_view_task( @@ -63,7 +63,7 @@ async fn test_on_prepare_view_not_exists_err() { .unwrap(); let task = new_drop_view_task("bar", view_id, false); - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropViewProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::TableNotFound); } @@ -87,7 +87,7 @@ async fn test_on_prepare_not_view_err() { .unwrap(); let task = new_drop_view_task(view_name, view_id, false); - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropViewProcedure::new(task, ddl_context); // It's not a view, expect error let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::InvalidArguments); @@ -117,12 +117,12 @@ async fn test_on_prepare_success() { let task = new_drop_view_task("bar", view_id, true); // Drop if exists - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = DropViewProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); let task = new_drop_view_task(view_name, view_id, false); // Prepare success - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropViewProcedure::new(task, ddl_context); procedure.on_prepare().await.unwrap(); assert_eq!(DropViewState::DeleteMetadata, procedure.state()); } @@ -161,7 +161,7 @@ async fn test_drop_view_success() { let task = new_drop_view_task(view_name, view_id, false); // Prepare success - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = DropViewProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; assert_eq!(DropViewState::InvalidateViewCache, procedure.state()); @@ -178,7 +178,7 @@ async fn test_drop_view_success() { // Drop again let task = new_drop_view_task(view_name, view_id, false); - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropViewProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::TableNotFound); } diff --git a/src/common/meta/src/ddl/tests/event/database.rs b/src/common/meta/src/ddl/tests/event/database.rs index 29cb48f7b80..8649982ed15 100644 --- a/src/common/meta/src/ddl/tests/event/database.rs +++ b/src/common/meta/src/ddl/tests/event/database.rs @@ -25,12 +25,12 @@ use common_event_recorder::event_table::{ PROCEDURE_ID_COLUMN as EVENT_TABLE_PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN as EVENT_TABLE_PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN as EVENT_TABLE_PROCEDURE_TRIGGER_COLUMN, - SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, + SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; -use super::test_util::{assert_event_filter, default_event_context_value, procedure_trigger_value}; +use super::test_util::assert_event_filter; use crate::ddl::alter_database::AlterDatabaseProcedure; use crate::ddl::create_database::CreateDatabaseProcedure; use crate::ddl::drop_database::DropDatabaseProcedure; @@ -39,7 +39,7 @@ use crate::ddl::event::database::{ DatabaseDdlEvent, }; use crate::rpc::ddl::{ - AlterDatabaseKind, EventContext, SetDatabaseOption, SetDatabaseOptions, UnsetDatabaseOption, + AlterDatabaseKind, SetDatabaseOption, SetDatabaseOptions, UnsetDatabaseOption, UnsetDatabaseOptions, }; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -50,13 +50,7 @@ fn test_create_database_submitted_event_contract() { ("password".to_string(), "do-not-record".to_string()), ("compaction.type".to_string(), "twcs".to_string()), ]); - let event = DatabaseDdlEvent::create_submitted( - "greptime", - "metrics", - true, - &options, - EventContext::default(), - ); + let event = DatabaseDdlEvent::create_submitted("greptime", "metrics", true, &options); assert_event_locator( &event, @@ -93,7 +87,6 @@ fn test_alter_database_set_and_unset_event_contracts() { SetDatabaseOption::Other("secret_token".to_string(), "hidden".to_string()), SetDatabaseOption::Ttl(std::time::Duration::from_secs(3600).into()), ])), - EventContext::default(), ); assert_event_locator( &set, @@ -120,7 +113,6 @@ fn test_alter_database_set_and_unset_event_contracts() { UnsetDatabaseOption::Ttl, UnsetDatabaseOption::Other("compaction.type".to_string()), ])), - EventContext::default(), ); assert_event_locator( &unset, @@ -140,8 +132,7 @@ fn test_alter_database_set_and_unset_event_contracts() { #[test] fn test_drop_database_submitted_event_contract() { - let event = - DatabaseDdlEvent::drop_submitted("greptime", "metrics", true, EventContext::default()); + let event = DatabaseDdlEvent::drop_submitted("greptime", "metrics", true); assert_event_locator( &event, @@ -177,14 +168,7 @@ fn test_database_lifecycle_events_preserve_locator_and_null_intent() { assert_eq!( DatabaseDdlEvent::create_lifecycle("greptime", "metrics").extra_schema(), - DatabaseDdlEvent::create_submitted( - "c", - "s", - false, - &HashMap::new(), - EventContext::default(), - ) - .extra_schema() + DatabaseDdlEvent::create_submitted("c", "s", false, &HashMap::new()).extra_schema() ); } @@ -198,7 +182,6 @@ fn test_database_events_preserve_procedure_envelope_contract() { "metrics", false, &HashMap::new(), - EventContext::default(), )), ProcedureState::Running, EventTrigger::Submitted, @@ -237,7 +220,6 @@ fn test_create_database_event_filter() { false, HashMap::new(), None, - EventContext::default(), new_ddl_context(Arc::new(MockDatanodeManager::new(()))), ); @@ -258,7 +240,6 @@ fn test_alter_database_event_filter() { )), }, }, - EventContext::default(), new_ddl_context(Arc::new(MockDatanodeManager::new(()))), ) .unwrap(); @@ -272,7 +253,6 @@ fn test_drop_database_event_filter() { "greptime".to_string(), "metrics".to_string(), false, - EventContext::default(), new_ddl_context(Arc::new(MockDatanodeManager::new(()))), ); @@ -285,14 +265,12 @@ fn assert_event_locator( catalog_name: Option<&str>, schema_name: Option<&str>, ) { - let has_event_context = !event.json_payload().unwrap().is_null(); assert_event_contract( event, event_type, &[ EVENT_TABLE_CATALOG_NAME_COLUMN.column_schema(), EVENT_TABLE_SCHEMA_NAME_COLUMN.column_schema(), - EVENT_CONTEXT_COLUMN.column_schema(), ], &[Row { values: vec![ @@ -302,11 +280,6 @@ fn assert_event_locator( Value { value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, - if has_event_context { - default_event_context_value() - } else { - Value { value_data: None } - }, ], }], ); @@ -320,7 +293,6 @@ fn assert_procedure_event_contract( catalog_name: Option<&str>, schema_name: Option<&str>, ) { - let has_event_context = procedure_trigger == "Submitted"; assert_event_contract( event, event_type, @@ -346,18 +318,14 @@ fn assert_procedure_event_contract( Value { value_data: Some(ValueData::StringValue(String::new())), }, - procedure_trigger_value(procedure_trigger), + jsonb_value(&serde_json::json!({"type": procedure_trigger})), Value { value_data: catalog_name.map(|value| ValueData::StringValue(value.to_string())), }, Value { value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, - if has_event_context { - default_event_context_value() - } else { - Value { value_data: None } - }, + Value { value_data: None }, ], }], ); diff --git a/src/common/meta/src/ddl/tests/event/flow.rs b/src/common/meta/src/ddl/tests/event/flow.rs index 39f24b74e51..aae3613a087 100644 --- a/src/common/meta/src/ddl/tests/event/flow.rs +++ b/src/common/meta/src/ddl/tests/event/flow.rs @@ -21,12 +21,13 @@ use common_event_recorder::Event; use common_event_recorder::event_table::{ CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, + jsonb_value, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; use table::table_name::TableName; -use super::test_util::{assert_event_filter, default_event_context_value, procedure_trigger_value}; +use super::test_util::assert_event_filter; use crate::ddl::create_flow::CreateFlowProcedure; use crate::ddl::drop_flow::DropFlowProcedure; use crate::ddl::event::flow::{ @@ -35,7 +36,6 @@ use crate::ddl::event::flow::{ use crate::ddl::test_util::flownode_handler::NaiveFlownodeHandler; use crate::ddl::tests::create_flow::{test_create_flow_task, test_query_context}; use crate::ddl::tests::drop_flow::test_drop_flow_task; -use crate::rpc::ddl::EventContext; use crate::test_util::{MockFlownodeManager, new_ddl_context}; #[test] @@ -49,7 +49,6 @@ fn test_flow_submitted_event_contracts() { expire_after: Some(300), eval_interval_secs: Some(60), }, - EventContext::default(), ); assert_event_contract( &create, @@ -60,7 +59,6 @@ fn test_flow_submitted_event_contracts() { ValueData::StringValue("greptime".to_string()).into(), ValueData::StringValue("metrics".to_string()).into(), Value { value_data: None }, - default_event_context_value(), ], }], ); @@ -75,8 +73,7 @@ fn test_flow_submitted_event_contracts() { }) ); - let drop = - FlowDdlEvent::drop_submitted("greptime", "metrics", 42, true, EventContext::default()); + let drop = FlowDdlEvent::drop_submitted("greptime", "metrics", 42, true); assert_event_contract( &drop, DROP_FLOW_EVENT_TYPE, @@ -86,7 +83,6 @@ fn test_flow_submitted_event_contracts() { ValueData::StringValue("greptime".to_string()).into(), ValueData::StringValue("metrics".to_string()).into(), ValueData::U32Value(42).into(), - default_event_context_value(), ], }], ); @@ -125,7 +121,6 @@ fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { } else { Value { value_data: None } }, - Value { value_data: None }, ], }], ); @@ -142,7 +137,6 @@ fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { ValueData::StringValue("greptime".to_string()).into(), ValueData::StringValue("metrics".to_string()).into(), ValueData::U32Value(42).into(), - Value { value_data: None }, ], }], ); @@ -162,7 +156,6 @@ fn test_flow_events_preserve_procedure_envelope_contract() { expire_after: None, eval_interval_secs: None, }, - EventContext::default(), )), ProcedureState::Running, EventTrigger::Submitted, @@ -212,7 +205,6 @@ fn test_create_flow_event_filter() { false, ), test_query_context(), - EventContext::default(), new_ddl_context(Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler))), ); assert_event_filter(&procedure, CREATE_FLOW_EVENT_TYPE); @@ -222,7 +214,6 @@ fn test_create_flow_event_filter() { fn test_drop_flow_event_filter() { let procedure = DropFlowProcedure::new( test_drop_flow_task("flow", 42, false), - EventContext::default(), new_ddl_context(Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler))), ); assert_event_filter(&procedure, DROP_FLOW_EVENT_TYPE); @@ -233,7 +224,6 @@ fn flow_schema() -> Vec { CATALOG_NAME_COLUMN.column_schema(), FLOW_NAME_COLUMN.column_schema(), FLOW_ID_COLUMN.column_schema(), - EVENT_CONTEXT_COLUMN.column_schema(), ] } @@ -257,6 +247,7 @@ fn assert_procedure_event_contract( PROCEDURE_TRIGGER_COLUMN.column_schema(), ]; schema.extend(flow_schema()); + schema.push(EVENT_CONTEXT_COLUMN.column_schema()); assert_event_contract( event, event_type, @@ -266,7 +257,7 @@ fn assert_procedure_event_contract( ValueData::StringValue(event.procedure_id.to_string()).into(), ValueData::StringValue(state.to_string()).into(), ValueData::StringValue(String::new()).into(), - procedure_trigger_value(trigger), + jsonb_value(&serde_json::json!({"type": trigger})), optional_string(locator.catalog_name), optional_string(locator.flow_name), locator @@ -274,11 +265,7 @@ fn assert_procedure_event_contract( .map(ValueData::U32Value) .map(Into::into) .unwrap_or(Value { value_data: None }), - if trigger == "Submitted" { - default_event_context_value() - } else { - Value { value_data: None } - }, + Value { value_data: None }, ], }], ); diff --git a/src/common/meta/src/ddl/tests/event/table.rs b/src/common/meta/src/ddl/tests/event/table.rs index 3326cb06808..24182fd19af 100644 --- a/src/common/meta/src/ddl/tests/event/table.rs +++ b/src/common/meta/src/ddl/tests/event/table.rs @@ -19,13 +19,13 @@ use api::v1::value::ValueData; use api::v1::{ColumnDataType, Repartition, SemanticType, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, - TABLE_ID_COLUMN, TABLE_NAME_COLUMN, jsonb_value, + CATALOG_NAME_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN, + TABLE_NAME_COLUMN, }; use common_event_recorder::{Event, EventTypeFilter}; use common_procedure::{ - ChildSubmissionOutcome, EventRuntimeContext, EventTrigger, Procedure, ProcedureId, - ProcedureState, RetryPhase, + ChildSubmissionOutcome, EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState, + RetryPhase, }; use common_time::Timestamp; use serde_json::{Value as JsonValue, json}; @@ -52,9 +52,9 @@ use crate::ddl::truncate_table::TruncateTableProcedure; use crate::ddl::undrop_table::UndropTableProcedure; use crate::key::DeserializedValueWithBytes; use crate::key::table_info::TableInfoValue; -use crate::rpc::ddl::{DropTableTask, EventContext, QueryContext, TruncateTableTask}; +use crate::rpc::ddl::{DropTableTask, TruncateTableTask}; #[cfg(feature = "enterprise")] -use crate::rpc::ddl::{PurgeDroppedTableTask, TriggerReason, UndropTableTask}; +use crate::rpc::ddl::{PurgeDroppedTableTask, UndropTableTask}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; struct EventCase { @@ -175,11 +175,7 @@ fn later_lifecycle_events_are_uniform() { .extra_rows() .unwrap() .into_iter() - .map(|mut row| { - // Lifecycle events retain locators but not submitted event context. - *row.values.last_mut().unwrap() = Value::default(); - row.values - }) + .map(|row| row.values) .collect::>(); for trigger in &triggers { @@ -205,13 +201,7 @@ fn later_lifecycle_events_are_uniform() { fn create_success_events_keep_allocated_ids() { let mut task = test_create_table_task("create_success"); task.table_info.ident.table_id = 7; - let create_table = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - test_context(), - ) - .unwrap(); + let create_table = CreateTableProcedure::new(task, test_context()).unwrap(); let state = ProcedureState::Done { output: Some(Arc::new(42_u32)), }; @@ -230,7 +220,6 @@ fn create_success_events_keep_allocated_ids() { test_create_logical_table_task("bar"), ], 1024, - EventContext::default(), test_context(), ); let state = ProcedureState::Done { @@ -269,14 +258,13 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "create"), true, "mito2", - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "create_if_not_exists": true, "engine": "mito2", }), - rows: vec![submitted_table_locator_values(Some("create"), None)], + rows: vec![table_locator_values(Some("create"), None)], }, EventCase { event_type: TableDdlEventType::CreateLogicalTables, @@ -288,15 +276,14 @@ fn event_cases() -> Vec { .with_physical_table_id(10), ], 2, - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "table_count": 2, }), rows: vec![ - submitted_logical_locator_values("logical1", None, 10), - submitted_logical_locator_values("logical2", None, 10), + logical_locator_values("logical1", None, 10), + logical_locator_values("logical2", None, 10), ], }, EventCase { @@ -305,13 +292,12 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "alter") .with_table_id(11), Some("drop_columns"), - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "kind": "drop_columns", }), - rows: vec![submitted_table_locator_values(Some("alter"), Some(11))], + rows: vec![table_locator_values(Some("alter"), Some(11))], }, EventCase { event_type: TableDdlEventType::AlterLogicalTables, @@ -324,7 +310,6 @@ fn event_cases() -> Vec { ], 2, ["rename_table", "add_columns", "add_columns"], - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, @@ -332,8 +317,8 @@ fn event_cases() -> Vec { "kinds": ["add_columns", "rename_table"], }), rows: vec![ - submitted_logical_locator_values("logical1", None, 10), - submitted_logical_locator_values("logical2", None, 10), + logical_locator_values("logical1", None, 10), + logical_locator_values("logical2", None, 10), ], }, EventCase { @@ -342,33 +327,26 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "drop") .with_table_id(12), true, - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "drop_if_exists": true, }), - rows: vec![submitted_table_locator_values(Some("drop"), Some(12))], + rows: vec![table_locator_values(Some("drop"), Some(12))], }, #[cfg(feature = "enterprise")] EventCase { event_type: TableDdlEventType::UndropTable, - event: TableDdlEvent::undrop_table_submitted( - TableDdlLocator::from_table_id(13), - EventContext::default(), - ), + event: TableDdlEvent::undrop_table_submitted(TableDdlLocator::from_table_id(13)), payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![submitted_table_locator_values(None, Some(13))], + rows: vec![table_locator_values(None, Some(13))], }, #[cfg(feature = "enterprise")] EventCase { event_type: TableDdlEventType::PurgeDroppedTable, - event: TableDdlEvent::purge_dropped_table_submitted( - TableDdlLocator::from_table_id(14), - EventContext::default(), - ), + event: TableDdlEvent::purge_dropped_table_submitted(TableDdlLocator::from_table_id(14)), payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![submitted_table_locator_values(None, Some(14))], + rows: vec![table_locator_values(None, Some(14))], }, EventCase { event_type: TableDdlEventType::TruncateTable, @@ -376,41 +354,29 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "truncate") .with_table_id(15), 4, - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "time_range_count": 4, }), - rows: vec![submitted_table_locator_values(Some("truncate"), Some(15))], + rows: vec![table_locator_values(Some("truncate"), Some(15))], }, ] } fn procedure_cases() -> Vec { - let create_table = CreateTableProcedure::new( - test_create_table_task("create"), - QueryContext::default(), - EventContext::default(), - test_context(), - ) - .unwrap(); + let create_table = + CreateTableProcedure::new(test_create_table_task("create"), test_context()).unwrap(); let create_logical_tables = CreateLogicalTablesProcedure::new( vec![ test_create_logical_table_task("logical1"), test_create_logical_table_task("logical2"), ], 41, - EventContext::default(), test_context(), ); - let alter_table = AlterTableProcedure::new( - 42, - test_alter_table_task("alter"), - EventContext::default(), - test_context(), - ) - .unwrap(); + let alter_table = + AlterTableProcedure::new(42, test_alter_table_task("alter"), test_context()).unwrap(); let alter_logical_tables = AlterLogicalTablesProcedure::new( vec![ make_alter_logical_table_add_column_task( @@ -425,7 +391,6 @@ fn procedure_cases() -> Vec { ), ], 43, - EventContext::default(), test_context(), ); let drop_table = DropTableProcedure::new( @@ -437,20 +402,12 @@ fn procedure_cases() -> Vec { drop_if_exists: true, }, test_context(), - EventContext::default(), ); #[cfg(feature = "enterprise")] - let undrop_table = UndropTableProcedure::new( - UndropTableTask { table_id: 45 }, - test_context(), - EventContext::default(), - ); + let undrop_table = UndropTableProcedure::new(UndropTableTask { table_id: 45 }, test_context()); #[cfg(feature = "enterprise")] - let purge_dropped_table = PurgeDroppedTableProcedure::new_if_expired( - PurgeDroppedTableTask { table_id: 46 }, - test_context(), - EventContext::new(TriggerReason::ScheduledGc), - ); + let purge_dropped_table = + PurgeDroppedTableProcedure::new(PurgeDroppedTableTask { table_id: 46 }, test_context()); let truncate_table = truncate_procedure(TruncateTableTask { catalog: DEFAULT_CATALOG_NAME.to_string(), schema: DEFAULT_SCHEMA_NAME.to_string(), @@ -471,7 +428,7 @@ fn procedure_cases() -> Vec { "create_if_not_exists": false, "engine": "mito2", }), - rows: vec![submitted_table_locator_values(Some("create"), None)], + rows: vec![table_locator_values(Some("create"), None)], }, ProcedureCase { procedure: Box::new(create_logical_tables), @@ -481,8 +438,8 @@ fn procedure_cases() -> Vec { "table_count": 2, }), rows: vec![ - submitted_logical_locator_values("logical1", None, 41), - submitted_logical_locator_values("logical2", None, 41), + logical_locator_values("logical1", None, 41), + logical_locator_values("logical2", None, 41), ], }, ProcedureCase { @@ -492,7 +449,7 @@ fn procedure_cases() -> Vec { "version": TABLE_DDL_PAYLOAD_VERSION, "kind": "drop_columns", }), - rows: vec![submitted_table_locator_values(Some("alter"), Some(42))], + rows: vec![table_locator_values(Some("alter"), Some(42))], }, ProcedureCase { procedure: Box::new(alter_logical_tables), @@ -503,8 +460,8 @@ fn procedure_cases() -> Vec { "kinds": ["add_columns"], }), rows: vec![ - submitted_logical_locator_values("logical1", None, 43), - submitted_logical_locator_values("logical2", None, 43), + logical_locator_values("logical1", None, 43), + logical_locator_values("logical2", None, 43), ], }, ProcedureCase { @@ -514,25 +471,21 @@ fn procedure_cases() -> Vec { "version": TABLE_DDL_PAYLOAD_VERSION, "drop_if_exists": true, }), - rows: vec![submitted_table_locator_values(Some("drop"), Some(44))], + rows: vec![table_locator_values(Some("drop"), Some(44))], }, #[cfg(feature = "enterprise")] ProcedureCase { procedure: Box::new(undrop_table), event_type: "undrop_table", payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![submitted_table_locator_values(None, Some(45))], + rows: vec![table_locator_values(None, Some(45))], }, #[cfg(feature = "enterprise")] ProcedureCase { procedure: Box::new(purge_dropped_table), event_type: "purge_dropped_table", payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![submitted_table_locator_values_with_event_context( - None, - Some(46), - EventContext::new(TriggerReason::ScheduledGc), - )], + rows: vec![table_locator_values(None, Some(46))], }, ProcedureCase { procedure: Box::new(truncate_table), @@ -541,7 +494,7 @@ fn procedure_cases() -> Vec { "version": TABLE_DDL_PAYLOAD_VERSION, "time_range_count": 1, }), - rows: vec![submitted_table_locator_values(Some("truncate"), Some(47))], + rows: vec![table_locator_values(Some("truncate"), Some(47))], }, ] } @@ -562,7 +515,6 @@ fn expected_schema(event_type: TableDdlEventType) -> Vec<(&'static str, i32)> { ColumnDataType::Uint32 as i32, )); } - schema.push((EVENT_CONTEXT_COLUMN.name(), ColumnDataType::Binary as i32)); schema } @@ -575,14 +527,12 @@ fn table_locator_values(table_name: Option<&str>, table_id: Option) -> Vec< } else { (Value::default(), Value::default()) }; - let mut values = vec![ + vec![ catalog_name, schema_name, table_name.map(string_value).unwrap_or_default(), table_id.map(table_id_value).unwrap_or_default(), - ]; - values.push(Value::default()); - values + ] } fn logical_locator_values( @@ -591,44 +541,10 @@ fn logical_locator_values( physical_table_id: u32, ) -> Vec { let mut values = table_locator_values(Some(table_name), table_id); - let event_context = values.pop().unwrap(); values.push(table_id_value(physical_table_id)); - values.push(event_context); values } -fn submitted_table_locator_values(table_name: Option<&str>, table_id: Option) -> Vec { - submitted_table_locator_values_with_event_context(table_name, table_id, EventContext::default()) -} - -fn submitted_table_locator_values_with_event_context( - table_name: Option<&str>, - table_id: Option, - event_context: EventContext, -) -> Vec { - let mut values = table_locator_values(table_name, table_id); - *values.last_mut().unwrap() = event_context_value_for(event_context); - values -} - -fn submitted_logical_locator_values( - table_name: &str, - table_id: Option, - physical_table_id: u32, -) -> Vec { - let mut values = logical_locator_values(table_name, table_id, physical_table_id); - *values.last_mut().unwrap() = event_context_value(); - values -} - -fn event_context_value() -> Value { - event_context_value_for(EventContext::default()) -} - -fn event_context_value_for(event_context: EventContext) -> Value { - jsonb_value(&serde_json::to_value(event_context).unwrap()) -} - fn string_value(value: &str) -> Value { ValueData::StringValue(value.to_string()).into() } @@ -643,7 +559,6 @@ fn truncate_procedure(task: TruncateTableTask) -> TruncateTableProcedure { task, DeserializedValueWithBytes::from_inner(TableInfoValue::new(table_info)), test_context(), - EventContext::default(), ) } @@ -661,11 +576,12 @@ fn event_for_state( lifecycle_state: &ProcedureState, ) -> Box { procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, event_type_filter: Arc::new(EventTypeFilter::All), + event_context: None, }) .unwrap() } diff --git a/src/common/meta/src/ddl/tests/event/test_util.rs b/src/common/meta/src/ddl/tests/event/test_util.rs index 281b5deb488..e20d294180e 100644 --- a/src/common/meta/src/ddl/tests/event/test_util.rs +++ b/src/common/meta/src/ddl/tests/event/test_util.rs @@ -15,32 +15,21 @@ use std::collections::HashSet; use std::sync::Arc; -use api::v1::Value; use common_event_recorder::EventTypeFilter; -use common_event_recorder::event_table::jsonb_value; -use common_procedure::{EventRuntimeContext, EventTrigger, Procedure, ProcedureId, ProcedureState}; - -use crate::rpc::ddl::EventContext; - -pub(crate) fn default_event_context_value() -> Value { - jsonb_value(&serde_json::to_value(EventContext::default()).unwrap()) -} - -pub(crate) fn procedure_trigger_value(trigger: &str) -> Value { - jsonb_value(&serde_json::json!({"type": trigger})) -} +use common_procedure::{EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState}; pub(crate) fn assert_event_filter(procedure: &dyn Procedure, event_type: &str) { let state = ProcedureState::Running; - let runtime_context = |event_type_filter| EventRuntimeContext { + let event_context = |event_type_filter| EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: Arc::new(event_type_filter), + event_context: None, }; let allowed = procedure - .event(&runtime_context(EventTypeFilter::Only(HashSet::from([ + .event(&event_context(EventTypeFilter::Only(HashSet::from([ event_type.to_string(), ])))) .unwrap(); @@ -49,7 +38,7 @@ pub(crate) fn assert_event_filter(procedure: &dyn Procedure, event_type: &str) { for denied in [HashSet::from(["other_event".to_string()]), HashSet::new()] { assert!( procedure - .event(&runtime_context(EventTypeFilter::Only(denied))) + .event(&event_context(EventTypeFilter::Only(denied))) .is_none() ); } diff --git a/src/common/meta/src/ddl/tests/event/view.rs b/src/common/meta/src/ddl/tests/event/view.rs index cabe7418ffb..5cb5a6bb7d9 100644 --- a/src/common/meta/src/ddl/tests/event/view.rs +++ b/src/common/meta/src/ddl/tests/event/view.rs @@ -19,16 +19,16 @@ use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::event_table::{ CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, - VIEW_NAME_COLUMN, + VIEW_NAME_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; use common_event_recorder::{Event, EventTypeFilter}; use common_procedure::{ - ChildSubmissionOutcome, EventRuntimeContext, EventTrigger, Output, Procedure, ProcedureEvent, + ChildSubmissionOutcome, EventContext, EventTrigger, Output, Procedure, ProcedureEvent, ProcedureId, ProcedureState, RetryPhase, }; -use super::test_util::{assert_event_filter, default_event_context_value, procedure_trigger_value}; +use super::test_util::assert_event_filter; use crate::ddl::create_view::CreateViewProcedure; use crate::ddl::drop_view::DropViewProcedure; use crate::ddl::event::view::{ @@ -36,7 +36,6 @@ use crate::ddl::event::view::{ }; use crate::ddl::tests::create_view::test_create_view_task; use crate::ddl::tests::drop_view::new_drop_view_task; -use crate::rpc::ddl::EventContext; use crate::test_util::{MockDatanodeManager, new_ddl_context}; #[test] @@ -44,7 +43,7 @@ fn test_view_submitted_event_contracts() { let mut task = test_create_view_task("v_metrics"); task.create_view.or_replace = true; task.create_view.create_if_not_exists = true; - let create = CreateViewProcedure::new(task, EventContext::default(), test_context()); + let create = CreateViewProcedure::new(task, test_context()); let event = event_for(&create, EventTrigger::Submitted); assert_view_event_contract( @@ -72,11 +71,7 @@ fn test_view_submitted_event_contracts() { assert!(!payload.contains(omitted)); } - let drop = DropViewProcedure::new( - new_drop_view_task("view_name", 42, true), - EventContext::default(), - test_context(), - ); + let drop = DropViewProcedure::new(new_drop_view_task("view_name", 42, true), test_context()); let event = event_for(&drop, EventTrigger::Submitted); assert_view_event_contract( @@ -140,16 +135,8 @@ fn test_view_lifecycle_event_contracts() { #[test] fn test_view_procedures_preserve_lifecycle_locators() { - let create = CreateViewProcedure::new( - test_create_view_task("view_name"), - EventContext::default(), - test_context(), - ); - let drop = DropViewProcedure::new( - new_drop_view_task("view_name", 42, false), - EventContext::default(), - test_context(), - ); + let create = CreateViewProcedure::new(test_create_view_task("view_name"), test_context()); + let drop = DropViewProcedure::new(new_drop_view_task("view_name", 42, false), test_context()); let triggers = [ EventTrigger::Recovered, EventTrigger::ChildSubmitted { @@ -199,11 +186,7 @@ fn test_view_procedures_preserve_lifecycle_locators() { #[test] fn test_create_view_succeeded_output_mapping() { - let procedure = CreateViewProcedure::new( - test_create_view_task("view_name"), - EventContext::default(), - test_context(), - ); + let procedure = CreateViewProcedure::new(test_create_view_task("view_name"), test_context()); let state = ProcedureState::Done { output: Some(Arc::new(84_u32)), }; @@ -240,21 +223,14 @@ fn test_create_view_succeeded_output_mapping() { #[test] fn test_create_view_event_filter() { - let procedure = CreateViewProcedure::new( - test_create_view_task("view_name"), - EventContext::default(), - test_context(), - ); + let procedure = CreateViewProcedure::new(test_create_view_task("view_name"), test_context()); assert_event_filter(&procedure, CREATE_VIEW_EVENT_TYPE); } #[test] fn test_drop_view_event_filter() { - let procedure = DropViewProcedure::new( - new_drop_view_task("view_name", 42, false), - EventContext::default(), - test_context(), - ); + let procedure = + DropViewProcedure::new(new_drop_view_task("view_name", 42, false), test_context()); assert_event_filter(&procedure, DROP_VIEW_EVENT_TYPE); } @@ -273,7 +249,6 @@ fn test_view_event_procedure_envelope_contract() { referenced_table_count: 1, column_count: 1, }, - EventContext::default(), )), ProcedureState::Running, EventTrigger::Submitted, @@ -344,17 +319,11 @@ fn view_schema() -> Vec { SCHEMA_NAME_COLUMN.column_schema(), VIEW_NAME_COLUMN.column_schema(), VIEW_ID_COLUMN.column_schema(), - EVENT_CONTEXT_COLUMN.column_schema(), ] } fn assert_view_event_contract(event: &dyn Event, event_type: &str, locator: ViewEventLocator<'_>) { - let mut values = locator.values(); - values.push(if event.json_payload().unwrap().is_null() { - Value { value_data: None } - } else { - default_event_context_value() - }); + let values = locator.values(); assert_event_contract(event, event_type, &view_schema(), &[Row { values }]); } @@ -372,19 +341,16 @@ fn assert_procedure_event_contract( PROCEDURE_TRIGGER_COLUMN.column_schema(), ]; schema.extend(view_schema()); + schema.push(EVENT_CONTEXT_COLUMN.column_schema()); let mut values = vec![ ValueData::StringValue(event.procedure_id.to_string()).into(), ValueData::StringValue(state.to_string()).into(), ValueData::StringValue(String::new()).into(), - procedure_trigger_value(trigger), + jsonb_value(&serde_json::json!({"type": trigger})), ]; values.extend(locator.values()); - values.push(if trigger == "Submitted" { - default_event_context_value() - } else { - Value { value_data: None } - }); + values.push(Value { value_data: None }); assert_event_contract(event, event_type, &schema, &[Row { values }]); } @@ -403,11 +369,12 @@ fn event_for_state( lifecycle_state: &ProcedureState, ) -> Box { procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, event_type_filter: Arc::new(EventTypeFilter::All), + event_context: None, }) .unwrap() } diff --git a/src/common/meta/src/ddl/truncate_table.rs b/src/common/meta/src/ddl/truncate_table.rs index 6a561b62afb..67d74361983 100644 --- a/src/common/meta/src/ddl/truncate_table.rs +++ b/src/common/meta/src/ddl/truncate_table.rs @@ -20,7 +20,7 @@ use api::v1::region::{ use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::debug; @@ -43,7 +43,7 @@ use crate::key::table_info::TableInfoValue; use crate::key::table_name::TableNameKey; use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; use crate::metrics; -use crate::rpc::ddl::{EventContext, TruncateTableTask}; +use crate::rpc::ddl::TruncateTableTask; use crate::rpc::router::{find_leader_regions, find_leaders}; pub struct TruncateTableProcedure { @@ -89,10 +89,7 @@ impl Procedure for TruncateTableProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::TruncateTable.as_str()) @@ -103,11 +100,9 @@ impl Procedure for TruncateTableProcedure { let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) .with_table_id(task.table_id); let event = match &ctx.trigger { - EventTrigger::Submitted => TableDdlEvent::truncate_table_submitted( - locator, - task.time_ranges.len(), - self.data.event_context.clone(), - ), + EventTrigger::Submitted => { + TableDdlEvent::truncate_table_submitted(locator, task.time_ranges.len()) + } _ => TableDdlEvent::lifecycle(TableDdlEventType::TruncateTable, [locator]), }; @@ -122,11 +117,10 @@ impl TruncateTableProcedure { task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, context: DdlContext, - event_context: EventContext, ) -> Self { Self { context, - data: TruncateTableData::new(task, table_info_value, event_context), + data: TruncateTableData::new(task, table_info_value), } } @@ -233,21 +227,17 @@ pub struct TruncateTableData { state: TruncateTableState, task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - #[serde(default)] - event_context: EventContext, } impl TruncateTableData { pub fn new( task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - event_context: EventContext, ) -> Self { Self { state: TruncateTableState::Prepare, task, table_info_value, - event_context, } } diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs index 3295b1163e2..73a361ffd0e 100644 --- a/src/common/meta/src/ddl/undrop_table.rs +++ b/src/common/meta/src/ddl/undrop_table.rs @@ -17,7 +17,7 @@ use api::v1::region::{ use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::tracing_context::TracingContext; @@ -43,7 +43,7 @@ use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; -use crate::rpc::ddl::{EventContext, UndropTableTask}; +use crate::rpc::ddl::UndropTableTask; use crate::rpc::router::{ RegionRoute, find_follower_regions, find_followers, find_leader_regions, find_leaders, }; @@ -56,17 +56,16 @@ pub struct UndropTableProcedure { impl UndropTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::UndropTable"; - pub fn new(task: UndropTableTask, context: DdlContext, event_context: EventContext) -> Self { - Self::new_with_original_table_name(task, context, None, event_context) + pub fn new(task: UndropTableTask, context: DdlContext) -> Self { + Self::new_with_original_table_name(task, context, None) } pub(crate) fn new_with_original_table_name( task: UndropTableTask, context: DdlContext, table_name: Option, - event_context: EventContext, ) -> Self { - let mut data = UndropTableData::new(task, event_context); + let mut data = UndropTableData::new(task); data.table_name = table_name; Self { context, data } } @@ -444,10 +443,7 @@ impl Procedure for UndropTableProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::UndropTable.as_str()) @@ -456,9 +452,7 @@ impl Procedure for UndropTableProcedure { } let locator = self.event_locator(); let event = match &ctx.trigger { - EventTrigger::Submitted => { - TableDdlEvent::undrop_table_submitted(locator, self.data.event_context.clone()) - } + EventTrigger::Submitted => TableDdlEvent::undrop_table_submitted(locator), _ => TableDdlEvent::lifecycle(TableDdlEventType::UndropTable, [locator]), }; @@ -565,12 +559,10 @@ pub struct UndropTableData { drop_generation: Option, #[serde(default)] tombstone_identity_loaded: bool, - #[serde(default)] - event_context: EventContext, } impl UndropTableData { - fn new(task: UndropTableTask, event_context: EventContext) -> Self { + fn new(task: UndropTableTask) -> Self { Self { state: UndropTableState::Prepare, task, @@ -582,7 +574,6 @@ impl UndropTableData { retention_expires_at: None, drop_generation: None, tombstone_identity_loaded: false, - event_context, } } @@ -647,11 +638,7 @@ mod tests { #[test] fn test_map_restore_metadata_error_without_table_name() { let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); - let procedure = UndropTableProcedure::new( - UndropTableTask { table_id: 42 }, - context, - EventContext::default(), - ); + let procedure = UndropTableProcedure::new(UndropTableTask { table_id: 42 }, context); let err = procedure.map_restore_metadata_error( error::TombstoneTargetAlreadyExistsSnafu { @@ -666,11 +653,8 @@ mod tests { #[test] fn test_recovered_legacy_snapshot_reloads_tombstone_identity() { let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); - let mut procedure = UndropTableProcedure::new( - UndropTableTask { table_id: 42 }, - context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(UndropTableTask { table_id: 42 }, context.clone()); procedure.data.state = UndropTableState::OpenRegions; let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); let data = data.as_object_mut().unwrap(); @@ -695,11 +679,7 @@ mod tests { let table_id = 1024; let region_id = RegionId::new(table_id, 1); - let mut procedure = UndropTableProcedure::new( - UndropTableTask { table_id }, - context, - EventContext::default(), - ); + let mut procedure = UndropTableProcedure::new(UndropTableTask { table_id }, context); procedure.data.table_route_value = Some(TableRouteValue::physical(vec![RegionRoute { region: Region::new_test(region_id), leader_peer: Some(Peer::empty(1)), @@ -738,11 +718,7 @@ mod tests { let table_id = 1024; let region_id = RegionId::new(table_id, 1); - let mut procedure = UndropTableProcedure::new( - UndropTableTask { table_id }, - context, - EventContext::default(), - ); + let mut procedure = UndropTableProcedure::new(UndropTableTask { table_id }, context); procedure.data.table_name = Some(TableName::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 137a16793d7..2eed4628ac3 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -78,8 +78,8 @@ use crate::rpc::ddl::trigger::DropTriggerTask; use crate::rpc::ddl::{ AlterDatabaseTask, AlterTableTask, CommentOnTask, CreateDatabaseTask, CreateFlowTask, CreateTableTask, CreateViewTask, DropDatabaseTask, DropFlowTask, DropTableTask, DropViewTask, - EventContext, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse, - TruncateTableTask, UndropTableTask, + PersistentEventContext, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, + SubmitDdlTaskResponse, TruncateTableTask, UndropTableTask, }; const MAX_REGION_ROUTE_CHANGE_RETRIES: usize = 3; @@ -188,7 +188,6 @@ pub trait RepartitionProcedureFactory: Send + Sync { source: RepartitionSource, to_exprs: Vec, timeout: Option, - event_context: EventContext, ) -> std::result::Result; fn register_loaders( @@ -342,7 +341,7 @@ impl DdlManager { repartition: Repartition, wait: bool, timeout: Duration, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -378,14 +377,14 @@ impl DdlManager { source, into_partition_exprs, Some(timeout), - event_context, ) .context(CreateRepartitionProcedureSnafu)?; self.repartition_procedure_factory .ensure_gc_requirement() .await .context(PersistRepartitionGcRequirementSnafu)?; - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); if wait { self.execute_procedure_and_wait(procedure_with_id).await } else { @@ -401,7 +400,7 @@ impl DdlManager { &self, table_id: TableId, alter_table_task: AlterTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ddl_options: DdlOptions, ) -> Result<(ProcedureId, Option)> { // make alter_table_task mutable so we can call .take() on its field @@ -456,12 +455,12 @@ impl DdlManager { let procedure = AlterTableProcedure::new_with_region_locks( table_id, alter_table_task.clone(), - event_context.clone(), region_ids_to_lock, context, )?; - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) + .with_event_context(event_context.clone()); let result = self.execute_procedure_and_wait(procedure_with_id).await?; if result .1 @@ -490,14 +489,18 @@ impl DdlManager { &self, create_table_task: CreateTableTask, query_context: QueryContext, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = - CreateTableProcedure::new(create_table_task, query_context, event_context, context)?; + let procedure = CreateTableProcedure::new_with_query_context( + create_table_task, + query_context, + context, + )?; - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -507,13 +510,14 @@ impl DdlManager { pub async fn submit_create_view_task( &self, create_view_task: CreateViewTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = CreateViewProcedure::new(create_view_task, event_context, context); + let procedure = CreateViewProcedure::new(create_view_task, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -524,18 +528,15 @@ impl DdlManager { &self, create_table_tasks: Vec, physical_table_id: TableId, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = CreateLogicalTablesProcedure::new( - create_table_tasks, - physical_table_id, - event_context, - context, - ); + let procedure = + CreateLogicalTablesProcedure::new(create_table_tasks, physical_table_id, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -546,18 +547,15 @@ impl DdlManager { &self, alter_table_tasks: Vec, physical_table_id: TableId, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = AlterLogicalTablesProcedure::new( - alter_table_tasks, - physical_table_id, - event_context, - context, - ); + let procedure = + AlterLogicalTablesProcedure::new(alter_table_tasks, physical_table_id, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -567,13 +565,14 @@ impl DdlManager { pub async fn submit_drop_table_task( &self, drop_table_task: DropTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = DropTableProcedure::new(drop_table_task, context, event_context); + let procedure = DropTableProcedure::new(drop_table_task, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -584,7 +583,7 @@ impl DdlManager { pub async fn submit_undrop_table_task( &self, undrop_table_task: UndropTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -610,9 +609,9 @@ impl DdlManager { undrop_table_task, context, Some(original_table_name), - event_context, ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) + .with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -624,7 +623,7 @@ impl DdlManager { pub async fn submit_purge_dropped_table_task( &self, purge_dropped_table_task: PurgeDroppedTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -638,9 +637,9 @@ impl DdlManager { #[cfg(feature = "enterprise")] { let context = self.create_context(); - let procedure = - PurgeDroppedTableProcedure::new(purge_dropped_table_task, context, event_context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = PurgeDroppedTableProcedure::new(purge_dropped_table_task, context); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) + .with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -666,12 +665,11 @@ impl DdlManager { #[cfg(feature = "enterprise")] { let context = self.create_context(); - let procedure = PurgeDroppedTableProcedure::new_if_expired( - purge_dropped_table_task, - context, - EventContext::new(TriggerReason::ScheduledGc), - ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let event_context = PersistentEventContext::new(TriggerReason::ScheduledGc); + let procedure = + PurgeDroppedTableProcedure::new_if_expired(purge_dropped_table_task, context); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) + .with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -688,7 +686,7 @@ impl DdlManager { options, creator, }: CreateDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateDatabaseProcedure::new( @@ -697,10 +695,10 @@ impl DdlManager { create_if_not_exists, options, creator, - event_context, context, ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -714,12 +712,12 @@ impl DdlManager { schema, drop_if_exists, }: DropDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = - DropDatabaseProcedure::new(catalog, schema, drop_if_exists, event_context, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = DropDatabaseProcedure::new(catalog, schema, drop_if_exists, context); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -727,11 +725,12 @@ impl DdlManager { pub async fn submit_alter_database( &self, alter_database_task: AlterDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = AlterDatabaseProcedure::new(alter_database_task, event_context, context)?; - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = AlterDatabaseProcedure::new(alter_database_task, context)?; + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -742,12 +741,12 @@ impl DdlManager { &self, create_flow: CreateFlowTask, query_context: QueryContext, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = - CreateFlowProcedure::new(create_flow, query_context, event_context, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = CreateFlowProcedure::new(create_flow, query_context, context); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -757,11 +756,12 @@ impl DdlManager { pub async fn submit_drop_flow_task( &self, drop_flow: DropFlowTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = DropFlowProcedure::new(drop_flow, event_context, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = DropFlowProcedure::new(drop_flow, context); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -771,11 +771,12 @@ impl DdlManager { pub async fn submit_drop_view_task( &self, drop_view: DropViewTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = DropViewProcedure::new(drop_view, event_context, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = DropViewProcedure::new(drop_view, context); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -786,17 +787,13 @@ impl DdlManager { &self, truncate_table_task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = TruncateTableProcedure::new( - truncate_table_task, - table_info_value, - context, - event_context, - ); + let procedure = TruncateTableProcedure::new(truncate_table_task, table_info_value, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -806,6 +803,7 @@ impl DdlManager { pub async fn submit_comment_on_task( &self, mut comment_on_task: CommentOnTask, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); comment_on_task @@ -815,7 +813,8 @@ impl DdlManager { ) .await?; let procedure = CommentOnProcedure::new(comment_on_task, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -869,7 +868,10 @@ impl DdlManager { timeout, task, } = request; - let event_context = EventContext::from_query_context(&query_context); + let event_context = ctx + .event_context + .clone() + .unwrap_or_else(|| crate::rpc::ddl::event_context_from_query_context(&query_context)); let ddl_options = DdlOptions { wait, timeout }; async move { debug!("Submitting Ddl task: {:?}", task); @@ -924,7 +926,9 @@ impl DdlManager { DropView(drop_view_task) => { handle_drop_view_task(self, drop_view_task, event_context).await } - CommentOn(comment_on_task) => handle_comment_on_task(self, comment_on_task).await, + CommentOn(comment_on_task) => { + handle_comment_on_task(self, comment_on_task, event_context).await + } #[cfg(feature = "enterprise")] CreateTrigger(create_trigger_task) => { handle_create_trigger_task(self, create_trigger_task, query_context).await @@ -943,7 +947,7 @@ impl DdlManager { async fn handle_truncate_table_task( ddl_manager: &DdlManager, truncate_table_task: TruncateTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let table_id = truncate_table_task.table_id; let table_metadata_manager = &ddl_manager.table_metadata_manager(); @@ -983,7 +987,7 @@ async fn handle_alter_table_task( ddl_manager: &DdlManager, alter_table_task: AlterTableTask, ddl_options: DdlOptions, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let table_ref = alter_table_task.table_ref(); @@ -1030,7 +1034,7 @@ async fn handle_alter_table_task( async fn handle_drop_table_task( ddl_manager: &DdlManager, drop_table_task: DropTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let table_id = drop_table_task.table_id; let (id, _) = ddl_manager @@ -1048,7 +1052,7 @@ async fn handle_drop_table_task( async fn handle_undrop_table_task( ddl_manager: &DdlManager, undrop_table_task: UndropTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let table_id = undrop_table_task.table_id; let (id, _) = ddl_manager @@ -1066,7 +1070,7 @@ async fn handle_undrop_table_task( async fn handle_purge_dropped_table_task( ddl_manager: &DdlManager, purge_dropped_table_task: PurgeDroppedTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager .submit_purge_dropped_table_task(purge_dropped_table_task, event_context) @@ -1084,7 +1088,7 @@ async fn handle_create_table_task( ddl_manager: &DdlManager, create_table_task: CreateTableTask, query_context: QueryContext, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, output) = ddl_manager .submit_create_table_task(create_table_task, query_context, event_context) @@ -1110,7 +1114,7 @@ async fn handle_create_table_task( async fn handle_create_logical_table_tasks( ddl_manager: &DdlManager, create_table_tasks: Vec, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { ensure!( !create_table_tasks.is_empty(), @@ -1155,7 +1159,7 @@ async fn handle_create_logical_table_tasks( async fn handle_create_database_task( ddl_manager: &DdlManager, create_database_task: CreateDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let catalog = create_database_task.catalog.clone(); let schema = create_database_task.schema.clone(); @@ -1178,7 +1182,7 @@ async fn handle_create_database_task( async fn handle_drop_database_task( ddl_manager: &DdlManager, drop_database_task: DropDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager .submit_drop_database(drop_database_task.clone(), event_context) @@ -1199,7 +1203,7 @@ async fn handle_drop_database_task( async fn handle_alter_database_task( ddl_manager: &DdlManager, alter_database_task: AlterDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager .submit_alter_database(alter_database_task.clone(), event_context) @@ -1221,7 +1225,7 @@ async fn handle_alter_database_task( async fn handle_drop_flow_task( ddl_manager: &DdlManager, drop_flow_task: DropFlowTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager .submit_drop_flow_task(drop_flow_task.clone(), event_context) @@ -1266,7 +1270,7 @@ async fn handle_drop_trigger_task( async fn handle_drop_view_task( ddl_manager: &DdlManager, drop_view_task: DropViewTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager .submit_drop_view_task(drop_view_task.clone(), event_context) @@ -1289,7 +1293,7 @@ async fn handle_create_flow_task( ddl_manager: &DdlManager, create_flow_task: CreateFlowTask, query_context: QueryContext, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, output) = ddl_manager .submit_create_flow_task(create_flow_task.clone(), query_context, event_context) @@ -1349,7 +1353,7 @@ async fn handle_create_trigger_task( async fn handle_alter_logical_table_tasks( ddl_manager: &DdlManager, alter_table_tasks: Vec, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { ensure!( !alter_table_tasks.is_empty(), @@ -1388,7 +1392,7 @@ async fn handle_alter_logical_table_tasks( async fn handle_create_view_task( ddl_manager: &DdlManager, create_view_task: CreateViewTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, output) = ddl_manager .submit_create_view_task(create_view_task, event_context) @@ -1414,9 +1418,10 @@ async fn handle_create_view_task( async fn handle_comment_on_task( ddl_manager: &DdlManager, comment_on_task: CommentOnTask, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager - .submit_comment_on_task(comment_on_task.clone()) + .submit_comment_on_task(comment_on_task.clone(), event_context) .await?; let procedure_id = id.to_string(); @@ -1469,7 +1474,7 @@ mod tests { use crate::procedure_executor::ExecutorContext; use crate::region_keeper::MemoryRegionKeeper; use crate::region_registry::LeaderRegionRegistry; - use crate::rpc::ddl::{CreatorGrantIntent, EventContext, UndropTableTask}; + use crate::rpc::ddl::{CreatorGrantIntent, PersistentEventContext, UndropTableTask}; #[cfg(not(feature = "enterprise"))] use crate::rpc::ddl::{DdlTask, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest}; use crate::sequence::SequenceBuilder; @@ -1506,7 +1511,6 @@ mod tests { _source: RepartitionSource, _to_exprs: Vec, _timeout: Option, - _event_context: EventContext, ) -> std::result::Result { unimplemented!() } @@ -1680,7 +1684,10 @@ mod tests { let ddl_manager = build_soft_drop_test_ddl_manager(); let err = ddl_manager - .submit_undrop_table_task(UndropTableTask { table_id: 1024 }, EventContext::default()) + .submit_undrop_table_task( + UndropTableTask { table_id: 1024 }, + PersistentEventContext::default(), + ) .await .unwrap_err(); @@ -1694,7 +1701,10 @@ mod tests { let ddl_manager = build_soft_drop_test_ddl_manager(); let err = ddl_manager - .submit_undrop_table_task(UndropTableTask { table_id: 1024 }, EventContext::default()) + .submit_undrop_table_task( + UndropTableTask { table_id: 1024 }, + PersistentEventContext::default(), + ) .await .unwrap_err(); assert!(matches!(err, crate::error::Error::Unsupported { .. })); @@ -1702,7 +1712,7 @@ mod tests { let err = ddl_manager .submit_purge_dropped_table_task( PurgeDroppedTableTask { table_id: 1024 }, - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap_err(); diff --git a/src/common/meta/src/procedure_executor.rs b/src/common/meta/src/procedure_executor.rs index c0165dd80f4..b812aeb6aac 100644 --- a/src/common/meta/src/procedure_executor.rs +++ b/src/common/meta/src/procedure_executor.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse}; +use common_event_recorder::PersistentEventContext; use common_procedure::{ProcedureId, ProcedureManagerRef}; use common_telemetry::tracing_context::W3cTrace; use snafu::{OptionExt, ResultExt}; @@ -30,9 +31,11 @@ use crate::rpc::procedure::{ }; /// The context of procedure executor. -#[derive(Debug, Default)] +#[derive(Debug, Clone, Default)] pub struct ExecutorContext { pub tracing_context: Option, + pub actor: Option, + pub event_context: Option, } /// The procedure executor that accepts ddl, region migration task etc. diff --git a/src/common/meta/src/rpc/ddl.rs b/src/common/meta/src/rpc/ddl.rs index 8aef3325e64..d7966804738 100644 --- a/src/common/meta/src/rpc/ddl.rs +++ b/src/common/meta/src/rpc/ddl.rs @@ -45,12 +45,12 @@ use base64::engine::general_purpose; use common_base::protocol::Channel; use common_catalog::{format_full_flow_name, format_full_table_name}; use common_error::ext::BoxedError; +pub use common_event_recorder::{PersistentEventContext, TriggerReason}; use common_time::{DatabaseTimeToLive, Timestamp}; use prost::Message; use serde::{Deserialize, Serialize}; use serde_with::{DefaultOnNull, serde_as}; use snafu::{OptionExt, ResultExt}; -use strum::{AsRefStr, EnumString}; use table::metadata::{TableId, TableInfo}; use table::requests::validate_database_option; use table::table_name::TableName; @@ -1668,77 +1668,19 @@ pub struct QueryContext { pub sst_min_sequences: HashMap, } -/// The stable context recorded for a procedure event. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct EventContext { - pub reason: TriggerReason, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub protocol: Option, - #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")] - pub extensions: serde_json::Map, -} - -impl EventContext { - /// Creates an event context with no additional extensions. - pub fn new(reason: TriggerReason) -> Self { - Self { - reason, - protocol: None, - extensions: Default::default(), - } - } - - /// Adds the protocol that originated the operation. - pub fn with_protocol(mut self, protocol: impl Into) -> Self { - self.protocol = Some(protocol.into()); - self - } - - /// Builds an event context from frontend query context metadata. - pub fn from_query_context(query_context: &QueryContext) -> Self { - let reason = query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) - .map(|reason| TriggerReason::from_extension(reason)) - .unwrap_or_default(); - let channel = Channel::from(u32::from(query_context.channel)); - let context = Self::new(reason); - if channel == Channel::Unknown { - context - } else { - context.with_protocol(channel.as_ref()) - } - } -} - -impl Default for EventContext { - fn default() -> Self { - Self::new(TriggerReason::default()) - } -} - -/// The stable classification of a procedure trigger. -#[derive( - Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, AsRefStr, EnumString, -)] -#[serde(rename_all = "snake_case")] -#[strum(serialize_all = "snake_case")] -pub enum TriggerReason { - Manual, - AutoCreate, - AutoAlter, - AutoRepartition, - AutoRebalance, - RegionFailover, - ScheduledGc, - #[default] - #[serde(other)] - Unknown, -} - -impl TriggerReason { - pub fn from_extension(value: &str) -> Self { - value.parse().unwrap_or_default() +/// Builds event context from frontend query context metadata. +pub fn event_context_from_query_context(query_context: &QueryContext) -> PersistentEventContext { + let reason = query_context + .extensions + .get(TRIGGER_REASON_EXTENSION_KEY) + .map(|reason| TriggerReason::from_extension(reason)) + .unwrap_or_default(); + let channel = Channel::from(u32::from(query_context.channel)); + let context = PersistentEventContext::new(reason); + if channel == Channel::Unknown { + context + } else { + context.with_protocol(channel.as_ref()) } } @@ -2323,7 +2265,7 @@ mod tests { #[test] fn test_event_context_serialization() { - let context = EventContext::new(TriggerReason::Manual).with_protocol("mysql"); + let context = PersistentEventContext::new(TriggerReason::Manual).with_protocol("mysql"); assert_eq!( serde_json::json!({ @@ -2335,7 +2277,7 @@ mod tests { assert_eq!( serde_json::json!({ "reason": "manual" }), - serde_json::to_value(EventContext::new(TriggerReason::Manual)).unwrap() + serde_json::to_value(PersistentEventContext::new(TriggerReason::Manual)).unwrap() ); } @@ -2348,7 +2290,7 @@ mod tests { ); query_context.channel = 4; - let event_context = EventContext::from_query_context(&query_context); + let event_context = event_context_from_query_context(&query_context); assert_eq!(event_context.reason, TriggerReason::AutoCreate); assert_eq!(event_context.protocol.as_deref(), Some("prometheus")); diff --git a/src/common/procedure-test/src/lib.rs b/src/common/procedure-test/src/lib.rs index f6d816381d1..2d342ed91ab 100644 --- a/src/common/procedure-test/src/lib.rs +++ b/src/common/procedure-test/src/lib.rs @@ -84,6 +84,7 @@ pub async fn execute_procedure_until_done(procedure: &mut dyn Procedure) -> Opti let ctx = Context { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; loop { @@ -110,6 +111,7 @@ pub async fn execute_procedure_once( let ctx = Context { procedure_id, provider: Arc::new(provider), + event_context: None, }; match procedure.execute(&ctx).await.unwrap() { @@ -137,6 +139,7 @@ pub async fn execute_until_suspended_or_done( let ctx = Context { procedure_id, provider: Arc::new(provider), + event_context: None, }; loop { @@ -155,6 +158,7 @@ pub fn new_test_procedure_context() -> Context { Context { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, } } diff --git a/src/common/procedure/src/event.rs b/src/common/procedure/src/event.rs index c3100b34cdd..a4a6c290bf3 100644 --- a/src/common/procedure/src/event.rs +++ b/src/common/procedure/src/event.rs @@ -19,13 +19,13 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, - jsonb_value, procedure_event_column_schemas, + EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, + PROCEDURE_TRIGGER_COLUMN, jsonb_value, nullable_json, procedure_event_column_schemas, }; use common_time::timestamp::{TimeUnit, Timestamp}; use snafu::ResultExt; -use crate::{EventTrigger, ProcedureId, ProcedureState}; +use crate::{EventTrigger, ProcedureContext, ProcedureId, ProcedureState}; pub const EVENTS_TABLE_PROCEDURE_ID_COLUMN_NAME: &str = PROCEDURE_ID_COLUMN.name(); pub const EVENTS_TABLE_PROCEDURE_STATE_COLUMN_NAME: &str = PROCEDURE_STATE_COLUMN.name(); @@ -43,6 +43,8 @@ pub struct ProcedureEvent { pub state: ProcedureState, /// The lifecycle trigger that caused the event to be emitted. pub trigger: EventTrigger, + /// Context associated with the root submission. + pub context: ProcedureContext, /// The event emitted by the procedure. It's generated by [Procedure::event]. pub internal_event: Box, } @@ -53,6 +55,22 @@ impl ProcedureEvent { internal_event: Box, state: ProcedureState, trigger: EventTrigger, + ) -> Self { + Self::new_with_context( + procedure_id, + internal_event, + state, + trigger, + ProcedureContext::default(), + ) + } + + pub fn new_with_context( + procedure_id: ProcedureId, + internal_event: Box, + state: ProcedureState, + trigger: EventTrigger, + context: ProcedureContext, ) -> Self { Self { procedure_id, @@ -60,6 +78,7 @@ impl ProcedureEvent { timestamp: Timestamp::current_time(TimeUnit::Nanosecond), state, trigger, + context, } } } @@ -79,7 +98,9 @@ impl Event for ProcedureEvent { fn extra_schema(&self) -> Vec { let mut schema = procedure_event_column_schemas(); - schema.append(&mut self.internal_event.extra_schema()); + let mut internal_schema = self.internal_event.extra_schema(); + schema.append(&mut internal_schema); + schema.push(EVENT_CONTEXT_COLUMN.column_schema()); schema } @@ -97,9 +118,15 @@ impl Event for ProcedureEvent { _ => String::new(), }; let trigger = serde_json::to_value(&self.trigger).context(SerializeEventSnafu)?; + let event_context = matches!(self.trigger, EventTrigger::Submitted) + .then_some(self.context.event_context.as_ref()) + .flatten() + .map(serde_json::to_value) + .transpose() + .context(SerializeEventSnafu)?; for internal_event_extra_row in internal_event_extra_rows.iter_mut() { - let mut values = Vec::with_capacity(4 + internal_event_extra_row.values.len()); + let mut values = Vec::with_capacity(5 + internal_event_extra_row.values.len()); values.extend([ ValueData::StringValue(procedure_id.clone()).into(), ValueData::StringValue(state.clone()).into(), @@ -107,6 +134,7 @@ impl Event for ProcedureEvent { jsonb_value(&trigger), ]); values.append(&mut internal_event_extra_row.values); + values.push(nullable_json(event_context.as_ref())); rows.push(Row { values }); } @@ -126,13 +154,15 @@ mod tests { use api::v1::{ColumnDataType, ColumnSchema, Row, SemanticType, Value}; use common_error::mock::MockError; use common_error::status_code::StatusCode; - use common_event_recorder::Event; - use common_event_recorder::event_table::{PROCEDURE_TRIGGER_COLUMN, jsonb_value}; + use common_event_recorder::event_table::{ + EVENT_CONTEXT_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, + }; + use common_event_recorder::{Event, PersistentEventContext, TriggerReason}; use serde_json::json; use crate::{ - ChildSubmissionOutcome, Error, EventTrigger, ProcedureEvent, ProcedureId, ProcedureState, - RetryPhase, + ChildSubmissionOutcome, Error, EventTrigger, ProcedureContext, ProcedureEvent, ProcedureId, + ProcedureState, RetryPhase, }; #[derive(Debug)] @@ -192,6 +222,7 @@ mod tests { ValueData::StringValue(String::new()).into(), jsonb_value(&json!({"type": "Submitted"})), ValueData::StringValue("test_event1".to_string()).into(), + Value { value_data: None }, ], }, Row { @@ -201,6 +232,7 @@ mod tests { ValueData::StringValue(String::new()).into(), jsonb_value(&json!({"type": "Submitted"})), ValueData::StringValue("test_event2".to_string()).into(), + Value { value_data: None }, ], }, Row { @@ -210,6 +242,7 @@ mod tests { ValueData::StringValue(String::new()).into(), jsonb_value(&json!({"type": "Submitted"})), Value { value_data: None }, + Value { value_data: None }, ], }, ] @@ -280,10 +313,41 @@ mod tests { semantic_type: SemanticType::Field.into(), ..Default::default() }, + EVENT_CONTEXT_COLUMN.column_schema(), ] ); } + #[test] + fn procedure_event_only_records_context_on_submission() { + let context = ProcedureContext::from_event_context(PersistentEventContext::new( + TriggerReason::AutoRepartition, + )); + let submitted = ProcedureEvent::new_with_context( + ProcedureId::random(), + Box::new(TestEvent {}), + ProcedureState::Running, + EventTrigger::Submitted, + context.clone(), + ); + let succeeded = ProcedureEvent::new_with_context( + ProcedureId::random(), + Box::new(TestEvent {}), + ProcedureState::Done { output: None }, + EventTrigger::Succeeded, + context, + ); + + assert_eq!( + submitted.extra_rows().unwrap()[0].values[5], + jsonb_value(&json!({"reason": "auto_repartition"})) + ); + assert_eq!( + succeeded.extra_rows().unwrap()[0].values[5], + Value { value_data: None } + ); + } + #[test] fn test_event_trigger_serialization() { let procedure_id = ProcedureId::parse_str("00000000-0000-0000-0000-000000000001").unwrap(); diff --git a/src/common/procedure/src/lib.rs b/src/common/procedure/src/lib.rs index 1d7a200ca58..a21e5d1c47a 100644 --- a/src/common/procedure/src/lib.rs +++ b/src/common/procedure/src/lib.rs @@ -30,8 +30,8 @@ pub use crate::error::{Error, Result}; pub use crate::event::ProcedureEvent; pub use crate::procedure::{ BoxedProcedure, BoxedProcedureLoader, ChildSubmissionOutcome, Context, ContextProvider, - ContextProviderRef, EventRuntimeContext, EventTrigger, LockKey, Output, ParseIdError, - PoisonKey, PoisonKeys, Procedure, ProcedureId, ProcedureInfo, ProcedureManager, + ContextProviderRef, EventContext, EventTrigger, LockKey, Output, ParseIdError, PoisonKey, + PoisonKeys, Procedure, ProcedureContext, ProcedureId, ProcedureInfo, ProcedureManager, ProcedureManagerRef, ProcedureState, ProcedureWithId, RetryPhase, Status, StringKey, }; pub use crate::watcher::Watcher; diff --git a/src/common/procedure/src/local.rs b/src/common/procedure/src/local.rs index 59eaee2d789..8fdd3167583 100644 --- a/src/common/procedure/src/local.rs +++ b/src/common/procedure/src/local.rs @@ -43,8 +43,8 @@ use crate::rwlock::{KeyRwLock, OwnedKeyRwLockGuard}; use crate::store::poison_store::PoisonStoreRef; use crate::store::{ProcedureMessage, ProcedureMessages, ProcedureStore, StateStoreRef}; use crate::{ - BoxedProcedure, ContextProvider, EventTrigger, LockKey, PoisonKey, ProcedureId, - ProcedureManager, ProcedureState, ProcedureWithId, StringKey, Watcher, + BoxedProcedure, ContextProvider, EventTrigger, LockKey, PoisonKey, ProcedureContext, + ProcedureId, ProcedureManager, ProcedureState, ProcedureWithId, StringKey, Watcher, }; /// The expired time of a procedure's metadata. @@ -81,6 +81,8 @@ pub(crate) struct ProcedureMeta { type_name: String, /// Parent procedure id. parent_id: Option, + /// Context associated with the root submission. + context: ProcedureContext, /// Notify to wait for subprocedures. child_notify: Notify, /// Lock required by this procedure. @@ -105,6 +107,7 @@ impl ProcedureMeta { id: ProcedureId, procedure_state: ProcedureState, parent_id: Option, + context: ProcedureContext, lock_key: LockKey, poison_keys: PoisonKeys, type_name: &str, @@ -113,6 +116,7 @@ impl ProcedureMeta { ProcedureMeta { id, parent_id, + context, child_notify: Notify::new(), lock_key, poison_keys, @@ -716,6 +720,7 @@ impl LocalManager { procedure_state: ProcedureState, step: u32, procedure: BoxedProcedure, + context: ProcedureContext, origin: RootSubmissionOrigin, ) -> Result { ensure!(self.manager_ctx.running(), ManagerNotStartSnafu); @@ -724,6 +729,7 @@ impl LocalManager { procedure_id, procedure_state, None, + context, procedure.lock_key(), procedure.poison_keys(), procedure.type_name(), @@ -837,6 +843,7 @@ impl LocalManager { procedure_state, loaded_procedure.step, loaded_procedure.procedure, + message.context.clone(), RootSubmissionOrigin::Recovery, ) { error!(e; "Failed to recover procedure {}", procedure_id); @@ -964,6 +971,7 @@ impl ProcedureManager for LocalManager { ProcedureState::Running, 0, procedure.procedure, + procedure.context, RootSubmissionOrigin::Fresh, ) } @@ -1021,6 +1029,7 @@ pub(crate) mod test_util { ProcedureId::random(), ProcedureState::Running, None, + ProcedureContext::default(), LockKey::default(), PoisonKeys::default(), "ProcedureAdapter", @@ -1043,7 +1052,10 @@ mod tests { use common_error::mock::MockError; use common_error::status_code::StatusCode; - use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef}; + use common_event_recorder::{ + Event, EventRecorder, EventTypeFilter, EventTypeFilterRef, PersistentEventContext, + TriggerReason, + }; use common_test_util::temp_dir::create_temp_dir; use tokio::sync::oneshot; use tokio::time::{sleep, timeout}; @@ -1053,8 +1065,8 @@ mod tests { use crate::store::state_store::ObjectStateStore; use crate::test_util::InMemoryPoisonStore; use crate::{ - ChildSubmissionOutcome, Context, EventRuntimeContext, EventTrigger, Procedure, - ProcedureEvent, Status, + ChildSubmissionOutcome, Context, EventContext, EventTrigger, Procedure, ProcedureEvent, + Status, }; fn new_test_manager_context() -> ManagerContext { @@ -1103,6 +1115,13 @@ mod tests { }) .collect() } + + fn procedure_context(&self, procedure_id: ProcedureId) -> Option { + self.events.lock().unwrap().iter().find_map(|event| { + let event = event.as_any().downcast_ref::().unwrap(); + (event.procedure_id == procedure_id).then(|| event.context.clone()) + }) + } } impl EventRecorder for CapturingEventRecorder { @@ -1296,7 +1315,7 @@ mod tests { self.poison_keys.clone() } - fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, _ctx: &EventContext<'_>) -> Option> { Some(Box::new(TestProcedureEvent)) } } @@ -1341,6 +1360,7 @@ mod tests { vec![ProcedureWithId { id: self.child_id, procedure: Box::new(ProcedureToLoad::new("child")), + context: ProcedureContext::default(), }], false, )) @@ -1354,7 +1374,7 @@ mod tests { LockKey::default() } - fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, _ctx: &EventContext<'_>) -> Option> { Some(Box::new(TestProcedureEvent)) } } @@ -1384,6 +1404,7 @@ mod tests { vec![ProcedureWithId { id: self.child_id, procedure: Box::new(ProcedureToLoad::new("child")), + context: ProcedureContext::default(), }], false, )) @@ -1397,7 +1418,7 @@ mod tests { LockKey::default() } - fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, _ctx: &EventContext<'_>) -> Option> { Some(Box::new(TestProcedureEvent)) } } @@ -1424,7 +1445,7 @@ mod tests { LockKey::default() } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { *self.captured_filter.lock().unwrap() = Some(ctx.event_type_filter.clone()); Some(Box::new(TestProcedureEvent)) } @@ -1458,6 +1479,7 @@ mod tests { procedure: Box::new(FilterCapturingProcedure { captured_filter: captured_filter.clone(), }), + context: ProcedureContext::default(), }) .await .unwrap(); @@ -1499,6 +1521,7 @@ mod tests { .submit(ProcedureWithId { id: ProcedureId::random(), procedure: Box::new(ProcedureToLoad::new("fresh submission")), + context: ProcedureContext::default(), }) .await .unwrap(); @@ -1533,6 +1556,7 @@ mod tests { ProcedureState::Running, 0, Box::new(ProcedureToLoad::new("submit root without runner")), + ProcedureContext::default(), RootSubmissionOrigin::Fresh, ) }); @@ -1584,6 +1608,9 @@ mod tests { child_id, child_submitted: false, }), + context: ProcedureContext::from_event_context(PersistentEventContext::new( + TriggerReason::AutoRepartition, + )), }) .await .unwrap(); @@ -1604,6 +1631,15 @@ mod tests { }, ))); assert!(procedure_events.contains(&(child_id, EventTrigger::Submitted))); + assert_eq!( + event_recorder + .procedure_context(child_id) + .unwrap() + .event_context + .unwrap() + .reason, + TriggerReason::AutoRepartition + ); } #[tokio::test] @@ -1637,6 +1673,7 @@ mod tests { proceed: Some(proceed_rx), child_submitted: false, }), + context: ProcedureContext::default(), }) .await .unwrap(); @@ -1711,12 +1748,15 @@ mod tests { let procedure = ProcedureToLoad::new("recovered submission"); let procedure_id = ProcedureId::random(); ProcedureStore::from_object_store(object_store) - .store_procedure( + .store_procedure_with_context( procedure_id, 0, procedure.type_name().to_string(), procedure.dump().unwrap(), None, + ProcedureContext::from_event_context(PersistentEventContext::new( + TriggerReason::AutoRebalance, + )), ) .await .unwrap(); @@ -1732,6 +1772,15 @@ mod tests { ); wait_for_trigger(&event_recorder, EventTrigger::Recovered).await; assert!(!event_recorder.triggers().contains(&EventTrigger::Submitted)); + assert_eq!( + event_recorder + .procedure_context(procedure_id) + .unwrap() + .event_context + .unwrap() + .reason, + TriggerReason::AutoRebalance + ); } #[tokio::test] @@ -1813,6 +1862,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .unwrap(); @@ -1965,6 +2015,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() @@ -1986,6 +2037,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(ProcedureToLoad::new("submit")), + context: ProcedureContext::default(), }) .await .unwrap_err(); @@ -2054,6 +2106,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .unwrap() @@ -2094,6 +2147,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .unwrap_err(), @@ -2126,6 +2180,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() @@ -2164,6 +2219,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() @@ -2193,6 +2249,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() @@ -2217,6 +2274,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() @@ -2265,6 +2323,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .unwrap_err(); @@ -2285,6 +2344,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() diff --git a/src/common/procedure/src/local/runner.rs b/src/common/procedure/src/local/runner.rs index a460b7376bf..30b775dbd7d 100644 --- a/src/common/procedure/src/local/runner.rs +++ b/src/common/procedure/src/local/runner.rs @@ -34,8 +34,8 @@ use crate::procedure::{Output, StringKey}; use crate::rwlock::OwnedKeyRwLockGuard; use crate::store::{ProcedureMessage, ProcedureStore}; use crate::{ - BoxedProcedure, ChildSubmissionOutcome, Context, Error, EventRuntimeContext, EventTrigger, - Procedure, ProcedureId, ProcedureState, ProcedureWithId, RetryPhase, Status, + BoxedProcedure, ChildSubmissionOutcome, Context, Error, EventContext, EventTrigger, Procedure, + ProcedureId, ProcedureState, ProcedureWithId, RetryPhase, Status, }; /// A guard to cleanup procedure state. @@ -227,6 +227,7 @@ impl Runner { let ctx = Context { procedure_id: self.meta.id, provider: self.manager_ctx.clone(), + event_context: self.meta.context.event_context.clone(), }; self.rolling_back = false; @@ -508,6 +509,7 @@ impl Runner { procedure_id, procedure_state, Some(self.meta.id), + self.meta.context.clone(), procedure.lock_key(), procedure.poison_keys(), procedure.type_name(), @@ -662,12 +664,13 @@ impl Runner { let data = self.procedure.dump()?; self.store - .store_procedure( + .store_procedure_with_context( self.meta.id, self.step, type_name, data, self.meta.parent_id, + self.meta.context.clone(), ) .await .map_err(|e| { @@ -708,6 +711,7 @@ impl Runner { parent_id: self.meta.parent_id, step: self.step, error: Some(error), + context: self.meta.context.clone(), }; self.store .rollback_procedure(self.meta.id, message) @@ -776,15 +780,22 @@ impl Runner { pub(crate) fn build_event(&self, trigger: EventTrigger) -> Option { let recorder = self.event_recorder.as_ref()?; let state = self.meta.state(); - let context = EventRuntimeContext { + let context = EventContext { procedure_id: self.meta.id, lifecycle_state: &state, trigger: trigger.clone(), event_type_filter: recorder.event_type_filter(), + event_context: self.meta.context.event_context.as_ref(), }; - self.procedure - .event(&context) - .map(|event| ProcedureEvent::new(self.meta.id, event, state, trigger)) + self.procedure.event(&context).map(|event| { + ProcedureEvent::new_with_context( + self.meta.id, + event, + state, + trigger, + self.meta.context.clone(), + ) + }) } /// Builds and dispatches an event from the live procedure. Delivery is best effort and is @@ -876,6 +887,7 @@ mod tests { Context { procedure_id, provider, + event_context: None, } } @@ -914,6 +926,7 @@ mod tests { Context { procedure_id, provider: Arc::new(MockProvider), + event_context: None, } } @@ -1100,6 +1113,7 @@ mod tests { ProcedureWithId { id: procedure_id, procedure: Box::new(child), + context: Default::default(), } } @@ -1480,6 +1494,7 @@ mod tests { subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(fail), + context: Default::default(), }], persist: true, }) @@ -1672,6 +1687,7 @@ mod tests { subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(fail), + context: Default::default(), }], persist: true, }) @@ -2192,6 +2208,7 @@ mod tests { subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(child), + context: Default::default(), }], persist: false, }) @@ -2248,6 +2265,7 @@ mod tests { subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(child), + context: Default::default(), }], persist: false, }) diff --git a/src/common/procedure/src/procedure.rs b/src/common/procedure/src/procedure.rs index a8af03e0da0..cd11aea9215 100644 --- a/src/common/procedure/src/procedure.rs +++ b/src/common/procedure/src/procedure.rs @@ -19,7 +19,7 @@ use std::str::FromStr; use std::sync::Arc; use async_trait::async_trait; -use common_event_recorder::{Event, EventTypeFilterRef}; +use common_event_recorder::{Event, EventTypeFilterRef, PersistentEventContext}; use serde::{Deserialize, Serialize}; use smallvec::{SmallVec, smallvec}; use snafu::{ResultExt, Snafu}; @@ -32,6 +32,26 @@ use crate::watcher::Watcher; pub type Output = Arc; +/// Context attached to a procedure submission and inherited by its children. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct ProcedureContext { + /// Context describing why and how the procedure was submitted. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub event_context: Option, +} + +impl ProcedureContext { + pub fn from_event_context(event_context: PersistentEventContext) -> Self { + Self { + event_context: Some(event_context), + } + } + + pub fn is_empty(&self) -> bool { + self.event_context.is_none() + } +} + /// Procedure execution status. #[derive(Debug)] pub enum Status { @@ -175,6 +195,8 @@ pub struct Context { pub procedure_id: ProcedureId, /// [ProcedureManager] context provider. pub provider: ContextProviderRef, + /// Event context inherited from the root submission. + pub event_context: Option, } impl Context { @@ -237,13 +259,13 @@ pub trait Procedure: Send { /// [`Event::extra_schema`] values. The event recorder batches events by type /// and rejects incompatible schemas; use a distinct event type for a /// different schema. - fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, _ctx: &EventContext<'_>) -> Option> { None } } /// Framework-owned context supplied when a procedure builds a lifecycle event. -pub struct EventRuntimeContext<'a> { +pub struct EventContext<'a> { /// Id of the procedure associated with the event. pub procedure_id: ProcedureId, /// Current framework state of the procedure. @@ -252,6 +274,8 @@ pub struct EventRuntimeContext<'a> { pub trigger: EventTrigger, /// Event types retained by the configured recorder. pub event_type_filter: EventTypeFilterRef, + /// Event context inherited from the root submission. + pub event_context: Option<&'a PersistentEventContext>, } /// Lifecycle action that causes the framework to invoke [`Procedure::event`]. @@ -336,7 +360,7 @@ impl Procedure for Box { (**self).poison_keys() } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { (**self).event(ctx) } } @@ -459,6 +483,8 @@ pub struct ProcedureWithId { /// Id of the procedure. pub id: ProcedureId, pub procedure: BoxedProcedure, + /// Context associated with this procedure submission. + pub context: ProcedureContext, } impl ProcedureWithId { @@ -468,8 +494,15 @@ impl ProcedureWithId { ProcedureWithId { id: ProcedureId::random(), procedure, + context: ProcedureContext::default(), } } + + /// Attaches event context to this procedure submission. + pub fn with_event_context(mut self, event_context: PersistentEventContext) -> Self { + self.context.event_context = Some(event_context); + self + } } impl fmt::Debug for ProcedureWithId { @@ -721,11 +754,12 @@ mod tests { #[test] fn test_default_procedure_event_hook() { let state = ProcedureState::Running; - let context = EventRuntimeContext { + let context = EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Succeeded, event_type_filter: Arc::new(common_event_recorder::EventTypeFilter::All), + event_context: None, }; assert!(DefaultEventProcedure.event(&context).is_none()); diff --git a/src/common/procedure/src/store.rs b/src/common/procedure/src/store.rs index 906682ed6e6..4867046f8de 100644 --- a/src/common/procedure/src/store.rs +++ b/src/common/procedure/src/store.rs @@ -20,9 +20,9 @@ use futures::TryStreamExt; use serde::{Deserialize, Serialize}; use snafu::ResultExt; -use crate::ProcedureId; use crate::error::{Result, ToJsonSnafu}; pub(crate) use crate::store::state_store::StateStoreRef; +use crate::{ProcedureContext, ProcedureId}; pub mod poison_store; pub mod state_store; @@ -55,6 +55,9 @@ pub struct ProcedureMessage { /// Errors raised during the procedure. #[serde(default, skip_serializing_if = "Option::is_none")] pub error: Option, + /// Context associated with the root submission. + #[serde(default, skip_serializing_if = "ProcedureContext::is_empty")] + pub context: ProcedureContext, } /// A collection of all procedures' messages. @@ -88,6 +91,7 @@ impl ProcedureStore { } /// Dump the `procedure` to the storage. + #[cfg(test)] pub(crate) async fn store_procedure( &self, procedure_id: ProcedureId, @@ -95,6 +99,26 @@ impl ProcedureStore { type_name: String, data: String, parent_id: Option, + ) -> Result<()> { + self.store_procedure_with_context( + procedure_id, + step, + type_name, + data, + parent_id, + ProcedureContext::default(), + ) + .await + } + + pub(crate) async fn store_procedure_with_context( + &self, + procedure_id: ProcedureId, + step: u32, + type_name: String, + data: String, + parent_id: Option, + context: ProcedureContext, ) -> Result<()> { let message = ProcedureMessage { type_name, @@ -102,6 +126,7 @@ impl ProcedureStore { parent_id, step, error: None, + context, }; let key = ParsedKey { prefix: &self.proc_path, @@ -356,6 +381,7 @@ mod tests { } use async_trait::async_trait; + use common_event_recorder::{PersistentEventContext, TriggerReason}; use common_test_util::temp_dir::{TempDir, create_temp_dir}; use object_store::services::Fs as Builder; @@ -462,6 +488,7 @@ mod tests { parent_id: None, step: 4, error: None, + context: ProcedureContext::default(), }; let json = serde_json::to_string(&message).unwrap(); @@ -477,6 +504,21 @@ mod tests { json, r#"{"type_name":"TestMessage","data":"no parent id","parent_id":"9f805a1f-05f7-490c-9f91-bd56e3cc54c1","step":4}"# ); + + message.context = ProcedureContext::from_event_context(PersistentEventContext::new( + TriggerReason::AutoRebalance, + )); + let json = serde_json::to_string(&message).unwrap(); + assert_eq!( + json, + r#"{"type_name":"TestMessage","data":"no parent id","parent_id":"9f805a1f-05f7-490c-9f91-bd56e3cc54c1","step":4,"context":{"event_context":{"reason":"auto_rebalance"}}}"# + ); + + let legacy: ProcedureMessage = serde_json::from_str( + r#"{"type_name":"TestMessage","data":"legacy","parent_id":null,"step":1}"#, + ) + .unwrap(); + assert_eq!(legacy.context, ProcedureContext::default()); } struct MockProcedure { @@ -541,6 +583,7 @@ mod tests { parent_id: None, step: 0, error: None, + context: ProcedureContext::default(), }; assert_eq!(expect, *msg); } @@ -589,6 +632,7 @@ mod tests { parent_id: None, step: 1, error: None, + context: ProcedureContext::default(), }; store .rollback_procedure(procedure_id, message) diff --git a/src/common/procedure/src/watcher.rs b/src/common/procedure/src/watcher.rs index 4ecd75ded0f..4717fcd0010 100644 --- a/src/common/procedure/src/watcher.rs +++ b/src/common/procedure/src/watcher.rs @@ -124,6 +124,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(MockProcedure { error: true }), + context: Default::default(), }) .await .unwrap(); diff --git a/src/meta-srv/src/event/gc.rs b/src/meta-srv/src/event/gc.rs index c467bc665c6..3575ac20c2f 100644 --- a/src/meta-srv/src/event/gc.rs +++ b/src/meta-srv/src/event/gc.rs @@ -21,10 +21,9 @@ use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, GC_REPORT_COLUMN, REGION_ID_COLUMN, REGION_NUMBER_COLUMN, - TABLE_ID_COLUMN, column_schemas, nullable_json, + GC_REPORT_COLUMN, REGION_ID_COLUMN, REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, column_schemas, + nullable_json, }; -use common_meta::rpc::ddl::EventContext; use serde::Serialize; use snafu::ResultExt; use store_api::storage::{GcReport, IndexVersion, RegionId}; @@ -67,7 +66,6 @@ struct BatchGcRegionRow { #[derive(Debug)] pub(crate) struct BatchGcEvent { payload: Option, - event_context: Option, // None emits a procedure-level lifecycle row with null Region dimensions. regions: Option>, } @@ -77,7 +75,6 @@ impl BatchGcEvent { regions: &[RegionId], full_file_listing: bool, timeout: Duration, - event_context: Option, ) -> Self { Self { payload: Some(BatchGcPayload { @@ -86,7 +83,6 @@ impl BatchGcEvent { full_file_listing, timeout, }), - event_context, regions: None, } } @@ -111,7 +107,6 @@ impl BatchGcEvent { (!regions.is_empty()).then_some(Self { payload: None, - event_context: None, regions: Some(regions), }) } @@ -136,22 +131,8 @@ impl Event for BatchGcEvent { } fn extra_rows(&self) -> Result> { - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; let Some(regions) = &self.regions else { - return Ok(vec![Row { - values: vec![ - Value { value_data: None }, - Value { value_data: None }, - Value { value_data: None }, - Value { value_data: None }, - nullable_json(event_context.as_ref()), - ], - }]); + return Ok(vec![null_row()]); }; regions @@ -164,7 +145,6 @@ impl Event for BatchGcEvent { ValueData::U32Value(region.region_id.table_id()).into(), ValueData::U32Value(region.region_id.region_number()).into(), nullable_json(Some(&report)), - nullable_json(None), ], }) }) @@ -182,11 +162,9 @@ fn schema() -> Vec { &TABLE_ID_COLUMN, ®ION_NUMBER_COLUMN, &GC_REPORT_COLUMN, - &EVENT_CONTEXT_COLUMN, ]) } -#[cfg(test)] fn null_row() -> Row { Row { values: (0..schema().len()) @@ -233,17 +211,17 @@ mod tests { use std::sync::Arc; use api::v1::ColumnSchema; - use common_event_recorder::EventTypeFilter; use common_event_recorder::event_table::{ - PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, + EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; + use common_event_recorder::{EventTypeFilter, PersistentEventContext, TriggerReason}; use common_meta::key::TableMetadataManager; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::sequence::SequenceBuilder; use common_procedure::{ - EventRuntimeContext, EventTrigger, Procedure, ProcedureEvent, ProcedureId, ProcedureState, + EventContext, EventTrigger, Procedure, ProcedureEvent, ProcedureId, ProcedureState, RetryPhase, }; use store_api::storage::FileId; @@ -256,12 +234,8 @@ mod tests { fn test_batch_gc_lifecycle_event_contract() { let first = RegionId::new(1024, 1); let second = RegionId::new(1024, 2); - let event = BatchGcEvent::with_config( - &[second, first, second], - true, - Duration::from_secs(10), - None, - ); + let event = + BatchGcEvent::with_config(&[second, first, second], true, Duration::from_secs(10)); assert_event_contract(&event, BATCH_GC_EVENT_TYPE, &schema(), &[null_row()]); assert_eq!( @@ -356,22 +330,35 @@ mod tests { fn test_batch_gc_event_filter() { let procedure = batch_gc_procedure(); let running = ProcedureState::Running; - let runtime_context = |trigger, lifecycle_state, event_type_filter| EventRuntimeContext { + let manual_context = PersistentEventContext::new(TriggerReason::Manual); + let event_context = |trigger, lifecycle_state, event_type_filter| EventContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, event_type_filter: Arc::new(event_type_filter), + event_context: None, }; assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Submitted, &running, EventTypeFilter::All, )) .is_none() ); + assert!( + procedure + .event(&EventContext { + procedure_id: ProcedureId::random(), + lifecycle_state: &running, + trigger: EventTrigger::Submitted, + event_type_filter: Arc::new(EventTypeFilter::All), + event_context: Some(&manual_context), + }) + .is_some() + ); let report = GcReport { processed_regions: HashSet::from([RegionId::new(1024, 1)]), @@ -382,7 +369,7 @@ mod tests { }; assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Succeeded, &done, EventTypeFilter::All, @@ -392,7 +379,7 @@ mod tests { assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Recovered, &running, EventTypeFilter::All, @@ -401,7 +388,7 @@ mod tests { ); let retrying = procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Retrying { phase: RetryPhase::Execute, attempt: 1, @@ -423,7 +410,7 @@ mod tests { assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Recovered, &running, EventTypeFilter::Only(HashSet::from(["another_event".to_string()])), @@ -432,7 +419,7 @@ mod tests { ); assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Recovered, &running, EventTypeFilter::Only(HashSet::new()), @@ -443,7 +430,7 @@ mod tests { let missing = ProcedureState::Done { output: None }; assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Succeeded, &missing, EventTypeFilter::All, @@ -455,7 +442,7 @@ mod tests { }; assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Succeeded, &wrong, EventTypeFilter::All, @@ -468,11 +455,12 @@ mod tests { fn test_batch_gc_terminal_event_with_report() { let mut procedure = batch_gc_procedure(); let running = ProcedureState::Running; - let runtime_context = |trigger| EventRuntimeContext { + let event_context = |trigger| EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &running, trigger, event_type_filter: Arc::new(EventTypeFilter::All), + event_context: None, }; let report = GcReport { need_retry_regions: HashSet::from([RegionId::new(1024, 2)]), @@ -481,7 +469,7 @@ mod tests { procedure.set_gc_report_for_test(report); let retrying = procedure - .event(&runtime_context(EventTrigger::Retrying { + .event(&event_context(EventTrigger::Retrying { phase: RetryPhase::Execute, attempt: 2, })) @@ -498,7 +486,7 @@ mod tests { assert_eq!(retrying.extra_rows().unwrap(), vec![null_row()]); for trigger in [EventTrigger::Failed, EventTrigger::Poisoned] { - let event = procedure.event(&runtime_context(trigger)).unwrap(); + let event = procedure.event(&event_context(trigger)).unwrap(); assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); assert_eq!( event.extra_rows().unwrap(), @@ -513,7 +501,7 @@ mod tests { let procedure = batch_gc_procedure(); let failed_without_report = procedure - .event(&runtime_context(EventTrigger::Failed)) + .event(&event_context(EventTrigger::Failed)) .unwrap(); assert_eq!( failed_without_report.json_payload().unwrap(), @@ -548,6 +536,7 @@ mod tests { ); let mut event_schema = procedure_schema(); event_schema.extend(schema()); + event_schema.push(EVENT_CONTEXT_COLUMN.column_schema()); let mut values = vec![ ValueData::StringValue(procedure_id.to_string()).into(), ValueData::StringValue("Done".to_string()).into(), @@ -564,6 +553,7 @@ mod tests { ) .values, ); + values.push(Value { value_data: None }); assert_event_contract( &event, @@ -589,7 +579,6 @@ mod tests { ValueData::U32Value(region_id.table_id()).into(), ValueData::U32Value(region_id.region_number()).into(), nullable_json(report.as_ref()), - nullable_json(None), ], } } @@ -607,7 +596,6 @@ mod tests { true, Duration::from_secs(10), HashMap::new(), - EventContext::default(), ) } } diff --git a/src/meta-srv/src/event/region_migration.rs b/src/meta-srv/src/event/region_migration.rs index 65b89df2515..cfbe114077e 100644 --- a/src/meta-srv/src/event/region_migration.rs +++ b/src/meta-srv/src/event/region_migration.rs @@ -17,14 +17,13 @@ use std::time::Duration; use api::v1::value::ValueData; use api::v1::{ColumnSchema, Row}; -use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, REGION_ID_COLUMN, REGION_MIGRATION_DST_NODE_ID_COLUMN, - REGION_MIGRATION_DST_PEER_ADDR_COLUMN, REGION_MIGRATION_SRC_NODE_ID_COLUMN, - REGION_MIGRATION_SRC_PEER_ADDR_COLUMN, REGION_MIGRATION_TRIGGER_REASON_COLUMN, - REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, column_schemas, nullable_json, + REGION_ID_COLUMN, REGION_MIGRATION_DST_NODE_ID_COLUMN, REGION_MIGRATION_DST_PEER_ADDR_COLUMN, + REGION_MIGRATION_SRC_NODE_ID_COLUMN, REGION_MIGRATION_SRC_PEER_ADDR_COLUMN, + REGION_MIGRATION_TRIGGER_REASON_COLUMN, REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, column_schemas, }; +use common_event_recorder::{Event, TriggerReason}; use serde::Serialize; use snafu::ResultExt; use store_api::storage::RegionId; @@ -50,7 +49,6 @@ pub(crate) struct RegionMigrationEvent { dst_peer_addr: String, // The timeout of the region migration. timeout: Duration, - event_context: Option, } #[derive(Debug, Serialize)] @@ -60,16 +58,15 @@ struct Payload { } impl RegionMigrationEvent { - pub fn from_persistent_ctx(ctx: &PersistentContext, include_event_context: bool) -> Self { + pub fn from_persistent_ctx(ctx: &PersistentContext, trigger_reason: TriggerReason) -> Self { Self { region_ids: ctx.region_ids.clone(), - trigger_reason: ctx.trigger_reason, + trigger_reason: RegionMigrationTriggerReason::from_trigger_reason(trigger_reason), src_node_id: ctx.from_peer.id, src_peer_addr: ctx.from_peer.addr.clone(), dst_node_id: ctx.to_peer.id, dst_peer_addr: ctx.to_peer.addr.clone(), timeout: ctx.timeout, - event_context: include_event_context.then(|| ctx.event_context.clone()), } } } @@ -89,7 +86,6 @@ impl Event for RegionMigrationEvent { ®ION_MIGRATION_SRC_PEER_ADDR_COLUMN, ®ION_MIGRATION_DST_NODE_ID_COLUMN, ®ION_MIGRATION_DST_PEER_ADDR_COLUMN, - &EVENT_CONTEXT_COLUMN, ]) } @@ -106,12 +102,6 @@ impl Event for RegionMigrationEvent { ValueData::StringValue(self.src_peer_addr.clone()).into(), ValueData::U64Value(self.dst_node_id).into(), ValueData::StringValue(self.dst_peer_addr.clone()).into(), - self.event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu) - .map(|value| nullable_json(value.as_ref()))?, ], }); } diff --git a/src/meta-srv/src/event/repartition.rs b/src/meta-srv/src/event/repartition.rs index 958bc4e11d8..1faf3f18212 100644 --- a/src/meta-srv/src/event/repartition.rs +++ b/src/meta-srv/src/event/repartition.rs @@ -21,11 +21,11 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PARENT_PROCEDURE_ID_COLUMN, - REPARTITION_GROUP_ID_COLUMN, SCHEMA_NAME_COLUMN, SOURCE_PARTITION_EXPR_COLUMN, - SOURCE_REGION_ID_COLUMN, SOURCE_REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, TABLE_NAME_COLUMN, - TARGET_PARTITION_EXPR_COLUMN, TARGET_REGION_ID_COLUMN, TARGET_REGION_NUMBER_COLUMN, - column_schemas, nullable_json, nullable_string, nullable_value, + CATALOG_NAME_COLUMN, PARENT_PROCEDURE_ID_COLUMN, REPARTITION_GROUP_ID_COLUMN, + SCHEMA_NAME_COLUMN, SOURCE_PARTITION_EXPR_COLUMN, SOURCE_REGION_ID_COLUMN, + SOURCE_REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, TABLE_NAME_COLUMN, TARGET_PARTITION_EXPR_COLUMN, + TARGET_REGION_ID_COLUMN, TARGET_REGION_NUMBER_COLUMN, column_schemas, nullable_string, + nullable_value, }; use serde::Serialize; use snafu::ResultExt; @@ -61,7 +61,6 @@ pub(crate) struct RepartitionEvent { table_name: Option, table_id: Option, payload: Option, - event_context: Option, } impl RepartitionEvent { @@ -91,7 +90,6 @@ impl RepartitionEvent { target_partition_columns: intent.target_partition_columns().map(ToOwned::to_owned), timeout: persistent_ctx.timeout, }), - event_context: Some(persistent_ctx.event_context.clone()), } } @@ -102,17 +100,10 @@ impl RepartitionEvent { table_name: Some(persistent_ctx.table_name.clone()), table_id: Some(persistent_ctx.table_id), payload: None, - event_context: None, } } fn schema() -> Vec { - let mut schema = Self::base_schema(); - schema.push(EVENT_CONTEXT_COLUMN.column_schema()); - schema - } - - fn base_schema() -> Vec { column_schemas([ &CATALOG_NAME_COLUMN, &SCHEMA_NAME_COLUMN, @@ -141,19 +132,12 @@ impl Event for RepartitionEvent { } fn extra_rows(&self) -> Result> { - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.schema_name.as_deref()), nullable_string(self.table_name.as_deref()), nullable_value(self.table_id.map(ValueData::U32Value)), - nullable_json(event_context.as_ref()), ], }]) } @@ -323,7 +307,7 @@ impl RepartitionGroupEvent { } fn schema() -> Vec { - let mut schema = RepartitionEvent::base_schema(); + let mut schema = RepartitionEvent::schema(); schema.extend(column_schemas([ &PARENT_PROCEDURE_ID_COLUMN, &REPARTITION_GROUP_ID_COLUMN, @@ -377,14 +361,13 @@ mod tests { use std::time::Duration; use api::v1::value::ValueData; - use api::v1::{ColumnSchema, Row}; + use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, + EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; - use common_meta::rpc::ddl::EventContext; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; use table::table_name::TableName; use uuid::Uuid; @@ -402,7 +385,6 @@ mod tests { TableName::new("greptime", "public", "repartition_events"), 1024, Some(Duration::from_secs(30)), - EventContext::default(), ) } @@ -438,9 +420,6 @@ mod tests { ValueData::StringValue("public".to_string()).into(), ValueData::StringValue("repartition_events".to_string()).into(), ValueData::U32Value(1024).into(), - jsonb_value(&serde_json::json!({ - "reason": "unknown", - })), ], }], ); @@ -654,7 +633,6 @@ mod tests { ValueData::StringValue("public".to_string()).into(), ValueData::StringValue("repartition_events".to_string()).into(), ValueData::U32Value(1024).into(), - Default::default(), ], }], ); @@ -690,6 +668,7 @@ mod tests { ); let mut schema = procedure_schema(); schema.extend(parent_schema()); + schema.push(EVENT_CONTEXT_COLUMN.column_schema()); assert_event_contract( &event, @@ -705,9 +684,7 @@ mod tests { ValueData::StringValue("public".to_string()).into(), ValueData::StringValue("repartition_events".to_string()).into(), ValueData::U32Value(1024).into(), - jsonb_value(&serde_json::json!({ - "reason": "unknown", - })), + Value { value_data: None }, ], }], ); @@ -719,12 +696,11 @@ mod tests { &SCHEMA_NAME_COLUMN, &TABLE_NAME_COLUMN, &TABLE_ID_COLUMN, - &EVENT_CONTEXT_COLUMN, ]) } fn group_schema() -> Vec { - let mut schema = RepartitionEvent::base_schema(); + let mut schema = parent_schema(); schema.extend(column_schemas([ &PARENT_PROCEDURE_ID_COLUMN, &REPARTITION_GROUP_ID_COLUMN, diff --git a/src/meta-srv/src/gc/ctx.rs b/src/meta-srv/src/gc/ctx.rs index 41027a65390..c9cc8049867 100644 --- a/src/meta-srv/src/gc/ctx.rs +++ b/src/meta-srv/src/gc/ctx.rs @@ -29,7 +29,7 @@ use common_meta::key::DroppedTableName; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; #[cfg(feature = "enterprise")] use common_meta::rpc::ddl::PurgeDroppedTableTask; use common_procedure::{ProcedureManagerRef, ProcedureWithId, watcher}; @@ -68,7 +68,7 @@ pub(crate) trait SchedulerCtx: Send + Sync { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result; #[cfg(feature = "enterprise")] @@ -256,7 +256,7 @@ impl SchedulerCtx for DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { self.gc_regions_inner( region_ids, @@ -310,7 +310,7 @@ impl DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { debug!( "Sending GC instruction for {} regions (full_file_listing: {})", @@ -326,9 +326,9 @@ impl DefaultGcSchedulerCtx { full_file_listing, timeout, region_routes_override, - event_context, ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); let id = procedure_with_id.id; diff --git a/src/meta-srv/src/gc/handler.rs b/src/meta-srv/src/gc/handler.rs index 4532e9a8ca3..16b29a39c7e 100644 --- a/src/meta-srv/src/gc/handler.rs +++ b/src/meta-srv/src/gc/handler.rs @@ -18,7 +18,7 @@ use std::time::Instant; use common_catalog::consts::MITO_ENGINE; use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_telemetry::tracing::Instrument as _; use common_telemetry::{debug, error, info, warn}; use futures::StreamExt; @@ -37,7 +37,10 @@ use crate::gc::tracker::RegionGcInfo; use crate::metrics::METRIC_META_GC_CANDIDATE_REGIONS; impl GcScheduler { - pub(crate) async fn trigger_gc(&self, event_context: EventContext) -> Result { + pub(crate) async fn trigger_gc( + &self, + event_context: PersistentEventContext, + ) -> Result { let start_time = Instant::now(); info!("Starting GC cycle"); @@ -212,7 +215,7 @@ impl GcScheduler { datanode_to_candidates: HashMap>, force_full_listing_by_peer: HashMap>, region_routes_override_by_peer: HashMap, - event_context: EventContext, + event_context: PersistentEventContext, ) -> GcJobReport { let mut per_datanode_reports = HashMap::new(); let mut failed_datanodes: HashMap<_, Vec<_>> = HashMap::new(); @@ -283,7 +286,7 @@ impl GcScheduler { candidates: Vec<(TableId, GcCandidate)>, force_full_listing: HashSet, region_routes_override: Region2Peers, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { info!( "Starting GC for datanode {} with {} candidate regions", diff --git a/src/meta-srv/src/gc/mock.rs b/src/meta-srv/src/gc/mock.rs index a941476ab44..94749dfb469 100644 --- a/src/meta-srv/src/gc/mock.rs +++ b/src/meta-srv/src/gc/mock.rs @@ -30,7 +30,7 @@ use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::{Region, RegionRoute}; use common_telemetry::debug; use ordered_float::OrderedFloat; @@ -75,7 +75,7 @@ pub struct MockSchedulerCtx { pub candidates: Arc>>>>, pub get_table_to_region_stats_calls: Arc>, pub gc_regions_calls: Arc>, - pub gc_event_contexts: Arc>>, + pub gc_event_contexts: Arc>>, // Error injection fields for testing pub get_table_to_region_stats_error: Arc>>, pub get_table_route_error: Arc>>, @@ -214,7 +214,7 @@ impl SchedulerCtx for MockSchedulerCtx { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { *self.gc_regions_calls.lock().unwrap() += 1; self.gc_event_contexts.lock().unwrap().push(event_context); diff --git a/src/meta-srv/src/gc/mock/basic.rs b/src/meta-srv/src/gc/mock/basic.rs index 45d64910561..ed051cc839e 100644 --- a/src/meta-srv/src/gc/mock/basic.rs +++ b/src/meta-srv/src/gc/mock/basic.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::Instant; use common_meta::peer::Peer; -use common_meta::rpc::ddl::{EventContext, TriggerReason}; +use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -36,7 +36,7 @@ async fn test_parallel_process_datanodes_empty() { HashMap::new(), HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -105,7 +105,7 @@ async fn test_parallel_process_datanodes_with_candidates() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -189,7 +189,7 @@ async fn test_handle_tick() { assert_eq!(*ctx.gc_regions_calls.lock().unwrap(), 1); assert_eq!( ctx.gc_event_contexts.lock().unwrap().as_slice(), - &[EventContext::new(TriggerReason::ScheduledGc)] + &[PersistentEventContext::new(TriggerReason::ScheduledGc)] ); let tracker = scheduler.region_gc_tracker.lock().await; @@ -249,6 +249,6 @@ async fn test_handle_manual_gc_without_regions_records_manual_event_context() { assert_eq!( ctx.gc_event_contexts.lock().unwrap().as_slice(), - &[EventContext::new(TriggerReason::Manual)] + &[PersistentEventContext::new(TriggerReason::Manual)] ); } diff --git a/src/meta-srv/src/gc/mock/concurrent.rs b/src/meta-srv/src/gc/mock/concurrent.rs index 7dca378b71d..5bce9a1be8b 100644 --- a/src/meta-srv/src/gc/mock/concurrent.rs +++ b/src/meta-srv/src/gc/mock/concurrent.rs @@ -18,7 +18,7 @@ use std::time::{Duration, Instant}; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::{Region, RegionRoute}; use common_telemetry::{info, init_default_ut_logging}; use store_api::region_engine::RegionRole; @@ -102,7 +102,7 @@ async fn test_concurrent_table_processing_limits() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -190,7 +190,7 @@ async fn test_datanode_processes_tables_with_partial_gc_failures() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -296,7 +296,7 @@ async fn test_region_gc_concurrency_limit() { candidates.into_iter().map(|c| (table_id, c)).collect(), HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); @@ -413,7 +413,7 @@ async fn test_region_gc_concurrency_with_partial_failures() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -557,7 +557,7 @@ async fn test_region_gc_concurrency_with_retryable_errors() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; diff --git a/src/meta-srv/src/gc/mock/full_list.rs b/src/meta-srv/src/gc/mock/full_list.rs index a5d5012c6d3..8193a9094e1 100644 --- a/src/meta-srv/src/gc/mock/full_list.rs +++ b/src/meta-srv/src/gc/mock/full_list.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -83,7 +83,7 @@ async fn test_full_file_listing_first_time_gc() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); @@ -157,7 +157,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); @@ -184,7 +184,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); @@ -260,7 +260,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); @@ -284,7 +284,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); diff --git a/src/meta-srv/src/gc/mock/misc.rs b/src/meta-srv/src/gc/mock/misc.rs index 9b3c35cbe9a..575a5c5fc08 100644 --- a/src/meta-srv/src/gc/mock/misc.rs +++ b/src/meta-srv/src/gc/mock/misc.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::Instant; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_telemetry::init_default_ut_logging; use store_api::storage::{FileRefsManifest, GcReport, RegionId}; @@ -76,7 +76,7 @@ async fn test_empty_file_refs_manifest() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -169,7 +169,7 @@ async fn test_multiple_regions_per_table() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; diff --git a/src/meta-srv/src/gc/procedure.rs b/src/meta-srv/src/gc/procedure.rs index d745e88e2c5..f9f597577ae 100644 --- a/src/meta-srv/src/gc/procedure.rs +++ b/src/meta-srv/src/gc/procedure.rs @@ -24,11 +24,11 @@ use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::lock_key::{RegionLock, TableLock}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::{EventContext, TriggerReason}; +use common_meta::rpc::ddl::TriggerReason; use common_procedure::error::ToJsonSnafu; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, - LockKey, Procedure, ProcedureState, Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, ProcedureState, Result as ProcedureResult, Status, }; use common_telemetry::tracing::Instrument as _; use common_telemetry::tracing_context::TracingContext; @@ -200,8 +200,6 @@ pub struct BatchGcData { /// mailbox timeout duration timeout: Duration, gc_report: Option, - #[serde(default)] - event_context: EventContext, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -219,7 +217,6 @@ pub enum State { impl BatchGcProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::BatchGcProcedure"; - #[allow(clippy::too_many_arguments)] pub fn new( mailbox: MailboxRef, table_metadata_manager: TableMetadataManagerRef, @@ -228,7 +225,6 @@ impl BatchGcProcedure { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: EventContext, ) -> Self { Self { mailbox, @@ -244,7 +240,6 @@ impl BatchGcProcedure { related_regions: HashMap::new(), file_refs: FileRefsManifest::default(), gc_report: None, - event_context, }, } } @@ -275,7 +270,6 @@ impl BatchGcProcedure { related_regions: HashMap::new(), file_refs, gc_report: Some(GcReport::default()), - event_context: EventContext::default(), }, } } @@ -1072,23 +1066,21 @@ impl Procedure for BatchGcProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(BATCH_GC_EVENT_TYPE) { return None; } let event = match &ctx.trigger { // Keep scheduled GC low-noise; record submitted manual requests for auditability. - EventTrigger::Submitted => (self.data.event_context.reason == TriggerReason::Manual) + EventTrigger::Submitted => ctx + .event_context + .is_some_and(|context| context.reason == TriggerReason::Manual) .then(|| { BatchGcEvent::with_config( &self.data.regions, self.data.full_file_listing, self.data.timeout, - Some(self.data.event_context.clone()), ) })?, EventTrigger::Recovered | EventTrigger::ChildSubmitted { .. } => return None, @@ -1106,7 +1098,6 @@ impl Procedure for BatchGcProcedure { &self.data.regions, self.data.full_file_listing, self.data.timeout, - None, ), EventTrigger::Failed | EventTrigger::Poisoned => self .data @@ -1118,7 +1109,6 @@ impl Procedure for BatchGcProcedure { &self.data.regions, self.data.full_file_listing, self.data.timeout, - None, ) }), }; @@ -1211,16 +1201,6 @@ mod tests { assert!(report.need_retry_regions.contains(®ion_id)); } - #[test] - fn test_batch_gc_data_defaults_missing_event_context() { - let mut data = serde_json::to_value(&batch_gc_procedure().data).unwrap(); - data.as_object_mut().unwrap().remove("event_context"); - - let recovered: BatchGcData = serde_json::from_value(data).unwrap(); - - assert_eq!(recovered.event_context, EventContext::default()); - } - #[tokio::test] async fn test_send_gc_instructions_preserves_partial_report() { let first_region = RegionId::new(1024, 1); @@ -1255,7 +1235,6 @@ mod tests { true, Duration::from_secs(10), HashMap::new(), - EventContext::default(), ); procedure.data.region_routes = HashMap::from([ (first_region, (first_peer, vec![])), @@ -1296,7 +1275,6 @@ mod tests { true, Duration::from_secs(10), HashMap::new(), - EventContext::default(), ) } } diff --git a/src/meta-srv/src/gc/scheduler.rs b/src/meta-srv/src/gc/scheduler.rs index 92b362ad947..b89c0d06be9 100644 --- a/src/meta-srv/src/gc/scheduler.rs +++ b/src/meta-srv/src/gc/scheduler.rs @@ -18,7 +18,7 @@ use std::time::{Duration, Instant}; use common_meta::DatanodeId; use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; -use common_meta::rpc::ddl::{EventContext, TriggerReason}; +use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; use common_telemetry::tracing::Instrument as _; use common_telemetry::{error, info}; use snafu::ResultExt; @@ -218,7 +218,7 @@ impl GcScheduler { } let span = common_telemetry::tracing::info_span!("meta_gc_handle_tick"); let report = self - .trigger_gc(EventContext::new(TriggerReason::ScheduledGc)) + .trigger_gc(PersistentEventContext::new(TriggerReason::ScheduledGc)) .instrument(span) .await?; @@ -305,7 +305,7 @@ impl GcScheduler { // No specific regions, use default tick behavior let Some(regions) = region_ids else { let report = self - .trigger_gc(EventContext::new(TriggerReason::Manual)) + .trigger_gc(PersistentEventContext::new(TriggerReason::Manual)) .await?; info!("Finished manual gc request"); return Ok(report); @@ -353,7 +353,7 @@ impl GcScheduler { full_listing, gc_timeout, Region2Peers::new(), - EventContext::new(TriggerReason::Manual), + PersistentEventContext::new(TriggerReason::Manual), ) .await?; combined_report.merge(report); @@ -367,7 +367,7 @@ impl GcScheduler { true, gc_timeout, dropped_routes_override, - EventContext::new(TriggerReason::Manual), + PersistentEventContext::new(TriggerReason::Manual), ) .await?; combined_report.merge(report); @@ -490,7 +490,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: EventContext, + _event_context: PersistentEventContext, ) -> Result { self.gc_regions_calls.fetch_add(1, Ordering::Relaxed); panic!("gc_regions should not be called in maintenance mode") @@ -564,7 +564,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: EventContext, + _event_context: PersistentEventContext, ) -> Result { Ok(GcReport::default()) } @@ -898,7 +898,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: EventContext, + _event_context: PersistentEventContext, ) -> Result { crate::error::UnexpectedSnafu { violated: "mock gc failure".to_string(), diff --git a/src/meta-srv/src/procedure/region_migration.rs b/src/meta-srv/src/procedure/region_migration.rs index 25fc7d6a309..ad1553c92d7 100644 --- a/src/meta-srv/src/procedure/region_migration.rs +++ b/src/meta-srv/src/procedure/region_migration.rs @@ -45,13 +45,12 @@ use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; use common_meta::lock_key::{CatalogLock, RegionLock, SchemaLock, TableLock}; use common_meta::peer::Peer; use common_meta::region_keeper::{MemoryRegionKeeperRef, OperatingRegionGuard}; -use common_meta::rpc::ddl::{EventContext, TriggerReason}; +use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; use common_procedure::error::{ Error as ProcedureError, FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu, }; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Status, - StringKey, + Context as ProcedureContext, EventContext, LockKey, Procedure, Status, StringKey, }; use common_telemetry::{debug, error, info}; use manager::RegionMigrationProcedureGuard; @@ -123,11 +122,9 @@ pub struct PersistentContext { /// The timeout for downgrading leader region and upgrading candidate region operations. #[serde(with = "humantime_serde", default = "default_timeout")] pub(crate) timeout: Duration, - /// The trigger reason of region migration. + /// The trigger reason persisted for compatibility with versions without procedure context. #[serde(default)] - pub(crate) trigger_reason: RegionMigrationTriggerReason, - #[serde(default)] - pub(crate) event_context: EventContext, + trigger_reason: RegionMigrationTriggerReason, } impl PersistentContext { @@ -137,7 +134,7 @@ impl PersistentContext { to_peer: Peer, region_ids: Vec, timeout: Duration, - event_context: EventContext, + trigger_reason: RegionMigrationTriggerReason, ) -> Self { #[allow(deprecated)] Self { @@ -148,17 +145,7 @@ impl PersistentContext { to_peer, region_ids, timeout, - trigger_reason: RegionMigrationTriggerReason::from_trigger_reason(event_context.reason), - event_context, - } - } - - fn normalize_trigger_reasons(&mut self) { - if self.trigger_reason != RegionMigrationTriggerReason::Unknown { - self.event_context.reason = self.trigger_reason.to_trigger_reason(); - } else if self.event_context.reason != TriggerReason::Unknown { - self.trigger_reason = - RegionMigrationTriggerReason::from_trigger_reason(self.event_context.reason); + trigger_reason, } } } @@ -173,7 +160,7 @@ impl RegionMigrationTriggerReason { } } - fn from_trigger_reason(reason: TriggerReason) -> Self { + pub(crate) fn from_trigger_reason(reason: TriggerReason) -> Self { match reason { TriggerReason::Manual => Self::Manual, TriggerReason::AutoRebalance => Self::AutoRebalance, @@ -457,6 +444,14 @@ pub struct Context { } impl Context { + pub(crate) fn trigger_reason( + &self, + event_context: Option<&PersistentEventContext>, + ) -> RegionMigrationTriggerReason { + event_context + .map(|ctx| RegionMigrationTriggerReason::from_trigger_reason(ctx.reason)) + .unwrap_or(self.persistent_ctx.trigger_reason) + } /// Returns the next operation's timeout. pub fn next_operation_timeout(&self) -> Option { self.persistent_ctx @@ -852,10 +847,9 @@ impl RegionMigrationProcedure { tracker: RegionMigrationProcedureTracker, ) -> ProcedureResult { let RegionMigrationDataOwned { - mut persistent_ctx, + persistent_ctx, state, } = serde_json::from_str(json).context(FromJsonSnafu)?; - persistent_ctx.normalize_trigger_reasons(); let guards = persistent_ctx .region_ids .iter() @@ -999,14 +993,16 @@ impl Procedure for RegionMigrationProcedure { LockKey::new(self.context.persistent_ctx.lock_key()) } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(REGION_MIGRATION_EVENT_TYPE) { return None; } Some(Box::new(RegionMigrationEvent::from_persistent_ctx( &self.context.persistent_ctx, - matches!(ctx.trigger, EventTrigger::Submitted), + self.context + .trigger_reason(ctx.event_context) + .to_trigger_reason(), ))) } } @@ -1063,7 +1059,7 @@ mod tests { let procedure = RegionMigrationProcedure::new(persistent_context, context, vec![]); let serialized = procedure.dump().unwrap(); - let expected = r#"{"persistent_ctx":{"catalog_and_schema":[["greptime","public"]],"from_peer":{"id":1,"addr":""},"to_peer":{"id":2,"addr":""},"region_ids":[4398046511105],"timeout":"10s","trigger_reason":"Unknown","event_context":{"reason":"unknown"}},"state":{"region_migration_state":"RegionMigrationStart"}}"#; + let expected = r#"{"persistent_ctx":{"catalog_and_schema":[["greptime","public"]],"from_peer":{"id":1,"addr":""},"to_peer":{"id":2,"addr":""},"region_ids":[4398046511105],"timeout":"10s","trigger_reason":"Unknown"},"state":{"region_migration_state":"RegionMigrationStart"}}"#; assert_eq!(expected, serialized); } @@ -1087,11 +1083,12 @@ mod tests { for trigger in triggers { let event = procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: common_procedure::ProcedureId::random(), lifecycle_state: &state, trigger, event_type_filter: Arc::new(common_event_recorder::EventTypeFilter::All), + event_context: None, }) .unwrap(); assert_eq!(event.event_type(), "region_migration"); @@ -1111,61 +1108,98 @@ mod tests { to_peer: Peer::empty(2), region_ids: vec![RegionId::new(1024, 1)], timeout: Duration::from_secs(10), - trigger_reason: RegionMigrationTriggerReason::default(), - event_context: EventContext::default(), + trigger_reason: RegionMigrationTriggerReason::Unknown, }; // NOTES: Changes it will break backward compatibility. let serialized = r#"{"catalog":"greptime","schema":"public","from_peer":{"id":1,"addr":""},"to_peer":{"id":2,"addr":""},"region_id":4398046511105}"#; let deserialized: PersistentContext = serde_json::from_str(serialized).unwrap(); assert_eq!(persistent_ctx, deserialized); - assert_eq!(deserialized.event_context, EventContext::default()); } #[test] - fn test_new_persistent_context_keeps_event_context() { - let mut event_context = - EventContext::new(TriggerReason::AutoRebalance).with_protocol("internal"); - event_context - .extensions - .insert("scheduler".into(), serde_json::json!("load_based")); - - let context = PersistentContext::new( - vec![("greptime".into(), "public".into())], - Peer::empty(1), - Peer::empty(2), - vec![RegionId::new(1024, 1)], - Duration::from_secs(10), - event_context.clone(), + fn test_legacy_trigger_reason_survives_recovery_and_repersistence() { + let serialized = r#"{"persistent_ctx":{"catalog":"greptime","schema":"public","from_peer":{"id":1,"addr":""},"to_peer":{"id":2,"addr":""},"region_id":4398046511105,"trigger_reason":"Failover"},"state":{"region_migration_state":"RegionMigrationStart"}}"#; + let env = TestingEnv::new(); + let procedure = RegionMigrationProcedure::from_json( + serialized, + env.context_factory(), + RegionMigrationProcedureTracker::default(), + ) + .unwrap(); + assert_eq!( + procedure.context.persistent_ctx.trigger_reason, + RegionMigrationTriggerReason::Failover + ); + let repersisted = procedure.dump().unwrap(); + assert_eq!( + serde_json::from_str::(&repersisted).unwrap()["persistent_ctx"]["trigger_reason"], + "Failover" ); + let recovered = RegionMigrationProcedure::from_json( + &repersisted, + env.context_factory(), + RegionMigrationProcedureTracker::default(), + ) + .unwrap(); assert_eq!( - context.trigger_reason, + recovered.context.trigger_reason(None), + RegionMigrationTriggerReason::Failover + ); + let event_context = PersistentEventContext::new(TriggerReason::AutoRebalance); + assert_eq!( + recovered.context.trigger_reason(Some(&event_context)), RegionMigrationTriggerReason::AutoRebalance ); - assert_eq!(context.event_context, event_context); + + let state = common_procedure::ProcedureState::Running; + let event = recovered + .event(&EventContext { + procedure_id: common_procedure::ProcedureId::random(), + lifecycle_state: &state, + trigger: common_procedure::EventTrigger::Recovered, + event_type_filter: Arc::new(common_event_recorder::EventTypeFilter::All), + event_context: None, + }) + .unwrap(); + assert_eq!( + event.extra_rows().unwrap()[0].values[3].value_data, + Some(api::v1::value::ValueData::StringValue( + "Failover".to_string() + )) + ); } #[test] - fn test_normalize_trigger_reasons() { - let mut context = new_persistent_context(); - context.trigger_reason = RegionMigrationTriggerReason::Manual; - context.event_context = EventContext::default(); - context.normalize_trigger_reasons(); - assert_eq!(context.event_context.reason, TriggerReason::Manual); + fn test_migration_reason_uses_event_context() { + let env = TestingEnv::new(); + let context = env.context_factory().new_context(new_persistent_context()); + let event_context = PersistentEventContext::new(TriggerReason::RegionFailover); - context.trigger_reason = RegionMigrationTriggerReason::Unknown; - context.event_context = EventContext::new(TriggerReason::RegionFailover); - context.normalize_trigger_reasons(); assert_eq!( - context.trigger_reason, + context.trigger_reason(Some(&event_context)), RegionMigrationTriggerReason::Failover ); - context.trigger_reason = RegionMigrationTriggerReason::Manual; - context.event_context = EventContext::new(TriggerReason::RegionFailover); - context.normalize_trigger_reasons(); - assert_eq!(context.event_context.reason, TriggerReason::Manual); + let procedure = + RegionMigrationProcedure::new(new_persistent_context(), env.context_factory(), vec![]); + let state = common_procedure::ProcedureState::Running; + let event = procedure + .event(&EventContext { + procedure_id: common_procedure::ProcedureId::random(), + lifecycle_state: &state, + trigger: common_procedure::EventTrigger::Submitted, + event_type_filter: Arc::new(common_event_recorder::EventTypeFilter::All), + event_context: Some(&event_context), + }) + .unwrap(); + assert_eq!( + event.extra_rows().unwrap()[0].values[3].value_data, + Some(api::v1::value::ValueData::StringValue( + "Failover".to_string() + )) + ); } #[derive(Debug, Serialize, Deserialize, Default)] diff --git a/src/meta-srv/src/procedure/region_migration/close_downgraded_region.rs b/src/meta-srv/src/procedure/region_migration/close_downgraded_region.rs index 5cdcad4503b..29cd1804d03 100644 --- a/src/meta-srv/src/procedure/region_migration/close_downgraded_region.rs +++ b/src/meta-srv/src/procedure/region_migration/close_downgraded_region.rs @@ -40,7 +40,7 @@ impl State for CloseDowngradedRegion { async fn next( &mut self, ctx: &mut Context, - _procedure_ctx: &ProcedureContext, + procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)> { if let Err(err) = self.close_downgraded_leader_region(ctx).await { let downgrade_leader_datanode = &ctx.persistent_ctx.from_peer; @@ -52,7 +52,7 @@ impl State for CloseDowngradedRegion { ctx.persistent_ctx.region_ids, ctx.persistent_ctx.from_peer, ctx.persistent_ctx.to_peer, - ctx.persistent_ctx.trigger_reason, + ctx.trigger_reason(procedure_ctx.event_context.as_ref()), ctx.volatile_ctx.metrics, ); Ok((Box::new(RegionMigrationEnd), Status::done())) diff --git a/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs b/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs index c58c1f8cd1f..97c3a8e2f31 100644 --- a/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs +++ b/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs @@ -382,7 +382,6 @@ mod tests { use common_meta::key::table_route::TableRouteValue; use common_meta::key::test_utils::new_test_table_info; use common_meta::peer::Peer; - use common_meta::rpc::ddl::{EventContext, TriggerReason}; use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::wal_provider::RegionWalOptions; use store_api::storage::RegionId; @@ -391,7 +390,9 @@ mod tests { use super::*; use crate::error::Error; use crate::procedure::region_migration::test_util::{TestingEnv, new_procedure_context}; - use crate::procedure::region_migration::{ContextFactory, PersistentContext}; + use crate::procedure::region_migration::{ + ContextFactory, PersistentContext, RegionMigrationTriggerReason, + }; use crate::procedure::test_util::{ new_close_region_reply, new_downgrade_region_reply, send_mock_reply, }; @@ -403,7 +404,7 @@ mod tests { Peer::empty(2), vec![RegionId::new(1024, 1)], Duration::from_millis(1000), - EventContext::new(TriggerReason::Manual), + RegionMigrationTriggerReason::Unknown, ) } diff --git a/src/meta-srv/src/procedure/region_migration/manager.rs b/src/meta-srv/src/procedure/region_migration/manager.rs index 3de79735e3b..e1d719d9747 100644 --- a/src/meta-srv/src/procedure/region_migration/manager.rs +++ b/src/meta-srv/src/procedure/region_migration/manager.rs @@ -21,7 +21,7 @@ use std::time::Duration; use common_meta::key::table_info::TableInfoValue; use common_meta::key::table_route::TableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::RegionRoute; use common_procedure::{ProcedureId, ProcedureManagerRef, ProcedureWithId, watcher}; use common_telemetry::{error, info, warn}; @@ -479,6 +479,7 @@ impl RegionMigrationManager { procedure_guards: Vec, catalog_and_schema: Vec<(String, String)>, ) -> Result { + let event_context = PersistentEventContext::new(task.trigger_reason.to_trigger_reason()); let procedure = RegionMigrationProcedure::new( PersistentContext::new( catalog_and_schema, @@ -486,12 +487,13 @@ impl RegionMigrationManager { task.to_peer.clone(), task.region_ids.clone(), task.timeout, - EventContext::new(task.trigger_reason.to_trigger_reason()), + task.trigger_reason, ), self.context_factory.clone(), procedure_guards, ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); let procedure_manager = self.procedure_manager.clone(); @@ -575,6 +577,7 @@ impl RegionMigrationManager { timeout, trigger_reason, } = task.clone(); + let event_context = PersistentEventContext::new(trigger_reason.to_trigger_reason()); let procedure = RegionMigrationProcedure::new( PersistentContext::new( vec![(catalog_name, schema_name)], @@ -582,12 +585,13 @@ impl RegionMigrationManager { to_peer, vec![region_id], timeout, - EventContext::new(trigger_reason.to_trigger_reason()), + trigger_reason, ), self.context_factory.clone(), vec![guard], ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); let procedure_manager = self.procedure_manager.clone(); diff --git a/src/meta-srv/src/procedure/region_migration/migration_abort.rs b/src/meta-srv/src/procedure/region_migration/migration_abort.rs index f3ad8052de5..981c2a20252 100644 --- a/src/meta-srv/src/procedure/region_migration/migration_abort.rs +++ b/src/meta-srv/src/procedure/region_migration/migration_abort.rs @@ -41,7 +41,7 @@ impl State for RegionMigrationAbort { async fn next( &mut self, ctx: &mut Context, - _procedure_ctx: &ProcedureContext, + procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)> { warn!( "Region migration is aborted: {}, regions: {:?}, from_peer: {}, to_peer: {}, trigger_reason: {}, {}", @@ -49,7 +49,7 @@ impl State for RegionMigrationAbort { ctx.persistent_ctx.region_ids, ctx.persistent_ctx.from_peer, ctx.persistent_ctx.to_peer, - ctx.persistent_ctx.trigger_reason, + ctx.trigger_reason(procedure_ctx.event_context.as_ref()), ctx.volatile_ctx.metrics, ); error::MigrationAbortSnafu { diff --git a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs index 51d7054c7c5..0cb1131e545 100644 --- a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs +++ b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs @@ -47,9 +47,12 @@ impl State for OpenCandidateRegion { async fn next( &mut self, ctx: &mut Context, - _procedure_ctx: &ProcedureContext, + procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)> { - let instruction = self.build_open_region_instruction(ctx).await?; + let trigger_reason = ctx.trigger_reason(procedure_ctx.event_context.as_ref()); + let instruction = self + .build_open_region_instruction(ctx, trigger_reason) + .await?; let now = Instant::now(); self.open_candidate_region(ctx, instruction).await?; ctx.update_open_candidate_region_elapsed(now); @@ -67,11 +70,15 @@ impl OpenCandidateRegion { /// /// Abort(non-retry): /// - Datanode Table is not found. - async fn build_open_region_instruction(&self, ctx: &mut Context) -> Result { + async fn build_open_region_instruction( + &self, + ctx: &mut Context, + trigger_reason: RegionMigrationTriggerReason, + ) -> Result { let region_ids = ctx.persistent_ctx.region_ids.clone(); let from_peer_id = ctx.persistent_ctx.from_peer.id; let to_peer_id = ctx.persistent_ctx.to_peer.id; - let reason = match ctx.persistent_ctx.trigger_reason { + let reason = match trigger_reason { RegionMigrationTriggerReason::Failover => OpenRegionReason::RegionFailover, _ => OpenRegionReason::RegionMigration, }; @@ -280,7 +287,7 @@ mod tests { let mut ctx = env.context_factory().new_context(persistent_context); let err = state - .build_open_region_instruction(&mut ctx) + .build_open_region_instruction(&mut ctx, RegionMigrationTriggerReason::Unknown) .await .unwrap_err(); @@ -291,7 +298,7 @@ mod tests { #[tokio::test] async fn test_build_open_region_instruction_reason() { let state = OpenCandidateRegion; - let mut persistent_context = new_persistent_context(); + let persistent_context = new_persistent_context(); let from_peer_id = persistent_context.from_peer.id; let region_id = persistent_context.region_ids[0]; let env = TestingEnv::new(); @@ -314,7 +321,10 @@ mod tests { let mut ctx = env .context_factory() .new_context(persistent_context.clone()); - let instruction = state.build_open_region_instruction(&mut ctx).await.unwrap(); + let instruction = state + .build_open_region_instruction(&mut ctx, RegionMigrationTriggerReason::Unknown) + .await + .unwrap(); let open_regions = instruction.into_open_regions().unwrap(); assert_eq!( Some(OpenRegionReason::RegionMigration), @@ -325,9 +335,11 @@ mod tests { open_regions[0].requirements ); - persistent_context.trigger_reason = RegionMigrationTriggerReason::Failover; let mut ctx = env.context_factory().new_context(persistent_context); - let instruction = state.build_open_region_instruction(&mut ctx).await.unwrap(); + let instruction = state + .build_open_region_instruction(&mut ctx, RegionMigrationTriggerReason::Failover) + .await + .unwrap(); let open_regions = instruction.into_open_regions().unwrap(); assert_eq!( Some(OpenRegionReason::RegionFailover), diff --git a/src/meta-srv/src/procedure/region_migration/test_util.rs b/src/meta-srv/src/procedure/region_migration/test_util.rs index 70342608646..44180dc6f8c 100644 --- a/src/meta-srv/src/procedure/region_migration/test_util.rs +++ b/src/meta-srv/src/procedure/region_migration/test_util.rs @@ -27,7 +27,6 @@ use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::peer::Peer; use common_meta::region_keeper::{MemoryRegionKeeper, MemoryRegionKeeperRef}; -use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::RegionRoute; use common_meta::sequence::SequenceBuilder; use common_meta::state_store::KvStateStore; @@ -53,7 +52,8 @@ use crate::procedure::region_migration::open_candidate_region::OpenCandidateRegi use crate::procedure::region_migration::update_metadata::UpdateMetadata; use crate::procedure::region_migration::upgrade_candidate_region::UpgradeCandidateRegion; use crate::procedure::region_migration::{ - Context, ContextFactory, DefaultContextFactory, PersistentContext, State, VolatileContext, + Context, ContextFactory, DefaultContextFactory, PersistentContext, + RegionMigrationTriggerReason, State, VolatileContext, }; use crate::procedure::test_util::{MailboxContext, send_mock_reply}; use crate::service::mailbox::Channel; @@ -157,6 +157,7 @@ impl TestingEnv { ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, } } @@ -190,7 +191,7 @@ pub fn new_persistent_context(from: u64, to: u64, region_id: RegionId) -> Persis Peer::empty(to), vec![region_id], Duration::from_secs(10), - EventContext::default(), + RegionMigrationTriggerReason::Unknown, ) } @@ -577,5 +578,6 @@ pub fn new_procedure_context() -> ProcedureContext { ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, } } diff --git a/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs index 16d0270354b..86e1c35c10c 100644 --- a/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs +++ b/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs @@ -356,7 +356,6 @@ mod tests { use common_meta::key::topic_name::TopicNameKey; use common_meta::key::topic_region::{ReplayCheckpoint, TopicRegionKey, TopicRegionValue}; use common_meta::peer::Peer; - use common_meta::rpc::ddl::{EventContext, TriggerReason}; use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::wal_provider::RegionWalOptions; use common_wal::options::KafkaWalOptions; @@ -365,7 +364,9 @@ mod tests { use super::*; use crate::error::Error; use crate::procedure::region_migration::test_util::{TestingEnv, new_procedure_context}; - use crate::procedure::region_migration::{ContextFactory, PersistentContext}; + use crate::procedure::region_migration::{ + ContextFactory, PersistentContext, RegionMigrationTriggerReason, + }; use crate::procedure::test_util::{ new_close_region_reply, new_upgrade_region_reply, send_mock_reply, }; @@ -377,7 +378,7 @@ mod tests { Peer::empty(2), vec![RegionId::new(1024, 1)], Duration::from_millis(1000), - EventContext::new(TriggerReason::Manual), + RegionMigrationTriggerReason::Unknown, ) } diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index d8b585b28ee..ee8b3ee749b 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -44,12 +44,11 @@ use common_meta::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use common_meta::node_manager::NodeManagerRef; use common_meta::region_keeper::{MemoryRegionKeeperRef, OperatingRegionGuard}; use common_meta::region_registry::LeaderRegionRegistryRef; -use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::{RegionRoute, operating_leader_region_roles}; use common_meta::wal_provider::RegionWalOptions; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - BoxedProcedure, Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, + BoxedProcedure, Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, Procedure, ProcedureManagerRef, Result as ProcedureResult, Status, StringKey, }; @@ -104,8 +103,6 @@ pub struct PersistentContext { #[serde(default)] /// Records table-level partition metadata updated by this repartition. pub partition_metadata_update: Option, - #[serde(default)] - pub event_context: EventContext, } fn default_timeout() -> Duration { @@ -113,7 +110,7 @@ fn default_timeout() -> Duration { } impl PersistentContext { - /// Creates a new [PersistentContext] with the table, timeout and event context. + /// Creates a new [PersistentContext] with the given table name, table id and timeout. /// /// If the timeout is not provided, the default timeout will be used. pub fn new( @@ -124,7 +121,6 @@ impl PersistentContext { }: TableName, table_id: TableId, timeout: Option, - event_context: EventContext, ) -> Self { Self { catalog_name, @@ -136,7 +132,6 @@ impl PersistentContext { unknown_procedures: vec![], timeout: timeout.unwrap_or_else(default_timeout), partition_metadata_update: None, - event_context, } } @@ -797,10 +792,7 @@ impl Procedure for RepartitionProcedure { LockKey::new(self.context.persistent_ctx.lock_key()) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(REPARTITION_EVENT_TYPE) { return None; } @@ -869,7 +861,6 @@ impl RepartitionProcedureFactory for GcDisabledRepartitionProcedureFactory { _source: RepartitionSource, _to_exprs: Vec, _timeout: Option, - _event_context: EventContext, ) -> std::result::Result { Err(BoxedError::new( error::InvalidArgumentsSnafu { @@ -908,9 +899,8 @@ impl RepartitionProcedureFactory for DefaultRepartitionProcedureFactory { source: RepartitionSource, to_exprs: Vec, timeout: Option, - event_context: EventContext, ) -> std::result::Result { - let persistent_ctx = PersistentContext::new(table_name, table_id, timeout, event_context); + let persistent_ctx = PersistentContext::new(table_name, table_id, timeout); let from = match source { RepartitionSource::Partitioned { exprs, @@ -1143,7 +1133,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); Context::new( @@ -1175,7 +1164,6 @@ mod tests { }, vec![], None, - EventContext::default(), ) .err() .expect("GC-disabled factory must reject repartition"); @@ -1277,53 +1265,58 @@ mod tests { let all = Arc::new(EventTypeFilter::All); let submitted = procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: all.clone(), + event_context: None, }) .unwrap(); assert_eq!(submitted.event_type(), REPARTITION_EVENT_TYPE); assert_ne!(submitted.json_payload().unwrap(), serde_json::Value::Null); let allowed = procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: Arc::new(EventTypeFilter::Only(HashSet::from([ REPARTITION_EVENT_TYPE.to_string(), ]))), + event_context: None, }) .unwrap(); assert_eq!(allowed.event_type(), REPARTITION_EVENT_TYPE); let succeeded = procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Succeeded, event_type_filter: all, + event_context: None, }) .unwrap(); assert_eq!(succeeded.json_payload().unwrap(), serde_json::Value::Null); - let filtered = procedure.event(&EventRuntimeContext { + let filtered = procedure.event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: Arc::new(EventTypeFilter::Only(HashSet::from([ "another_event".to_string() ]))), + event_context: None, }); assert!(filtered.is_none()); - let empty = procedure.event(&EventRuntimeContext { + let empty = procedure.event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: Arc::new(EventTypeFilter::Only(HashSet::new())), + event_context: None, }); assert!(empty.is_none()); } @@ -1382,7 +1375,7 @@ mod tests { } #[test] - fn test_persistent_context_serde_defaults_for_new_fields() { + fn test_persistent_context_partition_metadata_update_serde_default() { let json = r#"{ "catalog_name":"test_catalog", "schema_name":"test_schema", @@ -1394,10 +1387,7 @@ mod tests { let persistent_ctx: PersistentContext = serde_json::from_str(json).unwrap(); - assert!(persistent_ctx.failed_procedures.is_empty()); - assert!(persistent_ctx.unknown_procedures.is_empty()); assert!(persistent_ctx.partition_metadata_update.is_none()); - assert_eq!(persistent_ctx.event_context, EventContext::default()); } #[tokio::test] @@ -1474,7 +1464,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![with_rollback_metadata( test_plan(table_id), @@ -1541,7 +1530,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![test_plan(table_id)]; let context = Context::new( @@ -1595,7 +1583,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); let failed_plan = test_plan(table_id); let failed_plan = with_rollback_metadata( @@ -1729,7 +1716,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![failed_plan, succeeded_plan.clone()]; persistent_ctx.failed_procedures = vec![ProcedureMeta { @@ -1839,7 +1825,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![plan.clone()]; persistent_ctx.unknown_procedures = vec![ProcedureMeta { @@ -1901,7 +1886,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![with_rollback_metadata( test_plan(table_id), @@ -2030,7 +2014,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![failed_merge_plan, succeeded_split_plan.clone()]; persistent_ctx.failed_procedures = vec![ProcedureMeta { diff --git a/src/meta-srv/src/procedure/repartition/allocate_region.rs b/src/meta-srv/src/procedure/repartition/allocate_region.rs index aa36e95be23..857ac946708 100644 --- a/src/meta-srv/src/procedure/repartition/allocate_region.rs +++ b/src/meta-srv/src/procedure/repartition/allocate_region.rs @@ -777,6 +777,7 @@ mod tests { concurrent_region_route, region_wal_options: test_region_wal_options(&[1, 2]), }), + event_context: None, }; let mut state = ExecutePlan; diff --git a/src/meta-srv/src/procedure/repartition/collect.rs b/src/meta-srv/src/procedure/repartition/collect.rs index 56d9a007a16..b71c3d6650f 100644 --- a/src/meta-srv/src/procedure/repartition/collect.rs +++ b/src/meta-srv/src/procedure/repartition/collect.rs @@ -136,7 +136,6 @@ mod tests { use common_error::mock::MockError; use common_error::status_code::StatusCode; - use common_meta::rpc::ddl::EventContext; use common_meta::test_util::MockDatanodeManager; use common_procedure::{ Context as ProcedureContext, ContextProvider, Error as ProcedureError, ProcedureId, @@ -194,7 +193,6 @@ mod tests { table::table_name::TableName::new("test_catalog", "test_schema", "test_table"), 1024, None, - EventContext::default(), ); let mut ctx = crate::procedure::repartition::Context::new( &ddl_ctx, @@ -229,7 +227,6 @@ mod tests { table::table_name::TableName::new("test_catalog", "test_schema", "test_table"), 1024, None, - EventContext::default(), ); let mut ctx = crate::procedure::repartition::Context::new( &ddl_ctx, @@ -249,6 +246,7 @@ mod tests { receiver: rx, inner: MockContextProvider::default(), }), + event_context: None, }; let mut state = Collect { inflight_procedures: vec![ProcedureMeta { diff --git a/src/meta-srv/src/procedure/repartition/gc_requirement.rs b/src/meta-srv/src/procedure/repartition/gc_requirement.rs index 11e82d294ac..1ae24e40316 100644 --- a/src/meta-srv/src/procedure/repartition/gc_requirement.rs +++ b/src/meta-srv/src/procedure/repartition/gc_requirement.rs @@ -244,6 +244,7 @@ mod tests { block_after_persist: true, persisted_tx: Some(persisted_tx), }), + context: Default::default(), }) .await .unwrap(); diff --git a/src/meta-srv/src/procedure/repartition/group.rs b/src/meta-srv/src/procedure/repartition/group.rs index a3c38b59557..63e05bcb150 100644 --- a/src/meta-srv/src/procedure/repartition/group.rs +++ b/src/meta-srv/src/procedure/repartition/group.rs @@ -38,8 +38,8 @@ use common_meta::peer::Peer; use common_meta::rpc::router::RegionRoute; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, - LockKey, Procedure, ProcedureId, Result as ProcedureResult, Status, StringKey, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, ProcedureId, Result as ProcedureResult, Status, StringKey, }; use common_telemetry::{error, info}; use serde::{Deserialize, Serialize}; @@ -262,10 +262,7 @@ impl Procedure for RepartitionGroupProcedure { LockKey::new(self.context.persistent_ctx.lock_key()) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(REPARTITION_GROUP_EVENT_TYPE) { return None; } @@ -634,9 +631,7 @@ mod tests { use common_event_recorder::EventTypeFilter; use common_meta::key::TableMetadataManager; use common_meta::kv_backend::test_util::MockKvBackendBuilder; - use common_procedure::{ - EventRuntimeContext, EventTrigger, Procedure, ProcedureId, ProcedureState, - }; + use common_procedure::{EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState}; use crate::error::Error; use crate::event::repartition::REPARTITION_GROUP_EVENT_TYPE; @@ -713,15 +708,16 @@ mod tests { context: env.create_context(new_persistent_context(1024, vec![], vec![])), }; let state = ProcedureState::Running; - let runtime_context = |event_type_filter| EventRuntimeContext { + let event_context = |event_type_filter| EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: Arc::new(event_type_filter), + event_context: None, }; let allowed = procedure - .event(&runtime_context(EventTypeFilter::Only(HashSet::from([ + .event(&event_context(EventTypeFilter::Only(HashSet::from([ REPARTITION_GROUP_EVENT_TYPE.to_string(), ])))) .unwrap(); @@ -729,14 +725,14 @@ mod tests { assert!( procedure - .event(&runtime_context(EventTypeFilter::Only(HashSet::from([ + .event(&event_context(EventTypeFilter::Only(HashSet::from([ "another_event".to_string(), ])))) .is_none() ); assert!( procedure - .event(&runtime_context(EventTypeFilter::Only(HashSet::new()))) + .event(&event_context(EventTypeFilter::Only(HashSet::new()))) .is_none() ); } diff --git a/src/meta-srv/src/procedure/repartition/test_util.rs b/src/meta-srv/src/procedure/repartition/test_util.rs index a6f2330b83a..2ce494ee7c9 100644 --- a/src/meta-srv/src/procedure/repartition/test_util.rs +++ b/src/meta-srv/src/procedure/repartition/test_util.rs @@ -24,7 +24,6 @@ use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::node_manager::NodeManagerRef; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::sequence::SequenceBuilder; use common_meta::test_util::new_ddl_context_with_kv_backend; @@ -103,6 +102,7 @@ impl TestingEnv { ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, } } @@ -225,7 +225,6 @@ pub fn new_parent_context( TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); ParentContext::new( @@ -266,6 +265,7 @@ pub fn procedure_context_with_receivers( receivers, inner: MockContextProvider::default(), }), + event_context: None, } } diff --git a/src/meta-srv/src/procedure/tests.rs b/src/meta-srv/src/procedure/tests.rs index dd855c66aa0..4f70dcca95c 100644 --- a/src/meta-srv/src/procedure/tests.rs +++ b/src/meta-srv/src/procedure/tests.rs @@ -29,7 +29,7 @@ use common_meta::ddl::test_util::create_table::{ }; use common_meta::key::table_route::{PhysicalTableRouteValue, TableRouteValue}; use common_meta::node_manager::NodeManagerRef; -use common_meta::rpc::ddl::{CreateTableTask, EventContext, QueryContext}; +use common_meta::rpc::ddl::CreateTableTask; use common_meta::rpc::router::{RegionRoute, find_leaders}; use common_procedure::Status; use store_api::storage::RegionId; @@ -91,8 +91,6 @@ fn create_table_task(table_name: Option<&str>) -> CreateTableTask { fn test_region_request_builder() { let mut procedure = CreateTableProcedure::new( create_table_task(None), - QueryContext::default(), - EventContext::default(), test_data::new_ddl_context(Arc::new(NodeClients::default())), ) .unwrap(); @@ -190,8 +188,6 @@ async fn test_on_datanode_create_regions() { let mut procedure = CreateTableProcedure::new( create_table_task(None), - QueryContext::default(), - EventContext::default(), test_data::new_ddl_context(node_manager), ) .unwrap(); @@ -264,12 +260,8 @@ async fn test_on_datanode_create_logical_regions() { .unwrap() .0; let _ = kv_backend.txn(physical_route_txn).await.unwrap(); - let mut procedure = CreateLogicalTablesProcedure::new( - vec![task1, task2, task3], - physical_table_id, - EventContext::default(), - ctx, - ); + let mut procedure = + CreateLogicalTablesProcedure::new(vec![task1, task2, task3], physical_table_id, ctx); let expected_created_regions = Arc::new(Mutex::new(HashMap::from([(1, 3), (2, 3), (3, 3)]))); diff --git a/src/meta-srv/src/procedure/wal_prune.rs b/src/meta-srv/src/procedure/wal_prune.rs index 128fc58c44d..5aca85a51a6 100644 --- a/src/meta-srv/src/procedure/wal_prune.rs +++ b/src/meta-srv/src/procedure/wal_prune.rs @@ -25,8 +25,8 @@ use common_meta::lock_key::RemoteWalLock; use common_meta::region_registry::LeaderRegionRegistryRef; use common_procedure::error::ToJsonSnafu; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, - LockKey, Procedure, ProcedureState, Result as ProcedureResult, Status, StringKey, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, ProcedureState, Result as ProcedureResult, Status, StringKey, }; use common_telemetry::{info, warn}; use manager::{WalPruneProcedureGuard, WalPruneProcedureTracker}; @@ -213,10 +213,7 @@ impl Procedure for WalPruneProcedure { LockKey::new(vec![lock_key]) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(WAL_PRUNE_EVENT_TYPE) { return None; } @@ -506,11 +503,12 @@ mod tests { let mut procedure = WalPruneProcedure::new(context, None, "test_topic".to_string(), 42, false); let running = ProcedureState::Running; - let runtime_context = |trigger, lifecycle_state, event_type_filter| EventRuntimeContext { + let runtime_context = |trigger, lifecycle_state, event_type_filter| EventContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, event_type_filter: Arc::new(event_type_filter), + event_context: None, }; for trigger in [EventTrigger::Submitted, EventTrigger::Recovered] { diff --git a/src/meta-srv/src/service/procedure.rs b/src/meta-srv/src/service/procedure.rs index b20d7f64775..1b5c4c42dbc 100644 --- a/src/meta-srv/src/service/procedure.rs +++ b/src/meta-srv/src/service/procedure.rs @@ -107,6 +107,7 @@ impl procedure_service_server::ProcedureService for Metasrv { .submit_ddl_task( &ExecutorContext { tracing_context: Some(header.tracing_context), + ..Default::default() }, SubmitDdlTaskRequest { query_context, diff --git a/src/standalone/src/procedure.rs b/src/standalone/src/procedure.rs index 85973e9f643..ec63b11703b 100644 --- a/src/standalone/src/procedure.rs +++ b/src/standalone/src/procedure.rs @@ -20,7 +20,6 @@ use common_meta::ddl::DdlContext; use common_meta::ddl_manager::{RepartitionProcedureFactory, RepartitionSource}; use common_meta::key::runtime_switch::RuntimeSwitchManager; use common_meta::kv_backend::KvBackendRef; -use common_meta::rpc::ddl::EventContext; use common_meta::state_store::KvStateStore; use common_procedure::local::{EventRecorderHandle, LocalManager, ManagerConfig}; use common_procedure::options::ProcedureConfig; @@ -74,7 +73,6 @@ impl RepartitionProcedureFactory for StandaloneRepartitionProcedureFactory { _source: RepartitionSource, _to_exprs: Vec, _timeout: Option, - _event_context: EventContext, ) -> std::result::Result { Err(BoxedError::new(NoSupportRepartitionProcedureSnafu.build())) } diff --git a/tests-integration/src/tests/gc.rs b/tests-integration/src/tests/gc.rs index 5bc20e2e2cd..badaef5f7ac 100644 --- a/tests-integration/src/tests/gc.rs +++ b/tests-integration/src/tests/gc.rs @@ -17,7 +17,6 @@ use std::sync::Arc; use std::time::Duration; use common_meta::key::TableMetadataManagerRef; -use common_meta::rpc::ddl::EventContext; use common_procedure::ProcedureWithId; use common_telemetry::info; use common_test_util::recordbatch::check_output_stream; @@ -214,7 +213,6 @@ async fn test_gc_basic(store_type: &StorageType) { false, // full_file_listing Duration::from_secs(10), // timeout Default::default(), - EventContext::default(), ); // Submit the procedure to the procedure manager diff --git a/tests-integration/tests/gc_event.rs b/tests-integration/tests/gc_event.rs index b62ba1d3649..9ecedc77775 100644 --- a/tests-integration/tests/gc_event.rs +++ b/tests-integration/tests/gc_event.rs @@ -15,7 +15,7 @@ use std::sync::Arc; use std::time::Duration; -use common_meta::rpc::ddl::{EventContext, TriggerReason}; +use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; use common_procedure::{ProcedureId, ProcedureWithId, watcher}; use common_test_util::temp_dir::create_temp_dir; use meta_srv::gc::{BatchGcProcedure, GcSchedulerOptions}; @@ -129,7 +129,6 @@ async fn test_batch_gc_event() { false, Duration::from_secs(10), Default::default(), - EventContext::new(TriggerReason::Manual), ); let procedure_id = ProcedureId::parse_str("00000000-0000-0000-0000-00000000bac0").unwrap(); let mut watcher = cluster @@ -138,6 +137,9 @@ async fn test_batch_gc_event() { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: common_procedure::ProcedureContext::from_event_context( + PersistentEventContext::new(TriggerReason::Manual), + ), }) .await .unwrap(); diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index 3dac860066b..1990cc2c43d 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -19,7 +19,6 @@ use std::time::Duration; use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_error::root_source; use common_meta::key::table_name::TableNameKey; -use common_meta::rpc::ddl::EventContext; use common_procedure::{ProcedureWithId, watcher}; use common_query::Output; use common_telemetry::info; @@ -817,7 +816,6 @@ async fn trigger_table_gc(metasrv: &Arc, table_name: &str) { false, // full_file_listing Duration::from_secs(10), // timeout Default::default(), - EventContext::default(), ); // Submit the procedure to the procedure manager diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index 5c23e101de2..5194640880c 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -377,7 +377,7 @@ async fn test_table_ddl_procedure_events() { .ddl_manager() .submit_undrop_table_task( common_meta::rpc::ddl::UndropTableTask { table_id }, - common_meta::rpc::ddl::EventContext::new( + common_meta::rpc::ddl::PersistentEventContext::new( common_meta::rpc::ddl::TriggerReason::Manual, ), ) @@ -419,7 +419,7 @@ async fn test_table_ddl_procedure_events() { .ddl_manager() .submit_purge_dropped_table_task( common_meta::rpc::ddl::PurgeDroppedTableTask { table_id }, - common_meta::rpc::ddl::EventContext::new( + common_meta::rpc::ddl::PersistentEventContext::new( common_meta::rpc::ddl::TriggerReason::Manual, ), ) diff --git a/tests/conf/frontend-test.toml.template b/tests/conf/frontend-test.toml.template index 25d44ff6e4f..0042716bd80 100644 --- a/tests/conf/frontend-test.toml.template +++ b/tests/conf/frontend-test.toml.template @@ -1,3 +1,6 @@ [grpc] bind_addr = "{addrs.grpc_addr}" server_addr = "{addrs.grpc_addr}" + +[event_recorder] +event_types = [] From 943eee852f292bfffd7240867dc27248216c353f Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 11 Aug 2026 20:50:50 +0800 Subject: [PATCH 051/126] feat(event): record admin function executions (#8835) * feat(event): record admin function executions Signed-off-by: WenyXu * fix(event): handle admin function recording edge cases Signed-off-by: WenyXu * feat(event): record actor for admin functions Signed-off-by: WenyXu * fix(event): preserve admin function event values Signed-off-by: WenyXu * fix(event): preserve non-finite admin results Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 1 + config/config.md | 3 +- config/frontend.example.toml | 4 + config/standalone.example.toml | 2 +- src/cmd/tests/load_config_test.rs | 5 + src/common/event-recorder/src/event_table.rs | 63 ++ src/frontend/src/frontend.rs | 6 +- src/frontend/src/instance/builder.rs | 13 +- src/operator/Cargo.toml | 1 + src/operator/src/error.rs | 4 + src/operator/src/statement.rs | 16 + src/operator/src/statement/admin.rs | 116 +++- src/operator/src/statement/admin/event.rs | 569 ++++++++++++++++++ src/operator/src/statement/admin/layer.rs | 440 ++++++++++++++ .../tests/admin_function_event.rs | 193 ++++++ tests-integration/tests/main.rs | 1 + 16 files changed, 1421 insertions(+), 16 deletions(-) create mode 100644 src/operator/src/statement/admin/event.rs create mode 100644 src/operator/src/statement/admin/layer.rs create mode 100644 tests-integration/tests/admin_function_event.rs diff --git a/Cargo.lock b/Cargo.lock index 18248bd442d..4d0c271882b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9892,6 +9892,7 @@ dependencies = [ "common-catalog", "common-datasource", "common-error", + "common-event-recorder", "common-frontend", "common-function", "common-grpc", diff --git a/config/config.md b/config/config.md index 40272ddcfda..c7737a164b8 100644 --- a/config/config.md +++ b/config/config.md @@ -234,7 +234,7 @@ | `tracing.tokio_console_addr` | String | Unset | The tokio console address. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `memory` | -- | -- | The memory options. | | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | @@ -372,6 +372,7 @@ | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event type: `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | ### Metasrv diff --git a/config/frontend.example.toml b/config/frontend.example.toml index a5907545eba..260984a4ab3 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -405,3 +405,7 @@ enable_heap_profiling = true [event_recorder] ## TTL for the events table that will be used to store the events. Default is `90d`. ttl = "90d" +## Event types to record. Current available event type: `admin_function`. +## When omitted, all current and future event types are recorded. +## Set to an empty array to disable event recording. +#+ event_types = ["admin_function"] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index d256d681f9f..c8dff4944c1 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -922,7 +922,7 @@ ttl = "90d" ## `alter_database`, `drop_database`, `create_flow`, `drop_flow`, ## `create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`, ## `drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`, -## `create_view`, `drop_view`. +## `create_view`, `drop_view`, `admin_function`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. #+ event_types = ["create_database"] diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index aa1d98cf4b9..fd32680f4c6 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -445,6 +445,11 @@ fn test_load_event_types_from_env() { .component .event_recorder .event_types, + GreptimeOptions::::load_layered_options(None, env_prefix) + .unwrap() + .component + .event_recorder + .event_types, ] { assert!(event_types.allows("region_migration")); assert!(!event_types.allows("other_event")); diff --git a/src/common/event-recorder/src/event_table.rs b/src/common/event-recorder/src/event_table.rs index 0ec38c8eabc..1ea4369a00a 100644 --- a/src/common/event-recorder/src/event_table.rs +++ b/src/common/event-recorder/src/event_table.rs @@ -88,6 +88,9 @@ pub const TIMESTAMP_COLUMN: EventTableColumn = EventTableColumn::new( ColumnDataType::TimestampNanosecond, SemanticType::Timestamp, ); +/// The effective database user that initiated an event. +pub const ACTOR_COLUMN: EventTableColumn = + EventTableColumn::new("actor", ColumnDataType::String, SemanticType::Field); /// The canonical procedure identifier envelope column. pub const PROCEDURE_ID_COLUMN: EventTableColumn = EventTableColumn::new("procedure_id", ColumnDataType::String, SemanticType::Field); @@ -144,6 +147,25 @@ pub const LATEST_OFFSET_COLUMN: EventTableColumn = pub const GC_REPORT_COLUMN: EventTableColumn = EventTableColumn::json_binary("gc_report", ColumnDataType::Binary, SemanticType::Field); +/// The canonical ADMIN function name field. +pub const ADMIN_FUNCTION_NAME_COLUMN: EventTableColumn = EventTableColumn::new( + "admin_function_name", + ColumnDataType::String, + SemanticType::Field, +); +/// The canonical ADMIN function execution status field. +pub const ADMIN_FUNCTION_STATUS_COLUMN: EventTableColumn = EventTableColumn::new( + "admin_function_status", + ColumnDataType::String, + SemanticType::Field, +); +/// The canonical ADMIN function output field. +pub const ADMIN_FUNCTION_OUTPUT_COLUMN: EventTableColumn = EventTableColumn::json_binary( + "admin_function_output", + ColumnDataType::Binary, + SemanticType::Field, +); + /// The canonical table name field for table DDL events. pub const TABLE_NAME_COLUMN: EventTableColumn = EventTableColumn::new("table_name", ColumnDataType::String, SemanticType::Field); @@ -389,6 +411,47 @@ mod tests { ); } + #[test] + fn admin_function_schema_preserves_names_types_semantics_extensions_and_order() { + assert_eq!( + column_schemas([ + &ACTOR_COLUMN, + &ADMIN_FUNCTION_NAME_COLUMN, + &ADMIN_FUNCTION_STATUS_COLUMN, + &ADMIN_FUNCTION_OUTPUT_COLUMN, + ]), + vec![ + ColumnSchema { + column_name: "actor".to_string(), + datatype: ColumnDataType::String.into(), + semantic_type: SemanticType::Field.into(), + ..Default::default() + }, + ColumnSchema { + column_name: "admin_function_name".to_string(), + datatype: ColumnDataType::String.into(), + semantic_type: SemanticType::Field.into(), + ..Default::default() + }, + ColumnSchema { + column_name: "admin_function_status".to_string(), + datatype: ColumnDataType::String.into(), + semantic_type: SemanticType::Field.into(), + ..Default::default() + }, + ColumnSchema { + column_name: "admin_function_output".to_string(), + datatype: ColumnDataType::Binary.into(), + semantic_type: SemanticType::Field.into(), + datatype_extension: Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), + }), + ..Default::default() + }, + ] + ); + } + #[test] fn table_dimension_schema_preserves_names_types_semantics_and_order() { assert_eq!( diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 7a09b40e736..b8fea7b6043 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -116,7 +116,11 @@ impl Default for FrontendOptions { impl Configurable for FrontendOptions { fn env_list_keys() -> Option<&'static [&'static str]> { - Some(&["heartbeat_env_vars", "meta_client.metasrv_addrs"]) + Some(&[ + "heartbeat_env_vars", + "meta_client.metasrv_addrs", + "event_recorder.event_types", + ]) } } diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 85a2b8ff183..c49a8394beb 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -23,7 +23,7 @@ use catalog::process_manager::ProcessManagerRef; use catalog::system_schema::semantic_graph::EntityGraphProviderRef; use common_base::Plugins; use common_datasource::object_store::LocalFileAccess; -use common_event_recorder::EventRecorderImpl; +use common_event_recorder::{EventRecorderImpl, EventRecorderRef}; use common_meta::cache::{LayeredCacheRegistryRef, TableRouteCacheRef}; use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator}; use common_meta::key::TableMetadataManager; @@ -40,8 +40,8 @@ use operator::request::Requester; #[cfg(feature = "enterprise")] use operator::statement::CreateDatabaseHandlerRef; use operator::statement::{ - ExecutorConfigureContext, StatementExecutor, StatementExecutorConfiguratorRef, - StatementExecutorRef, + AdminEventRecorderHandle, AdminFunctionRecordingLayer, ExecutorConfigureContext, + StatementExecutor, StatementExecutorConfiguratorRef, StatementExecutorRef, }; use operator::table::TableMutationOperator; use partition::cache::PartitionInfoCacheRef; @@ -323,6 +323,10 @@ impl FrontendBuilder { statement_executor }; + let admin_event_recorder = AdminEventRecorderHandle::default(); + let statement_executor = statement_executor.with_admin_function_layer(Arc::new( + AdminFunctionRecordingLayer::new(admin_event_recorder.clone()), + )); let statement_executor = Arc::new(statement_executor); let pipeline_operator = Arc::new(PipelineOperator::new( @@ -337,13 +341,14 @@ impl FrontendBuilder { let slow_query_recorder = Arc::new(EventRecorderImpl::new(Box::new( EventHandlerImpl::new(statement_executor.clone(), self.options.slow_query.ttl), ))); - let event_recorder = Arc::new(EventRecorderImpl::with_event_type_filter( + let event_recorder: EventRecorderRef = Arc::new(EventRecorderImpl::with_event_type_filter( Box::new(EventHandlerImpl::new( statement_executor.clone(), self.options.event_recorder.ttl, )), self.options.event_recorder.event_types.clone(), )); + admin_event_recorder.install(&event_recorder); Ok(Instance { frontend_peer_addr, diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index bd34e21ed24..087ea461f10 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -33,6 +33,7 @@ common-base.workspace = true common-catalog.workspace = true common-datasource.workspace = true common-error.workspace = true +common-event-recorder.workspace = true common-frontend.workspace = true common-function.workspace = true common-grpc.workspace = true diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index b5e76f0f2cc..f1e94b2c616 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -65,6 +65,9 @@ pub enum Error { location: Location, }, + #[snafu(display("Admin function execution was cancelled"))] + AdminFunctionCancelled, + #[snafu(display("Expected {expected} args, but actual {actual}"))] FunctionArityMismatch { expected: usize, actual: usize }, @@ -1023,6 +1026,7 @@ impl ErrorExt for Error { | Error::EncodeJson { .. } | Error::DeserializePartitionExpr { .. } | Error::SerializePartitionExpr { .. } => StatusCode::Unexpected, + Error::AdminFunctionCancelled => StatusCode::Cancelled, Error::ViewNotFound { .. } | Error::ViewInfoNotFound { .. } | Error::TableNotFound { .. } => StatusCode::TableNotFound, diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index fea54bb27c3..478f06e83f7 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -81,6 +81,11 @@ use table::requests::{CopyDatabaseRequest, CopyDirection, CopyQueryToRequest, Co use table::table_name::TableName; use table::table_reference::TableReference; +pub use self::admin::{ + AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef, + AdminFunctionRecordingLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, + AdminFunctionServiceRef, +}; use self::set::{ set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding, }; @@ -139,6 +144,7 @@ pub struct StatementExecutor { inserter: InserterRef, process_manager: Option, origin_frontend_addr: String, + admin_function_service: AdminFunctionServiceRef, pub(crate) local_file_access: LocalFileAccess, #[cfg(feature = "enterprise")] create_database_handler: Option, @@ -180,6 +186,7 @@ impl StatementExecutor { origin_frontend_addr: String, local_file_access: LocalFileAccess, ) -> Self { + let admin_function_service = admin::new_admin_function_service(query_engine.clone()); Self { catalog_manager, query_engine, @@ -192,6 +199,7 @@ impl StatementExecutor { inserter, process_manager, origin_frontend_addr, + admin_function_service, local_file_access, #[cfg(feature = "enterprise")] create_database_handler: None, @@ -200,6 +208,14 @@ impl StatementExecutor { } } + /// Adds a layer around the ADMIN function execution service. + /// + /// The last added layer is the outermost layer. + pub fn with_admin_function_layer(mut self, layer: AdminFunctionLayerRef) -> Self { + self.admin_function_service = layer.layer(self.admin_function_service); + self + } + #[cfg(feature = "enterprise")] pub fn with_trigger_querier(mut self, querier: TriggerQuerierRef) -> Self { self.trigger_querier = Some(querier); diff --git a/src/operator/src/statement/admin.rs b/src/operator/src/statement/admin.rs index 7494977b214..2b4a909e3cd 100644 --- a/src/operator/src/statement/admin.rs +++ b/src/operator/src/statement/admin.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod event; +mod layer; + use std::sync::Arc; use common_function::function::FunctionContext; @@ -28,6 +31,10 @@ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::value::Value; use datatypes::vectors::VectorRef; +pub use layer::{ + AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef, + AdminFunctionRecordingLayer, +}; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use sql::ast::{Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Value as SqlValue}; @@ -38,14 +45,49 @@ use crate::statement::StatementExecutor; const DUMMY_COLUMN: &str = ""; -impl StatementExecutor { - /// Execute the [`Admin`] statement and returns the output. - #[tracing::instrument(skip_all)] - pub(super) async fn execute_admin_command( - &self, - stmt: Admin, - query_ctx: QueryContextRef, - ) -> Result { +/// A request to execute one ADMIN function statement. +#[derive(Clone)] +pub struct AdminFunctionRequest { + /// The parsed ADMIN statement to execute. + pub statement: Admin, + /// The query context of the request. + pub query_ctx: QueryContextRef, +} + +/// The client output and immediate typed result of one ADMIN function execution. +pub struct AdminFunctionResponse { + /// The output returned to the client. + pub output: Output, + /// The typed immediate result exposed to outer layers. + pub immediate_result: Option, +} + +/// Executes an ADMIN function request. +#[async_trait::async_trait] +pub trait AdminFunctionService: Send + Sync { + /// Executes an ADMIN function request. + async fn call(&self, request: AdminFunctionRequest) -> Result; +} + +/// A shared ADMIN function service. +pub type AdminFunctionServiceRef = Arc; + +#[derive(Clone)] +struct CoreAdminFunctionService { + query_engine: query::QueryEngineRef, +} + +impl CoreAdminFunctionService { + fn new(query_engine: query::QueryEngineRef) -> Self { + Self { query_engine } + } + + async fn execute(&self, request: AdminFunctionRequest) -> Result { + let AdminFunctionRequest { + statement: stmt, + query_ctx, + } = request; + let Admin::Func(func) = &stmt; // the function name should be in lower case. let func_name = func.name.to_string().to_lowercase(); @@ -141,6 +183,7 @@ impl StatementExecutor { (&result_columnar).try_into().context(CastSnafu)?; let result_vector: VectorRef = result_columnar.try_into_vector(1).context(CastSnafu)?; + let immediate_result = immediate_result(&result_vector); let column_schemas = vec![ColumnSchema::new( // Use statement as the result column name @@ -154,7 +197,46 @@ impl StatementExecutor { let batches = RecordBatches::try_new(schema, vec![batch]).context(error::BuildRecordBatchSnafu)?; - Ok(Output::new_with_record_batches(batches)) + Ok(AdminFunctionResponse { + output: Output::new_with_record_batches(batches), + immediate_result, + }) + } +} + +fn immediate_result(result_vector: &VectorRef) -> Option { + (!result_vector.is_empty()).then(|| result_vector.get(0)) +} + +#[async_trait::async_trait] +impl AdminFunctionService for CoreAdminFunctionService { + async fn call(&self, request: AdminFunctionRequest) -> Result { + self.execute(request).await + } +} + +/// Creates the core ADMIN function service. +pub(crate) fn new_admin_function_service( + query_engine: query::QueryEngineRef, +) -> AdminFunctionServiceRef { + Arc::new(CoreAdminFunctionService::new(query_engine)) +} + +impl StatementExecutor { + /// Executes the [`Admin`] statement and returns the output. + #[tracing::instrument(skip_all)] + pub(crate) async fn execute_admin_command( + &self, + stmt: Admin, + query_ctx: QueryContextRef, + ) -> Result { + self.admin_function_service + .call(AdminFunctionRequest { + statement: stmt, + query_ctx, + }) + .await + .map(|response| response.output) } } @@ -306,3 +388,19 @@ fn try_get_data_type_for_sql_value(value: &SqlValue) -> Result { .fail(), } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datatypes::vectors::{Int32Vector, VectorRef}; + + use crate::statement::admin::immediate_result; + + #[test] + fn empty_admin_function_result_has_no_immediate_value() { + let result_vector: VectorRef = Arc::new(Int32Vector::from(vec![])); + + assert_eq!(immediate_result(&result_vector), None); + } +} diff --git a/src/operator/src/statement/admin/event.rs b/src/operator/src/statement/admin/event.rs new file mode 100644 index 00000000000..bafa2ad8824 --- /dev/null +++ b/src/operator/src/statement/admin/event.rs @@ -0,0 +1,569 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::any::Any; + +use api::v1::value::ValueData; +use api::v1::{ColumnSchema, Row}; +use common_event_recorder::Event; +use common_event_recorder::error::Result as EventResult; +use common_event_recorder::event_table::{ + ACTOR_COLUMN, ADMIN_FUNCTION_NAME_COLUMN, ADMIN_FUNCTION_OUTPUT_COLUMN, + ADMIN_FUNCTION_STATUS_COLUMN, column_schemas, jsonb_value, +}; +use datatypes::value::Value; +use serde_json::{Value as JsonValue, json}; +use sql::ast::{Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Value as SqlValue}; +use sql::statements::admin::Admin; + +use crate::error::Error; +use crate::statement::admin::AdminFunctionRequest; + +/// Event type emitted for an ADMIN function execution. +pub(crate) const ADMIN_FUNCTION_EVENT_TYPE: &str = "admin_function"; + +const PAYLOAD_VERSION: u8 = 1; +const VERSION_KEY: &str = "version"; +const ARGUMENTS_KEY: &str = "arguments"; +const ARGUMENT_TYPE_KEY: &str = "type"; +const ARGUMENT_VALUE_KEY: &str = "value"; +const LIST_ARGUMENT_TYPE: &str = "list"; +const SUBQUERY_ARGUMENT_TYPE: &str = "subquery"; +const RESULT_KEY: &str = "result"; +const ERROR_KEY: &str = "error"; +const SUCCEEDED_STATUS: &str = "Succeeded"; +const FAILED_STATUS: &str = "Failed"; +const UNSUPPORTED: &str = ""; + +/// The ADMIN function input captured before execution. +#[derive(Debug, Clone)] +pub(crate) struct AdminFunctionEventInput { + actor: String, + function: String, + arguments: Option, +} + +impl AdminFunctionEventInput { + /// Creates an event input from an ADMIN function request. + pub(crate) fn from_request(request: &AdminFunctionRequest) -> Self { + let Admin::Func(function) = &request.statement; + let function_name = function.name.to_string().to_lowercase(); + + let arguments = match &function.args { + FunctionArguments::List(arguments) => Some(json!({ + (ARGUMENT_TYPE_KEY): LIST_ARGUMENT_TYPE, + (ARGUMENT_VALUE_KEY): arguments.args.iter().map(argument_to_json).collect::>(), + })), + FunctionArguments::Subquery(query) => Some(json!({ + (ARGUMENT_TYPE_KEY): SUBQUERY_ARGUMENT_TYPE, + (ARGUMENT_VALUE_KEY): query.to_string(), + })), + FunctionArguments::None => None, + }; + + Self { + actor: request.query_ctx.current_user().username().to_string(), + function: function_name, + arguments, + } + } +} + +/// An event describing the outcome of an ADMIN function execution. +#[derive(Debug)] +pub(crate) struct AdminFunctionEvent { + actor: String, + function_name: String, + status: &'static str, + output: JsonValue, + payload: JsonValue, +} + +impl AdminFunctionEvent { + /// Creates a successful ADMIN function event. + pub(crate) fn success(input: AdminFunctionEventInput, result: Option<&Value>) -> Self { + let AdminFunctionEventInput { + actor, + function, + arguments, + } = input; + Self { + actor, + function_name: function, + status: SUCCEEDED_STATUS, + output: result.map_or_else( + || json!({}), + |result| { + json!({ + (RESULT_KEY): value_to_json(result), + }) + }, + ), + payload: input_payload(arguments), + } + } + + /// Creates a failed ADMIN function event with the debug representation of the error. + pub(crate) fn failure(input: AdminFunctionEventInput, error: &Error) -> Self { + let AdminFunctionEventInput { + actor, + function, + arguments, + } = input; + Self { + actor, + function_name: function, + status: FAILED_STATUS, + output: json!({ + (ERROR_KEY): format!("{error:?}"), + }), + payload: input_payload(arguments), + } + } + + /// Creates a failed ADMIN function event for a cancelled execution. + pub(crate) fn cancelled(input: AdminFunctionEventInput) -> Self { + Self::failure(input, &Error::AdminFunctionCancelled) + } +} + +impl Event for AdminFunctionEvent { + fn event_type(&self) -> &str { + ADMIN_FUNCTION_EVENT_TYPE + } + + fn json_payload(&self) -> EventResult { + Ok(self.payload.clone()) + } + + fn extra_schema(&self) -> Vec { + column_schemas([ + &ACTOR_COLUMN, + &ADMIN_FUNCTION_NAME_COLUMN, + &ADMIN_FUNCTION_STATUS_COLUMN, + &ADMIN_FUNCTION_OUTPUT_COLUMN, + ]) + } + + fn extra_rows(&self) -> EventResult> { + Ok(vec![Row { + values: vec![ + ValueData::StringValue(self.actor.clone()).into(), + ValueData::StringValue(self.function_name.clone()).into(), + ValueData::StringValue(self.status.to_string()).into(), + jsonb_value(&self.output), + ], + }]) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +fn input_payload(arguments: Option) -> JsonValue { + let mut payload = json!({ VERSION_KEY: PAYLOAD_VERSION }); + if let Some(arguments) = arguments { + payload[ARGUMENTS_KEY] = arguments; + } + payload +} + +fn argument_to_json(argument: &FunctionArg) -> JsonValue { + match argument { + FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) => { + sql_value_to_json(&value.value) + } + _ => JsonValue::String(argument.to_string()), + } +} + +fn sql_value_to_json(value: &SqlValue) -> JsonValue { + match value { + SqlValue::Number(value, _) => { + serde_json::from_str(value).unwrap_or_else(|_| JsonValue::String(value.clone())) + } + SqlValue::Boolean(value) => JsonValue::Bool(*value), + SqlValue::Null => JsonValue::Null, + SqlValue::SingleQuotedString(value) | SqlValue::DoubleQuotedString(value) => { + JsonValue::String(value.clone()) + } + SqlValue::Placeholder(value) => JsonValue::String(value.clone()), + _ => value + .clone() + .into_string() + .map(JsonValue::String) + .unwrap_or_else(|| JsonValue::String(value.to_string())), + } +} + +fn value_to_json(value: &Value) -> JsonValue { + match value { + Value::Float32(value) if !value.0.is_finite() => JsonValue::String(value.to_string()), + Value::Float64(value) if !value.0.is_finite() => JsonValue::String(value.to_string()), + _ => JsonValue::try_from(value.clone()) + .unwrap_or_else(|_| JsonValue::String(UNSUPPORTED.to_string())), + } +} + +#[cfg(test)] +mod tests { + use api::v1::Row; + use api::v1::value::ValueData; + use common_event_recorder::Event; + use common_event_recorder::event_table::{ + ACTOR_COLUMN, ADMIN_FUNCTION_NAME_COLUMN, ADMIN_FUNCTION_OUTPUT_COLUMN, + ADMIN_FUNCTION_STATUS_COLUMN, column_schemas, jsonb_value, + }; + use datatypes::value::Value; + use serde_json::json; + use session::context::QueryContext; + use sql::ast::{ + Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, Value as SqlValue, + }; + use sql::dialect::GreptimeDbDialect; + use sql::parser::{ParseOptions, ParserContext}; + use sql::statements::admin::Admin; + use sql::statements::statement::Statement; + use sqlparser::ast::{DollarQuotedString, FunctionArgOperator, FunctionArgumentList}; + + use crate::error::Error; + use crate::statement::admin::AdminFunctionRequest; + use crate::statement::admin::event::{ + ADMIN_FUNCTION_EVENT_TYPE, AdminFunctionEvent, AdminFunctionEventInput, sql_value_to_json, + value_to_json, + }; + + fn request(sql: &str) -> AdminFunctionRequest { + let Statement::Admin(statement) = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .remove(0) + else { + panic!("expected ADMIN statement") + }; + AdminFunctionRequest { + statement, + query_ctx: QueryContext::arc(), + } + } + + fn request_with_arguments(arguments: FunctionArguments) -> AdminFunctionRequest { + let mut request = request("ADMIN plugin_function()"); + let Admin::Func(function) = &mut request.statement; + function.args = arguments; + request + } + + fn assert_admin_columns( + event: &AdminFunctionEvent, + function_name: &str, + status: &str, + output: serde_json::Value, + ) { + assert_eq!(event.event_type(), ADMIN_FUNCTION_EVENT_TYPE); + assert_eq!( + event.extra_schema(), + column_schemas([ + &ACTOR_COLUMN, + &ADMIN_FUNCTION_NAME_COLUMN, + &ADMIN_FUNCTION_STATUS_COLUMN, + &ADMIN_FUNCTION_OUTPUT_COLUMN, + ]) + ); + assert_eq!( + event.extra_rows().unwrap(), + vec![Row { + values: vec![ + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue(function_name.to_string()).into(), + ValueData::StringValue(status.to_string()).into(), + jsonb_value(&output), + ], + }] + ); + } + + #[test] + fn success_uses_separate_output_columns() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN flush_table('greptime.public.demo')", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(3))); + + assert_admin_columns(&event, "flush_table", "Succeeded", json!({"result": 3})); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["greptime.public.demo"]}, + }) + ); + } + + #[test] + fn successful_null_is_explicit() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN migrate_region(NULL, NULL, NULL)", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::Null)); + + assert_admin_columns( + &event, + "migrate_region", + "Succeeded", + json!({"result": null}), + ); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": [null, null, null]}, + }) + ); + } + + #[test] + fn unknown_function_values_are_recorded() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN plugin_function('plugin-value', NULL)", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::from("plugin-result"))); + + assert_admin_columns( + &event, + "plugin_function", + "Succeeded", + json!({"result": "plugin-result"}), + ); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["plugin-value", null]}, + }) + ); + } + + #[test] + fn failure_payload_has_error_only() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN flush_table('greptime.public.demo')", + )); + let error = Error::BuildAdminFunctionArgs { + msg: "invalid input".to_string(), + }; + let expected_error = format!("{error:?}"); + let event = AdminFunctionEvent::failure(input, &error); + + assert_admin_columns( + &event, + "flush_table", + "Failed", + json!({"error": expected_error}), + ); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["greptime.public.demo"]}, + }) + ); + } + + #[test] + fn cancellation_uses_debug_error_format() { + let input = AdminFunctionEventInput::from_request(&request("ADMIN flush_table('demo')")); + let event = AdminFunctionEvent::cancelled(input); + + assert_admin_columns( + &event, + "flush_table", + "Failed", + json!({"error": format!("{:?}", Error::AdminFunctionCancelled)}), + ); + } + + #[test] + fn subquery_arguments_use_display() { + let Statement::Query(query) = ParserContext::create_with_dialect( + "SELECT 1", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected query statement") + }; + let input = AdminFunctionEventInput::from_request(&request_with_arguments( + FunctionArguments::Subquery(Box::new(query.inner)), + )); + let event = AdminFunctionEvent::failure( + input, + &Error::BuildAdminFunctionArgs { + msg: "subquery is not executable".to_string(), + }, + ); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "subquery", "value": "SELECT 1"}, + }) + ); + } + + #[test] + fn no_arguments_omits_arguments_field_and_empty_result() { + let input = + AdminFunctionEventInput::from_request(&request_with_arguments(FunctionArguments::None)); + let event = AdminFunctionEvent::success(input, None); + + assert_admin_columns(&event, "plugin_function", "Succeeded", json!({})); + assert_eq!(event.json_payload().unwrap(), json!({"version": 1})); + } + + #[test] + fn empty_argument_list_is_recorded() { + let input = AdminFunctionEventInput::from_request(&request("ADMIN plugin_function()")); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(0))); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": []}, + }) + ); + } + + #[test] + fn records_extended_sql_values_and_non_literal_arguments() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN plugin_function(1, true, NULL, X'48656c6c6f', table_name)", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(0))); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": { + "type": "list", + "value": [1, true, null, "48656c6c6f", "table_name"], + }, + }) + ); + } + + #[test] + fn string_literals_preserve_logical_content() { + let values = [ + SqlValue::SingleQuotedString("single".to_string()), + SqlValue::DoubleQuotedString("double".to_string()), + SqlValue::DollarQuotedString(DollarQuotedString { + value: "dollar".to_string(), + tag: None, + }), + SqlValue::SingleQuotedRawStringLiteral("raw".to_string()), + SqlValue::SingleQuotedByteStringLiteral("bytes".to_string()), + SqlValue::NationalStringLiteral("national".to_string()), + SqlValue::UnicodeStringLiteral("unicode".to_string()), + SqlValue::HexStringLiteral("48656c6c6f".to_string()), + ]; + + let expected = [ + "single", + "double", + "dollar", + "raw", + "bytes", + "national", + "unicode", + "48656c6c6f", + ]; + for (value, expected) in values.iter().zip(expected) { + assert_eq!(sql_value_to_json(value), json!(expected)); + } + } + + #[test] + fn non_literal_arguments_use_sql_display() { + let arguments = FunctionArgumentList { + duplicate_treatment: None, + args: vec![ + FunctionArg::Named { + name: Ident::new("named"), + arg: FunctionArgExpr::Expr(Expr::Identifier(Ident::new("value"))), + operator: FunctionArgOperator::RightArrow, + }, + FunctionArg::Unnamed(FunctionArgExpr::Wildcard), + ], + clauses: vec![], + }; + let input = AdminFunctionEventInput::from_request(&request_with_arguments( + FunctionArguments::List(arguments), + )); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(0))); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["named => value", "*"]}, + }) + ); + assert_eq!( + sql_value_to_json(&SqlValue::Placeholder("$1".to_string())), + json!("$1") + ); + } + + #[test] + fn serializes_binary_result() { + assert_eq!( + value_to_json(&Value::from(vec![1_u8, 2, 3])), + json!([1, 2, 3]) + ); + } + + #[test] + fn serializes_non_finite_float_results_as_strings() { + assert_eq!( + value_to_json(&Value::Float32(f32::NAN.into())), + json!("NaN") + ); + assert_eq!( + value_to_json(&Value::Float32(f32::INFINITY.into())), + json!("inf") + ); + assert_eq!( + value_to_json(&Value::Float32(f32::NEG_INFINITY.into())), + json!("-inf") + ); + assert_eq!( + value_to_json(&Value::Float64(f64::NAN.into())), + json!("NaN") + ); + assert_eq!( + value_to_json(&Value::Float64(f64::INFINITY.into())), + json!("inf") + ); + assert_eq!( + value_to_json(&Value::Float64(f64::NEG_INFINITY.into())), + json!("-inf") + ); + } +} diff --git a/src/operator/src/statement/admin/layer.rs b/src/operator/src/statement/admin/layer.rs new file mode 100644 index 00000000000..195d47d48aa --- /dev/null +++ b/src/operator/src/statement/admin/layer.rs @@ -0,0 +1,440 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::{Arc, RwLock, Weak}; + +use common_event_recorder::{EventRecorder, EventRecorderRef}; + +use crate::error::Result; +use crate::statement::admin::event::{ + ADMIN_FUNCTION_EVENT_TYPE, AdminFunctionEvent, AdminFunctionEventInput, +}; +use crate::statement::admin::{ + AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, AdminFunctionServiceRef, +}; + +/// A composable layer over ADMIN function execution. +pub trait AdminFunctionLayer: Send + Sync { + /// Wraps `inner` and returns the resulting ADMIN function service. + fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef; +} + +/// A shared composable ADMIN function layer. +pub type AdminFunctionLayerRef = Arc; + +/// A late-bound weak reference to the event recorder. +/// +/// The weak reference avoids retaining the frontend event handler, which owns a +/// statement executor itself. +#[derive(Clone, Default)] +pub struct AdminEventRecorderHandle { + recorder: Arc>>>, +} + +impl AdminEventRecorderHandle { + /// Installs the event recorder as a weak reference. + pub fn install(&self, recorder: &EventRecorderRef) { + *self.recorder.write().unwrap_or_else(|e| e.into_inner()) = Some(Arc::downgrade(recorder)); + } + + fn upgrade(&self) -> Option { + self.recorder + .read() + .unwrap_or_else(|e| e.into_inner()) + .as_ref() + .and_then(Weak::upgrade) + } +} + +/// A layer that records ADMIN function inputs and outcomes as events. +#[derive(Clone)] +pub struct AdminFunctionRecordingLayer { + recorder: AdminEventRecorderHandle, +} + +impl AdminFunctionRecordingLayer { + /// Creates a recording layer backed by the late-bound recorder handle. + pub fn new(recorder: AdminEventRecorderHandle) -> Self { + Self { recorder } + } +} + +impl AdminFunctionLayer for AdminFunctionRecordingLayer { + fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef { + Arc::new(AdminFunctionRecordingService { + inner, + recorder: self.recorder.clone(), + }) + } +} + +#[derive(Clone)] +struct AdminFunctionRecordingService { + inner: AdminFunctionServiceRef, + recorder: AdminEventRecorderHandle, +} + +#[async_trait::async_trait] +impl AdminFunctionService for AdminFunctionRecordingService { + async fn call(&self, request: AdminFunctionRequest) -> Result { + // Avoid building the event input when this event type is disabled. Do + // not retain the recorder across ADMIN execution because its handler can + // own the statement executor that contains this layer. + let enabled = self.recorder.upgrade().is_some_and(|recorder| { + recorder + .event_type_filter() + .allows(ADMIN_FUNCTION_EVENT_TYPE) + }); + let mut recording = AdminFunctionRecordingGuard::new( + enabled.then(|| AdminFunctionEventInput::from_request(&request)), + self.recorder.clone(), + ); + let result = self.inner.call(request).await; + recording.record_result(&result); + result + } +} + +/// Records an ADMIN event after completion or when an in-flight call is dropped. +struct AdminFunctionRecordingGuard { + input: Option, + recorder: AdminEventRecorderHandle, +} + +impl AdminFunctionRecordingGuard { + fn new(input: Option, recorder: AdminEventRecorderHandle) -> Self { + Self { input, recorder } + } + + fn record_result(&mut self, result: &Result) { + let Some(input) = self.input.take() else { + return; + }; + let Some(recorder) = self.recorder.upgrade() else { + return; + }; + let event = match result { + Ok(response) => AdminFunctionEvent::success(input, response.immediate_result.as_ref()), + Err(error) => AdminFunctionEvent::failure(input, error), + }; + recorder.record(Box::new(event)); + } +} + +impl Drop for AdminFunctionRecordingGuard { + fn drop(&mut self) { + let Some(input) = self.input.take() else { + return; + }; + let Some(recorder) = self.recorder.upgrade() else { + return; + }; + recorder.record(Box::new(AdminFunctionEvent::cancelled(input))); + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::fmt::{Debug, Formatter}; + use std::sync::{Arc, Mutex}; + use std::time::Duration; + + use common_event_recorder::event_table::jsonb_value; + use common_event_recorder::{ + Event, EventRecorder, EventRecorderRef, EventTypeFilter, EventTypeFilterRef, + }; + use common_query::{Output, OutputData}; + use datatypes::value::Value; + use serde_json::json; + use session::context::QueryContext; + use sql::dialect::GreptimeDbDialect; + use sql::parser::{ParseOptions, ParserContext}; + use sql::statements::statement::Statement; + use tokio::sync::Notify; + + use crate::error::Result; + use crate::statement::admin::layer::{AdminEventRecorderHandle, AdminFunctionRecordingLayer}; + use crate::statement::admin::{ + AdminFunctionLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, + AdminFunctionServiceRef, + }; + + #[derive(Clone)] + struct TraceLayer { + trace: Arc>>, + before: &'static str, + after: &'static str, + } + + impl AdminFunctionLayer for TraceLayer { + fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef { + Arc::new(TraceService { + inner, + trace: self.trace.clone(), + before: self.before, + after: self.after, + }) + } + } + + struct TraceService { + inner: AdminFunctionServiceRef, + trace: Arc>>, + before: &'static str, + after: &'static str, + } + + #[async_trait::async_trait] + impl AdminFunctionService for TraceService { + async fn call(&self, request: AdminFunctionRequest) -> Result { + self.trace.lock().unwrap().push(self.before); + let result = self.inner.call(request).await; + self.trace.lock().unwrap().push(self.after); + result + } + } + + struct TestService { + trace: Option>>>, + immediate_result: Option, + } + + #[async_trait::async_trait] + impl AdminFunctionService for TestService { + async fn call(&self, _request: AdminFunctionRequest) -> Result { + if let Some(trace) = &self.trace { + trace.lock().unwrap().push("core"); + } + Ok(response(self.immediate_result.clone())) + } + } + + struct PendingService { + started: Arc, + } + + #[async_trait::async_trait] + impl AdminFunctionService for PendingService { + async fn call(&self, _request: AdminFunctionRequest) -> Result { + self.started.notify_one(); + std::future::pending().await + } + } + + struct TestRecorder { + events: Mutex>>, + filter: EventTypeFilterRef, + } + + impl TestRecorder { + fn new(filter: EventTypeFilter) -> Self { + Self { + events: Mutex::new(Vec::new()), + filter: Arc::new(filter), + } + } + } + + impl Debug for TestRecorder { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TestRecorder").finish() + } + } + + impl EventRecorder for TestRecorder { + fn record(&self, event: Box) { + self.events.lock().unwrap().push(event); + } + + fn event_type_filter(&self) -> EventTypeFilterRef { + self.filter.clone() + } + + fn close(&self) {} + } + + #[tokio::test] + async fn last_added_layer_runs_outermost() { + let trace = Arc::new(Mutex::new(Vec::new())); + let mut service: AdminFunctionServiceRef = Arc::new(TestService { + trace: Some(trace.clone()), + immediate_result: Some(Value::UInt64(0)), + }); + for (before, after) in [("a_before", "a_after"), ("b_before", "b_after")] { + service = TraceLayer { + trace: trace.clone(), + before, + after, + } + .layer(service); + } + + service.call(request()).await.unwrap(); + assert_eq!( + *trace.lock().unwrap(), + ["b_before", "a_before", "core", "a_after", "b_after"] + ); + } + + #[tokio::test] + async fn recording_respects_filter_and_weak_lifetime() { + let enabled = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let enabled_ref: EventRecorderRef = enabled.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&enabled_ref); + recording_service(handle.clone()) + .call(request()) + .await + .unwrap(); + assert_eq!(enabled.events.lock().unwrap().len(), 1); + + let disabled = Arc::new(TestRecorder::new(EventTypeFilter::Only(HashSet::new()))); + let disabled_ref: EventRecorderRef = disabled.clone(); + handle.install(&disabled_ref); + recording_service(handle.clone()) + .call(request()) + .await + .unwrap(); + assert!(disabled.events.lock().unwrap().is_empty()); + + drop(disabled_ref); + drop(disabled); + assert!(handle.upgrade().is_none()); + recording_service(handle).call(request()).await.unwrap(); + + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let service = recording_service(handle.clone()); + let future = service.call(request()); + drop(recorder_ref); + drop(recorder); + assert!(handle.upgrade().is_none()); + future.await.unwrap(); + } + + #[tokio::test] + async fn recording_does_not_retain_recorder_while_pending() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let started = Arc::new(Notify::new()); + let service = + AdminFunctionRecordingLayer::new(handle.clone()).layer(Arc::new(PendingService { + started: started.clone(), + })); + + let task = tokio::spawn(async move { service.call(request()).await }); + started.notified().await; + drop(recorder_ref); + drop(recorder); + + assert!(handle.upgrade().is_none()); + task.abort(); + assert!(matches!(task.await, Err(error) if error.is_cancelled())); + } + + #[tokio::test] + async fn dropping_pending_call_records_failure() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let started = Arc::new(Notify::new()); + let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(PendingService { + started: started.clone(), + })); + + let task = tokio::spawn(async move { service.call(request()).await }); + started.notified().await; + task.abort(); + assert!(matches!(task.await, Err(error) if error.is_cancelled())); + assert_eq!(recorder.events.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn timing_out_pending_call_records_failure() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(PendingService { + started: Arc::new(Notify::new()), + })); + + assert!( + tokio::time::timeout(Duration::from_millis(10), service.call(request())) + .await + .is_err() + ); + assert_eq!(recorder.events.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn empty_immediate_result_preserves_success_output_and_records_no_result() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(TestService { + trace: None, + immediate_result: None, + })); + + let response = service.call(request()).await.unwrap(); + assert!(matches!(response.output.data, OutputData::AffectedRows(0))); + assert_eq!(response.immediate_result, None); + + let events = recorder.events.lock().unwrap(); + assert_eq!(events.len(), 1); + assert_eq!( + events[0].extra_rows().unwrap()[0].values[3], + jsonb_value(&json!({})) + ); + } + + fn recording_service(handle: AdminEventRecorderHandle) -> AdminFunctionServiceRef { + let core = Arc::new(TestService { + trace: None, + immediate_result: Some(Value::UInt64(0)), + }); + AdminFunctionRecordingLayer::new(handle).layer(core) + } + + fn response(immediate_result: Option) -> AdminFunctionResponse { + AdminFunctionResponse { + output: Output::new_with_affected_rows(0), + immediate_result, + } + } + + fn request() -> AdminFunctionRequest { + let Statement::Admin(statement) = ParserContext::create_with_dialect( + "ADMIN flush_table('demo')", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected ADMIN statement") + }; + AdminFunctionRequest { + statement, + query_ctx: QueryContext::arc(), + } + } +} diff --git a/tests-integration/tests/admin_function_event.rs b/tests-integration/tests/admin_function_event.rs new file mode 100644 index 00000000000..d669d2900ad --- /dev/null +++ b/tests-integration/tests/admin_function_event.rs @@ -0,0 +1,193 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use common_test_util::temp_dir::create_temp_dir; +use frontend::instance::Instance; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::QueryContext; +use sql::ast::FunctionArguments; +use sql::dialect::GreptimeDbDialect; +use sql::parser::{ParseOptions, ParserContext}; +use sql::statements::admin::Admin; +use sql::statements::statement::Statement; +use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::{StorageType, get_test_store_config}; + +use crate::event_recorder_test_util::assert_eventually_eq; + +const TABLE_NAME: &str = "admin_function_event_test"; + +#[tokio::test(flavor = "multi_thread")] +async fn test_standalone_admin_function_events() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_admin_function_events") + .build() + .await; + assert_admin_function_events(standalone.fe_instance()).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_distributed_admin_function_events() { + let store_type = StorageType::File; + if !store_type.test_on() { + return; + } + + common_telemetry::init_default_ut_logging(); + let (store_config, _guard) = get_test_store_config(&store_type); + let home_dir = create_temp_dir("test_distributed_admin_function_events_data_home"); + let cluster = GreptimeDbClusterBuilder::new("test_distributed_admin_function_events") + .await + .with_datanodes(1) + .with_store_config(store_config) + .with_shared_home_dir(Arc::new(home_dir)) + .build(true) + .await; + assert_admin_function_events(cluster.fe_instance()).await; +} + +async fn assert_admin_function_events(instance: &Arc) { + execute( + instance, + &format!("CREATE TABLE {TABLE_NAME} (ts TIMESTAMP TIME INDEX)"), + ) + .await + .unwrap(); + execute(instance, &format!("ADMIN FLUSH_TABLE('{TABLE_NAME}')")) + .await + .unwrap(); + + assert_admin_event( + instance, + "flush_table", + "\ ++----------------+----------+---------------------+-----------+---------------------------------------------------------------------------------+--------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+---------------------+-----------+---------------------------------------------------------------------------------+--------------+ +| admin_function | greptime | flush_table | Succeeded | {\"arguments\":{\"type\":\"list\",\"value\":[\"admin_function_event_test\"]},\"version\":1} | {\"result\":0} | ++----------------+----------+---------------------+-----------+---------------------------------------------------------------------------------+--------------+", + ) + .await; + + let failure = execute(instance, "ADMIN missing_admin_function('missing-input')").await; + assert!(failure.is_err()); + + assert_admin_event( + instance, + "missing_admin_function", + "\ ++----------------+----------+------------------------+--------+---------------------------------------------------------------------+-----------------------------------------------------------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+------------------------+--------+---------------------------------------------------------------------+-----------------------------------------------------------------+ +| admin_function | greptime | missing_admin_function | Failed | {\"arguments\":{\"type\":\"list\",\"value\":[\"missing-input\"]},\"version\":1} | {\"error\":\"0: Admin function not found: missing_admin_function\"} | ++----------------+----------+------------------------+--------+---------------------------------------------------------------------+-----------------------------------------------------------------+", + ) + .await; + + let failure = + execute_with_arguments(instance, "missing_no_arguments", FunctionArguments::None).await; + assert!(failure.is_err()); + + assert_admin_event( + instance, + "missing_no_arguments", + "\ ++----------------+----------+----------------------+--------+---------------+---------------------------------------------------------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+----------------------+--------+---------------+---------------------------------------------------------------+ +| admin_function | greptime | missing_no_arguments | Failed | {\"version\":1} | {\"error\":\"0: Admin function not found: missing_no_arguments\"} | ++----------------+----------+----------------------+--------+---------------+---------------------------------------------------------------+", + ) + .await; + + let Statement::Query(query) = ParserContext::create_with_dialect( + "SELECT 1", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected query statement") + }; + let failure = execute_with_arguments( + instance, + "missing_subquery", + FunctionArguments::Subquery(Box::new(query.inner)), + ) + .await; + assert!(failure.is_err()); + + assert_admin_event( + instance, + "missing_subquery", + "\ ++----------------+----------+---------------------+--------+------------------------------------------------------------------+-----------------------------------------------------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+---------------------+--------+------------------------------------------------------------------+-----------------------------------------------------------+ +| admin_function | greptime | missing_subquery | Failed | {\"arguments\":{\"type\":\"subquery\",\"value\":\"SELECT 1\"},\"version\":1} | {\"error\":\"0: Admin function not found: missing_subquery\"} | ++----------------+----------+---------------------+--------+------------------------------------------------------------------+-----------------------------------------------------------+", + ) + .await; +} + +async fn assert_admin_event(instance: &Arc, function: &str, expected: &str) { + let query = format!( + r#"SELECT + type, + actor, + admin_function_name, + admin_function_status AS status, + json_to_string(payload) AS payload, + json_to_string(admin_function_output) AS output +FROM greptime_private.events +WHERE type = 'admin_function' + AND admin_function_name = '{function}' +ORDER BY timestamp DESC +LIMIT 1"# + ); + assert_eventually_eq(instance, &query, expected).await; +} + +async fn execute( + instance: &Arc, + sql: &str, +) -> servers::error::Result { + instance.do_query(sql, QueryContext::arc()).await.remove(0) +} + +async fn execute_with_arguments( + instance: &Arc, + function: &str, + arguments: FunctionArguments, +) -> operator::error::Result { + let Statement::Admin(mut statement) = ParserContext::create_with_dialect( + &format!("ADMIN {function}()"), + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected ADMIN statement") + }; + let Admin::Func(admin_function) = &mut statement; + admin_function.args = arguments; + + instance + .statement_executor() + .execute_sql(Statement::Admin(statement), QueryContext::arc()) + .await +} diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 1925b46d453..3112cb72c32 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -14,6 +14,7 @@ #![recursion_limit = "256"] +mod admin_function_event; mod database_ddl_event; mod event_recorder_test_util; #[macro_use] From 154f90b3652bf9b1b3e7eccc5f277f856a2f3079 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Wed, 12 Aug 2026 03:12:04 +0000 Subject: [PATCH 052/126] fix: harden permission checks and process visibility (#8852) Signed-off-by: shuiyisong --- src/auth/src/lib.rs | 4 +- src/auth/src/permission.rs | 2 + src/auth/src/user_info.rs | 75 +++++++----- src/auth/src/user_provider.rs | 10 +- .../src/user_provider/static_user_provider.rs | 4 + src/frontend/src/instance.rs | 112 +++++++++++++++++- src/frontend/src/instance/grpc.rs | 24 ++-- src/query/src/sql.rs | 6 +- 8 files changed, 195 insertions(+), 42 deletions(-) diff --git a/src/auth/src/lib.rs b/src/auth/src/lib.rs index d2a1c05b1cd..d107068940e 100644 --- a/src/auth/src/lib.rs +++ b/src/auth/src/lib.rs @@ -31,8 +31,8 @@ pub use common::{ }; pub use permission::{ ALL_ACTIONS, AccessMode, DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, - DefaultPermissionChecker, INFLUXDB_WRITE, JAEGER_QUERY, LOG_QUERY, LOG_WRITE, OPENTSDB_WRITE, - OTLP_WRITE, PIPELINE_DELETE, PIPELINE_INSERT, PIPELINE_QUERY, PROM_STORE_READ, + DefaultPermissionChecker, ICEBERG_READ, INFLUXDB_WRITE, JAEGER_QUERY, LOG_QUERY, LOG_WRITE, + OPENTSDB_WRITE, OTLP_WRITE, PIPELINE_DELETE, PIPELINE_INSERT, PIPELINE_QUERY, PROM_STORE_READ, PROM_STORE_WRITE, PROMQL_QUERY, PermissionAction, PermissionChecker, PermissionReq, PermissionResp, PermissionTableTarget, PermissionTableTargets, SEMANTIC_GRAPH_QUERY, }; diff --git a/src/auth/src/permission.rs b/src/auth/src/permission.rs index 901b935d6a2..984fe8ea4cf 100644 --- a/src/auth/src/permission.rs +++ b/src/auth/src/permission.rs @@ -139,6 +139,7 @@ pub const PROMQL_QUERY: PermissionAction = PermissionAction::read("promql.query" pub const LOG_QUERY: PermissionAction = PermissionAction::read("log.query"); pub const OPENTSDB_WRITE: PermissionAction = PermissionAction::write("opentsdb.write"); pub const INFLUXDB_WRITE: PermissionAction = PermissionAction::write("influxdb.write"); +pub const ICEBERG_READ: PermissionAction = PermissionAction::read("iceberg.read"); pub const PROM_STORE_WRITE: PermissionAction = PermissionAction::write("prom_store.write"); pub const PROM_STORE_READ: PermissionAction = PermissionAction::read("prom_store.read"); pub const OTLP_WRITE: PermissionAction = PermissionAction::write("otlp.write"); @@ -164,6 +165,7 @@ pub const ALL_ACTIONS: &[PermissionAction] = &[ LOG_QUERY, OPENTSDB_WRITE, INFLUXDB_WRITE, + ICEBERG_READ, PROM_STORE_WRITE, PROM_STORE_READ, OTLP_WRITE, diff --git a/src/auth/src/user_info.rs b/src/auth/src/user_info.rs index 6ce735c4e2f..ab7c7f32f46 100644 --- a/src/auth/src/user_info.rs +++ b/src/auth/src/user_info.rs @@ -21,6 +21,10 @@ use crate::UserInfoRef; pub trait UserInfo: Debug + Sync + Send { fn as_any(&self) -> &dyn Any; fn username(&self) -> &str; + + fn is_admin(&self) -> bool { + false + } } /// The user permission mode @@ -38,13 +42,14 @@ impl PermissionMode { /// - "rw", "readwrite", "read_write" => ReadWrite /// - "ro", "readonly", "read_only" => ReadOnly /// - "wo", "writeonly", "write_only" => WriteOnly - /// Returns None if the input string is not a valid permission mode. - pub fn from_str(s: &str) -> Self { + /// + /// Returns `None` if the input string is not a valid permission mode. + pub fn from_str(s: &str) -> Option { match s.to_lowercase().as_str() { - "readwrite" | "read_write" | "rw" => PermissionMode::ReadWrite, - "readonly" | "read_only" | "ro" => PermissionMode::ReadOnly, - "writeonly" | "write_only" | "wo" => PermissionMode::WriteOnly, - _ => PermissionMode::ReadWrite, + "readwrite" | "read_write" | "rw" => Some(PermissionMode::ReadWrite), + "readonly" | "read_only" | "ro" => Some(PermissionMode::ReadOnly), + "writeonly" | "write_only" | "wo" => Some(PermissionMode::WriteOnly), + _ => None, } } @@ -123,58 +128,72 @@ mod tests { // Test ReadWrite variants assert_eq!( PermissionMode::from_str("readwrite"), - PermissionMode::ReadWrite + Some(PermissionMode::ReadWrite) ); assert_eq!( PermissionMode::from_str("read_write"), - PermissionMode::ReadWrite + Some(PermissionMode::ReadWrite) + ); + assert_eq!( + PermissionMode::from_str("rw"), + Some(PermissionMode::ReadWrite) ); - assert_eq!(PermissionMode::from_str("rw"), PermissionMode::ReadWrite); assert_eq!( PermissionMode::from_str("ReadWrite"), - PermissionMode::ReadWrite + Some(PermissionMode::ReadWrite) + ); + assert_eq!( + PermissionMode::from_str("RW"), + Some(PermissionMode::ReadWrite) ); - assert_eq!(PermissionMode::from_str("RW"), PermissionMode::ReadWrite); // Test ReadOnly variants assert_eq!( PermissionMode::from_str("readonly"), - PermissionMode::ReadOnly + Some(PermissionMode::ReadOnly) ); assert_eq!( PermissionMode::from_str("read_only"), - PermissionMode::ReadOnly + Some(PermissionMode::ReadOnly) + ); + assert_eq!( + PermissionMode::from_str("ro"), + Some(PermissionMode::ReadOnly) ); - assert_eq!(PermissionMode::from_str("ro"), PermissionMode::ReadOnly); assert_eq!( PermissionMode::from_str("ReadOnly"), - PermissionMode::ReadOnly + Some(PermissionMode::ReadOnly) + ); + assert_eq!( + PermissionMode::from_str("RO"), + Some(PermissionMode::ReadOnly) ); - assert_eq!(PermissionMode::from_str("RO"), PermissionMode::ReadOnly); // Test WriteOnly variants assert_eq!( PermissionMode::from_str("writeonly"), - PermissionMode::WriteOnly + Some(PermissionMode::WriteOnly) ); assert_eq!( PermissionMode::from_str("write_only"), - PermissionMode::WriteOnly + Some(PermissionMode::WriteOnly) + ); + assert_eq!( + PermissionMode::from_str("wo"), + Some(PermissionMode::WriteOnly) ); - assert_eq!(PermissionMode::from_str("wo"), PermissionMode::WriteOnly); assert_eq!( PermissionMode::from_str("WriteOnly"), - PermissionMode::WriteOnly + Some(PermissionMode::WriteOnly) ); - assert_eq!(PermissionMode::from_str("WO"), PermissionMode::WriteOnly); - - // Test invalid inputs default to ReadWrite assert_eq!( - PermissionMode::from_str("invalid"), - PermissionMode::ReadWrite + PermissionMode::from_str("WO"), + Some(PermissionMode::WriteOnly) ); - assert_eq!(PermissionMode::from_str(""), PermissionMode::ReadWrite); - assert_eq!(PermissionMode::from_str("xyz"), PermissionMode::ReadWrite); + + for invalid in ["readonyl", "", "xyz"] { + assert_eq!(PermissionMode::from_str(invalid), None); + } } #[test] @@ -200,7 +219,7 @@ mod tests { for mode in modes { let str_repr = mode.as_str(); let parsed = PermissionMode::from_str(str_repr); - assert_eq!(mode, parsed); + assert_eq!(Some(mode), parsed); } } diff --git a/src/auth/src/user_provider.rs b/src/auth/src/user_provider.rs index 985e25978e8..bafbd4fffce 100644 --- a/src/auth/src/user_provider.rs +++ b/src/auth/src/user_provider.rs @@ -463,7 +463,7 @@ pub(crate) fn parse_credential_line( let (username_part, password) = (parts[0], parts[1]); let (username, permission_mode) = if let Some((user, perm)) = username_part.split_once(':') { - (user, PermissionMode::from_str(perm)) + (user, PermissionMode::from_str(perm)?) } else { (username_part, PermissionMode::default()) }; @@ -766,6 +766,14 @@ mod tests { let result = parse_credential_line("user=pass=word"); assert_eq!(result, None); + for line in [ + "user:readonyl=password", + "user:=password", + "user:arbitrary=password", + ] { + assert_eq!(parse_credential_line(line), None); + } + // Empty password let result = parse_credential_line("user="); assert_eq!( diff --git a/src/auth/src/user_provider/static_user_provider.rs b/src/auth/src/user_provider/static_user_provider.rs index 3e0dac8ee74..69d7fd63d06 100644 --- a/src/auth/src/user_provider/static_user_provider.rs +++ b/src/auth/src/user_provider/static_user_provider.rs @@ -147,6 +147,8 @@ pub mod test { let provider = StaticUserProvider::new("cmd:root=123456,admin=654321").unwrap(); test_authenticate(&provider, "root", "123456").await; test_authenticate(&provider, "admin", "654321").await; + + assert!(StaticUserProvider::new("cmd:user:readonyl=password").is_err()); } #[tokio::test] @@ -171,6 +173,7 @@ pub mod test { assert!( lw.write_all( b"root=123456 +invalid:readonyl=password admin=654321", ) .is_ok() @@ -182,5 +185,6 @@ admin=654321", let provider = StaticUserProvider::new(param.as_str()).unwrap(); test_authenticate(&provider, "root", "123456").await; test_authenticate(&provider, "admin", "654321").await; + test_authenticate_fails(&provider, "invalid", "password").await; } } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index bdf4d66c7e9..0c300835e0e 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1665,6 +1665,7 @@ fn should_track_plan_process(stmt: Option<&Statement>, plan: &LogicalPlan) -> bo #[cfg(test)] mod tests { + use std::any::Any; use std::collections::HashMap; use std::future::Future; use std::pin::Pin; @@ -1676,10 +1677,12 @@ mod tests { use api::prom_store::remote::{ Label, LabelMatcher, Query as RemoteQuery, ReadRequest, ReadResponse, Sample, }; + use api::v1::greptime_request::Request; use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse}; + use api::v1::query_request::Query; use auth::{ DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, JAEGER_QUERY, PIPELINE_DELETE, - PIPELINE_INSERT, PIPELINE_QUERY, PermissionAction, PermissionResp, UserInfoRef, + PIPELINE_INSERT, PIPELINE_QUERY, PermissionAction, PermissionResp, UserInfo, UserInfoRef, }; use catalog::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer}; use common_base::Plugins; @@ -1942,6 +1945,23 @@ mod tests { ) } + #[derive(Debug)] + struct AdminUserInfo; + + impl UserInfo for AdminUserInfo { + fn as_any(&self) -> &dyn Any { + self + } + + fn username(&self) -> &str { + "admin" + } + + fn is_admin(&self) -> bool { + true + } + } + struct RejectUnresolvedPermissionChecker; impl PermissionChecker for RejectUnresolvedPermissionChecker { @@ -3102,6 +3122,45 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_write_only_grpc_sql_is_checked_after_parsing() -> TestResult<()> { + let plugins = Plugins::new(); + plugins.insert::(Arc::new(WriteOnlyPermissionChecker)); + let instance = test_instance_with_plugins( + test_table(1024, "source")?, + test_table(1025, "target")?, + plugins, + ) + .await?; + + let insert = Request::Query(api::v1::QueryRequest { + query: Some(Query::Sql( + "INSERT INTO target SELECT * FROM source".to_string(), + )), + }); + servers::query_handler::grpc::GrpcQueryHandler::do_query( + &instance, + insert, + QueryContext::arc(), + ) + .await + .unwrap(); + + let select = Request::Query(api::v1::QueryRequest { + query: Some(Query::Sql("SELECT * FROM source".to_string())), + }); + assert_permission_denied( + servers::query_handler::grpc::GrpcQueryHandler::do_query( + &instance, + select, + QueryContext::arc(), + ) + .await, + ); + + Ok(()) + } + #[tokio::test] async fn test_target_independent_checker_skips_target_resolution() -> TestResult<()> { let physical_table = "physical_metric"; @@ -3438,6 +3497,57 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_show_processlist_catalog_scope() -> TestResult<()> { + let instance = + test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?) + .await?; + let _current_catalog = instance.process_manager().register_query( + "greptime".to_string(), + vec!["public".to_string()], + "current_catalog_query".to_string(), + String::new(), + None, + None, + ); + let _other_catalog = instance.process_manager().register_query( + "other".to_string(), + vec!["public".to_string()], + "other_catalog_query".to_string(), + String::new(), + None, + None, + ); + + for sql in ["SHOW PROCESSLIST", "SHOW FULL PROCESSLIST"] { + let output = execute_one_sql(&instance, sql, test_query_ctx(43)).await?; + let process_list = output.data.pretty_print().await; + assert!( + process_list.contains("current_catalog_query"), + "{process_list}" + ); + assert!( + !process_list.contains("other_catalog_query"), + "{process_list}" + ); + + let admin_ctx = test_query_ctx(44); + admin_ctx.set_current_user(Arc::new(AdminUserInfo)); + let output = execute_one_sql(&instance, sql, admin_ctx).await?; + let process_list = output.data.pretty_print().await; + assert!( + process_list.contains("current_catalog_query"), + "{process_list}" + ); + assert!( + process_list.contains("other_catalog_query"), + "{process_list}" + ); + } + + Ok(()) + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_kill_query_cancels_insert_select() -> TestResult<()> { assert_kill_cancels_insert_select("KILL QUERY 4242").await diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index acb3a79e049..7d4f0ee7959 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -71,15 +71,21 @@ impl GrpcQueryHandler for Instance { let interceptor = interceptor_ref.as_ref(); interceptor.pre_execute(&request, ctx.clone())?; - self.plugins - .get::() - .as_ref() - .check_permission_with_context( - ctx.current_user(), - PermissionReq::GrpcRequest(&request), - Some(&ctx.current_schema()), - ) - .context(PermissionSnafu)?; + if !matches!( + &request, + Request::Query(query_request) + if matches!(&query_request.query, Some(Query::Sql(_))) + ) { + self.plugins + .get::() + .as_ref() + .check_permission_with_context( + ctx.current_user(), + PermissionReq::GrpcRequest(&request), + Some(&ctx.current_schema()), + ) + .context(PermissionSnafu)?; + } let output = match request { Request::Inserts(requests) => self.handle_inserts(requests, ctx.clone()).await?, diff --git a/src/query/src/sql.rs b/src/query/src/sql.rs index fc755deb8a7..b2d979c6664 100644 --- a/src/query/src/sql.rs +++ b/src/query/src/sql.rs @@ -1360,7 +1360,11 @@ pub async fn show_processlist( ] }; - let filters = vec![]; + let filters = if query_ctx.current_user().is_admin() { + vec![] + } else { + vec![col(process_list::CATALOG).eq(lit(query_ctx.current_catalog()))] + }; let like_field = None; let sort = vec![col("id").sort(true, true)]; query_from_information_schema_table( From e778a72829d5638a6ea7892f6c4fa32ecede86a0 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 12 Aug 2026 04:06:31 +0000 Subject: [PATCH 053/126] feat: complete the derived-edge vocabulary of the entity graph (#8836) * feat(operator): pair calls edges across trace tables and derive virtual-node edges Union the normalized client and server spans of all trace tables before the join, so a client span pairs with a server span stored in a different table. A client span with no matching server span becomes an edge to a virtual node named by span attributes (peer.service / db.name / server.address), with confidence < 1.0 and attributes.connection_type; a window's real pairs win over virtual candidates for the same edge key. Signed-off-by: Dennis Zhuang * feat(operator): derive same-row co-declared edges from the built-in vocabulary A table declaring both entity types of a vocabulary pair witnesses the edge on every row carrying both identities: runs_on / contains / part_of for any declaring table (provenance 'attribute'), agent uses model / agent invoked tool only for trace sources (span-structure observations, provenance 'trace'). Signed-off-by: Dennis Zhuang * feat(operator): derive parent_agent-calls-agent edges from span structure Trace tables declaring an agent entity pair each span with its child span across tables (no span-kind filter), keep pairs whose agent identities differ, and aggregate RED metrics per window, anchored on the parent span like the service derivation is anchored on the client. Signed-off-by: Dennis Zhuang * feat(frontend): feed co-declared and agent sources into the relationships scan scan_relationships now passes every declaring table (with its trace-ness) to the co-declared branch and the trace tables' agent declarations to the agent-calls derivation. enumerate validates the fixed trace-v1 columns and derives around a malformed trace table instead of failing the whole scan. Signed-off-by: Dennis Zhuang * test: cover cross-table pairing, virtual nodes, co-declared and agent edges sqlness exercises the new derivations end to end (including a malformed trace-model table being skipped); the integration authorization test now also pins that a pair split across tables derives no edge when the caller cannot read one side. Signed-off-by: Dennis Zhuang * chore: update the relationships module doc for the new branches Signed-off-by: Dennis Zhuang * chore: import shared derivation helpers via crate paths The fmt CI gate rejects module-level 'use super::' imports. Signed-off-by: Dennis Zhuang * fix: fold co-declared duplicates, decouple agent calls, verify the trace time index Review findings: the co-declared branch lacked a cross-source DISTINCT, so two tables witnessing the same edge in one window emitted duplicate rows; the agent-calls derivation was gated on a usable service declaration; the trace schema guard accepted a table whose time index is not the column the derivations bucket by. The empty-trace-table test asserted a union invariant with no information and is dropped. Signed-off-by: Dennis Zhuang * fix: rename the agent-tool edge to invokes and track current OTel peer attributes The vocabulary's other relation names are present tense; semconv 1.39/1.26 replaced peer.service and db.name with service.peer.name and db.namespace, so the virtual-node candidates now check the current names first and keep the deprecated ones for existing telemetry. Signed-off-by: Dennis Zhuang * fix: trust the trace-v1 table option instead of matching the fixed schema The option is only ever stamped by the ingest path, which guarantees the fixed span columns; matching column types here couples the graph to every trace schema evolution (e.g. #8816) for a case that cannot occur. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/instance/entity_graph.rs | 85 +- src/operator/src/statement/semantic_graph.rs | 78 +- .../statement/semantic_graph/relationships.rs | 1462 ++++++++++++++--- tests-integration/src/instance.rs | 61 +- .../common/system/semantic_graph.result | 201 +++ .../common/system/semantic_graph.sql | 137 ++ 6 files changed, 1748 insertions(+), 276 deletions(-) diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs index e55e5f2b8d7..70c57846fd0 100644 --- a/src/frontend/src/instance/entity_graph.rs +++ b/src/frontend/src/instance/entity_graph.rs @@ -47,9 +47,9 @@ use datafusion::dataframe::DataFrame; use datafusion_expr::LogicalPlan; use futures::TryStreamExt; use operator::statement::semantic_graph::{ - CallsSource, DeclaredSource, EntityDeclaration, GraphQueryWindow, OBSERVED_AT_COLUMN, - RegistrySource, build_registry_plan, build_relationships_plan, - declared_relationships_schema_matches, + CallsSource, CoDeclaredSource, DeclaredSource, EntityDeclaration, GraphQueryWindow, + OBSERVED_AT_COLUMN, RegistrySource, RelationshipSources, build_registry_plan, + build_relationships_plan, declared_relationships_schema_matches, }; use query::QueryEngineRef; use session::context::{QueryContext, QueryContextBuilder, QueryContextRef}; @@ -73,6 +73,13 @@ pub struct EntityGraphProviderImpl { struct EntitySource { declarations: Vec, + is_trace: bool, + table: TableRef, +} + +struct TraceSource { + service: Option, + agent: Option, table: TableRef, } @@ -222,7 +229,7 @@ impl EntityGraphProviderImpl { &self, catalog: &str, query_ctx: Option<&QueryContext>, - ) -> Result<(Vec, Vec<(EntityDeclaration, TableRef)>), BoxedError> { + ) -> Result<(Vec, Vec), BoxedError> { let Some(catalog_manager) = self.catalog_manager.upgrade() else { return Ok((vec![], vec![])); }; @@ -288,26 +295,34 @@ impl EntityGraphProviderImpl { continue; } if is_trace { - // TODO(entity-graph): validate the complete fixed trace-v1 - // schema before adding the table to calls derivation, and - // skip malformed/stale tables instead of letting one poison - // the whole relationship scan. - match table_declarations - .iter() - .find(|d| d.entity_type == "service") - { - Some(service) => traces.push((service.clone(), table.clone())), + let find = |entity_type: &str| { + table_declarations + .iter() + .find(|d| d.entity_type == entity_type) + .cloned() + }; + let service = find("service"); + let agent = find("agent"); + if service.is_none() { // No usable service identity: the table cannot - // contribute calls edges (see declarations_for). - None => warn!( + // contribute service-calls edges (see declarations_for). + warn!( "Trace table `{}` has no usable service declaration; skipping calls derivation", table_info.name - ), + ); + } + if service.is_some() || agent.is_some() { + traces.push(TraceSource { + service, + agent, + table: table.clone(), + }); } } if !table_declarations.is_empty() { declarations.push(EntitySource { declarations: table_declarations, + is_trace, table, }); } @@ -480,21 +495,37 @@ impl EntityGraphProvider for EntityGraphProviderImpl { request: ScanRequest, query_ctx: Option, ) -> Result, BoxedError> { - let (_, traces) = self.enumerate(catalog, query_ctx.as_deref()).await?; - let mut scans = Vec::with_capacity(traces.len()); - for (service, trace) in traces { - scans.push(CallsSource { - service, - scan: self.read_table(trace)?, + let (sources, traces) = self.enumerate(catalog, query_ctx.as_deref()).await?; + let mut calls = Vec::with_capacity(traces.len()); + for trace in traces { + calls.push(CallsSource { + service: trace.service, + agent: trace.agent, + scan: self.read_table(trace.table)?, + }); + } + let mut co_declared = Vec::with_capacity(sources.len()); + for source in sources { + co_declared.push(CoDeclaredSource { + declarations: source.declarations, + is_trace: source.is_trace, + scan: self.read_table(source.table)?, }); } let declared = self.declared_source(catalog, query_ctx.as_deref()).await?; let Some(window) = Self::query_window(&request)? else { return Ok(None); }; - let Some(plan) = build_relationships_plan(scans, declared, &window) - .context(error::DataFusionSnafu) - .map_err(BoxedError::new)? + let Some(plan) = build_relationships_plan( + RelationshipSources { + traces: calls, + co_declared, + declared, + }, + &window, + ) + .context(error::DataFusionSnafu) + .map_err(BoxedError::new)? else { return Ok(None); }; @@ -529,6 +560,10 @@ mod tests { .iter() .map(|c| ColumnSchema::new(*c, ConcreteDataType::string_datatype(), true)), ); + assemble_table_info(column_schemas, extra) + } + + fn assemble_table_info(column_schemas: Vec, extra: &[(&str, &str)]) -> TableInfo { let schema = Arc::new( SchemaBuilder::try_from_columns(column_schemas) .unwrap() diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index af9108198d4..c14d9d0d04e 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -45,7 +45,9 @@ use datafusion::functions::{core as core_fns, datetime as datetime_fns, string a use datafusion::functions_nested::expr_fn::make_array; use datafusion_common::{Column, Result as DfResult, ScalarValue}; use datafusion_expr::{Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; -pub use relationships::{CallsSource, DeclaredSource, build_relationships_plan}; +pub use relationships::{ + CallsSource, CoDeclaredSource, DeclaredSource, RelationshipSources, build_relationships_plan, +}; use store_api::mito_engine_options::{APPEND_MODE_KEY, MERGE_MODE_KEY, TTL_KEY}; /// Whether an existing declared-edge table still matches the canonical @@ -513,11 +515,44 @@ const REGISTRY_COLUMNS: [&str; 10] = [ const REGISTRY_VALID_COLUMN: &str = "__entity_valid"; -/// Projects all entity declarations of one source table with one source scan. -/// -/// Each output field is first built as an array whose entries correspond to the -/// table's declarations. Unnesting the arrays expands one source row into one -/// row per declared entity without cloning the table scan under `UNION ALL`. +/// Expands one source row into one output row per entry of `rows` with a +/// single source scan: each output field is first built as an array whose +/// entries correspond to the rows, then unnested. Every row is `[valid, +/// values...]` aligned with `columns`; rows whose `valid` expression is false +/// are dropped, and the distinct output columns are `columns`. +fn unnest_rows( + df: DataFrame, + window_predicate: Expr, + valid_column: &'static str, + columns: &[&'static str], + rows: Vec>, +) -> DfResult { + let mut arrays = vec![Vec::with_capacity(rows.len()); columns.len() + 1]; + for row in rows { + debug_assert_eq!(row.len(), arrays.len()); + for (array, value) in arrays.iter_mut().zip(row) { + array.push(value); + } + } + let array_names = std::iter::once(valid_column) + .chain(columns.iter().copied()) + .collect::>(); + let array_projection = array_names + .iter() + .zip(arrays) + .map(|(name, values)| make_array(values).alias(*name)) + .collect::>(); + + df.filter(window_predicate)? + .select(array_projection)? + .unnest_columns(&array_names)? + .filter(ident(valid_column))? + .select(columns.iter().map(|c| ident(*c)).collect::>())? + .distinct() +} + +/// Projects all entity declarations of one source table with one source scan +/// via [`unnest_rows`]. fn registry_source( first: &EntityDeclaration, rest: &[EntityDeclaration], @@ -531,7 +566,7 @@ fn registry_source( .gt_eq(window.source_start()) .and(ts.lt(window.source_end())); - let mut arrays = vec![Vec::new(); REGISTRY_COLUMNS.len() + 1]; + let mut rows = Vec::with_capacity(1 + rest.len()); for decl in std::iter::once(first).chain(rest) { // CAST even a single-column id: id columns must be tags but not // necessarily strings, and the computed table declares entity_id @@ -575,7 +610,7 @@ fn registry_source( predicate.and(ident(id).is_not_null()) }); - let row = [ + rows.push(vec![ valid, bin.clone(), bin.clone(), @@ -587,27 +622,16 @@ fn registry_source( scope, descriptive, source_tables, - ]; - for (array, value) in arrays.iter_mut().zip(row) { - array.push(value); - } + ]); } - let array_names = std::iter::once(REGISTRY_VALID_COLUMN) - .chain(REGISTRY_COLUMNS) - .collect::>(); - let array_projection = array_names - .iter() - .zip(arrays) - .map(|(name, values)| make_array(values).alias(*name)) - .collect::>(); - - df.filter(window_predicate)? - .select(array_projection)? - .unnest_columns(&array_names)? - .filter(ident(REGISTRY_VALID_COLUMN))? - .select(REGISTRY_COLUMNS.into_iter().map(ident).collect::>())? - .distinct() + unnest_rows( + df, + window_predicate, + REGISTRY_VALID_COLUMN, + ®ISTRY_COLUMNS, + rows, + ) } /// A declaring table's scan paired with the entity declarations it carries — diff --git a/src/operator/src/statement/semantic_graph/relationships.rs b/src/operator/src/statement/semantic_graph/relationships.rs index 7ebf5ae86fd..416c8dc10e1 100644 --- a/src/operator/src/statement/semantic_graph/relationships.rs +++ b/src/operator/src/statement/semantic_graph/relationships.rs @@ -13,9 +13,10 @@ // limitations under the License. //! The `semantic_relationships` derivation: the plan builders for every edge -//! branch (trace-derived `calls`, the declared-edge union) behind the computed -//! table. Shared expression helpers and the entity registry live in the parent -//! module. +//! branch behind the computed table — trace-derived service `calls` (with +//! virtual-node edges for unmatched clients), agent `calls` from span +//! structure, same-row co-declared edges, and the declared-edge union. Shared +//! expression helpers and the entity registry live in the parent module. use common_catalog::consts::{ DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, @@ -25,14 +26,15 @@ use common_catalog::consts::{ use datafusion::arrow::datatypes::DataType; use datafusion::dataframe::DataFrame; use datafusion::functions::core as core_fns; -use datafusion::functions_aggregate::expr_fn::{count, sum}; +use datafusion::functions_aggregate::expr_fn::{bool_or, count, min, sum}; use datafusion::functions_window::expr_fn::row_number; -use datafusion_common::Result as DfResult; -use datafusion_expr::{ExprFunctionExt, JoinType, LogicalPlan, cast, ident, lit}; +use datafusion_common::{Result as DfResult, ScalarValue}; +use datafusion_expr::{Expr, ExprFunctionExt, JoinType, LogicalPlan, cast, ident, lit, when}; use crate::statement::semantic_graph::{ DECLARED_EDGE_IDENTITY_COLUMNS, EntityDeclaration, GraphQueryWindow, OBSERVED_AT_COLUMN, - bin_interval, bin_ms, entity_id_expr, interval, null_json, qcol, union_all, + bin_interval, bin_ms, entity_id_expr, interval, null_json, parse_json_expr, qcol, union_all, + unnest_rows, }; /// The projected columns of `semantic_relationships`, in order. Every derived @@ -67,10 +69,23 @@ const RELATIONSHIP_COLUMNS: [&str; 16] = [ const CHILD_SPAN_EARLY_NANOS: i64 = 5 * 60 * 1_000_000_000; const CHILD_SPAN_LATE_NANOS: i64 = 60 * 60 * 1_000_000_000; -/// A trace table's scan paired with the `service` declaration it derives -/// `calls` edges for — a unit of `build_relationships_plan`. +/// A trace table's scan paired with the entity declarations its derivations +/// key on — a unit of `build_relationships_plan`. The `service` declaration +/// feeds service calls, the `agent` declaration agent calls; the two are +/// independent, so a table whose service declaration is unusable still +/// derives agent edges (and vice versa). pub struct CallsSource { - pub service: EntityDeclaration, + pub service: Option, + pub agent: Option, + pub scan: DataFrame, +} + +/// A declaring table's scan paired with its entity declarations, from which +/// the same-row co-declaration rules derive edges. `is_trace` gates the +/// agent-edge vocabulary, which the RFC ties to span structure. +pub struct CoDeclaredSource { + pub declarations: Vec, + pub is_trace: bool, pub scan: DataFrame, } @@ -80,17 +95,29 @@ pub struct DeclaredSource { pub scan: DataFrame, } -/// Builds the `semantic_relationships` plan: the trace-derived `calls` branch -/// unioned with the declared-edge branch, re-projected to the 16-column -/// contract. Returns `None` when there is neither a trace table nor a -/// declared-edge table, so the computed table streams empty. +/// Everything `build_relationships_plan` derives edges from. +pub struct RelationshipSources { + pub traces: Vec, + pub co_declared: Vec, + pub declared: Option, +} + +/// Builds the `semantic_relationships` plan: the service-calls, agent-calls, +/// co-declared, and declared-edge branches unioned and re-projected to the +/// 16-column contract. Returns `None` when no source can contribute edges, so +/// the computed table streams empty. pub fn build_relationships_plan( - traces: Vec, - declared: Option, + sources: RelationshipSources, window: &GraphQueryWindow, ) -> DfResult> { - let mut union_df = calls_branch(traces, window)?; - if let Some(declared) = declared { + let mut union_df = calls_branch(&sources.traces, window)?; + if let Some(agent_calls) = agent_calls_branch(&sources.traces, window)? { + union_df = union_all(union_df, agent_calls)?; + } + if let Some(co_declared) = co_declared_branch(sources.co_declared, window)? { + union_df = union_all(union_df, co_declared)?; + } + if let Some(declared) = sources.declared { union_df = union_all(union_df, declared_edges(declared.scan, window)?)?; } union_df @@ -102,6 +129,112 @@ pub fn build_relationships_plan( .transpose() } +/// The same-row co-declaration vocabulary open to any declaring table: a row +/// carrying both identities witnesses the edge, and the built-in direction is +/// `src -> dst`. Rows are `(src_type, dst_type, rel_type)`. +const ATTRIBUTE_EDGE_VOCABULARY: [(&str, &str, &str); 6] = [ + ("service.instance", "host", "runs_on"), + ("process", "host", "runs_on"), + ("k8s.pod", "k8s.node", "runs_on"), + ("k8s.pod", "k8s.container", "contains"), + ("service.instance", "service", "part_of"), + ("k8s.pod", "k8s.workload", "part_of"), +]; + +/// The agent-edge vocabulary, applied only to trace sources: the RFC derives +/// `agent uses model` / `agent invokes tool` from span structure (an LLM- or +/// tool-call span carries both identities), so a non-trace table co-declaring +/// these types must not fabricate invocations. +const AGENT_EDGE_VOCABULARY: [(&str, &str, &str); 2] = + [("agent", "model", "uses"), ("agent", "tool", "invokes")]; + +const CO_DECLARED_VALID_COLUMN: &str = "__edge_valid"; + +/// The same-row co-declared branch: for each source table, every vocabulary +/// pair whose two entity types the table declares yields an edge projection, +/// expanded from one scan via `unnest_rows` (the registry idiom). Attribute +/// edges carry `provenance = 'attribute'`; agent edges are span-structure +/// observations and carry `provenance = 'trace'`. +fn co_declared_branch( + sources: Vec, + window: &GraphQueryWindow, +) -> DfResult> { + let mut union_df: Option = None; + for source in sources { + let find = |ty: &str| source.declarations.iter().find(|d| d.entity_type == ty); + let mut pairs = Vec::new(); + for (src_type, dst_type, rel_type) in ATTRIBUTE_EDGE_VOCABULARY { + if let (Some(src), Some(dst)) = (find(src_type), find(dst_type)) { + pairs.push((src, dst, rel_type, "attribute")); + } + } + if source.is_trace { + for (src_type, dst_type, rel_type) in AGENT_EDGE_VOCABULARY { + if let (Some(src), Some(dst)) = (find(src_type), find(dst_type)) { + pairs.push((src, dst, rel_type, "trace")); + } + } + } + let Some((first, _, _, _)) = pairs.first() else { + continue; + }; + + let ts = ident(&first.time_index); + let bin = bin_ms(ts.clone()); + let window_predicate = ts + .clone() + .gt_eq(window.source_start()) + .and(ts.lt(window.source_end())); + let null_i64 = || lit(ScalarValue::Int64(None)); + let rows = pairs + .iter() + .map(|(src, dst, rel_type, provenance)| { + // Both endpoints must identify something on the row for the + // row to witness the edge. + let valid = src + .id_columns + .iter() + .chain(&dst.id_columns) + .fold(lit(true), |predicate, id| { + predicate.and(ident(id).is_not_null()) + }); + vec![ + valid, + bin.clone(), + bin.clone(), + bin.clone() + bin_interval(), + bin.clone() + bin_interval(), + lit(src.entity_type.as_str()), + entity_id_expr(&src.id_columns, &|c| ident(c)), + lit(dst.entity_type.as_str()), + entity_id_expr(&dst.id_columns, &|c| ident(c)), + lit(*rel_type), + lit(*provenance), + lit(1.0_f64), + null_i64(), + null_i64(), + lit(ScalarValue::Float64(None)), + null_i64(), + null_json(), + ] + }) + .collect(); + + let branch = unnest_rows( + source.scan, + window_predicate, + CO_DECLARED_VALID_COLUMN, + &RELATIONSHIP_COLUMNS, + rows, + )?; + union_df = union_all(union_df, branch)?; + } + // The per-source DISTINCT inside `unnest_rows` cannot see other sources: + // two tables witnessing the same edge in the same window must still fold + // into one row per (window, edge). + union_df.map(DataFrame::distinct).transpose() +} + /// Ranking column used to pick the latest declared-edge revision per edge key. const DECLARED_REVISION_COLUMN: &str = "__declared_revision"; @@ -179,27 +312,354 @@ fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult, - window: &GraphQueryWindow, -) -> DfResult> { - let mut union_df: Option = None; +/// Span-attribute columns naming an uninstrumented peer, in precedence order, +/// each with the `connection_type` its match implies: current OTel names +/// first (`service.peer.name` and `db.namespace` replaced the deprecated +/// `peer.service` and `db.name` in semconv 1.39/1.26, kept for existing +/// telemetry), the bare network endpoint last. Attribute columns are dynamic +/// in `greptime_trace_v1` (created on first use), so only columns present in +/// a table's schema participate. +const VIRTUAL_DST_CANDIDATES: [(&str, &str); 5] = [ + ("span_attributes.service.peer.name", "virtual_node"), + ("span_attributes.peer.service", "virtual_node"), + ("span_attributes.db.namespace", "database"), + ("span_attributes.db.name", "database"), + ("span_attributes.server.address", "virtual_node"), +]; + +/// Confidence of a virtual-node edge: the peer was named by a client-side +/// attribute, not witnessed by a server span (the RFC requires `< 1.0`). +const VIRTUAL_NODE_CONFIDENCE: f64 = 0.5; + +/// The `calls` derivation: union the client spans and the server spans of all +/// trace tables (each projected to a normalized shape with its own table's +/// `service` identity), left-join clients to their child server spans on +/// `trace_id` + `parent_span_id`, and aggregate RED metrics per 60s window — +/// the plan form of the Tempo servicegraph connector. Unioning spans before +/// the join pairs a client with a server stored in a *different* trace table +/// (`x-greptime-trace-table-name` routing); with a single table it degenerates +/// to the previous self-join. Returns `None` when no trace table has a +/// usable `service` declaration. +/// +/// A client with no matching server span is an edge to a **virtual node** +/// named by [`VIRTUAL_DST_CANDIDATES`], with the client's own status/duration +/// and `confidence < 1.0`. Real pairs win: when a window's edge key holds any +/// pair, the edge reports only the pair population (the RFC pins RED metrics +/// to observed span pairs) and the unmatched clients are suppressed, so one +/// `(window, edge)` never yields two rows. +fn calls_branch(traces: &[CallsSource], window: &GraphQueryWindow) -> DfResult> { + let mut clients: Option = None; + let mut servers: Option = None; for trace in traces { - union_df = union_all(union_df, calls_pairs(&trace.service, trace.scan, window)?)?; + let Some(service) = &trace.service else { + continue; + }; + clients = union_all(clients, client_spans(service, &trace.scan, window)?)?; + servers = union_all(servers, server_spans(service, trace.scan.clone(), window)?)?; } - let Some(pairs) = union_df else { + let (Some(clients), Some(servers)) = (clients, servers) else { return Ok(None); }; - let df = pairs + let client = clients.alias("client")?; + let server = servers.alias("server")?; + let join_conditions = vec![ + qcol("client", TRACE_ID_COLUMN).eq(qcol("server", TRACE_ID_COLUMN)), + qcol("server", PARENT_SPAN_ID_COLUMN).eq(qcol("client", SPAN_ID_COLUMN)), + qcol("server", TRACE_TIMESTAMP_COLUMN) + .gt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) - interval(CHILD_SPAN_EARLY_NANOS)), + qcol("server", TRACE_TIMESTAMP_COLUMN) + .lt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) + interval(CHILD_SPAN_LATE_NANOS)), + ]; + + let observations = client + .join_on(server, JoinType::Left, join_conditions)? + .select(vec![ + bin_ms(qcol("client", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), + qcol("client", "src_id").alias("src_id"), + core_fns::coalesce() + .call(vec![ + qcol("server", "dst_id"), + qcol("client", "virtual_dst"), + ]) + .alias("dst_id"), + qcol("server", TRACE_ID_COLUMN) + .is_not_null() + .alias("paired"), + qcol("server", SPAN_STATUS_CODE_COLUMN).alias("server_status"), + qcol("server", DURATION_NANO_COLUMN).alias("server_duration_nano"), + qcol("client", SPAN_STATUS_CODE_COLUMN).alias("client_status"), + qcol("client", DURATION_NANO_COLUMN).alias("client_duration_nano"), + qcol("client", "virtual_conn").alias("virtual_conn"), + ])? + // No destination means no edge; a self-call is not an edge between + // two distinct entities. + .filter( + ident("dst_id") + .is_not_null() + .and(ident("src_id").not_eq(ident("dst_id"))), + )?; + + let paired = ident("paired"); + let df = observations + .aggregate( + vec![ident("observed_at"), ident("src_id"), ident("dst_id")], + vec![ + count(lit(1)) + .filter(paired.clone()) + .build()? + .alias("pair_count"), + count(lit(1)) + .filter( + paired + .clone() + .and(ident("server_status").eq(lit(SPAN_STATUS_ERROR))), + ) + .build()? + .alias("pair_errors"), + sum(ident("server_duration_nano")) + .filter(paired.clone()) + .build()? + .alias("pair_duration_nano"), + count(lit(1)) + .filter(!paired.clone()) + .build()? + .alias("unmatched_count"), + count(lit(1)) + .filter( + (!paired.clone()).and(ident("client_status").eq(lit(SPAN_STATUS_ERROR))), + ) + .build()? + .alias("unmatched_errors"), + sum(ident("client_duration_nano")) + .filter(!paired.clone()) + .build()? + .alias("unmatched_duration_nano"), + bool_or(paired.clone()).alias("has_pair"), + // `min` makes a mixed-provenance virtual edge deterministic + // (`database` sorts before `virtual_node`). + min(ident("virtual_conn")) + .filter(!paired) + .build()? + .alias("virtual_conn"), + ], + )? + .select(vec![ + ident("observed_at"), + ident("observed_at").alias("window_start"), + (ident("observed_at") + bin_interval()).alias("window_end"), + (ident("observed_at") + bin_interval()).alias("fresh_until"), + lit("service").alias("src_type"), + ident("src_id"), + lit("service").alias("dst_type"), + ident("dst_id"), + lit("calls").alias("rel_type"), + lit("trace").alias("provenance"), + real_wins(lit(1.0_f64), lit(VIRTUAL_NODE_CONFIDENCE))?.alias("confidence"), + real_wins(ident("pair_count"), ident("unmatched_count"))?.alias("request_count"), + real_wins(ident("pair_errors"), ident("unmatched_errors"))?.alias("error_count"), + // duration sums in nanoseconds; the contract column is seconds. + (cast( + real_wins( + ident("pair_duration_nano"), + ident("unmatched_duration_nano"), + )?, + DataType::Float64, + ) / lit(1e9_f64)) + .alias("duration_sum"), + real_wins(ident("pair_count"), ident("unmatched_count"))?.alias("duration_count"), + real_wins(null_json(), virtual_attrs_expr()?)?.alias("attributes"), + ])?; + Ok(Some(df)) +} + +/// `CASE WHEN has_pair THEN real ELSE virtual END` — the real-wins projection +/// over the mixed aggregate. +fn real_wins(real: Expr, r#virtual: Expr) -> DfResult { + when(ident("has_pair"), real).otherwise(r#virtual) +} + +/// The `attributes` JSON of a virtual edge, from the aggregated +/// `connection_type`. +fn virtual_attrs_expr() -> DfResult { + when( + ident("virtual_conn").eq(lit("database")), + parse_json_expr(lit(r#"{"connection_type":"database"}"#)), + ) + .otherwise(parse_json_expr(lit( + r#"{"connection_type":"virtual_node"}"#, + ))) +} + +/// The window + span-kind + identity predicate shared by both join sides. +/// `strict` bounds the scan to the source window; the non-strict side widens +/// by the join's time-proximity allowance (the join bounds reference the other +/// side's timestamp and cannot prune this side's scan on their own). +fn span_predicate(service: &EntityDeclaration, window: &GraphQueryWindow, strict: bool) -> Expr { + let ts = ident(TRACE_TIMESTAMP_COLUMN); + let mut predicate = if strict { + ts.clone() + .gt_eq(window.source_start()) + .and(ts.lt(window.source_end())) + } else { + ts.clone() + .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)) + .and(ts.lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS))) + }; + // A NULL identity component identifies nothing, on either endpoint. + for id in &service.id_columns { + predicate = predicate.and(ident(id).is_not_null()); + } + predicate +} + +/// One trace table's client spans, normalized to +/// `(timestamp, trace_id, span_id, src_id, status_code, duration_nano, +/// virtual_dst, virtual_conn)`. `src_id` is built from the table's `service` +/// declaration, so edges land on exactly the entity ids the registry emits (a +/// composite identity renders the same sorted `k=v` form); the per-table +/// projection is what lets tables with different declarations union. +fn client_spans( + service: &EntityDeclaration, + scan: &DataFrame, + window: &GraphQueryWindow, +) -> DfResult { + // NULLIF('') lets an empty value fall through to the next candidate. + let present: Vec<(Expr, &str)> = { + let schema = scan.schema(); + VIRTUAL_DST_CANDIDATES + .iter() + .filter(|(column, _)| schema.has_column_with_unqualified_name(column)) + .map(|(column, conn)| { + let value = + core_fns::nullif().call(vec![cast(ident(*column), DataType::Utf8), lit("")]); + (value, *conn) + }) + .collect() + }; + let null_utf8 = || lit(ScalarValue::Utf8(None)); + let virtual_dst = if present.is_empty() { + null_utf8() + } else { + core_fns::coalesce().call(present.iter().map(|(value, _)| value.clone()).collect()) + }; + let mut virtual_conn = null_utf8(); + for (value, conn) in present.into_iter().rev() { + virtual_conn = when(value.is_not_null(), lit(conn)).otherwise(virtual_conn)?; + } + + scan.clone() + .filter( + ident(SPAN_KIND_COLUMN) + .eq(lit(SPAN_KIND_CLIENT)) + .and(span_predicate(service, window, true)), + )? + .select(vec![ + ident(TRACE_TIMESTAMP_COLUMN), + ident(TRACE_ID_COLUMN), + ident(SPAN_ID_COLUMN), + // The cast inside entity_id_expr also normalizes tag columns, which + // come out of the storage engine dictionary-encoded. + entity_id_expr(&service.id_columns, &|c| ident(c)).alias("src_id"), + ident(SPAN_STATUS_CODE_COLUMN), + ident(DURATION_NANO_COLUMN), + virtual_dst.alias("virtual_dst"), + virtual_conn.alias("virtual_conn"), + ]) +} + +/// One trace table's server spans, normalized to +/// `(timestamp, trace_id, parent_span_id, dst_id, status_code, duration_nano)`. +fn server_spans( + service: &EntityDeclaration, + trace: DataFrame, + window: &GraphQueryWindow, +) -> DfResult { + trace + .filter( + ident(SPAN_KIND_COLUMN) + .eq(lit(SPAN_KIND_SERVER)) + .and(span_predicate(service, window, false)), + )? + .select(vec![ + ident(TRACE_TIMESTAMP_COLUMN), + ident(TRACE_ID_COLUMN), + ident(PARENT_SPAN_ID_COLUMN), + entity_id_expr(&service.id_columns, &|c| ident(c)).alias("dst_id"), + ident(SPAN_STATUS_CODE_COLUMN), + ident(DURATION_NANO_COLUMN), + ]) +} + +/// The `parent_agent calls agent` derivation over trace tables declaring an +/// `agent` entity: pair each span with its child span (no span-kind filter — +/// agent spans are typically INTERNAL) across all agent-declaring tables, keep +/// pairs whose agent identities differ, and aggregate RED metrics per 60s +/// window. Anchored like the service derivation on the caller: the parent side +/// is bounded to the source window and stamps `observed_at`, the child side +/// widens by the time-proximity allowance and supplies status/duration. +fn agent_calls_branch( + traces: &[CallsSource], + window: &GraphQueryWindow, +) -> DfResult> { + let mut parents: Option = None; + let mut children: Option = None; + for trace in traces { + let Some(agent) = &trace.agent else { + continue; + }; + let parent = trace + .scan + .clone() + .filter(span_predicate(agent, window, true))? + .select(vec![ + ident(TRACE_TIMESTAMP_COLUMN), + ident(TRACE_ID_COLUMN), + ident(SPAN_ID_COLUMN), + entity_id_expr(&agent.id_columns, &|c| ident(c)).alias("src_id"), + ])?; + let child = trace + .scan + .clone() + .filter(span_predicate(agent, window, false))? + .select(vec![ + ident(TRACE_TIMESTAMP_COLUMN), + ident(TRACE_ID_COLUMN), + ident(PARENT_SPAN_ID_COLUMN), + entity_id_expr(&agent.id_columns, &|c| ident(c)).alias("dst_id"), + ident(SPAN_STATUS_CODE_COLUMN), + ident(DURATION_NANO_COLUMN), + ])?; + parents = union_all(parents, parent)?; + children = union_all(children, child)?; + } + let (Some(parents), Some(children)) = (parents, children) else { + return Ok(None); + }; + + let parent = parents.alias("parent")?; + let child = children.alias("child")?; + let join_conditions = vec![ + qcol("parent", TRACE_ID_COLUMN).eq(qcol("child", TRACE_ID_COLUMN)), + qcol("child", PARENT_SPAN_ID_COLUMN).eq(qcol("parent", SPAN_ID_COLUMN)), + qcol("child", TRACE_TIMESTAMP_COLUMN) + .gt_eq(qcol("parent", TRACE_TIMESTAMP_COLUMN) - interval(CHILD_SPAN_EARLY_NANOS)), + qcol("child", TRACE_TIMESTAMP_COLUMN) + .lt_eq(qcol("parent", TRACE_TIMESTAMP_COLUMN) + interval(CHILD_SPAN_LATE_NANOS)), + ]; + + let df = parent + .join_on(child, JoinType::Inner, join_conditions)? + .select(vec![ + bin_ms(qcol("parent", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), + qcol("parent", "src_id").alias("src_id"), + qcol("child", "dst_id").alias("dst_id"), + qcol("child", SPAN_STATUS_CODE_COLUMN).alias("status_code"), + qcol("child", DURATION_NANO_COLUMN).alias("duration_nano"), + ])? + // A sub-span of the same agent is internal structure, not a call + // between two agents. + .filter(ident("src_id").not_eq(ident("dst_id")))? .aggregate( vec![ident("observed_at"), ident("src_id"), ident("dst_id")], vec![ @@ -216,16 +676,15 @@ fn calls_branch( ident("observed_at").alias("window_start"), (ident("observed_at") + bin_interval()).alias("window_end"), (ident("observed_at") + bin_interval()).alias("fresh_until"), - lit("service").alias("src_type"), + lit("agent").alias("src_type"), ident("src_id"), - lit("service").alias("dst_type"), + lit("agent").alias("dst_type"), ident("dst_id"), lit("calls").alias("rel_type"), lit("trace").alias("provenance"), lit(1.0_f64).alias("confidence"), ident("request_count"), ident("error_count"), - // duration_nano sums in nanoseconds; the contract column is seconds. (cast(ident("duration_nano_sum"), DataType::Float64) / lit(1e9_f64)) .alias("duration_sum"), ident("request_count").alias("duration_count"), @@ -234,74 +693,14 @@ fn calls_branch( Ok(Some(df)) } -/// One trace table's client/server span pairs, projected to the aggregation -/// inputs `(observed_at, src_id, dst_id, status_code, duration_nano)`. -/// Endpoint ids are built from the table's `service` entity declaration, so -/// edges land on exactly the entity ids the registry emits (a composite -/// service identity renders the same sorted `k=v` form on both sides). -fn calls_pairs( - service: &EntityDeclaration, - trace: DataFrame, - window: &GraphQueryWindow, -) -> DfResult { - let mut client_pred = ident(SPAN_KIND_COLUMN) - .eq(lit(SPAN_KIND_CLIENT)) - .and(ident(TRACE_TIMESTAMP_COLUMN).gt_eq(window.source_start())) - .and(ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end())); - let mut server_pred = ident(SPAN_KIND_COLUMN) - .eq(lit(SPAN_KIND_SERVER)) - // Static bounds implied by the window and the join's time-proximity - // conditions below; the join bounds reference client.timestamp and - // cannot prune the server-side scan. - .and( - ident(TRACE_TIMESTAMP_COLUMN) - .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)), - ) - .and( - ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS)), - ); - // A NULL identity component identifies nothing, on either endpoint. - for id in &service.id_columns { - client_pred = client_pred.and(ident(id).is_not_null()); - server_pred = server_pred.and(ident(id).is_not_null()); - } - - let client = trace.clone().filter(client_pred)?.alias("client")?; - let server = trace.filter(server_pred)?.alias("server")?; - - let join_conditions = vec![ - qcol("client", TRACE_ID_COLUMN).eq(qcol("server", TRACE_ID_COLUMN)), - qcol("server", PARENT_SPAN_ID_COLUMN).eq(qcol("client", SPAN_ID_COLUMN)), - qcol("server", TRACE_TIMESTAMP_COLUMN) - .gt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) - interval(CHILD_SPAN_EARLY_NANOS)), - qcol("server", TRACE_TIMESTAMP_COLUMN) - .lt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) + interval(CHILD_SPAN_LATE_NANOS)), - ]; - - client - .join_on(server, JoinType::Inner, join_conditions)? - .select(vec![ - bin_ms(qcol("client", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), - // The cast inside entity_id_expr also normalizes tag columns, which - // come out of the storage engine dictionary-encoded. - entity_id_expr(&service.id_columns, &|c| qcol("client", c)).alias("src_id"), - entity_id_expr(&service.id_columns, &|c| qcol("server", c)).alias("dst_id"), - qcol("server", SPAN_STATUS_CODE_COLUMN).alias("status_code"), - qcol("server", DURATION_NANO_COLUMN).alias("duration_nano"), - ])? - // Exclude self-calls on the composed identity: an edge needs two - // distinct service entities. - .filter(ident("src_id").not_eq(ident("dst_id"))) -} - #[cfg(test)] mod tests { use std::sync::Arc; use common_catalog::consts::{SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SERVICE_NAME_COLUMN}; use datafusion::arrow::array::{ - ArrayRef, BinaryArray, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, - TimestampNanosecondArray, UInt64Array, + Array, ArrayRef, BinaryArray, Float64Array, Int64Array, StringArray, + TimestampMillisecondArray, TimestampNanosecondArray, UInt64Array, }; use datafusion::arrow::datatypes::{Field, Schema, TimeUnit}; use datafusion::arrow::record_batch::RecordBatch; @@ -315,9 +714,29 @@ mod tests { GraphQueryWindow::from_observed(0, 10 * 60 * 1000) } - /// A trace-like table in the fixed `greptime_trace_v1` shape (ns timestamps). - fn trace_table_ctx() -> SessionContext { - let schema = Arc::new(Schema::new(vec![ + /// One span row of the fixed `greptime_trace_v1` shape: + /// `(ts_ms, trace_id, span_id, parent_span_id, span_kind, status_code, + /// service_name, duration_nano)`. + type Span<'a> = ( + i64, + &'a str, + &'a str, + Option<&'a str>, + &'a str, + &'a str, + &'a str, + u64, + ); + + /// Registers a trace-v1-shaped table (ns timestamps) plus optional dynamic + /// string columns (`extra`), one value per span. + fn register_trace_table( + ctx: &SessionContext, + name: &str, + extra: &[(&str, &[Option<&str>])], + spans: &[Span<'_>], + ) { + let mut fields = vec![ Field::new( TRACE_TIMESTAMP_COLUMN, DataType::Timestamp(TimeUnit::Nanosecond, None), @@ -329,84 +748,94 @@ mod tests { Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false), Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false), Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false), - Field::new("service_namespace", DataType::Utf8, false), Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false), - ])); + ]; + for (column, _) in extra { + fields.push(Field::new(*column, DataType::Utf8, true)); + } + let schema = Arc::new(Schema::new(fields)); const MS: i64 = 1_000_000; - // Two client->server pairs frontend->cart (one errored), one pair - // cart->cart (self-call, excluded), one unmatched client span. - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampNanosecondArray::from(vec![ - 1_000 * MS, // client frontend->cart - 1_010 * MS, // server cart - 2_000 * MS, // client frontend->cart (error) - 2_010 * MS, // server cart (error) - 3_000 * MS, // client cart->cart (self-call) - 3_010 * MS, // server cart - 4_000 * MS, // client with no matching server - ])) as ArrayRef, - Arc::new(StringArray::from(vec![ - "t1", "t1", "t2", "t2", "t3", "t3", "t4", - ])), - Arc::new(StringArray::from(vec![ - "c1", "s1", "c2", "s2", "c3", "s3", "c4", - ])), - Arc::new(StringArray::from(vec![ - None, - Some("c1"), - None, - Some("c2"), - None, - Some("c3"), - None, - ])), - Arc::new(StringArray::from(vec![ - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - ])), - Arc::new(StringArray::from(vec![ - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_ERROR", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - ])), - Arc::new(StringArray::from(vec![ - "frontend", "cart", "frontend", "cart", "cart", "cart", "frontend", - ])), - Arc::new(StringArray::from(vec!["ns1"; 7])), - Arc::new(UInt64Array::from(vec![ - 0, - 500_000_000, // 0.5s - 0, - 1_500_000_000, // 1.5s - 0, - 100, - 0, - ])), - ], - ) - .unwrap(); - let ctx = SessionContext::new(); + let mut columns: Vec = vec![ + Arc::new(TimestampNanosecondArray::from( + spans.iter().map(|s| s.0 * MS).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.1).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.2).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.3).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.4).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.5).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.6).collect::>(), + )), + Arc::new(UInt64Array::from( + spans.iter().map(|s| s.7).collect::>(), + )), + ]; + for (_, values) in extra { + columns.push(Arc::new(StringArray::from(values.to_vec()))); + } + let batch = RecordBatch::try_new(schema.clone(), columns).unwrap(); ctx.register_table( - "opentelemetry_traces", - Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch.clone()]]).unwrap()), - ) - .unwrap(); - ctx.register_table( - "opentelemetry_traces_2", + name, Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), ) .unwrap(); + } + + const UNSET: &str = "STATUS_CODE_UNSET"; + const ERROR: &str = "STATUS_CODE_ERROR"; + const CLIENT: &str = "SPAN_KIND_CLIENT"; + const SERVER: &str = "SPAN_KIND_SERVER"; + + /// Two client->server pairs frontend->cart (one errored), one pair + /// cart->cart (self-call, excluded), one unmatched client span. + const BASE_SPANS: [Span<'static>; 7] = [ + (1_000, "t1", "c1", None, CLIENT, UNSET, "frontend", 0), + ( + 1_010, + "t1", + "s1", + Some("c1"), + SERVER, + UNSET, + "cart", + 500_000_000, + ), + (2_000, "t2", "c2", None, CLIENT, UNSET, "frontend", 0), + ( + 2_010, + "t2", + "s2", + Some("c2"), + SERVER, + ERROR, + "cart", + 1_500_000_000, + ), + (3_000, "t3", "c3", None, CLIENT, UNSET, "cart", 0), + (3_010, "t3", "s3", Some("c3"), SERVER, UNSET, "cart", 100), + (4_000, "t4", "c4", None, CLIENT, UNSET, "frontend", 0), + ]; + + fn trace_table_ctx() -> SessionContext { + let ctx = SessionContext::new(); + let namespaces = [Some("ns1"); 7]; + register_trace_table( + &ctx, + "opentelemetry_traces", + &[("service_namespace", &namespaces)], + &BASE_SPANS, + ); ctx } @@ -428,7 +857,8 @@ mod tests { let trace = ctx.table("opentelemetry_traces").await.unwrap(); let plan = build_relationships_plan_for_test( vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + service: Some(trace_service_decl(&[SERVICE_NAME_COLUMN])), + agent: None, scan: trace, }], &test_window(), @@ -483,22 +913,272 @@ mod tests { assert!(json_texts(batch, 15)[0].is_none()); } + /// RED columns of one output row: `(request_count, error_count, + /// duration_sum, duration_count)`. + fn red_metrics(batch: &RecordBatch, row: usize) -> (i64, i64, f64, i64) { + let i64_at = |column: usize| { + batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap() + .value(row) + }; + let duration_sum = batch + .column(13) + .as_any() + .downcast_ref::() + .unwrap() + .value(row); + (i64_at(11), i64_at(12), duration_sum, i64_at(14)) + } + + fn confidence(batch: &RecordBatch, row: usize) -> f64 { + batch + .column(10) + .as_any() + .downcast_ref::() + .unwrap() + .value(row) + } + + fn sources(scans: &[DataFrame]) -> Vec { + scans + .iter() + .map(|scan| CallsSource { + service: Some(trace_service_decl(&[SERVICE_NAME_COLUMN])), + agent: None, + scan: scan.clone(), + }) + .collect() + } + #[tokio::test] async fn calls_plan_merges_edges_across_trace_tables() { - let ctx = trace_table_ctx(); - let t1 = ctx.table("opentelemetry_traces").await.unwrap(); - let t2 = ctx.table("opentelemetry_traces_2").await.unwrap(); - let plan = build_relationships_plan_for_test( - vec![ - CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: t1, - }, - CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: t2, - }, + let ctx = SessionContext::new(); + // Distinct traces observing the same frontend->cart edge, one per table. + register_trace_table( + &ctx, + "trace_a", + &[], + &[ + (1_000, "t1", "c1", None, CLIENT, UNSET, "frontend", 0), + ( + 1_010, + "t1", + "s1", + Some("c1"), + SERVER, + UNSET, + "cart", + 500_000_000, + ), ], + ); + register_trace_table( + &ctx, + "trace_b", + &[], + &[ + (2_000, "t2", "c2", None, CLIENT, UNSET, "frontend", 0), + ( + 2_010, + "t2", + "s2", + Some("c2"), + SERVER, + ERROR, + "cart", + 1_500_000_000, + ), + ], + ); + let a = ctx.table("trace_a").await.unwrap(); + let b = ctx.table("trace_b").await.unwrap(); + let plan = build_relationships_plan_for_test(sources(&[a, b]), &test_window()) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 1); + assert_eq!(red_metrics(&batches[0], 0), (2, 1, 2.0, 2)); + } + + #[tokio::test] + async fn calls_pair_split_across_trace_tables() { + let ctx = SessionContext::new(); + // The client span and its child server span land in different tables + // (per-table trace routing): the union-then-join pairing must find it. + register_trace_table( + &ctx, + "trace_a", + &[], + &[(1_000, "t1", "c1", None, CLIENT, UNSET, "frontend", 0)], + ); + register_trace_table( + &ctx, + "trace_b", + &[], + &[( + 1_010, + "t1", + "s1", + Some("c1"), + SERVER, + UNSET, + "cart", + 500_000_000, + )], + ); + let a = ctx.table("trace_a").await.unwrap(); + let b = ctx.table("trace_b").await.unwrap(); + let plan = build_relationships_plan_for_test(sources(&[a, b]), &test_window()) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 1); + let batch = &batches[0]; + assert_eq!(strings(batch, 5), vec!["frontend"]); + assert_eq!(strings(batch, 7), vec!["cart"]); + assert_eq!(confidence(batch, 0), 1.0); + assert_eq!(red_metrics(batch, 0), (1, 0, 0.5, 1)); + } + + #[tokio::test] + async fn virtual_node_edge_from_unmatched_client() { + let ctx = SessionContext::new(); + register_trace_table( + &ctx, + "trace_a", + &[( + "span_attributes.peer.service", + &[Some("redis"), None, Some("frontend")], + )], + &[ + // Unmatched client naming its peer: a virtual-node edge with + // the client's own status and duration. + ( + 1_000, + "t1", + "c1", + None, + CLIENT, + ERROR, + "frontend", + 250_000_000, + ), + // Unmatched client without a peer attribute: no edge. + (2_000, "t2", "c2", None, CLIENT, UNSET, "frontend", 0), + // Peer attribute naming the client's own service: excluded. + (3_000, "t3", "c3", None, CLIENT, UNSET, "frontend", 0), + ], + ); + let a = ctx.table("trace_a").await.unwrap(); + let plan = build_relationships_plan_for_test(sources(&[a]), &test_window()) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 1); + let batch = &batches[0]; + assert_eq!(strings(batch, 5), vec!["frontend"]); + assert_eq!(strings(batch, 7), vec!["redis"]); + assert_eq!(strings(batch, 8), vec!["calls"]); + assert_eq!(strings(batch, 9), vec!["trace"]); + assert_eq!(confidence(batch, 0), VIRTUAL_NODE_CONFIDENCE); + assert_eq!(red_metrics(batch, 0), (1, 1, 0.25, 1)); + assert_eq!( + json_texts(batch, 15), + vec![Some(r#"{"connection_type":"virtual_node"}"#.to_string())] + ); + } + + #[tokio::test] + async fn virtual_endpoint_fallback_and_connection_type() { + let ctx = SessionContext::new(); + register_trace_table( + &ctx, + "trace_a", + &[ + // An empty high-precedence value must fall through to the next + // candidate, whose match maps connection_type to `database`. + ("span_attributes.service.peer.name", &[Some("")]), + ("span_attributes.db.namespace", &[Some("mysql")]), + ], + &[(1_000, "t1", "c1", None, CLIENT, UNSET, "frontend", 100)], + ); + let a = ctx.table("trace_a").await.unwrap(); + let plan = build_relationships_plan_for_test(sources(&[a]), &test_window()) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let batch = &batches[0]; + assert_eq!(strings(batch, 7), vec!["mysql"]); + assert_eq!( + json_texts(batch, 15), + vec![Some(r#"{"connection_type":"database"}"#.to_string())] + ); + } + + #[tokio::test] + async fn agent_calls_from_parent_child_spans() { + let ctx = SessionContext::new(); + const INTERNAL: &str = "SPAN_KIND_INTERNAL"; + register_trace_table( + &ctx, + "agent_traces", + &[( + "agent_id", + &[ + Some("orchestrator"), + Some("researcher"), + Some("researcher"), + Some("researcher"), + ], + )], + &[ + // orchestrator delegates to researcher twice (one errored); + // the researcher's own sub-span is same-agent structure. + (1_000, "t1", "p1", None, INTERNAL, UNSET, "app", 0), + ( + 1_010, + "t1", + "a1", + Some("p1"), + INTERNAL, + ERROR, + "app", + 2_000_000_000, + ), + ( + 2_000, + "t1", + "a2", + Some("p1"), + INTERNAL, + UNSET, + "app", + 1_000_000_000, + ), + (1_020, "t1", "a3", Some("a1"), INTERNAL, UNSET, "app", 100), + ], + ); + let scan = ctx.table("agent_traces").await.unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![CallsSource { + service: None, + agent: Some(co_decl("agent", &["agent_id"])), + scan, + }], + co_declared: vec![], + declared: None, + }, &test_window(), ) .unwrap() @@ -506,28 +1186,343 @@ mod tests { let batches = collect(&ctx, plan).await; let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - // The same frontend->cart edge from both tables folds into one row with - // summed RED metrics. + // The same-agent child (a3 under a1) is excluded; agent edges need + // no service declaration. assert_eq!(total, 1); let batch = &batches[0]; - let request_count = batch - .column(11) - .as_any() - .downcast_ref::() + assert_eq!(strings(batch, 4), vec!["agent"]); + assert_eq!(strings(batch, 5), vec!["orchestrator"]); + assert_eq!(strings(batch, 6), vec!["agent"]); + assert_eq!(strings(batch, 7), vec!["researcher"]); + assert_eq!(strings(batch, 8), vec!["calls"]); + assert_eq!(strings(batch, 9), vec!["trace"]); + assert_eq!(confidence(batch, 0), 1.0); + assert_eq!(red_metrics(batch, 0), (2, 1, 3.0, 2)); + } + + #[tokio::test] + async fn agent_calls_anchor_on_the_parent_span() { + let ctx = SessionContext::new(); + const INTERNAL: &str = "SPAN_KIND_INTERNAL"; + // The parent is inside the queried window; the child starts after its + // end but within the proximity allowance and must still pair. + register_trace_table( + &ctx, + "agent_traces", + &[("agent_id", &[Some("orchestrator"), Some("researcher")])], + &[ + (599_000, "t1", "p1", None, INTERNAL, UNSET, "app", 0), + ( + 601_000, + "t1", + "a1", + Some("p1"), + INTERNAL, + UNSET, + "app", + 500_000_000, + ), + ], + ); + let scan = ctx.table("agent_traces").await.unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![CallsSource { + service: None, + agent: Some(co_decl("agent", &["agent_id"])), + scan, + }], + co_declared: vec![], + declared: None, + }, + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 1); + let batch = &batches[0]; + // observed_at is the parent's bin. + assert_eq!(ts_values(batch, 0), vec![540_000]); + assert_eq!(red_metrics(batch, 0), (1, 0, 0.5, 1)); + } + + /// A metric-like declaring table: ms timestamps, one row with a NULL + /// `instance` (witnessing nothing for instance edges), two rows duplicating + /// the same identities inside one bin (collapsed by DISTINCT). + fn co_declared_ctx() -> SessionContext { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("instance", DataType::Utf8, true), + Field::new("host", DataType::Utf8, false), + Field::new("service", DataType::Utf8, false), + Field::new("agent_id", DataType::Utf8, false), + Field::new("model_name", DataType::Utf8, false), + Field::new("tool_name", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 2_000, 3_000])) as ArrayRef, + Arc::new(StringArray::from(vec![Some("i-1"), Some("i-1"), None])), + Arc::new(StringArray::from(vec!["h-1", "h-1", "h-2"])), + Arc::new(StringArray::from(vec!["svc", "svc", "svc"])), + Arc::new(StringArray::from(vec!["agent-1"; 3])), + Arc::new(StringArray::from(vec!["gpt"; 3])), + Arc::new(StringArray::from(vec!["search"; 3])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + for name in ["co_metrics", "co_metrics_2"] { + ctx.register_table( + name, + Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch.clone()]]).unwrap()), + ) .unwrap(); - assert_eq!(request_count.value(0), 4); - let error_count = batch - .column(12) - .as_any() - .downcast_ref::() + } + ctx + } + + fn co_decl(entity_type: &str, id_columns: &[&str]) -> EntityDeclaration { + EntityDeclaration { + schema: "public".to_string(), + table: "co_metrics".to_string(), + time_index: "ts".to_string(), + entity_type: entity_type.to_string(), + id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + descriptive_columns: vec![], + scope_columns: vec![], + } + } + + async fn co_declared_edges( + ctx: &SessionContext, + declarations: Vec, + is_trace: bool, + ) -> Option> { + let scan = ctx.table("co_metrics").await.unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![], + co_declared: vec![CoDeclaredSource { + declarations, + is_trace, + scan, + }], + declared: None, + }, + &test_window(), + ) + .unwrap()?; + assert_eq!( + plan.schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(), + RELATIONSHIP_COLUMNS + ); + let batches = collect(ctx, plan).await; + let mut rows = vec![]; + for batch in &batches { + for i in 0..batch.num_rows() { + rows.push(( + strings(batch, 4)[i].clone(), + strings(batch, 5)[i].clone(), + strings(batch, 6)[i].clone(), + strings(batch, 7)[i].clone(), + strings(batch, 8)[i].clone(), + strings(batch, 9)[i].clone(), + )); + assert_eq!(confidence(batch, i), 1.0); + // Co-declared edges carry no RED metrics or attributes. + for column in 11..=15 { + assert!(batch.column(column).is_null(i)); + } + } + } + rows.sort(); + Some(rows) + } + + #[tokio::test] + async fn co_declared_direction_follows_vocabulary() { + let ctx = co_declared_ctx(); + // Declaration order must not matter: the vocabulary fixes direction. + let rows = co_declared_edges( + &ctx, + vec![ + co_decl("host", &["host"]), + co_decl("service", &["service"]), + co_decl("service.instance", &["instance"]), + ], + false, + ) + .await + .unwrap(); + // The two non-NULL-instance rows share one 60s bin, so DISTINCT folds + // them into one row per edge; the NULL-instance row witnesses nothing. + assert_eq!( + rows, + vec![ + ( + "service.instance".to_string(), + "i-1".to_string(), + "host".to_string(), + "h-1".to_string(), + "runs_on".to_string(), + "attribute".to_string(), + ), + ( + "service.instance".to_string(), + "i-1".to_string(), + "service".to_string(), + "svc".to_string(), + "part_of".to_string(), + "attribute".to_string(), + ), + ] + ); + } + + #[tokio::test] + async fn co_declared_edge_folds_across_sources() { + let ctx = co_declared_ctx(); + let declarations = || { + vec![ + co_decl("service.instance", &["instance"]), + co_decl("host", &["host"]), + ] + }; + let source = |scan| CoDeclaredSource { + declarations: declarations(), + is_trace: false, + scan, + }; + let a = ctx.table("co_metrics").await.unwrap(); + let b = ctx.table("co_metrics_2").await.unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![], + co_declared: vec![source(a), source(b)], + declared: None, + }, + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + // Both tables witness the same runs_on edge in the same window: one + // row per (window, edge), not one per source. + assert_eq!(total, 1); + } + + #[tokio::test] + async fn co_declared_ignores_types_outside_the_vocabulary() { + let ctx = co_declared_ctx(); + // (service, host) is not a vocabulary pair: no edge, no plan. + let rows = co_declared_edges( + &ctx, + vec![co_decl("service", &["service"]), co_decl("host", &["host"])], + false, + ) + .await; + assert!(rows.is_none()); + } + + #[tokio::test] + async fn agent_edges_require_a_trace_source() { + let ctx = co_declared_ctx(); + let declarations = || { + vec![ + co_decl("agent", &["agent_id"]), + co_decl("model", &["model_name"]), + co_decl("tool", &["tool_name"]), + ] + }; + // A non-trace table co-declaring agent/model/tool must not fabricate + // invocations. + assert!( + co_declared_edges(&ctx, declarations(), false) + .await + .is_none() + ); + + let rows = co_declared_edges(&ctx, declarations(), true).await.unwrap(); + assert_eq!( + rows, + vec![ + ( + "agent".to_string(), + "agent-1".to_string(), + "model".to_string(), + "gpt".to_string(), + "uses".to_string(), + "trace".to_string(), + ), + ( + "agent".to_string(), + "agent-1".to_string(), + "tool".to_string(), + "search".to_string(), + "invokes".to_string(), + "trace".to_string(), + ), + ] + ); + } + + #[tokio::test] + async fn real_pair_suppresses_virtual_edge() { + let ctx = SessionContext::new(); + register_trace_table( + &ctx, + "trace_a", + &[( + "span_attributes.peer.service", + &[Some("cart"), None, Some("cart")], + )], + &[ + // A sampled-out server: the client names the same peer an + // actual pair witnesses in the same window. + (1_000, "t1", "c1", None, CLIENT, ERROR, "frontend", 999), + ( + 2_010, + "t2", + "s2", + Some("c2"), + SERVER, + UNSET, + "cart", + 500_000_000, + ), + (2_000, "t2", "c2", None, CLIENT, UNSET, "frontend", 0), + ], + ); + let a = ctx.table("trace_a").await.unwrap(); + let plan = build_relationships_plan_for_test(sources(&[a]), &test_window()) + .unwrap() .unwrap(); - assert_eq!(error_count.value(0), 2); - let duration_sum = batch - .column(13) - .as_any() - .downcast_ref::() - .unwrap(); - assert!((duration_sum.value(0) - 4.0).abs() < 1e-9); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + // One (window, edge) row: the pair population wins; the unmatched + // client neither adds a second row nor inflates the pair RED metrics. + assert_eq!(total, 1); + let batch = &batches[0]; + assert_eq!(strings(batch, 7), vec!["cart"]); + assert_eq!(confidence(batch, 0), 1.0); + assert_eq!(red_metrics(batch, 0), (1, 0, 0.5, 1)); + assert!(json_texts(batch, 15)[0].is_none()); } #[tokio::test] @@ -536,7 +1531,11 @@ mod tests { let trace = ctx.table("opentelemetry_traces").await.unwrap(); let plan = build_relationships_plan_for_test( vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN, "service_namespace"]), + service: Some(trace_service_decl(&[ + SERVICE_NAME_COLUMN, + "service_namespace", + ])), + agent: None, scan: trace, }], &test_window(), @@ -564,7 +1563,14 @@ mod tests { traces: Vec, window: &GraphQueryWindow, ) -> DfResult> { - build_relationships_plan(traces, None, window) + build_relationships_plan( + RelationshipSources { + traces, + co_declared: vec![], + declared: None, + }, + window, + ) } /// Registers a declared-edge table holding: @@ -720,9 +1726,16 @@ mod tests { // Queried window [120s, 600s). let window = GraphQueryWindow::from_observed(120_000, 600_000); - let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) - .unwrap() - .unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![], + co_declared: vec![], + declared: Some(DeclaredSource { scan }), + }, + &window, + ) + .unwrap() + .unwrap(); let names = plan .schema() .fields() @@ -836,9 +1849,16 @@ mod tests { // the first revision was the edge's state then and must not be hidden // by the newer one. let window = GraphQueryWindow::from_observed(0, 150_000); - let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) - .unwrap() - .unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![], + co_declared: vec![], + declared: Some(DeclaredSource { scan }), + }, + &window, + ) + .unwrap() + .unwrap(); let batches = collect(&ctx, plan).await; let mut rows: Vec<(String, f64)> = vec![]; @@ -879,11 +1899,15 @@ mod tests { .unwrap(); let plan = build_relationships_plan( - vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: trace, - }], - Some(DeclaredSource { scan: declared }), + RelationshipSources { + traces: vec![CallsSource { + service: Some(trace_service_decl(&[SERVICE_NAME_COLUMN])), + agent: None, + scan: trace, + }], + co_declared: vec![], + declared: Some(DeclaredSource { scan: declared }), + }, &test_window(), ) .unwrap() diff --git a/tests-integration/src/instance.rs b/tests-integration/src/instance.rs index 547f64436f6..08d2a1aff15 100644 --- a/tests-integration/src/instance.rs +++ b/tests-integration/src/instance.rs @@ -563,7 +563,9 @@ mod tests { ) -> auth::error::Result { if matches!(req, auth::PermissionReq::Action(auth::SEMANTIC_GRAPH_QUERY)) && let auth::PermissionTableTargets::Resolved(targets) = targets - && targets.iter().any(|target| target.table == "secret_traces") + && targets + .iter() + .any(|target| target.table.starts_with("secret_")) { return Ok(auth::PermissionResp::Reject); } @@ -580,7 +582,7 @@ mod tests { .await; let instance = standalone.fe_instance().clone(); - for table in ["open_traces", "secret_traces"] { + for (table, trace_id) in [("open_traces", "t0"), ("secret_traces", "t1")] { let create = format!( r#"create table {table} ( "timestamp" timestamp(9) time index, @@ -597,8 +599,8 @@ mod tests { create_table(&instance, &create).await; let insert = format!( "insert into {table} values \ - (now(), 't0', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'client-{table}', 0), \ - (now(), 't0', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'server-{table}', 100)" + (now(), '{trace_id}', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'client-{table}', 0), \ + (now(), '{trace_id}', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'server-{table}', 100)" ); let output = query(&instance, &insert).await; let OutputData::AffectedRows(x) = output.data else { @@ -607,6 +609,51 @@ mod tests { assert_eq!(x, 2); } + // A pair split across tables (client here, server in a denied table): + // a join-derived edge needs read access to all of its input tables. + create_table( + &instance, + r#"create table cross_clients ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) + ) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true')"#, + ) + .await; + create_table( + &instance, + r#"create table secret_cross_servers ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) + ) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true')"#, + ) + .await; + for insert in [ + "insert into cross_clients values \ + (now(), 't9', 'c9', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'client-cross', 0)", + "insert into secret_cross_servers values \ + (now(), 't9', 's9', 'c9', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'server-cross', 100)", + ] { + let output = query(&instance, insert).await; + let OutputData::AffectedRows(x) = output.data else { + unreachable!() + }; + assert_eq!(x, 1); + } + let sql = "select src_id, dst_id from greptime_private.semantic_relationships \ order by src_id"; let output = query(&instance, sql).await; @@ -623,6 +670,10 @@ mod tests { !pretty_print.contains("secret_traces"), "denied source leaked into edges:\n{pretty_print}" ); + assert!( + !pretty_print.contains("client-cross"), + "edge with a denied join side leaked:\n{pretty_print}" + ); let sql = "select entity_id from greptime_private.semantic_entities order by entity_id"; let output = query(&instance, sql).await; @@ -636,7 +687,7 @@ mod tests { "allowed entity missing:\n{pretty_print}" ); assert!( - !pretty_print.contains("secret_traces"), + !pretty_print.contains("secret_traces") && !pretty_print.contains("server-cross"), "denied source leaked into entities:\n{pretty_print}" ); } diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index 73f00f57c6f..6617d9f839d 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -277,3 +277,204 @@ drop table greptime_private.semantic_relationships_declared; Affected Rows: 0 +-- Same-row co-declaration: a table declaring both entity types of a built-in +-- vocabulary pair witnesses the edge on every row carrying both identities; +-- the vocabulary fixes the direction. +create table graph_pod_metrics ( + ts timestamp time index, + instance string, + host string, + "service" string, + cpu double, + primary key (instance, host, "service") +) with ( + 'greptime.semantic.entity.service.instance.id' = 'instance', + 'greptime.semantic.entity.host.id' = 'host', + 'greptime.semantic.entity.service.id' = 'service' +); + +Affected Rows: 0 + +insert into graph_pod_metrics values (now(), 'cart-0', 'node-1', 'cart', 0.5); + +Affected Rows: 1 + +select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by rel_type, src_id; + ++------------------+--------+----------+--------+----------+------------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | ++------------------+--------+----------+--------+----------+------------+------------+ +| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | +| service.instance | cart-0 | host | node-1 | runs_on | attribute | 1.0 | ++------------------+--------+----------+--------+----------+------------+------------+ + +drop table graph_pod_metrics; + +Affected Rows: 0 + +-- Cross-table calls pairing: the client span and its child server span land in +-- different trace tables and must still pair into one edge. A trace-model +-- table that lost the fixed span columns is skipped instead of failing the +-- scan. +create table graph_traces_a ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +create table graph_traces_b ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +create table graph_traces_malformed ( + ts timestamp time index, + note string +) with ('table_data_model' = 'greptime_trace_v1'); + +Affected Rows: 0 + +insert into graph_traces_a values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0); + +Affected Rows: 1 + +insert into graph_traces_b values + (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000); + +Affected Rows: 1 + +insert into graph_traces_malformed values (now(), 'not a trace'); + +Affected Rows: 1 + +select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count +from greptime_private.semantic_relationships +order by src_id; + ++----------+--------+----------+------------+------------+---------------+-------------+ +| src_id | dst_id | rel_type | provenance | confidence | request_count | error_count | ++----------+--------+----------+------------+------------+---------------+-------------+ +| frontend | cart | calls | trace | 1.0 | 1 | 1 | ++----------+--------+----------+------------+------------+---------------+-------------+ + +drop table graph_traces_a; + +Affected Rows: 0 + +drop table graph_traces_b; + +Affected Rows: 0 + +drop table graph_traces_malformed; + +Affected Rows: 0 + +-- Virtual nodes: a client span with no matching server span names its peer +-- through span attributes; the edge carries confidence < 1.0 and the +-- connection type. +create table graph_traces_virtual ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "span_attributes.peer.service" string, + "span_attributes.db.name" string, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +insert into graph_traces_virtual values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), + (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); + +Affected Rows: 2 + +-- SQLNESS PROTOCOL MYSQL +select src_id, dst_id, rel_type, confidence, request_count, attributes +from greptime_private.semantic_relationships +order by dst_id; + ++----------+-----------+----------+------------+---------------+------------------------------------+ +| src_id | dst_id | rel_type | confidence | request_count | attributes | ++----------+-----------+----------+------------+---------------+------------------------------------+ +| frontend | orders-db | calls | 0.5 | 1 | {"connection_type":"database"} | +| frontend | redis | calls | 0.5 | 1 | {"connection_type":"virtual_node"} | ++----------+-----------+----------+------------+---------------+------------------------------------+ + +drop table graph_traces_virtual; + +Affected Rows: 0 + +-- Agent edges: span structure derives parent_agent calls agent, and span rows +-- co-declaring agent+model / agent+tool witness uses / invokes. +create table graph_agent_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + agent_id string, + model_name string, + tool_name string, + primary key (service_name, agent_id, model_name, tool_name) +) with ( + 'table_data_model' = 'greptime_trace_v1', + 'append_mode' = 'true', + 'greptime.semantic.entity.agent.id' = 'agent_id', + 'greptime.semantic.entity.model.id' = 'model_name', + 'greptime.semantic.entity.tool.id' = 'tool_name' +); + +Affected Rows: 0 + +insert into graph_agent_traces values + (now(), 't1', 'p1', NULL, 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 0, 'orchestrator', NULL, NULL), + (now(), 't1', 'a1', 'p1', 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 2000000000, 'researcher', 'gpt-5', NULL), + (now(), 't1', 'a2', 'a1', 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 500000000, 'researcher', NULL, 'web_search'); + +Affected Rows: 3 + +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, dst_id; + ++----------+--------------+----------+------------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++----------+--------------+----------+------------+----------+------------+ +| agent | orchestrator | agent | researcher | calls | trace | +| agent | researcher | tool | web_search | invokes | trace | +| agent | researcher | model | gpt-5 | uses | trace | ++----------+--------------+----------+------------+----------+------------+ + +drop table graph_agent_traces; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index c5a93d6c6b7..9acb003331a 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -161,3 +161,140 @@ values (now(), 'service', 'reborn', 'depends_on', 'service', 'db', 'declared', ' select src_id from greptime_private.semantic_relationships order by src_id; drop table greptime_private.semantic_relationships_declared; + +-- Same-row co-declaration: a table declaring both entity types of a built-in +-- vocabulary pair witnesses the edge on every row carrying both identities; +-- the vocabulary fixes the direction. +create table graph_pod_metrics ( + ts timestamp time index, + instance string, + host string, + "service" string, + cpu double, + primary key (instance, host, "service") +) with ( + 'greptime.semantic.entity.service.instance.id' = 'instance', + 'greptime.semantic.entity.host.id' = 'host', + 'greptime.semantic.entity.service.id' = 'service' +); + +insert into graph_pod_metrics values (now(), 'cart-0', 'node-1', 'cart', 0.5); + +select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by rel_type, src_id; + +drop table graph_pod_metrics; + +-- Cross-table calls pairing: the client span and its child server span land in +-- different trace tables and must still pair into one edge. A trace-model +-- table that lost the fixed span columns is skipped instead of failing the +-- scan. +create table graph_traces_a ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +create table graph_traces_b ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +create table graph_traces_malformed ( + ts timestamp time index, + note string +) with ('table_data_model' = 'greptime_trace_v1'); + +insert into graph_traces_a values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0); + +insert into graph_traces_b values + (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000); + +insert into graph_traces_malformed values (now(), 'not a trace'); + +select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count +from greptime_private.semantic_relationships +order by src_id; + +drop table graph_traces_a; + +drop table graph_traces_b; + +drop table graph_traces_malformed; + +-- Virtual nodes: a client span with no matching server span names its peer +-- through span attributes; the edge carries confidence < 1.0 and the +-- connection type. +create table graph_traces_virtual ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "span_attributes.peer.service" string, + "span_attributes.db.name" string, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +insert into graph_traces_virtual values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), + (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); + +-- SQLNESS PROTOCOL MYSQL +select src_id, dst_id, rel_type, confidence, request_count, attributes +from greptime_private.semantic_relationships +order by dst_id; + +drop table graph_traces_virtual; + +-- Agent edges: span structure derives parent_agent calls agent, and span rows +-- co-declaring agent+model / agent+tool witness uses / invokes. +create table graph_agent_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + agent_id string, + model_name string, + tool_name string, + primary key (service_name, agent_id, model_name, tool_name) +) with ( + 'table_data_model' = 'greptime_trace_v1', + 'append_mode' = 'true', + 'greptime.semantic.entity.agent.id' = 'agent_id', + 'greptime.semantic.entity.model.id' = 'model_name', + 'greptime.semantic.entity.tool.id' = 'tool_name' +); + +insert into graph_agent_traces values + (now(), 't1', 'p1', NULL, 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 0, 'orchestrator', NULL, NULL), + (now(), 't1', 'a1', 'p1', 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 2000000000, 'researcher', 'gpt-5', NULL), + (now(), 't1', 'a2', 'a1', 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 500000000, 'researcher', NULL, 'web_search'); + +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, dst_id; + +drop table graph_agent_traces; From 3b3a9032e06faff55819701e2836251ceb90206f Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Wed, 12 Aug 2026 04:59:10 +0000 Subject: [PATCH 054/126] feat(servers): expose native histograms over Prometheus HTTP (#8850) * feat(servers): expose native histograms over Prometheus HTTP Signed-off-by: shuiyisong * fix(servers): refine Prometheus HTTP metadata handling Signed-off-by: shuiyisong * test(servers): expand Prometheus metadata coverage Signed-off-by: shuiyisong * fix(servers): return OpenMetrics units from metadata API Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/servers/src/grpc/prom_query_gateway.rs | 9 +- src/servers/src/http.rs | 5 +- src/servers/src/http/prometheus.rs | 490 +++++++++++++++++- .../src/http/result/prometheus_resp.rs | 79 ++- src/servers/src/otlp/metrics.rs | 1 + src/servers/src/otlp/metrics/translator.rs | 14 + tests-integration/tests/grpc.rs | 1 + tests-integration/tests/http.rs | 218 ++++++++ 8 files changed, 795 insertions(+), 22 deletions(-) diff --git a/src/servers/src/grpc/prom_query_gateway.rs b/src/servers/src/grpc/prom_query_gateway.rs index 6f4151ce683..99940ad9a24 100644 --- a/src/servers/src/grpc/prom_query_gateway.rs +++ b/src/servers/src/grpc/prom_query_gateway.rs @@ -134,12 +134,19 @@ impl PrometheusGatewayService { } }; let (metric_name, mut result_type) = retrieve_metric_name_and_result_type(query.expr()); + let query_id = ctx.remote_query_id().map(str::to_string); let result = self.handler.do_query_parsed(query, ctx).await; // range query only returns matrix if is_range_query { result_type = ValueType::Matrix; }; - PrometheusJsonResponse::from_query_result(result, metric_name, result_type).await + PrometheusJsonResponse::from_query_result( + result, + metric_name, + result_type, + query_id.as_deref(), + ) + .await } } diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index d52bb307012..bf7478eacfd 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -60,8 +60,8 @@ use crate::http::influxdb::{influxdb_health, influxdb_ping, influxdb_write_v1, i use crate::http::otlp::OtlpState; use crate::http::prom_store::PromStoreState; use crate::http::prometheus::{ - build_info_query, format_query, instant_query, label_values_query, labels_query, parse_query, - range_query, series_query, + build_info_query, format_query, instant_query, label_values_query, labels_query, + metadata_query, parse_query, range_query, series_query, }; use crate::http::result::arrow_result::ArrowResponse; use crate::http::result::csv_result::CsvResponse; @@ -1341,6 +1341,7 @@ impl HttpServer { .route("/query", routing::post(instant_query).get(instant_query)) .route("/query_range", routing::post(range_query).get(range_query)) .route("/labels", routing::post(labels_query).get(labels_query)) + .route("/metadata", routing::get(metadata_query)) .route("/series", routing::post(series_query).get(series_query)) .route("/parse_query", routing::post(parse_query).get(parse_query)) .route( diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index 708eba652dc..218a5ae68f5 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -34,6 +34,7 @@ use common_catalog::parse_catalog_and_schema_from_db_string; use common_decimal::Decimal128; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; +use common_query::native_histogram::is_native_histogram_value_type; use common_query::{Output, OutputData}; use common_recordbatch::{RecordBatch, RecordBatches}; use common_telemetry::{debug, tracing}; @@ -64,6 +65,8 @@ use store_api::metric_engine_consts::{ DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY, }; use table::TableRef; +use table::metadata::TableInfo; +use table::requests::{SEMANTIC_METRIC_TEMPORALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT}; pub use super::result::prometheus_resp::PrometheusJsonResponse; use crate::error::{ @@ -71,6 +74,7 @@ use crate::error::{ NotSupportedSnafu, Result, TableNotFoundSnafu, UnexpectedResultSnafu, }; use crate::http::header::collect_plan_metrics; +use crate::otlp::metrics::ucum_to_openmetrics_unit; use crate::prom_store::{FIELD_NAME_LABEL, METRIC_NAME_LABEL, is_database_selection_label}; use crate::prometheus_handler::{ ParsedPromQuery, PrometheusHandlerRef, resolve_schema_from_matchers, @@ -136,6 +140,15 @@ pub struct PromData { pub result: PromQueryResult, } +/// Metadata for a Prometheus metric family. +#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct PromMetadata { + #[serde(rename = "type")] + pub metric_type: String, + pub unit: String, + pub help: String, +} + /// A "holder" for the reference([Arc]) to a column name, /// to help avoiding cloning [String]s when used as a [HashMap] key. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] @@ -153,6 +166,7 @@ pub enum PrometheusResponse { PromData(PromData), Labels(Vec), Series(Vec>), + Metadata(BTreeMap>), LabelValues(Vec), FormatQuery(String), BuildInfo(OwnedBuildInfo), @@ -166,7 +180,7 @@ impl PrometheusResponse { /// Append the other [`PrometheusResponse]`. /// # NOTE /// Only append matrix and vector results, otherwise just ignore the other response. - fn append(&mut self, other: PrometheusResponse) { + pub(super) fn append(&mut self, other: PrometheusResponse) { match (self, other) { ( PrometheusResponse::PromData(PromData { @@ -209,6 +223,14 @@ pub struct FormatQuery { query: Option, } +/// Query parameters for the Prometheus metric metadata endpoint. +#[derive(Debug, Default, Serialize, Deserialize)] +pub struct MetadataQuery { + db: Option, + limit: Option, + metric: Option, +} + #[axum_macros::debug_handler] #[tracing::instrument( skip_all, @@ -246,6 +268,67 @@ pub async fn build_info_query() -> PrometheusJsonResponse { PrometheusJsonResponse::success(PrometheusResponse::BuildInfo(build_info.into())) } +#[axum_macros::debug_handler] +#[tracing::instrument( + skip_all, + fields(protocol = "prometheus", request_type = "metadata_query") +)] +pub async fn metadata_query( + State(handler): State, + Query(params): Query, + Extension(mut query_ctx): Extension, +) -> PrometheusJsonResponse { + let (catalog, schema) = get_catalog_schema(¶ms.db, &query_ctx); + try_update_catalog_schema(&mut query_ctx, &catalog, &schema); + let query_ctx = Arc::new(query_ctx); + + if let Err(err) = handler.check_query_permission(&[], &query_ctx).await { + return PrometheusJsonResponse::error(err.status_code(), err.to_string()); + } + if let Some(metric) = ¶ms.metric + && let Err(err) = handler + .check_query_target_permission( + current_schema_metric_targets(&query_ctx, std::slice::from_ref(metric)), + &query_ctx, + ) + .await + { + return PrometheusJsonResponse::error(err.status_code(), err.to_string()); + } + + let mut metadata = + match retrieve_metric_metadata(&query_ctx, handler.catalog_manager(), ¶ms).await { + Ok(metadata) => metadata, + Err(err) => { + return PrometheusJsonResponse::error( + StatusCode::InvalidArguments, + err.to_string(), + ); + } + }; + let allowed_metric_names = match handler + .filter_metadata_metric_names( + metadata.keys().cloned().collect(), + query_ctx.current_schema().as_str(), + &query_ctx, + ) + .await + { + Ok(metric_names) => metric_names, + Err(err) => { + return PrometheusJsonResponse::error(err.status_code(), err.to_string()); + } + } + .into_iter() + .collect::>(); + metadata.retain(|metric, _| allowed_metric_names.contains(metric)); + if let Some(limit) = params.limit { + metadata = metadata.into_iter().take(limit).collect(); + } + + PrometheusJsonResponse::success(PrometheusResponse::Metadata(metadata)) +} + #[derive(Debug, Default, Serialize, Deserialize)] pub struct InstantQuery { query: Option, @@ -391,7 +474,7 @@ pub async fn instant_query( responses .into_iter() .reduce(|mut acc, resp| { - acc.data.append(resp.data); + acc.append_query_response(resp); acc }) .unwrap() @@ -407,8 +490,10 @@ async fn do_instant_query( query_ctx: QueryContextRef, ) -> PrometheusJsonResponse { let (metric_name, result_type) = retrieve_metric_name_and_result_type(prom_query.expr()); + let query_id = query_ctx.remote_query_id().map(str::to_string); let result = handler.do_query_parsed(prom_query, query_ctx).await; - PrometheusJsonResponse::from_query_result(result, metric_name, result_type).await + PrometheusJsonResponse::from_query_result(result, metric_name, result_type, query_id.as_deref()) + .await } #[derive(Debug, Default, Serialize, Deserialize)] @@ -519,7 +604,7 @@ pub async fn range_query( responses .into_iter() .reduce(|mut acc, resp| { - acc.data.append(resp.data); + acc.append_query_response(resp); acc }) .unwrap() @@ -535,8 +620,15 @@ async fn do_range_query( query_ctx: QueryContextRef, ) -> PrometheusJsonResponse { let (metric_name, _) = retrieve_metric_name_and_result_type(prom_query.expr()); + let query_id = query_ctx.remote_query_id().map(str::to_string); let result = handler.do_query_parsed(prom_query, query_ctx).await; - PrometheusJsonResponse::from_query_result(result, metric_name, ValueType::Matrix).await + PrometheusJsonResponse::from_query_result( + result, + metric_name, + ValueType::Matrix, + query_id.as_deref(), + ) + .await } #[derive(Debug, Default, Serialize)] @@ -1165,15 +1257,15 @@ fn record_batches_to_labels_name( labels: &mut HashSet, ) -> Result<()> { let mut column_indices = Vec::new(); - let mut field_column_indices = Vec::new(); + let mut value_column_indices = Vec::new(); for (i, column) in batches.schema().column_schemas().iter().enumerate() { - if let ConcreteDataType::Float64(_) = column.data_type { - field_column_indices.push(i); + if is_prometheus_value_column(&column.data_type) { + value_column_indices.push(i); } column_indices.push(i); } - if field_column_indices.is_empty() { + if value_column_indices.is_empty() { return Err(Error::Internal { err_msg: "no value column found".to_string(), }); @@ -1185,21 +1277,18 @@ fn record_batches_to_labels_name( .map(|c| batches.schema().column_name_by_index(*c).to_string()) .collect::>(); - let field_columns = field_column_indices + let value_columns = value_column_indices .iter() - .map(|i| { - let column = batch.column(*i); - column.as_primitive::() - }) + .map(|i| batch.column(*i)) .collect::>(); for row_index in 0..batch.num_rows() { - // if all field columns are null, skip this row - if field_columns.iter().all(|c| c.is_null(row_index)) { + // if all value columns are null, skip this row + if value_columns.iter().all(|c| c.is_null(row_index)) { continue; } - // if a field is not null, record the tag name and return + // if a value is not null, record the tag name and return names.iter().for_each(|name| { let _ = labels.insert(name.clone()); }); @@ -1209,6 +1298,10 @@ fn record_batches_to_labels_name( Ok(()) } +fn is_prometheus_value_column(data_type: &ConcreteDataType) -> bool { + matches!(data_type, ConcreteDataType::Float64(_)) || is_native_histogram_value_type(data_type) +} + pub(crate) fn retrieve_metric_name_and_result_type( promql_expr: &PromqlExpr, ) -> (Option, ValueType) { @@ -1869,6 +1962,107 @@ async fn retrieve_table_names( Ok(table_names) } +async fn retrieve_metric_metadata( + query_ctx: &QueryContext, + manager: CatalogManagerRef, + params: &MetadataQuery, +) -> Result>> { + let mut metadata = BTreeMap::new(); + if params.limit == Some(0) { + return Ok(metadata); + } + + let catalog = query_ctx.current_catalog(); + let schema = query_ctx.current_schema(); + + if let Some(metric) = ¶ms.metric { + let Some(table) = manager + .table(catalog, &schema, metric, Some(query_ctx)) + .await? + else { + return Ok(metadata); + }; + let table_info = table.table_info(); + if is_prometheus_metric_table(table_info.as_ref()) { + metadata.insert( + table_info.name.clone(), + vec![prometheus_metadata_from_table(table_info.as_ref())], + ); + } + return Ok(metadata); + } + + let mut tables_stream = manager.tables(catalog, &schema, Some(query_ctx)); + + while let Some(table) = tables_stream.next().await { + let table = table?; + let table_info = table.table_info(); + if !is_prometheus_metric_table(table_info.as_ref()) { + continue; + } + + metadata.insert( + table_info.name.clone(), + vec![prometheus_metadata_from_table(table_info.as_ref())], + ); + } + + Ok(metadata) +} + +fn is_prometheus_metric_table(table_info: &TableInfo) -> bool { + table_info + .meta + .options + .extra_options + .contains_key(LOGICAL_TABLE_METADATA_KEY) +} + +fn prometheus_metadata_from_table(table_info: &TableInfo) -> PromMetadata { + let options = &table_info.meta.options.extra_options; + let metric_type = match options.get(SEMANTIC_METRIC_TYPE) { + Some(metric_type) + if options + .get(SEMANTIC_METRIC_TEMPORALITY) + .is_some_and(|temporality| temporality == "delta") + && matches!( + metric_type.as_str(), + "counter" | "histogram" | "updown_counter" + ) => + { + "unknown".to_string() + } + Some(metric_type) => match metric_type.as_str() { + "updown_counter" => "gauge".to_string(), + "gauge_histogram" => "gaugehistogram".to_string(), + "mixed" => "unknown".to_string(), + metric_type => metric_type.to_string(), + }, + None if table_has_native_histogram_value(table_info) => "histogram".to_string(), + None => String::new(), + }; + let unit = options + .get(SEMANTIC_METRIC_UNIT) + .map(|unit| ucum_to_openmetrics_unit(unit)) + .unwrap_or_default(); + + PromMetadata { + metric_type, + unit, + // TODO: Persist and return Prometheus help text and OTLP metric descriptions. + help: String::new(), + } +} + +fn table_has_native_histogram_value(table_info: &TableInfo) -> bool { + table_info + .meta + .schema + .column_schemas() + .iter() + .any(|column| is_native_histogram_value_type(&column.data_type)) +} + async fn retrieve_field_names( query_ctx: &QueryContext, manager: CatalogManagerRef, @@ -2170,11 +2364,18 @@ mod tests { use catalog::memory::MemoryCatalogManager; use catalog::{RegisterSchemaRequest, RegisterTableRequest}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_query::native_histogram::{ + CounterResetHint, NativeHistogram, build_histogram_array, native_histogram_value_type, + }; + use common_query::prelude::greptime_native_histogram; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; + use datatypes::vectors::{StringVector, StructVector}; use promql_parser::parser::value::ValueType; use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType, TableVersion}; - use table::requests::TableOptions; + use table::requests::{ + SEMANTIC_METRIC_TEMPORALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, TableOptions, + }; use table::test_util::EmptyTable; use table::test_util::table_info::test_table_info; @@ -2191,6 +2392,7 @@ mod tests { struct TestPrometheusHandler { catalog_manager: CatalogManagerRef, + deny_operation: bool, denied_table: Option<&'static str>, metric_names: Vec, queries: Mutex>, @@ -2215,6 +2417,11 @@ mod tests { } async fn check_query_permission(&self, _: &[PromQuery], _: &QueryContextRef) -> Result<()> { + if self.deny_operation { + return auth::error::PermissionDeniedSnafu + .fail() + .context(crate::error::AuthSnafu); + } Ok(()) } @@ -2320,6 +2527,7 @@ mod tests { async fn test_promql_timer_records_parse_errors() { let handler: PrometheusHandlerRef = Arc::new(TestPrometheusHandler { catalog_manager: MemoryCatalogManager::new(), + deny_operation: false, denied_table: None, metric_names: Vec::new(), queries: Mutex::new(Vec::new()), @@ -2402,6 +2610,7 @@ mod tests { let response = label_values_query( State(Arc::new(TestPrometheusHandler { catalog_manager: manager, + deny_operation: false, denied_table: Some("denied"), metric_names: Vec::new(), queries: Mutex::new(Vec::new()), @@ -2444,6 +2653,7 @@ mod tests { let handler = Arc::new(TestPrometheusHandler { catalog_manager: manager, + deny_operation: false, denied_table: None, metric_names: vec!["cpu_user".to_string(), "cpu_system".to_string()], queries: Mutex::new(Vec::new()), @@ -3146,4 +3356,248 @@ mod tests { (PermissionTableTargets::resolved(Vec::new()), 2) ); } + + #[test] + fn test_record_batches_to_labels_name_accepts_native_histogram_value() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + greptime_native_histogram(), + native_histogram_value_type().clone(), + true, + ), + ])); + let batch = RecordBatch::new_empty(schema.clone()); + let batches = RecordBatches::try_new(schema.clone(), vec![batch]).unwrap(); + let mut labels = HashSet::new(); + + record_batches_to_labels_name(batches, &mut labels).unwrap(); + + assert!(labels.is_empty()); + + let histogram = NativeHistogram { + schema: 0, + zero_threshold: 0.001, + sum: 1.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: vec![], + positive_spans: vec![], + negative_spans: vec![], + count: 1.0, + zero_count: 1.0, + positive_buckets: vec![], + negative_buckets: vec![], + }; + let histogram_array = build_histogram_array(&[Some(histogram)]); + let histogram_array = histogram_array + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + let ConcreteDataType::Struct(histogram_type) = native_histogram_value_type().clone() else { + unreachable!("native histogram type must be a struct") + }; + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(StringVector::from(vec![Some("localhost")])) as _, + Arc::new(StructVector::try_new(histogram_type, histogram_array).unwrap()) as _, + ], + ) + .unwrap(); + let batches = RecordBatches::try_new(schema, vec![batch]).unwrap(); + + record_batches_to_labels_name(batches, &mut labels).unwrap(); + + assert_eq!( + labels, + HashSet::from(["host".to_string(), greptime_native_histogram().to_string(),]) + ); + } + + fn prometheus_metric_table_info(table_id: u32, name: &str) -> TableInfo { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ])); + let mut options = TableOptions::default(); + options.extra_options.insert( + LOGICAL_TABLE_METADATA_KEY.to_string(), + "greptime_physical_table".to_string(), + ); + options + .extra_options + .insert(SEMANTIC_METRIC_TYPE.to_string(), "counter".to_string()); + options + .extra_options + .insert(SEMANTIC_METRIC_UNIT.to_string(), "By".to_string()); + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![1]) + .engine("metric".to_string()) + .next_column_id(3) + .options(options) + .build() + .unwrap(); + + TableInfoBuilder::default() + .table_id(table_id) + .table_version(0 as TableVersion) + .name(name) + .catalog_name(DEFAULT_CATALOG_NAME) + .schema_name(DEFAULT_SCHEMA_NAME) + .table_type(TableType::Base) + .meta(meta) + .build() + .unwrap() + } + + #[tokio::test] + async fn test_retrieve_metric_metadata_uses_semantic_options() { + let table_info = prometheus_metric_table_info(1025, "http_requests_total"); + let manager: CatalogManagerRef = + MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&table_info)); + let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); + + let metadata = retrieve_metric_metadata(&query_ctx, manager, &MetadataQuery::default()) + .await + .unwrap(); + + assert_eq!( + metadata.get("http_requests_total"), + Some(&vec![PromMetadata { + metric_type: "counter".to_string(), + unit: "bytes".to_string(), + help: String::new(), + }]) + ); + + for (ucum, openmetrics) in [ + ("1", "ratios"), + ("ms", "milliseconds"), + ("m/s", "meters_per_second"), + ("USD", "USD"), + ] { + let mut table_info = table_info.clone(); + table_info + .meta + .options + .extra_options + .insert(SEMANTIC_METRIC_UNIT.to_string(), ucum.to_string()); + assert_eq!( + prometheus_metadata_from_table(&table_info).unit, + openmetrics + ); + } + + for (metric_type, temporality, expected) in [ + ("updown_counter", None, "gauge"), + ("gauge_histogram", None, "gaugehistogram"), + ("summary", None, "summary"), + ("mixed", None, "unknown"), + ("counter", Some("delta"), "unknown"), + ("histogram", Some("delta"), "unknown"), + ] { + let mut table_info = table_info.clone(); + table_info + .meta + .options + .extra_options + .insert(SEMANTIC_METRIC_TYPE.to_string(), metric_type.to_string()); + if let Some(temporality) = temporality { + table_info.meta.options.extra_options.insert( + SEMANTIC_METRIC_TEMPORALITY.to_string(), + temporality.to_string(), + ); + } + assert_eq!( + prometheus_metadata_from_table(&table_info).metric_type, + expected + ); + } + } + + #[tokio::test] + async fn test_metadata_query_checks_permissions_and_filters_tables() { + let allowed = prometheus_metric_table_info(1025, "allowed"); + let denied = prometheus_metric_table_info(1026, "denied"); + let manager = MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&allowed)); + manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: denied.name.clone(), + table_id: denied.table_id(), + table: EmptyTable::from_table_info(&denied), + }) + .unwrap(); + let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); + + let metadata = retrieve_metric_metadata( + &query_ctx, + manager.clone(), + &MetadataQuery { + metric: Some("allowed".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(metadata.len(), 1); + assert!(metadata.contains_key("allowed")); + + let response = metadata_query( + State(Arc::new(TestPrometheusHandler { + catalog_manager: manager.clone(), + deny_operation: true, + denied_table: None, + metric_names: Vec::new(), + queries: Mutex::new(Vec::new()), + })), + Query(MetadataQuery::default()), + Extension(query_ctx.clone()), + ) + .await; + assert_eq!(Some(StatusCode::PermissionDenied), response.status_code); + + let handler: PrometheusHandlerRef = Arc::new(TestPrometheusHandler { + catalog_manager: manager, + deny_operation: false, + denied_table: Some("denied"), + metric_names: Vec::new(), + queries: Mutex::new(Vec::new()), + }); + let response = metadata_query( + State(handler.clone()), + Query(MetadataQuery::default()), + Extension(query_ctx.clone()), + ) + .await; + assert!(response.status_code.is_none()); + let PrometheusResponse::Metadata(metadata) = response.data else { + panic!("expected metadata response"); + }; + assert_eq!( + metadata.keys().cloned().collect::>(), + vec!["allowed".to_string()] + ); + + let response = metadata_query( + State(handler), + Query(MetadataQuery { + metric: Some("denied".to_string()), + ..Default::default() + }), + Extension(query_ctx), + ) + .await; + assert_eq!(Some(StatusCode::PermissionDenied), response.status_code); + } } diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index 368d925c30c..da32103b79e 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -28,6 +28,9 @@ use common_query::native_histogram::{ NativeHistogram, is_native_histogram_value_type, read_histogram, }; use common_query::prometheus::{format_prometheus_float, is_prometheus_stale_nan}; +use common_query::promql_annotations::{ + PromqlAnnotationCollector, get_promql_annotation_collector, +}; use common_query::{Output, OutputData}; use common_recordbatch::RecordBatches; use datatypes::arrow_array::string_array_value_at_index; @@ -80,6 +83,8 @@ pub struct PrometheusJsonResponse { pub error_type: Option, #[serde(skip_serializing_if = "Option::is_none")] pub warnings: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub infos: Option>, #[serde(skip)] pub status_code: Option, @@ -124,6 +129,7 @@ impl PrometheusJsonResponse { error: Some(reason.into()), error_type: Some(error_type.to_string()), warnings: None, + infos: None, resp_metrics: Default::default(), status_code: Some(error_type), } @@ -136,17 +142,37 @@ impl PrometheusJsonResponse { error: None, error_type: None, warnings: None, + infos: None, resp_metrics: Default::default(), status_code: None, } } + /// Adds collected PromQL warnings and infos to the response. + fn append_promql_annotations(&mut self, collector: &PromqlAnnotationCollector) { + let mut warnings = self.warnings.take().unwrap_or_default(); + let mut infos = self.infos.take().unwrap_or_default(); + collector.append_to(&mut warnings, &mut infos); + self.warnings = (!warnings.is_empty()).then_some(warnings); + self.infos = (!infos.is_empty()).then_some(infos); + } + + /// Merges data and annotations from another expanded PromQL query response. + pub(crate) fn append_query_response(&mut self, mut other: Self) { + self.data.append(other.data); + merge_annotations(&mut self.warnings, other.warnings.take()); + merge_annotations(&mut self.infos, other.infos.take()); + } + /// Convert from `Result` pub async fn from_query_result( result: Result, metric_name: Option, result_type: ValueType, + query_id: Option<&str>, ) -> Self { + // Hold the collector while a streaming result is consumed. + let collector = query_id.and_then(get_promql_annotation_collector); let response: Result = try { let result = result?; let mut resp = @@ -187,7 +213,7 @@ impl PrometheusJsonResponse { let result_type_string = result_type.to_string(); - match response { + let mut response = match response { Ok(resp) => resp, Err(err) => { // Prometheus won't report error if querying nonexist label and metric @@ -202,7 +228,11 @@ impl PrometheusJsonResponse { Self::error(err.status_code(), err.output_msg()) } } + }; + if let Some(collector) = collector { + response.append_promql_annotations(&collector); } + response } /// Convert [RecordBatches] to [PromData] @@ -435,6 +465,16 @@ impl PrometheusJsonResponse { } } +fn merge_annotations(target: &mut Option>, source: Option>) { + let Some(source) = source else { + return; + }; + let target = target.get_or_insert_default(); + target.extend(source); + target.sort(); + target.dedup(); +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -444,6 +484,9 @@ mod tests { native_histogram_value_type, }; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; + use common_query::promql_annotations::{ + PromqlAnnotationCollector, promql_annotation_collector, + }; use common_recordbatch::{RecordBatch, RecordBatches}; use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; @@ -453,6 +496,40 @@ mod tests { use super::*; + #[tokio::test] + async fn query_response_preserves_and_merges_promql_annotations() { + let query_id = "query_response_preserves_and_merges_promql_annotations"; + let left = promql_annotation_collector(query_id); + left.record_warning("shared warning"); + left.record_info("left info"); + let mut response = PrometheusJsonResponse::from_query_result( + Ok(Output::new_with_record_batches(RecordBatches::empty())), + None, + ValueType::Vector, + Some(query_id), + ) + .await; + + let right = PromqlAnnotationCollector::default(); + right.record_warning("shared warning"); + right.record_info("right info"); + let mut other = PrometheusJsonResponse::success(PrometheusResponse::None); + other.append_promql_annotations(&right); + response.append_query_response(other); + + assert_eq!(response.warnings, Some(vec!["shared warning".to_string()])); + assert_eq!( + response.infos, + Some(vec!["left info".to_string(), "right info".to_string()]) + ); + let json = serde_json::to_value(response).unwrap(); + assert_eq!(json["warnings"], serde_json::json!(["shared warning"])); + assert_eq!( + json["infos"], + serde_json::json!(["left info", "right info"]) + ); + } + fn sample_histogram() -> NativeHistogram { NativeHistogram { schema: 0, diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index b43cc73338f..6d94b3a139f 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -38,6 +38,7 @@ use crate::semantic::{ mod translator; pub use translator::legacy_normalize_otlp_name; +pub(crate) use translator::ucum_to_openmetrics_unit; use translator::{translate_label_name, translate_metric_name}; /// the default column count for table writer diff --git a/src/servers/src/otlp/metrics/translator.rs b/src/servers/src/otlp/metrics/translator.rs index bb929cd6705..8b3b45d3274 100644 --- a/src/servers/src/otlp/metrics/translator.rs +++ b/src/servers/src/otlp/metrics/translator.rs @@ -219,6 +219,20 @@ fn build_clean_unit_suffix(unit: &str) -> (Option, Option) { ) } +/// Converts a UCUM unit to its Prometheus/OpenMetrics unit word. +pub(crate) fn ucum_to_openmetrics_unit(unit: &str) -> String { + if unit == "1" { + return "ratios".to_string(); + } + + match build_clean_unit_suffix(unit) { + (Some(main), Some(per)) => format!("{main}_per_{per}"), + (Some(main), None) => main, + (None, Some(per)) => format!("per_{per}"), + (None, None) => String::new(), + } +} + fn build_unit_suffixes(unit: &str) -> (String, String) { let (main, per) = unit.split_once('/').unwrap_or((unit, "")); let main_unit_suffix = unit_suffix(main, &UNIT_MAP); diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 792a2ce965d..23862870dba 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -980,6 +980,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { error: None, error_type: None, warnings: None, + infos: None, resp_metrics: Default::default(), status_code: None, }; diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index cd8af8768a4..042be385a6d 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2741,6 +2741,28 @@ pub async fn test_prometheus_remote_write_v2(store_type: StorageType) { ) .await; + let res = client + .get("/v1/prometheus/api/v1/metadata?metric=remote_write_v2_typed_total") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + serde_json::to_string_pretty(&body).unwrap(), + r#"{ + "status": "success", + "data": { + "remote_write_v2_typed_total": [ + { + "type": "counter", + "unit": "seconds", + "help": "" + } + ] + } +}"# + ); + guard.remove_all().await; } @@ -2882,6 +2904,202 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag ) .await; + let res = client + .get("/v1/prometheus/api/v1/metadata?metric=remote_write_v2_latency_seconds") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + serde_json::to_string_pretty(&body).unwrap(), + r#"{ + "status": "success", + "data": { + "remote_write_v2_latency_seconds": [ + { + "type": "histogram", + "unit": "", + "help": "" + } + ] + } +}"# + ); + + let res = client + .get("/v1/prometheus/api/v1/query?query=histogram_count(remote_write_v2_latency_seconds)&time=4") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + serde_json::to_string_pretty(&body).unwrap(), + r#"{ + "status": "success", + "data": { + "resultType": "vector", + "result": [ + { + "metric": { + "__name__": "remote_write_v2_latency_seconds", + "instance": "localhost:9090", + "job": "api" + }, + "value": [ + 4.0, + "6" + ] + } + ] + } +}"# + ); + + let res = client + .get("/v1/prometheus/api/v1/query?query=remote_write_v2_latency_seconds&time=4") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + serde_json::to_string_pretty(&body).unwrap(), + r#"{ + "status": "success", + "data": { + "resultType": "vector", + "result": [ + { + "metric": { + "__name__": "remote_write_v2_latency_seconds", + "instance": "localhost:9090", + "job": "api" + }, + "histogram": [ + 4.0, + { + "count": "6", + "sum": "20", + "buckets": [ + [ + 3, + "-0.002", + "0.002", + "0.5" + ], + [ + 0, + "1.4142135623730951", + "1.681792830507429", + "2" + ], + [ + 0, + "1.681792830507429", + "2", + "3.5" + ] + ] + } + ] + } + ] + } +}"# + ); + + let res = client + .get("/v1/prometheus/api/v1/query_range?query=remote_write_v2_latency_seconds&start=3&end=4&step=1") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + serde_json::to_string_pretty(&body).unwrap(), + r#"{ + "status": "success", + "data": { + "resultType": "matrix", + "result": [ + { + "metric": { + "__name__": "remote_write_v2_latency_seconds", + "instance": "localhost:9090", + "job": "api" + }, + "histograms": [ + [ + 3.0, + { + "count": "8", + "sum": "10", + "buckets": [ + [ + 1, + "-0.5", + "-0.3535533905932738", + "1" + ], + [ + 3, + "-0.001", + "0.001", + "1" + ], + [ + 0, + "0.7071067811865476", + "1", + "1" + ], + [ + 0, + "1", + "1.4142135623730951", + "3" + ], + [ + 0, + "1.4142135623730951", + "2", + "2" + ] + ] + } + ], + [ + 4.0, + { + "count": "6", + "sum": "20", + "buckets": [ + [ + 3, + "-0.002", + "0.002", + "0.5" + ], + [ + 0, + "1.4142135623730951", + "1.681792830507429", + "2" + ], + [ + 0, + "1.681792830507429", + "2", + "3.5" + ] + ] + } + ] + ] + } + ] + } +}"# + ); + validate_data( "prometheus_remote_write_v2_native_histogram_table_created", &client, From 9241e96fa2ea8ed3132e38dea93c655337d24040 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 12 Aug 2026 07:49:46 +0000 Subject: [PATCH 055/126] fix(metric-engine): handle Utf8View tag/label columns without panicking (#8772) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(metric-engine): handle Utf8View tag/label columns without panicking label_replace (planned as DataFusion regexp_replace) coerces to Utf8View, so label columns materialize as StringViewArray; build_tag_arrays' StringArray downcast then panicked ('tag column must be utf8') — e.g. for OTLP/json2 ingest. TSID computation, sparse-PK encoding and tag extraction now accept generic ArrayRef tag columns (Utf8/LargeUtf8/ Utf8View/Dictionary) via string_array_value_at_index, and build_tag_arrays errors instead of panicking on non-string columns. The mito2 time-series memtable string-field paths are hardened the same way. Adds label_replace_with_utf8view_labels_does_not_panic (issue #8732). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(metric-engine): add is_string_null_at helper for tag null checks Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(datatypes): use is_none_or to satisfy clippy unnecessary-map-or Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix: reject oversized string batches before memtable append Distinguish a full active string builder from a batch that cannot fit an empty Arrow string builder at all. Scan every string field so a later intrinsically oversized field cannot be skipped after an earlier field requests a freeze. Return InvalidBatch instead of reaching Arrow's offset overflow panic. Also cover Utf8View tags with nulls through the metric-engine tag, TSID, and sparse-primary-key path. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datatypes/src/arrow_array.rs | 37 ++++ src/metric-engine/src/batch_modifier.rs | 207 ++++++++++++++++-- src/mito2/src/memtable/time_series.rs | 202 ++++++++++++++--- .../src/http/result/prometheus_resp.rs | 52 +++++ 4 files changed, 447 insertions(+), 51 deletions(-) diff --git a/src/datatypes/src/arrow_array.rs b/src/datatypes/src/arrow_array.rs index e5ef7add827..848cc7c91f7 100644 --- a/src/datatypes/src/arrow_array.rs +++ b/src/datatypes/src/arrow_array.rs @@ -168,6 +168,43 @@ pub fn string_array_value_at_index(array: &ArrayRef, i: usize) -> Option<&str> { } } +/// Check whether the string value at index `i` is null for string or +/// dictionary-encoded string arrays. +/// +/// Returns `true` when the value is null or the array type is not a string +/// type, which corresponds to [`string_array_value_at_index`] returning `None`. +/// +/// # Panics +/// +/// If index `i` is out of bounds. +pub fn is_string_null_at(array: &ArrayRef, i: usize) -> bool { + match array.data_type() { + DataType::Utf8 => { + let array = array.as_string::(); + !array.is_valid(i) + } + DataType::LargeUtf8 => { + let array = array.as_string::(); + !array.is_valid(i) + } + DataType::Utf8View => { + let array = array.as_string_view(); + !array.is_valid(i) + } + DataType::Dictionary(key_type, value_type) + if key_type.is_integer() && value_type.is_string() => + { + downcast_dictionary_array! { + array => array + .key(i) + .is_none_or(|key| is_string_null_at(array.values(), key)), + _ => true, + } + } + _ => true, + } +} + /// Get the string value at index `i` for `Utf8`, `LargeUtf8`, or `Utf8View` arrays. /// /// Note: This method does not check for nulls and the value is arbitrary diff --git a/src/metric-engine/src/batch_modifier.rs b/src/metric-engine/src/batch_modifier.rs index 2162215d1b1..575dffff1a3 100644 --- a/src/metric-engine/src/batch_modifier.rs +++ b/src/metric-engine/src/batch_modifier.rs @@ -15,9 +15,10 @@ use std::hash::Hasher; use std::sync::Arc; -use datatypes::arrow::array::{Array, BinaryBuilder, StringArray, UInt64Array}; +use datatypes::arrow::array::{Array, ArrayRef, BinaryBuilder, UInt64Array}; use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::arrow_array::{is_string_null_at, string_array_value_at_index}; use fxhash::FxHasher; use mito_codec::row_converter::SparsePrimaryKeyCodec; use snafu::ResultExt; @@ -49,7 +50,7 @@ pub struct TagColumnInfo { pub fn compute_tsid_array( batch: &RecordBatch, sorted_tag_columns: &[TagColumnInfo], - tag_arrays: &[&StringArray], + tag_arrays: &[&ArrayRef], ) -> UInt64Array { let num_rows = batch.num_rows(); @@ -64,20 +65,22 @@ pub fn compute_tsid_array( let mut tsid_values = Vec::with_capacity(num_rows); for row in 0..num_rows { - let has_null = tag_arrays.iter().any(|arr| arr.is_null(row)); + let has_null = tag_arrays.iter().any(|arr| is_string_null_at(arr, row)); let tsid = if !has_null { let mut hasher = FxHasher::default(); hasher.write_u64(label_name_hash); for arr in tag_arrays { - hasher.write(arr.value(row).as_bytes()); - hasher.write_u8(0xff); + if let Some(value) = string_array_value_at_index(arr, row) { + hasher.write(value.as_bytes()); + hasher.write_u8(0xff); + } } hasher.finish() } else { let mut name_hasher = FxHasher::default(); for (tc, arr) in sorted_tag_columns.iter().zip(tag_arrays.iter()) { - if !arr.is_null(row) { + if !is_string_null_at(arr, row) { name_hasher.write(tc.name.as_bytes()); name_hasher.write_u8(0xff); } @@ -87,8 +90,8 @@ pub fn compute_tsid_array( let mut val_hasher = FxHasher::default(); val_hasher.write_u64(row_label_hash); for arr in tag_arrays { - if !arr.is_null(row) { - val_hasher.write(arr.value(row).as_bytes()); + if let Some(value) = string_array_value_at_index(arr, row) { + val_hasher.write(value.as_bytes()); val_hasher.write_u8(0xff); } } @@ -104,15 +107,26 @@ pub fn compute_tsid_array( fn build_tag_arrays<'a>( batch: &'a RecordBatch, sorted_tag_columns: &[TagColumnInfo], -) -> Vec<&'a StringArray> { +) -> Result> { sorted_tag_columns .iter() .map(|tc| { - batch - .column(tc.index) - .as_any() - .downcast_ref::() - .expect("tag column must be utf8") + let array = batch.column(tc.index); + match array.data_type() { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Ok(array), + DataType::Dictionary(key_type, value_type) + if key_type.is_integer() && value_type.is_string() => + { + Ok(array) + } + data_type => UnexpectedRequestSnafu { + reason: format!( + "Tag column '{}' must be a string, given: {data_type}", + tc.name + ), + } + .fail(), + } }) .collect() } @@ -144,7 +158,7 @@ pub fn modify_batch_sparse( ) -> Result { let num_rows = batch.num_rows(); let codec = SparsePrimaryKeyCodec::schemaless(); - let tag_arrays: Vec<&StringArray> = build_tag_arrays(&batch, sorted_tag_columns); + let tag_arrays = build_tag_arrays(&batch, sorted_tag_columns)?; let tsid_array = compute_tsid_array(&batch, sorted_tag_columns, &tag_arrays); let mut pk_builder = BinaryBuilder::with_capacity(num_rows, 0); @@ -158,8 +172,9 @@ pub fn modify_batch_sparse( let tags = sorted_tag_columns .iter() .zip(tag_arrays.iter()) - .filter(|(_, arr)| !arr.is_null(row)) - .map(|(tc, arr)| (tc.column_id, arr.value(row).as_bytes())); + .filter_map(|(tc, arr)| { + string_array_value_at_index(arr, row).map(|value| (tc.column_id, value.as_bytes())) + }); codec .encode_raw_tag_value(tags, &mut buffer) .context(EncodePrimaryKeySnafu)?; @@ -268,7 +283,7 @@ mod tests { column_id: 1, }, ]; - let tag_arrays = build_tag_arrays(&batch, &tag_columns); + let tag_arrays = build_tag_arrays(&batch, &tag_columns).unwrap(); let tsid_array = compute_tsid_array(&batch, &tag_columns, &tag_arrays); assert_eq!(tsid_array.value(0), 2721566936019240841); @@ -300,7 +315,7 @@ mod tests { column_id: 2, }, ]; - let tag_arrays_2 = build_tag_arrays(&batch_no_null, &tag_cols_2); + let tag_arrays_2 = build_tag_arrays(&batch_no_null, &tag_cols_2).unwrap(); let tsid_no_null = compute_tsid_array(&batch_no_null, &tag_cols_2, &tag_arrays_2); let schema3 = Arc::new(ArrowSchema::new(vec![ @@ -334,7 +349,7 @@ mod tests { column_id: 3, }, ]; - let tag_arrays_3 = build_tag_arrays(&batch_with_null, &tag_cols_3); + let tag_arrays_3 = build_tag_arrays(&batch_with_null, &tag_cols_3).unwrap(); let tsid_with_null = compute_tsid_array(&batch_with_null, &tag_cols_3, &tag_arrays_3); assert_eq!(tsid_no_null.value(0), tsid_with_null.value(0)); @@ -452,4 +467,156 @@ mod tests { .unwrap(); assert_eq!(actual_array.value(0), expected_pk.as_slice()); } + + #[test] + fn label_replace_with_utf8view_labels_does_not_panic() { + // Tag (label) columns may arrive as any string representation (`Utf8`, + // `LargeUtf8`, `Utf8View`, or dictionary-encoded). `modify_batch_sparse` + // must not assume the tag columns are plain `StringArray`s. + let tag_arrays: Vec = vec![ + Arc::new(StringArray::from(vec!["greptimedb"])), + Arc::new(datatypes::arrow::array::LargeStringArray::from(vec![ + "greptimedb", + ])), + Arc::new(datatypes::arrow::array::StringViewArray::from(vec![ + "greptimedb", + ])), + Arc::new(datatypes::arrow::array::DictionaryArray::< + datatypes::arrow::datatypes::UInt32Type, + >::new( + datatypes::arrow::array::UInt32Array::from(vec![0]), + Arc::new(StringArray::from(vec!["greptimedb"])), + )), + ]; + let tag_columns = vec![TagColumnInfo { + name: "namespace".to_string(), + index: 2, + column_id: 2, + }]; + let non_tag_indices = vec![0, 1]; + + let primary_keys = tag_arrays + .into_iter() + .map(|tag_array| { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("greptime_timestamp", DataType::Int64, false), + Field::new("greptime_value", DataType::Float64, true), + Field::new("namespace", tag_array.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1000])), + Arc::new(datatypes::arrow::array::Float64Array::from(vec![42.0])), + tag_array, + ], + ) + .unwrap(); + let modified = + modify_batch_sparse(batch, 1025, &tag_columns, &non_tag_indices).unwrap(); + modified + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + .to_vec() + }) + .collect::>(); + + assert!(primary_keys.windows(2).all(|keys| keys[0] == keys[1])); + } + + #[test] + fn test_compute_tsid_utf8view_with_nulls_matches_utf8() { + // Regression test for Utf8View tag/label columns: the whole + // `build_tag_arrays` -> `compute_tsid_array` -> sparse primary key path must handle + // `Utf8View` columns exactly like `Utf8`, including rows with null tags, instead of + // panicking or producing different TSIDs. + let namespace_values: Vec> = + vec![Some("ns-a"), Some("ns-b"), None, Some("ns-c")]; + let host_values: Vec> = vec![Some("host-1"), None, Some("host-2"), None]; + + let build_batch = |data_type: DataType| { + let (namespace, host): (ArrayRef, ArrayRef) = match &data_type { + DataType::Utf8 => ( + Arc::new(StringArray::from(namespace_values.clone())), + Arc::new(StringArray::from(host_values.clone())), + ), + DataType::Utf8View => ( + Arc::new(datatypes::arrow::array::StringViewArray::from( + namespace_values.clone(), + )), + Arc::new(datatypes::arrow::array::StringViewArray::from( + host_values.clone(), + )), + ), + _ => unreachable!(), + }; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("greptime_timestamp", DataType::Int64, false), + Field::new("greptime_value", DataType::Float64, true), + Field::new("namespace", data_type.clone(), true), + Field::new("host", data_type, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1000, 1001, 1002, 1003])), + Arc::new(datatypes::arrow::array::Float64Array::from(vec![ + 1.0, 2.0, 3.0, 4.0, + ])), + namespace, + host, + ], + ) + .unwrap() + }; + + let tag_columns = vec![ + TagColumnInfo { + name: "host".to_string(), + index: 3, + column_id: 3, + }, + TagColumnInfo { + name: "namespace".to_string(), + index: 2, + column_id: 2, + }, + ]; + + let utf8_batch = build_batch(DataType::Utf8); + let utf8view_batch = build_batch(DataType::Utf8View); + + // Both representations must be accepted by `build_tag_arrays`. + let utf8_tag_arrays = build_tag_arrays(&utf8_batch, &tag_columns).unwrap(); + let utf8view_tag_arrays = build_tag_arrays(&utf8view_batch, &tag_columns).unwrap(); + + // TSIDs must be identical across representations, including rows with null tags. + let utf8_tsids = compute_tsid_array(&utf8_batch, &tag_columns, &utf8_tag_arrays); + let utf8view_tsids = + compute_tsid_array(&utf8view_batch, &tag_columns, &utf8view_tag_arrays); + assert_eq!(utf8_tsids, utf8view_tsids); + + // The full sparse primary key (TSID + tag values) must also be identical. + let non_tag_indices = vec![0, 1]; + let modified_utf8 = + modify_batch_sparse(utf8_batch, 1025, &tag_columns, &non_tag_indices).unwrap(); + let modified_utf8view = + modify_batch_sparse(utf8view_batch, 1025, &tag_columns, &non_tag_indices).unwrap(); + let utf8_pks = modified_utf8 + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let utf8view_pks = modified_utf8view + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..utf8_pks.len() { + assert_eq!(utf8_pks.value(row), utf8view_pks.value(row)); + } + } } diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index f5b772b076f..7e89d569eeb 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/src/mito2/src/memtable/time_series.rs @@ -32,8 +32,8 @@ use datatypes::prelude::{ScalarVector, Vector, VectorRef}; use datatypes::types::TimestampType; use datatypes::value::{Value, ValueRef}; use datatypes::vectors::{ - Helper, TimestampMicrosecondVector, TimestampMillisecondVector, TimestampNanosecondVector, - TimestampSecondVector, UInt8Vector, UInt64Vector, + Helper, StringVector, TimestampMicrosecondVector, TimestampMillisecondVector, + TimestampNanosecondVector, TimestampSecondVector, UInt8Vector, UInt64Vector, }; use mito_codec::key_values::KeyValue; use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt}; @@ -72,6 +72,93 @@ const INITIAL_BUILDER_CAPACITY: usize = 4; /// Vector builder capacity. const BUILDER_CAPACITY: usize = 512; +fn checked_string_values_size(mut lengths: impl Iterator) -> Option { + lengths.try_fold(0_i32, |size, len| { + size.checked_add(i32::try_from(len).ok()?) + }) +} + +/// Checks whether a string field builder that currently holds `current_len` bytes can +/// additionally accommodate `space_needed` bytes of string data, given that the offset of a +/// single Arrow string array is limited to `limit` (`i32::MAX` in production). +/// +/// Returns: +/// - `Ok(true)` if `current_len + space_needed <= limit`; +/// - `Ok(false)` if `space_needed <= limit` but `current_len + space_needed > limit`, i.e. +/// the current builder is too full but an empty builder could accommodate the batch (the +/// caller may freeze the current builder and replay the batch onto an empty one); +/// - `Err(InvalidBatchSnafu)` if the batch's string data alone (`space_needed`, or the fact +/// that its total bytes cannot even be represented as `i32`) exceeds `limit`, i.e. the +/// batch can never be accommodated by any builder. +fn check_string_capacity(current_len: i32, space_needed: Option, limit: i32) -> Result { + let Some(space_needed) = space_needed.filter(|&space| space <= limit) else { + return error::InvalidBatchSnafu { + reason: format!( + "String data of the batch exceeds the Arrow string array offset limit ({limit}) and cannot be stored in a single column" + ), + } + .fail(); + }; + Ok(current_len + .checked_add(space_needed) + .is_some_and(|v| v <= limit)) +} + +/// Scans all string fields of a batch and checks whether they can be accommodated. +/// +/// Every string field is checked (no early return on `Ok(false)`) so that an intrinsic +/// oversize in *any* field surfaces as an error even when an earlier field only overflows +/// the current builder. +/// +/// Returns: +/// - `Ok(true)` if every string field fits into the current builders; +/// - `Ok(false)` if no field is intrinsically oversized but at least one field only fits +/// into an empty builder (the caller may freeze the current builder and replay the batch +/// onto an empty one); +/// - `Err(InvalidBatchSnafu)` if any single field's string data alone (`space_needed`, or +/// the fact that its total bytes cannot even be represented as `i32`) exceeds `limit`, +/// i.e. that field can never be accommodated by any builder. +fn scan_string_capacity( + fields: &[VectorRef], + field_builders: &[Option], + field_types: &[ConcreteDataType], + limit: i32, +) -> Result { + let mut can_fit_current = true; + for ((field_src, field_dest), field_type) in fields + .iter() + .zip(field_builders.iter()) + .zip(field_types.iter()) + { + if !matches!(field_type, ConcreteDataType::String(_)) { + continue; + } + let current_size = match field_dest { + Some(FieldBuilder::String(builder)) => builder.next_offset(), + None => 0, + Some(FieldBuilder::Other(_)) => unreachable!(), + }; + let array = field_src.to_arrow_array(); + let space_needed = if let Some(string_array) = array.as_any().downcast_ref::() + { + i32::try_from(string_array.value_data().len()).ok() + } else { + let string_vector = field_src + .as_any() + .downcast_ref::() + .with_context(|| error::InvalidBatchSnafu { + reason: format!( + "Field type mismatch, expecting String, given: {}", + field_src.data_type() + ), + })?; + checked_string_values_size(string_vector.iter_data().flatten().map(str::len)) + }; + can_fit_current &= check_string_capacity(current_size, space_needed, limit)?; + } + Ok(can_fit_current) +} + /// Builder to build [TimeSeriesMemtable]. #[derive(Debug, Default)] pub struct TimeSeriesMemtableBuilder { @@ -920,32 +1007,17 @@ impl ValueBuilder { } /// Checks if current value builder have sufficient space to accommodate `fields`. - /// Returns false if there is no space to accommodate fields due to offset overflow. + /// + /// Returns `Ok(false)` if the current builder lacks the remaining space to accommodate + /// the fields due to offset overflow, but an empty builder would be able to accommodate + /// the batch (the caller may freeze the current builder and replay the batch onto an + /// empty one). + /// + /// Returns `Err(InvalidBatchSnafu)` if the string data of a single batch itself exceeds + /// the Arrow string array offset limit and thus can never be accommodated, not even by an + /// empty builder. pub(crate) fn can_accommodate(&self, fields: &[VectorRef]) -> Result { - for (field_src, field_dest) in fields.iter().zip(self.fields.iter()) { - let Some(builder) = field_dest else { - continue; - }; - let FieldBuilder::String(builder) = builder else { - continue; - }; - let array = field_src.to_arrow_array(); - let string_array = array - .as_any() - .downcast_ref::() - .with_context(|| error::InvalidBatchSnafu { - reason: format!( - "Field type mismatch, expecting String, given: {}", - field_src.data_type() - ), - })?; - let space_needed = string_array.value_data().len() as i32; - // offset may overflow - if builder.next_offset().checked_add(space_needed).is_none() { - return Ok(false); - } - } - Ok(true) + scan_string_capacity(fields, &self.fields, &self.field_types, i32::MAX) } pub(crate) fn extend( @@ -1017,17 +1089,26 @@ impl ValueBuilder { match builder { FieldBuilder::String(builder) => { let array = field_src.to_arrow_array(); - let string_array = - array + if let Some(string_array) = array.as_any().downcast_ref::() { + builder.append_array(string_array); + } else { + let string_vector = field_src .as_any() - .downcast_ref::() + .downcast_ref::() .with_context(|| error::InvalidBatchSnafu { reason: format!( "Field type mismatch, expecting String, given: {}", field_src.data_type() ), })?; - builder.append_array(string_array); + for value in string_vector.iter_data() { + if let Some(value) = value { + builder.append(value); + } else { + builder.append_null(); + } + } + } } FieldBuilder::Other(builder) => { let len = field_src.len(); @@ -2296,4 +2377,63 @@ mod tests { assert!(iter.next().is_none()); } + + #[test] + fn test_can_accommodate_string_offset_overflow() { + // A batch whose string data alone exceeds the Arrow string array offset limit can + // never be accommodated, not even by an empty builder. `can_accommodate` must return + // a structured `InvalidBatch` error instead of `Ok(false)`: the latter would make + // `Series::extend` freeze the current builder and replay the same oversized batch + // onto an empty builder, which then panics on offset overflow in `StringBuilder`. + assert!(check_string_capacity(0, Some(101), 100).is_err()); + assert!(check_string_capacity(0, None, i32::MAX).is_err()); + + // A batch that fits in an empty builder but not in the current one reports + // `Ok(false)` so the caller can freeze and replay onto an empty builder. + assert!(!check_string_capacity(95, Some(10), 100).unwrap()); + assert!(!check_string_capacity(91, Some(10), 100).unwrap()); + + // The current builder can accommodate the batch when the combined size fits, + // including the case where it exactly reaches the limit. + assert!(check_string_capacity(0, Some(100), 100).unwrap()); + assert!(check_string_capacity(90, Some(10), 100).unwrap()); + } + + #[test] + fn test_can_accommodate_checks_all_string_fields() { + // Two string fields: the first only overflows the *current* builder (an empty + // builder would fit it), the second is intrinsically oversized. The scan must not + // early-return `Ok(false)` on the first field: it has to check every string field so + // the intrinsic oversize of the second field still surfaces as `InvalidBatch`. + let limit = 10; + // 8 bytes already in the current builder: 8 + 4 > 10, but 4 <= 10, so this field + // alone fits an empty builder. + let mut current = StringBuilder::with_capacity(1, 8); + current.append("12345678"); + let field_builders = [Some(FieldBuilder::String(current)), None]; + let field_types = [ + ConcreteDataType::string_datatype(), + ConcreteDataType::string_datatype(), + ]; + + let first = Arc::new(StringVector::from(StringArray::from(vec!["abcd"]))) as VectorRef; + // 15 bytes > limit: intrinsically oversized, can never fit any builder. + let second = Arc::new(StringVector::from(StringArray::from(vec![ + "123456789012345", + ]))) as VectorRef; + + let err = scan_string_capacity(&[first, second], &field_builders, &field_types, limit) + .unwrap_err(); + assert!( + matches!(err, error::Error::InvalidBatch { .. }), + "expected InvalidBatch, got {err:?}" + ); + + // Sanity: without the oversized field the result is just `Ok(false)` (current full, + // empty builder fits), which lets the caller freeze and replay onto an empty builder. + let first_only = Arc::new(StringVector::from(StringArray::from(vec!["abcd"]))) as VectorRef; + assert!( + !scan_string_capacity(&[first_only], &field_builders, &field_types, limit).unwrap() + ); + } } diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index da32103b79e..e79628fac0d 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -479,6 +479,7 @@ fn merge_annotations(target: &mut Option>, source: Option Date: Wed, 12 Aug 2026 07:50:54 +0000 Subject: [PATCH 056/126] feat(promql): support native histogram aggregations (#8848) * feat(promql): support native histogram aggregations Signed-off-by: shuiyisong * fix(promql): correct mixed native histogram aggregations Signed-off-by: shuiyisong * fix(promql): format mixed count_values labels consistently Signed-off-by: shuiyisong * fix(promql): preserve reset hint warnings for incompatible histograms Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/common/query/src/native_histogram.rs | 59 +- src/promql/src/functions.rs | 2 +- src/promql/src/functions/native_histogram.rs | 448 +++++++- src/query/src/promql/planner.rs | 1005 +++++++++++++++-- .../src/query_engine/default_serializer.rs | 39 +- 5 files changed, 1406 insertions(+), 147 deletions(-) diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index 2db873487bd..b090ac33fcd 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -247,6 +247,25 @@ pub struct NativeHistogram { pub negative_buckets: Vec, } +/// Formats a histogram component like Prometheus's default `%g` formatter. +fn format_promql_histogram_float(value: f64) -> String { + let abs = value.abs(); + if !value.is_finite() || value == 0.0 || (1e-4..1e6).contains(&abs) { + return format_prometheus_float(value); + } + + let scientific = format!("{value:e}"); + let Some((mantissa, exponent)) = scientific.rsplit_once('e') else { + return scientific; + }; + let (sign, exponent) = if let Some(exponent) = exponent.strip_prefix('-') { + ('-', exponent) + } else { + ('+', exponent.strip_prefix('+').unwrap_or(exponent)) + }; + format!("{mantissa}e{sign}{exponent:0>2}") +} + impl NativeHistogram { fn uses_custom_buckets(&self) -> bool { self.schema == CUSTOM_BUCKETS_SCHEMA @@ -378,7 +397,10 @@ impl NativeHistogram { /// Formats this histogram using PromQL native histogram sample notation. pub fn promql_string(&self) -> String { - let mut parts = vec![format!("count:{}", self.count), format!("sum:{}", self.sum)]; + let mut parts = vec![ + format!("count:{}", format_promql_histogram_float(self.count)), + format!("sum:{}", format_promql_histogram_float(self.sum)), + ]; if let Some(buckets) = self.all_buckets() { parts.extend( buckets @@ -392,7 +414,11 @@ impl NativeHistogram { }; format!( "{}{},{}{}:{}", - left, bucket.lower, bucket.upper, right, bucket.count + left, + format_promql_histogram_float(bucket.lower), + format_promql_histogram_float(bucket.upper), + right, + format_promql_histogram_float(bucket.count) ) }), ); @@ -1624,6 +1650,35 @@ mod tests { } } + #[test] + fn promql_string_matches_prometheus_float_format() { + assert_eq!(format_promql_histogram_float(0.0001), "0.0001"); + assert_eq!(format_promql_histogram_float(1_000_000.0), "1e+06"); + + let histogram = NativeHistogram { + schema: CUSTOM_BUCKETS_SCHEMA, + zero_threshold: 0.0, + sum: 2_349_209.324, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: vec![0.00001], + positive_spans: vec![Span { + offset: 0, + length: 2, + }], + negative_spans: Vec::new(), + count: 3.0, + zero_count: 0.0, + positive_buckets: vec![1.0, 2.0], + negative_buckets: Vec::new(), + }; + + assert_eq!( + histogram.promql_string(), + "{count:3, sum:2.349209324e+06, [-Inf,1e-05]:1, (1e-05,+Inf]:2}" + ); + } + #[test] fn span_rebuild_checks_offsets_and_preserves_empty_input() { assert_eq!( diff --git a/src/promql/src/functions.rs b/src/promql/src/functions.rs index b34c21c5569..22cd64b9147 100644 --- a/src/promql/src/functions.rs +++ b/src/promql/src/functions.rs @@ -54,7 +54,7 @@ pub use native_histogram::{ NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, - NativeHistogramToString, + NativeHistogramToString, PromqlFloatToString, }; pub use predict_linear::PredictLinear; pub use quantile::QuantileOverTime; diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index 5b633defa96..18363c8f771 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -20,6 +20,7 @@ use std::mem::size_of; use std::sync::Arc; use common_query::native_histogram::*; +use common_query::prometheus::format_prometheus_float; use common_query::promql_annotations::PromqlAnnotationCollector; use datafusion::arrow::array::{ Array, ArrayRef, BooleanArray, Float64Array, Float64Builder, Int64Array, StringBuilder, @@ -92,20 +93,28 @@ fn read_scalar_f64_arg( #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] enum AnnotationReturn { FloatNull, + BooleanTrue, BooleanFalse, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum AnnotationLevel { + Info, + Warning, +} + impl AnnotationReturn { fn data_type(self) -> DataType { match self { Self::FloatNull => DataType::Float64, - Self::BooleanFalse => DataType::Boolean, + Self::BooleanTrue | Self::BooleanFalse => DataType::Boolean, } } fn scalar_value(self) -> ScalarValue { match self { Self::FloatNull => ScalarValue::Float64(None), + Self::BooleanTrue => ScalarValue::Boolean(Some(true)), Self::BooleanFalse => ScalarValue::Boolean(Some(false)), } } @@ -116,6 +125,7 @@ struct NativeHistogramAnnotationUdf { name: &'static str, signature: Signature, return_kind: AnnotationReturn, + level: AnnotationLevel, message: String, collector: Option, } @@ -124,6 +134,7 @@ impl NativeHistogramAnnotationUdf { fn new( name: &'static str, return_kind: AnnotationReturn, + level: AnnotationLevel, message: String, collector: Option, ) -> Self { @@ -131,6 +142,7 @@ impl NativeHistogramAnnotationUdf { name, signature: Signature::variadic_any(Volatility::Volatile), return_kind, + level, message, collector, } @@ -141,6 +153,7 @@ impl PartialEq for NativeHistogramAnnotationUdf { fn eq(&self, other: &Self) -> bool { self.name == other.name && self.return_kind == other.return_kind + && self.level == other.level && self.message == other.message } } @@ -151,6 +164,7 @@ impl Hash for NativeHistogramAnnotationUdf { fn hash(&self, state: &mut H) { self.name.hash(state); self.return_kind.hash(state); + self.level.hash(state); self.message.hash(state); } } @@ -188,7 +202,10 @@ impl ScalarUDFImpl for NativeHistogramAnnotationUdf { .cloned() .or_else(|| self.collector.clone()) { - collector.record_info(self.message.clone()); + match self.level { + AnnotationLevel::Info => collector.record_info(self.message.clone()), + AnnotationLevel::Warning => collector.record_warning(self.message.clone()), + } } Ok(ColumnarValue::Scalar(self.return_kind.scalar_value())) } @@ -205,6 +222,10 @@ impl NativeHistogramDrop { "prom_native_histogram_drop_bool" } + const fn bool_true_name() -> &'static str { + "prom_native_histogram_keep_bool" + } + pub fn float_null_udf( message: String, collector: Option, @@ -212,6 +233,7 @@ impl NativeHistogramDrop { ScalarUDF::new_from_impl(NativeHistogramAnnotationUdf::new( Self::float_null_name(), AnnotationReturn::FloatNull, + AnnotationLevel::Info, message, collector, )) @@ -224,14 +246,39 @@ impl NativeHistogramDrop { ScalarUDF::new_from_impl(NativeHistogramAnnotationUdf::new( Self::bool_false_name(), AnnotationReturn::BooleanFalse, + AnnotationLevel::Info, + message, + collector, + )) + } + + pub fn bool_true_udf( + message: String, + collector: Option, + ) -> ScalarUDF { + ScalarUDF::new_from_impl(NativeHistogramAnnotationUdf::new( + Self::bool_true_name(), + AnnotationReturn::BooleanTrue, + AnnotationLevel::Info, + message, + collector, + )) + } + + pub fn warning_bool_false_udf( + message: String, + collector: Option, + ) -> ScalarUDF { + ScalarUDF::new_from_impl(NativeHistogramAnnotationUdf::new( + Self::bool_false_name(), + AnnotationReturn::BooleanFalse, + AnnotationLevel::Warning, message, collector, )) } } -// Synthetic drop UDFs only cover info-level invalid PromQL combinations. -// Warning annotations stay in the histogram UDFs that actually drop samples. fn record_info(collector: &Option, message: impl Into) { if let Some(collector) = collector { collector.record_info(message); @@ -358,10 +405,8 @@ fn histogram_pair_udf_with_collector( match (read_histogram(lhs, row)?, read_histogram(rhs, row)?) { (Some(lhs), Some(rhs)) => { record_custom_reconciliation(&collector, name, &lhs, &rhs); + record_counter_reset_contradiction(&collector, name, &lhs, &rhs); let result = op(&lhs, &rhs); - if result.is_some() { - record_counter_reset_contradiction(&collector, name, &lhs, &rhs); - } if result.is_none() { record_warning( &collector, @@ -697,6 +742,39 @@ impl NativeHistogramStdvar { } } +/// Formats float samples as PromQL label values. +pub struct PromqlFloatToString; + +impl PromqlFloatToString { + pub const fn name() -> &'static str { + "prom_float_to_string" + } + + pub fn scalar_udf() -> ScalarUDF { + create_udf( + Self::name(), + vec![DataType::Float64], + DataType::Utf8, + Volatility::Volatile, + Arc::new(|input: &[ColumnarValue]| { + let values = extract_array(&input[0])?; + let values = values + .as_any() + .downcast_ref::() + .expect("validated Float64 input"); + let mut result = StringBuilder::new(); + for value in values.iter() { + match value { + Some(value) => result.append_value(format_prometheus_float(value)), + None => result.append_null(), + } + } + Ok(ColumnarValue::Array(Arc::new(result.finish()))) + }), + ) + } +} + pub struct NativeHistogramToString; impl NativeHistogramToString { @@ -845,19 +923,38 @@ impl NativeHistogramAggregateAccumulator { } } - fn push_histogram(&mut self, histogram: NativeHistogram, count: u64) { - if self.dropped_incompatible { - return; + fn push_histogram(&mut self, histogram: NativeHistogram, count: u64) -> DfResult<()> { + if self.kind.needs_count() && count == 0 { + return Ok(()); } self.observe_reset_hints( histogram.reset_hint == COUNTER_RESET_HINT, histogram.reset_hint == NOT_COUNTER_RESET_HINT, ); + if self.dropped_incompatible { + return Ok(()); + } + let combined_count = if self.kind.needs_count() { + self.count.checked_add(count).ok_or_else(|| { + DataFusionError::Execution(format!( + "{}: native histogram sample count overflow", + self.kind.name() + )) + })? + } else { + self.count + }; let value = match self.value.take() { Some(value) => { record_custom_reconciliation(&self.collector, self.kind.name(), &value, &histogram); - match value.add(&histogram) { + let combined = match self.kind { + NativeHistogramAggregateKind::Sum => value.add(&histogram), + NativeHistogramAggregateKind::Avg => { + weighted_histogram_mean(value, self.count, histogram, count, combined_count) + } + }; + match combined { Some(value) => Some(value), None => { self.record_incompatible(); @@ -869,8 +966,9 @@ impl NativeHistogramAggregateAccumulator { }; if !self.dropped_incompatible { self.value = value; - self.count += count; + self.count = combined_count; } + Ok(()) } fn mark_incompatible(&mut self) { @@ -891,24 +989,48 @@ impl NativeHistogramAggregateAccumulator { } } -fn range_sum_histograms( +fn weighted_histogram_mean( + left: NativeHistogram, + left_count: u64, + right: NativeHistogram, + right_count: u64, + total_count: u64, +) -> Option { + let total_count = total_count as f64; + left.scale(left_count as f64 / total_count) + .add(&right.scale(right_count as f64 / total_count)) +} + +fn range_fold_histograms( samples: Vec, + kind: NativeHistogramAggregateKind, name: &'static str, collector: &Option, ) -> Option { + if samples + .iter() + .any(|histogram| histogram.reset_hint == COUNTER_RESET_HINT) + && samples + .iter() + .any(|histogram| histogram.reset_hint == NOT_COUNTER_RESET_HINT) + { + record_counter_reset_contradiction_warning(collector, name); + } + let mut value = None; - let mut counter_reset_seen = false; - let mut not_counter_reset_seen = false; + let mut count = 0u64; for histogram in samples { - counter_reset_seen |= histogram.reset_hint == COUNTER_RESET_HINT; - not_counter_reset_seen |= histogram.reset_hint == NOT_COUNTER_RESET_HINT; - if counter_reset_seen && not_counter_reset_seen { - record_counter_reset_contradiction_warning(collector, name); - } value = match value { Some(value) => { record_custom_reconciliation(collector, name, &value, &histogram); - match value.add(&histogram) { + let next_count = count.checked_add(1)?; + let combined = match kind { + NativeHistogramAggregateKind::Sum => value.add(&histogram), + NativeHistogramAggregateKind::Avg => { + weighted_histogram_mean(value, count, histogram, 1, next_count) + } + }; + match combined { Some(value) => Some(value), None => { record_warning( @@ -923,6 +1045,7 @@ fn range_sum_histograms( } None => Some(histogram), }; + count = count.checked_add(1)?; } value } @@ -1016,14 +1139,18 @@ fn native_histogram_range_histogram( continue; }; let histogram = match kind { - NativeHistogramRangeHistogramKind::Sum => { - range_sum_histograms(samples, func_name, &collector) - } - NativeHistogramRangeHistogramKind::Avg => { - let count = samples.len(); - range_sum_histograms(samples, func_name, &collector) - .map(|histogram| histogram.divide_by(count as f64)) - } + NativeHistogramRangeHistogramKind::Sum => range_fold_histograms( + samples, + NativeHistogramAggregateKind::Sum, + func_name, + &collector, + ), + NativeHistogramRangeHistogramKind::Avg => range_fold_histograms( + samples, + NativeHistogramAggregateKind::Avg, + func_name, + &collector, + ), NativeHistogramRangeHistogramKind::Last => samples.last().cloned(), }; result.push(histogram); @@ -2113,7 +2240,7 @@ impl DfAccumulator for NativeHistogramAggregateAccumulator { let Some(histogram) = read_histogram(histograms, row)? else { continue; }; - self.push_histogram(histogram, 1); + self.push_histogram(histogram, 1)?; } Ok(()) @@ -2125,7 +2252,7 @@ impl DfAccumulator for NativeHistogramAggregateAccumulator { (_, false, None) => None, (NativeHistogramAggregateKind::Sum, false, value) => value, (NativeHistogramAggregateKind::Avg, false, Some(value)) if self.count > 0 => { - Some(value.divide_by(self.count as f64)) + Some(value) } (NativeHistogramAggregateKind::Avg, _, _) => None, }; @@ -2214,22 +2341,21 @@ impl DfAccumulator for NativeHistogramAggregateAccumulator { })?; for row in 0..histograms.len() { - if dropped.value(row) { - self.mark_incompatible(); - continue; - } - if self.dropped_incompatible { - continue; - } self.observe_reset_hints( counter_reset_seen.value(row), not_counter_reset_seen.value(row), ); + if dropped.value(row) { + self.mark_incompatible(); + } + if self.dropped_incompatible { + continue; + } let Some(histogram) = read_histogram(histograms, row)? else { continue; }; let count = counts.map(|counts| counts.value(row)).unwrap_or(1); - self.push_histogram(histogram, count); + self.push_histogram(histogram, count)?; } Ok(()) @@ -2799,6 +2925,15 @@ mod tests { read_histogram(&array, 0).unwrap().unwrap() } + fn evaluated_histogram( + accumulator: &mut NativeHistogramAggregateAccumulator, + ) -> NativeHistogram { + let ScalarValue::Struct(array) = accumulator.evaluate().unwrap() else { + panic!("native histogram accumulator returned a non-struct value"); + }; + read_histogram(&array, 0).unwrap().unwrap() + } + fn histogram_range_input(values: Vec>) -> Vec { let timestamps = Arc::new(TimestampMillisecondArray::from_iter( (0..values.len()).map(|idx| Some((idx as i64 + 1) * 1000)), @@ -3195,6 +3330,80 @@ mod tests { assert_eq!(result.positive_buckets, vec![2.0]); } + #[test] + fn histogram_averages_avoid_sum_overflow() { + let large = sample_histogram(1.0e308, 1.0e308, vec![1.0e308]); + + let range_average = run_histogram_range_udf( + NativeHistogramAvgOverTime::scalar_udf(), + vec![large.clone(), large.clone()], + ); + assert_eq!(range_average.count, 1.0e308); + assert_eq!(range_average.sum, 1.0e308); + assert_eq!(range_average.positive_buckets, vec![1.0e308]); + + let mut aggregate = + NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Avg, None); + aggregate.push_histogram(large.clone(), 1).unwrap(); + aggregate.push_histogram(large, 1).unwrap(); + let aggregate_average = evaluated_histogram(&mut aggregate); + assert_eq!(aggregate_average.count, 1.0e308); + assert_eq!(aggregate_average.sum, 1.0e308); + assert_eq!(aggregate_average.positive_buckets, vec![1.0e308]); + } + + #[test] + fn histogram_average_partial_states_are_weighted() { + let mut first = + NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Avg, None); + first + .push_histogram(sample_histogram(1.0, 1.0, vec![1.0]), 1) + .unwrap(); + first + .push_histogram(sample_histogram(3.0, 3.0, vec![3.0]), 1) + .unwrap(); + let first_state = first + .state() + .unwrap() + .into_iter() + .map(|value| value.to_array_of_size(1).unwrap()) + .collect::>(); + + let mut second = + NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Avg, None); + second + .push_histogram(sample_histogram(8.0, 8.0, vec![8.0]), 1) + .unwrap(); + let second_state = second + .state() + .unwrap() + .into_iter() + .map(|value| value.to_array_of_size(1).unwrap()) + .collect::>(); + + let mut merged = + NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Avg, None); + merged.merge_batch(&first_state).unwrap(); + merged.merge_batch(&second_state).unwrap(); + let average = evaluated_histogram(&mut merged); + assert_eq!(average.count, 4.0); + assert_eq!(average.sum, 4.0); + assert_eq!(average.positive_buckets, vec![4.0]); + } + + #[test] + fn histogram_average_rejects_sample_count_overflow() { + let mut aggregate = + NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Avg, None); + aggregate.value = Some(sample_histogram(1.0, 1.0, vec![1.0])); + aggregate.count = u64::MAX; + + let error = aggregate + .push_histogram(sample_histogram(1.0, 1.0, vec![1.0]), 1) + .unwrap_err(); + assert!(error.to_string().contains("sample count overflow")); + } + #[test] fn presence_only_range_functions_preserve_null_semantics() { let first = sample_histogram(1.0, 1.0, vec![1.0]); @@ -3311,25 +3520,32 @@ mod tests { } #[test] - fn subtraction_records_reset_hint_contradictions() { - let mut left = sample_histogram(2.0, 2.0, vec![2.0]); - left.reset_hint = COUNTER_RESET_HINT; - let mut right = sample_histogram(1.0, 1.0, vec![1.0]); - right.reset_hint = NOT_COUNTER_RESET_HINT; - let collector = PromqlAnnotationCollector::default(); + fn subtraction_records_reset_hint_contradictions_for_incompatible_histograms() { + for incompatible in [false, true] { + let mut left = sample_histogram(2.0, 2.0, vec![2.0]); + left.reset_hint = COUNTER_RESET_HINT; + let mut right = sample_histogram(1.0, 1.0, vec![1.0]); + right.reset_hint = NOT_COUNTER_RESET_HINT; + if incompatible { + right.schema = CUSTOM_BUCKETS_SCHEMA; + right.custom_values = vec![1.0]; + } + let collector = PromqlAnnotationCollector::default(); - run_histogram_udf( - NativeHistogramSub::scalar_udf_with_collector(Some(collector.clone())), - vec![ - ColumnarValue::Array(build_histogram_array(&[Some(left)])), - ColumnarValue::Array(build_histogram_array(&[Some(right)])), - ], - ); + run_udf( + NativeHistogramSub::scalar_udf_with_collector(Some(collector.clone())), + vec![ + ColumnarValue::Array(build_histogram_array(&[Some(left)])), + ColumnarValue::Array(build_histogram_array(&[Some(right)])), + ], + native_histogram_arrow_type(), + ); - assert!(collected_warnings(&collector).contains(&format!( - "{}: native histogram counter reset hints contradict", - NativeHistogramSub::name() - ))); + assert!(collected_warnings(&collector).contains(&format!( + "{}: native histogram counter reset hints contradict", + NativeHistogramSub::name() + ))); + } } #[test] @@ -3343,8 +3559,9 @@ mod tests { let range_collector = PromqlAnnotationCollector::default(); assert!( - range_sum_histograms( + range_fold_histograms( vec![reset.clone(), unknown.clone(), not_reset.clone()], + NativeHistogramAggregateKind::Sum, NativeHistogramSumOverTime::name(), &Some(range_collector.clone()), ) @@ -3360,8 +3577,8 @@ mod tests { NativeHistogramAggregateKind::Avg, ] { let mut first_partial = NativeHistogramAggregateAccumulator::new(kind, None); - first_partial.push_histogram(reset.clone(), 1); - first_partial.push_histogram(unknown.clone(), 1); + first_partial.push_histogram(reset.clone(), 1).unwrap(); + first_partial.push_histogram(unknown.clone(), 1).unwrap(); let first_states = first_partial .state() .unwrap() @@ -3370,7 +3587,7 @@ mod tests { .collect::>(); let mut second_partial = NativeHistogramAggregateAccumulator::new(kind, None); - second_partial.push_histogram(not_reset.clone(), 1); + second_partial.push_histogram(not_reset.clone(), 1).unwrap(); let second_states = second_partial .state() .unwrap() @@ -3390,6 +3607,115 @@ mod tests { } } + #[test] + fn incompatible_aggregates_do_not_hide_reset_hint_contradictions() { + let mut reset = sample_histogram(1.0, 1.0, vec![1.0]); + reset.reset_hint = COUNTER_RESET_HINT; + let mut incompatible_reset = sample_histogram(2.0, 2.0, vec![2.0]); + incompatible_reset.schema = CUSTOM_BUCKETS_SCHEMA; + incompatible_reset.custom_values = vec![1.0]; + incompatible_reset.reset_hint = COUNTER_RESET_HINT; + let mut not_reset = sample_histogram(3.0, 3.0, vec![3.0]); + not_reset.reset_hint = NOT_COUNTER_RESET_HINT; + + for kind in [ + NativeHistogramAggregateKind::Sum, + NativeHistogramAggregateKind::Avg, + ] { + for histograms in [ + [reset.clone(), incompatible_reset.clone(), not_reset.clone()], + [reset.clone(), not_reset.clone(), incompatible_reset.clone()], + ] { + let collector = PromqlAnnotationCollector::default(); + let mut aggregate = + NativeHistogramAggregateAccumulator::new(kind, Some(collector.clone())); + for histogram in histograms { + aggregate.push_histogram(histogram, 1).unwrap(); + } + + assert!(aggregate.dropped_incompatible); + let warnings = collected_warnings(&collector); + assert!(warnings.contains(&format!( + "{}: dropped native histogram aggregate with incompatible schemas", + kind.name() + ))); + assert!(warnings.contains(&format!( + "{}: native histogram counter reset hints contradict", + kind.name() + ))); + } + + let mut dropped_partial = NativeHistogramAggregateAccumulator::new(kind, None); + dropped_partial.push_histogram(reset.clone(), 1).unwrap(); + dropped_partial + .push_histogram(incompatible_reset.clone(), 1) + .unwrap(); + let dropped_states = dropped_partial + .state() + .unwrap() + .into_iter() + .map(|value| value.to_array_of_size(1).unwrap()) + .collect::>(); + + let mut opposing_partial = NativeHistogramAggregateAccumulator::new(kind, None); + opposing_partial + .push_histogram(not_reset.clone(), 1) + .unwrap(); + let opposing_states = opposing_partial + .state() + .unwrap() + .into_iter() + .map(|value| value.to_array_of_size(1).unwrap()) + .collect::>(); + + for (first, second) in [ + (&dropped_states, &opposing_states), + (&opposing_states, &dropped_states), + ] { + let collector = PromqlAnnotationCollector::default(); + let mut merged = + NativeHistogramAggregateAccumulator::new(kind, Some(collector.clone())); + merged.merge_batch(first).unwrap(); + merged.merge_batch(second).unwrap(); + + assert!(merged.dropped_incompatible); + assert!(collected_warnings(&collector).contains(&format!( + "{}: native histogram counter reset hints contradict", + kind.name() + ))); + } + } + + for (kind, name) in [ + ( + NativeHistogramAggregateKind::Sum, + NativeHistogramSumOverTime::name(), + ), + ( + NativeHistogramAggregateKind::Avg, + NativeHistogramAvgOverTime::name(), + ), + ] { + let collector = PromqlAnnotationCollector::default(); + assert!( + range_fold_histograms( + vec![reset.clone(), incompatible_reset.clone(), not_reset.clone()], + kind, + name, + &Some(collector.clone()), + ) + .is_none() + ); + let warnings = collected_warnings(&collector); + assert!(warnings.contains(&format!( + "{name}: dropped native histogram range with incompatible schemas" + ))); + assert!(warnings.contains(&format!( + "{name}: native histogram counter reset hints contradict" + ))); + } + } + #[test] fn histogram_aggregate_accumulator_accounts_for_heap_allocations() { let histogram = sample_histogram(3.0, 3.0, vec![1.0, 2.0]); @@ -3402,7 +3728,7 @@ mod tests { NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Sum, None); let empty_size = accumulator.size(); - accumulator.push_histogram(histogram, 1); + accumulator.push_histogram(histogram, 1).unwrap(); assert!(heap_size > 0); assert_eq!(accumulator.size(), empty_size + heap_size); diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 0933bbcaca4..5dfb1a537a8 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -48,6 +48,7 @@ use datafusion::sql::TableReference; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; use datafusion_common::{DFSchema, NullEquality}; use datafusion_expr::expr::WindowFunctionParams; +use datafusion_expr::expr_fn::when; use datafusion_expr::utils::{conjunction, disjunction}; use datafusion_expr::{ ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit, @@ -64,17 +65,17 @@ use promql::extension_plan::{ use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, - NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAvg, - NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, - NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDivScalar, - NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, NativeHistogramIDelta, - NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, - NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq, - NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, - NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, - NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, PredictLinear, - PresentOverTime, QuantileOverTime, Rate, Resets, Round, StddevOverTime, StdvarOverTime, - SumOverTime, quantile_udaf, + NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, + NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, + NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, + NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, + NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, + NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg, + NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, + NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, + NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, + NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime, + Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; @@ -654,9 +655,21 @@ impl PromPlanner { // calculate columns to group by // Need to append time index column into group by columns let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?; + let mixed_sample_columns = + Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())); + // Aggregates over native histogram inputs may drop every sample in a group + // (e.g. `min` over histogram-only samples, or `sum` over histograms with + // incompatible schemas) and leave a NULL-valued group row behind. Compute this + // before `create_aggregate_exprs` mutates `ctx.field_columns`. + let preserve_any_value = mixed_sample_columns.is_some(); + let has_native_histogram = preserve_any_value + || self.all_field_columns_are_native_histograms(input.schema()); // convert op and value columns to aggregate exprs let (mut aggr_exprs, prev_field_exprs) = self.create_aggregate_exprs(*op, param, &input)?; + let prev_field_exprs = + normalize_cols(prev_field_exprs, &input).context(DataFusionPlanningSnafu)?; let keep_tsid = op.id() != token::T_COUNT_VALUES && input_has_tsid @@ -680,16 +693,28 @@ impl PromPlanner { let label = Self::get_param_value_as_str(*op, param)?; // `count_values` must be grouped by fields, // and project the fields to the new label. - group_exprs.extend(prev_field_exprs.clone()); + let count_value_exprs = prev_field_exprs.iter().map(|expr| { + match expr { + DfExpr::Column(column) => DfExpr::Column(column.clone()), + _ => DfExpr::Column(Column::from_name(expr.schema_name().to_string())), + } + .alias(label) + }); + let aggregate_group_exprs = group_exprs + .iter() + .cloned() + .chain(prev_field_exprs.clone()) + .collect::>(); + group_exprs.push(col(label)); let project_fields = self .create_field_column_exprs()? .into_iter() .chain(self.create_tag_column_exprs()?) .chain(Some(self.create_time_index_column_expr()?)) - .chain(prev_field_exprs.into_iter().map(|expr| expr.alias(label))); + .chain(count_value_exprs); builder - .aggregate(group_exprs.clone(), aggr_exprs) + .aggregate(aggregate_group_exprs, aggr_exprs) .context(DataFusionPlanningSnafu)? .project(project_fields) .context(DataFusionPlanningSnafu)? @@ -699,6 +724,59 @@ impl PromPlanner { .context(DataFusionPlanningSnafu)? }; + let builder = if let Some((float, histogram)) = mixed_sample_columns { + let builder = match op.id() { + token::T_SUM | token::T_AVG => builder + .filter(self.mixed_aggregate_filter_expr(*op, &float, &histogram)?) + .context(DataFusionPlanningSnafu)?, + token::T_MIN + | token::T_MAX + | token::T_STDDEV + | token::T_STDVAR + | token::T_QUANTILE => builder + .filter(self.mixed_ignored_histogram_filter_expr(*op, &histogram)?) + .context(DataFusionPlanningSnafu)?, + _ => builder, + }; + + match op.id() { + token::T_SUM + | token::T_AVG + | token::T_MIN + | token::T_MAX + | token::T_STDDEV + | token::T_STDVAR + | token::T_QUANTILE => { + let project_fields = self + .create_field_column_exprs()? + .into_iter() + .chain(self.create_tag_column_exprs()?) + .chain(self.ctx.use_tsid.then_some(DfExpr::Column( + Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME), + ))) + .chain(Some(self.create_time_index_column_expr()?)); + builder + .project(project_fields) + .context(DataFusionPlanningSnafu)? + } + _ => builder, + } + } else { + builder + }; + + // Drop group rows whose every aggregated sample was discarded (NULL), so that + // e.g. `group(min(native_histogram))` doesn't resurrect groups Prometheus + // considers unseen. For alternative float/histogram fields keep the row if any + // field survived. + let builder = if has_native_histogram { + builder + .filter(self.create_empty_values_filter_expr(preserve_any_value)?) + .context(DataFusionPlanningSnafu)? + } else { + builder + }; + let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false)); builder @@ -731,6 +809,32 @@ impl PromPlanner { let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?; + let mut input = input; + if let Some((float_column, histogram_column)) = + Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())) + { + let drop_histogram = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::bool_false_udf( + format!( + "{}: dropped native histogram samples because this aggregation is not supported for native histograms", + op + ), + self.promql_annotations.clone(), + )), + args: vec![col(&histogram_column)], + }); + let keep_float = when(col(&histogram_column).is_not_null(), drop_histogram) + .otherwise(col(&float_column).is_not_null()) + .context(DataFusionPlanningSnafu)?; + input = LogicalPlanBuilder::from(input) + .filter(keep_float) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + self.ctx.field_columns = vec![float_column]; + } + if self.all_field_columns_are_native_histograms(input.schema()) { let promql_annotations = self.promql_annotations.clone(); let input = self.projection_for_each_field_column(input, |col| { @@ -4123,9 +4227,11 @@ impl PromPlanner { param: &Option>, input_plan: &LogicalPlan, ) -> Result<(Vec, Vec)> { - let mut non_col_args = Vec::new(); + let mixed_sample_columns = + Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())); let is_group_agg = op.id() == token::T_GROUP; - if is_group_agg { + if is_group_agg && mixed_sample_columns.is_none() { ensure!( self.ctx.field_columns.len() == 1, MultiFieldsNotSupportedSnafu { @@ -4133,46 +4239,28 @@ impl PromPlanner { } ); } - let aggr = match op.id() { - token::T_SUM => sum_udaf(), - token::T_QUANTILE => { - let q = - Self::get_param_as_literal_expr(param, Some(op), Some(ArrowDataType::Float64))?; - non_col_args.push(q); - quantile_udaf() - } - token::T_AVG => avg_udaf(), - token::T_COUNT_VALUES | token::T_COUNT => count_udaf(), - token::T_MIN => min_udaf(), - token::T_MAX => max_udaf(), - // PromQL's `group()` aggregator produces 1 for each group. - // Use `max(1.0)` (per-group) to match semantics and output type (Float64). - token::T_GROUP => max_udaf(), - token::T_STDDEV => stddev_pop_udaf(), - token::T_STDVAR => var_pop_udaf(), - token::T_TOPK | token::T_BOTTOMK => UnsupportedExprSnafu { - name: format!("{op:?}"), - } - .fail()?, - _ => UnexpectedTokenSnafu { token: op }.fail()?, - }; + + if let Some((float, histogram)) = mixed_sample_columns { + return self.create_mixed_aggregate_exprs(op, param, &float, &histogram); + } + + if self.all_field_columns_are_native_histograms(input_plan.schema()) { + return self.create_native_histogram_aggregate_exprs(op, input_plan); + } // perform aggregate operation to each value column - let exprs: Vec = self + let exprs = self .ctx .field_columns .iter() .map(|col| { - if is_group_agg { - aggr.call(vec![lit(1_f64)]) - } else { - non_col_args.push(DfExpr::Column(Column::from_name(col))); - let expr = aggr.call(non_col_args.clone()); - non_col_args.pop(); - expr - } + Self::create_numeric_aggregate_expr( + op, + param, + DfExpr::Column(Column::from_name(col)), + ) }) - .collect::>(); + .collect::>>()?; // if the aggregator is `count_values`, it must be grouped by current fields. let prev_field_exprs = if op.id() == token::T_COUNT_VALUES { @@ -4208,6 +4296,273 @@ impl PromPlanner { Ok((exprs, prev_field_exprs)) } + fn create_numeric_aggregate_expr( + op: TokenType, + param: &Option>, + input: DfExpr, + ) -> Result { + let expr = match op.id() { + token::T_SUM => sum_udaf().call(vec![input]), + token::T_QUANTILE => { + let q = + Self::get_param_as_literal_expr(param, Some(op), Some(ArrowDataType::Float64))?; + quantile_udaf().call(vec![q, input]) + } + token::T_AVG => avg_udaf().call(vec![input]), + token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]), + token::T_MIN => min_udaf().call(vec![input]), + token::T_MAX => max_udaf().call(vec![input]), + // PromQL's `group()` aggregator produces 1 for each group. + // Use `max(1.0)` (per-group) to match semantics and output type (Float64). + token::T_GROUP => max_udaf().call(vec![lit(1_f64)]), + token::T_STDDEV => stddev_pop_udaf().call(vec![input]), + token::T_STDVAR => var_pop_udaf().call(vec![input]), + token::T_TOPK | token::T_BOTTOMK => { + return UnsupportedExprSnafu { + name: format!("{op:?}"), + } + .fail(); + } + _ => return UnexpectedTokenSnafu { token: op }.fail(), + }; + Ok(expr) + } + + fn create_mixed_aggregate_exprs( + &mut self, + op: TokenType, + param: &Option>, + float_column: &str, + histogram_column: &str, + ) -> Result<(Vec, Vec)> { + let float_input = DfExpr::Column(Column::from_name(float_column)); + let histogram_input = DfExpr::Column(Column::from_name(histogram_column)); + let float_count = count_udaf().call(vec![float_input.clone()]); + let histogram_count = count_udaf().call(vec![histogram_input.clone()]); + let mixed_sample_value = || { + DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: vec![ + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(PromqlFloatToString::scalar_udf()), + args: vec![float_input.clone()], + }), + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramToString::scalar_udf()), + args: vec![histogram_input.clone()], + }), + ], + }) + }; + + let (exprs, prev_field_exprs, field_columns) = match op.id() { + token::T_SUM | token::T_AVG => ( + vec![ + Self::create_numeric_aggregate_expr(op, param, float_input)? + .alias(float_column), + self.create_native_histogram_aggregate_expr(op, histogram_column)?, + float_count.alias(Self::mixed_sample_count_name(float_column)), + histogram_count.alias(Self::mixed_sample_count_name(histogram_column)), + ], + vec![], + vec![float_column.to_string(), histogram_column.to_string()], + ), + token::T_COUNT => { + let present = when( + float_input + .clone() + .is_not_null() + .or(histogram_input.clone().is_not_null()), + lit(1_i64), + ) + .otherwise(lit(ScalarValue::Int64(None))) + .context(DataFusionPlanningSnafu)?; + ( + vec![count_udaf().call(vec![present]).alias(float_column)], + vec![], + vec![float_column.to_string()], + ) + } + token::T_GROUP => ( + vec![max_udaf().call(vec![lit(1_f64)]).alias(float_column)], + vec![], + vec![float_column.to_string()], + ), + token::T_COUNT_VALUES => { + let value = mixed_sample_value(); + ( + vec![count_udaf().call(vec![value.clone()]).alias(float_column)], + vec![value], + vec![float_column.to_string()], + ) + } + token::T_MIN | token::T_MAX | token::T_STDDEV | token::T_STDVAR | token::T_QUANTILE => { + ( + vec![ + Self::create_numeric_aggregate_expr(op, param, float_input)? + .alias(float_column), + histogram_count.alias(Self::mixed_sample_count_name(histogram_column)), + ], + vec![], + vec![float_column.to_string()], + ) + } + token::T_TOPK | token::T_BOTTOMK => { + return UnsupportedExprSnafu { + name: format!("{op:?}"), + } + .fail(); + } + _ => return UnexpectedTokenSnafu { token: op }.fail(), + }; + + self.ctx.field_columns = field_columns; + Ok((exprs, prev_field_exprs)) + } + + fn mixed_sample_count_column(column: &str) -> DfExpr { + DfExpr::Column(Column::from_name(Self::mixed_sample_count_name(column))) + } + + fn mixed_sample_count_name(column: &str) -> String { + format!("__promql_sample_count({column})") + } + + fn mixed_aggregate_filter_expr( + &self, + op: TokenType, + float_column: &str, + histogram_column: &str, + ) -> Result { + let float_count = Self::mixed_sample_count_column(float_column); + let histogram_count = Self::mixed_sample_count_column(histogram_column); + let mixed = float_count + .clone() + .gt(lit(0_i64)) + .and(histogram_count.clone().gt(lit(0_i64))); + let drop_mixed = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::warning_bool_false_udf( + format!( + "{op}: dropped aggregation result containing both float and native histogram samples" + ), + self.promql_annotations.clone(), + )), + args: vec![float_count, histogram_count], + }); + + when(mixed, drop_mixed) + .otherwise(lit(true)) + .context(DataFusionPlanningSnafu) + } + + fn mixed_ignored_histogram_filter_expr( + &self, + op: TokenType, + histogram_column: &str, + ) -> Result { + let histogram_count = Self::mixed_sample_count_column(histogram_column); + let has_histograms = histogram_count.clone().gt(lit(0_i64)); + let record_info = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::bool_true_udf( + format!( + "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" + ), + self.promql_annotations.clone(), + )), + args: vec![histogram_count], + }); + + when(has_histograms, record_info) + .otherwise(lit(true)) + .context(DataFusionPlanningSnafu) + } + + fn create_native_histogram_aggregate_expr( + &self, + op: TokenType, + column: &str, + ) -> Result { + let input = DfExpr::Column(Column::from_name(column)); + let expr = match op.id() { + token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector( + self.promql_annotations.clone(), + )) + .call(vec![input]) + .alias(column), + token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector( + self.promql_annotations.clone(), + )) + .call(vec![input]) + .alias(column), + token::T_COUNT_VALUES | token::T_COUNT => { + count_udaf().call(vec![input]).alias(column) + } + token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(column), + token::T_MIN + | token::T_MAX + | token::T_STDDEV + | token::T_STDVAR + | token::T_QUANTILE + | token::T_TOPK + | token::T_BOTTOMK => sum_udaf() + .call(vec![DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::float_null_udf( + format!( + "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" + ), + self.promql_annotations.clone(), + )), + args: vec![input], + })]) + .alias(column), + _ => return UnexpectedTokenSnafu { token: op }.fail(), + }; + Ok(expr) + } + + fn create_native_histogram_aggregate_exprs( + &mut self, + op: TokenType, + input_plan: &LogicalPlan, + ) -> Result<(Vec, Vec)> { + let prev_field_exprs = if op.id() == token::T_COUNT_VALUES { + ensure!( + self.ctx.field_columns.len() == 1, + UnsupportedExprSnafu { + name: "count_values on multi-value input" + } + ); + self.ctx + .field_columns + .iter() + .map(|col| { + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramToString::scalar_udf()), + args: vec![DfExpr::Column(Column::from_name(col))], + }) + }) + .collect::>() + } else { + vec![] + }; + + let exprs = self + .ctx + .field_columns + .iter() + .map(|col| self.create_native_histogram_aggregate_expr(op, col)) + .collect::>>()?; + + let normalized_exprs = + normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?; + self.ctx.field_columns = normalized_exprs + .into_iter() + .map(|expr| expr.schema_name().to_string()) + .collect(); + + Ok((exprs, prev_field_exprs)) + } + fn get_param_value_as_str(op: TokenType, param: &Option>) -> Result<&str> { let param = param .as_deref() @@ -6346,7 +6701,7 @@ mod test { use common_base::Plugins; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_query::native_histogram::{ - CounterResetHint, NativeHistogram, build_histogram_array, + CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array, }; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; @@ -6880,6 +7235,65 @@ mod test { (planner, plan) } + async fn mixed_aggregate_input(histograms: Vec) -> (PromPlanner, LogicalPlan) { + let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0"); + let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0"); + let row_count = histograms.len() + 1; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new("k", ArrowDataType::Utf8, false), + Field::new(&float_field, ArrowDataType::Float64, true), + Field::new( + &histogram_field, + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let mut histogram_values = Vec::with_capacity(row_count); + histogram_values.push(None); + histogram_values.extend(histograms.into_iter().map(Some)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1; row_count])), + Arc::new(StringArray::from_iter_values( + (0..row_count).map(|row| format!("kind_{row}")), + )), + Arc::new(Float64Array::from_iter( + (0..row_count).map(|row| (row == 0).then_some(1.25)), + )), + build_histogram_array(&histogram_values), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let planner = PromPlanner { + table_provider, + ctx: PromPlannerContext { + table_name: Some("mixed".to_string()), + time_index_column: Some("ts".to_string()), + field_columns: vec![float_field, histogram_field], + tag_columns: vec!["k".to_string()], + ..Default::default() + }, + promql_annotations: None, + }; + (planner, plan) + } + fn assert_no_internal_or_keys(schema: &DFSchema) { assert!( schema @@ -6906,6 +7320,26 @@ mod test { .collect() } + fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + let values = datafusion::arrow::compute::cast( + batch.column_by_name(column).unwrap(), + &ArrowDataType::Float64, + ) + .unwrap(); + values + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .collect::>() + }) + .collect() + } + fn histograms(batches: &[RecordBatch], column: &str) -> Vec { batches .iter() @@ -7325,18 +7759,41 @@ mod test { ]; let schema = Arc::new(Schema::new(columns)); let table_meta = TableMetaBuilder::empty() - .schema(schema) + .schema(schema.clone()) .primary_key_indices(vec![0]) .value_indices(vec![2, 3]) .next_column_id(1024) .build() .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); + let table_info = Arc::new( + TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from(vec!["float", "histogram"])), + Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])), + build_histogram_array(&[None, Some(direct_or_histogram())]), + Arc::new(Float64Array::from(vec![Some(2.0), None])), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + table_name, + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 1024, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); assert!( catalog_list @@ -9812,6 +10269,50 @@ mod test { ); } + #[tokio::test] + async fn mixed_or_topk_bottomk_ignore_native_histograms() { + for op in ["topk", "bottomk"] { + let collector = PromqlAnnotationCollector::default(); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + assert!( + plan.schema() + .fields() + .iter() + .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()), + "{plan:?}" + ); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(values(&batches, &float_field), vec![2.0], "{op}"); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert_eq!( + infos, + vec![format!( + "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" + )] + ); + } + } + #[tokio::test] async fn native_histogram_scalar_is_ignored_before_scalar_calculate() { let plan = native_histogram_plan("scalar(some_metric)").await; @@ -10142,9 +10643,9 @@ mod test { #[tokio::test] async fn native_histogram_mixed_field_table_behaves() { - // Metric Engine exposes float and histogram samples as alternative nullable fields. - // Histogram functions must select the histogram field without adding a NULL float - // field that would make the empty-value filter reject every row. + // Exercise function planning after float and histogram samples have already been + // represented as alternative nullable fields. Histogram functions must select the + // histogram field without adding a NULL float field that would reject every row. let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; let plan = PromPlanner::stmt_to_plan( table_provider, @@ -11126,15 +11627,14 @@ Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo: PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) .await .unwrap(); - let expected = "Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.greptime_value ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]\ - \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series, prometheus_tsdb_head_series.greptime_value [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; + let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\ + \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); } @@ -11176,15 +11676,14 @@ Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo: .unwrap(); let expected = r#" Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, greptime_timestamp:Timestamp(ms)] - Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N] - Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.greptime_value ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N] - Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series, prometheus_tsdb_head_series.greptime_value [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N] - Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#; + Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N] + Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N] + Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64] + PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#; assert_eq!(format!("\n{}", plan.display_indent_schema()), expected); } @@ -11562,6 +12061,7 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom .contains("OR value fields have incompatible types") ); } + #[tokio::test] async fn test_or_with_histogram_quantile_missing_le_column() { let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#; @@ -11898,6 +12398,55 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom } } + #[tokio::test] + async fn test_histogram_only_min_drops_empty_aggregate_group() { + // `min` over native-histogram-only input drops every sample in the group, so the + // NULL-valued aggregate row must be filtered out. Otherwise an outer expression + // like `group()` resurrects the group Prometheus considers unseen. + let state = build_query_engine_state(); + for query in ["min(lh)", "group(min(lh))"] { + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "{query}" + ); + } + } + + #[tokio::test] + async fn test_mixed_min_drops_histogram_only_group() { + // With alternative float/histogram fields, `min by (tag)` keeps float-only groups + // (tag=a from `lf`) and drops histogram-only groups (tag=b from `lh`) instead of + // emitting a NULL-valued row for them. + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(values(&batches, &float_field), vec![2.0]); + } + #[tokio::test] async fn test_mixed_or_can_feed_another_or() { let state = build_query_engine_state(); @@ -12171,6 +12720,306 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom } } + #[tokio::test] + async fn test_native_histogram_sum_and_avg_execute_real_batches() { + for op_name in ["sum", "avg"] { + for incompatible in [false, true] { + let mut second = direct_or_histogram(); + if incompatible { + second.schema = CUSTOM_BUCKETS_SCHEMA; + second.custom_values = vec![1.0]; + } + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = + mixed_aggregate_input(vec![direct_or_histogram(), second]).await; + planner.promql_annotations = Some(collector.clone()); + let histogram_column = planner.ctx.field_columns[1].clone(); + planner.ctx.field_columns = vec![histogram_column.clone()]; + let input = LogicalPlanBuilder::from(input) + .project([col("ts"), col(&histogram_column)]) + .unwrap() + .build() + .unwrap(); + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = + parser::parse(&format!("{op_name}(mixed)")).unwrap() + else { + unreachable!() + }; + let (aggregate_exprs, _) = + planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter(planner.create_empty_values_filter_expr(false).unwrap()) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(infos.is_empty()); + if incompatible { + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + assert!(warnings.iter().any(|warning| { + warning + == &format!( + "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas" + ) + })); + } else { + let histograms = histograms(&batches, &histogram_column); + assert_eq!(histograms.len(), 1); + let expected = if op_name == "sum" { 2.0 } else { 1.0 }; + assert_eq!(histograms[0].count, expected); + assert_eq!(histograms[0].sum, expected); + assert!(warnings.is_empty()); + } + } + } + } + + #[tokio::test] + async fn test_canonical_mixed_count_group_and_count_values_execute() { + let state = build_query_engine_state(); + for (query, expected) in [ + ("count(some_metric)", vec![2.0]), + ("group(some_metric)", vec![1.0]), + (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]), + ] { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + assert!( + plan.schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_sample_count")), + "{query}: {plan:?}" + ); + let value_fields = plan + .schema() + .fields() + .iter() + .filter(|field| { + matches!( + field.data_type(), + ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64 + ) || field.data_type() == &native_histogram_value_type().as_arrow_type() + }) + .collect::>(); + assert_eq!(value_fields.len(), 1, "{query}: {plan:?}"); + assert_ne!( + value_fields[0].data_type(), + &native_histogram_value_type().as_arrow_type(), + "{query}: {plan:?}" + ); + let value_column = value_fields[0].name().clone(); + + let (_, batches) = execute(plan, &state).await; + let mut actual = numeric_values(&batches, &value_column); + actual.sort_by(f64::total_cmp); + assert_eq!(actual, expected, "{query}"); + + if query.starts_with("count_values") { + let mut sample_labels = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("sample") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .map(str::to_string) + }) + .collect::>(); + sample_labels.sort(); + let mut expected_labels = + vec!["2".to_string(), direct_or_histogram().promql_string()]; + expected_labels.sort(); + assert_eq!(sample_labels, expected_labels); + } + } + } + + #[tokio::test] + async fn test_mixed_or_sum_aggregates_each_sample_type() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = + parser::parse("sum(lhs)").unwrap() + else { + unreachable!() + }; + + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_direct_or(false).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts"), col("k")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([ + col(&float_column), + col(&histogram_column), + col("ts"), + col("k"), + ]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_column), vec![1.25]); + let histogram = batches + .iter() + .find_map(|batch| { + let values = batch + .column_by_name(&histogram_column)? + .as_any() + .downcast_ref::()?; + (0..values.len()).find_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .unwrap(); + assert_eq!(histogram.count, 1.0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_direct_or(false).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col(&histogram_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert_eq!( + warnings, + vec![ + "sum: dropped aggregation result containing both float and native histogram samples" + ] + ); + } + + #[tokio::test] + async fn test_mixed_or_sum_drops_incompatible_mixed_group() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = + parser::parse("sum(lhs)").unwrap() + else { + unreachable!() + }; + let mut custom = direct_or_histogram(); + custom.schema = CUSTOM_BUCKETS_SCHEMA; + custom.custom_values = vec![1.0]; + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col(&histogram_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.iter().any(|warning| { + warning + == "sum: dropped aggregation result containing both float and native histogram samples" + })); + } + + #[tokio::test] + async fn test_mixed_or_min_records_only_present_histograms() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = + parser::parse("min(lhs)").unwrap() + else { + unreachable!() + }; + let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms"; + + for (histograms, expected_infos) in [ + (vec![], vec![]), + (vec![direct_or_histogram()], vec![expected_info]), + ] { + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_aggregate_input(histograms).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_ignored_histogram_filter_expr(op, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_column), vec![1.25]); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert_eq!(infos, expected_infos); + } + } + #[tokio::test] async fn test_mixed_or_value_aliases_do_not_replace_labels() { let left = source( diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 030ed27b4e3..e009d80954f 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -40,8 +40,8 @@ use promql::functions::{ NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, - NativeHistogramToString, PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, Round, - StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, + NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime, + Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, }; use prost::Message; use session::context::QueryContextRef; @@ -188,6 +188,7 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { NativeHistogramDelta::scalar_udf(), NativeHistogramDivScalar::scalar_udf(), NativeHistogramDrop::bool_false_udf(String::new(), None), + NativeHistogramDrop::bool_true_udf(String::new(), None), NativeHistogramDrop::float_null_udf(String::new(), None), NativeHistogramEq::scalar_udf(), NativeHistogramFraction::scalar_udf(), @@ -211,6 +212,7 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { NativeHistogramSum::scalar_udf(), NativeHistogramSumOverTime::scalar_udf(), NativeHistogramToString::scalar_udf(), + PromqlFloatToString::scalar_udf(), ] { let _ = session_state.register_udf(Arc::new(udf)); } @@ -335,24 +337,47 @@ mod tests { None, ) .unwrap() + .aggregate( + Vec::::new(), + vec![ + Arc::new(NativeHistogramAggSum::aggregate_udf()) + .call(vec![col("histogram")]) + .alias("sum"), + Arc::new(NativeHistogramAggAvg::aggregate_udf()) + .call(vec![col("histogram")]) + .alias("avg"), + ], + ) + .unwrap() .project(vec![ Expr::ScalarFunction(ScalarFunction { func: Arc::new(NativeHistogramCount::scalar_udf()), - args: vec![col("histogram")], + args: vec![col("sum")], }), Expr::ScalarFunction(ScalarFunction { func: Arc::new(NativeHistogramDrop::bool_false_udf( "ignored annotation".to_string(), None, )), - args: vec![col("histogram")], + args: vec![col("sum")], + }), + Expr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::bool_true_udf( + "ignored annotation".to_string(), + None, + )), + args: vec![col("sum")], }), Expr::ScalarFunction(ScalarFunction { func: Arc::new(NativeHistogramDrop::float_null_udf( "ignored annotation".to_string(), None, )), - args: vec![col("histogram")], + args: vec![col("sum")], + }), + Expr::ScalarFunction(ScalarFunction { + func: Arc::new(PromqlFloatToString::scalar_udf()), + args: vec![lit(2.0)], }), ]) .unwrap() @@ -379,7 +404,11 @@ mod tests { let decoded = decoded.to_string(); assert!(decoded.contains("prom_native_histogram_count")); assert!(decoded.contains("prom_native_histogram_drop_bool")); + assert!(decoded.contains("prom_native_histogram_keep_bool")); assert!(decoded.contains("prom_native_histogram_drop_float")); + assert!(decoded.contains(NativeHistogramAggSum::name())); + assert!(decoded.contains(NativeHistogramAggAvg::name())); + assert!(decoded.contains(PromqlFloatToString::name())); let schema = Arc::new(Schema::new(vec![ Field::new( From 6249623cfb181b0222dfc03224945c625f326f1a Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Wed, 12 Aug 2026 08:15:12 +0000 Subject: [PATCH 057/126] chore: remove iceberg read (#8858) Signed-off-by: shuiyisong --- src/auth/src/lib.rs | 4 ++-- src/auth/src/permission.rs | 2 -- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/src/auth/src/lib.rs b/src/auth/src/lib.rs index d107068940e..d2a1c05b1cd 100644 --- a/src/auth/src/lib.rs +++ b/src/auth/src/lib.rs @@ -31,8 +31,8 @@ pub use common::{ }; pub use permission::{ ALL_ACTIONS, AccessMode, DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, - DefaultPermissionChecker, ICEBERG_READ, INFLUXDB_WRITE, JAEGER_QUERY, LOG_QUERY, LOG_WRITE, - OPENTSDB_WRITE, OTLP_WRITE, PIPELINE_DELETE, PIPELINE_INSERT, PIPELINE_QUERY, PROM_STORE_READ, + DefaultPermissionChecker, INFLUXDB_WRITE, JAEGER_QUERY, LOG_QUERY, LOG_WRITE, OPENTSDB_WRITE, + OTLP_WRITE, PIPELINE_DELETE, PIPELINE_INSERT, PIPELINE_QUERY, PROM_STORE_READ, PROM_STORE_WRITE, PROMQL_QUERY, PermissionAction, PermissionChecker, PermissionReq, PermissionResp, PermissionTableTarget, PermissionTableTargets, SEMANTIC_GRAPH_QUERY, }; diff --git a/src/auth/src/permission.rs b/src/auth/src/permission.rs index 984fe8ea4cf..901b935d6a2 100644 --- a/src/auth/src/permission.rs +++ b/src/auth/src/permission.rs @@ -139,7 +139,6 @@ pub const PROMQL_QUERY: PermissionAction = PermissionAction::read("promql.query" pub const LOG_QUERY: PermissionAction = PermissionAction::read("log.query"); pub const OPENTSDB_WRITE: PermissionAction = PermissionAction::write("opentsdb.write"); pub const INFLUXDB_WRITE: PermissionAction = PermissionAction::write("influxdb.write"); -pub const ICEBERG_READ: PermissionAction = PermissionAction::read("iceberg.read"); pub const PROM_STORE_WRITE: PermissionAction = PermissionAction::write("prom_store.write"); pub const PROM_STORE_READ: PermissionAction = PermissionAction::read("prom_store.read"); pub const OTLP_WRITE: PermissionAction = PermissionAction::write("otlp.write"); @@ -165,7 +164,6 @@ pub const ALL_ACTIONS: &[PermissionAction] = &[ LOG_QUERY, OPENTSDB_WRITE, INFLUXDB_WRITE, - ICEBERG_READ, PROM_STORE_WRITE, PROM_STORE_READ, OTLP_WRITE, From f2ffaef9dfee134fa46d54ba65e6f40b37982cef Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 12 Aug 2026 14:32:59 +0000 Subject: [PATCH 058/126] feat: update to pgwire 0.40.7 (#8860) --- Cargo.lock | 44 +++++++++++++++++++++++++------------------- 1 file changed, 25 insertions(+), 19 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4d0c271882b..5d8a2d1c6c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1216,6 +1216,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64-simd" version = "0.8.0" @@ -1614,7 +1620,7 @@ dependencies = [ "maybe-owned", "rustix 1.0.7", "rustix-linux-procfs", - "windows-sys 0.61.2", + "windows-sys 0.60.2", "winx", ] @@ -5120,7 +5126,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -5682,7 +5688,7 @@ checksum = "94e7099f6313ecacbe1256e8ff9d617b75d1bcb16a6fddef94866d225a01a14a" dependencies = [ "io-lifetimes 2.0.4", "rustix 1.0.7", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -6650,7 +6656,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.57.0", + "windows-core 0.61.2", ] [[package]] @@ -7078,7 +7084,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20fd6de4ccfcc187e38bc21cfa543cb5a302cb86a8b114eb7f0bf0dc9f8ac00f" dependencies = [ "io-lifetimes 3.0.1", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -7116,7 +7122,7 @@ checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -9087,7 +9093,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -9281,7 +9287,7 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77e878c846a8abae00dd069496dbe8751b16ac1c3d6bd2a7283a938e8228f90d" dependencies = [ - "proc-macro-crate 3.3.0", + "proc-macro-crate 1.3.1", "proc-macro2", "quote", "syn 2.0.117", @@ -10439,13 +10445,13 @@ dependencies = [ [[package]] name = "pgwire" -version = "0.40.1" +version = "0.40.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f690fba3b77370e278b606b75c6857d3f33f3e570781a75ef7516e76e44c0b19" +checksum = "8265901ede50d0879fe401c6fe282e7e4ff83ae10a48c1f8e4f89b92f7d5f604" dependencies = [ "async-trait", "aws-lc-rs", - "base64 0.22.1", + "base64 0.23.1", "bytes", "chrono", "derive-new", @@ -11762,7 +11768,7 @@ dependencies = [ "once_cell", "socket2 0.5.10", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -12711,7 +12717,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.4.15", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -12724,7 +12730,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.9.4", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -12815,7 +12821,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -13728,7 +13734,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -14114,7 +14120,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -14724,7 +14730,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.0.7", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -16916,7 +16922,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f3fd376f71958b862e7afb20cfe5a22830e1963462f3a17f49d82a6c1d1f42d" dependencies = [ "bitflags 2.12.1", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] From 546625c45af629b75f0ba28239a3a183e1183e19 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 13 Aug 2026 02:31:56 +0000 Subject: [PATCH 059/126] feat: embedded convention pack for the entity graph (prom/k8s, gen_ai naming) (#8854) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: embed the derivation conventions as data and adopt gen_ai entity naming Move the co-declared edge vocabulary, the agent-edge vocabulary and the virtual-destination candidates from Rust consts into an embedded conventions.yaml (include_str!), parsed once behind a LazyLock and validated against the entity-type grammar and the closed rel_type set; a broken file propagates as a plan error instead of panicking. The agent vocabulary entity types follow the GenAI semantic-convention namespace as written: gen_ai.agent / gen_ai.model / gen_ai.tool. Signed-off-by: Dennis Zhuang * feat: drop the tag requirement for entity identity columns Entity declarations no longer require id columns to be tag/primary-key columns; only column existence is validated. Trace pipelines flatten the identifying attributes (span_attributes.gen_ai.agent.id, ...) into field columns, so the tag rule locked real trace tables out of declaring entities while buying no correctness — the read-time derivation works on any column. Signed-off-by: Dennis Zhuang * feat: implicit declarations for well-known prometheus info metrics Tables stamped signal_type=metric + source=prometheus whose name matches the conventions.yaml whitelist (kube_pod_info, kube_node_info, kube_pod_owner, target_info) get implicit entity declarations: k8s.pod / k8s.node / k8s.workload with name-based identity and target_info's service / service.instance with the remaining tags as the descriptive snapshot. The existing co-declared vocabulary then derives runs_on and part_of from the same rows, so no new edge branch is needed. Explicit declarations of a type always suppress the implicit one, and the metric engine's physical table is excluded (it aggregates every logical table's columns). Signed-off-by: Dennis Zhuang * test: cover the prometheus conventions in sqlness and compact the graph cases Add the whitelisted-info-metric scenario (kube_pod_info, kube_pod_owner, target_info deriving runs_on / part_of, a non-whitelisted metric contributing nothing), fold the single-table calls, cross-table pairing and virtual-node cases into one trace scenario (they exercise the same union-before-join path), merge the two declaring-metric-table cases, and reuse one rename probe for both reserved names. Signed-off-by: Dennis Zhuang * fix: reject entity id columns without a stable string form Review follow-ups: the DDL check now validates against the schema and rejects binary-backed and nested types for identity columns (the derivation renders ids via CAST to Utf8, so the failure used to surface only when the graph was scanned); the agent sqlness case keeps its identity columns as fields to cover the relaxed tag rule end to end; stale tag-rule comments and a dangling const reference are cleaned up. Signed-off-by: Dennis Zhuang * fix: type-check every entity column role, not only ids The registry renders scope and descriptive values through the same CAST-to-string path as ids, so a binary-backed column in any role fails at scan time; the DDL check is now role-independent (and simpler). Signed-off-by: Dennis Zhuang * refactor: name the code-anchored vocabulary constants Entity types and edge attributes the derivation code itself anchors on (service, gen_ai.agent, calls, trace/attribute provenance) become constants in the conventions module; the rest of the vocabulary stays YAML-only data. ImplicitEntity is renamed PromImplicitEntity, and the implicit-declaration path logs each skip of a whitelisted info metric (wrong stamps, suppressed by an explicit declaration, missing id column) so a missing graph entity is diagnosable. Signed-off-by: Dennis Zhuang * refactor: single-source the graph constants The graph tables' column names move to common-catalog (the schemas catalog exposes and the plans operator builds must match column by column), and the conventions module now carries the complete built-in vocabulary — entity types, rel_types, provenances and connection types — with the embedded YAML validated by membership against it, so an edit drifting outside the vocabulary fails the conventions test instead of deriving nothing. Signed-off-by: Dennis Zhuang * fix: treat empty identity components as absent kube-state-metrics emits empty-string labels an entity id must not be built from: an unscheduled pod's node and an owner-less pod's owner_kind / owner_name. Standard Prometheus drops empty labels (they arrive as NULL and the existing predicate handles them), but other remote-write agents may keep them, which produced ghost entities with empty ids and false runs_on / part_of edges. Every identity predicate (registry, co-declared edges, span endpoints) now requires non-NULL and non-empty components through one shared helper. Signed-off-by: Dennis Zhuang * refactor: tighten the conventions DSL semantics Rename the co-declaration rule lists to what they are (co_declared_edges / trace_co_declared_edges — derivation rules, not a relation vocabulary), stop overstating the GenAI entity types (Greptime types derived from GenAI attributes; OTel defines no model/tool entities), move target_info's descriptive snapshot to service.instance (the remaining labels are the target's resource attributes, and instances would write conflicting snapshots onto the logical service), and extend the descriptor whitelist with the stable KSM sources: container info metrics (closing the k8s.pod contains k8s.container rule), kube_service_info (new k8s.service entity type) and the fuller descriptive label sets. Signed-off-by: Dennis Zhuang * fix: guard entity column types on ALTER as well ALTER MODIFY COLUMN could change a declared entity column to a type without a stable string form, deferring the failure to graph scan time; verify_alter now checks the post-alter schema. Dropping a declared column stays allowed — the read-time derivation skips the stale declaration, and semantic options cannot be altered off yet. Signed-off-by: Dennis Zhuang * feat: bridge traces and kube-state-metrics on the pod UID Trace-v1 tables now get implicit declarations from their flattened resource attributes (otlp_trace_entities in conventions.yaml): the service identity — replacing the hardcoded fallback — plus service.instance and k8s.pod, each applied only when its columns exist. A new co-declared rule derives service.instance runs_on k8s.pod, and the whitelisted kube-state-metrics pod identity switches from namespace+pod names to the UID, so the trace-side pod and every KSM descriptor land on one entity while names stay descriptive. This also removes pod identity from the multi-cluster same-name collision. The conventions rejection tests were passing for the wrong reason (a half-renamed fixture key failed deserialization before reaching any validation rule); they now assert the specific error each case targets. Sqlness covers the UID merge across descriptor tables, pod-contains- container, the k8s.service node, and the empty-uid/empty-node rows deriving nothing. Signed-off-by: Dennis Zhuang * test: cover the OTLP-to-graph chain end to end One real OTLP export must come out of semantic_relationships as the zero-configuration chain: service calls service, instance part_of service, instance runs_on pod (bridged by k8s.pod.uid). Resources without service.instance.id or k8s.pod.uid derive nothing extra. Signed-off-by: Dennis Zhuang * fix: identify k8s.service by UID Same reasoning as pods: a recreated same-name service must not merge into the old entity and same-named services across clusters must not collide; kube_service_info carries a stable uid and nothing joins on the service's name. Also drop a stale tag-rule mention from the option validation docs. Signed-off-by: Dennis Zhuang * chore: cut duplicated test coverage and redundant comments The trace service-fallback test collapsed into the resource-entities test (same synthesis path since the fallback moved to YAML; only the invalid-explicit-no-fallback case was distinct), role-duplicate and subsumed DDL cases are gone, the embedded-conventions test is just the parse (its assertions were decorative), and the YAML section comments no longer restate the struct docs. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- Cargo.lock | 2 + Cargo.toml | 1 + .../src/system_schema/semantic_graph.rs | 77 ++- src/common/catalog/src/consts.rs | 29 + src/frontend/src/instance/entity_graph.rs | 365 +++++++++-- src/operator/Cargo.toml | 2 + src/operator/src/error.rs | 8 - src/operator/src/statement/ddl.rs | 180 ++++-- src/operator/src/statement/semantic_graph.rs | 171 +++-- .../statement/semantic_graph/conventions.rs | 291 +++++++++ .../statement/semantic_graph/conventions.yaml | 115 ++++ .../statement/semantic_graph/relationships.rs | 323 +++++----- src/table/src/requests/semantic.rs | 6 +- tests-integration/tests/http.rs | 81 +++ .../common/system/semantic_graph.result | 586 ++++++++++-------- .../common/system/semantic_graph.sql | 346 ++++++----- 16 files changed, 1803 insertions(+), 780 deletions(-) create mode 100644 src/operator/src/statement/semantic_graph/conventions.rs create mode 100644 src/operator/src/statement/semantic_graph/conventions.yaml diff --git a/Cargo.lock b/Cargo.lock index 5d8a2d1c6c0..dcd23451407 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9936,7 +9936,9 @@ dependencies = [ "prost-types 0.14.1", "query", "regex", + "serde", "serde_json", + "serde_yaml_ng", "session", "snafu 0.8.6", "sql", diff --git a/Cargo.toml b/Cargo.toml index a138a880ec6..0e6da7888d3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -232,6 +232,7 @@ sea-query = "0.32" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } serde_with = "3" +serde_yaml_ng = "0.10" simd-json = "0.15" similar-asserts = "1.6.0" smallvec = { version = "1", features = ["serde"] } diff --git a/src/catalog/src/system_schema/semantic_graph.rs b/src/catalog/src/system_schema/semantic_graph.rs index 028ce79b906..b25dbb03289 100644 --- a/src/catalog/src/system_schema/semantic_graph.rs +++ b/src/catalog/src/system_schema/semantic_graph.rs @@ -35,9 +35,14 @@ use std::sync::{Arc, LazyLock, Weak}; use common_catalog::consts::{ - DEFAULT_PRIVATE_SCHEMA_NAME, SEMANTIC_ENTITIES_TABLE_ID, + CONFIDENCE_COLUMN, DEFAULT_PRIVATE_SCHEMA_NAME, DST_ID_COLUMN, DST_TYPE_COLUMN, + DURATION_COUNT_COLUMN, DURATION_SUM_COLUMN, EDGE_ATTRIBUTES_COLUMN, ENTITY_DESCRIPTIVE_COLUMN, + ENTITY_ID_ATTRS_COLUMN, ENTITY_ID_COLUMN, ENTITY_SCOPE_COLUMN, ENTITY_TYPE_COLUMN, + ERROR_COUNT_COLUMN, FRESH_UNTIL_COLUMN, OBSERVED_AT_COLUMN, PROVENANCE_COLUMN, REL_TYPE_COLUMN, + REQUEST_COUNT_COLUMN, SEMANTIC_ENTITIES_TABLE_ID, SEMANTIC_ENTITIES_TABLE_NAME as SEMANTIC_ENTITIES, SEMANTIC_RELATIONSHIPS_TABLE_ID, - SEMANTIC_RELATIONSHIPS_TABLE_NAME as SEMANTIC_RELATIONSHIPS, + SEMANTIC_RELATIONSHIPS_TABLE_NAME as SEMANTIC_RELATIONSHIPS, SOURCE_TABLES_COLUMN, + SRC_ID_COLUMN, SRC_TYPE_COLUMN, WINDOW_END_COLUMN, WINDOW_START_COLUMN, }; use common_error::ext::BoxedError; use common_recordbatch::adapter::AsyncRecordBatchStreamAdapter; @@ -190,16 +195,16 @@ fn json() -> ConcreteDataType { /// - `source_tables` — JSON array of the telemetry tables that contributed this entity. static ENTITIES_SCHEMA: LazyLock = LazyLock::new(|| { Arc::new(Schema::new(vec![ - ColumnSchema::new("observed_at", ts(), false).with_time_index(true), - ColumnSchema::new("window_start", ts(), true), - ColumnSchema::new("window_end", ts(), true), - ColumnSchema::new("fresh_until", ts(), true), - ColumnSchema::new("entity_type", string(), false), - ColumnSchema::new("entity_id", string(), false), - ColumnSchema::new("entity_id_attrs", json(), true), - ColumnSchema::new("scope", string(), true), - ColumnSchema::new("descriptive", json(), true), - ColumnSchema::new("source_tables", json(), true), + ColumnSchema::new(OBSERVED_AT_COLUMN, ts(), false).with_time_index(true), + ColumnSchema::new(WINDOW_START_COLUMN, ts(), true), + ColumnSchema::new(WINDOW_END_COLUMN, ts(), true), + ColumnSchema::new(FRESH_UNTIL_COLUMN, ts(), true), + ColumnSchema::new(ENTITY_TYPE_COLUMN, string(), false), + ColumnSchema::new(ENTITY_ID_COLUMN, string(), false), + ColumnSchema::new(ENTITY_ID_ATTRS_COLUMN, json(), true), + ColumnSchema::new(ENTITY_SCOPE_COLUMN, string(), true), + ColumnSchema::new(ENTITY_DESCRIPTIVE_COLUMN, json(), true), + ColumnSchema::new(SOURCE_TABLES_COLUMN, json(), true), ])) }); @@ -232,22 +237,38 @@ static ENTITIES_SCHEMA: LazyLock = LazyLock::new(|| { /// `db.system`, `peer.service`. static RELATIONSHIPS_SCHEMA: LazyLock = LazyLock::new(|| { Arc::new(Schema::new(vec![ - ColumnSchema::new("observed_at", ts(), false).with_time_index(true), - ColumnSchema::new("window_start", ts(), true), - ColumnSchema::new("window_end", ts(), true), - ColumnSchema::new("fresh_until", ts(), true), - ColumnSchema::new("src_type", string(), false), - ColumnSchema::new("src_id", string(), false), - ColumnSchema::new("dst_type", string(), false), - ColumnSchema::new("dst_id", string(), false), - ColumnSchema::new("rel_type", string(), false), - ColumnSchema::new("provenance", string(), false), - ColumnSchema::new("confidence", ConcreteDataType::float64_datatype(), true), - ColumnSchema::new("request_count", ConcreteDataType::int64_datatype(), true), - ColumnSchema::new("error_count", ConcreteDataType::int64_datatype(), true), - ColumnSchema::new("duration_sum", ConcreteDataType::float64_datatype(), true), - ColumnSchema::new("duration_count", ConcreteDataType::int64_datatype(), true), - ColumnSchema::new("attributes", json(), true), + ColumnSchema::new(OBSERVED_AT_COLUMN, ts(), false).with_time_index(true), + ColumnSchema::new(WINDOW_START_COLUMN, ts(), true), + ColumnSchema::new(WINDOW_END_COLUMN, ts(), true), + ColumnSchema::new(FRESH_UNTIL_COLUMN, ts(), true), + ColumnSchema::new(SRC_TYPE_COLUMN, string(), false), + ColumnSchema::new(SRC_ID_COLUMN, string(), false), + ColumnSchema::new(DST_TYPE_COLUMN, string(), false), + ColumnSchema::new(DST_ID_COLUMN, string(), false), + ColumnSchema::new(REL_TYPE_COLUMN, string(), false), + ColumnSchema::new(PROVENANCE_COLUMN, string(), false), + ColumnSchema::new( + CONFIDENCE_COLUMN, + ConcreteDataType::float64_datatype(), + true, + ), + ColumnSchema::new( + REQUEST_COUNT_COLUMN, + ConcreteDataType::int64_datatype(), + true, + ), + ColumnSchema::new(ERROR_COUNT_COLUMN, ConcreteDataType::int64_datatype(), true), + ColumnSchema::new( + DURATION_SUM_COLUMN, + ConcreteDataType::float64_datatype(), + true, + ), + ColumnSchema::new( + DURATION_COUNT_COLUMN, + ConcreteDataType::int64_datatype(), + true, + ), + ColumnSchema::new(EDGE_ATTRIBUTES_COLUMN, json(), true), ])) }); diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index 6053d5a67c8..b976cbc256e 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -219,4 +219,33 @@ pub const SEMANTIC_RELATIONSHIPS_TABLE_NAME: &str = "semantic_relationships"; /// Physical table of hand-declared edges, unioned into the computed /// `semantic_relationships`; see [`is_ddl_reserved_table`] for its lifecycle. pub const SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME: &str = "semantic_relationships_declared"; + +// Column names of the graph tables: `catalog` exposes these schemas and the +// read-time plans in `operator` must project exactly them. +pub const OBSERVED_AT_COLUMN: &str = "observed_at"; +pub const WINDOW_START_COLUMN: &str = "window_start"; +pub const WINDOW_END_COLUMN: &str = "window_end"; +pub const FRESH_UNTIL_COLUMN: &str = "fresh_until"; +pub const ENTITY_TYPE_COLUMN: &str = "entity_type"; +pub const ENTITY_ID_COLUMN: &str = "entity_id"; +pub const ENTITY_ID_ATTRS_COLUMN: &str = "entity_id_attrs"; +pub const ENTITY_SCOPE_COLUMN: &str = "scope"; +pub const ENTITY_DESCRIPTIVE_COLUMN: &str = "descriptive"; +pub const SOURCE_TABLES_COLUMN: &str = "source_tables"; +pub const SRC_TYPE_COLUMN: &str = "src_type"; +pub const SRC_ID_COLUMN: &str = "src_id"; +pub const DST_TYPE_COLUMN: &str = "dst_type"; +pub const DST_ID_COLUMN: &str = "dst_id"; +pub const REL_TYPE_COLUMN: &str = "rel_type"; +pub const PROVENANCE_COLUMN: &str = "provenance"; +pub const CONFIDENCE_COLUMN: &str = "confidence"; +pub const REQUEST_COUNT_COLUMN: &str = "request_count"; +pub const ERROR_COUNT_COLUMN: &str = "error_count"; +pub const DURATION_SUM_COLUMN: &str = "duration_sum"; +pub const DURATION_COUNT_COLUMN: &str = "duration_count"; +pub const EDGE_ATTRIBUTES_COLUMN: &str = "attributes"; +// Declared-edge table only. +pub const VALID_FROM_COLUMN: &str = "valid_from"; +pub const VALID_UNTIL_COLUMN: &str = "valid_until"; +pub const GENERATION_ID_COLUMN: &str = "generation_id"; // ---- End of entity relationship graph tables ---- diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs index 70c57846fd0..14f18ed1b79 100644 --- a/src/frontend/src/instance/entity_graph.rs +++ b/src/frontend/src/instance/entity_graph.rs @@ -34,8 +34,8 @@ use auth::{ use catalog::CatalogManager; use catalog::system_schema::semantic_graph::EntityGraphProvider; use common_catalog::consts::{ - DEFAULT_PRIVATE_SCHEMA_NAME, DEFAULT_SCHEMA_NAME, INFORMATION_SCHEMA_NAME, PG_CATALOG_NAME, - SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SERVICE_NAME_COLUMN, + DEFAULT_PRIVATE_SCHEMA_NAME, DEFAULT_SCHEMA_NAME, INFORMATION_SCHEMA_NAME, OBSERVED_AT_COLUMN, + PG_CATALOG_NAME, SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, }; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; @@ -47,9 +47,10 @@ use datafusion::dataframe::DataFrame; use datafusion_expr::LogicalPlan; use futures::TryStreamExt; use operator::statement::semantic_graph::{ - CallsSource, CoDeclaredSource, DeclaredSource, EntityDeclaration, GraphQueryWindow, - OBSERVED_AT_COLUMN, RegistrySource, RelationshipSources, build_registry_plan, - build_relationships_plan, declared_relationships_schema_matches, + CallsSource, CoDeclaredSource, Conventions, DeclaredSource, ENTITY_TYPE_GEN_AI_AGENT, + ENTITY_TYPE_SERVICE, EntityDeclaration, GraphQueryWindow, ImplicitEntity, RegistrySource, + RelationshipSources, build_registry_plan, build_relationships_plan, conventions, + declared_relationships_schema_matches, }; use query::QueryEngineRef; use session::context::{QueryContext, QueryContextBuilder, QueryContextRef}; @@ -59,7 +60,8 @@ use table::TableRef; use table::metadata::TableInfo; use table::predicate::{TimeRangeExtraction, extract_time_range_strict}; use table::requests::{ - EntityRole, is_trace_v1_table, parse_entity_columns, parse_entity_option_key, + EntityRole, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, SOURCE_PROMETHEUS, + is_trace_v1_table, parse_entity_columns, parse_entity_option_key, }; use crate::error; @@ -183,42 +185,127 @@ impl EntityGraphProviderImpl { .collect() } - /// All entity declarations of one table. Trace-v1 tables created before the - /// ingest-side auto-stamp carry no `entity.service.id` option; synthesize it - /// (their schema is fixed), so their `calls` edges have endpoint entities. - /// A table with an explicit service declaration never gets the synthesized - /// one — not even when the explicit declaration is invalid and skipped: - /// silently falling back would change entity identity behind the user's back. - fn declarations_for(table_info: &TableInfo) -> Vec { + /// All entity declarations of one table: the explicit options plus the + /// zero-configuration conventions (`otlp_trace_entities` for trace-v1 + /// tables — including the `service` identity of tables created before the + /// ingest-side auto-stamp — and the Prometheus descriptor whitelist). An + /// explicit declaration of a type always suppresses the implicit one, even + /// when the explicit declaration is invalid and skipped: silently falling + /// back would change entity identity behind the user's back. + fn declarations_for( + table_info: &TableInfo, + conventions: &Conventions, + ) -> Vec { let mut declarations = Self::parse_declarations(table_info); - let has_explicit_service = table_info.meta.options.extra_options.keys().any(|key| { + if is_trace_v1_table(table_info) { + Self::extend_with_implicit_entities( + table_info, + &conventions.otlp_trace_entities, + &mut declarations, + ); + } + Self::extend_with_prometheus_conventions(table_info, conventions, &mut declarations); + declarations + } + + /// Whether the table carries an explicit `entity..id` option. + fn explicitly_declares(table_info: &TableInfo, entity_type: &str) -> bool { + table_info.meta.options.extra_options.keys().any(|key| { parse_entity_option_key(key) - .is_some_and(|(ty, role)| ty == "service" && role == EntityRole::Id) - }); - if is_trace_v1_table(table_info) - && !has_explicit_service - && table_info - .meta - .schema - .column_schema_by_name(SERVICE_NAME_COLUMN) - .is_some() - && let Some(time_index) = table_info - .meta - .schema - .timestamp_column() - .map(|c| c.name.clone()) + .is_some_and(|(ty, role)| ty == entity_type && role == EntityRole::Id) + }) + } + + /// Implicit declarations of the well-known Prometheus entity-descriptor + /// metrics (the `prometheus_info_metrics` whitelist of `conventions.yaml`), + /// gated on the ingest-stamped `signal_type=metric` + `source=prometheus` + /// options and keyed by table name. The metric engine's physical table + /// aggregates every logical table's columns and must not contribute a + /// duplicate source. + fn extend_with_prometheus_conventions( + table_info: &TableInfo, + conventions: &Conventions, + declarations: &mut Vec, + ) { + let Some(implicit_entities) = conventions.prometheus_info_metrics.get(&table_info.name) + else { + return; + }; + let options = &table_info.meta.options.extra_options; + if options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str) != Some(SIGNAL_TYPE_METRIC) + || options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(SOURCE_PROMETHEUS) + || table_info.is_physical_table() { + debug!( + "Table `{}` matches the info-metric whitelist but is not a prometheus logical \ + metric table; skipping its implicit declarations", + table_info.name + ); + return; + } + Self::extend_with_implicit_entities(table_info, implicit_entities, declarations); + } + + /// Synthesizes the applicable subset of `entities` on `table_info`: + /// explicit declarations win, every id column must exist (no guessing), + /// descriptive columns are filtered to those present. + fn extend_with_implicit_entities( + table_info: &TableInfo, + entities: &[ImplicitEntity], + declarations: &mut Vec, + ) { + let schema = &table_info.meta.schema; + let Some(time_index) = schema.timestamp_column().map(|c| c.name.clone()) else { + debug!( + "Table `{}` has no time index; skipping its implicit declarations", + table_info.name + ); + return; + }; + for implicit in entities { + if Self::explicitly_declares(table_info, &implicit.entity) { + debug!( + "Table `{}` explicitly declares `{}`; the implicit declaration is suppressed", + table_info.name, implicit.entity + ); + continue; + } + if let Some(missing) = implicit + .id + .iter() + .find(|c| schema.column_schema_by_name(c).is_none()) + { + debug!( + "Table `{}` lacks the id column `{}`; skipping the implicit `{}` declaration", + table_info.name, missing, implicit.entity + ); + continue; + } + let descriptive_columns = if implicit.descriptive_rest { + table_info + .meta + .row_key_column_names() + .filter(|c| !implicit.id.contains(c)) + .cloned() + .collect() + } else { + implicit + .descriptive + .iter() + .filter(|c| schema.column_schema_by_name(c).is_some()) + .cloned() + .collect() + }; declarations.push(EntityDeclaration { schema: table_info.schema_name.clone(), table: table_info.name.clone(), - time_index, - entity_type: "service".to_string(), - id_columns: vec![SERVICE_NAME_COLUMN.to_string()], - descriptive_columns: vec![], + time_index: time_index.clone(), + entity_type: implicit.entity.clone(), + id_columns: implicit.id.clone(), + descriptive_columns, scope_columns: vec![], }); } - declarations } /// Enumerates entity declarations and trace tables across a catalog. Trace @@ -233,6 +320,10 @@ impl EntityGraphProviderImpl { let Some(catalog_manager) = self.catalog_manager.upgrade() else { return Ok((vec![], vec![])); }; + let conventions = conventions() + .map_err(datafusion::error::DataFusionError::Internal) + .context(error::DataFusionSnafu) + .map_err(BoxedError::new)?; // A target-blind checker (e.g. the default mode-based one) answers the // same for every table: ask once up front instead of per table. @@ -273,7 +364,7 @@ impl EntityGraphProviderImpl { let mut tables = catalog_manager.tables(catalog, &schema, query_ctx); while let Some(table) = tables.try_next().await.map_err(BoxedError::new)? { let table_info = table.table_info(); - let table_declarations = Self::declarations_for(&table_info); + let table_declarations = Self::declarations_for(&table_info, conventions); let is_trace = is_trace_v1_table(&table_info); // Authorize only tables that would contribute rows. if per_table_auth @@ -301,8 +392,8 @@ impl EntityGraphProviderImpl { .find(|d| d.entity_type == entity_type) .cloned() }; - let service = find("service"); - let agent = find("agent"); + let service = find(ENTITY_TYPE_SERVICE); + let agent = find(ENTITY_TYPE_GEN_AI_AGENT); if service.is_none() { // No usable service identity: the table cannot // contribute service-calls edges (see declarations_for). @@ -541,6 +632,7 @@ mod tests { use common_catalog::consts::{DEFAULT_CATALOG_NAME, MITO_ENGINE}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, SchemaBuilder}; + use store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY; use table::metadata::{TableInfoBuilder, TableMeta, TableType}; use table::requests::{TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1, TableOptions}; @@ -564,6 +656,15 @@ mod tests { } fn assemble_table_info(column_schemas: Vec, extra: &[(&str, &str)]) -> TableInfo { + named_table_info("t1", column_schemas, vec![], extra) + } + + fn named_table_info( + name: &str, + column_schemas: Vec, + primary_key_indices: Vec, + extra: &[(&str, &str)], + ) -> TableInfo { let schema = Arc::new( SchemaBuilder::try_from_columns(column_schemas) .unwrap() @@ -579,7 +680,7 @@ mod tests { }; let meta = TableMeta { schema, - primary_key_indices: vec![], + primary_key_indices, value_indices: vec![], engine: MITO_ENGINE.to_string(), next_column_id: 1, @@ -591,7 +692,7 @@ mod tests { }; TableInfoBuilder::default() .table_id(1) - .name("t1") + .name(name) .catalog_name(DEFAULT_CATALOG_NAME) .schema_name(DEFAULT_SCHEMA_NAME) .table_version(0) @@ -601,6 +702,24 @@ mod tests { .unwrap() } + /// A metric-engine-logical-table shape: every label column is a tag. + fn prom_table_info(name: &str, tags: &[&str], extra: &[(&str, &str)]) -> TableInfo { + let mut column_schemas = vec![ + ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ]; + column_schemas.extend( + tags.iter() + .map(|c| ColumnSchema::new(*c, ConcreteDataType::string_datatype(), true)), + ); + let primary_key_indices = (1..=tags.len()).collect(); + named_table_info(name, column_schemas, primary_key_indices, extra) + } + #[test] fn future_lower_only_window_matches_nothing() { let future_ms = common_time::util::current_time_millis() + 86_400_000; @@ -634,7 +753,7 @@ mod tests { ("greptime.semantic.entity.host.id", "gone"), ], ); - let declarations = EntityGraphProviderImpl::declarations_for(&info); + let declarations = EntityGraphProviderImpl::declarations_for(&info, conventions().unwrap()); assert_eq!(declarations.len(), 1); assert_eq!(declarations[0].entity_type, "service"); @@ -645,45 +764,161 @@ mod tests { ("greptime.semantic.entity.service.descriptive", "gone"), ], ); - assert!(EntityGraphProviderImpl::declarations_for(&info).is_empty()); + assert!( + EntityGraphProviderImpl::declarations_for(&info, conventions().unwrap()).is_empty() + ); + } + + const PROM_STAMPS: &[(&str, &str)] = &[ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, SOURCE_PROMETHEUS), + ]; + + fn sorted_declarations(info: &TableInfo) -> Vec { + let mut declarations = + EntityGraphProviderImpl::declarations_for(info, conventions().unwrap()); + declarations.sort_by(|a, b| a.entity_type.cmp(&b.entity_type)); + declarations } #[test] - fn trace_v1_table_gets_implicit_service_declaration() { - let info = table_info( - &["service_name"], + fn prometheus_info_metric_gets_implicit_declarations() { + let info = prom_table_info( + "kube_pod_info", + &["namespace", "pod", "uid", "node", "job", "instance"], + PROM_STAMPS, + ); + let declarations = sorted_declarations(&info); + assert_eq!(declarations.len(), 2); + assert_eq!(declarations[0].entity_type, "k8s.node"); + assert_eq!(declarations[0].id_columns, vec!["node"]); + assert_eq!(declarations[1].entity_type, "k8s.pod"); + assert_eq!(declarations[1].id_columns, vec!["uid"]); + // host_ip/pod_ip/created_by_* are absent from this table; descriptive + // shrinks to the present columns. + assert_eq!( + declarations[1].descriptive_columns, + vec!["namespace", "pod", "node"] + ); + assert_eq!(declarations[1].time_index, "greptime_timestamp"); + } + + #[test] + fn trace_table_gets_implicit_resource_entities() { + let full = table_info( + &[ + "service_name", + "resource_attributes.service.instance.id", + "resource_attributes.k8s.pod.uid", + "resource_attributes.k8s.pod.name", + ], &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], ); - let declarations = EntityGraphProviderImpl::declarations_for(&info); - assert_eq!(declarations.len(), 1); - let decl = &declarations[0]; - assert_eq!(decl.entity_type, "service"); - assert_eq!(decl.id_columns, vec![SERVICE_NAME_COLUMN.to_string()]); - assert_eq!(decl.time_index, "ts"); - - // An explicit service declaration wins; no duplicate is synthesized. - let info = table_info( - &["service_name"], - &[ - (TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1), - ("greptime.semantic.entity.service.id", "service_name"), - ], + let declarations = sorted_declarations(&full); + let types: Vec<&str> = declarations + .iter() + .map(|d| d.entity_type.as_str()) + .collect(); + assert_eq!(types, vec!["k8s.pod", "service", "service.instance"]); + assert_eq!( + declarations[0].id_columns, + vec!["resource_attributes.k8s.pod.uid"] + ); + assert_eq!( + declarations[0].descriptive_columns, + vec!["resource_attributes.k8s.pod.name"] + ); + assert_eq!( + declarations[2].id_columns, + vec!["service_name", "resource_attributes.service.instance.id"] ); - assert_eq!(EntityGraphProviderImpl::declarations_for(&info).len(), 1); - // A trace-model table without the fixed column synthesizes nothing. - let info = table_info(&[], &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)]); - assert!(EntityGraphProviderImpl::declarations_for(&info).is_empty()); + // Missing uid column: no pod entity synthesized, no name-based guess. + let no_uid = table_info( + &["service_name", "resource_attributes.k8s.pod.name"], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + let types: Vec = sorted_declarations(&no_uid) + .into_iter() + .map(|d| d.entity_type) + .collect(); + assert_eq!(types, vec!["service"]); - // An explicit but invalid declaration must not silently fall back to - // the synthesized one: identity must not change behind the user's back. - let info = table_info( + // An explicit declaration suppresses the implicit one even when it is + // invalid and skipped: identity must not change behind the user's back. + let invalid_explicit = table_info( &["service_name"], &[ (TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1), ("greptime.semantic.entity.service.id", "gone"), ], ); - assert!(EntityGraphProviderImpl::declarations_for(&info).is_empty()); + assert!(sorted_declarations(&invalid_explicit).is_empty()); + } + + #[test] + fn prometheus_implicit_declarations_are_gated() { + let labels: &[&str] = &["namespace", "pod", "node"]; + // A non-Prometheus source. + assert!( + sorted_declarations(&prom_table_info( + "kube_pod_info", + labels, + &[ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, "opentelemetry"), + ], + )) + .is_empty() + ); + // Not a whitelisted metric name. + assert!( + sorted_declarations(&prom_table_info("http_requests_total", labels, PROM_STAMPS)) + .is_empty() + ); + let mut stamps = PROM_STAMPS.to_vec(); + stamps.push((PHYSICAL_TABLE_METADATA_KEY, "true")); + assert!(sorted_declarations(&prom_table_info("kube_pod_info", labels, &stamps)).is_empty()); + + // A missing id column (uid) drops that entity, not the whole table. + let info = prom_table_info("kube_pod_info", &["namespace", "pod", "node"], PROM_STAMPS); + let declarations = sorted_declarations(&info); + assert_eq!(declarations.len(), 1); + assert_eq!(declarations[0].entity_type, "k8s.node"); + } + + #[test] + fn explicit_declaration_suppresses_the_implicit_one() { + let mut stamps = PROM_STAMPS.to_vec(); + stamps.push(("greptime.semantic.entity.k8s.pod.id", "pod")); + let info = prom_table_info("kube_pod_info", &["namespace", "pod", "node"], &stamps); + let declarations = sorted_declarations(&info); + assert_eq!(declarations.len(), 2); + assert_eq!(declarations[0].entity_type, "k8s.node"); + assert_eq!(declarations[1].entity_type, "k8s.pod"); + // The explicit identity wins over the conventional [namespace, pod]. + assert_eq!(declarations[1].id_columns, vec!["pod"]); + } + + #[test] + fn target_info_descriptive_rest_covers_remaining_tags() { + let info = prom_table_info( + "target_info", + &["job", "instance", "k8s_cluster_name", "service_version"], + PROM_STAMPS, + ); + let declarations = sorted_declarations(&info); + assert_eq!(declarations.len(), 2); + assert_eq!(declarations[0].entity_type, "service"); + assert_eq!(declarations[0].id_columns, vec!["job"]); + assert!(declarations[0].descriptive_columns.is_empty()); + // The remaining labels are the target's resource attributes: they + // describe the instance, not the logical service. + assert_eq!(declarations[1].entity_type, "service.instance"); + assert_eq!(declarations[1].id_columns, vec!["job", "instance"]); + assert_eq!( + declarations[1].descriptive_columns, + vec!["k8s_cluster_name", "service_version"] + ); } } diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index 087ea461f10..c9258df2c0f 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -69,7 +69,9 @@ prost.workspace = true prost-types = { workspace = true, optional = true } query.workspace = true regex.workspace = true +serde.workspace = true serde_json.workspace = true +serde_yaml_ng.workspace = true session.workspace = true snafu.workspace = true sql.workspace = true diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index f1e94b2c616..0ed3c74bbed 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -235,13 +235,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Invalid entity semantic option, reason: {}", reason))] - InvalidEntitySemanticOption { - reason: String, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Table not found: {}", table_name))] TableNotFound { table_name: String }, @@ -957,7 +950,6 @@ impl ErrorExt for Error { | Error::InvalidConfigValue { .. } | Error::InvalidInsertRequest { .. } | Error::InvalidDeleteRequest { .. } - | Error::InvalidEntitySemanticOption { .. } | Error::IllegalPrimaryKeysDef { .. } | Error::SchemaNotFound { .. } | Error::SchemaExists { .. } diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 3f637ebebcf..256b849123e 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -99,8 +99,8 @@ use table::TableRef; use table::dist_table::DistTable; use table::metadata::{self, TableId, TableInfo, TableMeta, TableType}; use table::requests::{ - AlterKind, AlterTableRequest, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, EntityRole, - REPARTITION_COLUMN_HINT_KEY, TableOptions, parse_entity_columns, parse_entity_option_key, + AlterKind, AlterTableRequest, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, REPARTITION_COLUMN_HINT_KEY, + TableOptions, parse_entity_columns, parse_entity_option_key, }; use table::table_name::TableName; use table::table_reference::TableReference; @@ -109,12 +109,12 @@ use crate::error::{ self, AlterExprToRequestSnafu, BuildDfLogicalPlanSnafu, CatalogSnafu, ColumnDataTypeSnafu, ColumnNotFoundSnafu, ConvertSchemaSnafu, CreateLogicalTablesSnafu, DeserializePartitionExprSnafu, EmptyDdlExprSnafu, ExternalSnafu, ExtractTableNamesSnafu, - FlowNotFoundSnafu, InvalidEntitySemanticOptionSnafu, InvalidPartitionRuleSnafu, - InvalidPartitionSnafu, InvalidSqlSnafu, InvalidTableNameSnafu, InvalidViewNameSnafu, - InvalidViewStmtSnafu, NotSupportedSnafu, PartitionExprToPbSnafu, Result, SchemaInUseSnafu, - SchemaNotFoundSnafu, SchemaReadOnlySnafu, SerializePartitionExprSnafu, SubstraitCodecSnafu, - TableAlreadyExistsSnafu, TableDdlReservedSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu, - TableReadOnlySnafu, UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu, + FlowNotFoundSnafu, InvalidPartitionRuleSnafu, InvalidPartitionSnafu, InvalidSqlSnafu, + InvalidTableNameSnafu, InvalidViewNameSnafu, InvalidViewStmtSnafu, NotSupportedSnafu, + PartitionExprToPbSnafu, Result, SchemaInUseSnafu, SchemaNotFoundSnafu, SchemaReadOnlySnafu, + SerializePartitionExprSnafu, SubstraitCodecSnafu, TableAlreadyExistsSnafu, + TableDdlReservedSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu, TableReadOnlySnafu, + UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu, }; use crate::expr_helper::{self, RepartitionRequest, RepartitionSource}; use crate::statement::StatementExecutor; @@ -2371,6 +2371,19 @@ pub fn verify_alter( .context(error::BuildTableMetaSnafu { table_name })?; validate_json2_columns_append_mode(&new_meta.schema, &new_meta.options)?; + // A column referenced by an entity declaration may be dropped (the + // read-time derivation skips the stale declaration), but must not change + // to a type without a stable string form. + for (key, value) in &new_meta.options.extra_options { + if parse_entity_option_key(key).is_none() { + continue; + } + for col in &parse_entity_columns(value) { + if let Some(column) = new_meta.schema.column_schema_by_name(col) { + ensure_entity_column_renders_as_string(key, col, &column.data_type)?; + } + } + } Ok(true) } @@ -2429,11 +2442,7 @@ pub fn create_table_info( &create_table.time_index, )?; - validate_entity_semantic_options( - &table_options, - &column_name_to_index_map, - &primary_key_indices, - )?; + validate_entity_semantic_options(&table_options, &schema)?; let meta = TableMeta { schema, @@ -2581,32 +2590,54 @@ fn ensure_table_definition_writable(schema: &str, table: &str) -> Result<()> { Ok(()) } -/// Validates `greptime.semantic.entity..{id|descriptive|scope}` options -/// against the table schema: every named column must exist, and `id` columns must -/// additionally be tag (primary-key) columns so entity identity stays indexable -/// and joinable. See `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. -fn validate_entity_semantic_options( - table_options: &TableOptions, - column_name_to_index_map: &HashMap, - primary_key_indices: &[usize], +/// Whether `CAST(column AS Utf8)` yields a stable entity-id string — the +/// binary-backed and nested types do not, and without the DDL check the +/// failure would surface only when the graph is scanned. +fn has_stable_string_form(data_type: &datatypes::prelude::ConcreteDataType) -> bool { + use datatypes::prelude::ConcreteDataType; + !matches!( + data_type, + ConcreteDataType::Binary(_) + | ConcreteDataType::Json(_) + | ConcreteDataType::Vector(_) + | ConcreteDataType::List(_) + | ConcreteDataType::Struct(_) + | ConcreteDataType::Dictionary(_) + | ConcreteDataType::Null(_) + ) +} + +fn ensure_entity_column_renders_as_string( + key: &str, + col: &str, + data_type: &datatypes::prelude::ConcreteDataType, ) -> Result<()> { + ensure!( + has_stable_string_form(data_type), + InvalidSqlSnafu { + err_msg: format!( + "entity column `{col}` (option `{key}`) has type `{data_type}`, which cannot \ + render as a string", + ), + } + ); + Ok(()) +} + +/// Validates `greptime.semantic.entity..{id|descriptive|scope}` options +/// against the table schema: every named column must exist (tag or field) and +/// have a stable string form — the derivation renders id, scope and +/// descriptive values as strings. +fn validate_entity_semantic_options(table_options: &TableOptions, schema: &Schema) -> Result<()> { for (key, value) in &table_options.extra_options { - let Some((_, role)) = parse_entity_option_key(key) else { + if parse_entity_option_key(key).is_none() { continue; }; - let is_id_role = role == EntityRole::Id; for col in &parse_entity_columns(value) { - let idx = column_name_to_index_map - .get(col) + let column = schema + .column_schema_by_name(col) .context(ColumnNotFoundSnafu { msg: col })?; - ensure!( - !is_id_role || primary_key_indices.contains(idx), - InvalidEntitySemanticOptionSnafu { - reason: format!( - "entity id column `{col}` (option `{key}`) must be a tag/primary-key column" - ), - } - ); + ensure_entity_column_renders_as_string(key, col, &column.data_type)?; } } Ok(()) @@ -3142,13 +3173,12 @@ mod test { #[test] fn test_validate_entity_semantic_options() { - let col_map = HashMap::from([ - ("service_name".to_string(), 0usize), - ("host_id".to_string(), 1usize), - ("value".to_string(), 2usize), + let schema = Schema::new(vec![ + ColumnSchema::new("service_name", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("host_id", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new("payload", ConcreteDataType::binary_datatype(), true), ]); - // service_name (0) and host_id (1) are tags; value (2) is a field. - let pk = vec![0usize, 1usize]; let opts = |pairs: &[(&str, &str)]| TableOptions { extra_options: pairs .iter() @@ -3157,39 +3187,30 @@ mod test { ..Default::default() }; - // id on a tag column, composite id on tags, and descriptive on a field - // column are all accepted. + // Any existing column with a string form may be an id, tag or field. for pairs in [ - &[("greptime.semantic.entity.service.id", "service_name")][..], &[( "greptime.semantic.entity.process.id", "service_name,host_id", )][..], - &[("greptime.semantic.entity.service.descriptive", "value")][..], + &[("greptime.semantic.entity.service.id", "value")][..], ] { - assert!(validate_entity_semantic_options(&opts(pairs), &col_map, &pk).is_ok()); + assert!(validate_entity_semantic_options(&opts(pairs), &schema).is_ok()); } - // id pointing at a field column is rejected. - let err = validate_entity_semantic_options( - &opts(&[("greptime.semantic.entity.service.id", "value")]), - &col_map, - &pk, - ) - .unwrap_err(); - assert!(matches!( - err, - error::Error::InvalidEntitySemanticOption { .. } - )); - - // id pointing at a missing column is rejected. - let err = validate_entity_semantic_options( + let missing = validate_entity_semantic_options( &opts(&[("greptime.semantic.entity.service.id", "nope")]), - &col_map, - &pk, + &schema, ) .unwrap_err(); - assert!(matches!(err, error::Error::ColumnNotFound { .. })); + assert!(matches!(missing, error::Error::ColumnNotFound { .. })); + + let binary_id = validate_entity_semantic_options( + &opts(&[("greptime.semantic.entity.service.id", "payload")]), + &schema, + ) + .unwrap_err(); + assert!(matches!(binary_id, error::Error::InvalidSql { .. })); } #[test] @@ -3555,6 +3576,43 @@ WITH ('repartition.column.hint' = ' host ')", ); } + #[test] + fn test_verify_alter_guards_entity_column_types() { + let expr = create_expr_from_sql( + "CREATE TABLE t (svc STRING, ts TIMESTAMP TIME INDEX) \ + WITH ('greptime.semantic.entity.service.id'='svc');", + ); + let table_info = Arc::new(create_table_info(&expr, vec![]).unwrap()); + let alter = |kind| AlterTableExpr { + catalog_name: "greptime".to_string(), + schema_name: "public".to_string(), + table_name: "t".to_string(), + kind: Some(kind), + }; + + let modify = alter(Kind::ModifyColumnTypes(api::v1::ModifyColumnTypes { + modify_column_types: vec![api::v1::ModifyColumnType { + column_name: "svc".to_string(), + target_type: api::v1::ColumnDataType::Binary as i32, + target_type_extension: None, + }], + })); + let err = verify_alter(1, table_info.clone(), modify).unwrap_err(); + assert!( + err.to_string().contains("cannot render as a string"), + "{err}" + ); + + // Dropping a declared column stays allowed: the derivation skips the + // stale declaration. + let drop = alter(Kind::DropColumns(api::v1::DropColumns { + drop_columns: vec![api::v1::DropColumn { + name: "svc".to_string(), + }], + })); + assert!(verify_alter(1, table_info, drop).unwrap()); + } + #[test] fn test_json2_requires_append_mode() { let cases = [ diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index c14d9d0d04e..777c6f677ea 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -25,6 +25,7 @@ //! including filter pushdown into the source table scans. See //! `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. +mod conventions; mod relationships; use std::sync::{Arc, LazyLock}; @@ -35,10 +36,25 @@ use api::v1::{ SemanticType, }; use common_catalog::consts::{ - DEFAULT_PRIVATE_SCHEMA_NAME, SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + CONFIDENCE_COLUMN, DEFAULT_CATALOG_NAME, DEFAULT_PRIVATE_SCHEMA_NAME, DST_ID_COLUMN, + DST_TYPE_COLUMN, DURATION_COUNT_COLUMN, DURATION_SUM_COLUMN, EDGE_ATTRIBUTES_COLUMN, + ENTITY_DESCRIPTIVE_COLUMN, ENTITY_ID_ATTRS_COLUMN, ENTITY_ID_COLUMN, ENTITY_SCOPE_COLUMN, + ENTITY_TYPE_COLUMN, ERROR_COUNT_COLUMN, FRESH_UNTIL_COLUMN, GENERATION_ID_COLUMN, + OBSERVED_AT_COLUMN, PROVENANCE_COLUMN, REL_TYPE_COLUMN, REQUEST_COUNT_COLUMN, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN, SRC_ID_COLUMN, + SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, WINDOW_START_COLUMN, }; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; +pub use conventions::{ + Conventions, ENTITY_TYPE_GEN_AI_AGENT, ENTITY_TYPE_GEN_AI_MODEL, ENTITY_TYPE_GEN_AI_TOOL, + ENTITY_TYPE_HOST, ENTITY_TYPE_K8S_CONTAINER, ENTITY_TYPE_K8S_NODE, ENTITY_TYPE_K8S_POD, + ENTITY_TYPE_K8S_WORKLOAD, ENTITY_TYPE_PROCESS, ENTITY_TYPE_SERVICE, + ENTITY_TYPE_SERVICE_INSTANCE, ImplicitEntity, PROVENANCE_AGENT, PROVENANCE_ATTRIBUTE, + PROVENANCE_DECLARED, PROVENANCE_TRACE, REL_TYPE_CALLS, REL_TYPE_CONTAINS, REL_TYPE_DEPENDS_ON, + REL_TYPE_INVOKES, REL_TYPE_OWNS, REL_TYPE_PART_OF, REL_TYPE_RUNS_ON, REL_TYPE_USES, + conventions, +}; use datafusion::arrow::datatypes::{DataType, TimeUnit}; use datafusion::dataframe::DataFrame; use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns}; @@ -86,7 +102,7 @@ pub fn declared_relationships_schema_matches(table_info: &table::metadata::Table return false; } - let canonical = build_declared_relationships_expr("greptime"); + let canonical = build_declared_relationships_expr(DEFAULT_CATALOG_NAME); if schema.column_schemas().len() != canonical.column_defs.len() { return false; } @@ -110,9 +126,6 @@ pub fn declared_relationships_schema_matches(table_info: &table::metadata::Table /// service-graph convention. const BIN_NANOS: i64 = 60 * 1_000_000_000; -/// Time index of the graph tables: when an observation was recorded. -pub const OBSERVED_AT_COLUMN: &str = "observed_at"; - /// Default retention for the declared-edge table; expiry slides the topology window. const DEFAULT_DECLARED_RELATIONSHIPS_TTL: &str = "90d"; /// Environment variable overriding the declared-edge table's TTL at creation @@ -133,14 +146,14 @@ fn declared_relationships_ttl() -> String { /// `provenance` and `generation_id` are in the key so a declared edge and a /// (future) derived edge for the same pair coexist without clobbering. pub const DECLARED_PRIMARY_KEY_COLUMNS: [&str; 8] = [ - "src_type", - "src_id", - "rel_type", - "dst_type", - "dst_id", - "provenance", - "scope", - "generation_id", + SRC_TYPE_COLUMN, + SRC_ID_COLUMN, + REL_TYPE_COLUMN, + DST_TYPE_COLUMN, + DST_ID_COLUMN, + PROVENANCE_COLUMN, + ENTITY_SCOPE_COLUMN, + GENERATION_ID_COLUMN, ]; /// The externally visible edge identity: the primary key minus `scope` and @@ -148,12 +161,12 @@ pub const DECLARED_PRIMARY_KEY_COLUMNS: [&str; 8] = [ /// uses this identity, or assertions differing only in those two columns would /// surface as indistinguishable duplicate rows. const DECLARED_EDGE_IDENTITY_COLUMNS: [&str; 6] = [ - "src_type", - "src_id", - "rel_type", - "dst_type", - "dst_id", - "provenance", + SRC_TYPE_COLUMN, + SRC_ID_COLUMN, + REL_TYPE_COLUMN, + DST_TYPE_COLUMN, + DST_ID_COLUMN, + PROVENANCE_COLUMN, ]; fn column( @@ -210,32 +223,32 @@ pub fn build_declared_relationships_expr(catalog: &str) -> CreateTableExpr { SemanticType::Timestamp, false, ), - field("window_start", ColumnDataType::TimestampMillisecond), - field("window_end", ColumnDataType::TimestampMillisecond), - field("fresh_until", ColumnDataType::TimestampMillisecond), + field(WINDOW_START_COLUMN, ColumnDataType::TimestampMillisecond), + field(WINDOW_END_COLUMN, ColumnDataType::TimestampMillisecond), + field(FRESH_UNTIL_COLUMN, ColumnDataType::TimestampMillisecond), // Declared-only business validity. NULL valid_from = valid since the // declaration; NULL valid_until = valid for as long as the row exists // (TTL expiry retires the edge with the row). - field("valid_from", ColumnDataType::TimestampMillisecond), - field("valid_until", ColumnDataType::TimestampMillisecond), + field(VALID_FROM_COLUMN, ColumnDataType::TimestampMillisecond), + field(VALID_UNTIL_COLUMN, ColumnDataType::TimestampMillisecond), // Endpoints + edge identity (all tags, in primary-key order). - tag("src_type"), - tag("src_id"), - tag("rel_type"), - tag("dst_type"), - tag("dst_id"), - tag("provenance"), - tag("scope"), - tag("generation_id"), + tag(SRC_TYPE_COLUMN), + tag(SRC_ID_COLUMN), + tag(REL_TYPE_COLUMN), + tag(DST_TYPE_COLUMN), + tag(DST_ID_COLUMN), + tag(PROVENANCE_COLUMN), + tag(ENTITY_SCOPE_COLUMN), + tag(GENERATION_ID_COLUMN), // Confidence + RED metrics (populated for derived edges; usually NULL here). - field("confidence", ColumnDataType::Float64), - field("request_count", ColumnDataType::Int64), - field("error_count", ColumnDataType::Int64), - field("duration_sum", ColumnDataType::Float64), - field("duration_count", ColumnDataType::Int64), + field(CONFIDENCE_COLUMN, ColumnDataType::Float64), + field(REQUEST_COUNT_COLUMN, ColumnDataType::Int64), + field(ERROR_COUNT_COLUMN, ColumnDataType::Int64), + field(DURATION_SUM_COLUMN, ColumnDataType::Float64), + field(DURATION_COUNT_COLUMN, ColumnDataType::Int64), // JSONB, so the union matches the computed table's json column // without a per-scan parse. - json_field("attributes"), + json_field(EDGE_ATTRIBUTES_COLUMN), ]; let table_options = [(TTL_KEY.to_string(), declared_relationships_ttl())] @@ -408,6 +421,16 @@ fn cast_string_or_empty(column: &str) -> Expr { /// a single column, the sorted `k=v,k=v` rendering for a composite. `col` /// constructs the column reference (unqualified for registry branches, /// join-side-qualified for the calls derivation). +/// A row identifies an entity only when every identity component is present +/// and non-empty: kube-state-metrics descriptors emit empty-string labels (an +/// unscheduled pod's `node`, an owner-less pod's `owner_*`), and an empty +/// string is never a meaningful entity id. +pub(crate) fn identifies(column: &str) -> Expr { + ident(column) + .is_not_null() + .and(cast(ident(column), DataType::Utf8).not_eq(lit(""))) +} + fn entity_id_expr(id_columns: &[String], col: &dyn Fn(&str) -> Expr) -> Expr { if let [id] = id_columns { cast(col(id), DataType::Utf8) @@ -501,16 +524,16 @@ fn sorted_kv_expr_with(sorted_cols: &[String], nullable: bool, col: &dyn Fn(&str } const REGISTRY_COLUMNS: [&str; 10] = [ - "observed_at", - "window_start", - "window_end", - "fresh_until", - "entity_type", - "entity_id", - "entity_id_attrs", - "scope", - "descriptive", - "source_tables", + OBSERVED_AT_COLUMN, + WINDOW_START_COLUMN, + WINDOW_END_COLUMN, + FRESH_UNTIL_COLUMN, + ENTITY_TYPE_COLUMN, + ENTITY_ID_COLUMN, + ENTITY_ID_ATTRS_COLUMN, + ENTITY_SCOPE_COLUMN, + ENTITY_DESCRIPTIVE_COLUMN, + SOURCE_TABLES_COLUMN, ]; const REGISTRY_VALID_COLUMN: &str = "__entity_valid"; @@ -568,10 +591,10 @@ fn registry_source( let mut rows = Vec::with_capacity(1 + rest.len()); for decl in std::iter::once(first).chain(rest) { - // CAST even a single-column id: id columns must be tags but not - // necessarily strings, and the computed table declares entity_id - // STRING. Composite ids additionally carry a JSON object of the id - // columns in entity_id_attrs. + // CAST even a single-column id: id columns need not be strings, and + // the computed table declares entity_id STRING. Composite ids + // additionally carry a JSON object of the id columns in + // entity_id_attrs. let entity_id = entity_id_expr(&decl.id_columns, &|c| ident(c)); let entity_id_attrs = if decl.id_columns.len() == 1 { null_json() @@ -603,12 +626,12 @@ fn registry_source( json_quote(&format!("{}.{}", decl.schema, decl.table)) ))); - // Tag columns may still be nullable; a NULL identity component - // identifies nothing. Keep this predicate per declaration so a NULL - // identity for one entity does not remove other entities on the row. - let valid = decl.id_columns.iter().fold(lit(true), |predicate, id| { - predicate.and(ident(id).is_not_null()) - }); + // Keep this predicate per declaration so an absent identity for one + // entity does not remove other entities on the row. + let valid = decl + .id_columns + .iter() + .fold(lit(true), |predicate, id| predicate.and(identifies(id))); rows.push(vec![ valid, @@ -961,8 +984,31 @@ mod tests { } #[tokio::test] - async fn registry_skips_null_identity_rows() { - let ctx = metric_table_ctx(); + async fn registry_skips_absent_identity_rows() { + // NULL identifies nothing, and so does a kube-state-metrics-style + // empty label (an unscheduled pod's `node` arrives as ""). + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("host", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 2_000, 3_000])) as ArrayRef, + Arc::new(StringArray::from(vec![Some("h2"), None, Some("")])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "app_latency", + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); let df = ctx.table("app_latency").await.unwrap(); let plan = build_registry_plan( vec![RegistrySource { @@ -974,9 +1020,8 @@ mod tests { .unwrap() .unwrap(); let batches = collect(&ctx, plan).await; - let mut ids: Vec = batches.iter().flat_map(|b| strings(b, 5)).collect(); - ids.sort(); - assert_eq!(ids, vec!["h2", r#"we"ird\host"#]); + let ids: Vec = batches.iter().flat_map(|b| strings(b, 5)).collect(); + assert_eq!(ids, vec!["h2"]); } #[tokio::test] diff --git a/src/operator/src/statement/semantic_graph/conventions.rs b/src/operator/src/statement/semantic_graph/conventions.rs new file mode 100644 index 00000000000..36d058921b5 --- /dev/null +++ b/src/operator/src/statement/semantic_graph/conventions.rs @@ -0,0 +1,291 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The built-in derivation conventions of the entity graph, embedded as +//! [`conventions.yaml`](./conventions.yaml): the co-declared edge vocabulary, +//! the virtual-destination candidates, and the implicit declarations of +//! well-known Prometheus entity-descriptor metrics. The file is data shipped +//! with the binary, not an operator-editable configuration surface; explicit +//! `greptime.semantic.entity.*` declarations always override it. + +use std::collections::{BTreeMap, HashSet}; +use std::sync::LazyLock; + +use serde::Deserialize; + +/// A same-row co-declaration rule: a source-table row carrying both entity +/// identities witnesses `src -rel-> dst`. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct EdgeRule { + pub src: String, + pub dst: String, + pub rel: String, +} + +/// A span-attribute column that may name an uninstrumented peer, and the +/// `connection_type` its match implies. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct VirtualDstCandidate { + pub column: String, + pub connection_type: String, +} + +/// One implicit entity declaration: of a whitelisted Prometheus info metric, +/// or of a trace-v1 table's flattened resource attributes. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ImplicitEntity { + pub entity: String, + /// Identifying label columns; every one must exist on the table for the + /// declaration to apply. + pub id: Vec, + /// Descriptive label columns, filtered to those present (kube-state-metrics + /// label sets vary across versions). + #[serde(default)] + pub descriptive: Vec, + /// Snapshot every tag column except the id columns as descriptive — + /// `target_info`-style enrichment over an open attribute set. + #[serde(default)] + pub descriptive_rest: bool, +} + +/// The parsed, validated conventions file. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Conventions { + pub co_declared_edges: Vec, + pub trace_co_declared_edges: Vec, + pub virtual_dst_candidates: Vec, + pub otlp_trace_entities: Vec, + pub prometheus_info_metrics: BTreeMap>, +} + +/// The built-in entity-type vocabulary. User-declared types are open-ended; +/// the embedded conventions file must stay inside this set. +pub const ENTITY_TYPE_SERVICE: &str = "service"; +pub const ENTITY_TYPE_SERVICE_INSTANCE: &str = "service.instance"; +pub const ENTITY_TYPE_HOST: &str = "host"; +pub const ENTITY_TYPE_PROCESS: &str = "process"; +pub const ENTITY_TYPE_K8S_POD: &str = "k8s.pod"; +pub const ENTITY_TYPE_K8S_NODE: &str = "k8s.node"; +pub const ENTITY_TYPE_K8S_CONTAINER: &str = "k8s.container"; +pub const ENTITY_TYPE_K8S_WORKLOAD: &str = "k8s.workload"; +pub const ENTITY_TYPE_K8S_SERVICE: &str = "k8s.service"; +pub const ENTITY_TYPE_GEN_AI_AGENT: &str = "gen_ai.agent"; +pub const ENTITY_TYPE_GEN_AI_MODEL: &str = "gen_ai.model"; +pub const ENTITY_TYPE_GEN_AI_TOOL: &str = "gen_ai.tool"; + +const ENTITY_TYPES: [&str; 12] = [ + ENTITY_TYPE_SERVICE, + ENTITY_TYPE_SERVICE_INSTANCE, + ENTITY_TYPE_HOST, + ENTITY_TYPE_PROCESS, + ENTITY_TYPE_K8S_POD, + ENTITY_TYPE_K8S_NODE, + ENTITY_TYPE_K8S_CONTAINER, + ENTITY_TYPE_K8S_WORKLOAD, + ENTITY_TYPE_K8S_SERVICE, + ENTITY_TYPE_GEN_AI_AGENT, + ENTITY_TYPE_GEN_AI_MODEL, + ENTITY_TYPE_GEN_AI_TOOL, +]; + +/// The relationship vocabulary (the RFC's rel_type table) and the edge +/// provenances. +pub const REL_TYPE_CALLS: &str = "calls"; +pub const REL_TYPE_RUNS_ON: &str = "runs_on"; +pub const REL_TYPE_CONTAINS: &str = "contains"; +pub const REL_TYPE_PART_OF: &str = "part_of"; +pub const REL_TYPE_USES: &str = "uses"; +pub const REL_TYPE_INVOKES: &str = "invokes"; +pub const REL_TYPE_DEPENDS_ON: &str = "depends_on"; +pub const REL_TYPE_OWNS: &str = "owns"; +pub const PROVENANCE_TRACE: &str = "trace"; +pub const PROVENANCE_ATTRIBUTE: &str = "attribute"; +pub const PROVENANCE_DECLARED: &str = "declared"; +pub const PROVENANCE_AGENT: &str = "agent"; + +const REL_TYPES: [&str; 8] = [ + REL_TYPE_CALLS, + REL_TYPE_RUNS_ON, + REL_TYPE_CONTAINS, + REL_TYPE_PART_OF, + REL_TYPE_USES, + REL_TYPE_INVOKES, + REL_TYPE_DEPENDS_ON, + REL_TYPE_OWNS, +]; + +/// The `connection_type` values a virtual-destination candidate may imply; +/// the calls derivation branches on them when building edge attributes. +pub const CONNECTION_TYPE_DATABASE: &str = "database"; +pub const CONNECTION_TYPE_VIRTUAL_NODE: &str = "virtual_node"; + +const CONNECTION_TYPES: [&str; 2] = [CONNECTION_TYPE_DATABASE, CONNECTION_TYPE_VIRTUAL_NODE]; + +static CONVENTIONS: LazyLock> = + LazyLock::new(|| parse(include_str!("conventions.yaml"))); + +/// The embedded conventions. `Err` means the embedded file is broken — pinned +/// by unit test, and propagated by the derivation paths rather than panicking. +pub fn conventions() -> Result<&'static Conventions, String> { + CONVENTIONS.as_ref().map_err(Clone::clone) +} + +fn parse(yaml: &str) -> Result { + let conventions: Conventions = + serde_yaml_ng::from_str(yaml).map_err(|e| format!("malformed conventions: {e}"))?; + validate(&conventions)?; + Ok(conventions) +} + +fn validate(conventions: &Conventions) -> Result<(), String> { + let mut seen_edges = HashSet::new(); + for rule in conventions + .co_declared_edges + .iter() + .chain(&conventions.trace_co_declared_edges) + { + for ty in [&rule.src, &rule.dst] { + if !ENTITY_TYPES.contains(&ty.as_str()) { + return Err(format!("unknown entity type `{ty}` in edge vocabulary")); + } + } + if !REL_TYPES.contains(&rule.rel.as_str()) { + return Err(format!( + "unknown rel_type `{}` in edge vocabulary", + rule.rel + )); + } + if !seen_edges.insert((&rule.src, &rule.dst, &rule.rel)) { + return Err(format!( + "duplicate edge rule `{} -{}-> {}`", + rule.src, rule.rel, rule.dst + )); + } + } + + let mut seen_columns = HashSet::new(); + for candidate in &conventions.virtual_dst_candidates { + if candidate.column.is_empty() { + return Err("empty virtual destination column".to_string()); + } + if !CONNECTION_TYPES.contains(&candidate.connection_type.as_str()) { + return Err(format!( + "unknown connection_type `{}` for virtual destination `{}`", + candidate.connection_type, candidate.column + )); + } + if !seen_columns.insert(&candidate.column) { + return Err(format!( + "duplicate virtual destination column `{}`", + candidate.column + )); + } + } + + let per_table = conventions + .prometheus_info_metrics + .iter() + .map(|(table, entities)| (table.as_str(), entities)) + .chain(std::iter::once(( + "otlp traces", + &conventions.otlp_trace_entities, + ))); + for (table, entities) in per_table { + let mut seen_types = HashSet::new(); + for implicit in entities { + if !ENTITY_TYPES.contains(&implicit.entity.as_str()) { + return Err(format!( + "unknown entity type `{}` for info metric `{table}`", + implicit.entity + )); + } + if !seen_types.insert(&implicit.entity) { + return Err(format!( + "duplicate entity type `{}` for info metric `{table}`", + implicit.entity + )); + } + if implicit.id.is_empty() || implicit.id.iter().any(String::is_empty) { + return Err(format!( + "entity `{}` of info metric `{table}` needs non-empty id columns", + implicit.entity + )); + } + if implicit.descriptive_rest && !implicit.descriptive.is_empty() { + return Err(format!( + "entity `{}` of info metric `{table}` sets both descriptive and \ + descriptive_rest", + implicit.entity + )); + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn embedded_conventions_parse_and_validate() { + conventions().unwrap(); + } + + /// Each case must fail on exactly the rule it names, so the shared + /// boilerplate is valid and the mutated part is inside the vocabulary. + fn broken(edges: &str, info_metrics: &str) -> String { + format!( + "co_declared_edges: [{edges}]\ntrace_co_declared_edges: []\n\ + virtual_dst_candidates: []\notlp_trace_entities: []\n\ + prometheus_info_metrics: {{{info_metrics}}}" + ) + } + + #[test] + fn validation_rejects_broken_conventions() { + let err = |edges, info| parse(&broken(edges, info)).unwrap_err(); + + assert!(err("{src: host, dst: service, rel: pets}", "").contains("unknown rel_type")); + assert!( + err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "").contains("unknown entity type") + ); + assert!( + err( + "{src: host, dst: service, rel: uses}, {src: host, dst: service, rel: uses}", + "" + ) + .contains("duplicate edge rule") + ); + assert!( + err( + "", + "t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]" + ) + .contains("descriptive_rest") + ); + // Unknown YAML keys are rejected, catching typos in the embedded file. + assert!( + parse(&broken( + "{src: host, dst: service, rel: uses, direction: down}", + "" + )) + .is_err() + ); + } +} diff --git a/src/operator/src/statement/semantic_graph/conventions.yaml b/src/operator/src/statement/semantic_graph/conventions.yaml new file mode 100644 index 00000000000..9fc37fe2a02 --- /dev/null +++ b/src/operator/src/statement/semantic_graph/conventions.yaml @@ -0,0 +1,115 @@ +# The built-in derivation conventions of the entity graph. Embedded and +# versioned with the binary (include_str!), not an operator-editable surface; +# non-standard deployments override per table with explicit +# greptime.semantic.entity.* declarations, which always win. + +co_declared_edges: + - { src: service.instance, dst: host, rel: runs_on } + - { src: process, dst: host, rel: runs_on } + - { src: k8s.pod, dst: k8s.node, rel: runs_on } + - { src: k8s.pod, dst: k8s.container, rel: contains } + - { src: service.instance, dst: service, rel: part_of } + - { src: k8s.pod, dst: k8s.workload, rel: part_of } + - { src: service.instance, dst: k8s.pod, rel: runs_on } + +# Applied only to trace sources; the endpoints are Greptime entity types +# derived from GenAI semantic-convention attributes. +trace_co_declared_edges: + - { src: gen_ai.agent, dst: gen_ai.model, rel: uses } + - { src: gen_ai.agent, dst: gen_ai.tool, rel: invokes } + +# Span-attribute columns naming an uninstrumented peer, in precedence order: +# current OTel names first (service.peer.name / db.namespace replaced the +# deprecated peer.service / db.name in semconv 1.39/1.26, kept for existing +# telemetry), the bare network endpoint last. +virtual_dst_candidates: + - { column: span_attributes.service.peer.name, connection_type: virtual_node } + - { column: span_attributes.peer.service, connection_type: virtual_node } + - { column: span_attributes.db.namespace, connection_type: database } + - { column: span_attributes.db.name, connection_type: database } + - { column: span_attributes.server.address, connection_type: virtual_node } + +# Implicit declarations for greptime_trace_v1 tables; the pod UID bridges +# trace-side pods onto the kube-state-metrics entities below. +otlp_trace_entities: + - entity: service + id: [service_name] + - entity: service.instance + id: [service_name, resource_attributes.service.instance.id] + - entity: k8s.pod + id: [resource_attributes.k8s.pod.uid] + descriptive: + - resource_attributes.k8s.pod.name + - resource_attributes.k8s.namespace.name + +# Implicit declarations for Prometheus entity-descriptor metrics, keyed by +# table name. Pod and service identities are the UIDs (names stay +# descriptive); k8s.workload is a Greptime abstraction with the owner kind +# encoded in its identity. +prometheus_info_metrics: + kube_pod_info: + - entity: k8s.pod + id: [uid] + descriptive: + - namespace + - pod + - node + - host_ip + - pod_ip + - created_by_kind + - created_by_name + - priority_class + - host_network + - entity: k8s.node + id: [node] + kube_node_info: + - entity: k8s.node + id: [node] + descriptive: + - kernel_version + - os_image + - container_runtime_version + - kubelet_version + - pod_cidr + - provider_id + - system_uuid + - internal_ip + kube_pod_owner: + - entity: k8s.pod + id: [uid] + descriptive: [namespace, pod] + - entity: k8s.workload + id: [namespace, owner_kind, owner_name] + descriptive: [owner_is_controller] + kube_pod_container_info: + - entity: k8s.pod + id: [uid] + descriptive: [namespace, pod] + - entity: k8s.container + id: [uid, container] + descriptive: [image, image_id, image_spec, container_id] + kube_pod_init_container_info: + - entity: k8s.pod + id: [uid] + descriptive: [namespace, pod] + - entity: k8s.container + id: [uid, container] + descriptive: [image, image_id, image_spec, container_id, restart_policy] + kube_service_info: + - entity: k8s.service + id: [uid] + descriptive: + - namespace + - service + - cluster_ip + - external_name + - load_balancer_ip + - external_traffic_policy + # The remaining labels are the target's resource attributes: the snapshot + # belongs to the instance, not the logical service. + target_info: + - entity: service + id: [job] + - entity: service.instance + id: [job, instance] + descriptive_rest: true diff --git a/src/operator/src/statement/semantic_graph/relationships.rs b/src/operator/src/statement/semantic_graph/relationships.rs index 416c8dc10e1..3b5c6d37b78 100644 --- a/src/operator/src/statement/semantic_graph/relationships.rs +++ b/src/operator/src/statement/semantic_graph/relationships.rs @@ -19,9 +19,13 @@ //! expression helpers and the entity registry live in the parent module. use common_catalog::consts::{ - DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, - SPAN_KIND_COLUMN, SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, - TRACE_ID_COLUMN, TRACE_TIMESTAMP_COLUMN, + CONFIDENCE_COLUMN, DST_ID_COLUMN, DST_TYPE_COLUMN, DURATION_COUNT_COLUMN, DURATION_NANO_COLUMN, + DURATION_SUM_COLUMN, EDGE_ATTRIBUTES_COLUMN, ENTITY_SCOPE_COLUMN, ERROR_COUNT_COLUMN, + FRESH_UNTIL_COLUMN, GENERATION_ID_COLUMN, OBSERVED_AT_COLUMN, PARENT_SPAN_ID_COLUMN, + PROVENANCE_COLUMN, REL_TYPE_COLUMN, REQUEST_COUNT_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, + SPAN_KIND_COLUMN, SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, SRC_ID_COLUMN, + SRC_TYPE_COLUMN, TRACE_ID_COLUMN, TRACE_TIMESTAMP_COLUMN, VALID_FROM_COLUMN, + VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, WINDOW_START_COLUMN, }; use datafusion::arrow::datatypes::DataType; use datafusion::dataframe::DataFrame; @@ -31,12 +35,21 @@ use datafusion::functions_window::expr_fn::row_number; use datafusion_common::{Result as DfResult, ScalarValue}; use datafusion_expr::{Expr, ExprFunctionExt, JoinType, LogicalPlan, cast, ident, lit, when}; +use crate::statement::semantic_graph::conventions::{ + CONNECTION_TYPE_DATABASE, CONNECTION_TYPE_VIRTUAL_NODE, Conventions, ENTITY_TYPE_GEN_AI_AGENT, + ENTITY_TYPE_SERVICE, PROVENANCE_ATTRIBUTE, PROVENANCE_TRACE, REL_TYPE_CALLS, +}; use crate::statement::semantic_graph::{ - DECLARED_EDGE_IDENTITY_COLUMNS, EntityDeclaration, GraphQueryWindow, OBSERVED_AT_COLUMN, - bin_interval, bin_ms, entity_id_expr, interval, null_json, parse_json_expr, qcol, union_all, + DECLARED_EDGE_IDENTITY_COLUMNS, EntityDeclaration, GraphQueryWindow, bin_interval, bin_ms, + conventions, entity_id_expr, identifies, interval, null_json, parse_json_expr, qcol, union_all, unnest_rows, }; +/// The embedded conventions, with a broken file surfaced as a plan error. +fn builtin() -> DfResult<&'static Conventions> { + conventions().map_err(datafusion_common::DataFusionError::Internal) +} + /// The projected columns of `semantic_relationships`, in order. Every derived /// branch and the declared-edge branch must project exactly these so the /// top-level `UNION ALL` type-aligns; `build_relationships_plan` re-selects @@ -44,22 +57,22 @@ use crate::statement::semantic_graph::{ /// additionally stores `valid_from`/`valid_until`, which feed the validity /// filter and the projected window columns.) const RELATIONSHIP_COLUMNS: [&str; 16] = [ - "observed_at", - "window_start", - "window_end", - "fresh_until", - "src_type", - "src_id", - "dst_type", - "dst_id", - "rel_type", - "provenance", - "confidence", - "request_count", - "error_count", - "duration_sum", - "duration_count", - "attributes", + OBSERVED_AT_COLUMN, + WINDOW_START_COLUMN, + WINDOW_END_COLUMN, + FRESH_UNTIL_COLUMN, + SRC_TYPE_COLUMN, + SRC_ID_COLUMN, + DST_TYPE_COLUMN, + DST_ID_COLUMN, + REL_TYPE_COLUMN, + PROVENANCE_COLUMN, + CONFIDENCE_COLUMN, + REQUEST_COUNT_COLUMN, + ERROR_COUNT_COLUMN, + DURATION_SUM_COLUMN, + DURATION_COUNT_COLUMN, + EDGE_ATTRIBUTES_COLUMN, ]; /// A child server span starts no earlier than 5 minutes before its client span @@ -129,25 +142,6 @@ pub fn build_relationships_plan( .transpose() } -/// The same-row co-declaration vocabulary open to any declaring table: a row -/// carrying both identities witnesses the edge, and the built-in direction is -/// `src -> dst`. Rows are `(src_type, dst_type, rel_type)`. -const ATTRIBUTE_EDGE_VOCABULARY: [(&str, &str, &str); 6] = [ - ("service.instance", "host", "runs_on"), - ("process", "host", "runs_on"), - ("k8s.pod", "k8s.node", "runs_on"), - ("k8s.pod", "k8s.container", "contains"), - ("service.instance", "service", "part_of"), - ("k8s.pod", "k8s.workload", "part_of"), -]; - -/// The agent-edge vocabulary, applied only to trace sources: the RFC derives -/// `agent uses model` / `agent invokes tool` from span structure (an LLM- or -/// tool-call span carries both identities), so a non-trace table co-declaring -/// these types must not fabricate invocations. -const AGENT_EDGE_VOCABULARY: [(&str, &str, &str); 2] = - [("agent", "model", "uses"), ("agent", "tool", "invokes")]; - const CO_DECLARED_VALID_COLUMN: &str = "__edge_valid"; /// The same-row co-declared branch: for each source table, every vocabulary @@ -159,19 +153,22 @@ fn co_declared_branch( sources: Vec, window: &GraphQueryWindow, ) -> DfResult> { + let vocabulary = builtin()?; let mut union_df: Option = None; for source in sources { let find = |ty: &str| source.declarations.iter().find(|d| d.entity_type == ty); let mut pairs = Vec::new(); - for (src_type, dst_type, rel_type) in ATTRIBUTE_EDGE_VOCABULARY { - if let (Some(src), Some(dst)) = (find(src_type), find(dst_type)) { - pairs.push((src, dst, rel_type, "attribute")); + for rule in &vocabulary.co_declared_edges { + if let (Some(src), Some(dst)) = (find(&rule.src), find(&rule.dst)) { + pairs.push((src, dst, rule.rel.as_str(), PROVENANCE_ATTRIBUTE)); } } if source.is_trace { - for (src_type, dst_type, rel_type) in AGENT_EDGE_VOCABULARY { - if let (Some(src), Some(dst)) = (find(src_type), find(dst_type)) { - pairs.push((src, dst, rel_type, "trace")); + // Agent edges are tied to span structure by the RFC: a non-trace + // table co-declaring these types must not fabricate invocations. + for rule in &vocabulary.trace_co_declared_edges { + if let (Some(src), Some(dst)) = (find(&rule.src), find(&rule.dst)) { + pairs.push((src, dst, rule.rel.as_str(), PROVENANCE_TRACE)); } } } @@ -195,9 +192,7 @@ fn co_declared_branch( .id_columns .iter() .chain(&dst.id_columns) - .fold(lit(true), |predicate, id| { - predicate.and(ident(id).is_not_null()) - }); + .fold(lit(true), |predicate, id| predicate.and(identifies(id))); vec![ valid, bin.clone(), @@ -249,9 +244,9 @@ const DECLARED_REVISION_COLUMN: &str = "__declared_revision"; /// above the computed table, and the physical revision time would fail them. fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult { let eff_valid_from = - core_fns::coalesce().call(vec![ident("valid_from"), ident(OBSERVED_AT_COLUMN)]); + core_fns::coalesce().call(vec![ident(VALID_FROM_COLUMN), ident(OBSERVED_AT_COLUMN)]); let eff_valid_until = - core_fns::coalesce().call(vec![ident("valid_until"), window.observed_end()]); + core_fns::coalesce().call(vec![ident(VALID_UNTIL_COLUMN), window.observed_end()]); // As-of the queried window: a revision recorded after its end, or whose // validity starts after it, was not the edge's state inside the window and @@ -271,17 +266,17 @@ fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult DfResult DfResult DfResult DfResult DfResult DfResult { /// `connection_type`. fn virtual_attrs_expr() -> DfResult { when( - ident("virtual_conn").eq(lit("database")), - parse_json_expr(lit(r#"{"connection_type":"database"}"#)), + ident("virtual_conn").eq(lit(CONNECTION_TYPE_DATABASE)), + parse_json_expr(lit(format!( + r#"{{"connection_type":"{CONNECTION_TYPE_DATABASE}"}}"# + ))), ) - .otherwise(parse_json_expr(lit( - r#"{"connection_type":"virtual_node"}"#, - ))) + .otherwise(parse_json_expr(lit(format!( + r#"{{"connection_type":"{CONNECTION_TYPE_VIRTUAL_NODE}"}}"# + )))) } /// The window + span-kind + identity predicate shared by both join sides. @@ -506,9 +493,9 @@ fn span_predicate(service: &EntityDeclaration, window: &GraphQueryWindow, strict .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)) .and(ts.lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS))) }; - // A NULL identity component identifies nothing, on either endpoint. + // An absent identity component identifies nothing, on either endpoint. for id in &service.id_columns { - predicate = predicate.and(ident(id).is_not_null()); + predicate = predicate.and(identifies(id)); } predicate } @@ -524,16 +511,22 @@ fn client_spans( scan: &DataFrame, window: &GraphQueryWindow, ) -> DfResult { - // NULLIF('') lets an empty value fall through to the next candidate. + // Attribute columns are dynamic in `greptime_trace_v1` (created on first + // use), so only candidate columns present in the table's schema + // participate. NULLIF('') lets an empty value fall through to the next + // candidate. let present: Vec<(Expr, &str)> = { let schema = scan.schema(); - VIRTUAL_DST_CANDIDATES + builtin()? + .virtual_dst_candidates .iter() - .filter(|(column, _)| schema.has_column_with_unqualified_name(column)) - .map(|(column, conn)| { - let value = - core_fns::nullif().call(vec![cast(ident(*column), DataType::Utf8), lit("")]); - (value, *conn) + .filter(|candidate| schema.has_column_with_unqualified_name(&candidate.column)) + .map(|candidate| { + let value = core_fns::nullif().call(vec![ + cast(ident(&candidate.column), DataType::Utf8), + lit(""), + ]); + (value, candidate.connection_type.as_str()) }) .collect() }; @@ -560,7 +553,7 @@ fn client_spans( ident(SPAN_ID_COLUMN), // The cast inside entity_id_expr also normalizes tag columns, which // come out of the storage engine dictionary-encoded. - entity_id_expr(&service.id_columns, &|c| ident(c)).alias("src_id"), + entity_id_expr(&service.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), ident(DURATION_NANO_COLUMN), virtual_dst.alias("virtual_dst"), @@ -585,7 +578,7 @@ fn server_spans( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&service.id_columns, &|c| ident(c)).alias("dst_id"), + entity_id_expr(&service.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), ident(DURATION_NANO_COLUMN), ]) @@ -616,7 +609,7 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias("src_id"), + entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), ])?; let child = trace .scan @@ -626,7 +619,7 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias("dst_id"), + entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), ident(DURATION_NANO_COLUMN), ])?; @@ -651,44 +644,48 @@ fn agent_calls_branch( let df = parent .join_on(child, JoinType::Inner, join_conditions)? .select(vec![ - bin_ms(qcol("parent", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), - qcol("parent", "src_id").alias("src_id"), - qcol("child", "dst_id").alias("dst_id"), + bin_ms(qcol("parent", TRACE_TIMESTAMP_COLUMN)).alias(OBSERVED_AT_COLUMN), + qcol("parent", SRC_ID_COLUMN).alias(SRC_ID_COLUMN), + qcol("child", DST_ID_COLUMN).alias(DST_ID_COLUMN), qcol("child", SPAN_STATUS_CODE_COLUMN).alias("status_code"), - qcol("child", DURATION_NANO_COLUMN).alias("duration_nano"), + qcol("child", DURATION_NANO_COLUMN).alias(DURATION_NANO_COLUMN), ])? // A sub-span of the same agent is internal structure, not a call // between two agents. - .filter(ident("src_id").not_eq(ident("dst_id")))? + .filter(ident(SRC_ID_COLUMN).not_eq(ident(DST_ID_COLUMN)))? .aggregate( - vec![ident("observed_at"), ident("src_id"), ident("dst_id")], vec![ - count(lit(1)).alias("request_count"), + ident(OBSERVED_AT_COLUMN), + ident(SRC_ID_COLUMN), + ident(DST_ID_COLUMN), + ], + vec![ + count(lit(1)).alias(REQUEST_COUNT_COLUMN), count(lit(1)) .filter(ident("status_code").eq(lit(SPAN_STATUS_ERROR))) .build()? - .alias("error_count"), - sum(ident("duration_nano")).alias("duration_nano_sum"), + .alias(ERROR_COUNT_COLUMN), + sum(ident(DURATION_NANO_COLUMN)).alias("duration_nano_sum"), ], )? .select(vec![ - ident("observed_at"), - ident("observed_at").alias("window_start"), - (ident("observed_at") + bin_interval()).alias("window_end"), - (ident("observed_at") + bin_interval()).alias("fresh_until"), - lit("agent").alias("src_type"), - ident("src_id"), - lit("agent").alias("dst_type"), - ident("dst_id"), - lit("calls").alias("rel_type"), - lit("trace").alias("provenance"), - lit(1.0_f64).alias("confidence"), - ident("request_count"), - ident("error_count"), + ident(OBSERVED_AT_COLUMN), + ident(OBSERVED_AT_COLUMN).alias(WINDOW_START_COLUMN), + (ident(OBSERVED_AT_COLUMN) + bin_interval()).alias(WINDOW_END_COLUMN), + (ident(OBSERVED_AT_COLUMN) + bin_interval()).alias(FRESH_UNTIL_COLUMN), + lit(ENTITY_TYPE_GEN_AI_AGENT).alias(SRC_TYPE_COLUMN), + ident(SRC_ID_COLUMN), + lit(ENTITY_TYPE_GEN_AI_AGENT).alias(DST_TYPE_COLUMN), + ident(DST_ID_COLUMN), + lit(REL_TYPE_CALLS).alias(REL_TYPE_COLUMN), + lit(PROVENANCE_TRACE).alias(PROVENANCE_COLUMN), + lit(1.0_f64).alias(CONFIDENCE_COLUMN), + ident(REQUEST_COUNT_COLUMN), + ident(ERROR_COUNT_COLUMN), (cast(ident("duration_nano_sum"), DataType::Float64) / lit(1e9_f64)) - .alias("duration_sum"), - ident("request_count").alias("duration_count"), - null_json().alias("attributes"), + .alias(DURATION_SUM_COLUMN), + ident(REQUEST_COUNT_COLUMN).alias(DURATION_COUNT_COLUMN), + null_json().alias(EDGE_ATTRIBUTES_COLUMN), ])?; Ok(Some(df)) } @@ -1173,7 +1170,7 @@ mod tests { RelationshipSources { traces: vec![CallsSource { service: None, - agent: Some(co_decl("agent", &["agent_id"])), + agent: Some(co_decl("gen_ai.agent", &["agent_id"])), scan, }], co_declared: vec![], @@ -1190,9 +1187,9 @@ mod tests { // no service declaration. assert_eq!(total, 1); let batch = &batches[0]; - assert_eq!(strings(batch, 4), vec!["agent"]); + assert_eq!(strings(batch, 4), vec!["gen_ai.agent"]); assert_eq!(strings(batch, 5), vec!["orchestrator"]); - assert_eq!(strings(batch, 6), vec!["agent"]); + assert_eq!(strings(batch, 6), vec!["gen_ai.agent"]); assert_eq!(strings(batch, 7), vec!["researcher"]); assert_eq!(strings(batch, 8), vec!["calls"]); assert_eq!(strings(batch, 9), vec!["trace"]); @@ -1229,7 +1226,7 @@ mod tests { RelationshipSources { traces: vec![CallsSource { service: None, - agent: Some(co_decl("agent", &["agent_id"])), + agent: Some(co_decl("gen_ai.agent", &["agent_id"])), scan, }], co_declared: vec![], @@ -1444,9 +1441,9 @@ mod tests { let ctx = co_declared_ctx(); let declarations = || { vec![ - co_decl("agent", &["agent_id"]), - co_decl("model", &["model_name"]), - co_decl("tool", &["tool_name"]), + co_decl("gen_ai.agent", &["agent_id"]), + co_decl("gen_ai.model", &["model_name"]), + co_decl("gen_ai.tool", &["tool_name"]), ] }; // A non-trace table co-declaring agent/model/tool must not fabricate @@ -1462,17 +1459,17 @@ mod tests { rows, vec![ ( - "agent".to_string(), + "gen_ai.agent".to_string(), "agent-1".to_string(), - "model".to_string(), + "gen_ai.model".to_string(), "gpt".to_string(), "uses".to_string(), "trace".to_string(), ), ( - "agent".to_string(), + "gen_ai.agent".to_string(), "agent-1".to_string(), - "tool".to_string(), + "gen_ai.tool".to_string(), "search".to_string(), "invokes".to_string(), "trace".to_string(), diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index cb9d6aa7476..5c0e4c3caed 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -209,9 +209,9 @@ pub fn is_semantic_option_key(key: &str) -> bool { /// non-empty string. Closed-domain keys accept a fixed set, plus the `unknown` /// sentinel, plus `mixed` for the keys where one long-lived table can /// legitimately see multiple values. Entity keys ([`is_entity_option_key`]) take a -/// comma-separated column-name list (each token non-empty); the "id columns must -/// be tag columns" invariant is enforced later against the table schema at DDL -/// time, not here. Keys that are neither whitelisted nor a well-formed entity key +/// comma-separated column-name list (each token non-empty); column existence and +/// the stable-string-form rule for entity columns are enforced later against +/// the table schema at DDL time, not here. Keys that are neither whitelisted nor a well-formed entity key /// are rejected. pub fn validate_semantic_option(key: &str, value: &str) -> bool { if is_entity_option_key(key) { diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 042be385a6d..b2c8fe394e2 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -161,6 +161,7 @@ macro_rules! http_tests { test_otlp_metric_translation_strategies, test_otlp_traces_v0, test_otlp_traces_v1, + test_otlp_traces_v1_entity_graph, test_otlp_logs, test_loki_pb_logs, test_loki_pb_logs_with_pipeline, @@ -6886,6 +6887,86 @@ pub async fn test_otlp_traces_v0(store_type: StorageType) { guard.remove_all().await; } +/// One real OTLP export must come out of `semantic_relationships` as the +/// zero-configuration chain; resources missing an identity column derive +/// nothing extra. +pub async fn test_otlp_traces_v1_entity_graph(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_otlp_traces_entity_graph").await; + let client = TestClient::new(app).await; + + let span = |trace_id: &str, span_id: &str, parent: &str, kind: u32| { + json!({ + "traceId": trace_id, + "spanId": span_id, + "parentSpanId": parent, + "name": "op", + "kind": kind, + "startTimeUnixNano": "1736480942444376000", + "endTimeUnixNano": "1736480942444499000", + "attributes": [], + "status": { "message": "", "code": 0 } + }) + }; + let resource_spans = |attributes: Vec, spans: Vec| { + json!({ + "resource": { "attributes": attributes }, + "scopeSpans": [{ "scope": { "name": "graph-tests" }, "spans": spans }], + "schemaUrl": "https://opentelemetry.io/schemas/1.4.0" + }) + }; + + // frontend: full identity (instance id + pod uid); cart: bare service — + // no instance entity; batch: instance id but no pod uid — no runs_on. + let req: ExportTraceServiceRequest = serde_json::from_value(json!({ + "resourceSpans": [ + resource_spans( + vec![ + make_string_attr("service.name", "frontend"), + make_string_attr("service.instance.id", "inst-1"), + make_string_attr("k8s.pod.uid", "uid-9"), + make_string_attr("k8s.pod.name", "api-1"), + make_string_attr("k8s.namespace.name", "default"), + ], + vec![span("c05d7a4ec8e1f231f02ed6e8da8655b4", "d24f921c75f68e23", "", 3)], + ), + resource_spans( + vec![make_string_attr("service.name", "cart")], + vec![span( + "c05d7a4ec8e1f231f02ed6e8da8655b4", + "9630f2916e2f7909", + "d24f921c75f68e23", + 2, + )], + ), + resource_spans( + vec![ + make_string_attr("service.name", "batch"), + make_string_attr("service.instance.id", "binst-1"), + ], + vec![span("cc9e0991a2e63d274984bd44ee669203", "8f847259b0f6e1ab", "", 1)], + ), + ] + })) + .unwrap(); + let res = send_trace_v1_req(&client, "graph_traces", req, false).await; + assert_eq!(StatusCode::OK, res.status()); + + let expected = r#"[["service","frontend","service","cart","calls"],["service.instance","resource_attributes.service.instance.id=binst-1,service_name=batch","service","batch","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","service","frontend","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","k8s.pod","uid-9","runs_on"]]"#; + validate_data( + "otlp_traces_entity_graph", + &client, + "select src_type, src_id, dst_type, dst_id, rel_type \ + from greptime_private.semantic_relationships \ + where observed_at >= '2025-01-01 00:00:00' order by rel_type, src_id;", + expected, + ) + .await; + + guard.remove_all().await; +} + pub async fn test_otlp_traces_v1(store_type: StorageType) { // init common_telemetry::init_default_ut_logging(); diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index 6617d9f839d..ed6a3d3f16f 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -44,47 +44,8 @@ drop table greptime_private.semantic_relationships; Error: 1001(Unsupported), Table `semantic_relationships` is read-only --- Read-time derivation: declare entity identities on a user table and observe --- them in the computed registry (single-column and composite ids, scope and --- descriptive columns, JSON output columns). -create table graph_app_latency ( - ts timestamp time index, - service_name string, - host string, - env string, - latency double, - primary key (service_name, host, env) -) with ( - 'greptime.semantic.entity.service.id' = 'service_name', - 'greptime.semantic.entity.service.scope' = 'env', - 'greptime.semantic.entity.process.id' = 'service_name,host', - 'greptime.semantic.entity.process.descriptive' = 'env' -); - -Affected Rows: 0 - -insert into graph_app_latency values (now(), 'cart', 'h1', 'us-east', 3.5); - -Affected Rows: 1 - --- SQLNESS PROTOCOL MYSQL -select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_tables -from greptime_private.semantic_entities -order by entity_type, entity_id; - -+-------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | -+-------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| process | host=h1,service_name=cart | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_latency"] | -| service | cart | | us-east | | ["public.graph_app_latency"] | -+-------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ - -drop table graph_app_latency; - -Affected Rows: 0 - --- Renaming a physical table INTO a reserved computed-table name would let the --- overlay shadow it, orphaning its data. +-- Renaming a physical table INTO a reserved name would let the overlay shadow +-- it (or squat the declared table's canonical schema), orphaning its data. create table greptime_private.graph_rename_probe (ts timestamp time index); Affected Rows: 0 @@ -93,13 +54,96 @@ alter table greptime_private.graph_rename_probe rename semantic_entities; Error: 1001(Unsupported), Table `semantic_entities` is read-only +alter table greptime_private.graph_rename_probe rename semantic_relationships_declared; + +Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) + drop table greptime_private.graph_rename_probe; Affected Rows: 0 --- Calls derivation: a trace-v1-model table with client/server span pairs. --- Trace tables need no entity declaration: the service entities are implicit. -create table graph_traces ( +-- Read-time derivation from one declaring metric table: entity identities in +-- the computed registry (single-column and composite ids, scope and +-- descriptive columns, JSON output columns), and the same rows witnessing the +-- co-declared vocabulary edges (runs_on / part_of, direction built in). +create table graph_app_metrics ( + ts timestamp time index, + service_name string, + instance string, + host string, + env string, + latency double, + primary key (service_name, instance, host, env) +) with ( + 'greptime.semantic.entity.service.id' = 'service_name', + 'greptime.semantic.entity.service.scope' = 'env', + 'greptime.semantic.entity.service.instance.id' = 'instance', + 'greptime.semantic.entity.host.id' = 'host', + 'greptime.semantic.entity.process.id' = 'service_name,host', + 'greptime.semantic.entity.process.descriptive' = 'env' +); + +Affected Rows: 0 + +insert into graph_app_metrics values (now(), 'cart', 'cart-0', 'h1', 'us-east', 3.5); + +Affected Rows: 1 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id; + ++------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ +| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | ++------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ +| host | h1 | | | | ["public.graph_app_metrics"] | +| process | host=h1,service_name=cart | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_metrics"] | +| service | cart | | us-east | | ["public.graph_app_metrics"] | +| service.instance | cart-0 | | | | ["public.graph_app_metrics"] | ++------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ + +select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by rel_type, src_id; + ++------------------+---------------------------+----------+--------+----------+------------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | ++------------------+---------------------------+----------+--------+----------+------------+------------+ +| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | +| service.instance | cart-0 | host | h1 | runs_on | attribute | 1.0 | +| process | host=h1,service_name=cart | host | h1 | runs_on | attribute | 1.0 | ++------------------+---------------------------+----------+--------+----------+------------+------------+ + +drop table graph_app_metrics; + +Affected Rows: 0 + +-- Calls derivation over trace-v1 tables, all branches in one scan: the client +-- span and its child server span land in different tables and still pair into +-- one edge with RED metrics; the epoch-timestamped pair falls outside the +-- default one-hour window; unmatched clients become virtual-node edges named +-- from span attributes (confidence < 1.0, connection type); a trace-model +-- table that lost the fixed span columns is skipped instead of failing the +-- scan. Trace tables need no entity declaration: service entities are +-- implicit. +create table graph_traces_a ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "span_attributes.peer.service" string, + "span_attributes.db.name" string, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +create table graph_traces_b ( "timestamp" timestamp(9) time index, trace_id string, span_id string, @@ -113,43 +157,68 @@ create table graph_traces ( Affected Rows: 0 --- The epoch-timestamped pair falls outside the default one-hour derivation --- window and must not appear in the results. -insert into graph_traces values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), - (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000), - (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), - (now(), 't2', 's2', 'c2', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'cart', 500000000), - (0, 't0', 'c0', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'stale-src', 0), - (0, 't0', 's0', 'c0', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'stale-dst', 100); +create table graph_traces_malformed ( + ts timestamp time index, + note string +) with ('table_data_model' = 'greptime_trace_v1'); + +Affected Rows: 0 + +insert into graph_traces_a values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0, NULL, NULL), + (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0, NULL, NULL), + (0, 't0', 'c0', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'stale-src', 0, NULL, NULL), + (0, 't0', 's0', 'c0', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'stale-dst', 100, NULL, NULL), + (now(), 't3', 'c3', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), + (now(), 't4', 'c4', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); Affected Rows: 6 +insert into graph_traces_b values + (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000), + (now(), 't2', 's2', 'c2', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'cart', 500000000); + +Affected Rows: 2 + +insert into graph_traces_malformed values (now(), 'not a trace'); + +Affected Rows: 1 + -- SQLNESS PROTOCOL MYSQL -select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence, +select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count, duration_sum, duration_count, attributes from greptime_private.semantic_relationships -order by src_id; +order by dst_id; -+----------+----------+----------+--------+----------+------------+------------+---------------+-------------+--------------+----------------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | request_count | error_count | duration_sum | duration_count | attributes | -+----------+----------+----------+--------+----------+------------+------------+---------------+-------------+--------------+----------------+------------+ -| service | frontend | service | cart | calls | trace | 1 | 2 | 1 | 2 | 2 | | -+----------+----------+----------+--------+----------+------------+------------+---------------+-------------+--------------+----------------+------------+ ++----------+-----------+----------+------------+------------+---------------+-------------+--------------+----------------+------------------------------------+ +| src_id | dst_id | rel_type | provenance | confidence | request_count | error_count | duration_sum | duration_count | attributes | ++----------+-----------+----------+------------+------------+---------------+-------------+--------------+----------------+------------------------------------+ +| frontend | cart | calls | trace | 1 | 2 | 1 | 2 | 2 | | +| frontend | orders-db | calls | trace | 0.5 | 1 | 0 | 0.1 | 1 | {"connection_type":"database"} | +| frontend | redis | calls | trace | 0.5 | 1 | 0 | 0.25 | 1 | {"connection_type":"virtual_node"} | ++----------+-----------+----------+------------+------------+---------------+-------------+--------------+----------------+------------------------------------+ -- SQLNESS PROTOCOL MYSQL select entity_type, entity_id, entity_id_attrs, scope, source_tables from greptime_private.semantic_entities order by entity_id; -+-------------+-----------+-----------------+-------+-------------------------+ -| entity_type | entity_id | entity_id_attrs | scope | source_tables | -+-------------+-----------+-----------------+-------+-------------------------+ -| service | cart | | | ["public.graph_traces"] | -| service | frontend | | | ["public.graph_traces"] | -+-------------+-----------+-----------------+-------+-------------------------+ ++-------------+-----------+-----------------+-------+---------------------------+ +| entity_type | entity_id | entity_id_attrs | scope | source_tables | ++-------------+-----------+-----------------+-------+---------------------------+ +| service | cart | | | ["public.graph_traces_b"] | +| service | frontend | | | ["public.graph_traces_a"] | ++-------------+-----------+-----------------+-------+---------------------------+ -drop table graph_traces; +drop table graph_traces_a; + +Affected Rows: 0 + +drop table graph_traces_b; + +Affected Rows: 0 + +drop table graph_traces_malformed; Affected Rows: 0 @@ -222,8 +291,7 @@ where observed_at < '2001-01-02 00:00:00'; Error: 3001(EngineExecuteQuery), Invalid SQL, error: the observed_at filter has no lower bound; the graph cannot derive over unbounded history — add e.g. observed_at >= '2026-01-01 00:00:00' -- The declared-edge table's definition is system-owned: user CREATE/ALTER are --- rejected and a physical table cannot be renamed into the name, while plain --- DML stays allowed. +-- rejected, while plain DML stays allowed. create table greptime_private.semantic_relationships_declared (ts timestamp time index); Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) @@ -232,18 +300,6 @@ alter table greptime_private.semantic_relationships_declared add column extra st Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) -create table greptime_private.declared_rename_probe (ts timestamp time index); - -Affected Rows: 0 - -alter table greptime_private.declared_rename_probe rename semantic_relationships_declared; - -Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) - -drop table greptime_private.declared_rename_probe; - -Affected Rows: 0 - delete from greptime_private.semantic_relationships_declared; Affected Rows: 4 @@ -277,161 +333,9 @@ drop table greptime_private.semantic_relationships_declared; Affected Rows: 0 --- Same-row co-declaration: a table declaring both entity types of a built-in --- vocabulary pair witnesses the edge on every row carrying both identities; --- the vocabulary fixes the direction. -create table graph_pod_metrics ( - ts timestamp time index, - instance string, - host string, - "service" string, - cpu double, - primary key (instance, host, "service") -) with ( - 'greptime.semantic.entity.service.instance.id' = 'instance', - 'greptime.semantic.entity.host.id' = 'host', - 'greptime.semantic.entity.service.id' = 'service' -); - -Affected Rows: 0 - -insert into graph_pod_metrics values (now(), 'cart-0', 'node-1', 'cart', 0.5); - -Affected Rows: 1 - -select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence -from greptime_private.semantic_relationships -order by rel_type, src_id; - -+------------------+--------+----------+--------+----------+------------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | -+------------------+--------+----------+--------+----------+------------+------------+ -| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | -| service.instance | cart-0 | host | node-1 | runs_on | attribute | 1.0 | -+------------------+--------+----------+--------+----------+------------+------------+ - -drop table graph_pod_metrics; - -Affected Rows: 0 - --- Cross-table calls pairing: the client span and its child server span land in --- different trace tables and must still pair into one edge. A trace-model --- table that lost the fixed span columns is skipped instead of failing the --- scan. -create table graph_traces_a ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -Affected Rows: 0 - -create table graph_traces_b ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -Affected Rows: 0 - -create table graph_traces_malformed ( - ts timestamp time index, - note string -) with ('table_data_model' = 'greptime_trace_v1'); - -Affected Rows: 0 - -insert into graph_traces_a values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0); - -Affected Rows: 1 - -insert into graph_traces_b values - (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000); - -Affected Rows: 1 - -insert into graph_traces_malformed values (now(), 'not a trace'); - -Affected Rows: 1 - -select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count -from greptime_private.semantic_relationships -order by src_id; - -+----------+--------+----------+------------+------------+---------------+-------------+ -| src_id | dst_id | rel_type | provenance | confidence | request_count | error_count | -+----------+--------+----------+------------+------------+---------------+-------------+ -| frontend | cart | calls | trace | 1.0 | 1 | 1 | -+----------+--------+----------+------------+------------+---------------+-------------+ - -drop table graph_traces_a; - -Affected Rows: 0 - -drop table graph_traces_b; - -Affected Rows: 0 - -drop table graph_traces_malformed; - -Affected Rows: 0 - --- Virtual nodes: a client span with no matching server span names its peer --- through span attributes; the edge carries confidence < 1.0 and the --- connection type. -create table graph_traces_virtual ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - "span_attributes.peer.service" string, - "span_attributes.db.name" string, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -Affected Rows: 0 - -insert into graph_traces_virtual values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), - (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); - -Affected Rows: 2 - --- SQLNESS PROTOCOL MYSQL -select src_id, dst_id, rel_type, confidence, request_count, attributes -from greptime_private.semantic_relationships -order by dst_id; - -+----------+-----------+----------+------------+---------------+------------------------------------+ -| src_id | dst_id | rel_type | confidence | request_count | attributes | -+----------+-----------+----------+------------+---------------+------------------------------------+ -| frontend | orders-db | calls | 0.5 | 1 | {"connection_type":"database"} | -| frontend | redis | calls | 0.5 | 1 | {"connection_type":"virtual_node"} | -+----------+-----------+----------+------------+---------------+------------------------------------+ - -drop table graph_traces_virtual; - -Affected Rows: 0 - -- Agent edges: span structure derives parent_agent calls agent, and span rows --- co-declaring agent+model / agent+tool witness uses / invokes. +-- co-declaring gen_ai.agent+gen_ai.model / gen_ai.agent+gen_ai.tool witness +-- uses / invokes. The identity columns are fields, not tags. create table graph_agent_traces ( "timestamp" timestamp(9) time index, trace_id string, @@ -444,13 +348,13 @@ create table graph_agent_traces ( agent_id string, model_name string, tool_name string, - primary key (service_name, agent_id, model_name, tool_name) + primary key (service_name) ) with ( 'table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true', - 'greptime.semantic.entity.agent.id' = 'agent_id', - 'greptime.semantic.entity.model.id' = 'model_name', - 'greptime.semantic.entity.tool.id' = 'tool_name' + 'greptime.semantic.entity.gen_ai.agent.id' = 'agent_id', + 'greptime.semantic.entity.gen_ai.model.id' = 'model_name', + 'greptime.semantic.entity.gen_ai.tool.id' = 'tool_name' ); Affected Rows: 0 @@ -466,15 +370,205 @@ select src_type, src_id, dst_type, dst_id, rel_type, provenance from greptime_private.semantic_relationships order by rel_type, dst_id; -+----------+--------------+----------+------------+----------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | -+----------+--------------+----------+------------+----------+------------+ -| agent | orchestrator | agent | researcher | calls | trace | -| agent | researcher | tool | web_search | invokes | trace | -| agent | researcher | model | gpt-5 | uses | trace | -+----------+--------------+----------+------------+----------+------------+ ++--------------+--------------+--------------+------------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++--------------+--------------+--------------+------------+----------+------------+ +| gen_ai.agent | orchestrator | gen_ai.agent | researcher | calls | trace | +| gen_ai.agent | researcher | gen_ai.tool | web_search | invokes | trace | +| gen_ai.agent | researcher | gen_ai.model | gpt-5 | uses | trace | ++--------------+--------------+--------------+------------+----------+------------+ drop table graph_agent_traces; Affected Rows: 0 +-- Prometheus conventions: whitelisted entity-descriptor metrics (stamped +-- signal_type=metric + source=prometheus by the remote-write path) get +-- implicit declarations, and the co-declared rules derive runs_on / part_of / +-- contains from their rows. Pod identity is the UID, so every KSM descriptor +-- lands on one entity; the empty labels kube-state-metrics emits (unscheduled +-- pod's node, ownerless pod's owner_*) identify nothing. A non-whitelisted +-- metric contributes nothing. +create table kube_pod_info ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + pod string, + node string, + greptime_value double, + primary key (uid, namespace, pod, node) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +create table kube_pod_owner ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + pod string, + owner_kind string, + owner_name string, + greptime_value double, + primary key (uid, namespace, pod, owner_kind, owner_name) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +create table kube_pod_container_info ( + greptime_timestamp timestamp(3) time index, + uid string, + container string, + image string, + greptime_value double, + primary key (uid, container, image) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +create table kube_service_info ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + "service" string, + cluster_ip string, + greptime_value double, + primary key (uid, namespace, "service", cluster_ip) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +create table target_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + k8s_cluster_name string, + greptime_value double, + primary key (job, instance, k8s_cluster_name) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +create table http_requests_total ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + greptime_value double, + primary key (job, instance) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +insert into kube_pod_info values + (now(), 'uid-1', 'default', 'api-1', 'node-a', 1), + (now(), 'uid-2', 'default', 'api-2', 'node-a', 1), + (now(), 'uid-3', 'default', 'pending-1', '', 1), + (now(), '', 'default', 'ghost', 'node-a', 1); + +Affected Rows: 4 + +insert into kube_pod_owner values + (now(), 'uid-1', 'default', 'api-1', 'ReplicaSet', 'api-rs', 1), + (now(), 'uid-2', 'default', 'api-2', 'ReplicaSet', 'api-rs', 1); + +Affected Rows: 2 + +insert into kube_pod_container_info values + (now(), 'uid-1', 'main', 'nginx:1.25', 1); + +Affected Rows: 1 + +insert into kube_service_info values + (now(), 'svc-uid-1', 'default', 'api-svc', '10.0.0.1', 1); + +Affected Rows: 1 + +insert into target_info values + (now(), 'shop/api', 'inst-1', 'prod', 1); + +Affected Rows: 1 + +insert into http_requests_total values + (now(), 'shop/api', 'inst-1', 42); + +Affected Rows: 1 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + ++------------------+-----------------------------------------------------------+------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+-----------------------------------------------------------+------------------------------------+ +| k8s.container | container=main,uid=uid-1 | ["public.kube_pod_container_info"] | +| k8s.node | node-a | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_owner"] | +| k8s.pod | uid-1 | ["public.kube_pod_container_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_owner"] | +| k8s.pod | uid-3 | ["public.kube_pod_info"] | +| k8s.service | svc-uid-1 | ["public.kube_service_info"] | +| k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | ["public.kube_pod_owner"] | +| service | shop/api | ["public.target_info"] | +| service.instance | instance=inst-1,job=shop/api | ["public.target_info"] | ++------------------+-----------------------------------------------------------+------------------------------------+ + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + ++------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ +| k8s.pod | uid-1 | k8s.container | container=main,uid=uid-1 | contains | attribute | +| service.instance | instance=inst-1,job=shop/api | service | shop/api | part_of | attribute | +| k8s.pod | uid-1 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | +| k8s.pod | uid-2 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | +| k8s.pod | uid-1 | k8s.node | node-a | runs_on | attribute | +| k8s.pod | uid-2 | k8s.node | node-a | runs_on | attribute | ++------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ + +drop table kube_pod_info; + +Affected Rows: 0 + +drop table kube_pod_owner; + +Affected Rows: 0 + +drop table kube_pod_container_info; + +Affected Rows: 0 + +drop table kube_service_info; + +Affected Rows: 0 + +drop table target_info; + +Affected Rows: 0 + +drop table http_requests_total; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index 9acb003331a..181a8da66a9 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -22,43 +22,73 @@ drop table greptime_private.semantic_entities; drop table greptime_private.semantic_relationships; --- Read-time derivation: declare entity identities on a user table and observe --- them in the computed registry (single-column and composite ids, scope and --- descriptive columns, JSON output columns). -create table graph_app_latency ( +-- Renaming a physical table INTO a reserved name would let the overlay shadow +-- it (or squat the declared table's canonical schema), orphaning its data. +create table greptime_private.graph_rename_probe (ts timestamp time index); + +alter table greptime_private.graph_rename_probe rename semantic_entities; + +alter table greptime_private.graph_rename_probe rename semantic_relationships_declared; + +drop table greptime_private.graph_rename_probe; + +-- Read-time derivation from one declaring metric table: entity identities in +-- the computed registry (single-column and composite ids, scope and +-- descriptive columns, JSON output columns), and the same rows witnessing the +-- co-declared vocabulary edges (runs_on / part_of, direction built in). +create table graph_app_metrics ( ts timestamp time index, service_name string, + instance string, host string, env string, latency double, - primary key (service_name, host, env) + primary key (service_name, instance, host, env) ) with ( 'greptime.semantic.entity.service.id' = 'service_name', 'greptime.semantic.entity.service.scope' = 'env', + 'greptime.semantic.entity.service.instance.id' = 'instance', + 'greptime.semantic.entity.host.id' = 'host', 'greptime.semantic.entity.process.id' = 'service_name,host', 'greptime.semantic.entity.process.descriptive' = 'env' ); -insert into graph_app_latency values (now(), 'cart', 'h1', 'us-east', 3.5); +insert into graph_app_metrics values (now(), 'cart', 'cart-0', 'h1', 'us-east', 3.5); -- SQLNESS PROTOCOL MYSQL select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_tables from greptime_private.semantic_entities order by entity_type, entity_id; -drop table graph_app_latency; +select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by rel_type, src_id; --- Renaming a physical table INTO a reserved computed-table name would let the --- overlay shadow it, orphaning its data. -create table greptime_private.graph_rename_probe (ts timestamp time index); +drop table graph_app_metrics; -alter table greptime_private.graph_rename_probe rename semantic_entities; +-- Calls derivation over trace-v1 tables, all branches in one scan: the client +-- span and its child server span land in different tables and still pair into +-- one edge with RED metrics; the epoch-timestamped pair falls outside the +-- default one-hour window; unmatched clients become virtual-node edges named +-- from span attributes (confidence < 1.0, connection type); a trace-model +-- table that lost the fixed span columns is skipped instead of failing the +-- scan. Trace tables need no entity declaration: service entities are +-- implicit. +create table graph_traces_a ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "span_attributes.peer.service" string, + "span_attributes.db.name" string, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); -drop table greptime_private.graph_rename_probe; - --- Calls derivation: a trace-v1-model table with client/server span pairs. --- Trace tables need no entity declaration: the service entities are implicit. -create table graph_traces ( +create table graph_traces_b ( "timestamp" timestamp(9) time index, trace_id string, span_id string, @@ -70,28 +100,41 @@ create table graph_traces ( primary key (service_name) ) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); --- The epoch-timestamped pair falls outside the default one-hour derivation --- window and must not appear in the results. -insert into graph_traces values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), +create table graph_traces_malformed ( + ts timestamp time index, + note string +) with ('table_data_model' = 'greptime_trace_v1'); + +insert into graph_traces_a values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0, NULL, NULL), + (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0, NULL, NULL), + (0, 't0', 'c0', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'stale-src', 0, NULL, NULL), + (0, 't0', 's0', 'c0', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'stale-dst', 100, NULL, NULL), + (now(), 't3', 'c3', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), + (now(), 't4', 'c4', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); + +insert into graph_traces_b values (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000), - (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), - (now(), 't2', 's2', 'c2', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'cart', 500000000), - (0, 't0', 'c0', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'stale-src', 0), - (0, 't0', 's0', 'c0', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'stale-dst', 100); + (now(), 't2', 's2', 'c2', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'cart', 500000000); + +insert into graph_traces_malformed values (now(), 'not a trace'); -- SQLNESS PROTOCOL MYSQL -select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence, +select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count, duration_sum, duration_count, attributes from greptime_private.semantic_relationships -order by src_id; +order by dst_id; -- SQLNESS PROTOCOL MYSQL select entity_type, entity_id, entity_id_attrs, scope, source_tables from greptime_private.semantic_entities order by entity_id; -drop table graph_traces; +drop table graph_traces_a; + +drop table graph_traces_b; + +drop table graph_traces_malformed; -- Declared edges: the physical table is created by the system on the first -- INSERT with its canonical schema. Re-asserting an edge stores a new revision; @@ -134,18 +177,11 @@ select src_id from greptime_private.semantic_relationships where observed_at < '2001-01-02 00:00:00'; -- The declared-edge table's definition is system-owned: user CREATE/ALTER are --- rejected and a physical table cannot be renamed into the name, while plain --- DML stays allowed. +-- rejected, while plain DML stays allowed. create table greptime_private.semantic_relationships_declared (ts timestamp time index); alter table greptime_private.semantic_relationships_declared add column extra string; -create table greptime_private.declared_rename_probe (ts timestamp time index); - -alter table greptime_private.declared_rename_probe rename semantic_relationships_declared; - -drop table greptime_private.declared_rename_probe; - delete from greptime_private.semantic_relationships_declared; select src_id from greptime_private.semantic_relationships order by src_id; @@ -162,111 +198,9 @@ select src_id from greptime_private.semantic_relationships order by src_id; drop table greptime_private.semantic_relationships_declared; --- Same-row co-declaration: a table declaring both entity types of a built-in --- vocabulary pair witnesses the edge on every row carrying both identities; --- the vocabulary fixes the direction. -create table graph_pod_metrics ( - ts timestamp time index, - instance string, - host string, - "service" string, - cpu double, - primary key (instance, host, "service") -) with ( - 'greptime.semantic.entity.service.instance.id' = 'instance', - 'greptime.semantic.entity.host.id' = 'host', - 'greptime.semantic.entity.service.id' = 'service' -); - -insert into graph_pod_metrics values (now(), 'cart-0', 'node-1', 'cart', 0.5); - -select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence -from greptime_private.semantic_relationships -order by rel_type, src_id; - -drop table graph_pod_metrics; - --- Cross-table calls pairing: the client span and its child server span land in --- different trace tables and must still pair into one edge. A trace-model --- table that lost the fixed span columns is skipped instead of failing the --- scan. -create table graph_traces_a ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -create table graph_traces_b ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -create table graph_traces_malformed ( - ts timestamp time index, - note string -) with ('table_data_model' = 'greptime_trace_v1'); - -insert into graph_traces_a values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0); - -insert into graph_traces_b values - (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000); - -insert into graph_traces_malformed values (now(), 'not a trace'); - -select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count -from greptime_private.semantic_relationships -order by src_id; - -drop table graph_traces_a; - -drop table graph_traces_b; - -drop table graph_traces_malformed; - --- Virtual nodes: a client span with no matching server span names its peer --- through span attributes; the edge carries confidence < 1.0 and the --- connection type. -create table graph_traces_virtual ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - "span_attributes.peer.service" string, - "span_attributes.db.name" string, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -insert into graph_traces_virtual values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), - (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); - --- SQLNESS PROTOCOL MYSQL -select src_id, dst_id, rel_type, confidence, request_count, attributes -from greptime_private.semantic_relationships -order by dst_id; - -drop table graph_traces_virtual; - -- Agent edges: span structure derives parent_agent calls agent, and span rows --- co-declaring agent+model / agent+tool witness uses / invokes. +-- co-declaring gen_ai.agent+gen_ai.model / gen_ai.agent+gen_ai.tool witness +-- uses / invokes. The identity columns are fields, not tags. create table graph_agent_traces ( "timestamp" timestamp(9) time index, trace_id string, @@ -279,13 +213,13 @@ create table graph_agent_traces ( agent_id string, model_name string, tool_name string, - primary key (service_name, agent_id, model_name, tool_name) + primary key (service_name) ) with ( 'table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true', - 'greptime.semantic.entity.agent.id' = 'agent_id', - 'greptime.semantic.entity.model.id' = 'model_name', - 'greptime.semantic.entity.tool.id' = 'tool_name' + 'greptime.semantic.entity.gen_ai.agent.id' = 'agent_id', + 'greptime.semantic.entity.gen_ai.model.id' = 'model_name', + 'greptime.semantic.entity.gen_ai.tool.id' = 'tool_name' ); insert into graph_agent_traces values @@ -298,3 +232,129 @@ from greptime_private.semantic_relationships order by rel_type, dst_id; drop table graph_agent_traces; + +-- Prometheus conventions: whitelisted entity-descriptor metrics (stamped +-- signal_type=metric + source=prometheus by the remote-write path) get +-- implicit declarations, and the co-declared rules derive runs_on / part_of / +-- contains from their rows. Pod identity is the UID, so every KSM descriptor +-- lands on one entity; the empty labels kube-state-metrics emits (unscheduled +-- pod's node, ownerless pod's owner_*) identify nothing. A non-whitelisted +-- metric contributes nothing. +create table kube_pod_info ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + pod string, + node string, + greptime_value double, + primary key (uid, namespace, pod, node) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +create table kube_pod_owner ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + pod string, + owner_kind string, + owner_name string, + greptime_value double, + primary key (uid, namespace, pod, owner_kind, owner_name) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +create table kube_pod_container_info ( + greptime_timestamp timestamp(3) time index, + uid string, + container string, + image string, + greptime_value double, + primary key (uid, container, image) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +create table kube_service_info ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + "service" string, + cluster_ip string, + greptime_value double, + primary key (uid, namespace, "service", cluster_ip) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +create table target_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + k8s_cluster_name string, + greptime_value double, + primary key (job, instance, k8s_cluster_name) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +create table http_requests_total ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + greptime_value double, + primary key (job, instance) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +insert into kube_pod_info values + (now(), 'uid-1', 'default', 'api-1', 'node-a', 1), + (now(), 'uid-2', 'default', 'api-2', 'node-a', 1), + (now(), 'uid-3', 'default', 'pending-1', '', 1), + (now(), '', 'default', 'ghost', 'node-a', 1); + +insert into kube_pod_owner values + (now(), 'uid-1', 'default', 'api-1', 'ReplicaSet', 'api-rs', 1), + (now(), 'uid-2', 'default', 'api-2', 'ReplicaSet', 'api-rs', 1); + +insert into kube_pod_container_info values + (now(), 'uid-1', 'main', 'nginx:1.25', 1); + +insert into kube_service_info values + (now(), 'svc-uid-1', 'default', 'api-svc', '10.0.0.1', 1); + +insert into target_info values + (now(), 'shop/api', 'inst-1', 'prod', 1); + +insert into http_requests_total values + (now(), 'shop/api', 'inst-1', 42); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + +drop table kube_pod_info; + +drop table kube_pod_owner; + +drop table kube_pod_container_info; + +drop table kube_service_info; + +drop table target_info; + +drop table http_requests_total; From d0fecdd6b0876b91f2730b70cdb291a6899d5d34 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 13 Aug 2026 06:57:46 +0000 Subject: [PATCH 060/126] fix(mito2): publish committed sequence only after rows are installed (#8862) The committed-sequence watermark must never cover rows that are not yet physically visible. Previously write_memtable() published next_sequence - 1 before bulk parts were installed, so a scan opening a snapshot could bind H to invisible sequences and permanently miss rows after checkpoint advance. Publish once, after both ordinary and bulk memtable writes complete, in the single-region fast path, the multi-region spawned tasks, and WAL replay; skip publication for contexts whose WAL entry could not be built. Add a deterministic worker-level race test using a cfg(test) bulk-install barrier proving the committed sequence stays put until installation. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/region/opener.rs | 3 + src/mito2/src/region_write_ctx.rs | 166 ++++++++++++++++++++++++++- src/mito2/src/worker/handle_write.rs | 164 +++++++++++++++++++++++++- 3 files changed, 326 insertions(+), 7 deletions(-) diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 5d2d36df107..23695ca424c 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -908,6 +908,9 @@ where region_write_ctx.set_next_entry_id(last_entry_id + 1); region_write_ctx.write_memtable().await; region_write_ctx.write_bulk().await; + // Publish the replayed sequences only after all rows (including bulk + // parts) are installed, matching the write path ordering. + region_write_ctx.publish_sequence_and_entry_id(); } // TODO(weny): We need to update `flushed_entry_id` in the region manifest diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index b71e3e5ab17..f64f9810f03 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -199,6 +199,12 @@ impl RegionWriteCtx { &self.version } + /// Returns the version control of the region. + #[cfg(test)] + pub(crate) fn version_control(&self) -> &VersionControlRef { + &self.version_control + } + /// Returns whether writes in this context should skip WAL. pub(crate) fn skip_wal(&self) -> bool { self.provider == Provider::Noop || self.version.options.skip_wal @@ -286,10 +292,6 @@ impl RegionWriteCtx { let bytes = new_memory_usage.saturating_sub(prev_memory_usage.unwrap_or_default()); written_bytes.fetch_add(bytes as u64, Ordering::Relaxed); } - // Updates region sequence and entry id. Since we stores last sequence and entry id in region, we need - // to decrease `next_sequence` and `next_entry_id` by 1. - self.version_control - .set_sequence_and_entry_id(self.next_sequence - 1, self.next_entry_id - 1); } pub(crate) fn push_bulk( @@ -324,6 +326,8 @@ impl RegionWriteCtx { if self.failed || self.bulk_parts.is_empty() { return; } + #[cfg(test)] + test_hooks::pause_before_bulk_install(self.region_id).await; let _timer = metrics::REGION_WORKER_HANDLE_WRITE_ELAPSED .with_label_values(&["write_bulk"]) .start_timer(); @@ -369,11 +373,165 @@ impl RegionWriteCtx { let bytes = new_memory_usage.saturating_sub(prev_memory_usage.unwrap_or_default()); written_bytes.fetch_add(bytes as u64, Ordering::Relaxed); } + } + + /// Publishes the sequences and entry id assigned by this context to the + /// region's committed watermark. + /// + /// Must be called after both [`write_memtable`](Self::write_memtable) and + /// [`write_bulk`](Self::write_bulk) have completed, so the committed + /// sequence never covers rows that are not yet physically installed in the + /// memtable (bulk part sequences are assigned in [`push_bulk`](Self::push_bulk) + /// before installation). Since we store the last sequence and entry id in + /// the region, we decrease `next_sequence` and `next_entry_id` by 1. + /// + /// If the write operation failed (e.g. the WAL entry could not be built), + /// no rows were installed and nothing must be published: advancing the + /// committed watermark here would make it cover rows that were never + /// written. + pub(crate) fn publish_sequence_and_entry_id(&self) { + if self.failed { + return; + } self.version_control .set_sequence_and_entry_id(self.next_sequence - 1, self.next_entry_id - 1); } } +/// Test-only hooks to make write ordering races deterministic. +#[cfg(test)] +pub(crate) mod test_hooks { + use std::sync::Mutex; + use std::sync::atomic::{AtomicU64, Ordering}; + + use store_api::storage::RegionId; + use tokio::sync::watch; + + /// Channels of an armed bulk-install barrier. + /// + /// `reached` signals that a bulk write paused between ordinary-memtable + /// handling and bulk installation; `release` unblocks it. Dropping the + /// senders (by disarming the barrier) also unblocks paused writes because + /// their `wait_for` fails on a closed channel. + struct ActiveBarrier { + id: u64, + /// Only bulk writes for this region pause at the barrier; writes for + /// other regions pass through untouched, so concurrently running tests + /// can never satisfy (or be blocked by) each other's barrier. + target_region_id: RegionId, + reached: watch::Sender, + release: watch::Sender, + } + + /// The currently armed barrier, if any. Wrapped in a `Mutex` so a test can + /// arm a fresh barrier after a previous one was released (or dropped), and + /// so a barrier can never leak past the test that owns it. + static ACTIVE_BARRIER: Mutex> = Mutex::new(None); + static NEXT_BARRIER_ID: AtomicU64 = AtomicU64::new(1); + + fn lock_active_barrier() -> std::sync::MutexGuard<'static, Option> { + // Never let a poisoned mutex (e.g. a panic in another test while + // holding the lock) hang or break unrelated tests. + ACTIVE_BARRIER + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + /// RAII guard for the bulk-install barrier. + /// + /// The guard owns the armed barrier. Releasing it — or dropping the guard, + /// e.g. when a test panics — unblocks any write paused at the barrier and + /// disarms it, so a paused write can never hang and later tests can arm a + /// fresh barrier. + pub(crate) struct BulkInstallBarrier { + id: u64, + reached_rx: watch::Receiver, + release_tx: watch::Sender, + released: bool, + } + + impl BulkInstallBarrier { + /// Waits until a bulk write paused at the barrier. + pub(crate) async fn wait_until_reached(&mut self) { + if !*self.reached_rx.borrow() { + let _ = self.reached_rx.wait_for(|reached| *reached).await; + } + } + + /// Unblocks the paused write and disarms the barrier. + pub(crate) fn release(&mut self) { + if self.released { + return; + } + self.released = true; + // Flip the release value so writes paused on this barrier proceed. + let _ = self.release_tx.send(true); + disarm_barrier(self.id); + } + } + + impl Drop for BulkInstallBarrier { + fn drop(&mut self) { + self.release(); + } + } + + /// Arms the bulk-install barrier for `target_region_id` and returns a + /// guard that owns it. + /// + /// Any previously armed barrier is replaced: its senders are dropped, which + /// unblocks writes that were paused on it instead of letting them hang. + pub(crate) fn arm_bulk_install_barrier(target_region_id: RegionId) -> BulkInstallBarrier { + let (reached_tx, reached_rx) = watch::channel(false); + let (release_tx, _release_rx) = watch::channel(false); + let id = NEXT_BARRIER_ID.fetch_add(1, Ordering::Relaxed); + let mut active = lock_active_barrier(); + *active = Some(ActiveBarrier { + id, + target_region_id, + reached: reached_tx, + release: release_tx.clone(), + }); + BulkInstallBarrier { + id, + reached_rx, + release_tx, + released: false, + } + } + + /// Removes the barrier with `id` from the statics if it is still active. + fn disarm_barrier(id: u64) { + let mut active = lock_active_barrier(); + if active.as_ref().is_some_and(|barrier| barrier.id == id) { + *active = None; + } + } + + /// Pauses a bulk write for `region_id` before installing its parts until + /// the test releases the barrier (or it is disarmed). Bulk writes for + /// other regions return immediately. + pub(crate) async fn pause_before_bulk_install(region_id: RegionId) { + let (reached_tx, release_rx) = { + let active = lock_active_barrier(); + match active.as_ref() { + Some(barrier) if barrier.target_region_id == region_id => { + (barrier.reached.clone(), barrier.release.subscribe()) + } + _ => return, + } + }; + // Signal that a bulk write reached the pause point. + let _ = reached_tx.send(true); + let mut release_rx = release_rx; + if !*release_rx.borrow() { + // The sender is dropped when the barrier is disarmed, which makes + // `wait_for` return an error instead of hanging forever. + let _ = release_rx.wait_for(|released| *released).await; + } + } +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 3be40553145..0c2eecdea20 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -128,6 +128,10 @@ impl RegionWorkerLoop { let mut region_ctx = region_ctxs.into_values().next().unwrap(); region_ctx.write_memtable().await; region_ctx.write_bulk().await; + // Publish only after all rows (including bulk parts) are + // physically installed, so a scan opening on the committed + // sequence can never bind H to invisible rows. + region_ctx.publish_sequence_and_entry_id(); put_rows += region_ctx.put_num; delete_rows += region_ctx.delete_num; } else { @@ -138,6 +142,9 @@ impl RegionWorkerLoop { common_runtime::spawn_global(async move { region_ctx.write_memtable().await; region_ctx.write_bulk().await; + // The spawned task owns the moved ctx, so publish + // inside the task after the memtable writes. + region_ctx.publish_sequence_and_entry_id(); (region_ctx.put_num, region_ctx.delete_num) }) }) @@ -734,7 +741,12 @@ fn check_partition_expr_version( #[cfg(test)] mod tests { - use api::v1::{Row, Rows}; + use api::v1::helper::{tag_column_schema, time_index_column_schema}; + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row, Rows}; + use common_recordbatch::DfRecordBatch; + use datatypes::arrow::array::{ArrayRef, StringArray, TimestampMillisecondArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; use futures::stream; use log_store::error::{ Error as LogStoreError, IllegalStateSnafu, InvalidProviderSnafu, Result as LogStoreResult, @@ -746,9 +758,42 @@ mod tests { use tokio::sync::oneshot; use super::*; + use crate::memtable::bulk::part::BulkPart; use crate::request::OptionOutputTx; + use crate::test_util::ts_ms_value; use crate::test_util::version_util::VersionControlBuilder; + /// Creates a bulk part with `num_rows` rows. The schema carries the + /// builder metadata's columns (`tag_0` primary key + `ts` time index) so + /// the bulk-install conversion succeeds; `timestamp_index` points at the + /// `ts` column. + fn new_bulk_part(num_rows: i64) -> BulkPart { + let schema = Arc::new(Schema::new(vec![ + Field::new("tag_0", DataType::Utf8, true), + Field::new( + "ts", + DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + ])); + let tag = Arc::new(StringArray::from_iter_values( + (0..num_rows).map(|value| value.to_string()), + )) as ArrayRef; + let ts = Arc::new(TimestampMillisecondArray::from( + (0..num_rows).collect::>(), + )) as ArrayRef; + let batch = DfRecordBatch::try_new(schema, vec![tag, ts]).unwrap(); + + BulkPart { + batch, + max_timestamp: num_rows - 1, + min_timestamp: 0, + sequence: 0, + timestamp_index: 1, + raw_data: None, + } + } + /// A log store that fails to build entries for `failing_region` and fails the /// whole batch when `fail_append` is true. #[derive(Debug, Default)] @@ -857,8 +902,18 @@ mod tests { ctx.push_mutation( OpType::Put as i32, Some(Rows { - schema: vec![], - rows: vec![Row { values: vec![] }], + schema: vec![ + time_index_column_schema("ts", ColumnDataType::TimestampMillisecond), + tag_column_schema("tag_0", ColumnDataType::String), + ], + rows: vec![Row { + values: vec![ + ts_ms_value(0), + api::v1::Value { + value_data: Some(ValueData::StringValue("a".to_string())), + }, + ], + }], }), None, OptionOutputTx::from(tx), @@ -878,8 +933,10 @@ mod tests { let mut region_ctxs = HashMap::new(); let (ctx, failing_rx) = new_region_ctx(failing_region); + let failing_committed_sequence = ctx.version_control().committed_sequence(); region_ctxs.insert(failing_region, ctx); let (ctx, ok_rx) = new_region_ctx(ok_region); + let ok_committed_sequence = ctx.version_control().committed_sequence(); region_ctxs.insert(ok_region, ctx); let entry_id = region_ctxs[&ok_region].next_entry_id(); @@ -890,12 +947,113 @@ mod tests { assert!(!region_ctxs[&ok_region].is_failed()); assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); + // Run the memtable and publication phase the worker runs after `write_wal`. + // The failed region installed no rows, so its committed sequence must not + // advance; the successful sibling region advances normally. + for region_ctx in region_ctxs.values_mut() { + region_ctx.write_memtable().await; + region_ctx.write_bulk().await; + region_ctx.publish_sequence_and_entry_id(); + } + + assert_eq!( + failing_committed_sequence, + region_ctxs[&failing_region] + .version_control() + .committed_sequence() + ); + assert_eq!( + ok_committed_sequence + 1, + region_ctxs[&ok_region] + .version_control() + .committed_sequence() + ); + // Waiters of the failed region get the error while others get the result. drop(region_ctxs); assert!(failing_rx.await.unwrap().is_err()); assert_eq!(1, ok_rx.await.unwrap().unwrap()); } + // The committed sequence must not be published before the bulk part's rows + // are physically installed in the memtable: publication happens strictly + // after `write_bulk` returns. Armed with the bulk-install test barrier, + // this deterministically pauses the worker's memtable phase between the + // ordinary-memtable handling and the bulk installation and asserts the + // region's committed sequence is still its initial value (0) — read via + // the `version_control()` accessor, never through a scanner API. After the + // barrier is released, the committed sequence must advance to cover the + // bulk rows. + #[tokio::test] + async fn test_bulk_write_sequence_not_committed_before_install_worker_level() { + let region_id = RegionId::new(1, 1); + let version_control = Arc::new(VersionControlBuilder::new().build()); + assert_eq!( + 0, + version_control.committed_sequence(), + "the builder must start at sequence 0" + ); + + let mut region_ctxs = HashMap::new(); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + let (tx, rx) = oneshot::channel(); + // Push a 3-row bulk part without an explicit sequence, exactly like the + // worker's `process_bulk_requests`. + assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(3), None)); + region_ctxs.insert(region_id, ctx); + + // Run the WAL phase the worker runs before the memtable phase. + let wal = Wal::new(Arc::new(MockLogStore::default())); + assert!(write_wal(&wal, &mut region_ctxs).await); + assert!(!region_ctxs[®ion_id].is_failed()); + + // Arm the bulk-install barrier: `write_bulk` pauses right before the + // parts are physically installed into the memtable. + let mut barrier = crate::region_write_ctx::test_hooks::arm_bulk_install_barrier(region_id); + + // Run the worker's memtable and publication phases in the background. + let write_handle = tokio::spawn(async move { + let mut region_ctx = region_ctxs.remove(®ion_id).unwrap(); + region_ctx.write_memtable().await; + region_ctx.write_bulk().await; + region_ctx.publish_sequence_and_entry_id(); + }); + + // Wait until the write paused at the barrier: deterministic, no sleeps. + tokio::time::timeout( + std::time::Duration::from_secs(10), + barrier.wait_until_reached(), + ) + .await + .expect("bulk write never reached the install barrier"); + + // The committed sequence must not have advanced yet: publication + // happens strictly after the bulk part is installed. + assert_eq!( + 0, + version_control.committed_sequence(), + "committed sequence leaked before the bulk part was installed" + ); + + // Release the barrier: the bulk part installs and the committed + // sequence advances to cover all 3 bulk rows. + barrier.release(); + write_handle.await.expect("bulk write should complete"); + assert_eq!( + 3, + version_control.committed_sequence(), + "committed sequence must cover the installed bulk rows" + ); + + // The bulk waiter is notified with the number of installed rows. + assert_eq!(3, rx.await.unwrap().unwrap()); + } + #[tokio::test] async fn test_write_wal_all_regions_failed_to_build_entries() { let failing_region = RegionId::new(1, 1); From 1af4c335242e7e5e06ee91bc428a71bb631d0798 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 13 Aug 2026 07:29:00 +0000 Subject: [PATCH 061/126] refactor(event): separate procedure submission context (#8856) * refactor(event): separate procedure submission context Signed-off-by: WenyXu * fix(event): map extensions and forward GC context Signed-off-by: WenyXu * fix(gc): initialize integration test context Signed-off-by: WenyXu * fix(test): pass procedure context to DDL helpers Signed-off-by: WenyXu * refactor: simplify procedure submission contexts Signed-off-by: WenyXu * refactor(event): separate procedure and query contexts Signed-off-by: WenyXu * fix(event): clarify procedure context propagation Signed-off-by: WenyXu * fix(event): preserve procedure submission context Signed-off-by: WenyXu * refactor(event): move DDL context by value Signed-off-by: WenyXu * fix(test): retain manual GC event context Signed-off-by: WenyXu * refactor(event): tighten procedure context API Signed-off-by: WenyXu * chore: update greptime-proto Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 3 +- Cargo.toml | 2 +- src/common/event-recorder/src/context.rs | 119 ++++++ src/common/event-recorder/src/lib.rs | 2 +- src/common/function/src/admin/gc.rs | 52 ++- .../function/src/admin/migrate_region.rs | 17 +- src/common/function/src/admin/purge_table.rs | 14 +- src/common/function/src/handlers.rs | 20 +- src/common/function/src/state.rs | 15 +- src/common/meta/src/ddl/tests/event/flow.rs | 12 +- src/common/meta/src/ddl/tests/event/view.rs | 12 +- src/common/meta/src/ddl_manager.rs | 388 +++++++++++++----- src/common/meta/src/procedure_executor.rs | 16 +- src/common/meta/src/rpc/ddl.rs | 146 +------ src/common/procedure/src/local.rs | 83 +++- src/common/procedure/src/procedure.rs | 12 +- src/common/procedure/src/store.rs | 8 + src/frontend/src/instance.rs | 4 +- src/meta-client/Cargo.toml | 1 + src/meta-client/src/client.rs | 80 ++-- src/meta-client/src/client/procedure.rs | 227 ++++++++-- src/meta-client/src/error.rs | 7 + src/meta-srv/src/gc/ctx.rs | 13 +- src/meta-srv/src/gc/handler.rs | 18 +- src/meta-srv/src/gc/mock.rs | 10 +- src/meta-srv/src/gc/mock/basic.rs | 21 +- src/meta-srv/src/gc/mock/concurrent.rs | 12 +- src/meta-srv/src/gc/mock/full_list.rs | 12 +- src/meta-srv/src/gc/mock/misc.rs | 6 +- src/meta-srv/src/gc/scheduler.rs | 35 +- .../src/procedure/region_migration.rs | 4 +- .../src/procedure/region_migration/manager.rs | 103 ++++- src/meta-srv/src/service/procedure.rs | 195 +++++++-- src/operator/src/procedure.rs | 73 ++-- src/operator/src/statement/comment.rs | 19 +- src/operator/src/statement/ddl.rs | 227 +++++----- src/operator/src/utils.rs | 63 ++- tests-integration/tests/gc_event.rs | 10 +- tests-integration/tests/region_migration.rs | 174 +++++--- tests-integration/tests/repartition.rs | 6 +- tests-integration/tests/table_ddl_event.rs | 16 +- 41 files changed, 1487 insertions(+), 770 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index dcd23451407..97b2fc76a79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6039,7 +6039,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=e7be20ff855b7522efe235383d82238d56e676c3#e7be20ff855b7522efe235383d82238d56e676c3" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=32f467fa2ba2b3588a58381a24af83de09fbb00a#32f467fa2ba2b3588a58381a24af83de09fbb00a" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", @@ -8312,6 +8312,7 @@ dependencies = [ "async-trait", "common-base", "common-error", + "common-event-recorder", "common-grpc", "common-macro", "common-meta", diff --git a/Cargo.toml b/Cargo.toml index 0e6da7888d3..63dbc76f383 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -159,7 +159,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "e7be20ff855b7522efe235383d82238d56e676c3" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "32f467fa2ba2b3588a58381a24af83de09fbb00a" } hex = "0.4" http = "1" humantime = "2.1" diff --git a/src/common/event-recorder/src/context.rs b/src/common/event-recorder/src/context.rs index 601276967e3..1220c25419f 100644 --- a/src/common/event-recorder/src/context.rs +++ b/src/common/event-recorder/src/context.rs @@ -12,9 +12,51 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; + +use api::v1::meta::ProcedureEventContext as PbProcedureEventContext; use serde::{Deserialize, Serialize}; use strum::{AsRefStr, EnumString}; +/// Event metadata supplied when a procedure is submitted. +/// +/// Protocol is deliberately absent. The trusted submission adapter derives it +/// from the typed query channel before serializing the protobuf request. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ProcedureEventInput { + pub reason: TriggerReason, + pub extensions: HashMap, +} + +impl ProcedureEventInput { + /// Creates procedure event input with a trigger reason and no extensions. + pub fn new(reason: TriggerReason) -> Self { + Self { + reason, + extensions: Default::default(), + } + } +} + +impl From<&ProcedureEventInput> for PbProcedureEventContext { + fn from(input: &ProcedureEventInput) -> Self { + Self { + reason: input.reason.as_ref().to_string(), + protocol: String::new(), + extensions: input.extensions.clone(), + } + } +} + +impl From for ProcedureEventInput { + fn from(context: PbProcedureEventContext) -> Self { + Self { + reason: TriggerReason::from_extension(&context.reason), + extensions: context.extensions, + } + } +} + /// Stable context recorded for a procedure event. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct PersistentEventContext { @@ -42,12 +84,40 @@ impl PersistentEventContext { } } +impl From<(ProcedureEventInput, Option)> for PersistentEventContext { + fn from((input, protocol): (ProcedureEventInput, Option)) -> Self { + Self { + reason: input.reason, + protocol, + extensions: input + .extensions + .into_iter() + .map(|(key, value)| (key, serde_json::Value::String(value))) + .collect(), + } + } +} + impl Default for PersistentEventContext { fn default() -> Self { Self::new(TriggerReason::default()) } } +impl From for PersistentEventContext { + fn from(context: PbProcedureEventContext) -> Self { + Self { + reason: TriggerReason::from_extension(&context.reason), + protocol: (!context.protocol.is_empty()).then_some(context.protocol), + extensions: context + .extensions + .into_iter() + .map(|(key, value)| (key, serde_json::Value::String(value))) + .collect(), + } + } +} + /// Stable classification of a procedure trigger. #[derive( Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, AsRefStr, EnumString, @@ -96,6 +166,55 @@ mod tests { ); } + #[test] + fn test_event_context_from_protobuf() { + let protobuf = PbProcedureEventContext { + reason: "auto_create".to_string(), + protocol: "postgres".to_string(), + extensions: HashMap::from([ + ("source".to_string(), "sql".to_string()), + ("tenant".to_string(), "a".to_string()), + ]), + }; + assert_eq!( + PersistentEventContext::from(protobuf), + PersistentEventContext { + reason: TriggerReason::AutoCreate, + protocol: Some("postgres".to_string()), + extensions: serde_json::Map::from_iter([ + ("source".to_string(), json!("sql")), + ("tenant".to_string(), json!("a")), + ]), + } + ); + } + + #[test] + fn test_event_input_protobuf_has_no_protocol() { + let input = ProcedureEventInput { + reason: TriggerReason::AutoCreate, + extensions: HashMap::from([("source".to_string(), "sql".to_string())]), + }; + + assert_eq!( + PbProcedureEventContext::from(&input), + PbProcedureEventContext { + reason: "auto_create".to_string(), + protocol: String::new(), + extensions: input.extensions.clone(), + } + ); + + assert_eq!( + ProcedureEventInput::from(PbProcedureEventContext { + reason: "auto_create".to_string(), + protocol: "untrusted".to_string(), + extensions: input.extensions.clone(), + }), + input + ); + } + #[test] fn test_trigger_reason_deserializes_unknown_value() { let reason: TriggerReason = serde_json::from_str("\"future_reason\"").unwrap(); diff --git a/src/common/event-recorder/src/lib.rs b/src/common/event-recorder/src/lib.rs index 5b9cb4a7b97..7ee71722229 100644 --- a/src/common/event-recorder/src/lib.rs +++ b/src/common/event-recorder/src/lib.rs @@ -22,5 +22,5 @@ pub mod recorder; #[cfg(any(test, feature = "testing"))] pub mod testing; -pub use context::{PersistentEventContext, TriggerReason}; +pub use context::{PersistentEventContext, ProcedureEventInput, TriggerReason}; pub use recorder::*; diff --git a/src/common/function/src/admin/gc.rs b/src/common/function/src/admin/gc.rs index e9ba33a050b..0090140a460 100644 --- a/src/common/function/src/admin/gc.rs +++ b/src/common/function/src/admin/gc.rs @@ -38,17 +38,20 @@ const DEFAULT_FULL_FILE_LISTING: bool = false; )] pub(crate) async fn gc_regions( procedure_service_handler: &ProcedureServiceHandlerRef, - _ctx: &QueryContextRef, + query_ctx: &QueryContextRef, params: &[ValueRef<'_>], ) -> Result { let (region_ids, full_file_listing) = parse_gc_regions_params(params)?; let resp = procedure_service_handler - .gc_regions(GcRegionsRequest { - region_ids, - full_file_listing, - timeout: None, - }) + .gc_regions( + query_ctx.clone(), + GcRegionsRequest { + region_ids, + full_file_listing, + timeout: None, + }, + ) .await?; Ok(Value::from(resp.processed_regions)) @@ -69,13 +72,16 @@ pub(crate) async fn gc_table( parse_gc_table_params(params, query_ctx)?; let resp = procedure_service_handler - .gc_table(GcTableRequest { - catalog_name, - schema_name, - table_name, - full_file_listing, - timeout: None, - }) + .gc_table( + query_ctx.clone(), + GcTableRequest { + catalog_name, + schema_name, + table_name, + full_file_listing, + timeout: None, + }, + ) .await?; Ok(Value::from(resp.processed_regions)) @@ -268,13 +274,17 @@ mod tests { impl ProcedureServiceHandler for MockProcedureServiceHandler { async fn purge_table( &self, - _table_name: table::table_name::TableName, _query_ctx: QueryContextRef, + _table_name: table::table_name::TableName, ) -> Result<()> { unreachable!() } - async fn migrate_region(&self, _request: MigrateRegionRequest) -> Result> { + async fn migrate_region( + &self, + _query_ctx: QueryContextRef, + _request: MigrateRegionRequest, + ) -> Result> { unreachable!() } @@ -297,12 +307,20 @@ mod tests { unreachable!() } - async fn gc_regions(&self, request: GcRegionsRequest) -> Result { + async fn gc_regions( + &self, + _query_ctx: QueryContextRef, + request: GcRegionsRequest, + ) -> Result { *self.gc_regions_request.lock().unwrap() = Some(request); Ok(GcResponse::default()) } - async fn gc_table(&self, request: GcTableRequest) -> Result { + async fn gc_table( + &self, + _query_ctx: QueryContextRef, + request: GcTableRequest, + ) -> Result { *self.gc_table_request.lock().unwrap() = Some(request); Ok(GcResponse::default()) } diff --git a/src/common/function/src/admin/migrate_region.rs b/src/common/function/src/admin/migrate_region.rs index d958b3cce95..5d258e14946 100644 --- a/src/common/function/src/admin/migrate_region.rs +++ b/src/common/function/src/admin/migrate_region.rs @@ -47,7 +47,7 @@ const DEFAULT_TIMEOUT_SECS: u64 = 300; )] pub(crate) async fn migrate_region( procedure_service_handler: &ProcedureServiceHandlerRef, - _ctx: &QueryContextRef, + query_ctx: &QueryContextRef, params: &[ValueRef<'_>], ) -> Result { let (region_id, from_peer, to_peer, timeout) = match params.len() { @@ -82,12 +82,15 @@ pub(crate) async fn migrate_region( match (region_id, from_peer, to_peer, timeout) { (Some(region_id), Some(from_peer), Some(to_peer), Some(timeout)) => { let pid = procedure_service_handler - .migrate_region(MigrateRegionRequest { - region_id, - from_peer, - to_peer, - timeout: Duration::from_secs(timeout), - }) + .migrate_region( + query_ctx.clone(), + MigrateRegionRequest { + region_id, + from_peer, + to_peer, + timeout: Duration::from_secs(timeout), + }, + ) .await?; match pid { diff --git a/src/common/function/src/admin/purge_table.rs b/src/common/function/src/admin/purge_table.rs index 59613e155f0..7e5fcd84371 100644 --- a/src/common/function/src/admin/purge_table.rs +++ b/src/common/function/src/admin/purge_table.rs @@ -64,8 +64,8 @@ pub(crate) async fn purge_table( procedure_service_handler .purge_table( - TableName::new(catalog_name, schema_name, table_name), query_ctx.clone(), + TableName::new(catalog_name, schema_name, table_name), ) .await?; Ok(Value::from(0_u64)) @@ -131,8 +131,8 @@ mod tests { impl ProcedureServiceHandler for RecordingHandler { async fn purge_table( &self, - table_name: TableName, query_ctx: QueryContextRef, + table_name: TableName, ) -> Result<()> { if self.fail { return InvalidFuncArgsSnafu { @@ -144,7 +144,11 @@ mod tests { Ok(()) } - async fn migrate_region(&self, _: MigrateRegionRequest) -> Result> { + async fn migrate_region( + &self, + _: QueryContextRef, + _: MigrateRegionRequest, + ) -> Result> { unreachable!() } async fn reconcile(&self, _: ReconcileRequest) -> Result> { @@ -159,10 +163,10 @@ mod tests { fn catalog_manager(&self) -> &CatalogManagerRef { unreachable!() } - async fn gc_regions(&self, _: GcRegionsRequest) -> Result { + async fn gc_regions(&self, _: QueryContextRef, _: GcRegionsRequest) -> Result { unreachable!() } - async fn gc_table(&self, _: GcTableRequest) -> Result { + async fn gc_table(&self, _: QueryContextRef, _: GcTableRequest) -> Result { unreachable!() } } diff --git a/src/common/function/src/handlers.rs b/src/common/function/src/handlers.rs index 24ce115a0d6..6ee8064d6c7 100644 --- a/src/common/function/src/handlers.rs +++ b/src/common/function/src/handlers.rs @@ -89,10 +89,14 @@ pub trait TableMutationHandler: Send + Sync { #[async_trait] pub trait ProcedureServiceHandler: Send + Sync { /// Permanently purge a dropped table. - async fn purge_table(&self, table_name: TableName, query_ctx: QueryContextRef) -> Result<()>; + async fn purge_table(&self, query_ctx: QueryContextRef, table_name: TableName) -> Result<()>; /// Migrate a region from source peer to target peer, returns the procedure id if success. - async fn migrate_region(&self, request: MigrateRegionRequest) -> Result>; + async fn migrate_region( + &self, + query_ctx: QueryContextRef, + request: MigrateRegionRequest, + ) -> Result>; /// Reconcile a table, database or catalog, returns the procedure id if success. async fn reconcile(&self, request: ReconcileRequest) -> Result>; @@ -107,10 +111,18 @@ pub trait ProcedureServiceHandler: Send + Sync { fn catalog_manager(&self) -> &CatalogManagerRef; /// Manually trigger GC for specific regions. - async fn gc_regions(&self, request: MetaGcRegionsRequest) -> Result; + async fn gc_regions( + &self, + query_ctx: QueryContextRef, + request: MetaGcRegionsRequest, + ) -> Result; /// Manually trigger GC for a table. - async fn gc_table(&self, request: MetaGcTableRequest) -> Result; + async fn gc_table( + &self, + query_ctx: QueryContextRef, + request: MetaGcTableRequest, + ) -> Result; } /// This flow service handler is only use for flush flow for now. diff --git a/src/common/function/src/state.rs b/src/common/function/src/state.rs index b1802b62578..2b9f3a517a5 100644 --- a/src/common/function/src/state.rs +++ b/src/common/function/src/state.rs @@ -60,14 +60,15 @@ impl FunctionState { impl ProcedureServiceHandler for MockProcedureServiceHandler { async fn purge_table( &self, - _table_name: table::table_name::TableName, _query_ctx: QueryContextRef, + _table_name: table::table_name::TableName, ) -> Result<()> { Ok(()) } async fn migrate_region( &self, + _ctx: QueryContextRef, _request: MigrateRegionRequest, ) -> Result> { Ok(Some("test_pid".to_string())) @@ -92,7 +93,11 @@ impl FunctionState { Ok(()) } - async fn gc_regions(&self, _request: GcRegionsRequest) -> Result { + async fn gc_regions( + &self, + _context: QueryContextRef, + _request: GcRegionsRequest, + ) -> Result { Ok(GcResponse { processed_regions: 1, need_retry_regions: vec![], @@ -101,7 +106,11 @@ impl FunctionState { }) } - async fn gc_table(&self, _request: GcTableRequest) -> Result { + async fn gc_table( + &self, + _context: QueryContextRef, + _request: GcTableRequest, + ) -> Result { Ok(GcResponse { processed_regions: 1, need_retry_regions: vec![], diff --git a/src/common/meta/src/ddl/tests/event/flow.rs b/src/common/meta/src/ddl/tests/event/flow.rs index aae3613a087..cec0162adf8 100644 --- a/src/common/meta/src/ddl/tests/event/flow.rs +++ b/src/common/meta/src/ddl/tests/event/flow.rs @@ -19,9 +19,8 @@ use api::v1::{ColumnSchema, Row, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, - PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, - jsonb_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, jsonb_value, + procedure_event_column_schemas, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; @@ -240,12 +239,7 @@ fn assert_procedure_event_contract( trigger: &str, locator: FlowEventLocator<'_>, ) { - let mut schema = vec![ - PROCEDURE_ID_COLUMN.column_schema(), - PROCEDURE_STATE_COLUMN.column_schema(), - PROCEDURE_ERROR_COLUMN.column_schema(), - PROCEDURE_TRIGGER_COLUMN.column_schema(), - ]; + let mut schema = procedure_event_column_schemas(); schema.extend(flow_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); assert_event_contract( diff --git a/src/common/meta/src/ddl/tests/event/view.rs b/src/common/meta/src/ddl/tests/event/view.rs index 5cb5a6bb7d9..6ebb62cda82 100644 --- a/src/common/meta/src/ddl/tests/event/view.rs +++ b/src/common/meta/src/ddl/tests/event/view.rs @@ -17,9 +17,8 @@ use std::sync::Arc; use api::v1::value::ValueData; use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, - PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, - VIEW_NAME_COLUMN, jsonb_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, + VIEW_NAME_COLUMN, jsonb_value, procedure_event_column_schemas, }; use common_event_recorder::testing::assert_event_contract; use common_event_recorder::{Event, EventTypeFilter}; @@ -334,12 +333,7 @@ fn assert_procedure_event_contract( trigger: &str, locator: ViewEventLocator<'_>, ) { - let mut schema = vec![ - PROCEDURE_ID_COLUMN.column_schema(), - PROCEDURE_STATE_COLUMN.column_schema(), - PROCEDURE_ERROR_COLUMN.column_schema(), - PROCEDURE_TRIGGER_COLUMN.column_schema(), - ]; + let mut schema = procedure_event_column_schemas(); schema.extend(view_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 2eed4628ac3..7ba408ca2d9 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -19,9 +19,12 @@ use api::v1::Repartition; use api::v1::alter_table_expr::Kind; use api::v1::repartition::Source as PbRepartitionSource; use common_error::ext::BoxedError; +use common_event_recorder::PersistentEventContext; +#[cfg(feature = "enterprise")] +use common_event_recorder::TriggerReason; use common_procedure::{ - BoxedProcedure, BoxedProcedureLoader, Output, ProcedureId, ProcedureManagerRef, - ProcedureWithId, watcher, + BoxedProcedure, BoxedProcedureLoader, Output, ProcedureContext, ProcedureId, + ProcedureManagerRef, ProcedureWithId, watcher, }; use common_telemetry::tracing_context::{FutureExt, TracingContext}; use common_telemetry::{debug, info, tracing}; @@ -54,7 +57,7 @@ use crate::error::{ PersistRepartitionGcRequirementSnafu, ProcedureOutputSnafu, RegisterProcedureLoaderSnafu, RegisterRepartitionProcedureLoaderSnafu, Result, SubmitProcedureSnafu, TableInfoNotFoundSnafu, TableNotFoundSnafu, TableRouteNotFoundSnafu, UnexpectedLogicalRouteTableSnafu, - WaitProcedureSnafu, + UnsupportedSnafu, WaitProcedureSnafu, }; use crate::key::table_info::TableInfoValue; use crate::key::table_name::TableNameKey; @@ -70,16 +73,14 @@ use crate::rpc::ddl::DdlTask::{ DropTable, DropView, PurgeDroppedTable, TruncateTable, UndropTable, }; #[cfg(feature = "enterprise")] -use crate::rpc::ddl::TriggerReason; -#[cfg(feature = "enterprise")] use crate::rpc::ddl::trigger::CreateTriggerTask; #[cfg(feature = "enterprise")] use crate::rpc::ddl::trigger::DropTriggerTask; use crate::rpc::ddl::{ AlterDatabaseTask, AlterTableTask, CommentOnTask, CreateDatabaseTask, CreateFlowTask, CreateTableTask, CreateViewTask, DropDatabaseTask, DropFlowTask, DropTableTask, DropViewTask, - PersistentEventContext, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, - SubmitDdlTaskResponse, TruncateTableTask, UndropTableTask, + PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse, + TruncateTableTask, UndropTableTask, }; const MAX_REGION_ROUTE_CHANGE_RETRIES: usize = 3; @@ -122,6 +123,7 @@ pub trait TriggerDdlManager: Send + Sync { procedure_manager: ProcedureManagerRef, ddl_context: DdlContext, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result; async fn drop_trigger( @@ -130,6 +132,7 @@ pub trait TriggerDdlManager: Send + Sync { procedure_manager: ProcedureManagerRef, ddl_context: DdlContext, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result; fn as_any(&self) -> &dyn std::any::Any; @@ -341,7 +344,7 @@ impl DdlManager { repartition: Repartition, wait: bool, timeout: Duration, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -384,7 +387,7 @@ impl DdlManager { .await .context(PersistRepartitionGcRequirementSnafu)?; let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); if wait { self.execute_procedure_and_wait(procedure_with_id).await } else { @@ -400,7 +403,7 @@ impl DdlManager { &self, table_id: TableId, alter_table_task: AlterTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ddl_options: DdlOptions, ) -> Result<(ProcedureId, Option)> { // make alter_table_task mutable so we can call .take() on its field @@ -421,7 +424,7 @@ impl DdlManager { repartition, ddl_options.wait, ddl_options.timeout, - event_context, + procedure_context, ) .await; } @@ -460,7 +463,7 @@ impl DdlManager { )?; let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context.clone()); + .with_context(procedure_context.clone()); let result = self.execute_procedure_and_wait(procedure_with_id).await?; if result .1 @@ -489,7 +492,7 @@ impl DdlManager { &self, create_table_task: CreateTableTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -500,7 +503,7 @@ impl DdlManager { )?; let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -510,14 +513,14 @@ impl DdlManager { pub async fn submit_create_view_task( &self, create_view_task: CreateViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateViewProcedure::new(create_view_task, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -528,7 +531,7 @@ impl DdlManager { &self, create_table_tasks: Vec, physical_table_id: TableId, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -536,7 +539,7 @@ impl DdlManager { CreateLogicalTablesProcedure::new(create_table_tasks, physical_table_id, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -547,7 +550,7 @@ impl DdlManager { &self, alter_table_tasks: Vec, physical_table_id: TableId, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -555,7 +558,7 @@ impl DdlManager { AlterLogicalTablesProcedure::new(alter_table_tasks, physical_table_id, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -565,14 +568,14 @@ impl DdlManager { pub async fn submit_drop_table_task( &self, drop_table_task: DropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropTableProcedure::new(drop_table_task, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -583,7 +586,7 @@ impl DdlManager { pub async fn submit_undrop_table_task( &self, undrop_table_task: UndropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -611,7 +614,7 @@ impl DdlManager { Some(original_table_name), ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context); + .with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -623,7 +626,7 @@ impl DdlManager { pub async fn submit_purge_dropped_table_task( &self, purge_dropped_table_task: PurgeDroppedTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -639,7 +642,7 @@ impl DdlManager { let context = self.create_context(); let procedure = PurgeDroppedTableProcedure::new(purge_dropped_table_task, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context); + .with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -665,11 +668,13 @@ impl DdlManager { #[cfg(feature = "enterprise")] { let context = self.create_context(); - let event_context = PersistentEventContext::new(TriggerReason::ScheduledGc); + let procedure_context = ProcedureContext::from_event_context( + PersistentEventContext::new(TriggerReason::ScheduledGc), + ); let procedure = PurgeDroppedTableProcedure::new_if_expired(purge_dropped_table_task, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context); + .with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -686,7 +691,7 @@ impl DdlManager { options, creator, }: CreateDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateDatabaseProcedure::new( @@ -698,7 +703,7 @@ impl DdlManager { context, ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -712,12 +717,12 @@ impl DdlManager { schema, drop_if_exists, }: DropDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropDatabaseProcedure::new(catalog, schema, drop_if_exists, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -725,12 +730,12 @@ impl DdlManager { pub async fn submit_alter_database( &self, alter_database_task: AlterDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = AlterDatabaseProcedure::new(alter_database_task, context)?; let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -741,12 +746,12 @@ impl DdlManager { &self, create_flow: CreateFlowTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateFlowProcedure::new(create_flow, query_context, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -756,12 +761,12 @@ impl DdlManager { pub async fn submit_drop_flow_task( &self, drop_flow: DropFlowTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropFlowProcedure::new(drop_flow, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -771,12 +776,12 @@ impl DdlManager { pub async fn submit_drop_view_task( &self, drop_view: DropViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropViewProcedure::new(drop_view, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -787,13 +792,13 @@ impl DdlManager { &self, truncate_table_task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = TruncateTableProcedure::new(truncate_table_task, table_info_value, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -803,7 +808,7 @@ impl DdlManager { pub async fn submit_comment_on_task( &self, mut comment_on_task: CommentOnTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); comment_on_task @@ -814,7 +819,7 @@ impl DdlManager { .await?; let procedure = CommentOnProcedure::new(comment_on_task, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -853,89 +858,125 @@ impl DdlManager { pub async fn submit_ddl_task( &self, - ctx: &ExecutorContext, + context: ExecutorContext, request: SubmitDdlTaskRequest, ) -> Result { - let span = ctx - .tracing_context + let ExecutorContext { + tracing_context, + query_context, + actor, + event_input, + } = context; + let query_context = query_context.context(UnsupportedSnafu { + operation: "submit_ddl_task without query context", + })?; + let procedure_context = ProcedureContext { + actor, + event_context: event_input + .map(|input| PersistentEventContext::from((input, query_context.protocol()))), + }; + let span = tracing_context .as_ref() .map(TracingContext::from_w3c) .unwrap_or_else(TracingContext::from_current_span) .attach(tracing::info_span!("DdlManager::submit_ddl_task")); let SubmitDdlTaskRequest { - query_context, wait, timeout, task, } = request; - let event_context = ctx - .event_context - .clone() - .unwrap_or_else(|| crate::rpc::ddl::event_context_from_query_context(&query_context)); let ddl_options = DdlOptions { wait, timeout }; async move { debug!("Submitting Ddl task: {:?}", task); match task { CreateTable(create_table_task) => { - handle_create_table_task(self, create_table_task, query_context, event_context) - .await + handle_create_table_task( + self, + create_table_task, + query_context, + procedure_context, + ) + .await } DropTable(drop_table_task) => { - handle_drop_table_task(self, drop_table_task, event_context).await + handle_drop_table_task(self, drop_table_task, procedure_context).await } UndropTable(undrop_table_task) => { - handle_undrop_table_task(self, undrop_table_task, event_context).await + handle_undrop_table_task(self, undrop_table_task, procedure_context).await } PurgeDroppedTable(purge_dropped_table_task) => { - handle_purge_dropped_table_task(self, purge_dropped_table_task, event_context) - .await + handle_purge_dropped_table_task( + self, + purge_dropped_table_task, + procedure_context, + ) + .await } AlterTable(alter_table_task) => { - handle_alter_table_task(self, alter_table_task, ddl_options, event_context) + handle_alter_table_task(self, alter_table_task, ddl_options, procedure_context) .await } TruncateTable(truncate_table_task) => { - handle_truncate_table_task(self, truncate_table_task, event_context).await + handle_truncate_table_task(self, truncate_table_task, procedure_context).await } CreateLogicalTables(create_table_tasks) => { - handle_create_logical_table_tasks(self, create_table_tasks, event_context).await + handle_create_logical_table_tasks(self, create_table_tasks, procedure_context) + .await } AlterLogicalTables(alter_table_tasks) => { - handle_alter_logical_table_tasks(self, alter_table_tasks, event_context).await + handle_alter_logical_table_tasks(self, alter_table_tasks, procedure_context) + .await } DropLogicalTables(_) => todo!(), CreateDatabase(create_database_task) => { - handle_create_database_task(self, create_database_task, event_context).await + handle_create_database_task(self, create_database_task, procedure_context).await } DropDatabase(drop_database_task) => { - handle_drop_database_task(self, drop_database_task, event_context).await + handle_drop_database_task(self, drop_database_task, procedure_context).await } AlterDatabase(alter_database_task) => { - handle_alter_database_task(self, alter_database_task, event_context).await + handle_alter_database_task(self, alter_database_task, procedure_context).await } CreateFlow(create_flow_task) => { - handle_create_flow_task(self, create_flow_task, query_context, event_context) - .await + handle_create_flow_task( + self, + create_flow_task, + query_context, + procedure_context, + ) + .await } DropFlow(drop_flow_task) => { - handle_drop_flow_task(self, drop_flow_task, event_context).await + handle_drop_flow_task(self, drop_flow_task, procedure_context).await } CreateView(create_view_task) => { - handle_create_view_task(self, create_view_task, event_context).await + handle_create_view_task(self, create_view_task, procedure_context).await } DropView(drop_view_task) => { - handle_drop_view_task(self, drop_view_task, event_context).await + handle_drop_view_task(self, drop_view_task, procedure_context).await } CommentOn(comment_on_task) => { - handle_comment_on_task(self, comment_on_task, event_context).await + handle_comment_on_task(self, comment_on_task, procedure_context).await } #[cfg(feature = "enterprise")] CreateTrigger(create_trigger_task) => { - handle_create_trigger_task(self, create_trigger_task, query_context).await + handle_create_trigger_task( + self, + create_trigger_task, + query_context, + procedure_context, + ) + .await } #[cfg(feature = "enterprise")] DropTrigger(drop_trigger_task) => { - handle_drop_trigger_task(self, drop_trigger_task, query_context).await + handle_drop_trigger_task( + self, + drop_trigger_task, + query_context, + procedure_context, + ) + .await } } } @@ -947,7 +988,7 @@ impl DdlManager { async fn handle_truncate_table_task( ddl_manager: &DdlManager, truncate_table_task: TruncateTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_id = truncate_table_task.table_id; let table_metadata_manager = &ddl_manager.table_metadata_manager(); @@ -972,7 +1013,7 @@ async fn handle_truncate_table_task( ); let (id, _) = ddl_manager - .submit_truncate_table_task(truncate_table_task, table_info_value, event_context) + .submit_truncate_table_task(truncate_table_task, table_info_value, procedure_context) .await?; info!("Table: {table_id} is truncated via procedure_id {id:?}"); @@ -987,7 +1028,7 @@ async fn handle_alter_table_task( ddl_manager: &DdlManager, alter_table_task: AlterTableTask, ddl_options: DdlOptions, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_ref = alter_table_task.table_ref(); @@ -1020,7 +1061,7 @@ async fn handle_alter_table_task( ); let (id, _) = ddl_manager - .submit_alter_table_task(table_id, alter_table_task, event_context, ddl_options) + .submit_alter_table_task(table_id, alter_table_task, procedure_context, ddl_options) .await?; info!("Table: {table_id} is altered via procedure_id {id:?}"); @@ -1034,11 +1075,11 @@ async fn handle_alter_table_task( async fn handle_drop_table_task( ddl_manager: &DdlManager, drop_table_task: DropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_id = drop_table_task.table_id; let (id, _) = ddl_manager - .submit_drop_table_task(drop_table_task, event_context) + .submit_drop_table_task(drop_table_task, procedure_context) .await?; info!("Table: {table_id} is dropped via procedure_id {id:?}"); @@ -1052,11 +1093,11 @@ async fn handle_drop_table_task( async fn handle_undrop_table_task( ddl_manager: &DdlManager, undrop_table_task: UndropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_id = undrop_table_task.table_id; let (id, _) = ddl_manager - .submit_undrop_table_task(undrop_table_task, event_context) + .submit_undrop_table_task(undrop_table_task, procedure_context) .await?; info!("Table: {table_id} is undropped via procedure_id {id:?}"); @@ -1070,10 +1111,10 @@ async fn handle_undrop_table_task( async fn handle_purge_dropped_table_task( ddl_manager: &DdlManager, purge_dropped_table_task: PurgeDroppedTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_purge_dropped_table_task(purge_dropped_table_task, event_context) + .submit_purge_dropped_table_task(purge_dropped_table_task, procedure_context) .await?; info!("Dropped table is purged via procedure_id {id:?}"); @@ -1088,10 +1129,10 @@ async fn handle_create_table_task( ddl_manager: &DdlManager, create_table_task: CreateTableTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_table_task(create_table_task, query_context, event_context) + .submit_create_table_task(create_table_task, query_context, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1114,7 +1155,7 @@ async fn handle_create_table_task( async fn handle_create_logical_table_tasks( ddl_manager: &DdlManager, create_table_tasks: Vec, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { ensure!( !create_table_tasks.is_empty(), @@ -1130,7 +1171,7 @@ async fn handle_create_logical_table_tasks( let num_logical_tables = create_table_tasks.len(); let (id, output) = ddl_manager - .submit_create_logical_table_tasks(create_table_tasks, physical_table_id, event_context) + .submit_create_logical_table_tasks(create_table_tasks, physical_table_id, procedure_context) .await?; info!( @@ -1159,12 +1200,12 @@ async fn handle_create_logical_table_tasks( async fn handle_create_database_task( ddl_manager: &DdlManager, create_database_task: CreateDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let catalog = create_database_task.catalog.clone(); let schema = create_database_task.schema.clone(); let (id, _) = ddl_manager - .submit_create_database(create_database_task, event_context) + .submit_create_database(create_database_task, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1182,10 +1223,10 @@ async fn handle_create_database_task( async fn handle_drop_database_task( ddl_manager: &DdlManager, drop_database_task: DropDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_database(drop_database_task.clone(), event_context) + .submit_drop_database(drop_database_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1203,10 +1244,10 @@ async fn handle_drop_database_task( async fn handle_alter_database_task( ddl_manager: &DdlManager, alter_database_task: AlterDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_alter_database(alter_database_task.clone(), event_context) + .submit_alter_database(alter_database_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1225,10 +1266,10 @@ async fn handle_alter_database_task( async fn handle_drop_flow_task( ddl_manager: &DdlManager, drop_flow_task: DropFlowTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_flow_task(drop_flow_task.clone(), event_context) + .submit_drop_flow_task(drop_flow_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1248,6 +1289,7 @@ async fn handle_drop_trigger_task( ddl_manager: &DdlManager, drop_trigger_task: DropTriggerTask, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result { let Some(m) = ddl_manager.trigger_ddl_manager.as_ref() else { use crate::error::UnsupportedSnafu; @@ -1263,6 +1305,7 @@ async fn handle_drop_trigger_task( ddl_manager.procedure_manager.clone(), ddl_manager.ddl_context.clone(), query_context, + procedure_context, ) .await } @@ -1270,10 +1313,10 @@ async fn handle_drop_trigger_task( async fn handle_drop_view_task( ddl_manager: &DdlManager, drop_view_task: DropViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_view_task(drop_view_task.clone(), event_context) + .submit_drop_view_task(drop_view_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1293,10 +1336,10 @@ async fn handle_create_flow_task( ddl_manager: &DdlManager, create_flow_task: CreateFlowTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_flow_task(create_flow_task.clone(), query_context, event_context) + .submit_create_flow_task(create_flow_task.clone(), query_context, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1331,6 +1374,7 @@ async fn handle_create_trigger_task( ddl_manager: &DdlManager, create_trigger_task: CreateTriggerTask, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result { let Some(m) = ddl_manager.trigger_ddl_manager.as_ref() else { use crate::error::UnsupportedSnafu; @@ -1346,6 +1390,7 @@ async fn handle_create_trigger_task( ddl_manager.procedure_manager.clone(), ddl_manager.ddl_context.clone(), query_context, + procedure_context, ) .await } @@ -1353,7 +1398,7 @@ async fn handle_create_trigger_task( async fn handle_alter_logical_table_tasks( ddl_manager: &DdlManager, alter_table_tasks: Vec, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { ensure!( !alter_table_tasks.is_empty(), @@ -1373,7 +1418,7 @@ async fn handle_alter_logical_table_tasks( let num_logical_tables = alter_table_tasks.len(); let (id, _) = ddl_manager - .submit_alter_logical_table_tasks(alter_table_tasks, physical_table_id, event_context) + .submit_alter_logical_table_tasks(alter_table_tasks, physical_table_id, procedure_context) .await?; info!( @@ -1392,10 +1437,10 @@ async fn handle_alter_logical_table_tasks( async fn handle_create_view_task( ddl_manager: &DdlManager, create_view_task: CreateViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_view_task(create_view_task, event_context) + .submit_create_view_task(create_view_task, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1418,10 +1463,10 @@ async fn handle_create_view_task( async fn handle_comment_on_task( ddl_manager: &DdlManager, comment_on_task: CommentOnTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_comment_on_task(comment_on_task.clone(), event_context) + .submit_comment_on_task(comment_on_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1439,16 +1484,22 @@ async fn handle_comment_on_task( #[cfg(test)] mod tests { use std::sync::Arc; + #[cfg(feature = "enterprise")] + use std::sync::Mutex; use std::time::Duration; + #[cfg(feature = "enterprise")] + use common_base::protocol::Channel; use common_error::ext::BoxedError; #[cfg(feature = "enterprise")] use common_error::ext::ErrorExt; #[cfg(feature = "enterprise")] use common_error::status_code::StatusCode; + #[cfg(feature = "enterprise")] + use common_event_recorder::{PersistentEventContext, ProcedureEventInput, TriggerReason}; use common_procedure::local::LocalManager; use common_procedure::test_util::InMemoryPoisonStore; - use common_procedure::{BoxedProcedure, ProcedureManagerRef}; + use common_procedure::{BoxedProcedure, ProcedureContext, ProcedureManagerRef}; use store_api::storage::TableId; use table::table_name::TableName; @@ -1470,13 +1521,16 @@ mod tests { use crate::kv_backend::memory::MemoryKvBackend; use crate::node_manager::{DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef}; use crate::peer::Peer; - #[cfg(not(feature = "enterprise"))] use crate::procedure_executor::ExecutorContext; use crate::region_keeper::MemoryRegionKeeper; use crate::region_registry::LeaderRegionRegistry; - use crate::rpc::ddl::{CreatorGrantIntent, PersistentEventContext, UndropTableTask}; + #[cfg(feature = "enterprise")] + use crate::rpc::ddl::trigger::{CreateTriggerTask, DropTriggerTask}; + use crate::rpc::ddl::{CreatorGrantIntent, UndropTableTask}; #[cfg(not(feature = "enterprise"))] use crate::rpc::ddl::{DdlTask, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest}; + #[cfg(feature = "enterprise")] + use crate::rpc::ddl::{DdlTask, QueryContext, SubmitDdlTaskRequest}; use crate::sequence::SequenceBuilder; use crate::state_store::KvStateStore; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -1543,6 +1597,50 @@ mod tests { } } + #[cfg(feature = "enterprise")] + #[derive(Default)] + struct RecordingTriggerDdlManager { + procedure_contexts: Mutex>, + } + + #[cfg(feature = "enterprise")] + #[async_trait::async_trait] + impl super::TriggerDdlManager for RecordingTriggerDdlManager { + async fn create_trigger( + &self, + _create_trigger_task: CreateTriggerTask, + _procedure_manager: ProcedureManagerRef, + _ddl_context: DdlContext, + _query_context: QueryContext, + procedure_context: ProcedureContext, + ) -> crate::error::Result { + self.procedure_contexts + .lock() + .unwrap() + .push(procedure_context); + Ok(Default::default()) + } + + async fn drop_trigger( + &self, + _drop_trigger_task: DropTriggerTask, + _procedure_manager: ProcedureManagerRef, + _ddl_context: DdlContext, + _query_context: QueryContext, + procedure_context: ProcedureContext, + ) -> crate::error::Result { + self.procedure_contexts + .lock() + .unwrap() + .push(procedure_context); + Ok(Default::default()) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + #[test] fn test_generic_loader_captures_configured_committer() { let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); @@ -1678,6 +1776,67 @@ mod tests { ) } + #[cfg(feature = "enterprise")] + #[tokio::test] + async fn test_trigger_ddl_forwards_procedure_context() { + let trigger_ddl_manager = Arc::new(RecordingTriggerDdlManager::default()); + let ddl_manager = build_soft_drop_test_ddl_manager() + .with_trigger_ddl_manager(trigger_ddl_manager.clone()); + let procedure_context = ProcedureContext { + actor: Some("test-user".to_string()), + event_context: Some( + PersistentEventContext::new(TriggerReason::Manual).with_protocol("mysql"), + ), + }; + let executor_context = || ExecutorContext { + query_context: Some(QueryContext { + channel: Channel::Mysql as u8, + ..Default::default() + }), + actor: Some("test-user".to_string()), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + + ddl_manager + .submit_ddl_task( + executor_context(), + SubmitDdlTaskRequest::new(DdlTask::CreateTrigger(CreateTriggerTask { + catalog_name: "greptime".to_string(), + trigger_name: "test_trigger".to_string(), + if_not_exists: false, + sql: "SELECT 1".to_string(), + channels: vec![], + labels: Default::default(), + annotations: Default::default(), + interval: Duration::from_secs(1), + raw_interval_expr: None, + r#for: None, + for_raw_expr: None, + keep_firing_for: None, + keep_firing_for_raw_expr: None, + })), + ) + .await + .unwrap(); + ddl_manager + .submit_ddl_task( + executor_context(), + SubmitDdlTaskRequest::new(DdlTask::DropTrigger(DropTriggerTask { + catalog_name: "greptime".to_string(), + trigger_name: "test_trigger".to_string(), + drop_if_exists: false, + })), + ) + .await + .unwrap(); + + assert_eq!( + *trigger_ddl_manager.procedure_contexts.lock().unwrap(), + vec![procedure_context.clone(), procedure_context] + ); + } + #[cfg(feature = "enterprise")] #[tokio::test] async fn test_submit_undrop_missing_tombstone_returns_table_not_found_directly() { @@ -1686,7 +1845,7 @@ mod tests { let err = ddl_manager .submit_undrop_table_task( UndropTableTask { table_id: 1024 }, - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap_err(); @@ -1703,7 +1862,7 @@ mod tests { let err = ddl_manager .submit_undrop_table_task( UndropTableTask { table_id: 1024 }, - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap_err(); @@ -1712,7 +1871,7 @@ mod tests { let err = ddl_manager .submit_purge_dropped_table_task( PurgeDroppedTableTask { table_id: 1024 }, - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap_err(); @@ -1730,8 +1889,11 @@ mod tests { ] { let err = ddl_manager .submit_ddl_task( - &ExecutorContext::default(), - SubmitDdlTaskRequest::new(QueryContext::default(), task), + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(task), ) .await .unwrap_err(); diff --git a/src/common/meta/src/procedure_executor.rs b/src/common/meta/src/procedure_executor.rs index b812aeb6aac..a873c26b39a 100644 --- a/src/common/meta/src/procedure_executor.rs +++ b/src/common/meta/src/procedure_executor.rs @@ -15,7 +15,7 @@ use std::sync::Arc; use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse}; -use common_event_recorder::PersistentEventContext; +use common_event_recorder::ProcedureEventInput; use common_procedure::{ProcedureId, ProcedureManagerRef}; use common_telemetry::tracing_context::W3cTrace; use snafu::{OptionExt, ResultExt}; @@ -24,7 +24,7 @@ use crate::ddl_manager::DdlManagerRef; use crate::error::{ ParseProcedureIdSnafu, ProcedureNotFoundSnafu, QueryProcedureSnafu, Result, UnsupportedSnafu, }; -use crate::rpc::ddl::{SubmitDdlTaskRequest, SubmitDdlTaskResponse}; +use crate::rpc::ddl::{QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; use crate::rpc::procedure::{ self, GcRegionsRequest, GcResponse, GcTableRequest, ManageRegionFollowerRequest, MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, @@ -34,8 +34,14 @@ use crate::rpc::procedure::{ #[derive(Debug, Clone, Default)] pub struct ExecutorContext { pub tracing_context: Option, + /// Query execution data available at the frontend/standalone submission boundary. + /// + /// Only DDL serializes the full context. Migration and GC use its typed + /// channel to derive protocol without sending the query context itself. + pub query_context: Option, pub actor: Option, - pub event_context: Option, + /// Caller-supplied event metadata. Protocol is derived by the submission adapter. + pub event_input: Option, } /// The procedure executor that accepts ddl, region migration task etc. @@ -44,7 +50,7 @@ pub trait ProcedureExecutor: Send + Sync { /// Submit a ddl task async fn submit_ddl_task( &self, - ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> Result; @@ -129,7 +135,7 @@ impl LocalProcedureExecutor { impl ProcedureExecutor for LocalProcedureExecutor { async fn submit_ddl_task( &self, - ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> Result { self.ddl_manager.submit_ddl_task(ctx, request).await diff --git a/src/common/meta/src/rpc/ddl.rs b/src/common/meta/src/rpc/ddl.rs index d7966804738..47bd537b9ba 100644 --- a/src/common/meta/src/rpc/ddl.rs +++ b/src/common/meta/src/rpc/ddl.rs @@ -45,7 +45,7 @@ use base64::engine::general_purpose; use common_base::protocol::Channel; use common_catalog::{format_full_flow_name, format_full_table_name}; use common_error::ext::BoxedError; -pub use common_event_recorder::{PersistentEventContext, TriggerReason}; +pub use common_event_recorder::TriggerReason; use common_time::{DatabaseTimeToLive, Timestamp}; use prost::Message; use serde::{Deserialize, Serialize}; @@ -68,8 +68,6 @@ use crate::key::table_name::{TableNameKey, TableNameManager}; /// Reserved query-context extension key for the frontend peer address that submitted a DDL request. pub const ORIGIN_FRONTEND_ADDR_EXTENSION_KEY: &str = "__greptime_origin_frontend.addr"; -/// Reserved query-context extension key for the trigger reason supplied by frontend. -pub const TRIGGER_REASON_EXTENSION_KEY: &str = "__greptime_event.trigger_reason"; /// Reserved query-context extension key for the authenticated database creator. pub const CREATE_DATABASE_CREATOR_EXTENSION_KEY: &str = "__greptime_create_database.creator"; /// Internal gRPC metadata key for the authenticated database creator. @@ -336,7 +334,6 @@ impl TryFrom for DdlTask { #[derive(Clone)] pub struct SubmitDdlTaskRequest { - pub query_context: QueryContext, pub wait: bool, pub timeout: Duration, pub task: DdlTask, @@ -344,9 +341,8 @@ pub struct SubmitDdlTaskRequest { impl SubmitDdlTaskRequest { /// The default constructor for [`SubmitDdlTaskRequest`]. - pub fn new(query_context: QueryContext, task: DdlTask) -> Self { + pub fn new(task: DdlTask) -> Self { Self { - query_context, wait: Self::default_wait(), timeout: Self::default_timeout(), task, @@ -377,26 +373,11 @@ impl TryFrom for PbDdlTaskRequest { fn try_from(request: SubmitDdlTaskRequest) -> Result { let SubmitDdlTaskRequest { - mut query_context, wait, timeout, task, } = request; - query_context - .extensions - .remove(CREATE_DATABASE_CREATOR_EXTENSION_KEY); - if let DdlTask::CreateDatabase(CreateDatabaseTask { - creator: Some(creator), - .. - }) = &task - { - query_context.extensions.insert( - CREATE_DATABASE_CREATOR_EXTENSION_KEY.to_string(), - serde_json::to_string(creator).context(error::SerdeJsonSnafu)?, - ); - } - let task = match task { DdlTask::CreateTable(task) => Task::CreateTableTask(task.try_into()?), DdlTask::DropTable(task) => Task::DropTableTask(task.into()), @@ -444,10 +425,12 @@ impl TryFrom for PbDdlTaskRequest { Ok(Self { header: None, - query_context: Some(query_context.into()), + query_context: None, timeout_secs: ddl_timeout_secs(timeout), wait, task: Some(task), + event_context: None, + actor: None, }) } } @@ -1668,22 +1651,6 @@ pub struct QueryContext { pub sst_min_sequences: HashMap, } -/// Builds event context from frontend query context metadata. -pub fn event_context_from_query_context(query_context: &QueryContext) -> PersistentEventContext { - let reason = query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) - .map(|reason| TriggerReason::from_extension(reason)) - .unwrap_or_default(); - let channel = Channel::from(u32::from(query_context.channel)); - let context = PersistentEventContext::new(reason); - if channel == Channel::Unknown { - context - } else { - context.with_protocol(channel.as_ref()) - } -} - impl QueryContext { /// Get the current catalog pub fn current_catalog(&self) -> &str { @@ -1710,6 +1677,12 @@ impl QueryContext { self.channel } + /// Returns the protocol derived from the typed query channel. + pub fn protocol(&self) -> Option { + let channel = Channel::from(u32::from(self.channel)); + (channel != Channel::Unknown).then(|| channel.as_ref().to_string()) + } + pub fn snapshot_seqs(&self) -> &HashMap { &self.snapshot_seqs } @@ -1922,10 +1895,7 @@ mod tests { #[test] fn test_undrop_table_task_pb_roundtrip() { let expected = UndropTableTask { table_id: 1024 }; - let request = SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::UndropTable(expected.clone()), - ); + let request = SubmitDdlTaskRequest::new(DdlTask::UndropTable(expected.clone())); let pb = PbDdlTaskRequest::try_from(request).unwrap(); let pb_task = pb.task.unwrap(); @@ -1937,10 +1907,7 @@ mod tests { #[test] fn test_purge_dropped_table_task_pb_roundtrip() { let expected = PurgeDroppedTableTask { table_id: 1024 }; - let request = SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::PurgeDroppedTable(expected.clone()), - ); + let request = SubmitDdlTaskRequest::new(DdlTask::PurgeDroppedTable(expected.clone())); let pb = PbDdlTaskRequest::try_from(request).unwrap(); let pb_task = pb.task.unwrap(); @@ -1949,53 +1916,6 @@ mod tests { assert!(matches!(de, DdlTask::PurgeDroppedTable(task) if task == expected)); } - #[test] - fn test_create_database_creator_transport_sanitizes_extension() { - let request = |creator| { - let mut query_context = QueryContext::default(); - query_context.extensions.insert( - CREATE_DATABASE_CREATOR_EXTENSION_KEY.to_string(), - "spoofed".to_string(), - ); - SubmitDdlTaskRequest::new( - query_context, - DdlTask::new_create_database( - "greptime".to_string(), - "test".to_string(), - false, - HashMap::new(), - creator, - ), - ) - }; - let creator = CreatorGrantIntent { - username: "alice".to_string(), - created_at_ns: 42, - }; - - let mut pb = PbDdlTaskRequest::try_from(request(Some(creator.clone()))).unwrap(); - let encoded = pb - .query_context - .as_ref() - .unwrap() - .extensions - .get(CREATE_DATABASE_CREATOR_EXTENSION_KEY) - .unwrap(); - assert_eq!(creator, serde_json::from_str(encoded).unwrap()); - assert!(matches!( - DdlTask::try_from(pb.task.take().unwrap()).unwrap(), - DdlTask::CreateDatabase(CreateDatabaseTask { creator: None, .. }) - )); - - let pb = PbDdlTaskRequest::try_from(request(None)).unwrap(); - assert!( - !pb.query_context - .unwrap() - .extensions - .contains_key(CREATE_DATABASE_CREATOR_EXTENSION_KEY) - ); - } - #[test] fn test_undrop_table_task_json_roundtrip() { let task = UndropTableTask { table_id: 1024 }; @@ -2262,44 +2182,4 @@ mod tests { let reason: TriggerReason = serde_json::from_str("\"future_reason\"").unwrap(); assert_eq!(TriggerReason::Unknown, reason); } - - #[test] - fn test_event_context_serialization() { - let context = PersistentEventContext::new(TriggerReason::Manual).with_protocol("mysql"); - - assert_eq!( - serde_json::json!({ - "reason": "manual", - "protocol": "mysql", - }), - serde_json::to_value(context).unwrap() - ); - - assert_eq!( - serde_json::json!({ "reason": "manual" }), - serde_json::to_value(PersistentEventContext::new(TriggerReason::Manual)).unwrap() - ); - } - - #[test] - fn test_event_context_from_query_context_preserves_extensions() { - let mut query_context = QueryContext::default(); - query_context.extensions.insert( - TRIGGER_REASON_EXTENSION_KEY.to_string(), - TriggerReason::AutoCreate.as_ref().to_string(), - ); - query_context.channel = 4; - - let event_context = event_context_from_query_context(&query_context); - - assert_eq!(event_context.reason, TriggerReason::AutoCreate); - assert_eq!(event_context.protocol.as_deref(), Some("prometheus")); - assert_eq!( - query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) - .map(String::as_str), - Some("auto_create") - ); - } } diff --git a/src/common/procedure/src/local.rs b/src/common/procedure/src/local.rs index 8fdd3167583..c870cc39a67 100644 --- a/src/common/procedure/src/local.rs +++ b/src/common/procedure/src/local.rs @@ -1062,7 +1062,7 @@ mod tests { use super::*; use crate::error::{self, Error}; - use crate::store::state_store::ObjectStateStore; + use crate::store::state_store::{ObjectStateStore, StateStore}; use crate::test_util::InMemoryPoisonStore; use crate::{ ChildSubmissionOutcome, Context, EventContext, EventTrigger, Procedure, ProcedureEvent, @@ -1721,7 +1721,63 @@ mod tests { } #[tokio::test] - async fn test_recovery_emits_recovered_event() { + async fn test_recovery_loads_legacy_procedure_message() { + let dir = create_temp_dir("legacy_procedure_message_recovery"); + let object_store = test_util::new_object_store(&dir); + let state_store = Arc::new(ObjectStateStore::new(object_store)); + let poison_manager = Arc::new(InMemoryPoisonStore::new()); + let event_recorder = Arc::new(CapturingEventRecorder::default()); + let manager = LocalManager::new( + ManagerConfig { + parent_path: "data/".to_string(), + ..Default::default() + }, + state_store.clone(), + poison_manager, + None, + None, + ); + manager + .event_recorder_handle() + .install(event_recorder.clone()); + manager.manager_ctx.start(); + manager + .register_loader("ProcedureToLoad", ProcedureToLoad::loader()) + .unwrap(); + + let procedure = ProcedureToLoad::new("legacy recovered submission"); + let procedure_id = ProcedureId::random(); + let key = format!("data/procedure/{procedure_id}/0000000000.step"); + // This is the persisted JSON shape written before ProcedureMessage had `context`. + let legacy_message = serde_json::json!({ + "type_name": procedure.type_name(), + "data": procedure.dump().unwrap(), + "parent_id": null, + "step": 0, + }); + state_store + .put(&key, serde_json::to_vec(&legacy_message).unwrap()) + .await + .unwrap(); + + manager.recover().await.unwrap(); + + assert!( + manager + .procedure_state(procedure_id) + .await + .unwrap() + .is_some() + ); + wait_for_trigger(&event_recorder, EventTrigger::Recovered).await; + assert_eq!( + event_recorder.procedure_context(procedure_id), + Some(ProcedureContext::default()) + ); + } + + #[tokio::test] + async fn test_recovery_preserves_procedure_context() { let dir = create_temp_dir("recovery_submission_event"); let object_store = test_util::new_object_store(&dir); let state_store = Arc::new(ObjectStateStore::new(object_store.clone())); @@ -1747,6 +1803,16 @@ mod tests { let procedure = ProcedureToLoad::new("recovered submission"); let procedure_id = ProcedureId::random(); + let mut event_context = + PersistentEventContext::new(TriggerReason::AutoRebalance).with_protocol("mysql"); + event_context.extensions.insert( + "source".to_string(), + serde_json::Value::String("test".to_string()), + ); + let procedure_context = ProcedureContext { + actor: Some("alice".to_string()), + event_context: Some(event_context), + }; ProcedureStore::from_object_store(object_store) .store_procedure_with_context( procedure_id, @@ -1754,9 +1820,7 @@ mod tests { procedure.type_name().to_string(), procedure.dump().unwrap(), None, - ProcedureContext::from_event_context(PersistentEventContext::new( - TriggerReason::AutoRebalance, - )), + procedure_context.clone(), ) .await .unwrap(); @@ -1773,13 +1837,8 @@ mod tests { wait_for_trigger(&event_recorder, EventTrigger::Recovered).await; assert!(!event_recorder.triggers().contains(&EventTrigger::Submitted)); assert_eq!( - event_recorder - .procedure_context(procedure_id) - .unwrap() - .event_context - .unwrap() - .reason, - TriggerReason::AutoRebalance + event_recorder.procedure_context(procedure_id), + Some(procedure_context) ); } diff --git a/src/common/procedure/src/procedure.rs b/src/common/procedure/src/procedure.rs index cd11aea9215..ba6a5d275cb 100644 --- a/src/common/procedure/src/procedure.rs +++ b/src/common/procedure/src/procedure.rs @@ -35,6 +35,9 @@ pub type Output = Arc; /// Context attached to a procedure submission and inherited by its children. #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct ProcedureContext { + /// Effective user that submitted the procedure. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub actor: Option, /// Context describing why and how the procedure was submitted. #[serde(default, skip_serializing_if = "Option::is_none")] pub event_context: Option, @@ -43,12 +46,13 @@ pub struct ProcedureContext { impl ProcedureContext { pub fn from_event_context(event_context: PersistentEventContext) -> Self { Self { + actor: None, event_context: Some(event_context), } } pub fn is_empty(&self) -> bool { - self.event_context.is_none() + self.actor.is_none() && self.event_context.is_none() } } @@ -498,9 +502,9 @@ impl ProcedureWithId { } } - /// Attaches event context to this procedure submission. - pub fn with_event_context(mut self, event_context: PersistentEventContext) -> Self { - self.context.event_context = Some(event_context); + /// Attaches a complete persisted context to this procedure submission. + pub fn with_context(mut self, context: ProcedureContext) -> Self { + self.context = context; self } } diff --git a/src/common/procedure/src/store.rs b/src/common/procedure/src/store.rs index 4867046f8de..b2f19b2f94c 100644 --- a/src/common/procedure/src/store.rs +++ b/src/common/procedure/src/store.rs @@ -514,6 +514,14 @@ mod tests { r#"{"type_name":"TestMessage","data":"no parent id","parent_id":"9f805a1f-05f7-490c-9f91-bd56e3cc54c1","step":4,"context":{"event_context":{"reason":"auto_rebalance"}}}"# ); + message.context.actor = Some("alice".to_string()); + let json = serde_json::to_string(&message).unwrap(); + assert!(json.contains(r#""actor":"alice""#)); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + message + ); + let legacy: ProcedureMessage = serde_json::from_str( r#"{"type_name":"TestMessage","data":"legacy","parent_id":null,"step":1}"#, ) diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 0c300835e0e..78b701a3442 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -2281,7 +2281,7 @@ mod tests { impl ProcedureExecutor for NoopProcedureExecutor { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + _ctx: ExecutorContext, _request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { common_meta::error::UnsupportedSnafu { @@ -2358,7 +2358,7 @@ mod tests { impl ProcedureExecutor for MockProcedureExecutor { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + _ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { self.submitted.lock().unwrap().push(request.task.clone()); diff --git a/src/meta-client/Cargo.toml b/src/meta-client/Cargo.toml index 28c03b3f0dd..984f434cc51 100644 --- a/src/meta-client/Cargo.toml +++ b/src/meta-client/Cargo.toml @@ -30,6 +30,7 @@ tokio-stream = { workspace = true, features = ["net"] } tonic.workspace = true [dev-dependencies] +common-event-recorder.workspace = true datatypes.workspace = true futures.workspace = true meta-srv = { workspace = true, features = ["mock"] } diff --git a/src/meta-client/src/client.rs b/src/meta-client/src/client.rs index 4fc6020fffb..e0932af5a69 100644 --- a/src/meta-client/src/client.rs +++ b/src/meta-client/src/client.rs @@ -49,7 +49,10 @@ use common_meta::peer::PeerDiscovery; use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor}; use common_meta::range_stream::PaginationStream; use common_meta::rpc::KeyValue; -use common_meta::rpc::ddl::{SubmitDdlTaskRequest, SubmitDdlTaskResponse}; +use common_meta::rpc::ddl::{ + CREATE_DATABASE_CREATOR_EXTENSION_KEY, CreateDatabaseTask, DdlTask, SubmitDdlTaskRequest, + SubmitDdlTaskResponse, +}; use common_meta::rpc::procedure::{ AddRegionFollowerRequest, AddTableFollowerRequest, GcRegionsRequest, GcResponse, GcTableRequest, ManageRegionFollowerRequest, MigrateRegionRequest, MigrateRegionResponse, @@ -66,7 +69,7 @@ use common_time::util::DefaultSystemTimer; use config::Client as ConfigClient; use futures::TryStreamExt; use heartbeat::{Client as HeartbeatClient, HeartbeatConfig}; -use procedure::Client as ProcedureClient; +use procedure::{Client as ProcedureClient, procedure_actor, procedure_event_context}; use serde::de::DeserializeOwned; use snafu::{OptionExt, ResultExt}; use store::Client as StoreClient; @@ -75,7 +78,7 @@ pub use self::heartbeat::{HeartbeatSender, HeartbeatStream}; use crate::client::ask_leader::{LeaderProviderFactoryImpl, LeaderProviderFactoryRef}; use crate::error::{ ConvertMetaConfigSnafu, ConvertMetaRequestSnafu, ConvertMetaResponseSnafu, Error, - GetFlowStatSnafu, NotStartedSnafu, Result, + GetFlowStatSnafu, MissingQueryContextSnafu, NotStartedSnafu, Result, }; pub type Id = u64; @@ -334,10 +337,10 @@ pub trait RegionFollowerClient: Sync + Send + Debug { impl ProcedureExecutor for MetaClient { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> MetaResult { - self.submit_ddl_task(request) + MetaClient::submit_ddl_task(self, ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -345,10 +348,10 @@ impl ProcedureExecutor for MetaClient { async fn migrate_region( &self, - _ctx: &ExecutorContext, + ctx: &ExecutorContext, request: MigrateRegionRequest, ) -> MetaResult { - self.migrate_region(request) + self.migrate_region(ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -418,10 +421,10 @@ impl ProcedureExecutor for MetaClient { async fn gc_regions( &self, - _ctx: &ExecutorContext, + ctx: &ExecutorContext, request: GcRegionsRequest, ) -> MetaResult { - self.gc_regions(request) + self.gc_regions(ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -429,10 +432,10 @@ impl ProcedureExecutor for MetaClient { async fn gc_table( &self, - _ctx: &ExecutorContext, + ctx: &ExecutorContext, request: GcTableRequest, ) -> MetaResult { - self.gc_table(request) + self.gc_table(ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -807,10 +810,12 @@ impl MetaClient { /// Submit a region migration task. pub async fn migrate_region( &self, + context: &ExecutorContext, request: MigrateRegionRequest, ) -> Result { self.procedure_client()? .migrate_region( + context, request.region_id, request.from_peer, request.to_peer, @@ -825,28 +830,57 @@ impl MetaClient { } /// Manually trigger GC for specific regions. - pub async fn gc_regions(&self, request: GcRegionsRequest) -> Result { - self.procedure_client()?.gc_regions(request).await + pub async fn gc_regions( + &self, + context: &ExecutorContext, + request: GcRegionsRequest, + ) -> Result { + self.procedure_client()?.gc_regions(context, request).await } /// Manually trigger GC for a table (all its regions). - pub async fn gc_table(&self, request: GcTableRequest) -> Result { - self.procedure_client()?.gc_table(request).await + pub async fn gc_table( + &self, + context: &ExecutorContext, + request: GcTableRequest, + ) -> Result { + self.procedure_client()?.gc_table(context, request).await } - /// Submit a DDL task + /// Submit a DDL task. pub async fn submit_ddl_task( &self, - req: SubmitDdlTaskRequest, + context: ExecutorContext, + request: SubmitDdlTaskRequest, ) -> Result { - let res = self - .procedure_client()? - .submit_ddl_task(req.try_into().context(ConvertMetaRequestSnafu)?) + let event_context = procedure_event_context(&context); + let actor = procedure_actor(&context); + let mut query_context = context.query_context.context(MissingQueryContextSnafu)?; + query_context + .extensions + .remove(CREATE_DATABASE_CREATOR_EXTENSION_KEY); + if let DdlTask::CreateDatabase(CreateDatabaseTask { + creator: Some(creator), + .. + }) = &request.task + { + query_context.extensions.insert( + CREATE_DATABASE_CREATOR_EXTENSION_KEY.to_string(), + serde_json::to_string(creator).context(ConvertMetaConfigSnafu)?, + ); + } + + let mut request: api::v1::meta::DdlTaskRequest = + request.try_into().context(ConvertMetaRequestSnafu)?; + request.query_context = Some(api::v1::QueryContext::from(query_context)); + request.event_context = event_context; + request.actor = actor; + + self.procedure_client()? + .submit_ddl_task(request) .await? .try_into() - .context(ConvertMetaResponseSnafu)?; - - Ok(res) + .context(ConvertMetaResponseSnafu) } pub fn heartbeat_client(&self) -> Result { diff --git a/src/meta-client/src/client/procedure.rs b/src/meta-client/src/client/procedure.rs index d7d5ccfaff6..ccf2c28188f 100644 --- a/src/meta-client/src/client/procedure.rs +++ b/src/meta-client/src/client/procedure.rs @@ -20,11 +20,13 @@ use api::v1::meta::ddl_task_request::Task; use api::v1::meta::procedure_service_client::ProcedureServiceClient; use api::v1::meta::{ DdlTaskRequest, DdlTaskResponse, GcRegionsRequest, GcRegionsResponse, GcTableRequest, - GcTableResponse, MigrateRegionRequest, MigrateRegionResponse, ProcedureDetailRequest, - ProcedureDetailResponse, ProcedureId, ProcedureStateResponse, QueryProcedureRequest, - ReconcileRequest, ReconcileResponse, RequestHeader, ResponseHeader, Role, + GcTableResponse, MigrateRegionRequest, MigrateRegionResponse, ProcedureActor, + ProcedureDetailRequest, ProcedureDetailResponse, ProcedureEventContext, ProcedureId, + ProcedureStateResponse, QueryProcedureRequest, ReconcileRequest, ReconcileResponse, + RequestHeader, ResponseHeader, Role, }; use common_grpc::channel_manager::ChannelManager; +use common_meta::procedure_executor::ExecutorContext; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, }; @@ -43,6 +45,33 @@ use crate::client::{Id, LeaderProviderRef, util}; use crate::error; use crate::error::Result; +/// Builds the event context transported by a procedure RPC. +/// +/// The caller can only supply reason/extensions. Protocol is derived here from +/// the trusted, typed query channel held locally in the executor context. +pub(crate) fn procedure_event_context(context: &ExecutorContext) -> Option { + context.event_input.as_ref().map(|input| { + let mut event_context = ProcedureEventContext::from(input); + event_context.protocol = context + .query_context + .as_ref() + .and_then(|query_context| query_context.protocol()) + .unwrap_or_default(); + event_context + }) +} + +/// Builds the optional procedure actor transported by a procedure RPC. +pub(crate) fn procedure_actor(context: &ExecutorContext) -> Option { + context + .actor + .as_deref() + .filter(|username| !username.is_empty()) + .map(|username| ProcedureActor { + username: username.to_string(), + }) +} + #[derive(Clone, Debug)] pub struct Client { inner: Arc>, @@ -92,6 +121,7 @@ impl Client { /// - `timeout`: timeout for downgrading region and upgrading region operations pub async fn migrate_region( &self, + context: &ExecutorContext, region_id: u64, from_peer: u64, to_peer: u64, @@ -99,7 +129,7 @@ impl Client { ) -> Result { let inner = self.inner.read().await; inner - .migrate_region(region_id, from_peer, to_peer, timeout) + .migrate_region(context, region_id, from_peer, to_peer, timeout) .await } @@ -114,14 +144,22 @@ impl Client { inner.list_procedures().await } - pub async fn gc_regions(&self, request: MetaGcRegionsRequest) -> Result { + pub async fn gc_regions( + &self, + context: &ExecutorContext, + request: MetaGcRegionsRequest, + ) -> Result { let inner = self.inner.read().await; - inner.gc_regions(request).await + inner.gc_regions(context, request).await } - pub async fn gc_table(&self, request: MetaGcTableRequest) -> Result { + pub async fn gc_table( + &self, + context: &ExecutorContext, + request: MetaGcTableRequest, + ) -> Result { let inner = self.inner.read().await; - inner.gc_table(request).await + inner.gc_table(context, request).await } } @@ -225,6 +263,7 @@ impl Inner { async fn migrate_region( &self, + context: &ExecutorContext, region_id: u64, from_peer: u64, to_peer: u64, @@ -235,6 +274,8 @@ impl Inner { from_peer, to_peer, timeout_secs: timeout.as_secs() as u32, + event_context: procedure_event_context(context), + actor: procedure_actor(context), ..Default::default() }; @@ -278,7 +319,11 @@ impl Inner { .await } - async fn gc_regions(&self, request: MetaGcRegionsRequest) -> Result { + async fn gc_regions( + &self, + context: &ExecutorContext, + request: MetaGcRegionsRequest, + ) -> Result { let timeout = request.timeout; let req = GcRegionsRequest { header: Some(RequestHeader { @@ -290,6 +335,8 @@ impl Inner { region_ids: request.region_ids, full_file_listing: request.full_file_listing, timeout_secs: gc_timeout_secs(timeout), + event_context: procedure_event_context(context), + actor: procedure_actor(context), }; let resp: GcRegionsResponse = self @@ -315,7 +362,11 @@ impl Inner { }) } - async fn gc_table(&self, request: MetaGcTableRequest) -> Result { + async fn gc_table( + &self, + context: &ExecutorContext, + request: MetaGcTableRequest, + ) -> Result { let timeout = request.timeout; let req = GcTableRequest { header: Some(RequestHeader { @@ -329,6 +380,8 @@ impl Inner { table_name: request.table_name, full_file_listing: request.full_file_listing, timeout_secs: gc_timeout_secs(timeout), + event_context: procedure_event_context(context), + actor: procedure_actor(context), }; let resp: GcTableResponse = self @@ -460,11 +513,14 @@ mod tests { ReconcileResponse, ResponseHeader, Role, }; use async_trait::async_trait; + use common_base::protocol::Channel; use common_error::status_code::StatusCode; + use common_event_recorder::{PersistentEventContext, ProcedureEventInput}; + use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor}; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, CommentObjectType, CommentOnTask, CreatorGrantIntent, DdlTask, QueryContext, - SubmitDdlTaskRequest, + SubmitDdlTaskRequest, TriggerReason, }; use common_telemetry::common_error::ext::ErrorExt; use common_telemetry::info; @@ -474,8 +530,8 @@ mod tests { use tonic::codec::CompressionEncoding; use tonic::{Request, Response, Status}; - use super::gc_timeout_secs; use crate::client::MetaClientBuilder; + use crate::client::procedure::{gc_timeout_secs, procedure_actor, procedure_event_context}; #[test] fn test_gc_timeout_secs() { @@ -486,6 +542,57 @@ mod tests { assert_eq!(gc_timeout_secs(Some(Duration::from_secs(10))), 10); } + #[test] + fn test_procedure_event_context_derives_protocol_from_query_context() { + let context = ExecutorContext { + query_context: Some(QueryContext { + channel: Channel::Postgres as u8, + ..Default::default() + }), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + + assert_eq!( + procedure_event_context(&context).map(PersistentEventContext::from), + Some(PersistentEventContext::new(TriggerReason::Manual).with_protocol("postgres")) + ); + + let automatic_context = ExecutorContext { + event_input: Some(ProcedureEventInput::new(TriggerReason::ScheduledGc)), + ..Default::default() + }; + assert_eq!( + procedure_event_context(&automatic_context).map(PersistentEventContext::from), + Some(PersistentEventContext::new(TriggerReason::ScheduledGc)) + ); + + let unknown_channel_context = ExecutorContext { + query_context: Some(QueryContext::default()), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + assert_eq!( + procedure_event_context(&unknown_channel_context).map(PersistentEventContext::from), + Some(PersistentEventContext::new(TriggerReason::Manual)) + ); + } + + #[test] + fn test_procedure_actor() { + let context = ExecutorContext { + actor: Some(String::new()), + ..Default::default() + }; + assert!(procedure_actor(&context).is_none()); + + let context = ExecutorContext { + actor: Some("alice".to_string()), + ..Default::default() + }; + assert_eq!(procedure_actor(&context).unwrap().username, "alice"); + } + #[derive(Clone)] struct MockHeartbeat { leader_addr: String, @@ -621,19 +728,28 @@ mod tests { username: "alice".to_string(), created_at_ns: 42, }; - client - .submit_ddl_task(SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::new_create_database( - "greptime".to_string(), - "metrics".to_string(), - false, - Default::default(), - Some(creator.clone()), - ), - )) - .await - .unwrap(); + let executor_context = |actor: String| ExecutorContext { + query_context: Some(QueryContext { + channel: Channel::Postgres as u8, + ..Default::default() + }), + actor: Some(actor), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + ProcedureExecutor::submit_ddl_task( + &client, + executor_context("effective-user".to_string()), + SubmitDdlTaskRequest::new(DdlTask::new_create_database( + "greptime".to_string(), + "metrics".to_string(), + false, + Default::default(), + Some(creator.clone()), + )), + ) + .await + .unwrap(); let request = request_rx.recv().await.unwrap(); let encoded = serde_json::to_string(&creator).unwrap(); @@ -647,9 +763,36 @@ mod tests { .as_ref(), encoded.as_bytes() ); - let extensions = &request.into_inner().query_context.unwrap().extensions; + let request = request.into_inner(); + assert_eq!(request.actor.unwrap().username, "effective-user"); + assert_eq!( + PersistentEventContext::from(request.event_context.unwrap()), + PersistentEventContext::new(TriggerReason::Manual).with_protocol("postgres") + ); + let extensions = &request.query_context.unwrap().extensions; assert_eq!(extensions[CREATE_DATABASE_CREATOR_EXTENSION_KEY], encoded); + ProcedureExecutor::submit_ddl_task( + &client, + executor_context(String::new()), + SubmitDdlTaskRequest::new(DdlTask::new_drop_database( + "greptime".to_string(), + "metrics".to_string(), + false, + )), + ) + .await + .unwrap(); + assert!( + request_rx + .recv() + .await + .unwrap() + .into_inner() + .actor + .is_none() + ); + server_handle.abort(); } @@ -689,22 +832,28 @@ mod tests { .build(); client.start(&[addr_str.as_str()]).await.unwrap(); - let mut request = SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::new_comment_on(CommentOnTask { - catalog_name: "greptime".to_string(), - schema_name: "public".to_string(), - object_type: CommentObjectType::Table, - object_name: "test_table".to_string(), - column_name: None, - object_id: None, - comment: Some("timeout".to_string()), - }), - ); + let mut request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(CommentOnTask { + catalog_name: "greptime".to_string(), + schema_name: "public".to_string(), + object_type: CommentObjectType::Table, + object_name: "test_table".to_string(), + column_name: None, + object_id: None, + comment: Some("timeout".to_string()), + })); request.timeout = Duration::from_secs(1); let now = Instant::now(); - let err = client.submit_ddl_task(request).await.unwrap_err(); + let err = client + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + request, + ) + .await + .unwrap_err(); let elapsed = now.elapsed(); // The request should be cancelled within 1 second. assert!(elapsed < Duration::from_secs(2)); diff --git a/src/meta-client/src/error.rs b/src/meta-client/src/error.rs index 5121bba847a..4623cfee33b 100644 --- a/src/meta-client/src/error.rs +++ b/src/meta-client/src/error.rs @@ -67,6 +67,12 @@ pub enum Error { location: Location, }, + #[snafu(display("Procedure submission requires query context"))] + MissingQueryContext { + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to send heartbeat: {}", err_msg))] SendHeartbeat { err_msg: String, @@ -141,6 +147,7 @@ impl ErrorExt for Error { | Error::NoLeader { .. } | Error::AskLeaderTimeout { .. } | Error::NotStarted { .. } + | Error::MissingQueryContext { .. } | Error::SendHeartbeat { .. } | Error::CreateHeartbeatStream { .. } | Error::CreateChannel { .. } diff --git a/src/meta-srv/src/gc/ctx.rs b/src/meta-srv/src/gc/ctx.rs index c9cc8049867..4b46a941e25 100644 --- a/src/meta-srv/src/gc/ctx.rs +++ b/src/meta-srv/src/gc/ctx.rs @@ -29,10 +29,9 @@ use common_meta::key::DroppedTableName; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; -use common_meta::rpc::ddl::PersistentEventContext; #[cfg(feature = "enterprise")] use common_meta::rpc::ddl::PurgeDroppedTableTask; -use common_procedure::{ProcedureManagerRef, ProcedureWithId, watcher}; +use common_procedure::{ProcedureContext, ProcedureManagerRef, ProcedureWithId, watcher}; use common_telemetry::debug; use snafu::{OptionExt as _, ResultExt as _}; use store_api::storage::{GcReport, RegionId}; @@ -68,7 +67,7 @@ pub(crate) trait SchedulerCtx: Send + Sync { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result; #[cfg(feature = "enterprise")] @@ -256,14 +255,14 @@ impl SchedulerCtx for DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { self.gc_regions_inner( region_ids, full_file_listing, timeout, region_routes_override, - event_context, + procedure_context, ) .await } @@ -310,7 +309,7 @@ impl DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { debug!( "Sending GC instruction for {} regions (full_file_listing: {})", @@ -328,7 +327,7 @@ impl DefaultGcSchedulerCtx { region_routes_override, ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let id = procedure_with_id.id; diff --git a/src/meta-srv/src/gc/handler.rs b/src/meta-srv/src/gc/handler.rs index 16b29a39c7e..eda87fbb144 100644 --- a/src/meta-srv/src/gc/handler.rs +++ b/src/meta-srv/src/gc/handler.rs @@ -18,7 +18,7 @@ use std::time::Instant; use common_catalog::consts::MITO_ENGINE; use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; +use common_procedure::ProcedureContext; use common_telemetry::tracing::Instrument as _; use common_telemetry::{debug, error, info, warn}; use futures::StreamExt; @@ -39,7 +39,7 @@ use crate::metrics::METRIC_META_GC_CANDIDATE_REGIONS; impl GcScheduler { pub(crate) async fn trigger_gc( &self, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let start_time = Instant::now(); info!("Starting GC cycle"); @@ -111,7 +111,7 @@ impl GcScheduler { datanode_to_candidates, dropped_assignment.force_full_listing, dropped_assignment.region_routes_override, - event_context, + procedure_context, ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_dispatch_to_datanodes" @@ -215,7 +215,7 @@ impl GcScheduler { datanode_to_candidates: HashMap>, force_full_listing_by_peer: HashMap>, region_routes_override_by_peer: HashMap, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> GcJobReport { let mut per_datanode_reports = HashMap::new(); let mut failed_datanodes: HashMap<_, Vec<_>> = HashMap::new(); @@ -237,7 +237,7 @@ impl GcScheduler { .get(&peer) .cloned() .unwrap_or_default(); - let event_context = event_context.clone(); + let procedure_context = procedure_context.clone(); async move { ( peer, @@ -247,7 +247,7 @@ impl GcScheduler { candidates, force_full_listing, region_routes_override, - event_context, + procedure_context, ) .await, ) @@ -286,7 +286,7 @@ impl GcScheduler { candidates: Vec<(TableId, GcCandidate)>, force_full_listing: HashSet, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { info!( "Starting GC for datanode {} with {} candidate regions", @@ -336,7 +336,7 @@ impl GcScheduler { false, self.config.mailbox_timeout, region_routes_override.clone(), - event_context.clone(), + procedure_context.clone(), ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_call_datanode", @@ -369,7 +369,7 @@ impl GcScheduler { true, self.config.mailbox_timeout, region_routes_override, - event_context, + procedure_context, ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_call_datanode", diff --git a/src/meta-srv/src/gc/mock.rs b/src/meta-srv/src/gc/mock.rs index 94749dfb469..dcf913462c2 100644 --- a/src/meta-srv/src/gc/mock.rs +++ b/src/meta-srv/src/gc/mock.rs @@ -26,12 +26,13 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use common_catalog::consts::MITO_ENGINE; +use common_event_recorder::PersistentEventContext; use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::{Region, RegionRoute}; +use common_procedure::ProcedureContext; use common_telemetry::debug; use ordered_float::OrderedFloat; use store_api::region_engine::RegionRole; @@ -214,10 +215,13 @@ impl SchedulerCtx for MockSchedulerCtx { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { *self.gc_regions_calls.lock().unwrap() += 1; - self.gc_event_contexts.lock().unwrap().push(event_context); + self.gc_event_contexts + .lock() + .unwrap() + .push(procedure_context.event_context.unwrap_or_default()); // Check per-region error injection first (for any region) for ®ion_id in region_ids { diff --git a/src/meta-srv/src/gc/mock/basic.rs b/src/meta-srv/src/gc/mock/basic.rs index ed051cc839e..70f7a10b1a1 100644 --- a/src/meta-srv/src/gc/mock/basic.rs +++ b/src/meta-srv/src/gc/mock/basic.rs @@ -16,8 +16,10 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::Instant; +use common_event_recorder::PersistentEventContext; use common_meta::peer::Peer; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; +use common_meta::rpc::ddl::TriggerReason; +use common_procedure::ProcedureContext; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -36,7 +38,7 @@ async fn test_parallel_process_datanodes_empty() { HashMap::new(), HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -105,7 +107,7 @@ async fn test_parallel_process_datanodes_with_candidates() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -245,7 +247,18 @@ async fn test_handle_manual_gc_without_regions_records_manual_event_context() { last_tracker_cleanup: Arc::new(tokio::sync::Mutex::new(Instant::now())), }; - scheduler.handle_manual_gc(None, None, None).await.unwrap(); + scheduler + .handle_manual_gc( + None, + None, + None, + ProcedureContext { + event_context: Some(PersistentEventContext::new(TriggerReason::Manual)), + ..Default::default() + }, + ) + .await + .unwrap(); assert_eq!( ctx.gc_event_contexts.lock().unwrap().as_slice(), diff --git a/src/meta-srv/src/gc/mock/concurrent.rs b/src/meta-srv/src/gc/mock/concurrent.rs index 5bce9a1be8b..bfb7ab1d26c 100644 --- a/src/meta-srv/src/gc/mock/concurrent.rs +++ b/src/meta-srv/src/gc/mock/concurrent.rs @@ -18,8 +18,8 @@ use std::time::{Duration, Instant}; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::{Region, RegionRoute}; +use common_procedure::ProcedureContext; use common_telemetry::{info, init_default_ut_logging}; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -102,7 +102,7 @@ async fn test_concurrent_table_processing_limits() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -190,7 +190,7 @@ async fn test_datanode_processes_tables_with_partial_gc_failures() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -296,7 +296,7 @@ async fn test_region_gc_concurrency_limit() { candidates.into_iter().map(|c| (table_id, c)).collect(), HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -413,7 +413,7 @@ async fn test_region_gc_concurrency_with_partial_failures() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -557,7 +557,7 @@ async fn test_region_gc_concurrency_with_retryable_errors() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; diff --git a/src/meta-srv/src/gc/mock/full_list.rs b/src/meta-srv/src/gc/mock/full_list.rs index 8193a9094e1..843e45134da 100644 --- a/src/meta-srv/src/gc/mock/full_list.rs +++ b/src/meta-srv/src/gc/mock/full_list.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; +use common_procedure::ProcedureContext; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -83,7 +83,7 @@ async fn test_full_file_listing_first_time_gc() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -157,7 +157,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -184,7 +184,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -260,7 +260,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -284,7 +284,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); diff --git a/src/meta-srv/src/gc/mock/misc.rs b/src/meta-srv/src/gc/mock/misc.rs index 575a5c5fc08..fabf5bcae4b 100644 --- a/src/meta-srv/src/gc/mock/misc.rs +++ b/src/meta-srv/src/gc/mock/misc.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::Instant; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; +use common_procedure::ProcedureContext; use common_telemetry::init_default_ut_logging; use store_api::storage::{FileRefsManifest, GcReport, RegionId}; @@ -76,7 +76,7 @@ async fn test_empty_file_refs_manifest() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -169,7 +169,7 @@ async fn test_multiple_regions_per_table() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; diff --git a/src/meta-srv/src/gc/scheduler.rs b/src/meta-srv/src/gc/scheduler.rs index b89c0d06be9..857d4deb57c 100644 --- a/src/meta-srv/src/gc/scheduler.rs +++ b/src/meta-srv/src/gc/scheduler.rs @@ -16,9 +16,11 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::{Duration, Instant}; +use common_event_recorder::PersistentEventContext; use common_meta::DatanodeId; use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; +use common_meta::rpc::ddl::TriggerReason; +use common_procedure::ProcedureContext; use common_telemetry::tracing::Instrument as _; use common_telemetry::{error, info}; use snafu::ResultExt; @@ -104,6 +106,7 @@ pub enum Event { full_file_listing: Option, /// Optional override for timeout. If None, uses scheduler config. timeout: Option, + procedure_context: ProcedureContext, }, } @@ -182,12 +185,13 @@ impl GcScheduler { region_ids, full_file_listing, timeout, + procedure_context, } => { info!("Received manually gc request"); let span = common_telemetry::tracing::info_span!("meta_gc_tick", trigger = "manual"); let result = self - .handle_manual_gc(region_ids, full_file_listing, timeout) + .handle_manual_gc(region_ids, full_file_listing, timeout, procedure_context) .instrument(span) .await; if let Err(e) = &result { @@ -218,7 +222,10 @@ impl GcScheduler { } let span = common_telemetry::tracing::info_span!("meta_gc_handle_tick"); let report = self - .trigger_gc(PersistentEventContext::new(TriggerReason::ScheduledGc)) + .trigger_gc(ProcedureContext { + actor: None, + event_context: Some(PersistentEventContext::new(TriggerReason::ScheduledGc)), + }) .instrument(span) .await?; @@ -294,6 +301,7 @@ impl GcScheduler { region_ids: Option>, full_file_listing: Option, timeout: Option, + procedure_context: ProcedureContext, ) -> Result { info!("Start to handle manual gc request"); @@ -304,9 +312,7 @@ impl GcScheduler { // No specific regions, use default tick behavior let Some(regions) = region_ids else { - let report = self - .trigger_gc(PersistentEventContext::new(TriggerReason::Manual)) - .await?; + let report = self.trigger_gc(procedure_context).await?; info!("Finished manual gc request"); return Ok(report); }; @@ -353,7 +359,7 @@ impl GcScheduler { full_listing, gc_timeout, Region2Peers::new(), - PersistentEventContext::new(TriggerReason::Manual), + procedure_context.clone(), ) .await?; combined_report.merge(report); @@ -367,7 +373,7 @@ impl GcScheduler { true, gc_timeout, dropped_routes_override, - PersistentEventContext::new(TriggerReason::Manual), + procedure_context, ) .await?; combined_report.merge(report); @@ -490,7 +496,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: PersistentEventContext, + _procedure_context: ProcedureContext, ) -> Result { self.gc_regions_calls.fetch_add(1, Ordering::Relaxed); panic!("gc_regions should not be called in maintenance mode") @@ -564,7 +570,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: PersistentEventContext, + _procedure_context: ProcedureContext, ) -> Result { Ok(GcReport::default()) } @@ -861,7 +867,10 @@ mod tests { *ctx.dropped_tables.lock().unwrap() = vec![dropped_table(1, Some(i64::MIN))]; let scheduler = soft_drop_scheduler(ctx.clone()); - scheduler.handle_manual_gc(None, None, None).await.unwrap(); + scheduler + .handle_manual_gc(None, None, None, ProcedureContext::default()) + .await + .unwrap(); assert!(ctx.purge_attempts.lock().unwrap().is_empty()); } @@ -898,7 +907,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: PersistentEventContext, + _procedure_context: ProcedureContext, ) -> Result { crate::error::UnexpectedSnafu { violated: "mock gc failure".to_string(), @@ -949,6 +958,7 @@ mod tests { Some(vec![RegionId::new(1, 0)]), Some(false), Some(Duration::from_secs(1)), + ProcedureContext::default(), ) .await; @@ -977,6 +987,7 @@ mod tests { Some(vec![RegionId::new(1, 0)]), Some(false), Some(Duration::from_secs(1)), + ProcedureContext::default(), ) .await; diff --git a/src/meta-srv/src/procedure/region_migration.rs b/src/meta-srv/src/procedure/region_migration.rs index ad1553c92d7..75ab5c79397 100644 --- a/src/meta-srv/src/procedure/region_migration.rs +++ b/src/meta-srv/src/procedure/region_migration.rs @@ -32,7 +32,7 @@ use std::fmt::{Debug, Display}; use std::time::Duration; use common_error::ext::BoxedError; -use common_event_recorder::Event; +use common_event_recorder::{Event, PersistentEventContext}; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::ddl::RegionFailureDetectorControllerRef; use common_meta::instruction::CacheIdent; @@ -45,7 +45,7 @@ use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; use common_meta::lock_key::{CatalogLock, RegionLock, SchemaLock, TableLock}; use common_meta::peer::Peer; use common_meta::region_keeper::{MemoryRegionKeeperRef, OperatingRegionGuard}; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; +use common_meta::rpc::ddl::TriggerReason; use common_procedure::error::{ Error as ProcedureError, FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu, }; diff --git a/src/meta-srv/src/procedure/region_migration/manager.rs b/src/meta-srv/src/procedure/region_migration/manager.rs index e1d719d9747..c21de506243 100644 --- a/src/meta-srv/src/procedure/region_migration/manager.rs +++ b/src/meta-srv/src/procedure/region_migration/manager.rs @@ -18,12 +18,14 @@ use std::fmt::Display; use std::sync::{Arc, RwLock}; use std::time::Duration; +use common_event_recorder::PersistentEventContext; use common_meta::key::table_info::TableInfoValue; use common_meta::key::table_route::TableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::RegionRoute; -use common_procedure::{ProcedureId, ProcedureManagerRef, ProcedureWithId, watcher}; +use common_procedure::{ + ProcedureContext, ProcedureId, ProcedureManagerRef, ProcedureWithId, watcher, +}; use common_telemetry::{error, info, warn}; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt, ensure}; @@ -458,8 +460,20 @@ impl RegionMigrationManager { } let submitting_region_ids = task.region_ids.clone(); + // TODO(weny): Pass the actor when region migration task submission propagates it. + let procedure_context = ProcedureContext { + actor: None, + event_context: Some(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + }; let procedure_id = self - .submit_procedure_inner(task, procedure_guards, catalog_and_schema) + .submit_procedure_inner( + procedure_context, + task, + procedure_guards, + catalog_and_schema, + ) .await?; Ok(SubmitRegionMigrationTaskResult { migrated, @@ -475,11 +489,11 @@ impl RegionMigrationManager { async fn submit_procedure_inner( &self, + procedure_context: ProcedureContext, task: RegionMigrationTaskBatch, procedure_guards: Vec, catalog_and_schema: Vec<(String, String)>, ) -> Result { - let event_context = PersistentEventContext::new(task.trigger_reason.to_trigger_reason()); let procedure = RegionMigrationProcedure::new( PersistentContext::new( catalog_and_schema, @@ -493,7 +507,7 @@ impl RegionMigrationManager { procedure_guards, ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); let procedure_manager = self.procedure_manager.clone(); @@ -529,8 +543,14 @@ impl RegionMigrationManager { /// Submits a new region migration procedure. pub async fn submit_procedure( &self, + procedure_context: ProcedureContext, mut task: RegionMigrationProcedureTask, ) -> Result> { + if let Some(event_context) = procedure_context.event_context.as_ref() { + task.trigger_reason = + RegionMigrationTriggerReason::from_trigger_reason(event_context.reason); + } + let Some(guard) = self.insert_running_procedure(&task) else { return error::MigrationRunningSnafu { region_id: task.region_id, @@ -577,7 +597,6 @@ impl RegionMigrationManager { timeout, trigger_reason, } = task.clone(); - let event_context = PersistentEventContext::new(trigger_reason.to_trigger_reason()); let procedure = RegionMigrationProcedure::new( PersistentContext::new( vec![(catalog_name, schema_name)], @@ -591,7 +610,7 @@ impl RegionMigrationManager { vec![guard], ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); let procedure_manager = self.procedure_manager.clone(); @@ -655,7 +674,15 @@ mod test { .unwrap() .insert(region_id, task.clone()); - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::MigrationRunning { .. }); } @@ -673,7 +700,15 @@ mod test { trigger_reason: RegionMigrationTriggerReason::Manual, }; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::InvalidArguments { .. }); } @@ -691,7 +726,15 @@ mod test { trigger_reason: RegionMigrationTriggerReason::Manual, }; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::TableRouteNotFound { .. }); } @@ -719,7 +762,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::RegionRouteNotFound { .. }); } @@ -747,7 +798,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::LeaderPeerChanged { .. }); assert_eq!( err.to_string(), @@ -780,7 +839,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::InvalidArguments { .. }); assert_eq!( err.to_string(), @@ -813,7 +880,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::RegionMigrated { .. }); } diff --git a/src/meta-srv/src/service/procedure.rs b/src/meta-srv/src/service/procedure.rs index 1b5c4c42dbc..9e226bc18fe 100644 --- a/src/meta-srv/src/service/procedure.rs +++ b/src/meta-srv/src/service/procedure.rs @@ -18,10 +18,12 @@ use api::v1::meta::reconcile_request::Target; use api::v1::meta::{ DdlTaskRequest as PbDdlTaskRequest, DdlTaskResponse as PbDdlTaskResponse, GcRegionsRequest, GcRegionsResponse, GcStats, GcTableRequest, GcTableResponse, MigrateRegionRequest, - MigrateRegionResponse, ProcedureDetailRequest, ProcedureDetailResponse, ProcedureStateResponse, + MigrateRegionResponse, ProcedureActor, ProcedureDetailRequest, ProcedureDetailResponse, + ProcedureEventContext as PbProcedureEventContext, ProcedureStateResponse, QueryProcedureRequest, ReconcileCatalog, ReconcileDatabase, ReconcileRequest, ReconcileResponse, ReconcileTable, ResolveStrategy, Role, procedure_service_server, }; +use common_event_recorder::{PersistentEventContext, ProcedureEventInput}; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_name::TableNameKey; use common_meta::procedure_executor::ExecutorContext; @@ -33,6 +35,7 @@ use common_meta::rpc::procedure::{ self, GcRegionsRequest as MetaGcRegionsRequest, GcResponse, GcTableRequest as MetaGcTableRequest, }; +use common_procedure::ProcedureContext; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; use table::table_reference::TableReference; @@ -47,6 +50,45 @@ use crate::procedure::region_migration::manager::{ use crate::service::GrpcResult; use crate::{check_leader, error, gc}; +struct ProcedureSubmission { + actor: Option, + event_context: Option, +} + +impl TryFrom<(Option, Option)> for ProcedureSubmission { + type Error = Status; + + fn try_from( + (actor, event_context): (Option, Option), + ) -> Result { + let actor = actor + .map(|actor| { + if actor.username.is_empty() { + Err(Status::invalid_argument( + "Procedure actor must not be empty", + )) + } else { + Ok(actor.username) + } + }) + .transpose()?; + + Ok(Self { + actor, + event_context, + }) + } +} + +impl From for ProcedureContext { + fn from(context: ProcedureSubmission) -> Self { + Self { + actor: context.actor, + event_context: context.event_context.map(PersistentEventContext::from), + } + } +} + #[async_trait::async_trait] impl procedure_service_server::ProcedureService for Metasrv { async fn query( @@ -89,9 +131,15 @@ impl procedure_service_server::ProcedureService for Metasrv { task, wait, timeout_secs, + event_context, + actor, } = request; let header = header.context(error::MissingRequestHeaderSnafu)?; + let ProcedureSubmission { + actor, + event_context, + } = ProcedureSubmission::try_from((actor, event_context))?; let mut query_context = query_context .context(error::MissingRequiredParameterSnafu { param: "query_context", @@ -102,15 +150,17 @@ impl procedure_service_server::ProcedureService for Metasrv { .try_into() .context(error::ConvertProtoDataSnafu)?; restore_create_database_creator(&metadata, header.role, &mut task, &mut query_context)?; + let executor_context = ExecutorContext { + tracing_context: Some(header.tracing_context), + query_context: Some(query_context), + actor, + event_input: event_context.map(ProcedureEventInput::from), + }; let resp = self .ddl_manager() .submit_ddl_task( - &ExecutorContext { - tracing_context: Some(header.tracing_context), - ..Default::default() - }, + executor_context, SubmitDdlTaskRequest { - query_context, wait, timeout: Duration::from_secs(timeout_secs.into()), task, @@ -135,9 +185,13 @@ impl procedure_service_server::ProcedureService for Metasrv { from_peer, to_peer, timeout_secs, + event_context, + actor, } = request.into_inner(); let _header = header.context(error::MissingRequestHeaderSnafu)?; + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); let from_peer = self .lookup_datanode_peer(from_peer) .await? @@ -149,13 +203,16 @@ impl procedure_service_server::ProcedureService for Metasrv { let pid = self .region_migration_manager() - .submit_procedure(RegionMigrationProcedureTask { - region_id: region_id.into(), - from_peer, - to_peer, - timeout: Duration::from_secs(timeout_secs.into()), - trigger_reason: RegionMigrationTriggerReason::Manual, - }) + .submit_procedure( + procedure_context, + RegionMigrationProcedureTask { + region_id: region_id.into(), + from_peer, + to_peer, + timeout: Duration::from_secs(timeout_secs.into()), + trigger_reason: RegionMigrationTriggerReason::Manual, + }, + ) .await? .map(procedure::pid_to_pb_pid); @@ -266,16 +323,23 @@ impl procedure_service_server::ProcedureService for Metasrv { region_ids, full_file_listing, timeout_secs, + event_context, + actor, } = request.into_inner(); let _header = header.context(error::MissingRequestHeaderSnafu)?; + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); let response = self - .handle_gc_regions(MetaGcRegionsRequest { - region_ids, - full_file_listing, - timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), - }) + .handle_gc_regions( + procedure_context, + MetaGcRegionsRequest { + region_ids, + full_file_listing, + timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), + }, + ) .await?; Ok(Response::new(gc_response_to_regions_pb(response))) @@ -291,18 +355,25 @@ impl procedure_service_server::ProcedureService for Metasrv { table_name, full_file_listing, timeout_secs, + event_context, + actor, } = request.into_inner(); let _header = header.context(error::MissingRequestHeaderSnafu)?; + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); let response = self - .handle_gc_table(MetaGcTableRequest { - catalog_name, - schema_name, - table_name, - full_file_listing, - timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), - }) + .handle_gc_table( + procedure_context, + MetaGcTableRequest { + catalog_name, + schema_name, + table_name, + full_file_listing, + timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), + }, + ) .await?; Ok(Response::new(gc_response_to_table_pb(response))) @@ -371,17 +442,30 @@ impl Metasrv { } } - async fn handle_gc_regions(&self, request: MetaGcRegionsRequest) -> error::Result { + async fn handle_gc_regions( + &self, + procedure_context: ProcedureContext, + request: MetaGcRegionsRequest, + ) -> error::Result { let region_ids: Vec = request .region_ids .into_iter() .map(RegionId::from_u64) .collect(); - self.trigger_gc_for_regions(region_ids, request.full_file_listing, request.timeout) - .await + self.trigger_gc_for_regions( + procedure_context, + region_ids, + request.full_file_listing, + request.timeout, + ) + .await } - async fn handle_gc_table(&self, request: MetaGcTableRequest) -> error::Result { + async fn handle_gc_table( + &self, + procedure_context: ProcedureContext, + request: MetaGcTableRequest, + ) -> error::Result { let table_name_key = TableNameKey::new( &request.catalog_name, &request.schema_name, @@ -406,13 +490,19 @@ impl Metasrv { .context(TableMetadataManagerSnafu)?; let region_ids: Vec = route.region_routes.iter().map(|r| r.region.id).collect(); - self.trigger_gc_for_regions(region_ids, request.full_file_listing, request.timeout) - .await + self.trigger_gc_for_regions( + procedure_context, + region_ids, + request.full_file_listing, + request.timeout, + ) + .await } /// Triggers manual GC for specified regions and returns the GC response. async fn trigger_gc_for_regions( &self, + procedure_context: ProcedureContext, region_ids: Vec, full_file_listing: bool, timeout: Option, @@ -429,6 +519,7 @@ impl Metasrv { region_ids: Some(region_ids), full_file_listing: Some(full_file_listing), timeout, + procedure_context, }) .await .map_err(|_| { @@ -502,14 +593,15 @@ fn gc_response_to_table_pb(resp: GcResponse) -> GcTableResponse { mod tests { use std::time::Duration; - use api::v1::meta::Role; + use api::v1::meta::{ProcedureActor, ProcedureEventContext, Role}; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, CreatorGrantIntent, DdlTask, QueryContext, }; + use common_procedure::ProcedureContext; use tonic::metadata::{MetadataMap, MetadataValue}; - use super::{Metasrv, restore_create_database_creator}; + use super::{Metasrv, ProcedureSubmission, restore_create_database_creator}; fn create_database_task() -> DdlTask { DdlTask::new_create_database( @@ -646,4 +738,41 @@ mod tests { assert_eq!(task.creator, Some(creator)); assert!(trusted.extensions.is_empty()); } + + #[test] + fn test_procedure_submission_rejects_empty_actor() { + let submission = ProcedureSubmission::try_from(( + Some(ProcedureActor { + username: "alice".to_string(), + }), + Some(ProcedureEventContext { + reason: "manual".to_string(), + protocol: "postgres".to_string(), + extensions: Default::default(), + }), + )) + .unwrap(); + assert_eq!(submission.actor.as_deref(), Some("alice")); + assert_eq!( + submission + .event_context + .as_ref() + .map(|context| context.protocol.as_str()), + Some("postgres") + ); + + assert!( + ProcedureSubmission::try_from(( + Some(ProcedureActor { + username: String::new(), + }), + None, + )) + .is_err() + ); + + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((None, None)).unwrap()); + assert_eq!(procedure_context.event_context, None); + } } diff --git a/src/operator/src/procedure.rs b/src/operator/src/procedure.rs index 1a849afaf7a..03fb4f909c9 100644 --- a/src/operator/src/procedure.rs +++ b/src/operator/src/procedure.rs @@ -27,11 +27,12 @@ use common_meta::rpc::procedure::{ }; use common_query::error as query_error; use common_query::error::Result as QueryResult; +use session::context::QueryContextRef; use snafu::ResultExt; use table::table_name::TableName; use crate::error; -use crate::utils::to_meta_query_context_with_trigger_reason; +use crate::utils::to_executor_context; /// The operator for procedures which implements [`ProcedureServiceHandler`]. #[derive(Clone)] @@ -59,8 +60,8 @@ impl ProcedureServiceOperator { impl ProcedureServiceHandler for ProcedureServiceOperator { async fn purge_table( &self, + query_ctx: QueryContextRef, table_name: TableName, - query_ctx: session::context::QueryContextRef, ) -> QueryResult<()> { let dropped = self .table_metadata_manager @@ -76,24 +77,25 @@ impl ProcedureServiceHandler for ProcedureServiceOperator { }) .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu)?; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_purge_dropped_table(dropped.table_id), - ); + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_purge_dropped_table(dropped.table_id)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu)?; Ok(()) } - async fn migrate_region(&self, request: MigrateRegionRequest) -> QueryResult> { + async fn migrate_region( + &self, + query_ctx: QueryContextRef, + request: MigrateRegionRequest, + ) -> QueryResult> { + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); Ok(self .procedure_executor - .migrate_region(&ExecutorContext::default(), request) + .migrate_region(&executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu)? @@ -135,17 +137,27 @@ impl ProcedureServiceHandler for ProcedureServiceOperator { &self.catalog_manager } - async fn gc_regions(&self, request: MetaGcRegionsRequest) -> QueryResult { + async fn gc_regions( + &self, + query_ctx: QueryContextRef, + request: MetaGcRegionsRequest, + ) -> QueryResult { + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); self.procedure_executor - .gc_regions(&ExecutorContext::default(), request) + .gc_regions(&executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu) } - async fn gc_table(&self, request: MetaGcTableRequest) -> QueryResult { + async fn gc_table( + &self, + query_ctx: QueryContextRef, + request: MetaGcTableRequest, + ) -> QueryResult { + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); self.procedure_executor - .gc_table(&ExecutorContext::default(), request) + .gc_table(&executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu) @@ -162,10 +174,10 @@ mod tests { use common_meta::key::table_route::TableRouteValue; use common_meta::key::test_utils::new_test_table_info_with_name; use common_meta::kv_backend::memory::MemoryKvBackend; - use common_meta::procedure_executor::{ProcedureExecutor, ProcedureExecutorRef}; - use common_meta::rpc::ddl::{ - DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TRIGGER_REASON_EXTENSION_KEY, + use common_meta::procedure_executor::{ + ExecutorContext, ProcedureExecutor, ProcedureExecutorRef, }; + use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; use common_meta::rpc::procedure::{MigrateRegionResponse, ProcedureStateResponse}; use session::context::QueryContextBuilder; use table::table_name::TableName; @@ -175,6 +187,7 @@ mod tests { #[derive(Default)] struct RecordingProcedureExecutor { requests: Mutex>, + contexts: Mutex>, fail: bool, } @@ -182,9 +195,10 @@ mod tests { impl ProcedureExecutor for RecordingProcedureExecutor { async fn submit_ddl_task( &self, - _: &ExecutorContext, + context: ExecutorContext, request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { + self.contexts.lock().unwrap().push(context); self.requests.lock().unwrap().push(request); if self.fail { return common_meta::error::UnsupportedSnafu { @@ -273,20 +287,19 @@ mod tests { .build() .into(); - operator.purge_table(name.clone(), query_ctx).await.unwrap(); + operator.purge_table(query_ctx, name.clone()).await.unwrap(); let requests = executor.requests.lock().unwrap(); assert!( matches!(&requests[0].task, DdlTask::PurgeDroppedTable(task) if task.table_id == 42) ); - assert_eq!(requests[0].query_context.current_catalog, "catalog"); - assert_eq!(requests[0].query_context.current_schema, "schema"); + let contexts = executor.contexts.lock().unwrap(); + let query_context = contexts[0].query_context.as_ref().unwrap(); + assert_eq!(query_context.current_catalog, "catalog"); + assert_eq!(query_context.current_schema, "schema"); assert_eq!( - Some(&TriggerReason::Manual.as_ref().to_string()), - requests[0] - .query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) + contexts[0].event_input.as_ref().map(|input| input.reason), + Some(TriggerReason::Manual) ); } @@ -303,12 +316,12 @@ mod tests { ); let error = operator .purge_table( - TableName::new("catalog", "schema", "missing"), QueryContextBuilder::default() .current_catalog("catalog".to_string()) .current_schema("schema".to_string()) .build() .into(), + TableName::new("catalog", "schema", "missing"), ) .await .unwrap_err(); @@ -326,12 +339,12 @@ mod tests { assert!( operator .purge_table( - name, QueryContextBuilder::default() .current_catalog("catalog".to_string()) .current_schema("schema".to_string()) .build() - .into() + .into(), + name, ) .await .is_err() diff --git a/src/operator/src/statement/comment.rs b/src/operator/src/statement/comment.rs index 824287fc80b..c55e1be9f74 100644 --- a/src/operator/src/statement/comment.rs +++ b/src/operator/src/statement/comment.rs @@ -15,7 +15,6 @@ use api::v1::CommentOnExpr; use common_error::ext::BoxedError; use common_meta::cache_invalidator::Context; -use common_meta::procedure_executor::ExecutorContext; use common_meta::rpc::ddl::{ CommentObjectType, CommentOnTask, DdlTask, SubmitDdlTaskRequest, TriggerReason, }; @@ -30,7 +29,7 @@ use crate::error::{ self, ExecuteDdlSnafu, ExternalSnafu, InvalidSqlSnafu, Result, TableMetadataManagerSnafu, }; use crate::statement::StatementExecutor; -use crate::utils::to_meta_query_context_with_trigger_reason; +use crate::utils::to_executor_context; impl StatementExecutor { /// Adds a comment to a database object (table, column, or flow). @@ -54,13 +53,11 @@ impl StatementExecutor { .context(TableMetadataManagerSnafu)?; let cache_idents = comment_on_task.cache_idents(); - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); - let request = - SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_comment_on(comment_on_task)); + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(comment_on_task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(ExecuteDdlSnafu)?; @@ -88,13 +85,11 @@ impl StatementExecutor { .context(TableMetadataManagerSnafu)?; let cache_idents = comment_on_task.cache_idents(); - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); - let request = - SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_comment_on(comment_on_task)); + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(comment_on_task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(ExecuteDdlSnafu)?; diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 256b849123e..3c10186ca4e 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -47,7 +47,6 @@ use common_meta::instruction::CacheIdent; #[cfg(feature = "enterprise")] use common_meta::key::TableMetadataManagerRef; use common_meta::key::schema_name::{SchemaName, SchemaNameKey}; -use common_meta::procedure_executor::ExecutorContext; #[cfg(feature = "enterprise")] use common_meta::procedure_executor::ProcedureExecutorRef; #[cfg(feature = "enterprise")] @@ -119,10 +118,7 @@ use crate::error::{ use crate::expr_helper::{self, RepartitionRequest, RepartitionSource}; use crate::statement::StatementExecutor; use crate::statement::show::create_partitions_stmt; -use crate::utils::{ - to_meta_query_context_with_origin_frontend_and_trigger_reason, - to_meta_query_context_with_trigger_reason, -}; +use crate::utils::{to_executor_context, to_executor_context_with_origin_frontend}; #[derive(Debug, Clone, Copy)] struct DdlSubmitOptions { @@ -729,13 +725,11 @@ impl StatementExecutor { }) .context(error::InvalidExprSnafu)?; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_create_trigger(task), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_trigger(task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -793,13 +787,11 @@ impl StatementExecutor { create_flow: Some(expr), }) .context(error::InvalidExprSnafu)?; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_create_flow(task), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_flow(task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1109,14 +1101,12 @@ impl StatementExecutor { table_type: TableType::View, }; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(ctx, TriggerReason::Manual), - DdlTask::new_create_view(expr, view_info.clone()), - ); + let executor_context = to_executor_context(ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_view(expr, view_info.clone())); let resp = self .procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu)?; @@ -1195,13 +1185,11 @@ impl StatementExecutor { expr: DropFlowTask, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_flow(expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_flow(expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1230,13 +1218,11 @@ impl StatementExecutor { expr: DropTriggerTask, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_trigger(expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_trigger(expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1311,13 +1297,11 @@ impl StatementExecutor { expr: DropViewTask, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_view(expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_view(expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1780,17 +1764,13 @@ impl StatementExecutor { source: Some(source), ..Default::default() }; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context.clone(), TriggerReason::Manual); - let mut req = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_alter_table(AlterTableExpr { - catalog_name: request.catalog_name.clone(), - schema_name: request.schema_name.clone(), - table_name: request.table_name.clone(), - kind: Some(Kind::Repartition(repartition)), - }), - ); + let executor_context = to_executor_context(query_context.clone(), TriggerReason::Manual); + let mut req = SubmitDdlTaskRequest::new(DdlTask::new_alter_table(AlterTableExpr { + catalog_name: request.catalog_name.clone(), + schema_name: request.schema_name.clone(), + table_name: request.table_name.clone(), + kind: Some(Kind::Repartition(repartition)), + })); req.wait = ddl_options.wait; req.timeout = ddl_options.timeout; @@ -1806,7 +1786,7 @@ impl StatementExecutor { let response = self .procedure_executor - .submit_ddl_task(&ExecutorContext::default(), req) + .submit_ddl_task(executor_context, req) .await .context(error::ExecuteDdlSnafu)?; @@ -1888,12 +1868,11 @@ impl StatementExecutor { .await .context(TableMetadataManagerSnafu)?; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, trigger_reason); + let executor_context = to_executor_context(query_context, trigger_reason); let (req, invalidate_keys) = if physical_table_id == table_id { // This is physical table - let req = SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_alter_table(expr)); + let req = SubmitDdlTaskRequest::new(DdlTask::new_alter_table(expr)); let invalidate_keys = vec![ CacheIdent::TableId(table_id), @@ -1903,10 +1882,7 @@ impl StatementExecutor { (req, invalidate_keys) } else { // This is logical table - let req = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_alter_logical_tables(vec![expr]), - ); + let req = SubmitDdlTaskRequest::new(DdlTask::new_alter_logical_tables(vec![expr])); let mut invalidate_keys = vec![ CacheIdent::TableId(physical_table_id), @@ -1934,7 +1910,7 @@ impl StatementExecutor { }; self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), req) + .submit_ddl_task(executor_context, req) .await .context(error::ExecuteDdlSnafu)?; @@ -2025,17 +2001,19 @@ impl StatementExecutor { .map(|expr| expr.as_pb_partition().context(PartitionExprToPbSnafu)) .collect::>>()?; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_origin_frontend_and_trigger_reason( - query_context, - &self.origin_frontend_addr, - trigger_reason, - ), - DdlTask::new_create_table(create_table, partitions, table_info), + let executor_context = to_executor_context_with_origin_frontend( + query_context, + &self.origin_frontend_addr, + trigger_reason, ); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_table( + create_table, + partitions, + table_info, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2046,17 +2024,15 @@ impl StatementExecutor { query_context: QueryContextRef, trigger_reason: TriggerReason, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_origin_frontend_and_trigger_reason( - query_context, - &self.origin_frontend_addr, - trigger_reason, - ), - DdlTask::new_create_logical_tables(tables_data), + let executor_context = to_executor_context_with_origin_frontend( + query_context, + &self.origin_frontend_addr, + trigger_reason, ); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_logical_tables(tables_data)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2067,15 +2043,11 @@ impl StatementExecutor { query_context: QueryContextRef, trigger_reason: TriggerReason, ) -> Result { - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, trigger_reason); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_alter_logical_tables(tables_data), - ); + let executor_context = to_executor_context(query_context, trigger_reason); + let request = SubmitDdlTaskRequest::new(DdlTask::new_alter_logical_tables(tables_data)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2087,19 +2059,17 @@ impl StatementExecutor { drop_if_exists: bool, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_table( - table_name.catalog_name.clone(), - table_name.schema_name.clone(), - table_name.table_name.clone(), - table_id, - drop_if_exists, - ), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_table( + table_name.catalog_name.clone(), + table_name.schema_name.clone(), + table_name.table_name.clone(), + table_id, + drop_if_exists, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2111,13 +2081,12 @@ impl StatementExecutor { drop_if_exists: bool, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_database(catalog, schema, drop_if_exists), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = + SubmitDdlTaskRequest::new(DdlTask::new_drop_database(catalog, schema, drop_if_exists)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2127,13 +2096,11 @@ impl StatementExecutor { alter_expr: AlterDatabaseExpr, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_alter_database(alter_expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_alter_database(alter_expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2145,21 +2112,17 @@ impl StatementExecutor { time_ranges: Vec<(Timestamp, Timestamp)>, query_context: QueryContextRef, ) -> Result { - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_truncate_table( - table_name.catalog_name.clone(), - table_name.schema_name.clone(), - table_name.table_name.clone(), - table_id, - time_ranges, - ), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_truncate_table( + table_name.catalog_name.clone(), + table_name.schema_name.clone(), + table_name.table_name.clone(), + table_id, + time_ranges, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2239,13 +2202,17 @@ impl StatementExecutor { query_context: QueryContextRef, creator: Option, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_create_database(catalog, database, create_if_not_exists, options, creator), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_database( + catalog, + database, + create_if_not_exists, + options, + creator, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2876,14 +2843,10 @@ async fn execute_undrop_table( table_name: table_name.to_string(), })?; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_undrop_table(dropped.table_id), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_undrop_table(dropped.table_id)); procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu)?; @@ -2931,9 +2894,7 @@ mod test { ExecutorContext, ProcedureExecutor, ProcedureExecutorRef, }; #[cfg(feature = "enterprise")] - use common_meta::rpc::ddl::{ - DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TRIGGER_REASON_EXTENSION_KEY, - }; + use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; #[cfg(feature = "enterprise")] use common_meta::rpc::procedure::{ MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, @@ -2951,6 +2912,7 @@ mod test { #[derive(Default)] struct RecordingProcedureExecutor { requests: Mutex>, + contexts: Mutex>, fail: bool, } @@ -2959,9 +2921,10 @@ mod test { impl ProcedureExecutor for RecordingProcedureExecutor { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { + self.contexts.lock().unwrap().push(ctx); self.requests.lock().unwrap().push(request); if self.fail { return common_meta::error::UnsupportedSnafu { @@ -3078,12 +3041,10 @@ mod test { let requests = procedure.requests.lock().unwrap(); assert!(matches!(&requests[0].task, DdlTask::UndropTable(task) if task.table_id == 42)); + let contexts = procedure.contexts.lock().unwrap(); assert_eq!( - requests[0] - .query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY), - Some(&TriggerReason::Manual.as_ref().to_string()) + contexts[0].event_input.as_ref().map(|input| input.reason), + Some(TriggerReason::Manual) ); assert_eq!( cache.invalidations.lock().unwrap()[0], diff --git a/src/operator/src/utils.rs b/src/operator/src/utils.rs index 03649713b0b..160d7120726 100644 --- a/src/operator/src/utils.rs +++ b/src/operator/src/utils.rs @@ -14,9 +14,9 @@ use std::sync::{Arc, RwLock}; -use common_meta::rpc::ddl::{ - ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TRIGGER_REASON_EXTENSION_KEY, TriggerReason, -}; +use common_event_recorder::ProcedureEventInput; +use common_meta::procedure_executor::ExecutorContext; +use common_meta::rpc::ddl::{ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TriggerReason}; use common_time::Timezone; use session::context::{QueryContextBuilder, QueryContextRef}; use snafu::ResultExt; @@ -37,31 +37,34 @@ pub fn to_meta_query_context( } } -/// Converts a session query context and records the DDL trigger reason. -pub fn to_meta_query_context_with_trigger_reason( +/// Builds the execution envelope for a frontend-originated procedure. +pub(crate) fn to_executor_context( query_context: QueryContextRef, trigger_reason: TriggerReason, -) -> common_meta::rpc::ddl::QueryContext { - let mut meta_query_context = to_meta_query_context(query_context); - meta_query_context.extensions.insert( - TRIGGER_REASON_EXTENSION_KEY.to_string(), - trigger_reason.as_ref().to_string(), - ); - meta_query_context +) -> ExecutorContext { + let actor = query_context.current_user().username().to_string(); + ExecutorContext { + query_context: Some(to_meta_query_context(query_context)), + actor: Some(actor), + event_input: Some(ProcedureEventInput::new(trigger_reason)), + ..Default::default() + } } -pub fn to_meta_query_context_with_origin_frontend_and_trigger_reason( +/// Builds a frontend-originated procedure envelope with the trusted frontend origin. +pub(crate) fn to_executor_context_with_origin_frontend( query_context: QueryContextRef, origin_frontend_addr: &str, trigger_reason: TriggerReason, -) -> common_meta::rpc::ddl::QueryContext { - let mut meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, trigger_reason); - meta_query_context.extensions.insert( - ORIGIN_FRONTEND_ADDR_EXTENSION_KEY.to_string(), - origin_frontend_addr.to_string(), - ); - meta_query_context +) -> ExecutorContext { + let mut executor_context = to_executor_context(query_context, trigger_reason); + if let Some(query_context) = &mut executor_context.query_context { + query_context.extensions.insert( + ORIGIN_FRONTEND_ADDR_EXTENSION_KEY.to_string(), + origin_frontend_addr.to_string(), + ); + } + executor_context } pub fn try_to_session_query_context( @@ -87,14 +90,12 @@ mod tests { use std::collections::HashMap; use std::sync::{Arc, RwLock}; - use common_meta::rpc::ddl::{ - ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TRIGGER_REASON_EXTENSION_KEY, TriggerReason, - }; + use common_meta::rpc::ddl::{ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TriggerReason}; use common_time::Timezone; use session::context::QueryContextBuilder; use super::{ - to_meta_query_context, to_meta_query_context_with_origin_frontend_and_trigger_reason, + to_executor_context_with_origin_frontend, to_meta_query_context, try_to_session_query_context, }; @@ -122,7 +123,7 @@ mod tests { } #[test] - fn test_meta_query_context_with_origin_frontend_overrides_reserved_key() { + fn test_executor_context_with_origin_frontend_overrides_reserved_key() { let session_ctx = Arc::new( QueryContextBuilder::default() .set_extension( @@ -132,11 +133,12 @@ mod tests { .build(), ); - let meta_ctx = to_meta_query_context_with_origin_frontend_and_trigger_reason( + let executor_context = to_executor_context_with_origin_frontend( session_ctx, "127.0.0.1:4000", TriggerReason::Manual, ); + let meta_ctx = executor_context.query_context.as_ref().unwrap(); assert_eq!( meta_ctx @@ -146,11 +148,8 @@ mod tests { Some("127.0.0.1:4000") ); assert_eq!( - meta_ctx - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) - .map(String::as_str), - Some("manual") + executor_context.event_input.unwrap().reason, + TriggerReason::Manual ); } } diff --git a/tests-integration/tests/gc_event.rs b/tests-integration/tests/gc_event.rs index 9ecedc77775..54a0787d12a 100644 --- a/tests-integration/tests/gc_event.rs +++ b/tests-integration/tests/gc_event.rs @@ -15,8 +15,8 @@ use std::sync::Arc; use std::time::Duration; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; -use common_procedure::{ProcedureId, ProcedureWithId, watcher}; +use common_event_recorder::{PersistentEventContext, TriggerReason}; +use common_procedure::{ProcedureContext, ProcedureId, ProcedureWithId, watcher}; use common_test_util::temp_dir::create_temp_dir; use meta_srv::gc::{BatchGcProcedure, GcSchedulerOptions}; use mito2::gc::GcConfig; @@ -137,9 +137,9 @@ async fn test_batch_gc_event() { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), - context: common_procedure::ProcedureContext::from_event_context( - PersistentEventContext::new(TriggerReason::Manual), - ), + context: ProcedureContext::from_event_context(PersistentEventContext::new( + TriggerReason::Manual, + )), }) .await .unwrap(); diff --git a/tests-integration/tests/region_migration.rs b/tests-integration/tests/region_migration.rs index 3d27dabc3f7..d67180249c9 100644 --- a/tests-integration/tests/region_migration.rs +++ b/tests-integration/tests/region_migration.rs @@ -23,10 +23,11 @@ use common_event_recorder::event_table::{ }; use common_event_recorder::{ DEFAULT_EVENTS_TABLE_NAME, DEFAULT_FLUSH_INTERVAL_SECONDS, EVENTS_TABLE_TIMESTAMP_COLUMN_NAME, - EVENTS_TABLE_TYPE_COLUMN_NAME, + EVENTS_TABLE_TYPE_COLUMN_NAME, PersistentEventContext, TriggerReason, }; use common_meta::key::{RegionDistribution, RegionRoleSet, TableMetadataManagerRef}; use common_meta::peer::Peer; +use common_procedure::ProcedureContext; use common_procedure::event::{ EVENTS_TABLE_PROCEDURE_ID_COLUMN_NAME, EVENTS_TABLE_PROCEDURE_STATE_COLUMN_NAME, }; @@ -183,13 +184,18 @@ pub async fn test_region_migration(store_type: StorageType, endpoints: Vec Date: Thu, 13 Aug 2026 07:52:44 +0000 Subject: [PATCH 062/126] perf(mito2): optimize dictionary primary key sorting (#8767) Signed-off-by: lyang24 --- src/mito2/benches/memtable_bench.rs | 117 +++++++++++- src/mito2/src/memtable/bulk/part.rs | 285 +++++++++++++++++++++++++++- 2 files changed, 394 insertions(+), 8 deletions(-) diff --git a/src/mito2/benches/memtable_bench.rs b/src/mito2/benches/memtable_bench.rs index c9151eea4ca..80331941865 100644 --- a/src/mito2/benches/memtable_bench.rs +++ b/src/mito2/benches/memtable_bench.rs @@ -22,10 +22,17 @@ use std::sync::Arc; -use criterion::{Criterion, criterion_group, criterion_main}; +use common_recordbatch::DfRecordBatch; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use datatypes::arrow::array::{ + ArrayRef, BinaryDictionaryBuilder, Int64Array, TimestampMillisecondArray, UInt8Array, + UInt64Array, +}; +use datatypes::arrow::compute::{SortColumn, SortOptions}; +use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type}; use mito_codec::row_converter::DensePrimaryKeyCodec; use mito2::memtable::bulk::context::BulkIterContext; -use mito2::memtable::bulk::part::BulkPartConverter; +use mito2::memtable::bulk::part::{BulkPartConverter, sort_primary_key_record_batch}; use mito2::memtable::bulk::part_reader::BulkPartBatchIter; use mito2::memtable::bulk::{BulkMemtable, BulkMemtableConfig}; use mito2::memtable::time_series::TimeSeriesMemtable; @@ -39,6 +46,111 @@ use mito2::test_util::memtable_util; const DEFAULT_BATCH_SIZE: usize = 8 * 1024; +fn primary_key_sort_batch( + series_count: usize, + samples_per_series: usize, + sorted_within_series: bool, +) -> DfRecordBatch { + let num_rows = series_count * samples_per_series; + let primary_keys = (0..series_count) + .map(|series| format!("series_{series:08}").into_bytes()) + .collect::>(); + let mut primary_key_builder = BinaryDictionaryBuilder::::new(); + let mut timestamps = Vec::with_capacity(num_rows); + let mut sequences = Vec::with_capacity(num_rows); + + for series_order in 0..series_count { + // 17 is coprime to the power-of-two cardinalities used below and gives a deterministic + // non-lexical series order. + let series = series_order * 17 % series_count; + for sample_order in 0..samples_per_series { + let sample = if sorted_within_series { + sample_order + } else { + sample_order * 37 % samples_per_series + }; + primary_key_builder + .append(primary_keys[series].as_slice()) + .unwrap(); + timestamps.push(sample as i64); + sequences.push((series * samples_per_series + sample_order) as u64); + } + } + + let columns: Vec = vec![ + Arc::new(Int64Array::from_iter_values(0..num_rows as i64)), + Arc::new(TimestampMillisecondArray::from(timestamps)), + Arc::new(primary_key_builder.finish()), + Arc::new(UInt64Array::from(sequences)), + Arc::new(UInt8Array::from_value(1, num_rows)), + ]; + let schema = Arc::new(Schema::new(vec![ + Field::new("value", DataType::Int64, false), + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new( + "__primary_key", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Binary)), + false, + ), + Field::new("__sequence", DataType::UInt64, false), + Field::new("__op_type", DataType::UInt8, false), + ])); + DfRecordBatch::try_new(schema, columns).unwrap() +} + +fn lexsort_primary_key_record_batch(batch: &DfRecordBatch) -> DfRecordBatch { + let total_columns = batch.num_columns(); + let sort_columns = vec![ + SortColumn { + values: batch.column(total_columns - 3).clone(), + options: Some(SortOptions { + descending: false, + nulls_first: true, + }), + }, + SortColumn { + values: batch.column(total_columns - 4).clone(), + options: Some(SortOptions { + descending: false, + nulls_first: true, + }), + }, + SortColumn { + values: batch.column(total_columns - 2).clone(), + options: Some(SortOptions { + descending: true, + nulls_first: true, + }), + }, + ]; + let indices = datatypes::arrow::compute::lexsort_to_indices(&sort_columns, None).unwrap(); + datatypes::arrow::compute::take_record_batch(batch, &indices).unwrap() +} + +fn primary_key_sort(c: &mut Criterion) { + let mut group = c.benchmark_group("primary_key_sort"); + group.sample_size(20); + + for (name, series_count, samples_per_series, sorted_within_series) in [ + ("sorted_runs", 512, 16, true), + ("shuffled_runs", 512, 16, false), + ("high_cardinality", 8192, 1, true), + ] { + let batch = primary_key_sort_batch(series_count, samples_per_series, sorted_within_series); + group.throughput(Throughput::Elements(batch.num_rows() as u64)); + group.bench_function(format!("{name}/arrow_lexsort"), |b| { + b.iter(|| lexsort_primary_key_record_batch(&batch)); + }); + group.bench_function(format!("{name}/rank_scatter"), |b| { + b.iter(|| sort_primary_key_record_batch(&batch).unwrap()); + }); + } +} + /// Writes rows. fn write_rows(c: &mut Criterion) { let metadata = Arc::new(memtable_util::metadata_with_primary_key(vec![1, 0], true)); @@ -391,6 +503,7 @@ fn bulk_part_record_batch_iter_filter(c: &mut Criterion) { criterion_group!( benches, + primary_key_sort, write_rows, full_scan, filter_1_host, diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 20482794079..cc2c37a7853 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -30,11 +30,13 @@ use datafusion_common::pruning::PruningStatistics; use datafusion_expr::utils::expr_to_columns; use datatypes::arrow; use datatypes::arrow::array::{ - Array, ArrayRef, BinaryArray, BooleanArray, StringDictionaryBuilder, UInt8Array, UInt64Array, + Array, ArrayRef, BinaryArray, BooleanArray, DictionaryArray, StringDictionaryBuilder, + TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, + TimestampSecondArray, UInt8Array, UInt32Array, UInt64Array, }; use datatypes::arrow::compute::{SortColumn, SortOptions, concat_batches}; use datatypes::arrow::datatypes::{ - DataType as ArrowDataType, Field, Schema, SchemaRef, UInt32Type, + DataType as ArrowDataType, Field, Schema, SchemaRef, TimeUnit, UInt32Type, }; use datatypes::data_type::DataType; use datatypes::extension::json::is_json2_extension_type; @@ -771,6 +773,147 @@ fn new_primary_key_column_builders( /// Sorts the record batch with primary key format. pub fn sort_primary_key_record_batch(batch: &RecordBatch) -> Result { + let indices = if let Some(indices) = try_sort_primary_key_indices(batch)? { + indices + } else { + lexsort_primary_key_indices(batch)? + }; + + datatypes::arrow::compute::take_record_batch(batch, &indices).context(ComputeArrowSnafu) +} + +/// Sorts the known, non-null primary-key layout without comparing encoded keys for every row. +/// +/// Returns `None` if the batch does not have the exact layout supported by this fast path. +fn try_sort_primary_key_indices(batch: &RecordBatch) -> Result> { + let total_columns = batch.num_columns(); + let timestamp = batch.column(total_columns - 4); + let primary_key = batch.column(total_columns - 3); + let sequence = batch.column(total_columns - 2); + + let Some(primary_key) = primary_key + .as_any() + .downcast_ref::>() + else { + return Ok(None); + }; + let Some(sequence) = sequence.as_any().downcast_ref::() else { + return Ok(None); + }; + + if batch.num_rows() > u32::MAX as usize + || primary_key.null_count() != 0 + || primary_key.values().data_type() != &ArrowDataType::Binary + || primary_key.values().null_count() != 0 + || primary_key.values().len() > batch.num_rows() + || timestamp.null_count() != 0 + || sequence.null_count() != 0 + { + return Ok(None); + } + + let indices = match timestamp.data_type() { + ArrowDataType::Timestamp(TimeUnit::Second, _) => timestamp + .as_any() + .downcast_ref::() + .map(|timestamp| sort_primary_key_indices(primary_key, timestamp.values(), sequence)), + ArrowDataType::Timestamp(TimeUnit::Millisecond, _) => timestamp + .as_any() + .downcast_ref::() + .map(|timestamp| sort_primary_key_indices(primary_key, timestamp.values(), sequence)), + ArrowDataType::Timestamp(TimeUnit::Microsecond, _) => timestamp + .as_any() + .downcast_ref::() + .map(|timestamp| sort_primary_key_indices(primary_key, timestamp.values(), sequence)), + ArrowDataType::Timestamp(TimeUnit::Nanosecond, _) => timestamp + .as_any() + .downcast_ref::() + .map(|timestamp| sort_primary_key_indices(primary_key, timestamp.values(), sequence)), + _ => None, + }; + + indices.transpose() +} + +/// Computes sorted row indices by ranking dictionary values, scattering rows into primary-key +/// buckets, and sorting only buckets that are not already ordered by time and sequence. +fn sort_primary_key_indices( + primary_key: &PrimaryKeyArray, + timestamps: &[i64], + sequences: &UInt64Array, +) -> Result { + debug_assert_eq!(primary_key.len(), timestamps.len()); + debug_assert_eq!(primary_key.len(), sequences.len()); + debug_assert_eq!(primary_key.null_count(), 0); + debug_assert_eq!(primary_key.values().null_count(), 0); + debug_assert_eq!(sequences.null_count(), 0); + + // Ranks compare in the same order as the dictionary values. Equal dictionary values receive + // the same rank, which is required because Arrow dictionary concatenation only deduplicates + // values on a best-effort basis. + let value_ranks = datatypes::arrow::compute::rank( + primary_key.values().as_ref(), + Some(SortOptions { + descending: false, + nulls_first: true, + }), + ) + .context(ComputeArrowSnafu)?; + + // A rank is at most the number of dictionary values. It is not necessarily dense when the + // dictionary contains duplicate values, so keep one extra slot and allow empty buckets. + let mut counts = vec![0usize; value_ranks.len() + 1]; + for &key in primary_key.keys().values() { + counts[value_ranks[key as usize] as usize] += 1; + } + + let mut bucket_offsets = vec![0usize; counts.len()]; + let mut offset = 0; + for (bucket_offset, &count) in bucket_offsets.iter_mut().zip(&counts) { + *bucket_offset = offset; + offset += count; + } + + // Scatter in input order. This preserves already sorted runs within each primary key. + let mut next_offsets = bucket_offsets.clone(); + let mut indices = vec![0u32; primary_key.len()]; + for (row, &key) in primary_key.keys().values().iter().enumerate() { + let rank = value_ranks[key as usize] as usize; + indices[next_offsets[rank]] = row as u32; + next_offsets[rank] += 1; + } + + let sequences = sequences.values(); + for (&start, count) in bucket_offsets.iter().zip(counts) { + let end = start + count; + let bucket = &mut indices[start..end]; + if !bucket.is_sorted_by(|left, right| { + compare_time_sequence(timestamps, sequences, *left, *right).is_le() + }) { + bucket.sort_unstable_by(|left, right| { + compare_time_sequence(timestamps, sequences, *left, *right) + }); + } + } + + Ok(UInt32Array::from(indices)) +} + +#[inline] +fn compare_time_sequence( + timestamps: &[i64], + sequences: &[u64], + left: u32, + right: u32, +) -> std::cmp::Ordering { + let left = left as usize; + let right = right as usize; + timestamps[left] + .cmp(×tamps[right]) + .then_with(|| sequences[right].cmp(&sequences[left])) +} + +fn lexsort_primary_key_indices(batch: &RecordBatch) -> Result { let total_columns = batch.num_columns(); let sort_columns = vec![ // Primary key column (ascending) @@ -799,10 +942,7 @@ pub fn sort_primary_key_record_batch(batch: &RecordBatch) -> Result }, ]; - let indices = datatypes::arrow::compute::lexsort_to_indices(&sort_columns, None) - .context(ComputeArrowSnafu)?; - - datatypes::arrow::compute::take_record_batch(batch, &indices).context(ComputeArrowSnafu) + datatypes::arrow::compute::lexsort_to_indices(&sort_columns, None).context(ComputeArrowSnafu) } /// Converts a `BulkPart` that is unordered and without encoded primary keys into a `BulkPart` @@ -1674,6 +1814,8 @@ mod tests { use datatypes::prelude::{ConcreteDataType, Value}; use datatypes::schema::ColumnSchema; use mito_codec::row_converter::build_primary_key_codec; + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::storage::RegionId; use store_api::storage::consts::ReservedColumnId; @@ -1692,6 +1834,137 @@ mod tests { sequence: u64, } + #[test] + fn test_sort_primary_key_record_batch_with_duplicate_dictionary_values() { + // Dictionary keys 1 and 2 both represent "series_a". This can happen after Arrow + // concatenates dictionaries because dictionary deduplication is best-effort. + let primary_key = DictionaryArray::try_new( + UInt32Array::from(vec![0, 1, 0, 2, 2]), + Arc::new(BinaryArray::from_vec(vec![ + b"series_b".as_slice(), + b"series_a".as_slice(), + b"series_a".as_slice(), + ])), + ) + .unwrap(); + let timestamps = TimestampMillisecondArray::from(vec![20, 30, 10, 30, 10]); + let sequences = UInt64Array::from(vec![5, 6, 4, 8, 9]); + let row_ids = UInt32Array::from_iter_values(0..5); + let op_types = UInt8Array::from_value(OpType::Put as u8, 5); + let schema = Arc::new(Schema::new(vec![ + Field::new("row_id", ArrowDataType::UInt32, false), + Field::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new( + PRIMARY_KEY_COLUMN_NAME, + ArrowDataType::Dictionary( + Box::new(ArrowDataType::UInt32), + Box::new(ArrowDataType::Binary), + ), + false, + ), + Field::new("__sequence", ArrowDataType::UInt64, false), + Field::new("__op_type", ArrowDataType::UInt8, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(row_ids), + Arc::new(timestamps), + Arc::new(primary_key), + Arc::new(sequences), + Arc::new(op_types), + ], + ) + .unwrap(); + + let expected_indices = lexsort_primary_key_indices(&batch).unwrap(); + let expected = + datatypes::arrow::compute::take_record_batch(&batch, &expected_indices).unwrap(); + let actual = sort_primary_key_record_batch(&batch).unwrap(); + + let expected_row_ids = expected + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let actual_row_ids = actual + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(expected_row_ids, actual_row_ids); + assert_eq!(actual_row_ids.values(), &[4, 3, 1, 2, 0]); + } + + #[test] + fn test_sort_primary_key_record_batch_against_lexsort() { + let mut rng = StdRng::seed_from_u64(0x5eed); + for _ in 0..100 { + let num_rows = rng.random_range(0..128); + let keys = UInt32Array::from_iter_values((0..num_rows).map(|_| rng.random_range(0..8))); + let primary_key = DictionaryArray::try_new( + keys, + Arc::new(BinaryArray::from_vec(vec![ + b"d".as_slice(), + b"a".as_slice(), + b"c".as_slice(), + b"b".as_slice(), + b"a".as_slice(), + b"d".as_slice(), + b"b".as_slice(), + b"c".as_slice(), + ])), + ) + .unwrap(); + let timestamps = TimestampMillisecondArray::from_iter_values( + (0..num_rows).map(|_| rng.random_range(-1000..1000)), + ); + // Unique sequences ensure that the oracle and fast path have no fully equal sort keys. + let sequences = UInt64Array::from_iter_values(0..num_rows as u64); + let row_ids = UInt32Array::from_iter_values(0..num_rows as u32); + let op_types = UInt8Array::from_value(OpType::Put as u8, num_rows); + let schema = Arc::new(Schema::new(vec![ + Field::new("row_id", ArrowDataType::UInt32, false), + Field::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new( + PRIMARY_KEY_COLUMN_NAME, + ArrowDataType::Dictionary( + Box::new(ArrowDataType::UInt32), + Box::new(ArrowDataType::Binary), + ), + false, + ), + Field::new("__sequence", ArrowDataType::UInt64, false), + Field::new("__op_type", ArrowDataType::UInt8, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(row_ids), + Arc::new(timestamps), + Arc::new(primary_key), + Arc::new(sequences), + Arc::new(op_types), + ], + ) + .unwrap(); + + let expected_indices = lexsort_primary_key_indices(&batch).unwrap(); + let expected = + datatypes::arrow::compute::take_record_batch(&batch, &expected_indices).unwrap(); + let actual = sort_primary_key_record_batch(&batch).unwrap(); + assert_eq!(expected.column(0), actual.column(0)); + } + } + #[test] fn test_unordered_part_tracks_estimated_bytes() { let mut part = UnorderedPart::new(); From 764c93bf43bd82da6eae8a5fa671fc56c39cc156 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 13 Aug 2026 10:51:53 +0000 Subject: [PATCH 063/126] perf(query): choose bounded CTE as hash join build side (#8807) * fix(query): choose bounded CTE as hash join build side Signed-off-by: jeremyhi * perf(query): remove join estimate cap * test(compat): accept repartition in analyze plan Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- Cargo.lock | 68 +++--- Cargo.toml | 30 +-- src/query/src/dist_plan/merge_scan.rs | 219 +++++++++++++++++- src/query/src/optimizer/pass_distribution.rs | 4 +- src/standalone/src/options.rs | 10 + .../distributed/explain/step_aggr.result | 8 + .../explain/step_aggr_advance.result | 111 ++++++--- .../explain/step_aggr_basic.result | 10 + .../distributed/explain/subqueries.result | 13 +- .../distributed/flow-tql/tsid_on_phy.result | 4 +- .../cases/distributed/optimizer/count.result | 5 + .../optimizer/first_value_advance.result | 8 + .../optimizer/last_value_advance.result | 8 + .../optimizer/range_select_projection.result | 2 +- .../common/aggregate/distinct.result | 10 +- .../common/aggregate/multi_regions.result | 2 + .../create/metric_engine_partition.result | 11 +- .../common/cte/cte_join_build_side.result | 116 ++++++++++ .../common/cte/cte_join_build_side.sql | 46 ++++ .../filter/hash_join_topk_dyn_filter.result | 12 +- .../common/order/order_by_exceptions.result | 7 +- .../promql/histogram_multi_partition.result | 1 + .../promql/tsid_binary_join_regression.result | 20 +- .../tsid_histogram_quantile_regression.result | 2 +- .../cases/standalone/common/range/nest.result | 4 +- .../standalone/common/tql/partition.result | 5 +- .../standalone/common/tql/tql-cte.result | 24 +- tests/cases/standalone/optimizer/count.result | 5 + .../optimizer/first_value_advance.result | 8 + .../optimizer/last_value_advance.result | 8 + .../optimizer/rewrite_set_comparison.result | 22 +- .../tql-explain-analyze/analyze.result | 1 + .../tql-explain-analyze/tsid_column.result | 8 +- .../verify.result | 1 + tests/conf/datanode-test.toml.template | 3 + tests/conf/frontend-test.toml.template | 3 + tests/conf/standalone-test.toml.template | 3 + 37 files changed, 674 insertions(+), 148 deletions(-) create mode 100644 tests/cases/standalone/common/cte/cte_join_build_side.result create mode 100644 tests/cases/standalone/common/cte/cte_join_build_side.sql diff --git a/Cargo.lock b/Cargo.lock index 97b2fc76a79..a7a7ff9b877 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3733,7 +3733,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "arrow-schema 58.3.0", @@ -3787,7 +3787,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3811,7 +3811,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3833,7 +3833,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -3857,7 +3857,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "futures", "log", @@ -3867,7 +3867,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-compression", @@ -3901,7 +3901,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "arrow-ipc 58.3.0", @@ -3924,7 +3924,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3946,7 +3946,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3969,7 +3969,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3998,12 +3998,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" [[package]] name = "datafusion-execution" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4025,7 +4025,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4047,7 +4047,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4059,7 +4059,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4090,7 +4090,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4111,7 +4111,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4123,7 +4123,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "arrow-ord 58.3.0", @@ -4147,7 +4147,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4162,7 +4162,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4179,7 +4179,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4188,7 +4188,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "datafusion-doc", "quote", @@ -4198,7 +4198,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "chrono", @@ -4247,7 +4247,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4270,7 +4270,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4284,7 +4284,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4300,7 +4300,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4318,7 +4318,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4349,7 +4349,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "chrono", @@ -4376,7 +4376,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4386,7 +4386,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4402,7 +4402,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "async-trait", "datafusion-common", @@ -4415,7 +4415,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "bigdecimal 0.4.8", @@ -4433,7 +4433,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 63dbc76f383..bfd6153bf5b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -346,21 +346,21 @@ git = "https://github.com/GreptimeTeam/greptime-meter.git" rev = "5618e779cf2bb4755b499c630fba4c35e91898cb" [patch.crates-io] -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "2aefa08a8d69c96eec2d6d6703598a009bba6e4c" } # on branch v0.61.x [profile.release] diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index aafc17c469f..8d3580ea452 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -42,8 +42,9 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, SendableRecordBatchStream, }; -use datafusion_common::{Column as ColumnExpr, DataFusionError, Result}; -use datafusion_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion_common::stats::Precision; +use datafusion_common::{Column as ColumnExpr, DataFusionError, Result, Statistics}; +use datafusion_expr::{Expr, Extension, FetchType, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::{Distribution, EquivalenceProperties, PhysicalSortExpr}; use datatypes::extension::json::{ @@ -84,6 +85,51 @@ fn query_engine_state_from_task_context(context: &TaskContext) -> Option Option { + match plan { + LogicalPlan::Limit(limit) => { + let input_bound = remote_plan_row_bound(&limit.input); + match limit.get_fetch_type() { + Ok(FetchType::Literal(Some(fetch))) => { + Some(input_bound.map_or(fetch, |bound| bound.min(fetch))) + } + _ => input_bound, + } + } + LogicalPlan::Sort(sort) => { + let input_bound = remote_plan_row_bound(&sort.input); + sort.fetch + .map(|fetch| input_bound.map_or(fetch, |bound| bound.min(fetch))) + .or(input_bound) + } + LogicalPlan::Projection(projection) => remote_plan_row_bound(&projection.input), + LogicalPlan::Filter(filter) => remote_plan_row_bound(&filter.input), + LogicalPlan::SubqueryAlias(alias) => remote_plan_row_bound(&alias.input), + LogicalPlan::Window(window) => remote_plan_row_bound(&window.input), + LogicalPlan::Repartition(repartition) => remote_plan_row_bound(&repartition.input), + LogicalPlan::Distinct(distinct) => remote_plan_row_bound(distinct.input()), + LogicalPlan::Aggregate(aggregate) => { + if aggregate + .group_expr + .iter() + .any(|expr| matches!(expr, Expr::GroupingSet(_))) + { + None + } else if aggregate.group_expr.is_empty() { + Some(1) + } else { + remote_plan_row_bound(&aggregate.input) + } + } + _ => None, + } +} + fn remote_dyn_filter_enabled(query_ctx: &QueryContextRef) -> Result { remote_dyn_filter_pushdown_enabled_from_extensions(&query_ctx.extensions()) .map_err(|err| DataFusionError::External(Box::new(err))) @@ -429,6 +475,7 @@ impl MergeScanExec { enable_per_region_metrics: bool, ) -> Result { let arrow_schema = maybe_amend_json2_field(arrow_schema); + let output_partition_count = Self::output_partition_count(regions.len(), target_partition); // States the output ordering of the plan. // @@ -438,7 +485,7 @@ impl MergeScanExec { // // Otherwise, we need to use the default ordering. let eq_properties = if let LogicalPlan::Sort(sort) = &plan - && target_partition >= regions.len() + && output_partition_count >= regions.len() { let lex_ordering = sort .expr @@ -477,7 +524,7 @@ impl MergeScanExec { } }) .collect(); - let partitioning = Partitioning::Hash(partition_exprs, target_partition); + let partitioning = Partitioning::Hash(partition_exprs, output_partition_count); let properties = Arc::new(PlanProperties::new( eq_properties, @@ -504,6 +551,28 @@ impl MergeScanExec { }) } + /// Conservative row-count upper bound for all selected regions. + /// + /// This is not an expected cardinality: DataFusion receives it as an + /// inexact estimate only because `Statistics` has no upper-bound precision. + fn estimated_num_rows(&self) -> Precision { + if self.regions.is_empty() { + return Precision::Inexact(0); + } + + let Some(rows_per_region) = remote_plan_row_bound(&self.plan) else { + return Precision::Absent; + }; + rows_per_region + .checked_mul(self.regions.len()) + .map_or(Precision::Absent, Precision::Inexact) + } + + /// Number of partitions populated by the region striping in [`Self::to_stream`]. + fn output_partition_count(num_regions: usize, target_partition: usize) -> usize { + num_regions.max(1).min(target_partition.max(1)) + } + pub fn to_stream( &self, context: Arc, @@ -518,7 +587,8 @@ impl MergeScanExec { let sub_stage_metrics_moved = self.sub_stage_metrics.clone(); let partition_metrics_moved = self.partition_metrics.clone(); let plan = self.plan.clone(); - let target_partition = self.target_partition; + let target_partition = + Self::output_partition_count(self.regions.len(), self.target_partition); let remote_dyn_filter_enabled = remote_dyn_filter_enabled(&self.query_ctx)?; let captured_remote_dyn_filters = if remote_dyn_filter_enabled { self.captured_remote_dyn_filters() @@ -855,7 +925,7 @@ impl MergeScanExec { metric: self.metric.clone(), properties: Arc::new(PlanProperties::new( self.properties.eq_properties.clone(), - Partitioning::Hash(overlaps, self.target_partition), + Partitioning::Hash(overlaps, self.partition_count()), self.properties.emission_type, self.properties.boundedness, )), @@ -906,7 +976,7 @@ impl MergeScanExec { } pub fn partition_count(&self) -> usize { - self.target_partition + Self::output_partition_count(self.regions.len(), self.target_partition) } pub fn region_count(&self) -> usize { @@ -1174,6 +1244,16 @@ impl ExecutionPlan for MergeScanExec { Some(self.metric.clone_inner()) } + fn partition_statistics(&self, partition: Option) -> Result { + if partition.is_some() { + return Ok(Statistics::new_unknown(&self.arrow_schema)); + } + + let mut statistics = Statistics::new_unknown(&self.arrow_schema); + statistics.num_rows = self.estimated_num_rows(); + Ok(statistics) + } + fn name(&self) -> &str { "MergeScanExec" } @@ -1404,7 +1484,6 @@ mod tests { region_count: u64, target_partition: usize, ) -> MergeScanExec { - let session_state = SessionStateBuilder::new().build(); let plan = LogicalPlanBuilder::empty(true) .project(vec![lit(1i64).alias("ts")]) .unwrap() @@ -1412,11 +1491,21 @@ mod tests { .unwrap() .build() .unwrap(); - let schema = plan.schema().as_arrow().clone(); let regions = (0..region_count) .map(|region_number| RegionId::new(1024, region_number as u32)) .collect(); + merge_scan_exec_with_plan(regions, plan, target_partition) + } + + fn merge_scan_exec_with_plan( + regions: Vec, + plan: LogicalPlan, + target_partition: usize, + ) -> MergeScanExec { + let session_state = SessionStateBuilder::new().build(); + let schema = plan.schema().as_arrow().clone(); + MergeScanExec::new( &session_state, // The table name is not relevant to these ordering metadata tests; @@ -1555,6 +1644,118 @@ mod tests { assert!(exec.properties().output_ordering().is_some()); } + #[test] + fn merge_scan_reports_populated_partition_count() { + let cases = [(0, 10, 1), (1, 10, 1), (3, 2, 2), (5, 10, 5), (3, 0, 1)]; + + for (region_count, target, expected) in cases { + let exec = merge_scan_exec_with_sorted_input(region_count, target); + assert_eq!(exec.partition_count(), expected); + assert_eq!( + exec.properties().output_partitioning().partition_count(), + expected + ); + } + } + + #[test] + fn merge_scan_reports_only_deterministic_plan_bounds() { + use datafusion::functions_aggregate::expr_fn::count; + use datafusion_expr::GroupingSet; + + let regions = vec![RegionId::new(1024, 1), RegionId::new(1024, 2)]; + let limited = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .limit(0, Some(50)) + .unwrap() + .build() + .unwrap(); + assert_eq!( + merge_scan_exec_with_plan(regions.clone(), limited, 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Inexact(100) + ); + + let large_bound = i32::MAX as usize + 1; + let large_limit = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .limit(0, Some(large_bound)) + .unwrap() + .build() + .unwrap(); + assert_eq!( + merge_scan_exec_with_plan(vec![RegionId::new(1024, 1)], large_limit, 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Inexact(large_bound) + ); + + let uncapped = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .build() + .unwrap(); + assert_eq!( + merge_scan_exec_with_plan(regions.clone(), uncapped.clone(), 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Absent + ); + assert_eq!( + merge_scan_exec_with_plan(Vec::new(), uncapped, 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Inexact(0) + ); + + let global_aggregate = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .limit(0, Some(0)) + .unwrap() + .aggregate(Vec::::new(), vec![count(lit(1))]) + .unwrap() + .build() + .unwrap(); + assert_eq!( + merge_scan_exec_with_plan(regions.clone(), global_aggregate, 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Inexact(2) + ); + + let grouping_sets = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .limit(0, Some(50)) + .unwrap() + .aggregate( + vec![Expr::GroupingSet(GroupingSet::GroupingSets(vec![ + vec![], + vec![col("col")], + ]))], + Vec::::new(), + ) + .unwrap() + .build() + .unwrap(); + assert_eq!( + merge_scan_exec_with_plan(regions, grouping_sets, 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Absent + ); + } + #[test] fn sub_stage_metrics_are_sorted_by_region_id() { let exec = merge_scan_exec_with_sorted_input(0, 1); diff --git a/src/query/src/optimizer/pass_distribution.rs b/src/query/src/optimizer/pass_distribution.rs index 37274286482..c45d4695d93 100644 --- a/src/query/src/optimizer/pass_distribution.rs +++ b/src/query/src/optimizer/pass_distribution.rs @@ -277,8 +277,8 @@ mod tests { panic!("expected right merge scan hash partitioning"); }; - assert_eq!(*left_count, 32); - assert_eq!(*right_count, 32); + assert_eq!(*left_count, 2); + assert_eq!(*right_count, 2); assert_eq!( column_names(left_exprs), vec![DATA_SCHEMA_TSID_COLUMN_NAME, "greptime_timestamp"] diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 637715013af..bedf9eebc70 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -157,6 +157,7 @@ impl StandaloneOptions { meta_client: None, logging: cloned_opts.logging, user_provider: cloned_opts.user_provider, + query: cloned_opts.query, slow_query: cloned_opts.slow_query, event_recorder: cloned_opts.event_recorder, heartbeat_env_vars: cloned_opts.heartbeat_env_vars.clone(), @@ -227,4 +228,13 @@ mod tests { &frontend_options.event_recorder.event_types, )); } + + #[test] + fn test_query_options_propagated_to_components() { + let mut options = StandaloneOptions::default(); + options.query.parallelism = 4; + + assert_eq!(options.frontend_options().query.parallelism, 4); + assert_eq!(options.datanode_options().query.parallelism, 4); + } } diff --git a/tests/cases/distributed/explain/step_aggr.result b/tests/cases/distributed/explain/step_aggr.result index 78218d514da..89634a3cc6d 100644 --- a/tests/cases/distributed/explain/step_aggr.result +++ b/tests/cases/distributed/explain/step_aggr.result @@ -67,6 +67,7 @@ FROM |_|_AggregateExec: mode=Final, gby=[], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -96,6 +97,7 @@ FROM |_|_|_AggregateExec: mode=Final, gby=[], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED @@ -148,6 +150,7 @@ FROM | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[avg(integers.i)]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[avg(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -172,6 +175,7 @@ FROM | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[avg(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[avg(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__avg_state(integers.i)] REDACTED @@ -249,6 +253,7 @@ ORDER BY |_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -285,6 +290,7 @@ ORDER BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED @@ -362,6 +368,7 @@ ORDER BY |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -398,6 +405,7 @@ ORDER BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED diff --git a/tests/cases/distributed/explain/step_aggr_advance.result b/tests/cases/distributed/explain/step_aggr_advance.result index a78a177bc3d..9fc8dfe5a1c 100644 --- a/tests/cases/distributed/explain/step_aggr_advance.result +++ b/tests/cases/distributed/explain/step_aggr_advance.result @@ -108,8 +108,11 @@ tql explain (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op | | ]] | | physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, greptime_timestamp@2 ASC NULLS LAST] | | | SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, greptime_timestamp@2 ASC NULLS LAST], preserve_partitioning=[true] | -| | AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -127,7 +130,10 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op +-+-+-+ | 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, greptime_timestamp@2 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, greptime_timestamp@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[__sum_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED @@ -177,8 +183,11 @@ tql explain (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim | | ]] | | physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST, greptime_timestamp@1 ASC NULLS LAST] | | | SortExec: expr=[a@0 ASC NULLS LAST, greptime_timestamp@1 ASC NULLS LAST], preserve_partitioning=[true] | -| | AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -196,7 +205,10 @@ tql analyze (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim +-+-+-+ | 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, greptime_timestamp@1 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, greptime_timestamp@1 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[__avg_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED @@ -311,8 +323,11 @@ tql explain (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr | | ]] | | physical_plan | SortPreservingMergeExec: [b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | | | SortExec: expr=[b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] | -| | AggregateExec: mode=SinglePartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -330,7 +345,10 @@ tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr +-+-+-+ | 0_| 0_|_SortPreservingMergeExec: [b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED @@ -380,7 +398,8 @@ tql explain sum(aggr_optimize_not); | | AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(aggr_optimize_not.greptime_value)] | | | RepartitionExec: partitioning=REDACTED | | AggregateExec: mode=Partial, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(aggr_optimize_not.greptime_value)] | -| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -401,6 +420,7 @@ tql analyze sum(aggr_optimize_not); |_|_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[__sum_state(aggr_optimize_not.greptime_value)] REDACTED @@ -477,14 +497,17 @@ tql explain (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize | physical_plan | ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp, sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))@5 / sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))@4 as aggr_optimize_not.sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000))) / aggr_optimize_not_count.sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] | | | REDACTED | | CoalescePartitionsExec | -| | AggregateExec: mode=SinglePartitioned, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] | -| | FilterExec: prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000))@1 IS NOT NULL | -| | ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_rate(greptime_timestamp_range@6, greptime_value@5, greptime_timestamp@4, 120000) as prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)), a@0 as a, b@1 as b, c@2 as c] | -| | PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] | -| | PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] | -| | PromSeriesDivideExec: tags=["a", "b", "c", "d"] | -| | SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true] | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] | +| | FilterExec: prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000))@1 IS NOT NULL | +| | ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_rate(greptime_timestamp_range@6, greptime_value@5, greptime_timestamp@4, 120000) as prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)), a@0 as a, b@1 as b, c@2 as c] | +| | PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] | +| | PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] | +| | PromSeriesDivideExec: tags=["a", "b", "c", "d"] | +| | SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | RepartitionExec: partitioning=REDACTED | | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | | | SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] | @@ -507,13 +530,16 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize | 0_| 0_|_ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp, sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))@5 / sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))@4 as aggr_optimize_not.sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000))) / aggr_optimize_not_count.sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED |_|_|_REDACTED |_|_|_CoalescePartitionsExec REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED |_|_|_FilterExec: prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000))@1 IS NOT NULL REDACTED |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_rate(greptime_timestamp_range@6, greptime_value@5, greptime_timestamp@4, 120000) as prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)), a@0 as a, b@1 as b, c@2 as c] REDACTED |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED |_|_|_SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED @@ -710,9 +736,11 @@ GROUP BY | | TableScan: aggr_optimize_not | | | ]] | | physical_plan | ProjectionExec: expr=[min(aggr_optimize_not.greptime_value)@2 as min(aggr_optimize_not.greptime_value)] | -| | AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] | -| | CooperativeExec | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -736,8 +764,10 @@ GROUP BY | stage | node | plan_| +-+-+-+ | 0_| 0_|_ProjectionExec: expr=[min(aggr_optimize_not.greptime_value)@2 as min(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_CooperativeExec REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED @@ -776,9 +806,11 @@ GROUP BY | | TableScan: aggr_optimize_not | | | ]] | | physical_plan | ProjectionExec: expr=[min(aggr_optimize_not.greptime_value)@2 + max(aggr_optimize_not.greptime_value)@3 as min(aggr_optimize_not.greptime_value) + max(aggr_optimize_not.greptime_value)] | -| | AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] | -| | CooperativeExec | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -802,8 +834,10 @@ GROUP BY | stage | node | plan_| +-+-+-+ | 0_| 0_|_ProjectionExec: expr=[min(aggr_optimize_not.greptime_value)@2 + max(aggr_optimize_not.greptime_value)@3 as min(aggr_optimize_not.greptime_value) + max(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_CooperativeExec REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value), __max_state(aggr_optimize_not.greptime_value)] REDACTED @@ -851,9 +885,11 @@ GROUP BY | | Projection: aggr_optimize_not.a, aggr_optimize_not.b, aggr_optimize_not.greptime_value | | | TableScan: aggr_optimize_not | | | ]] | -| physical_plan | AggregateExec: mode=SinglePartitioned, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] | -| | CooperativeExec | -| | MergeScanExec: REDACTED +| physical_plan | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -886,8 +922,10 @@ GROUP BY +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_CooperativeExec REDACTED +| 0_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED @@ -1017,6 +1055,7 @@ EXPLAIN SELECT COUNT(DISTINCT val_col_1) FROM step_aggr_extended; |_|_AggregateExec: mode=FinalPartitioned, gby=[alias1@0 as alias1], aggr=[]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[val_col_1@0 as alias1], aggr=[]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_ProjectionExec: expr=[val_col_1@2 as val_col_1]_| |_|_MergeScanExec: REDACTED |_|_| @@ -1060,6 +1099,7 @@ EXPLAIN SELECT pk_col_2, sum(val_col_1) FROM step_aggr_extended GROUP BY pk_col_ |_|_AggregateExec: mode=FinalPartitioned, gby=[pk_col_2@0 as pk_col_2], aggr=[sum(step_aggr_extended.val_col_1)]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[pk_col_2@0 as pk_col_2], aggr=[sum(step_aggr_extended.val_col_1)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -1096,6 +1136,7 @@ EXPLAIN SELECT SUM(val_col_3), COUNT(val_col_2), COUNT(val_col_3), COUNT(*) FROM |_|_AggregateExec: mode=Final, gby=[], aggr=[sum(step_aggr_extended.val_col_3), count(step_aggr_extended.val_col_2), count(step_aggr_extended.val_col_3), count(Int64(1))]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[sum(step_aggr_extended.val_col_3), count(step_aggr_extended.val_col_2), count(step_aggr_extended.val_col_3), count(Int64(1))]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -1130,6 +1171,7 @@ EXPLAIN SELECT MIN(pk_col_1), MAX(val_col_2) FROM step_aggr_extended; | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[min(step_aggr_extended.pk_col_1), max(step_aggr_extended.val_col_2)]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[min(step_aggr_extended.pk_col_1), max(step_aggr_extended.val_col_2)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -1167,6 +1209,7 @@ EXPLAIN SELECT SUM(val_col_1), COUNT(*) FROM step_aggr_extended WHERE pk_col_1 = |_|_AggregateExec: mode=Final, gby=[], aggr=[sum(step_aggr_extended.val_col_1), count(Int64(1))]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[sum(step_aggr_extended.val_col_1), count(Int64(1))]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/distributed/explain/step_aggr_basic.result b/tests/cases/distributed/explain/step_aggr_basic.result index 6a002115438..cf11b7f9ef0 100644 --- a/tests/cases/distributed/explain/step_aggr_basic.result +++ b/tests/cases/distributed/explain/step_aggr_basic.result @@ -60,6 +60,7 @@ FROM | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[count(integers.i)]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[count(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -85,6 +86,7 @@ FROM | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(integers.i)] REDACTED @@ -156,6 +158,7 @@ ORDER BY |_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(integers.i)]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -189,6 +192,7 @@ ORDER BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[__count_state(integers.i)] REDACTED @@ -262,6 +266,7 @@ ORDER BY |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[count(integers.i)]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[count(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -296,6 +301,7 @@ ORDER BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[count(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[count(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[__count_state(integers.i)] REDACTED @@ -377,6 +383,7 @@ ORDER BY |_|_AggregateExec: mode=FinalPartitioned, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[count(integers.i)] | |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[count(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -413,6 +420,7 @@ ORDER BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[count(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[count(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[__count_state(integers.i)] REDACTED @@ -500,6 +508,7 @@ FROM |_|_AggregateExec: mode=Final, gby=[], aggr=[uddsketch_merge(Int64(128),Float64(0.01),sink_table.udd_state), hll_merge(sink_table.hll_state)]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[uddsketch_merge(Int64(128),Float64(0.01),sink_table.udd_state), hll_merge(sink_table.hll_state)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -527,6 +536,7 @@ FROM |_|_|_AggregateExec: mode=Final, gby=[], aggr=[uddsketch_merge(Int64(128),Float64(0.01),sink_table.udd_state), hll_merge(sink_table.hll_state)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[uddsketch_merge(Int64(128),Float64(0.01),sink_table.udd_state), hll_merge(sink_table.hll_state)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__uddsketch_merge_state(Int64(128),Float64(0.01),sink_table.udd_state), __hll_merge_state(sink_table.hll_state)] REDACTED diff --git a/tests/cases/distributed/explain/subqueries.result b/tests/cases/distributed/explain/subqueries.result index e6f7286535e..d7ad6ed920c 100644 --- a/tests/cases/distributed/explain/subqueries.result +++ b/tests/cases/distributed/explain/subqueries.result @@ -101,9 +101,9 @@ order by t.i desc; | physical_plan | SortPreservingMergeExec: [i@0 DESC]_| |_|_SortExec: expr=[i@0 DESC], preserve_partitioning=[true]_| |_|_CrossJoinExec_| -|_|_CoalescePartitionsExec_| |_|_ProjectionExec: expr=[i@0 as i]_| |_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_ProjectionExec: expr=[]_| |_|_MergeScanExec: REDACTED |_|_| @@ -179,7 +179,9 @@ EXPLAIN SELECT * FROM integers i1 WHERE EXISTS(SELECT i FROM integers WHERE i=i1 | physical_plan | SortPreservingMergeExec: [i@0 ASC NULLS LAST]_| |_|_SortExec: expr=[i@0 ASC NULLS LAST], preserve_partitioning=[true]_| |_|_REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -212,7 +214,9 @@ EXPLAIN SELECT * FROM integers i1 WHERE EXISTS(SELECT count(i) FROM integers WHE | physical_plan | SortPreservingMergeExec: [i@0 ASC NULLS LAST]_| |_|_SortExec: expr=[i@0 ASC NULLS LAST], preserve_partitioning=[true]_| |_|_REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_ProjectionExec: expr=[i@1 as i]_| |_|_MergeScanExec: REDACTED |_|_| @@ -363,7 +367,10 @@ EXPLAIN SELECT DISTINCT x FROM (SELECT a AS x FROM t) sq ORDER BY x; |_| ]]_| | physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_| |_|_SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[true]_| -|_|_AggregateExec: mode=SinglePartitioned, gby=[x@0 as x], aggr=[]_| +|_|_AggregateExec: mode=FinalPartitioned, gby=[x@0 as x], aggr=[]_| +|_|_RepartitionExec: partitioning=REDACTED +|_|_AggregateExec: mode=Partial, gby=[x@0 as x], aggr=[]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -560,7 +567,9 @@ ORDER BY u1.a; | physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST]_| |_|_SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true]_| |_|_REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/distributed/flow-tql/tsid_on_phy.result b/tests/cases/distributed/flow-tql/tsid_on_phy.result index 857e45d16ad..68dd01374c1 100644 --- a/tests/cases/distributed/flow-tql/tsid_on_phy.result +++ b/tests/cases/distributed/flow-tql/tsid_on_phy.result @@ -119,8 +119,8 @@ TQL EXPLAIN ( | | TableScan: phy projection=[ts, v, tag1, tag2, le, tag4, tag5, tag6, tag7, tag8, __table_id, __tsid], partial_filters=[phy.ts >= TimestampMillisecond(1769137200001, None), phy.ts <= TimestampMillisecond(1769139900000, None), phy.__table_id=UInt32(REDACTED)] | | | ]] | | physical_plan | HistogramFoldExec: le=@0, field=@4, quantile=0.5 | -| | SortExec: expr=[tag4@1 ASC NULLS LAST, tag5@2 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[true] | -| | RepartitionExec: REDACTED +| | RepartitionExec: REDACTED +| | SortExec: expr=[tag4@1 ASC NULLS LAST, tag5@2 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] | | | MergeScanExec: REDACTED | | | +---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ diff --git a/tests/cases/distributed/optimizer/count.result b/tests/cases/distributed/optimizer/count.result index a339bc3bf42..7368ad6bbd7 100644 --- a/tests/cases/distributed/optimizer/count.result +++ b/tests/cases/distributed/optimizer/count.result @@ -281,6 +281,7 @@ select count(1) from count_where_bug; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[{count[count]:REDACTED} as __count_state(count_where_bug.ts)] REDACTED @@ -316,6 +317,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -342,6 +344,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -378,6 +381,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -417,6 +421,7 @@ select count(1) from count_where_bug where num != 3; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED diff --git a/tests/cases/distributed/optimizer/first_value_advance.result b/tests/cases/distributed/optimizer/first_value_advance.result index 08e1b8a4112..519ec9d9c15 100644 --- a/tests/cases/distributed/optimizer/first_value_advance.result +++ b/tests/cases/distributed/optimizer/first_value_advance.result @@ -325,6 +325,7 @@ explain select first_value(ts order by ts) from t; | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -348,6 +349,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -414,6 +416,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -447,6 +450,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -525,6 +529,7 @@ explain | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -548,6 +553,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED @@ -711,6 +717,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -744,6 +751,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/distributed/optimizer/last_value_advance.result b/tests/cases/distributed/optimizer/last_value_advance.result index 7b131d44f37..6269c5d397f 100644 --- a/tests/cases/distributed/optimizer/last_value_advance.result +++ b/tests/cases/distributed/optimizer/last_value_advance.result @@ -325,6 +325,7 @@ explain select last_value(ts order by ts) from t; | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -348,6 +349,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -414,6 +416,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -447,6 +450,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -525,6 +529,7 @@ explain | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -548,6 +553,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED @@ -711,6 +717,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -744,6 +751,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/distributed/optimizer/range_select_projection.result b/tests/cases/distributed/optimizer/range_select_projection.result index c32878abbf0..451e743da9c 100644 --- a/tests/cases/distributed/optimizer/range_select_projection.result +++ b/tests/cases/distributed/optimizer/range_select_projection.result @@ -48,7 +48,7 @@ ORDER BY station, "channel", ts; | 0_| 0_|_SortExec: expr=[station@1 ASC NULLS LAST, channel@2 ASC NULLS LAST, ts@0 ASC NULLS LAST], preserve_REDACTED |_|_|_ProjectionExec: expr=[ts@1 as ts, station@2 as station, channel@3 as channel, avg(range_select_projection.value_a + range_select_projection.value_b) RANGE 10s@0 as avg_value] REDACTED |_|_|_RangeSelectExec: range_expr=[avg(range_select_projection.value_a + range_select_projection.value_b) RANGE 10s], align=5000ms, align_to=0ms, align_by=[station@1, channel@2], time_index=ts REDACTED -|_|_|_CoalescePartitionsExec REDACTED +|_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED diff --git a/tests/cases/standalone/common/aggregate/distinct.result b/tests/cases/standalone/common/aggregate/distinct.result index 5a721b7a481..5c10ee22239 100644 --- a/tests/cases/standalone/common/aggregate/distinct.result +++ b/tests/cases/standalone/common/aggregate/distinct.result @@ -169,7 +169,10 @@ EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; +-+-+-+ | 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a], aggr=[] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[] REDACTED @@ -199,7 +202,10 @@ EXPLAIN ANALYZE SELECT DISTINCT a, b FROM test ORDER BY a; +-+-+-+ | 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED diff --git a/tests/cases/standalone/common/aggregate/multi_regions.result b/tests/cases/standalone/common/aggregate/multi_regions.result index 4b104df849d..8627b83c153 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.result +++ b/tests/cases/standalone/common/aggregate/multi_regions.result @@ -62,6 +62,7 @@ select sum(val) from t; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[sum(t.val)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[sum(t.val)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__sum_state(t.val)] REDACTED @@ -96,6 +97,7 @@ select sum(val) from t group by idc; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[idc@0 as idc], aggr=[sum(t.val)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[idc@0 as idc], aggr=[sum(t.val)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[idc@0 as idc], aggr=[__sum_state(t.val)] REDACTED diff --git a/tests/cases/standalone/common/create/metric_engine_partition.result b/tests/cases/standalone/common/create/metric_engine_partition.result index 0a5899bd9b3..ff3fd40e940 100644 --- a/tests/cases/standalone/common/create/metric_engine_partition.result +++ b/tests/cases/standalone/common/create/metric_engine_partition.result @@ -148,7 +148,10 @@ select host, count(*) from logical_table_2 GROUP BY host ORDER BY host; | physical_plan | SortPreservingMergeExec: [host@0 ASC NULLS LAST]_| |_|_SortExec: expr=[host@0 ASC NULLS LAST], preserve_partitioning=[true]_| |_|_ProjectionExec: expr=[host@0 as host, count(Int64(1))@1 as count(*)]_| -|_|_AggregateExec: mode=SinglePartitioned, gby=[host@0 as host], aggr=[count(Int64(1))]_| +|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED +|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -188,6 +191,7 @@ select ts, count(*) from logical_table_2 GROUP BY ts ORDER BY ts; |_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(Int64(1))]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -283,6 +287,7 @@ select a, count(*) from logical_table_3 GROUP BY a ORDER BY a; |_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(Int64(1))]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -351,7 +356,7 @@ EXPLAIN select count(*) from logical_table_4; | plan_type_| plan_| +-+-+ | logical_plan_| Projection: count(Int64(1)) AS count(*)_| -|_|_Aggregate: groupBy=[[]], aggr=[[__count_merge(__count_state(logical_table_4.ts)) AS count(Int64(1))]] | +|_|_Aggregate: groupBy=[[]], aggr=[[__count_merge(__count_state(logical_table_4.ts)) AS count(Int64(1))]]_| |_|_MergeScan [is_placeholder=false, remote_input=[_| |_| Aggregate: groupBy=[[]], aggr=[[__count_state(logical_table_4.ts)]]_| |_|_TableScan: logical_table_4_| @@ -360,6 +365,7 @@ EXPLAIN select count(*) from logical_table_4; |_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -399,6 +405,7 @@ select ts, count(*) from logical_table_4 GROUP BY ts ORDER BY ts; |_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(Int64(1))]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/standalone/common/cte/cte_join_build_side.result b/tests/cases/standalone/common/cte/cte_join_build_side.result new file mode 100644 index 00000000000..51ee4d98a4d --- /dev/null +++ b/tests/cases/standalone/common/cte/cte_join_build_side.result @@ -0,0 +1,116 @@ +-- Regression test for issue 6874: a bounded CTE derived from the same table +-- must become the hash-join build side in both standalone and distributed mode. +CREATE TABLE cte_join_logs ( + ts TIMESTAMP TIME INDEX, + "db" STRING, + val DOUBLE, +) WITH (append_mode = 'true'); + +Affected Rows: 0 + +INSERT INTO cte_join_logs VALUES + ('2024-01-01 00:00:00', 'a', 1.0), + ('2024-01-01 00:00:01', 'a', 2.0), + ('2024-01-01 00:00:02', 'b', 3.0), + ('2024-01-01 00:00:03', 'b', 4.0), + ('2024-01-01 00:00:04', 'c', 5.0), + ('2024-01-01 00:00:05', 'c', 6.0); + +Affected Rows: 6 + +-- Control: the plain Top-N aggregation should remain below MergeScan. +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +EXPLAIN SELECT "db", count(*) AS cnt +FROM cte_join_logs +GROUP BY "db" +ORDER BY cnt DESC +LIMIT 5; + ++---------------+----------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+----------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Limit: skip=0, fetch=5 | +| | Sort: cnt DESC NULLS FIRST | +| | Projection: cte_join_logs.db, count(Int64(1)) AS count(*) AS cnt | +| | Aggregate: groupBy=[[cte_join_logs.db]], aggr=[[count(cte_join_logs.ts) AS count(Int64(1))]] | +| | TableScan: cte_join_logs | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+----------------------------------------------------------------------------------------------------+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +EXPLAIN WITH top_dbs AS ( + SELECT "db", count(*) AS cnt FROM cte_join_logs + GROUP BY "db" ORDER BY cnt DESC LIMIT 5 +) +SELECT t."db", count(*) AS c +FROM cte_join_logs AS t +JOIN top_dbs td ON t."db" = td."db" +GROUP BY t."db" ORDER BY c DESC, t."db"; + ++---------------+----------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+----------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: c DESC NULLS FIRST, t.db ASC NULLS LAST | +| | Projection: t.db, count(Int64(1)) AS count(*) AS c | +| | Aggregate: groupBy=[[t.db]], aggr=[[count(Int64(1))]] | +| | Projection: t.db | +| | Inner Join: t.db = td.db | +| | Projection: t.db | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: t | +| | Filter: cte_join_logs.db IS NOT NULL | +| | TableScan: cte_join_logs, partial_filters=[cte_join_logs.db IS NOT NULL] | +| | ]] | +| | Projection: td.db | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: td | +| | SubqueryAlias: top_dbs | +| | Filter: cte_join_logs.db IS NOT NULL | +| | Limit: skip=0, fetch=5 | +| | Sort: cnt DESC NULLS FIRST | +| | Projection: cte_join_logs.db, count(Int64(1)) AS count(*) AS cnt | +| | Aggregate: groupBy=[[cte_join_logs.db]], aggr=[[count(cte_join_logs.ts) AS count(Int64(1))]] | +| | TableScan: cte_join_logs | +| | ]] | +| physical_plan | SortPreservingMergeExec: [c@1 DESC, db@0 ASC NULLS LAST] | +| | SortExec: expr=[c@1 DESC, db@0 ASC NULLS LAST], preserve_REDACTED +| | ProjectionExec: expr=[db@0 as db, count(Int64(1))@1 as c] | +| | AggregateExec: mode=FinalPartitioned, gby=[db@0 as db], aggr=[count(Int64(1))] | +| | RepartitionExec: REDACTED +| | AggregateExec: mode=Partial, gby=[db@0 as db], aggr=[count(Int64(1))] | +| | HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(db@0, db@0)], projection=[db@1] | +| | ProjectionExec: expr=[db@0 as db] | +| | MergeScanExec: REDACTED +| | RepartitionExec: REDACTED +| | ProjectionExec: expr=[db@1 as db] | +| | MergeScanExec: REDACTED +| | | ++---------------+----------------------------------------------------------------------------------------------------------+ + +WITH top_dbs AS ( + SELECT "db", count(*) AS cnt FROM cte_join_logs + GROUP BY "db" ORDER BY cnt DESC LIMIT 5 +) +SELECT t."db", count(*) AS c +FROM cte_join_logs AS t +JOIN top_dbs td ON t."db" = td."db" +GROUP BY t."db" ORDER BY c DESC, t."db"; + ++----+---+ +| db | c | ++----+---+ +| a | 2 | +| b | 2 | +| c | 2 | ++----+---+ + +DROP TABLE cte_join_logs; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/cte/cte_join_build_side.sql b/tests/cases/standalone/common/cte/cte_join_build_side.sql new file mode 100644 index 00000000000..f609b826042 --- /dev/null +++ b/tests/cases/standalone/common/cte/cte_join_build_side.sql @@ -0,0 +1,46 @@ +-- Regression test for issue 6874: a bounded CTE derived from the same table +-- must become the hash-join build side in both standalone and distributed mode. +CREATE TABLE cte_join_logs ( + ts TIMESTAMP TIME INDEX, + "db" STRING, + val DOUBLE, +) WITH (append_mode = 'true'); + +INSERT INTO cte_join_logs VALUES + ('2024-01-01 00:00:00', 'a', 1.0), + ('2024-01-01 00:00:01', 'a', 2.0), + ('2024-01-01 00:00:02', 'b', 3.0), + ('2024-01-01 00:00:03', 'b', 4.0), + ('2024-01-01 00:00:04', 'c', 5.0), + ('2024-01-01 00:00:05', 'c', 6.0); + +-- Control: the plain Top-N aggregation should remain below MergeScan. +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +EXPLAIN SELECT "db", count(*) AS cnt +FROM cte_join_logs +GROUP BY "db" +ORDER BY cnt DESC +LIMIT 5; + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +EXPLAIN WITH top_dbs AS ( + SELECT "db", count(*) AS cnt FROM cte_join_logs + GROUP BY "db" ORDER BY cnt DESC LIMIT 5 +) +SELECT t."db", count(*) AS c +FROM cte_join_logs AS t +JOIN top_dbs td ON t."db" = td."db" +GROUP BY t."db" ORDER BY c DESC, t."db"; + +WITH top_dbs AS ( + SELECT "db", count(*) AS cnt FROM cte_join_logs + GROUP BY "db" ORDER BY cnt DESC LIMIT 5 +) +SELECT t."db", count(*) AS c +FROM cte_join_logs AS t +JOIN top_dbs td ON t."db" = td."db" +GROUP BY t."db" ORDER BY c DESC, t."db"; + +DROP TABLE cte_join_logs; diff --git a/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result b/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result index a81fc4918bb..5798ddd07c1 100644 --- a/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result +++ b/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result @@ -67,11 +67,11 @@ WHERE c.tier IN ('gold', 'bronze'); | | Filter: customers.customer_id IS NOT NULL AND (customers.tier = Utf8("gold") OR customers.tier = Utf8("bronze")) | | | TableScan: customers, partial_filters=[customers.tier = Utf8("gold") OR customers.tier = Utf8("bronze"), customers.customer_id IS NOT NULL] | | | ]] | -| physical_plan | HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@1, customer_id@0)], projection=[id@0, amount@2, name@4, tier@5] | -| | RepartitionExec: partitioning=REDACTED -| | ProjectionExec: expr=[id@0 as id, customer_id@1 as customer_id, amount@2 as amount] | +| physical_plan | HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(customer_id@1, customer_id@0)], projection=[id@0, amount@2, name@4, tier@5] | +| | ProjectionExec: expr=[id@0 as id, customer_id@1 as customer_id, amount@2 as amount] | +| | CooperativeExec | | | MergeScanExec: REDACTED -| | RepartitionExec: partitioning=REDACTED +| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 | | | ProjectionExec: expr=[customer_id@0 as customer_id, name@1 as name, tier@2 as tier] | | | MergeScanExec: REDACTED | | | @@ -107,9 +107,9 @@ WHERE c.tier IN ('gold', 'bronze'); +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@1, customer_id@0)], projection=[id@0, amount@2, name@4, tier@5] metrics=REDACTED_| -|_|_|_RepartitionExec: partitioning=REDACTED +| 0_| 0_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(customer_id@1, customer_id@0)], projection=[id@0, amount@2, name@4, tier@5] metrics=REDACTED_| |_|_|_ProjectionExec: expr=[id@0 as id, customer_id@1 as customer_id, amount@2 as amount] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| |_|_|_MergeScanExec: REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_ProjectionExec: expr=[customer_id@0 as customer_id, name@1 as name, tier@2 as tier] metrics=REDACTED_| diff --git a/tests/cases/standalone/common/order/order_by_exceptions.result b/tests/cases/standalone/common/order/order_by_exceptions.result index bde25ea927c..bdf215f8f26 100644 --- a/tests/cases/standalone/common/order/order_by_exceptions.result +++ b/tests/cases/standalone/common/order/order_by_exceptions.result @@ -82,9 +82,10 @@ EXPLAIN SELECT a % 2, b FROM test UNION SELECT a % 2 AS k, b FROM test ORDER BY | | AggregateExec: mode=FinalPartitioned, gby=[test.a % Int64(2)@0 as test.a % Int64(2), b@1 as b], aggr=[] | | | RepartitionExec: REDACTED | | AggregateExec: mode=Partial, gby=[test.a % Int64(2)@0 as test.a % Int64(2), b@1 as b], aggr=[] | -| | InterleaveExec | -| | MergeScanExec: REDACTED -| | MergeScanExec: REDACTED +| | RepartitionExec: REDACTED +| | InterleaveExec | +| | MergeScanExec: REDACTED +| | MergeScanExec: REDACTED | | | +---------------+-----------------------------------------------------------------------------------------------------------+ diff --git a/tests/cases/standalone/common/promql/histogram_multi_partition.result b/tests/cases/standalone/common/promql/histogram_multi_partition.result index 091ed7bfcc3..3d5490f38a1 100644 --- a/tests/cases/standalone/common/promql/histogram_multi_partition.result +++ b/tests/cases/standalone/common/promql/histogram_multi_partition.result @@ -47,6 +47,7 @@ tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bu |_|_|_AggregateExec: mode=FinalPartitioned, gby=[le@0 as le, ts@1 as ts], aggr=[sum(histogram_gap_bucket.val)] REDACTED |_|_|_RepartitionExec: partitioning=Hash([le@0, ts@1],REDACTED |_|_|_AggregateExec: mode=Partial, gby=[le@0 as le, ts@1 as ts], aggr=[sum(histogram_gap_bucket.val)] REDACTED +|_|_|_RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2 REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[le@0 as le, ts@1 as ts], aggr=[__sum_state(histogram_gap_bucket.val)] REDACTED diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result index 75aaf8176c9..72a6dc9c5b4 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result @@ -117,11 +117,11 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; +-+-+-+ | 0_| 0_|_ProjectionExec: expr=[host@2 as host, job@3 as job, ts@5 as ts, __tsid@4 as __tsid, greptime_value@0 / greptime_value@1 as tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value] REDACTED |_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@2, ts@4)], projection=[greptime_value@0, greptime_value@3, host@4, job@5, __tsid@6, ts@7], NullsEqual: true REDACTED -|_|_|_CoalescePartitionsExec REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, __tsid@3 as __tsid, ts@4 as ts] REDACTED -|_|_|_MergeScanExec: REDACTED |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED @@ -363,11 +363,11 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; +-+-+-+ | 0_| 0_|_ProjectionExec: expr=[host@2 as host, job@3 as job, ts@5 as ts, __tsid@4 as __tsid, CAST(greptime_value@1 < greptime_value@0 AS Float64) as tsid_binary_join_left.greptime_value > tsid_binary_join_right.greptime_value] REDACTED |_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@2, ts@4)], projection=[greptime_value@0, greptime_value@3, host@4, job@5, __tsid@6, ts@7], NullsEqual: true REDACTED -|_|_|_CoalescePartitionsExec REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, __tsid@3 as __tsid, ts@4 as ts] REDACTED -|_|_|_MergeScanExec: REDACTED |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED @@ -411,7 +411,7 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsi |_|_|_RepartitionExec: partitioning=Hash([REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, __tsid@3 as __tsid, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED -|_|_|_CooperativeExec REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED @@ -466,7 +466,7 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) |_|_|_RepartitionExec: partitioning=Hash([REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, __tsid@3 as __tsid, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED -|_|_|_CooperativeExec REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED @@ -516,12 +516,12 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@0, ts@4)], projection=[greptime_value@2, greptime_value@3, host@4, job@5, __tsid@6, ts@7], NullsEqual: true REDACTED |_|_|_ProjectionExec: expr=[ts@0 as ts, __tsid@1 as __tsid, greptime_value@2 as greptime_value] REDACTED |_|_|_UnionDistinctOnExec: on col=[5, 6], ts_col=0 REDACTED -|_|_|_CoalescePartitionsExec REDACTED -|_|_|_MergeScanExec: REDACTED -|_|_|_CoalescePartitionsExec REDACTED +|_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[ts@2 as ts, __tsid@3 as __tsid, greptime_value@4 as greptime_value, host@5 as host, job@6 as job, CASE WHEN __common_expr_1@0 IS NOT NULL THEN __common_expr_1@0 ELSE_END as __promql_or_match_0, CASE WHEN __common_expr_2@1 IS NOT NULL THEN __common_expr_2@1 ELSE_END as __promql_or_match_1] REDACTED |_|_|_ProjectionExec: expr=[CAST(host@1 AS Utf8) as __common_expr_1, CAST(job@2 AS Utf8) as __common_expr_2, ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job] REDACTED @@ -575,7 +575,7 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid |_|_|_RepartitionExec: partitioning=Hash([REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, __tsid@3 as __tsid, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED -|_|_|_CooperativeExec REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED diff --git a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result index fc883d96e71..ba2e97bcc2d 100644 --- a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result +++ b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result @@ -109,8 +109,8 @@ TQL ANALYZE (0, 10, '5s') histogram_quantile(0.5, tsid_no_aggr_histogram_bucket) | stage | node | plan_| +-+-+-+ | 0_| 0_|_HistogramFoldExec: le=@2, field=@0, quantile=0.5 REDACTED -|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@2 AS Float64) ASC NULLS LAST], preserve_partitioning=[true] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@2 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[val@0 as val, job@1 as job, le@2 as le, ts@4 as ts] REDACTED diff --git a/tests/cases/standalone/common/range/nest.result b/tests/cases/standalone/common/range/nest.result index 74742d5d862..184b27545d0 100644 --- a/tests/cases/standalone/common/range/nest.result +++ b/tests/cases/standalone/common/range/nest.result @@ -138,7 +138,7 @@ EXPLAIN SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s'; |_|_TableScan: host_| |_| ]]_| | physical_plan | RangeSelectExec: range_expr=[min(host.val) RANGE 5s], align=5000ms, align_to=0ms, align_by=[host@1], time_index=ts | -|_|_CoalescePartitionsExec_| +|_|_CooperativeExec_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -154,7 +154,7 @@ EXPLAIN ANALYZE SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s'; | stage | node | plan_| +-+-+-+ | 0_| 0_|_RangeSelectExec: range_expr=[min(host.val) RANGE 5s], align=5000ms, align_to=0ms, align_by=[host@1], time_index=ts REDACTED -|_|_|_CoalescePartitionsExec REDACTED +|_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED diff --git a/tests/cases/standalone/common/tql/partition.result b/tests/cases/standalone/common/tql/partition.result index 45927f8abbd..f1072edbb32 100644 --- a/tests/cases/standalone/common/tql/partition.result +++ b/tests/cases/standalone/common/tql/partition.result @@ -68,7 +68,10 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_SortPreservingMergeExec: [k@0 ASC NULLS LAST, j@1 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[k@0 ASC NULLS LAST, j@1 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[k@0 as k, j@1 as j], aggr=[avg(prom_irate(j_range,i))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[k@0 as k, j@1 as j], aggr=[avg(prom_irate(j_range,i))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[k@0 as k, j@1 as j], aggr=[avg(prom_irate(j_range,i))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[k@0 as k, j@1 as j], aggr=[__avg_state(prom_irate(j_range,i))] REDACTED diff --git a/tests/cases/standalone/common/tql/tql-cte.result b/tests/cases/standalone/common/tql/tql-cte.result index 0547754b483..8efe983763b 100644 --- a/tests/cases/standalone/common/tql/tql-cte.result +++ b/tests/cases/standalone/common/tql/tql-cte.result @@ -487,8 +487,8 @@ ORDER BY ts; | | ]] | | physical_plan | ProjectionExec: expr=[ts@0 as ts, val@1 as val, lag(tql_base.val,Int64(1)) ORDER BY [tql_base.ts ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as prev_value] | | | BoundedWindowAggExec: wdw=[lag(tql_base.val,Int64(1)) ORDER BY [tql_base.ts ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "lag(tql_base.val,Int64(1)) ORDER BY [tql_base.ts ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Float64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] | -| | SortPreservingMergeExec: [ts@0 ASC NULLS LAST] | -| | SortExec: expr=[ts@0 ASC NULLS LAST], preserve_REDACTED +| | SortExec: expr=[ts@0 ASC NULLS LAST], preserve_REDACTED +| | CooperativeExec | | | MergeScanExec: REDACTED | | | +---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -765,11 +765,15 @@ LIMIT 5; | | ProjectionExec: expr=[ts@0 as ts, cpu@1 as avg_value, host@2 as host] | | | HashJoinExec: mode=Partitioned, join_type=Inner, on=[(date_trunc(Utf8("second"),t.ts)@2, date_trunc(Utf8("second"),l.ts)@2)], projection=[ts@0, cpu@1, host@4] | | | RepartitionExec: REDACTED -| | ProjectionExec: expr=[ts@0 as ts, cpu@2 as cpu, date_trunc(second, ts@0) as date_trunc(Utf8("second"),t.ts)] | -| | MergeScanExec: REDACTED +| | ProjectionExec: expr=[ts@0 as ts, cpu@1 as cpu, date_trunc(second, ts@0) as date_trunc(Utf8("second"),t.ts)] | +| | RepartitionExec: REDACTED +| | ProjectionExec: expr=[ts@0 as ts, cpu@2 as cpu] | +| | MergeScanExec: REDACTED | | RepartitionExec: REDACTED | | ProjectionExec: expr=[ts@0 as ts, host@1 as host, date_trunc(second, ts@0) as date_trunc(Utf8("second"),l.ts)] | -| | MergeScanExec: REDACTED +| | RepartitionExec: REDACTED +| | ProjectionExec: expr=[ts@0 as ts, host@1 as host] | +| | MergeScanExec: REDACTED | | | +---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -843,11 +847,15 @@ LIMIT 5; | | ProjectionExec: expr=[ts@1 as ts, cpu@0 as avg_value, host@2 as host] | | | HashJoinExec: mode=Partitioned, join_type=Inner, on=[(date_trunc(Utf8("second"),t.ts)@2, date_trunc(Utf8("second"),l.ts)@2)], projection=[cpu@0, ts@1, host@4] | | | RepartitionExec: REDACTED -| | ProjectionExec: expr=[cpu@0 as cpu, ts@2 as ts, date_trunc(second, ts@2) as date_trunc(Utf8("second"),t.ts)] | -| | MergeScanExec: REDACTED +| | ProjectionExec: expr=[cpu@0 as cpu, ts@1 as ts, date_trunc(second, ts@1) as date_trunc(Utf8("second"),t.ts)] | +| | RepartitionExec: REDACTED +| | ProjectionExec: expr=[cpu@0 as cpu, ts@2 as ts] | +| | MergeScanExec: REDACTED | | RepartitionExec: REDACTED | | ProjectionExec: expr=[ts@0 as ts, host@1 as host, date_trunc(second, ts@0) as date_trunc(Utf8("second"),l.ts)] | -| | MergeScanExec: REDACTED +| | RepartitionExec: REDACTED +| | ProjectionExec: expr=[ts@0 as ts, host@1 as host] | +| | MergeScanExec: REDACTED | | | +---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------+ diff --git a/tests/cases/standalone/optimizer/count.result b/tests/cases/standalone/optimizer/count.result index 99f910136fd..8a128a92000 100644 --- a/tests/cases/standalone/optimizer/count.result +++ b/tests/cases/standalone/optimizer/count.result @@ -278,6 +278,7 @@ select count(1) from count_where_bug; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[{count[count]:REDACTED} as __count_state(count_where_bug.ts)] REDACTED @@ -313,6 +314,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -339,6 +341,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -375,6 +378,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -414,6 +418,7 @@ select count(1) from count_where_bug where num != 3; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED diff --git a/tests/cases/standalone/optimizer/first_value_advance.result b/tests/cases/standalone/optimizer/first_value_advance.result index afcefae4cff..c543a8df2a2 100644 --- a/tests/cases/standalone/optimizer/first_value_advance.result +++ b/tests/cases/standalone/optimizer/first_value_advance.result @@ -322,6 +322,7 @@ explain select first_value(ts order by ts) from t; | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -345,6 +346,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -411,6 +413,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -444,6 +447,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -522,6 +526,7 @@ explain | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -545,6 +550,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED @@ -705,6 +711,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -738,6 +745,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/standalone/optimizer/last_value_advance.result b/tests/cases/standalone/optimizer/last_value_advance.result index a1882a75887..3199692a73d 100644 --- a/tests/cases/standalone/optimizer/last_value_advance.result +++ b/tests/cases/standalone/optimizer/last_value_advance.result @@ -322,6 +322,7 @@ explain select last_value(ts order by ts) from t; | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -345,6 +346,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -411,6 +413,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -444,6 +447,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -522,6 +526,7 @@ explain | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -545,6 +550,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED @@ -705,6 +711,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -738,6 +745,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/standalone/optimizer/rewrite_set_comparison.result b/tests/cases/standalone/optimizer/rewrite_set_comparison.result index eb41d5f78e8..6d9e04458c7 100644 --- a/tests/cases/standalone/optimizer/rewrite_set_comparison.result +++ b/tests/cases/standalone/optimizer/rewrite_set_comparison.result @@ -81,11 +81,12 @@ EXPLAIN SELECT v FROM sc_t WHERE v > ANY(SELECT v FROM sc_s) ORDER BY v; | | CoalescePartitionsExec | | | FilterExec: mark@1 OR NOT mark@1 AND NULL | | | NestedLoopJoinExec: join_type=LeftMark, filter=(v@0 > v@1) IS NOT DISTINCT FROM true | -| | CoalescePartitionsExec | -| | ProjectionExec: expr=[v@1 as v] | -| | MergeScanExec: REDACTED -| | MergeScanExec: REDACTED -| | MergeScanExec: REDACTED +| | ProjectionExec: expr=[v@1 as v] | +| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 | +| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 | +| | MergeScanExec: REDACTED | | | +---------------+---------------------------------------------------------------------------------------------------------------+ @@ -135,11 +136,12 @@ EXPLAIN SELECT v FROM sc_t WHERE v != ALL(SELECT v FROM sc_s) ORDER BY v; | | CoalescePartitionsExec | | | FilterExec: NOT mark@1, projection=[v@0] | | | NestedLoopJoinExec: join_type=LeftMark, filter=(v@0 != v@1) IS NOT DISTINCT FROM false | -| | CoalescePartitionsExec | -| | ProjectionExec: expr=[v@1 as v] | -| | MergeScanExec: REDACTED -| | MergeScanExec: REDACTED -| | MergeScanExec: REDACTED +| | ProjectionExec: expr=[v@1 as v] | +| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 | +| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 | +| | MergeScanExec: REDACTED | | | +---------------+----------------------------------------------------------------------------------------------------+ diff --git a/tests/cases/standalone/tql-explain-analyze/analyze.result b/tests/cases/standalone/tql-explain-analyze/analyze.result index 1753d59da91..203debf90bf 100644 --- a/tests/cases/standalone/tql-explain-analyze/analyze.result +++ b/tests/cases/standalone/tql-explain-analyze/analyze.result @@ -288,6 +288,7 @@ TQL ANALYZE sum(test2); |_|_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(test2.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(test2.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[__sum_state(test2.greptime_value)] REDACTED diff --git a/tests/cases/standalone/tql-explain-analyze/tsid_column.result b/tests/cases/standalone/tql-explain-analyze/tsid_column.result index 4a7a8750605..8d27a55ab38 100644 --- a/tests/cases/standalone/tql-explain-analyze/tsid_column.result +++ b/tests/cases/standalone/tql-explain-analyze/tsid_column.result @@ -102,10 +102,10 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid | 0_| 0_|_ProjectionExec: expr=[ts@1 as ts, sum(prom_irate(ts_range,val))@2 / scalar(count(count(tsid_metric.val)))@0 as lhs.sum(prom_irate(ts_range,val)) / rhs.scalar(count(count(tsid_metric.val)))] REDACTED |_|_|_REDACTED |_|_|_ScalarCalculateExec: tags=[] REDACTED -|_|_|_CoalescePartitionsExec REDACTED -|_|_|_MergeScanExec: REDACTED |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -153,10 +153,10 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m | 0_| 0_|_ProjectionExec: expr=[ts@1 as ts, sum(prom_irate(ts_range,val))@2 / scalar(count(sum(tsid_metric.val)))@0 as lhs.sum(prom_irate(ts_range,val)) / rhs.scalar(count(sum(tsid_metric.val)))] REDACTED |_|_|_REDACTED |_|_|_ScalarCalculateExec: tags=[] REDACTED -|_|_|_CoalescePartitionsExec REDACTED -|_|_|_MergeScanExec: REDACTED |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED diff --git a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result index 07d0b8c764c..45aac8d9c39 100644 --- a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result +++ b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result @@ -17,6 +17,7 @@ EXPLAIN ANALYZE VERBOSE SELECT count(*) FROM t_analyze_verbose_remote_metrics_ex |_|_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(t_analyze_verbose_remote_REDACTED diff --git a/tests/conf/datanode-test.toml.template b/tests/conf/datanode-test.toml.template index 63d3be2f8c4..907fcee4ddf 100644 --- a/tests/conf/datanode-test.toml.template +++ b/tests/conf/datanode-test.toml.template @@ -3,6 +3,9 @@ mode = 'distributed' require_lease_before_startup = true rpc_runtime_size = 8 +[query] +parallelism = 4 + [[region_engine]] [region_engine.mito] {{ if enable_flat_format }} diff --git a/tests/conf/frontend-test.toml.template b/tests/conf/frontend-test.toml.template index 0042716bd80..632947a3cde 100644 --- a/tests/conf/frontend-test.toml.template +++ b/tests/conf/frontend-test.toml.template @@ -2,5 +2,8 @@ bind_addr = "{addrs.grpc_addr}" server_addr = "{addrs.grpc_addr}" +[query] +parallelism = 4 + [event_recorder] event_types = [] diff --git a/tests/conf/standalone-test.toml.template b/tests/conf/standalone-test.toml.template index 88e3b1018cd..ebe25f82487 100644 --- a/tests/conf/standalone-test.toml.template +++ b/tests/conf/standalone-test.toml.template @@ -2,6 +2,9 @@ mode = 'standalone' enable_memory_catalog = false require_lease_before_startup = true +[query] +parallelism = 4 + [[region_engine]] [region_engine.mito] {{ if enable_flat_format }} From 354921c80eaf13c975b38e8cdd6f6767001f5a8e Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 13 Aug 2026 10:58:29 +0000 Subject: [PATCH 064/126] chore: update mysql test drivers and lru (#8868) * chore: update mysql test drivers and lru * fix: test --- Cargo.lock | 192 +++++++++++-------- src/servers/Cargo.toml | 5 +- src/servers/tests/mysql/mysql_server_test.rs | 2 +- tests-integration/Cargo.toml | 2 +- tests/runner/Cargo.toml | 2 +- 5 files changed, 123 insertions(+), 80 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a7a7ff9b877..ece381b7b84 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -424,7 +424,7 @@ checksum = "0c6cd424c2693bcdbc150d843dc9d4d137dd2de4782ce6df491ad11a3a0416c0" dependencies = [ "bytes", "half", - "num-bigint", + "num-bigint 0.4.6", "num-traits", ] @@ -1254,7 +1254,7 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6773ddc0eafc0e509fb60e48dff7f450f8e674a0686ae8605e8d9901bd5eefa" dependencies = [ - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", "serde", @@ -1268,7 +1268,7 @@ checksum = "1a22f228ab7a1b23027ccc6c350b72868017af7ea8356fbdf19f8d991c690013" dependencies = [ "autocfg", "libm", - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", ] @@ -1517,6 +1517,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "btoi" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b5ab9db53bcda568284df0fd39f6eac24ad6f7ba7ff1168b9e76eba6576b976" +dependencies = [ + "num-traits", +] + [[package]] name = "bufstream" version = "0.1.4" @@ -1620,7 +1629,7 @@ dependencies = [ "maybe-owned", "rustix 1.0.7", "rustix-linux-procfs", - "windows-sys 0.60.2", + "windows-sys 0.61.2", "winx", ] @@ -2232,7 +2241,7 @@ checksum = "fe6d2e5af09e8c8ad56c969f2157a3d4238cebc7c55f0a517728c38f7b200f81" dependencies = [ "serde", "termcolor", - "unicode-width 0.1.14", + "unicode-width 0.2.1", ] [[package]] @@ -2243,7 +2252,7 @@ checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" dependencies = [ "serde", "termcolor", - "unicode-width 0.1.14", + "unicode-width 0.2.1", ] [[package]] @@ -2259,7 +2268,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -5126,7 +5135,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -6197,6 +6206,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" dependencies = [ "allocator-api2", + "equivalent", + "foldhash 0.2.0", ] [[package]] @@ -6539,7 +6550,7 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", - "webpki-roots 1.0.1", + "webpki-roots 1.0.9", ] [[package]] @@ -6585,7 +6596,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.4", "tokio", "tower-service", "tracing", @@ -7084,7 +7095,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20fd6de4ccfcc187e38bc21cfa543cb5a302cb86a8b114eb7f0bf0dc9f8ac00f" dependencies = [ "io-lifetimes 3.0.1", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -7825,7 +7836,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.48.5", + "windows-targets 0.53.5", ] [[package]] @@ -8099,15 +8110,6 @@ dependencies = [ "vob", ] -[[package]] -name = "lru" -version = "0.12.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" -dependencies = [ - "hashbrown 0.15.4", -] - [[package]] name = "lru" version = "0.16.4" @@ -8117,6 +8119,15 @@ dependencies = [ "hashbrown 0.16.1", ] +[[package]] +name = "lru" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d2f2f9b4ba7e6b24d95e7e899329d35be83bcded72c8540cdd5368932d1d90a" +dependencies = [ + "hashbrown 0.17.1", +] + [[package]] name = "lru-slab" version = "0.1.2" @@ -8734,28 +8745,29 @@ checksum = "2195bf6aa996a481483b29d62a7663eed3fe39600c460e323f8ff41e90bdd89b" [[package]] name = "mysql" -version = "26.0.0" +version = "28.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64453aedc258ac8c720b46c8264302fad39cef6c02483f68adbad4bcd22d6fab" +checksum = "4a732193888328fc060ab901c0ed1355521267a51ffbfd9a0b3786434c6b8e7f" dependencies = [ "bufstream", "bytes", "crossbeam-queue", + "crossbeam-utils", "flate2", "io-enum", "libc", - "lru 0.12.5", - "mysql_common 0.34.1", + "lru 0.16.4", + "mysql_common 0.37.3", "named_pipe", "pem", "percent-encoding", "rustls", "rustls-pemfile", - "socket2 0.5.10", + "socket2 0.6.4", "twox-hash", "url", "webpki", - "webpki-roots 0.26.11", + "webpki-roots 1.0.9", ] [[package]] @@ -8766,7 +8778,7 @@ checksum = "63c3512cf11487168e0e9db7157801bf5273be13055a9cc95356dc9e0035e49c" dependencies = [ "darling 0.20.11", "heck 0.5.0", - "num-bigint", + "num-bigint 0.4.6", "proc-macro-crate 3.3.0", "proc-macro-error2", "proc-macro2", @@ -8784,7 +8796,7 @@ checksum = "66f62cad7623a9cb6f8f64037f0c4f69c8db8e82914334a83c9788201c2c1bfa" dependencies = [ "darling 0.20.11", "heck 0.5.0", - "num-bigint", + "num-bigint 0.4.6", "proc-macro-crate 3.3.0", "proc-macro-error2", "proc-macro2", @@ -8796,36 +8808,32 @@ dependencies = [ [[package]] name = "mysql_async" -version = "0.35.1" +version = "0.37.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d14cf024116ba8fef4a7fec5abf0bd5de89b9fb29a7e55818a119ac5ec745077" +checksum = "3519e91b0d254ac1ffa495bc42053286cb2172ad7241d5b3b1b9f8a891f21ee2" dependencies = [ "bytes", - "crossbeam", + "crossbeam-queue", + "crossbeam-utils", "flate2", "futures-core", "futures-sink", "futures-util", "keyed_priority_queue", - "lru 0.12.5", - "mysql_common 0.34.1", - "pem", + "lru 0.18.2", + "mysql_common 0.37.3", "percent-encoding", - "pin-project", - "rand 0.8.5", + "rand 0.10.1", "rustls", - "rustls-pemfile", "serde", - "serde_json", - "socket2 0.5.10", + "socket2 0.6.4", "thiserror 2.0.17", "tokio", "tokio-rustls", "tokio-util", "twox-hash", "url", - "webpki", - "webpki-roots 0.26.11", + "webpki-roots 1.0.9", ] [[package]] @@ -8839,7 +8847,7 @@ dependencies = [ "bindgen", "bitflags 2.12.1", "bitvec", - "btoi", + "btoi 0.4.3", "byteorder", "bytes", "cc", @@ -8850,7 +8858,7 @@ dependencies = [ "frunk", "lazy_static", "mysql-common-derive 0.31.2", - "num-bigint", + "num-bigint 0.4.6", "num-traits", "rand 0.8.5", "regex", @@ -8870,35 +8878,58 @@ dependencies = [ [[package]] name = "mysql_common" -version = "0.34.1" +version = "0.37.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34a9141e735d5bb02414a7ac03add09522466d4db65bdd827069f76ae0850e58" +checksum = "0f27695f286b461da077b8c2f72f47feaa04ce3c3f9c0976257410e90e21208a" dependencies = [ "base64 0.22.1", "bitflags 2.12.1", - "btoi", + "btoi 0.5.0", "byteorder", "bytes", - "cc", - "cmake", "crc32fast", "flate2", - "lazy_static", + "getrandom 0.3.4", "mysql-common-derive 0.32.1", - "num-bigint", + "num-bigint 0.4.6", "num-traits", - "rand 0.8.5", "regex", "saturating", "serde", "serde_json", "sha1 0.10.6", "sha2 0.10.9", - "subprocess", - "thiserror 1.0.69", + "thiserror 2.0.17", "time", "uuid", - "zstd", +] + +[[package]] +name = "mysql_common" +version = "0.38.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d11189bb0e285f58b71744ee23aa8b943bd5120ef0877555813d2107374d3970" +dependencies = [ + "base64 0.22.1", + "bitflags 2.12.1", + "btoi 0.5.0", + "byteorder", + "bytes", + "crc32fast", + "flate2", + "getrandom 0.4.1", + "hex", + "mysql-common-derive 0.32.1", + "num-bigint 0.5.1", + "num-traits", + "regex", + "saturating", + "serde", + "serde_json", + "sha1 0.11.0", + "sha2 0.11.0", + "thiserror 2.0.17", + "uuid", ] [[package]] @@ -9094,7 +9125,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -9103,7 +9134,7 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" dependencies = [ - "num-bigint", + "num-bigint 0.4.6", "num-complex", "num-integer", "num-iter", @@ -9121,6 +9152,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-bigint-dig" version = "0.8.6" @@ -9205,7 +9246,7 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" dependencies = [ - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", ] @@ -9288,7 +9329,7 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77e878c846a8abae00dd069496dbe8751b16ac1c3d6bd2a7283a938e8228f90d" dependencies = [ - "proc-macro-crate 1.3.1", + "proc-macro-crate 3.3.0", "proc-macro2", "quote", "syn 2.0.117", @@ -10202,7 +10243,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "lz4_flex 0.13.1", - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", "object_store", @@ -11223,7 +11264,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.4.1", + "heck 0.5.0", "itertools 0.12.1", "log", "multimap", @@ -11243,8 +11284,8 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.4.1", - "itertools 0.10.5", + "heck 0.5.0", + "itertools 0.14.0", "log", "multimap", "once_cell", @@ -11292,7 +11333,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -11305,7 +11346,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -12300,7 +12341,7 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams 0.4.2", "web-sys", - "webpki-roots 1.0.1", + "webpki-roots 1.0.9", ] [[package]] @@ -12824,7 +12865,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -13364,7 +13405,7 @@ dependencies = [ "mime_guess", "moka", "mysql_async", - "mysql_common 0.34.1", + "mysql_common 0.38.2", "notify", "object-pool", "once_cell", @@ -13397,6 +13438,7 @@ dependencies = [ "rustls-pki-types", "serde", "serde_json", + "servers", "session", "simd-json", "simdutf8", @@ -13708,7 +13750,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -13737,7 +13779,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -14733,7 +14775,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.0.7", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -16414,14 +16456,14 @@ version = "0.26.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" dependencies = [ - "webpki-roots 1.0.1", + "webpki-roots 1.0.9", ] [[package]] name = "webpki-roots" -version = "1.0.1" +version = "1.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8782dd5a41a24eed3a4f40b606249b3e236ca61adf1f25ea4d45c73de122b502" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" dependencies = [ "rustls-pki-types", ] @@ -16481,7 +16523,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index c5e98d789a6..90dd1890222 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -83,7 +83,7 @@ log-query.workspace = true loki-proto.workspace = true metric-engine.workspace = true mime_guess = "2.0" -mysql_common = "0.34" +mysql_common = "0.38" notify.workspace = true object-pool = "0.5" once_cell.workspace = true @@ -157,12 +157,13 @@ common-test-util.workspace = true criterion = "0.5" json5 = "0.4" moka.workspace = true -mysql_async = { version = "0.35", default-features = false, features = [ +mysql_async = { version = "0.37", default-features = false, features = [ "default-rustls", ] } permutation = "0.4" rand.workspace = true serde_json.workspace = true +servers = { workspace = true, features = ["testing"] } session = { workspace = true, features = ["testing"] } table.workspace = true tempfile = "3.0.0" diff --git a/src/servers/tests/mysql/mysql_server_test.rs b/src/servers/tests/mysql/mysql_server_test.rs index e71369df09b..9ed5ac041e8 100644 --- a/src/servers/tests/mysql/mysql_server_test.rs +++ b/src/servers/tests/mysql/mysql_server_test.rs @@ -961,7 +961,7 @@ async fn test_query_prepared() -> Result<()> { .await { Err(mysql_async::Error::Server(e)) => assert_eq!( - "ERROR HY000 (1210): (InvalidArguments): Invalid prepare statement: Invalid SQL syntax: sql parser error: INTERVAL requires a unit after the literal value", + "ERROR 1210 (HY000): (InvalidArguments): Invalid prepare statement: Invalid SQL syntax: sql parser error: INTERVAL requires a unit after the literal value", e.to_string() ), _ => unreachable!(), diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 1456693704a..5376d7520ab 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -105,7 +105,7 @@ hex.workspace = true http.workspace = true itertools.workspace = true jsonb.workspace = true -mysql_async = { version = "0.35", default-features = false, features = [ +mysql_async = { version = "0.37", default-features = false, features = [ "time", "default-rustls-ring", ] } diff --git a/tests/runner/Cargo.toml b/tests/runner/Cargo.toml index 9801f6bc0bb..8d07c2ad1b3 100644 --- a/tests/runner/Cargo.toml +++ b/tests/runner/Cargo.toml @@ -19,7 +19,7 @@ datatypes = { workspace = true } flate2 = "1.0" hex = "0.4" local-ip-address = "0.6" -mysql = { version = "26", default-features = false, features = ["minimal", "rustls-tls"] } +mysql = { version = "28", default-features = false, features = ["minimal", "rustls-tls"] } num_cpus = "1.16" regex.workspace = true reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "stream"] } From 7539e60139657203ab0ae48fa92cff1b9ebb171f Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 13 Aug 2026 12:10:49 +0000 Subject: [PATCH 065/126] refactor: remove trivial tests (#8877) * refactor: remove trivial tests Signed-off-by: Lei, HUANG * fix: remove unused trace test import Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/auth/src/user_info.rs | 20 ------------------- .../src/instance/otlp/trace_ingest/tests.rs | 9 --------- src/meta-srv/src/gc/dropped.rs | 16 --------------- src/mito2/src/manifest/action.rs | 10 ---------- src/puffin/src/partial_reader/async.rs | 7 ------- .../src/tests/reconcile_table.rs | 3 --- 6 files changed, 65 deletions(-) diff --git a/src/auth/src/user_info.rs b/src/auth/src/user_info.rs index ab7c7f32f46..908f7bbb71a 100644 --- a/src/auth/src/user_info.rs +++ b/src/auth/src/user_info.rs @@ -223,26 +223,6 @@ mod tests { } } - #[test] - fn test_default_user_info_with_name() { - let user_info = DefaultUserInfo::with_name("test_user"); - assert_eq!(user_info.username(), "test_user"); - } - - #[test] - fn test_default_user_info_with_name_and_permission() { - let user_info = - DefaultUserInfo::with_name_and_permission("test_user", PermissionMode::ReadOnly); - assert_eq!(user_info.username(), "test_user"); - - // Cast to DefaultUserInfo to access permission_mode - let default_user = user_info - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(default_user.permission_mode, PermissionMode::ReadOnly); - } - #[test] fn test_user_info_as_any() { let user_info = DefaultUserInfo::with_name("test_user"); diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index 0f18a7194be..b3a0570db51 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -32,7 +32,6 @@ use servers::otlp::trace::SERVICE_NAME_COLUMN; use servers::otlp::trace::attributes::Attributes; use servers::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; use servers::otlp::trace::v1::{TraceBinaryType, TraceRetryColumn}; -use servers::query_handler::TraceIngestOutcome; use super::{ ChunkFailureReaction, Instance, TraceChunkRetry, TraceChunkSchemaState, TraceRequestSchema, @@ -101,14 +100,6 @@ fn test_classify_trace_prewrite_failure() { ); } -#[test] -fn test_add_trace_write_cost() { - let mut outcome = TraceIngestOutcome::default(); - Instance::add_trace_write_cost(&mut outcome, 3); - Instance::add_trace_write_cost(&mut outcome, 5); - assert_eq!(outcome.write_cost, 8); -} - #[test] fn test_finish_trace_failure_message() { let message = Instance::finish_trace_failure_message( diff --git a/src/meta-srv/src/gc/dropped.rs b/src/meta-srv/src/gc/dropped.rs index b4793ee5561..d0c7b1d2d64 100644 --- a/src/meta-srv/src/gc/dropped.rs +++ b/src/meta-srv/src/gc/dropped.rs @@ -284,19 +284,3 @@ impl<'a> DroppedRegionCollector<'a> { Ok(assignment) } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_dropped_region_info() { - let info = DroppedRegionInfo { - region_id: RegionId::new(1, 1), - table_id: 1, - dst_regions: HashSet::new(), - }; - assert_eq!(info.region_id, RegionId::new(1, 1)); - assert_eq!(info.table_id, 1); - } -} diff --git a/src/mito2/src/manifest/action.rs b/src/mito2/src/manifest/action.rs index 2cde150782b..2d0c17137a8 100644 --- a/src/mito2/src/manifest/action.rs +++ b/src/mito2/src/manifest/action.rs @@ -679,16 +679,6 @@ mod tests { .unwrap(); } - #[test] - fn test_region_manifest_builder() { - // TODO(ruihang): port this test case - } - - #[test] - fn test_encode_decode_region_checkpoint() { - // TODO(ruihang): port this test case - } - #[test] fn test_region_manifest_compatibility() { // Test deserializing RegionManifest from old schema where FileId is a UUID string diff --git a/src/puffin/src/partial_reader/async.rs b/src/puffin/src/partial_reader/async.rs index 0b9a5c54c71..0525628ea34 100644 --- a/src/puffin/src/partial_reader/async.rs +++ b/src/puffin/src/partial_reader/async.rs @@ -110,11 +110,4 @@ mod tests { // seeking past the portion returns an error assert!(reader.read(31..32).await.is_err()); } - - #[tokio::test] - async fn is_eof_returns_true_at_end_of_portion() { - let data: Vec = (0..100).collect(); - let reader = PartialReader::new(data, 10, 30); - let _ = reader.read(0..20).await.unwrap(); - } } diff --git a/tests-integration/src/tests/reconcile_table.rs b/tests-integration/src/tests/reconcile_table.rs index c2feecf30ed..d31845be461 100644 --- a/tests-integration/src/tests/reconcile_table.rs +++ b/tests-integration/src/tests/reconcile_table.rs @@ -541,9 +541,6 @@ async fn test_renamed_table() { check_output_stream(output.data, expected).await; } -#[tokio::test] -async fn test_rename_table() {} - async fn unwrap_err(output: OutputData) -> String { let error = match output { OutputData::Stream(stream) => collect_batches(stream).await.unwrap_err(), From 3c80df043a46f0511bd81a64a072ca18436efac3 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 13 Aug 2026 14:17:09 +0000 Subject: [PATCH 066/126] fix(mito2): keep deletion markers when compacting part of a window (#8872) `TwcsPicker::find_inputs` decides `filter_deleted` from the shape of the whole time window, but the compaction inputs are only a subset of it: `reduce_runs` and `merge_seq_files` narrow the selection down and the max input file num limit narrows it further. When a deletion marker lands in the compacted set while the file holding the row it masks stays behind, the marker is dropped from the output and the old row becomes visible again. Re-check the final selection against the rest of the window and stop filtering deleted rows whenever something left behind still overlaps the inputs. The check compares ranges inclusively, so it also covers files that share only a boundary timestamp: run detection treats those as non-overlapping, which is how a single timestamp delete file ends up in the same run as the file it deletes rows from. Signed-off-by: jeremyhi --- src/mito2/src/compaction/twcs.rs | 117 ++++++++++++++++++- src/mito2/src/engine/compaction_test.rs | 149 ++++++++++++++++++++++++ 2 files changed, 265 insertions(+), 1 deletion(-) diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index bbfc1171987..342426bbefe 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -169,7 +169,7 @@ impl TwcsPicker { let found_runs = sorted_runs.len(); // We only remove deletion markers if we found less than 2 runs and not in append mode. // because after compaction there will be no overlapping files. - let filter_deleted = + let mut filter_deleted = found_runs <= 2 && !self.append_mode && !window_has_overlap(files, windows); if found_runs == 0 { return (vec![], filter_deleted); @@ -210,6 +210,14 @@ impl TwcsPicker { ); } + // The inputs are only a subset of the window: `reduce_runs` and `merge_seq_files` both + // narrow the selection down and the file num limit above narrows it further. A deletion + // marker may only be dropped when every file that can hold a row it masks is compacted + // along with it, so re-check the final selection against what stays in the window. + if filter_deleted && overlaps_files_left_behind(&inputs, &files_to_merge) { + filter_deleted = false; + } + if inputs.len() > 1 { // If we have more than one group to compact. log_pick_result( @@ -227,6 +235,20 @@ impl TwcsPicker { } } +/// Returns whether any file group of `window_files` that is not part of `inputs` overlaps +/// `inputs`. Such a group keeps rows a deletion marker in `inputs` masks, so the compaction +/// must not filter deleted rows out. +/// +/// Overlapping is inclusive on both boundaries here: two files that only share a boundary +/// timestamp can still hold the same row. +fn overlaps_files_left_behind(inputs: &[FileGroup], window_files: &[FileGroup]) -> bool { + let picked: HashSet<_> = inputs.iter().flat_map(|fg| fg.file_ids()).collect(); + window_files + .iter() + .filter(|fg| !fg.file_ids().iter().any(|id| picked.contains(id))) + .any(|fg| inputs.iter().any(|input| input.overlap_inclusive(fg))) +} + #[allow(clippy::too_many_arguments)] fn log_pick_result( region_id: RegionId, @@ -1456,6 +1478,99 @@ mod tests { assert_eq!(output[0].inputs.len(), 32); } + #[tokio::test] + async fn test_limit_max_input_files_keeps_deletion_markers() { + common_telemetry::init_default_ut_logging(); + + // One large file group spanning the whole window plus 32 small ones nested inside + // it. That is exactly 2 runs, so the window on its own allows filtering deletions. + let mut files = vec![new_file_handle_with_size_and_sequence( + FileId::random(), + 0, + 3_000_000, + 0, + 1, + 1024 * 1024 * 1024, + )]; + files.extend((0..32).map(|idx: i64| { + new_file_handle_with_size_and_sequence( + FileId::random(), + (idx + 1) * 10_000, + (idx + 1) * 10_000 + 1_000, + 0, + (idx + 2) as u64, + 1024, + ) + })); + + let windows = assign_to_windows(files.iter(), 3600); + + let picker = TwcsPicker { + trigger_file_num: 4, + time_window_seconds: Some(3600), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let active_window = find_latest_window_in_seconds(files.iter(), 3600); + let output = picker + .build_output_with_time_range(RegionId::from_u64(123), windows, active_window, None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + // The input file num limit picks the smallest groups first, so the large group is + // left behind while the small groups that overlap it are compacted. + assert_eq!(32, output[0].inputs.len()); + assert!( + !output[0].filter_deleted, + "deletion markers must be kept once the file num limit drops files they may mask" + ); + } + + #[tokio::test] + async fn test_limit_max_input_files_still_filters_without_overlap() { + common_telemetry::init_default_ut_logging(); + + // 40 file groups with disjoint time ranges, i.e. a single run. Nothing the file num + // limit leaves behind can hold a row masked by a deletion marker we compact. + let files: Vec<_> = (0..40i64) + .map(|idx| { + new_file_handle_with_size_and_sequence( + FileId::random(), + (idx + 1) * 10_000, + (idx + 1) * 10_000 + 1_000, + 0, + (idx + 1) as u64, + 1024, + ) + }) + .collect(); + + let windows = assign_to_windows(files.iter(), 3600); + + let picker = TwcsPicker { + trigger_file_num: 4, + time_window_seconds: Some(3600), + max_output_file_size: Some(1024 * 1024 * 1024), + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let active_window = find_latest_window_in_seconds(files.iter(), 3600); + let output = picker + .build_output_with_time_range(RegionId::from_u64(123), windows, active_window, None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + assert_eq!(32, output[0].inputs.len()); + assert!(output[0].filter_deleted); + } + #[tokio::test] async fn test_newer_windows_have_priority() { let older_file_ids = [FileId::random(), FileId::random()]; diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index 4404b8d46a1..1fc3a8f785d 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -1342,6 +1342,155 @@ async fn test_compaction_region_with_overlapping_delete_all_with_format(flat_for assert!(vec.is_empty()); } +#[tokio::test] +async fn test_compaction_input_limit_keeps_rows_deleted() { + test_compaction_input_limit_keeps_rows_deleted_with_format(false).await; + test_compaction_input_limit_keeps_rows_deleted_with_format(true).await; +} + +/// Creates a region that only compacts when asked, so that a test can build an exact file +/// layout with `put_and_flush` and `delete_and_flush`. +async fn env_for_manual_compaction( + env: &mut TestEnv, + region_id: RegionId, + flat_format: bool, +) -> (MitoEngine, Vec) { + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }) + .await; + + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let request = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.time_window", "1h") + .build(); + let column_schemas = request + .column_metadatas + .iter() + .map(column_metadata_to_column_schema) + .collect::>(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + (engine, column_schemas) +} + +/// The picker caps a compaction at 32 input files and drops the largest file groups to get +/// there. A deletion marker among the picked files must not be filtered out while the file +/// holding the rows it masks stays behind, otherwise those rows become visible again. +async fn test_compaction_input_limit_keeps_rows_deleted_with_format(flat_format: bool) { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let region_id = RegionId::new(1, 1); + let (engine, column_schemas) = + env_for_manual_compaction(&mut env, region_id, flat_format).await; + + // One large file spanning the whole time window. + put_and_flush(&engine, region_id, &column_schemas, 0..3000).await; + // Deletes 6 rows of that file. The markers land in a tiny file that overlaps it. + delete_and_flush(&engine, region_id, &column_schemas, 10..16).await; + // 31 more tiny files that overlap the large one but not each other, so the window holds + // 33 file groups forming 2 runs. + for i in 2..33 { + put_and_flush(&engine, region_id, &column_schemas, i * 10..i * 10 + 6).await; + } + + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!( + 33, + scanner.num_files(), + "unexpected files: {:?}", + scanner.file_ids() + ); + + compact(&engine, region_id).await; + + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + // The 32 tiny files are merged into one; the large file exceeds the input file num limit + // and is left behind. + assert_eq!( + 2, + scanner.num_files(), + "unexpected files: {:?}", + scanner.file_ids() + ); + let stream = scanner.scan().await.unwrap(); + let vec = collect_stream_ts(stream).await; + assert!( + !(10..16).any(|ts| vec.contains(&(ts * 1000))), + "deleted rows are visible again after compaction" + ); + assert_eq!(2994, vec.len()); +} + +#[tokio::test] +async fn test_compaction_of_part_of_a_run_keeps_rows_deleted() { + test_compaction_of_part_of_a_run_keeps_rows_deleted_with_format(false).await; + test_compaction_of_part_of_a_run_keeps_rows_deleted_with_format(true).await; +} + +/// A run is supposed to hold no overlapping files, which is what lets `merge_seq_files` +/// compact part of a run and still filter deleted rows. Run detection compares time ranges +/// exclusively though, so a file covering a single timestamp lands in the same run as the +/// file it deletes rows from. The rows must stay deleted when only one of the two is picked. +async fn test_compaction_of_part_of_a_run_keeps_rows_deleted_with_format(flat_format: bool) { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let region_id = RegionId::new(1, 1); + let (engine, column_schemas) = + env_for_manual_compaction(&mut env, region_id, flat_format).await; + + // One large file spanning the whole time window. + put_and_flush(&engine, region_id, &column_schemas, 0..3000).await; + // Deletes a single row of that file, so the marker lands in a file covering one timestamp. + delete_and_flush(&engine, region_id, &column_schemas, 1..2).await; + // 31 more single row files, each holding another key at its own timestamp. + for ts in 1..32 { + let rows = Rows { + schema: column_schemas.clone(), + rows: build_rows_for_key("b", ts * 10, ts * 10 + 1, 0), + }; + put_rows(&engine, region_id, rows).await; + flush(&engine, region_id).await; + } + + compact(&engine, region_id).await; + + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + let stream = scanner.scan().await.unwrap(); + let vec = collect_stream_ts(stream).await; + assert!( + !vec.contains(&1000), + "deleted row is visible again after compaction" + ); + assert_eq!(3030, vec.len()); +} + // For issue https://github.com/GreptimeTeam/greptimedb/issues/3633 #[tokio::test] async fn test_readonly_during_compaction() { From f692ac32f3f92584ee81abe04cd7fc0e669f3a3d Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 14 Aug 2026 04:21:46 +0000 Subject: [PATCH 067/126] test(mito2): isolate sequence publication barrier (#8876) * test(mito2): isolate sequence publication barrier Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(mito2): trim sequence barrier test noise Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/region_write_ctx.rs | 84 +++++++++++----------------- src/mito2/src/worker/handle_write.rs | 47 ++-------------- 2 files changed, 39 insertions(+), 92 deletions(-) diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index f64f9810f03..400f7db1fc4 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -199,7 +199,6 @@ impl RegionWriteCtx { &self.version } - /// Returns the version control of the region. #[cfg(test)] pub(crate) fn version_control(&self) -> &VersionControlRef { &self.version_control @@ -327,7 +326,7 @@ impl RegionWriteCtx { return; } #[cfg(test)] - test_hooks::pause_before_bulk_install(self.region_id).await; + test_hooks::pause_before_bulk_install(self.region_id, &self.version_control).await; let _timer = metrics::REGION_WORKER_HANDLE_WRITE_ELAPSED .with_label_values(&["write_bulk"]) .start_timer(); @@ -375,20 +374,10 @@ impl RegionWriteCtx { } } - /// Publishes the sequences and entry id assigned by this context to the - /// region's committed watermark. - /// - /// Must be called after both [`write_memtable`](Self::write_memtable) and - /// [`write_bulk`](Self::write_bulk) have completed, so the committed - /// sequence never covers rows that are not yet physically installed in the - /// memtable (bulk part sequences are assigned in [`push_bulk`](Self::push_bulk) - /// before installation). Since we store the last sequence and entry id in - /// the region, we decrease `next_sequence` and `next_entry_id` by 1. - /// - /// If the write operation failed (e.g. the WAL entry could not be built), - /// no rows were installed and nothing must be published: advancing the - /// committed watermark here would make it cover rows that were never - /// written. + /// Publishes the assigned sequences and entry id to the region's committed + /// watermark. Only call after both [`write_memtable`](Self::write_memtable) + /// and [`write_bulk`](Self::write_bulk) have completed; a failed context + /// must not publish. pub(crate) fn publish_sequence_and_entry_id(&self) { if self.failed { return; @@ -407,25 +396,20 @@ pub(crate) mod test_hooks { use store_api::storage::RegionId; use tokio::sync::watch; - /// Channels of an armed bulk-install barrier. - /// - /// `reached` signals that a bulk write paused between ordinary-memtable - /// handling and bulk installation; `release` unblocks it. Dropping the - /// senders (by disarming the barrier) also unblocks paused writes because - /// their `wait_for` fails on a closed channel. + use crate::region::version::VersionControlRef; + + /// Channels of an armed bulk-install barrier; dropping the senders (by + /// disarming) unblocks writes paused on it. struct ActiveBarrier { id: u64, - /// Only bulk writes for this region pause at the barrier; writes for - /// other regions pass through untouched, so concurrently running tests - /// can never satisfy (or be blocked by) each other's barrier. + /// Only bulk writes for this region on this version control (Arc + /// identity) pause at the barrier. target_region_id: RegionId, + target_version_control: VersionControlRef, reached: watch::Sender, release: watch::Sender, } - /// The currently armed barrier, if any. Wrapped in a `Mutex` so a test can - /// arm a fresh barrier after a previous one was released (or dropped), and - /// so a barrier can never leak past the test that owns it. static ACTIVE_BARRIER: Mutex> = Mutex::new(None); static NEXT_BARRIER_ID: AtomicU64 = AtomicU64::new(1); @@ -437,12 +421,8 @@ pub(crate) mod test_hooks { .unwrap_or_else(|poisoned| poisoned.into_inner()) } - /// RAII guard for the bulk-install barrier. - /// - /// The guard owns the armed barrier. Releasing it — or dropping the guard, - /// e.g. when a test panics — unblocks any write paused at the barrier and - /// disarms it, so a paused write can never hang and later tests can arm a - /// fresh barrier. + /// RAII guard: releasing (or dropping) it unblocks a paused write and + /// disarms the barrier. pub(crate) struct BulkInstallBarrier { id: u64, reached_rx: watch::Receiver, @@ -451,20 +431,17 @@ pub(crate) mod test_hooks { } impl BulkInstallBarrier { - /// Waits until a bulk write paused at the barrier. pub(crate) async fn wait_until_reached(&mut self) { if !*self.reached_rx.borrow() { let _ = self.reached_rx.wait_for(|reached| *reached).await; } } - /// Unblocks the paused write and disarms the barrier. pub(crate) fn release(&mut self) { if self.released { return; } self.released = true; - // Flip the release value so writes paused on this barrier proceed. let _ = self.release_tx.send(true); disarm_barrier(self.id); } @@ -476,12 +453,12 @@ pub(crate) mod test_hooks { } } - /// Arms the bulk-install barrier for `target_region_id` and returns a - /// guard that owns it. - /// - /// Any previously armed barrier is replaced: its senders are dropped, which - /// unblocks writes that were paused on it instead of letting them hang. - pub(crate) fn arm_bulk_install_barrier(target_region_id: RegionId) -> BulkInstallBarrier { + /// Arms the bulk-install barrier and returns the owning guard; any + /// previously armed barrier is replaced. + pub(crate) fn arm_bulk_install_barrier( + target_region_id: RegionId, + target_version_control: VersionControlRef, + ) -> BulkInstallBarrier { let (reached_tx, reached_rx) = watch::channel(false); let (release_tx, _release_rx) = watch::channel(false); let id = NEXT_BARRIER_ID.fetch_add(1, Ordering::Relaxed); @@ -489,6 +466,7 @@ pub(crate) mod test_hooks { *active = Some(ActiveBarrier { id, target_region_id, + target_version_control, reached: reached_tx, release: release_tx.clone(), }); @@ -500,7 +478,6 @@ pub(crate) mod test_hooks { } } - /// Removes the barrier with `id` from the statics if it is still active. fn disarm_barrier(id: u64) { let mut active = lock_active_barrier(); if active.as_ref().is_some_and(|barrier| barrier.id == id) { @@ -508,20 +485,27 @@ pub(crate) mod test_hooks { } } - /// Pauses a bulk write for `region_id` before installing its parts until - /// the test releases the barrier (or it is disarmed). Bulk writes for - /// other regions return immediately. - pub(crate) async fn pause_before_bulk_install(region_id: RegionId) { + /// Pauses a bulk write before installing its parts until the barrier is + /// released or disarmed. + pub(crate) async fn pause_before_bulk_install( + region_id: RegionId, + version_control: &VersionControlRef, + ) { let (reached_tx, release_rx) = { let active = lock_active_barrier(); match active.as_ref() { - Some(barrier) if barrier.target_region_id == region_id => { + Some(barrier) + if barrier.target_region_id == region_id + && std::sync::Arc::ptr_eq( + &barrier.target_version_control, + version_control, + ) => + { (barrier.reached.clone(), barrier.release.subscribe()) } _ => return, } }; - // Signal that a bulk write reached the pause point. let _ = reached_tx.send(true); let mut release_rx = release_rx; if !*release_rx.borrow() { diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 0c2eecdea20..96940756248 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -128,9 +128,6 @@ impl RegionWorkerLoop { let mut region_ctx = region_ctxs.into_values().next().unwrap(); region_ctx.write_memtable().await; region_ctx.write_bulk().await; - // Publish only after all rows (including bulk parts) are - // physically installed, so a scan opening on the committed - // sequence can never bind H to invisible rows. region_ctx.publish_sequence_and_entry_id(); put_rows += region_ctx.put_num; delete_rows += region_ctx.delete_num; @@ -142,8 +139,6 @@ impl RegionWorkerLoop { common_runtime::spawn_global(async move { region_ctx.write_memtable().await; region_ctx.write_bulk().await; - // The spawned task owns the moved ctx, so publish - // inside the task after the memtable writes. region_ctx.publish_sequence_and_entry_id(); (region_ctx.put_num, region_ctx.delete_num) }) @@ -763,10 +758,6 @@ mod tests { use crate::test_util::ts_ms_value; use crate::test_util::version_util::VersionControlBuilder; - /// Creates a bulk part with `num_rows` rows. The schema carries the - /// builder metadata's columns (`tag_0` primary key + `ts` time index) so - /// the bulk-install conversion succeeds; `timestamp_index` points at the - /// `ts` column. fn new_bulk_part(num_rows: i64) -> BulkPart { let schema = Arc::new(Schema::new(vec![ Field::new("tag_0", DataType::Utf8, true), @@ -887,7 +878,6 @@ mod tests { } } - /// Creates a write context for `region_id` with one pending mutation of one row. fn new_region_ctx( region_id: RegionId, ) -> (RegionWriteCtx, oneshot::Receiver>) { @@ -940,16 +930,13 @@ mod tests { region_ctxs.insert(ok_region, ctx); let entry_id = region_ctxs[&ok_region].next_entry_id(); - // The failed region must not fail the batch or panic the worker. + // The failed region must not fail the batch. assert!(write_wal(&wal, &mut region_ctxs).await); assert!(region_ctxs[&failing_region].is_failed()); assert!(!region_ctxs[&ok_region].is_failed()); assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); - // Run the memtable and publication phase the worker runs after `write_wal`. - // The failed region installed no rows, so its committed sequence must not - // advance; the successful sibling region advances normally. for region_ctx in region_ctxs.values_mut() { region_ctx.write_memtable().await; region_ctx.write_bulk().await; @@ -969,30 +956,15 @@ mod tests { .committed_sequence() ); - // Waiters of the failed region get the error while others get the result. drop(region_ctxs); assert!(failing_rx.await.unwrap().is_err()); assert_eq!(1, ok_rx.await.unwrap().unwrap()); } - // The committed sequence must not be published before the bulk part's rows - // are physically installed in the memtable: publication happens strictly - // after `write_bulk` returns. Armed with the bulk-install test barrier, - // this deterministically pauses the worker's memtable phase between the - // ordinary-memtable handling and the bulk installation and asserts the - // region's committed sequence is still its initial value (0) — read via - // the `version_control()` accessor, never through a scanner API. After the - // barrier is released, the committed sequence must advance to cover the - // bulk rows. #[tokio::test] async fn test_bulk_write_sequence_not_committed_before_install_worker_level() { let region_id = RegionId::new(1, 1); let version_control = Arc::new(VersionControlBuilder::new().build()); - assert_eq!( - 0, - version_control.committed_sequence(), - "the builder must start at sequence 0" - ); let mut region_ctxs = HashMap::new(); let mut ctx = RegionWriteCtx::new( @@ -1002,21 +974,18 @@ mod tests { None, ); let (tx, rx) = oneshot::channel(); - // Push a 3-row bulk part without an explicit sequence, exactly like the - // worker's `process_bulk_requests`. assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(3), None)); region_ctxs.insert(region_id, ctx); - // Run the WAL phase the worker runs before the memtable phase. let wal = Wal::new(Arc::new(MockLogStore::default())); assert!(write_wal(&wal, &mut region_ctxs).await); assert!(!region_ctxs[®ion_id].is_failed()); - // Arm the bulk-install barrier: `write_bulk` pauses right before the - // parts are physically installed into the memtable. - let mut barrier = crate::region_write_ctx::test_hooks::arm_bulk_install_barrier(region_id); + let mut barrier = crate::region_write_ctx::test_hooks::arm_bulk_install_barrier( + region_id, + version_control.clone(), + ); - // Run the worker's memtable and publication phases in the background. let write_handle = tokio::spawn(async move { let mut region_ctx = region_ctxs.remove(®ion_id).unwrap(); region_ctx.write_memtable().await; @@ -1024,7 +993,6 @@ mod tests { region_ctx.publish_sequence_and_entry_id(); }); - // Wait until the write paused at the barrier: deterministic, no sleeps. tokio::time::timeout( std::time::Duration::from_secs(10), barrier.wait_until_reached(), @@ -1032,16 +1000,12 @@ mod tests { .await .expect("bulk write never reached the install barrier"); - // The committed sequence must not have advanced yet: publication - // happens strictly after the bulk part is installed. assert_eq!( 0, version_control.committed_sequence(), "committed sequence leaked before the bulk part was installed" ); - // Release the barrier: the bulk part installs and the committed - // sequence advances to cover all 3 bulk rows. barrier.release(); write_handle.await.expect("bulk write should complete"); assert_eq!( @@ -1050,7 +1014,6 @@ mod tests { "committed sequence must cover the installed bulk rows" ); - // The bulk waiter is notified with the number of installed rows. assert_eq!(3, rx.await.unwrap().unwrap()); } From b1263fc65ffb95a80353177f9fef01f551879f68 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Fri, 14 Aug 2026 06:15:42 +0000 Subject: [PATCH 068/126] perf: optimize Prometheus remote write v2 decoding (#8873) * test: add bench for prom decode Signed-off-by: shuiyisong * chore: merge v2 decode manually and add to bench Signed-off-by: shuiyisong * refactor: update v2 path Signed-off-by: shuiyisong * refactor: update v2 path Signed-off-by: shuiyisong * fix: use constant Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/servers/Cargo.toml | 5 + src/servers/benches/prom_decode.rs | 247 ++++- src/servers/src/http/prom_store.rs | 33 +- src/servers/src/prom_remote_write/README.md | 33 +- src/servers/src/prom_remote_write/v2.rs | 998 +++++++++++++++--- src/servers/src/row_writer.rs | 45 +- .../tests/prom_remote_write_v2_test.rs | 2 +- 7 files changed, 1147 insertions(+), 216 deletions(-) diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 90dd1890222..d64a1bd5a5d 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -180,6 +180,11 @@ common-version.workspace = true name = "bench_prom" harness = false +[[bench]] +name = "prom_decode" +harness = false +required-features = ["testing"] + [[bench]] name = "to_http_output" harness = false diff --git a/src/servers/benches/prom_decode.rs b/src/servers/benches/prom_decode.rs index e9db0ad7716..06b0d02e940 100644 --- a/src/servers/benches/prom_decode.rs +++ b/src/servers/benches/prom_decode.rs @@ -12,21 +12,137 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::time::Duration; -use api::prom_store::remote::WriteRequest; -use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; +use api::greptime_proto::io::prometheus::write::v2 as write_v2; +use api::prom_store::remote::{self as write_v1, WriteRequest}; +use criterion::{ + BatchSize, BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main, +}; use prost::Message; use servers::prom_remote_write::decode::{PromSeriesProcessor, PromWriteRequest}; +use servers::prom_remote_write::v2::test_util as remote_write_v2; use servers::prom_remote_write::validation::{PromValidationMode, validate_label_name}; use servers::prom_store::to_grpc_row_insert_requests; -fn bench_decode_prom_request(c: &mut Criterion) { +fn load_fixture_v1_bytes() -> Vec { let mut d = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")); d.push("benches"); d.push("write_request.pb.data"); + std::fs::read(d).expect("read write_request.pb.data fixture") +} - let data = std::fs::read(d).unwrap(); +/// Convert a PRW v1 `WriteRequest` into an equivalent PRW v2 `Request`. +/// +/// Labels are interned into the v2 symbol table so payloads with repeated label +/// names/values exercise the main on-wire advantage of v2. +fn write_request_v1_to_v2(v1: &WriteRequest) -> write_v2::Request { + let mut symbols = vec![String::new()]; + let mut symbol_index: HashMap = HashMap::new(); + symbol_index.insert(String::new(), 0); + + let mut intern = |value: &str| -> u32 { + if let Some(&idx) = symbol_index.get(value) { + return idx; + } + let idx = symbols.len() as u32; + symbols.push(value.to_string()); + symbol_index.insert(value.to_string(), idx); + idx + }; + + let timeseries = v1 + .timeseries + .iter() + .map(|series| { + let mut labels: Vec<&write_v1::Label> = series.labels.iter().collect(); + // PRW v2 requires lexicographically sorted label pairs. + labels.sort_by(|a, b| a.name.cmp(&b.name).then_with(|| a.value.cmp(&b.value))); + + let mut labels_refs = Vec::with_capacity(labels.len() * 2); + for label in labels { + labels_refs.push(intern(&label.name)); + labels_refs.push(intern(&label.value)); + } + + write_v2::TimeSeries { + labels_refs, + samples: series + .samples + .iter() + .map(|sample| write_v2::Sample { + value: sample.value, + timestamp: sample.timestamp, + start_timestamp: 0, + }) + .collect(), + histograms: Vec::new(), + exemplars: Vec::new(), + metadata: None, + } + }) + .collect(); + + write_v2::Request { + symbols, + timeseries, + } +} + +/// Synthetic sample workload with controllable label reuse. +/// +/// `shared_labels` keeps common label names/values across series (v2-friendly). +/// Each series also gets a unique `series_id` label so cardinality stays high. +fn generate_sample_workload( + series_count: usize, + samples_per_series: usize, + shared_label_count: usize, +) -> (WriteRequest, write_v2::Request) { + let shared_labels: Vec = (0..shared_label_count) + .map(|i| write_v1::Label { + name: format!("label_{i}"), + value: format!("value_{}", i % 16), + }) + .collect(); + + let timeseries = (0..series_count) + .map(|series_idx| { + let mut labels = Vec::with_capacity(shared_label_count + 2); + labels.push(write_v1::Label { + name: "__name__".to_string(), + value: format!("metric_{}", series_idx % 32), + }); + labels.extend(shared_labels.iter().cloned()); + labels.push(write_v1::Label { + name: "series_id".to_string(), + value: series_idx.to_string(), + }); + + write_v1::TimeSeries { + labels, + samples: (0..samples_per_series) + .map(|sample_idx| write_v1::Sample { + value: (series_idx * 1000 + sample_idx) as f64, + timestamp: 1_700_000_000_000 + sample_idx as i64 * 1_000, + }) + .collect(), + exemplars: Vec::new(), + histograms: Vec::new(), + } + }) + .collect(); + + let v1 = WriteRequest { + timeseries, + metadata: Vec::new(), + }; + let v2 = write_request_v1_to_v2(&v1); + (v1, v2) +} + +fn bench_decode_prom_request(c: &mut Criterion) { + let data = load_fixture_v1_bytes(); let mut group = c.benchmark_group("decode_prom_request"); group.measurement_time(Duration::from_secs(3)); @@ -64,6 +180,128 @@ fn bench_decode_prom_request(c: &mut Criterion) { group.finish(); } +/// Compare Prometheus remote-write v1 vs v2 decoding efficiency. +/// +/// Two layers are measured on equivalent logical payloads: +/// 1. protobuf decode only (`WriteRequest` / v2 `Request`) +/// 2. full Greptime path to row-insert requests +/// +/// Workloads: +/// - `fixture`: existing `write_request.pb.data` converted v1 -> v2 +/// - `synthetic_*`: generated series with shared labels (symbol-table friendly) +#[allow(clippy::print_stderr)] +fn bench_prom_v1_vs_v2_decode(c: &mut Criterion) { + let fixture_v1_bytes = load_fixture_v1_bytes(); + let fixture_v1 = WriteRequest::decode(fixture_v1_bytes.as_slice()).unwrap(); + let fixture_v2 = write_request_v1_to_v2(&fixture_v1); + let fixture_v2_bytes = fixture_v2.encode_to_vec(); + + let synthetic = [ + ( + "synthetic_1k_series_1_sample", + generate_sample_workload(1_000, 1, 8), + ), + ( + "synthetic_1k_series_10_samples", + generate_sample_workload(1_000, 10, 8), + ), + ( + "synthetic_5k_series_1_sample", + generate_sample_workload(5_000, 1, 8), + ), + ]; + + let mut workloads: Vec<(&str, Vec, Vec)> = + vec![("fixture", fixture_v1_bytes, fixture_v2_bytes)]; + for (name, (v1, v2)) in &synthetic { + workloads.push((*name, v1.encode_to_vec(), v2.encode_to_vec())); + } + + eprintln!("\nprom remote-write v1 vs v2 payload sizes:"); + for (name, v1_bytes, v2_bytes) in &workloads { + let ratio = v2_bytes.len() as f64 / v1_bytes.len() as f64; + eprintln!( + " {name}: v1={} B, v2={} B, v2/v1={ratio:.3}", + v1_bytes.len(), + v2_bytes.len() + ); + } + + // --- protobuf decode only --- + { + let mut group = c.benchmark_group("prom_rw_protobuf_decode"); + group.measurement_time(Duration::from_secs(3)); + + for (name, v1_bytes, v2_bytes) in &workloads { + group.throughput(Throughput::Bytes(v1_bytes.len() as u64)); + group.bench_with_input(BenchmarkId::new("v1", name), v1_bytes, |b, bytes| { + b.iter(|| { + black_box(WriteRequest::decode(black_box(bytes.as_slice())).unwrap()); + }); + }); + + group.throughput(Throughput::Bytes(v2_bytes.len() as u64)); + group.bench_with_input(BenchmarkId::new("v2", name), v2_bytes, |b, bytes| { + b.iter(|| { + black_box(write_v2::Request::decode(black_box(bytes.as_slice())).unwrap()); + }); + }); + } + + group.finish(); + } + + // --- full path: protobuf -> Greptime row inserts --- + { + let mut group = c.benchmark_group("prom_rw_decode_to_rows"); + group.sample_size(50); + group.measurement_time(Duration::from_secs(5)); + + for (name, v1_bytes, v2_bytes) in &workloads { + group.throughput(Throughput::Bytes(v1_bytes.len() as u64)); + group.bench_with_input(BenchmarkId::new("v1_custom", name), v1_bytes, |b, bytes| { + let mut prom_request = PromWriteRequest::default(); + let mut processor = PromSeriesProcessor::default_processor(); + b.iter_batched( + || bytes.clone(), + |bytes| { + prom_request + .decode(black_box(bytes), PromValidationMode::Strict, &mut processor) + .unwrap(); + let rows = prom_request.as_row_insert_requests(); + black_box(&rows); + }, + BatchSize::LargeInput, + ); + }); + + // Baseline: stock prost WriteRequest + existing converter. + group.throughput(Throughput::Bytes(v1_bytes.len() as u64)); + group.bench_with_input(BenchmarkId::new("v1_prost", name), v1_bytes, |b, bytes| { + b.iter(|| { + let request = WriteRequest::decode(black_box(bytes.as_slice())).unwrap(); + black_box(to_grpc_row_insert_requests(&request).unwrap()); + }); + }); + + group.throughput(Throughput::Bytes(v2_bytes.len() as u64)); + group.bench_with_input(BenchmarkId::new("v2_manual", name), v2_bytes, |b, bytes| { + b.iter(|| { + black_box( + remote_write_v2::decode_uncompressed_write_requests( + black_box(bytes.as_slice()), + true, + ) + .unwrap(), + ); + }); + }); + } + + group.finish(); + } +} + /// Benchmark comparing UTF-8 string validation (`decode_string`) vs /// direct byte-level Prometheus label name validation (`decode_label_name`). fn bench_label_name_validation(c: &mut Criterion) { @@ -155,6 +393,7 @@ fn bench_utf8_validation(c: &mut Criterion) { criterion_group!( benches, bench_decode_prom_request, + bench_prom_v1_vs_v2_decode, bench_label_name_validation, bench_utf8_validation ); diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index ac4f39628d1..ed960264b0d 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -48,7 +48,7 @@ use crate::http::header::{ use crate::pending_rows_batcher::PendingRowsBatcher; use crate::prom_remote_write::decode::PromSeriesProcessor; use crate::prom_remote_write::decode_remote_write_request; -use crate::prom_remote_write::v2::{decode_remote_write_v2_request, into_write_requests}; +use crate::prom_remote_write::v2::decode_remote_write_v2; use crate::prom_remote_write::validation::PromValidationMode; use crate::prom_store::snappy_decompress; use crate::query_handler::{PipelineHandlerRef, PromStoreProtocolHandlerRef, PromStoreResponse}; @@ -235,23 +235,11 @@ async fn remote_write_v2( let (db, mut query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V2_VERSION); - let request = match decode_remote_write_v2_request(is_zstd, body) { - Ok(request) => request, - Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), - }; - if !experimental_enable_prometheus_native_histogram && request_has_native_histograms(&request) { - return Ok(remote_write_v2_error_response( - error::InvalidPromRemoteRequestSnafu { - msg: "prometheus remote write v2 native histogram ingestion is experimental; set prom_store.experimental_enable_prometheus_native_histogram = true to enable it" - .to_string(), - } - .build(), - 0, - 0, - 0, - )); - } - let req = match into_write_requests(request) { + let req = match decode_remote_write_v2( + is_zstd, + body, + experimental_enable_prometheus_native_histogram, + ) { Ok(req) => req, Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), }; @@ -302,15 +290,6 @@ async fn remote_write_v2( Ok((StatusCode::NO_CONTENT, headers).into_response()) } -fn request_has_native_histograms( - request: &api::greptime_proto::io::prometheus::write::v2::Request, -) -> bool { - request - .timeseries - .iter() - .any(|series| !series.histograms.is_empty()) -} - fn vm_proto_version_response(params: &RemoteWriteQuery) -> Option { params .get_vm_proto_version diff --git a/src/servers/src/prom_remote_write/README.md b/src/servers/src/prom_remote_write/README.md index d21eead1901..9e86a8753dd 100644 --- a/src/servers/src/prom_remote_write/README.md +++ b/src/servers/src/prom_remote_write/README.md @@ -11,28 +11,33 @@ Remote write v2 enters through `remote_write_v2` in ```mermaid flowchart TD A["HTTP /v1/prometheus/write"] --> B["remote_write_v2"] - B --> C["decode_remote_write_v2_request"] - C --> D["into_write_requests"] + B --> C["decode_remote_write_v2"] + C --> D["borrow symbols and time-series bytes"] + D --> E["decode leaf messages and build rows"] - D --> E["samples ContextReq"] - D --> F["histograms ContextReq"] + E --> F["samples ContextReq"] + E --> G["histograms ContextReq"] - E --> G["write_prometheus_rows_with_progress"] - G --> H["metric engine / pending batcher when enabled"] + F --> H["write_prometheus_rows_with_progress"] + H --> I["metric engine / pending batcher when enabled"] - F --> I["histogram ContextOpt"] - I --> J["write_prometheus_rows_with_progress"] - J --> K["same metric-engine flag as samples, no batcher"] - K --> L["table: "] + G --> J["histogram ContextOpt"] + J --> K["write_prometheus_rows_with_progress"] + K --> L["same metric-engine flag as samples, no batcher"] + L --> M["table: "] - L --> M["field: configured native-histogram Struct"] - M --> N["struct children: counts, spans, buckets, sum, schema"] - H --> P["written headers and counters"] - K --> P + M --> N["field: configured native-histogram Struct"] + N --> O["struct children: counts, spans, buckets, sum, schema"] + I --> P["written headers and counters"] + L --> P ``` The conversion step splits one v2 request into two `ContextReq`s: +The `decode` codec metric covers decompression and the borrowed request +envelope. The `convert` metric covers per-series scans, leaf prost decoding, +validation, and row construction. + - samples keep the existing sample table name and can use the metric-engine physical table path and pending rows batcher; - native histograms keep the existing metric table name. diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 5a25feb566a..31da7ac1f35 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -17,19 +17,24 @@ use std::collections::hash_map::Entry; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; use api::greptime_proto::io::prometheus::write::v2::{ - BucketSpan, Histogram, Request, Sample, TimeSeries, metadata, + BucketSpan, Exemplar, Histogram, Metadata, Sample, metadata, }; #[cfg(test)] -use api::greptime_proto::io::prometheus::write::v2::{Exemplar, Metadata}; +use api::greptime_proto::io::prometheus::write::v2::{Request, TimeSeries}; use api::helper::ColumnDataTypeWrapper; use api::v1::value::ValueData; -use api::v1::{ColumnSchema, ListValue, RowInsertRequest, Rows, SemanticType, Value}; -use bytes::Bytes; +use api::v1::{ + ColumnDataType, ColumnSchema, ListValue, RowInsertRequest, Rows, SemanticType, Value, +}; +use bytes::{Buf, Bytes}; use common_grpc::precision::Precision; use common_query::native_histogram::*; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use pipeline::{ContextOpt, ContextReq}; -use prost::Message; +use prost::encoding::{ + DecodeContext, WireType, decode_key, decode_varint, message, skip_field, uint32, +}; +use prost::{DecodeError, Message}; use snafu::{OptionExt, ResultExt, ensure}; use table::requests::{ METADATA_QUALITY_DECLARED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_TYPE, @@ -52,23 +57,60 @@ use crate::semantic::{ openmetrics_unit_to_ucum, }; -type PromTags = Vec<(String, String)>; -type ResolvedSeriesLabels = (PromCtx, String, PromTags); +type PromTags<'a> = Vec<(&'a str, String)>; +type ResolvedSeriesLabels<'a> = (PromCtx, String, PromTags<'a>); const MAX_REMOTE_WRITE_V2_SCHEMA: i32 = 8; const MAX_REDUCIBLE_REMOTE_WRITE_V2_SCHEMA: i32 = 52; +const TIME_SERIES_LABELS_REFS_TAG: u32 = 1; +const TIME_SERIES_SAMPLES_TAG: u32 = 2; +const TIME_SERIES_HISTOGRAMS_TAG: u32 = 3; +const TIME_SERIES_EXEMPLARS_TAG: u32 = 4; +const TIME_SERIES_METADATA_TAG: u32 = 5; -pub(crate) fn decode_remote_write_v2_request(is_zstd: bool, body: Bytes) -> Result { - let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_DECODE_ELAPSED.start_timer(); +struct BorrowedRequest<'a> { + symbols: Vec<&'a str>, + timeseries: Vec<&'a [u8]>, +} - // Match the v1 decoder's VictoriaMetrics fallback: some clients may send a - // mismatched content-encoding header, so try the other compression on failure. - let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { - buf - } else { - try_decompress(!is_zstd, &body[..])? - }; +impl<'a> BorrowedRequest<'a> { + fn decode(mut buf: &'a [u8]) -> std::result::Result { + let mut symbols = Vec::new(); + let mut timeseries = Vec::new(); - Request::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu) + while buf.has_remaining() { + let (tag, wire_type) = decode_key(&mut buf)?; + match tag { + 4 => { + let value = + take_length_delimited(wire_type, &mut buf).map_err(|mut error| { + error.push("Request", "symbols"); + error + })?; + let symbol = std::str::from_utf8(value).map_err(|_| { + let mut error = + DecodeError::new("invalid string value: data is not UTF-8 encoded"); + error.push("Request", "symbols"); + error + })?; + symbols.push(symbol); + } + 5 => { + let series = + take_length_delimited(wire_type, &mut buf).map_err(|mut error| { + error.push("Request", "timeseries"); + error + })?; + timeseries.push(series); + } + _ => skip_field(wire_type, tag, &mut buf, DecodeContext::default())?, + } + } + + Ok(Self { + symbols, + timeseries, + }) + } } pub(crate) struct RemoteWriteV2WriteRequests { @@ -81,19 +123,33 @@ pub(crate) struct RemoteWriteV2WriteRequests { pub semantic_index: SemanticIndexes, } -/// Converts a PRW v2 request into normal sample writes and native histogram writes. -/// -/// A metric name may appear in only one payload kind per request because the -/// metric-engine logical table is either a float metric or a native histogram. -pub(crate) fn into_write_requests(request: Request) -> Result { - let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CONVERT_ELAPSED.start_timer(); - let Request { - symbols, - timeseries, - } = request; +pub(crate) fn decode_remote_write_v2( + is_zstd: bool, + body: Bytes, + native_histograms_enabled: bool, +) -> Result { + let decode_timer = crate::metrics::METRIC_HTTP_PROM_STORE_DECODE_ELAPSED.start_timer(); + // Match the v1 decoder's VictoriaMetrics fallback: some clients may send a + // mismatched content-encoding header, so try the other compression on failure. + let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { + buf + } else { + try_decompress(!is_zstd, &body[..])? + }; + let request = BorrowedRequest::decode(&buf).context(error::DecodePromRemoteRequestSnafu)?; + drop(decode_timer); + + let _convert_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CONVERT_ELAPSED.start_timer(); + convert_remote_write_v2(request, native_histograms_enabled) +} + +fn convert_remote_write_v2( + request: BorrowedRequest<'_>, + native_histograms_enabled: bool, +) -> Result { ensure!( - symbols.first().map(|s| s.as_str()) == Some(""), + request.symbols.first().copied() == Some(""), error::InvalidPromRemoteRequestSnafu { msg: "remote write v2 symbols must start with an empty string".to_string(), } @@ -104,28 +160,40 @@ pub(crate) fn into_write_requests(request: Request) -> Result 0 { - histogram_count > 0 + let has_other_value_type = if counts.samples > 0 { + counts.histograms > 0 || histogram_tables .get(&prom_ctx) .is_some_and(|tables| tables.contains_key(&table_name)) @@ -143,42 +211,40 @@ pub(crate) fn into_write_requests(request: Request) -> Result 0 && histogram_count == 0 { - // Fast path for regular sample-only series. Move the resolved labels into - // the sample writer instead of cloning them for a histogram path we won't use. - let table_data = get_or_create_table_data( - &mut sample_tables, - prom_ctx, - table_name, - tags.len() + 2, - sample_count, - ); + let column_count = + tags.len() + .checked_add(2) + .with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 series has too many labels".to_string(), + })?; + let (writer, row_count) = if counts.samples > 0 { + ( + SeriesWriter::Samples(get_or_create_table_data( + &mut sample_tables, + prom_ctx, + table_name, + column_count, + counts.samples, + )), + counts.samples, + ) + } else { + ( + SeriesWriter::Histograms(get_or_create_table_data( + &mut histogram_tables, + prom_ctx, + table_name, + column_count, + counts.histograms, + )), + counts.histograms, + ) + }; - write_samples(table_data, series.samples, tags)?; - sample_count_total += sample_count as u64; - // The owned labels were moved above, so skip the mixed-series path below. - continue; - } - - if histogram_count > 0 { - let table_data = get_or_create_table_data( - &mut histogram_tables, - prom_ctx, - table_name, - tags.len() + 2, - histogram_count, - ); - - let mut histograms = series.histograms; - let Some(last_histogram) = histograms.pop() else { - continue; - }; - for histogram in &histograms { - write_native_histogram(table_data, histogram, tags.iter().cloned())?; - } - write_native_histogram(table_data, &last_histogram, tags.into_iter())?; - histogram_count_total += histogram_count as u64; - } + decode_series_leaves(series, Some(writer), tags, row_count, &mut scratch)?; + sample_count_total = checked_total(sample_count_total, counts.samples, "sample")?; + histogram_count_total = + checked_total(histogram_count_total, counts.histograms, "histogram")?; } Ok(RemoteWriteV2WriteRequests { @@ -201,14 +267,11 @@ pub(crate) fn into_write_requests(request: Request) -> Result Result<()> { - let Some(series_metadata) = &series.metadata else { - return Ok(()); - }; // Symbol 0 is the mandatory empty string: no help / no unit. if series_metadata.help_ref != 0 { symbol_ref(symbols, series_metadata.help_ref, "metadata help")?; @@ -258,6 +321,256 @@ fn metric_type_value(wire_type: i32) -> Option<&'static str> { } } +#[derive(Default)] +struct SeriesCounts { + samples: usize, + histograms: usize, +} + +fn scan_series( + mut buf: &[u8], + labels_refs: &mut Vec, + metadata: &mut Metadata, +) -> std::result::Result { + labels_refs.clear(); + metadata.clear(); + let mut counts = SeriesCounts::default(); + + while buf.has_remaining() { + let (tag, wire_type) = decode_key(&mut buf)?; + match tag { + TIME_SERIES_LABELS_REFS_TAG => { + uint32::merge_repeated(wire_type, labels_refs, &mut buf, DecodeContext::default()) + .map_err(|mut error| { + error.push("TimeSeries", "labels_refs"); + error + })? + } + TIME_SERIES_SAMPLES_TAG => { + take_length_delimited(wire_type, &mut buf).map_err(|mut error| { + error.push("TimeSeries", "samples"); + error + })?; + counts.samples = counts.samples.checked_add(1).ok_or_else(|| { + DecodeError::new("remote write v2 sample count overflows usize") + })?; + } + TIME_SERIES_HISTOGRAMS_TAG => { + take_length_delimited(wire_type, &mut buf).map_err(|mut error| { + error.push("TimeSeries", "histograms"); + error + })?; + counts.histograms = counts.histograms.checked_add(1).ok_or_else(|| { + DecodeError::new("remote write v2 histogram count overflows usize") + })?; + } + TIME_SERIES_EXEMPLARS_TAG => { + take_length_delimited(wire_type, &mut buf)?; + } + TIME_SERIES_METADATA_TAG => { + message::merge(wire_type, metadata, &mut buf, DecodeContext::default()).map_err( + |mut error| { + error.push("TimeSeries", "metadata"); + error + }, + )? + } + _ => skip_field(wire_type, tag, &mut buf, DecodeContext::default())?, + } + } + + Ok(counts) +} + +#[derive(Default)] +struct LeafScratch { + sample: Sample, + histogram: Histogram, + exemplar: Exemplar, +} + +enum SeriesWriter<'a> { + Samples(&'a mut TableData), + Histograms(&'a mut TableData), +} + +fn decode_series_leaves( + mut buf: &[u8], + mut writer: Option>, + mut tags: PromTags<'_>, + mut rows_remaining: usize, + scratch: &mut LeafScratch, +) -> Result<()> { + let mut sample_row_template = None; + + while buf.has_remaining() { + let (tag, wire_type) = decode_key(&mut buf).context(error::DecodePromRemoteRequestSnafu)?; + match tag { + TIME_SERIES_LABELS_REFS_TAG => { + skip_field(wire_type, tag, &mut buf, DecodeContext::default()) + .context(error::DecodePromRemoteRequestSnafu)? + } + TIME_SERIES_SAMPLES_TAG => { + scratch.sample.clear(); + message::merge( + wire_type, + &mut scratch.sample, + &mut buf, + DecodeContext::default(), + ) + .map_err(|mut error| { + error.push("TimeSeries", "samples"); + error + }) + .context(error::DecodePromRemoteRequestSnafu)?; + rows_remaining = rows_remaining.checked_sub(1).with_context(|| { + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 sample count changed between scans".to_string(), + } + })?; + if let Some(SeriesWriter::Samples(table_data)) = &mut writer { + if sample_row_template.is_none() { + let timestamp_index = table_data.ensure_column_by_name( + greptime_timestamp(), + ColumnDataType::TimestampMillisecond, + SemanticType::Timestamp, + )?; + let value_index = table_data.ensure_column_by_name( + greptime_value(), + ColumnDataType::Float64, + SemanticType::Field, + )?; + let mut row = table_data.alloc_one_row(); + row_writer::write_tags( + table_data, + std::mem::take(&mut tags).into_iter(), + &mut row, + )?; + sample_row_template = Some((row, timestamp_index, value_index)); + } + + if let Some((row, timestamp_index, value_index)) = sample_row_template.as_mut() + { + row[*timestamp_index].value_data = Some( + ValueData::TimestampMillisecondValue(scratch.sample.timestamp), + ); + row[*value_index].value_data = + Some(ValueData::F64Value(scratch.sample.value)); + let row = if rows_remaining == 0 { + std::mem::take(row) + } else { + row.clone() + }; + table_data.add_row(row); + } + } + } + TIME_SERIES_HISTOGRAMS_TAG => { + scratch.histogram.clear(); + message::merge( + wire_type, + &mut scratch.histogram, + &mut buf, + DecodeContext::default(), + ) + .map_err(|mut error| { + error.push("TimeSeries", "histograms"); + error + }) + .context(error::DecodePromRemoteRequestSnafu)?; + rows_remaining = rows_remaining.checked_sub(1).with_context(|| { + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 histogram count changed between scans".to_string(), + } + })?; + if let Some(SeriesWriter::Histograms(table_data)) = &mut writer { + if rows_remaining == 0 { + write_native_histogram( + table_data, + &scratch.histogram, + std::mem::take(&mut tags).into_iter(), + )?; + } else { + write_native_histogram( + table_data, + &scratch.histogram, + tags.iter().cloned(), + )?; + } + } + } + TIME_SERIES_EXEMPLARS_TAG => { + scratch.exemplar.clear(); + message::merge( + wire_type, + &mut scratch.exemplar, + &mut buf, + DecodeContext::default(), + ) + .map_err(|mut error| { + error.push("TimeSeries", "exemplars"); + error + }) + .context(error::DecodePromRemoteRequestSnafu)?; + } + TIME_SERIES_METADATA_TAG => { + take_length_delimited(wire_type, &mut buf) + .map_err(|mut error| { + error.push("TimeSeries", "metadata"); + error + }) + .context(error::DecodePromRemoteRequestSnafu)?; + } + _ => skip_field(wire_type, tag, &mut buf, DecodeContext::default()) + .context(error::DecodePromRemoteRequestSnafu)?, + } + } + + ensure!( + rows_remaining == 0, + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 row count changed between scans".to_string(), + } + ); + + Ok(()) +} + +fn take_length_delimited<'a>( + wire_type: WireType, + buf: &mut &'a [u8], +) -> std::result::Result<&'a [u8], DecodeError> { + if wire_type != WireType::LengthDelimited { + return Err(DecodeError::new(format!( + "invalid wire type: {wire_type:?} (expected LengthDelimited)" + ))); + } + + let len = decode_varint(buf)?; + let len = + usize::try_from(len).map_err(|_| DecodeError::new("length delimiter exceeds usize"))?; + if len > buf.len() { + return Err(DecodeError::new("buffer underflow")); + } + let (value, remaining) = buf.split_at(len); + *buf = remaining; + Ok(value) +} + +fn checked_total(total: u64, count: usize, name: &str) -> Result { + let count = + u64::try_from(count) + .ok() + .with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {name} count exceeds u64"), + })?; + total + .checked_add(count) + .with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {name} count overflows u64"), + }) +} + fn get_or_create_table_data( tables: &mut HashMap>, prom_ctx: PromCtx, @@ -275,46 +588,10 @@ fn get_or_create_table_data( } } -fn write_samples( - table_data: &mut TableData, - mut samples: Vec, - tags: PromTags, -) -> Result<()> { - let Some(last_sample) = samples.pop() else { - return Ok(()); - }; - - for sample in &samples { - write_sample(table_data, sample, tags.iter().cloned())?; - } - - write_sample(table_data, &last_sample, tags.into_iter()) -} - -fn write_sample( - table_data: &mut TableData, - sample: &Sample, - tags: impl Iterator, -) -> Result<()> { - let mut row = table_data.alloc_one_row(); - row_writer::write_ts_to_millis( - table_data, - greptime_timestamp(), - Some(sample.timestamp), - Precision::Millisecond, - &mut row, - )?; - row_writer::write_f64(table_data, greptime_value(), sample.value, &mut row)?; - row_writer::write_tags(table_data, tags, &mut row)?; - table_data.add_row(row); - - Ok(()) -} - -fn write_native_histogram( +fn write_native_histogram<'a>( table_data: &mut TableData, histogram: &Histogram, - tags: impl Iterator, + tags: impl Iterator, ) -> Result<()> { // Persist both int and float families into the logical table schema. Only one // family is populated per row; the other is written as NULL so PromQL can @@ -579,7 +856,7 @@ fn validate_native_histogram_spans( let span_len = spans .iter() .try_fold(0usize, |sum, span| sum.checked_add(span.length as usize)) - .context(error::InvalidPromRemoteRequestSnafu { + .with_context(|| error::InvalidPromRemoteRequestSnafu { msg: format!("remote write v2 native histogram {name} spans overflow"), })?; ensure!( @@ -606,13 +883,13 @@ fn validate_native_histogram_spans( current_index = if span_index == 0 { span.offset } else { - current_index.checked_add(span.offset).context( + current_index.checked_add(span.offset).with_context(|| { error::InvalidPromRemoteRequestSnafu { msg: format!( "remote write v2 native histogram {name} span index overflows i32" ), - }, - )? + } + })? }; for _ in 0..span.length { @@ -701,7 +978,7 @@ fn validate_integer_native_histogram_counts( .try_fold(zero_count, |total, bucket| { total.checked_add(*bucket as u64) }) - .context(error::InvalidPromRemoteRequestSnafu { + .with_context(|| error::InvalidPromRemoteRequestSnafu { msg: "remote write v2 native histogram bucket total overflows u64".to_string(), })?; ensure!( @@ -824,11 +1101,12 @@ fn bucket_counts_from_deltas(deltas: &[i64]) -> Result> { let mut buckets = Vec::with_capacity(deltas.len()); for delta in deltas { - count = count - .checked_add(*delta) - .context(error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram bucket count overflows i64".to_string(), - })?; + count = + count + .checked_add(*delta) + .with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 native histogram bucket count overflows i64".to_string(), + })?; ensure!( count >= 0, error::InvalidPromRemoteRequestSnafu { @@ -841,11 +1119,11 @@ fn bucket_counts_from_deltas(deltas: &[i64]) -> Result> { Ok(buckets) } -fn ensure_no_internal_histogram_labels(tags: &PromTags) -> Result<()> { +fn ensure_no_internal_histogram_labels(tags: &PromTags<'_>) -> Result<()> { // The histogram field column is generated from the protobuf payload. for (name, _) in tags { ensure!( - name != greptime_native_histogram() && name != NATIVE_HISTOGRAM_FIELD, + *name != greptime_native_histogram() && *name != NATIVE_HISTOGRAM_FIELD, error::InvalidPromRemoteRequestSnafu { msg: format!( "remote write v2 label `{name}` conflicts with an internal native histogram label" @@ -858,12 +1136,12 @@ fn ensure_no_internal_histogram_labels(tags: &PromTags) -> Result<()> { } fn resolve_series_labels<'a>( - symbols: &'a [String], - series: &TimeSeries, + symbols: &'a [&str], + labels_refs: &[u32], label_names: &mut HashSet<&'a str>, -) -> Result { +) -> Result> { ensure!( - series.labels_refs.len().is_multiple_of(2), + labels_refs.len().is_multiple_of(2), error::InvalidPromRemoteRequestSnafu { msg: "remote write v2 labels_refs must contain name/value pairs".to_string(), } @@ -871,10 +1149,10 @@ fn resolve_series_labels<'a>( let mut prom_ctx = PromCtx::default(); let mut table_name = None; - let mut tags = Vec::with_capacity(series.labels_refs.len() / 2); + let mut tags = Vec::with_capacity(labels_refs.len() / 2); label_names.clear(); - for pair in series.labels_refs.chunks_exact(2) { + for pair in labels_refs.chunks_exact(2) { let name = symbol_ref(symbols, pair[0], "label name")?; let value = symbol_ref(symbols, pair[1], "label value")?; validate_label(name)?; @@ -893,10 +1171,10 @@ fn resolve_series_labels<'a>( continue; } - tags.push((name.to_string(), value.to_string())); + tags.push((name, value.to_string())); } - let table_name = table_name.context(error::InvalidPromRemoteRequestSnafu { + let table_name = table_name.with_context(|| error::InvalidPromRemoteRequestSnafu { msg: "missing '__name__' label in time-series".to_string(), })?; ensure!( @@ -920,10 +1198,15 @@ fn validate_label(name: &str) -> Result<()> { Ok(()) } -fn symbol_ref<'a>(symbols: &'a [String], idx: u32, field: &str) -> Result<&'a str> { +fn symbol_ref<'a>(symbols: &'a [&str], idx: u32, field: &str) -> Result<&'a str> { + let idx = usize::try_from(idx) + .ok() + .with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {field} symbol reference exceeds usize"), + })?; symbols - .get(idx as usize) - .map(String::as_str) + .get(idx) + .copied() .with_context(|| error::InvalidPromRemoteRequestSnafu { msg: format!( "remote write v2 {field} symbol reference {idx} is out of range, symbols len: {}", @@ -994,8 +1277,12 @@ pub mod test_util { use api::greptime_proto::io::prometheus::write::v2::{Histogram, Request, Sample, TimeSeries}; use api::v1::RowInsertRequest; use bytes::Bytes; + use prost::Message; + use snafu::ResultExt; - use crate::error::Result; + use crate::error::{self, Result}; + use crate::prom_remote_write::try_decompress; + use crate::prom_store::snappy_compress; pub fn request_with_labels_and_samples( labels: Vec<(&str, &str)>, @@ -1012,13 +1299,42 @@ pub mod test_util { } pub fn decode_request(is_zstd: bool, body: Bytes) -> Result { - super::decode_remote_write_v2_request(is_zstd, body) + let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { + buf + } else { + try_decompress(!is_zstd, &body[..])? + }; + Request::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu) } pub fn write_requests( request: Request, ) -> Result<(Vec, Vec, u64, u64)> { - let requests = super::into_write_requests(request)?; + let body = Bytes::from(snappy_compress(&request.encode_to_vec())?); + decode_write_requests(false, body, true) + } + + pub fn decode_write_requests( + is_zstd: bool, + body: Bytes, + native_histograms_enabled: bool, + ) -> Result<(Vec, Vec, u64, u64)> { + let requests = super::decode_remote_write_v2(is_zstd, body, native_histograms_enabled)?; + Ok(( + requests.samples.all_req().collect(), + requests.histograms.all_req().collect(), + requests.sample_count, + requests.histogram_count, + )) + } + + pub fn decode_uncompressed_write_requests( + body: &[u8], + native_histograms_enabled: bool, + ) -> Result<(Vec, Vec, u64, u64)> { + let request = + super::BorrowedRequest::decode(body).context(error::DecodePromRemoteRequestSnafu)?; + let requests = super::convert_remote_write_v2(request, native_histograms_enabled)?; Ok(( requests.samples.all_req().collect(), requests.histograms.all_req().collect(), @@ -1109,7 +1425,7 @@ mod tests { let body = Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); - let decoded = decode_remote_write_v2_request(false, body).unwrap(); + let decoded = test_util::decode_request(false, body.clone()).unwrap(); assert_eq!(decoded.symbols, request.symbols); assert_eq!(decoded.timeseries.len(), 1); @@ -1117,11 +1433,291 @@ mod tests { assert_eq!(decoded.timeseries[0].samples.len(), 1); assert_eq!(decoded.timeseries[0].samples[0].value, 42.0); assert_eq!(decoded.timeseries[0].metadata.as_ref().unwrap().r#type, 1); + assert_eq!( + decode_remote_write_v2(true, body, true) + .unwrap() + .sample_count, + 1 + ); + } + + #[test] + fn test_fused_decoder_accepts_arbitrary_field_order_and_split_labels() { + let mut first_sample = Sample { + value: 42.0, + timestamp: 1000, + start_timestamp: 500, + } + .encode_to_vec(); + first_sample.extend(varint_field(90, 1)); + let second_sample = Sample { + value: 43.0, + timestamp: 2000, + start_timestamp: 1000, + } + .encode_to_vec(); + + let mut series = encoded_message_field(2, &first_sample); + series.extend(packed_u32_field(1, &[1])); + series.extend(varint_field(90, 1)); + series.extend(encoded_message_field(2, &second_sample)); + series.extend(varint_field(1, 2)); + + let mut wire = encoded_message_field(5, &series); + wire.extend(string_field(4, b"")); + wire.extend(varint_field(90, 1)); + wire.extend(string_field(4, METRIC_NAME_LABEL.as_bytes())); + wire.extend(encoded_message_field(5, &packed_u32_field(1, &[99]))); + wire.extend(string_field(4, b"http_requests_total")); + + let requests = decode_wire(&wire, true).unwrap(); + assert_eq!(requests.sample_count, 2); + assert_eq!(requests.histogram_count, 0); + let rows = requests.samples.all_req().next().unwrap().rows.unwrap(); + assert_eq!(rows.rows.len(), 2); + assert_eq!( + rows.schema + .iter() + .map(|column| column.column_name.as_str()) + .collect::>(), + vec![greptime_timestamp(), greptime_value()] + ); + assert_eq!( + rows.rows[0].values[1].value_data, + Some(ValueData::F64Value(42.0)) + ); + assert_eq!( + rows.rows[1].values[1].value_data, + Some(ValueData::F64Value(43.0)) + ); + } + + #[test] + fn test_fused_decoder_accepts_histogram_before_labels_and_unknown_fields() { + let mut histogram = Histogram { + count: Some(Count::CountInt(0)), + zero_count: Some(ZeroCount::ZeroCountInt(0)), + timestamp: 2000, + start_timestamp: 1000, + ..Default::default() + } + .encode_to_vec(); + histogram.extend(varint_field(90, 1)); + + let mut series = encoded_message_field(3, &histogram); + series.extend(packed_u32_field(1, &[1, 2])); + let wire = request_wire(&["", METRIC_NAME_LABEL, "metric"], &[series]); + + let requests = decode_wire(&wire, true).unwrap(); + assert_eq!(requests.histogram_count, 1); + let rows = requests.histograms.all_req().next().unwrap().rows.unwrap(); + assert_eq!( + histogram_field_value(&rows, 0, START_TIMESTAMP_FIELD), + Some(ValueData::TimestampMillisecondValue(1000)) + ); + } + + #[test] + fn test_fused_decoder_rejects_malformed_wire() { + let mut wrong_request_wire = varint_field(4, 0); + let invalid_utf8 = string_field(4, &[0xff]); + + let mut truncated_request = Vec::new(); + prost::encoding::encode_key(5, WireType::LengthDelimited, &mut truncated_request); + prost::encoding::encode_varint(2, &mut truncated_request); + truncated_request.push(0); + + let mut oversized_request = Vec::new(); + prost::encoding::encode_key(4, WireType::LengthDelimited, &mut oversized_request); + prost::encoding::encode_varint(u64::MAX, &mut oversized_request); + + let malformed_varint = vec![0x80; 10]; + + let mut wrong_labels_wire = Vec::new(); + prost::encoding::encode_key(1, WireType::SixtyFourBit, &mut wrong_labels_wire); + wrong_labels_wire.extend([0; 8]); + wrong_labels_wire = request_wire(&[""], &[std::mem::take(&mut wrong_labels_wire)]); + + let wrong_series_wire = request_wire(&[""], &[varint_field(2, 0)]); + + let mut truncated_series = Vec::new(); + prost::encoding::encode_key(2, WireType::LengthDelimited, &mut truncated_series); + prost::encoding::encode_varint(2, &mut truncated_series); + truncated_series.push(0x08); + let truncated_series = request_wire(&[""], &[truncated_series]); + + let mut oversized_series = Vec::new(); + prost::encoding::encode_key(3, WireType::LengthDelimited, &mut oversized_series); + prost::encoding::encode_varint(u64::MAX, &mut oversized_series); + let oversized_series = request_wire(&[""], &[oversized_series]); + + let invalid_sample = series_wire(&[1, 2], 2, &varint_field(1, 1)); + let invalid_sample = request_wire(&["", METRIC_NAME_LABEL, "metric"], &[invalid_sample]); + let invalid_histogram = series_wire(&[1, 2], 3, &varint_field(3, 1)); + let invalid_histogram = + request_wire(&["", METRIC_NAME_LABEL, "metric"], &[invalid_histogram]); + + for (name, wire) in [ + ( + "wrong request wire type", + std::mem::take(&mut wrong_request_wire), + ), + ("invalid utf8", invalid_utf8), + ("truncated request", truncated_request), + ("oversized request length", oversized_request), + ("malformed varint", malformed_varint), + ("wrong labels wire type", wrong_labels_wire), + ("wrong series wire type", wrong_series_wire), + ("truncated series", truncated_series), + ("oversized series length", oversized_series), + ("invalid sample", invalid_sample), + ("invalid histogram", invalid_histogram), + ] { + let error = decode_wire_error(&wire, true, name); + assert!( + matches!(error, error::Error::DecodePromRemoteRequest { .. }), + "{name}: {error}" + ); + } + } + + #[test] + fn test_fused_decoder_rejects_malformed_ignored_messages() { + for tag in [4, 5] { + let mut series = series_wire(&[1, 2], 2, &Sample::default().encode_to_vec()); + series.extend(encoded_message_field(tag, &[0x08])); + let wire = request_wire(&["", METRIC_NAME_LABEL, "metric"], &[series]); + + let error = decode_wire_error(&wire, true, "malformed ignored message"); + assert!(matches!( + error, + error::Error::DecodePromRemoteRequest { .. } + )); + } + } + + #[test] + fn test_fused_decoder_ignores_exemplar_symbol_refs() { + let request = Request { + symbols: vec![ + String::new(), + METRIC_NAME_LABEL.to_string(), + "metric".to_string(), + ], + timeseries: vec![TimeSeries { + labels_refs: vec![1, 2], + samples: vec![Sample::default()], + exemplars: vec![Exemplar { + labels_refs: vec![99], + ..Default::default() + }], + ..Default::default() + }], + }; + + assert_eq!(decode_test_request(request).unwrap().sample_count, 1); + } + + #[test] + fn test_fused_decoder_preserves_empty_request_and_series_behavior() { + let request = Request { + symbols: vec![ + String::new(), + METRIC_NAME_LABEL.to_string(), + "metric".to_string(), + "job".to_string(), + ], + timeseries: vec![ + TimeSeries { + labels_refs: vec![99, 99], + ..Default::default() + }, + TimeSeries { + labels_refs: vec![1], + ..Default::default() + }, + TimeSeries { + labels_refs: vec![3, 2, 3, 2], + ..Default::default() + }, + TimeSeries::default(), + ], + }; + + let requests = decode_test_request(request).unwrap(); + assert_eq!(requests.sample_count, 0); + assert_eq!(requests.histogram_count, 0); + + let requests = decode_test_request(Request { + symbols: vec![String::new()], + timeseries: Vec::new(), + }) + .unwrap(); + assert_eq!(requests.sample_count, 0); + assert!(decode_wire(&[], true).is_err()); + assert!(decode_wire(&[0x0a, 0x00], true).is_err()); + } + + #[test] + fn test_fused_decoder_rejects_duplicate_resolved_label_names() { + let request = Request { + symbols: vec![ + String::new(), + METRIC_NAME_LABEL.to_string(), + "metric".to_string(), + "job".to_string(), + "job".to_string(), + "api".to_string(), + "worker".to_string(), + ], + timeseries: vec![TimeSeries { + labels_refs: vec![1, 2, 3, 5, 4, 6], + samples: vec![Sample::default()], + ..Default::default() + }], + }; + + assert_invalid( + "duplicate resolved labels", + request, + "label name `job` is repeated", + ); + } + + #[test] + fn test_fused_decoder_pins_experimental_error_precedence() { + let histogram = series_wire(&[1, 2], 3, &Histogram::default().encode_to_vec()); + let malformed_sample = series_wire(&[1, 2], 2, &[0x08]); + + let wire = request_wire( + &["", METRIC_NAME_LABEL, "metric"], + &[histogram.clone(), malformed_sample.clone()], + ); + let error = decode_wire_error(&wire, false, "histogram before malformed series"); + assert!(error.to_string().contains("ingestion is experimental")); + + let wire = request_wire( + &["", METRIC_NAME_LABEL, "metric"], + &[malformed_sample, histogram.clone()], + ); + let error = decode_wire_error(&wire, false, "malformed series before histogram"); + assert!(matches!( + error, + error::Error::DecodePromRemoteRequest { .. } + )); + + let missing_name = series_wire(&[3, 4], 2, &Sample::default().encode_to_vec()); + let wire = request_wire( + &["", METRIC_NAME_LABEL, "metric", "job", "api"], + &[missing_name, histogram], + ); + let error = decode_wire_error(&wire, false, "conversion error before histogram"); + assert!(error.to_string().contains("missing '__name__'")); } #[test] fn test_into_context_req_samples() { - let ctx_req = into_write_requests(test_util::request_with_labels_and_samples( + let ctx_req = decode_test_request(test_util::request_with_labels_and_samples( vec![ (METRIC_NAME_LABEL, "http_requests_total"), ("job", "api"), @@ -1183,11 +1779,19 @@ mod tests { rows.rows[1].values[1].value_data, Some(ValueData::F64Value(43.0)) ); + assert_eq!( + rows.rows[1].values[2].value_data, + Some(ValueData::StringValue("api".to_string())) + ); + assert_eq!( + rows.rows[1].values[3].value_data, + Some(ValueData::StringValue("localhost:9090".to_string())) + ); } #[test] fn test_into_context_req_special_labels() { - let ctx_req = into_write_requests(test_util::request_with_labels_and_samples( + let ctx_req = decode_test_request(test_util::request_with_labels_and_samples( vec![ (METRIC_NAME_LABEL, "cpu_usage"), (DATABASE_LABEL, "tenant_a"), @@ -1518,7 +2122,7 @@ mod tests { #[test] fn test_into_context_req_allows_nan_observations_outside_buckets() { - into_write_requests(request_with_histogram(Histogram { + decode_test_request(request_with_histogram(Histogram { count: Some(Count::CountInt(2)), sum: f64::NAN, positive_spans: vec![BucketSpan { @@ -1533,7 +2137,7 @@ mod tests { #[test] fn test_into_context_req_allows_empty_label_values() { - let ctx_req = into_write_requests(test_util::request_with_labels_and_samples( + let ctx_req = decode_test_request(test_util::request_with_labels_and_samples( vec![(METRIC_NAME_LABEL, "metric"), ("job", "")], vec![Sample { value: 1.0, @@ -1649,7 +2253,7 @@ mod tests { }]; histogram.negative_deltas = vec![1]; } - into_write_requests(request_with_histogram(histogram.clone())).unwrap(); + decode_test_request(request_with_histogram(histogram.clone())).unwrap(); let beyond = max_index + 1; if positive { @@ -1694,7 +2298,7 @@ mod tests { ], }; - let ctx_req = into_write_requests(request).unwrap(); + let ctx_req = decode_test_request(request).unwrap(); assert_eq!(ctx_req.sample_count, 1); assert_eq!(ctx_req.histogram_count, 1); @@ -1708,7 +2312,7 @@ mod tests { test_util::request_with_labels_and_samples(vec![(METRIC_NAME_LABEL, "metric")], vec![]); request.timeseries[0].histograms.push(Histogram::default()); - let ctx_req = into_write_requests(request).unwrap(); + let ctx_req = decode_test_request(request).unwrap(); assert_eq!(ctx_req.sample_count, 0); assert_eq!(ctx_req.histogram_count, 1); @@ -1744,7 +2348,7 @@ mod tests { #[test] fn test_into_context_req_preserves_histogram_start_timestamp() { - let ctx_req = into_write_requests(test_util::request_with_labels_and_histograms( + let ctx_req = decode_test_request(test_util::request_with_labels_and_histograms( vec![(METRIC_NAME_LABEL, "metric")], vec![Histogram { timestamp: 2000, @@ -1774,7 +2378,7 @@ mod tests { ); request.timeseries[0].histograms.push(Histogram::default()); - let err = match into_write_requests(request) { + let err = match decode_test_request(request) { Ok(_) => panic!("expected invalid request error"), Err(err) => err, }; @@ -1790,7 +2394,7 @@ mod tests { assert_eq!(greptime_native_histogram(), "custom_native_histogram"); let err = ensure_no_internal_histogram_labels(&vec![( - NATIVE_HISTOGRAM_FIELD.to_string(), + NATIVE_HISTOGRAM_FIELD, "user_value".to_string(), )]) .unwrap_err(); @@ -1837,7 +2441,7 @@ mod tests { ], }; - let ctx_req = into_write_requests(request).unwrap(); + let ctx_req = decode_test_request(request).unwrap(); assert_eq!(ctx_req.histogram_count, 2); let mut inserts = ctx_req.histograms.all_req().collect::>(); @@ -1883,6 +2487,65 @@ mod tests { )); } + fn decode_wire( + wire: &[u8], + native_histograms_enabled: bool, + ) -> Result { + let body = Bytes::from(crate::prom_store::snappy_compress(wire).unwrap()); + decode_remote_write_v2(false, body, native_histograms_enabled) + } + + fn decode_wire_error(wire: &[u8], native_histograms_enabled: bool, name: &str) -> error::Error { + match decode_wire(wire, native_histograms_enabled) { + Ok(_) => panic!("{name}: expected decoder error"), + Err(error) => error, + } + } + + fn request_wire(symbols: &[&str], series: &[Vec]) -> Vec { + let mut wire = Vec::new(); + for symbol in symbols { + wire.extend(string_field(4, symbol.as_bytes())); + } + for series in series { + wire.extend(encoded_message_field(5, series)); + } + wire + } + + fn series_wire(labels_refs: &[u32], leaf_tag: u32, leaf: &[u8]) -> Vec { + let mut series = packed_u32_field(1, labels_refs); + series.extend(encoded_message_field(leaf_tag, leaf)); + series + } + + fn string_field(tag: u32, value: &[u8]) -> Vec { + encoded_message_field(tag, value) + } + + fn encoded_message_field(tag: u32, value: &[u8]) -> Vec { + let mut field = Vec::new(); + prost::encoding::encode_key(tag, WireType::LengthDelimited, &mut field); + prost::encoding::encode_varint(u64::try_from(value.len()).unwrap(), &mut field); + field.extend_from_slice(value); + field + } + + fn varint_field(tag: u32, value: u64) -> Vec { + let mut field = Vec::new(); + prost::encoding::encode_key(tag, WireType::Varint, &mut field); + prost::encoding::encode_varint(value, &mut field); + field + } + + fn packed_u32_field(tag: u32, values: &[u32]) -> Vec { + let mut packed = Vec::new(); + for value in values { + prost::encoding::encode_varint(u64::from(*value), &mut packed); + } + encoded_message_field(tag, &packed) + } + fn request_with_sample(labels: Vec<(&str, &str)>) -> Request { test_util::request_with_labels_and_samples( labels, @@ -1901,8 +2564,21 @@ mod tests { ) } + fn decode_test_request(request: Request) -> Result { + decode_test_request_with_histograms(request, true) + } + + fn decode_test_request_with_histograms( + request: Request, + native_histograms_enabled: bool, + ) -> Result { + let body = + Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); + decode_remote_write_v2(false, body, native_histograms_enabled) + } + fn assert_invalid(name: &str, request: Request, expected: &str) { - let err = match into_write_requests(request) { + let err = match decode_test_request(request) { Ok(_) => panic!("{name}: expected invalid request error"), Err(err) => err, }; @@ -1993,7 +2669,7 @@ mod tests { >; fn decoded_index(request: Request) -> DecodedIndex { - let encoded = into_write_requests(request) + let encoded = decode_test_request(request) .unwrap() .semantic_index .encode("public") @@ -2045,7 +2721,7 @@ mod tests { assert!(!untyped.contains_key(SEMANTIC_METRIC_TYPE)); assert!(!untyped.contains_key(SEMANTIC_METRIC_METADATA_QUALITY)); - let requests = into_write_requests(metadata_request(vec![( + let requests = decode_test_request(metadata_request(vec![( vec![(METRIC_NAME_LABEL, "untyped_unitless_total")], metadata::MetricType::Unspecified as i32, None, @@ -2054,7 +2730,7 @@ mod tests { assert!(requests.semantic_index.is_empty()); // No metadata at all behaves the same. - let requests = into_write_requests(test_util::request_with_labels_and_samples( + let requests = decode_test_request(test_util::request_with_labels_and_samples( vec![(METRIC_NAME_LABEL, "bare_total")], vec![Sample { value: 1.0, @@ -2123,7 +2799,7 @@ mod tests { None, )]); request.timeseries[0].metadata.as_mut().unwrap().unit_ref = 999; - let err = into_write_requests(request).err().unwrap(); + let err = decode_test_request(request).err().unwrap(); assert!(err.to_string().contains("out of range"), "{err}"); // unit_ref must be validated even when the type is UNSPECIFIED @@ -2134,7 +2810,7 @@ mod tests { None, )]); request.timeseries[0].metadata.as_mut().unwrap().unit_ref = 999; - let err = into_write_requests(request).err().unwrap(); + let err = decode_test_request(request).err().unwrap(); assert!(err.to_string().contains("out of range"), "{err}"); // help_ref is validated although help is never persisted. @@ -2144,7 +2820,7 @@ mod tests { None, )]); request.timeseries[0].metadata.as_mut().unwrap().help_ref = 999; - let err = into_write_requests(request).err().unwrap(); + let err = decode_test_request(request).err().unwrap(); assert!(err.to_string().contains("out of range"), "{err}"); } } diff --git a/src/servers/src/row_writer.rs b/src/servers/src/row_writer.rs index 45afcbca6de..61fd9d77e5b 100644 --- a/src/servers/src/row_writer.rs +++ b/src/servers/src/row_writer.rs @@ -92,6 +92,26 @@ impl TableData { Ok(index) } + /// Ensures a default column schema without allocating its name when it already exists. + pub(crate) fn ensure_column_by_name( + &mut self, + name: &str, + datatype: ColumnDataType, + semantic_type: SemanticType, + ) -> Result { + if let Some(index) = self.column_indexes.get(name).copied() { + check_schema(datatype, semantic_type, &self.schema[index])?; + return Ok(index); + } + + self.ensure_column(ColumnSchema { + column_name: name.to_string(), + datatype: datatype as i32, + semantic_type: semantic_type as i32, + ..Default::default() + }) + } + #[allow(dead_code)] pub fn columns(&self) -> &Vec { &self.schema @@ -211,11 +231,14 @@ impl MultiTableData { } /// Write data as tags into the table data. -pub fn write_tags( +pub fn write_tags( table_data: &mut TableData, - tags: impl Iterator, + tags: impl Iterator, one_row: &mut Vec, -) -> Result<()> { +) -> Result<()> +where + K: AsRef + Into, +{ let ktv_iter = tags.map(|(k, v)| (k, ColumnDataType::String, Some(ValueData::StringValue(v)))); write_by_semantic_type(table_data, SemanticType::Tag, ktv_iter, one_row) } @@ -326,12 +349,15 @@ pub(crate) fn write_by_schema( Ok(()) } -fn write_by_semantic_type( +fn write_by_semantic_type( table_data: &mut TableData, semantic_type: SemanticType, - ktv_iter: impl Iterator)>, + ktv_iter: impl Iterator)>, one_row: &mut Vec, -) -> Result<()> { +) -> Result<()> +where + K: AsRef + Into, +{ let TableData { schema, column_indexes, @@ -339,12 +365,13 @@ fn write_by_semantic_type( } = table_data; for (name, datatype, value) in ktv_iter { - let index = column_indexes.get(&name); + let index = column_indexes.get(name.as_ref()).copied(); if let Some(index) = index { - check_schema(datatype, semantic_type, &schema[*index])?; - one_row[*index].value_data = value; + check_schema(datatype, semantic_type, &schema[index])?; + one_row[index].value_data = value; } else { let index = schema.len(); + let name = name.into(); schema.push(ColumnSchema { column_name: name.clone(), datatype: datatype as i32, diff --git a/src/servers/tests/prom_remote_write_v2_test.rs b/src/servers/tests/prom_remote_write_v2_test.rs index 94d260d3fac..75fee046f7c 100644 --- a/src/servers/tests/prom_remote_write_v2_test.rs +++ b/src/servers/tests/prom_remote_write_v2_test.rs @@ -77,7 +77,7 @@ fn test_decode_remote_write_v2_native_histogram_dump() { assert_eq!(histogram.timestamp, 1782358160412); let (sample_inserts, histogram_inserts, sample_count, histogram_count) = - remote_write_v2::write_requests(decoded).unwrap(); + remote_write_v2::decode_write_requests(false, Bytes::from_static(BODY), true).unwrap(); assert!(sample_inserts.is_empty()); assert_eq!(sample_count, 0); assert_eq!(histogram_count, 1); From 6538db6d618279f640958c40cdbb5a1796fb40cd Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Fri, 14 Aug 2026 06:21:17 +0000 Subject: [PATCH 069/126] refactor(json2): push down json2 type hints to parquet reads (#8833) * refactor(mito): push down json2 type hints to parquet reads Signed-off-by: fys * refactor(mito): share json2 target types with arc Signed-off-by: fys * refactor(mito): derive json2 output schema from target types Signed-off-by: fys * refactor(mito): simplify read columns construction Signed-off-by: fys * fix: cargo check Signed-off-by: fys * fix(mito): reject JSON hints for non-JSON2 read columns Signed-off-by: fys * fix: do not pushdown json type hint of non-json2-col Signed-off-by: fys * fix: unit test Signed-off-by: fys * refactor(query): simplify JSON type hint application Signed-off-by: fys * refactor: clean up JSON2 type hint handling Signed-off-by: fys * refactor(mito2): keep JSON2 hints with flat read format Signed-off-by: fys * fix(mito): use raw parquet projection for output schema Signed-off-by: fys * fix(query): note JSON2 hint scope limitation Signed-off-by: fys * refactor(mito): store JSON target types as native types Signed-off-by: fys * test(json2): cover join hint qualifier limitation Signed-off-by: fys * refactor(mito): remove JSON2 fallback from compat cast Signed-off-by: fys * docs(mito): document ReadColumns ordering contract Signed-off-by: fys * fix: cargo clippy Signed-off-by: fys --------- Signed-off-by: fys --- src/cmd/src/datanode/parquetbench.rs | 2 +- src/mito2/src/cache.rs | 20 +- src/mito2/src/compaction/reader.rs | 33 +- src/mito2/src/engine/scan_test.rs | 50 +-- src/mito2/src/memtable/bulk/context.rs | 4 +- src/mito2/src/read/compat.rs | 50 +-- src/mito2/src/read/flat_projection.rs | 83 ++--- src/mito2/src/read/last_row.rs | 15 +- src/mito2/src/read/projection.rs | 3 +- src/mito2/src/read/range_cache.rs | 4 +- src/mito2/src/read/read_columns.rs | 125 +++---- src/mito2/src/read/scan_region.rs | 327 ++++++++++++------ src/mito2/src/sst/parquet/file_range.rs | 9 +- src/mito2/src/sst/parquet/flat_format.rs | 79 ++++- src/mito2/src/sst/parquet/format.rs | 154 +++++---- src/mito2/src/sst/parquet/json_align/mod.rs | 2 - .../src/sst/parquet/json_align/schema.rs | 250 ------------- src/mito2/src/sst/parquet/prefilter.rs | 30 +- src/mito2/src/sst/parquet/reader.rs | 16 +- .../src/optimizer/json_type_concretize.rs | 91 +++-- src/table/src/table/adapter.rs | 5 + .../common/types/json/json2_limit.result | 69 ++++ .../common/types/json/json2_limit.sql | 39 +++ 23 files changed, 729 insertions(+), 731 deletions(-) delete mode 100644 src/mito2/src/sst/parquet/json_align/schema.rs diff --git a/src/cmd/src/datanode/parquetbench.rs b/src/cmd/src/datanode/parquetbench.rs index 301c6440ec1..2b46a775f76 100644 --- a/src/cmd/src/datanode/parquetbench.rs +++ b/src/cmd/src/datanode/parquetbench.rs @@ -554,7 +554,7 @@ async fn run_flat_prune_iteration( let reader_builder = ParquetReaderBuilder::new(table_dir, path_type, file_handle, object_store) .expected_metadata(Some(region_meta)) .cache(CacheStrategy::Disabled) - .projection(projection.map(ReadColumns::from_deduped_column_ids)); + .projection(projection.map(ReadColumns::new)); let mut reader_metrics = ReaderMetrics::default(); let start = Instant::now(); let mut stats = IterationStats::default(); diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 27db7daa5ab..ac1f5b66918 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -49,7 +49,7 @@ use puffin::puffin_manager::cache::{PuffinMetadataCache, PuffinMetadataCacheRef} use smallvec::SmallVec; use snafu::{OptionExt, ResultExt}; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; -use store_api::storage::{ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector}; +use store_api::storage::{ColumnId, ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector}; use crate::cache::cache_size::parquet_meta_size; use crate::cache::file_cache::{FileType, IndexKey}; @@ -65,6 +65,7 @@ use crate::memtable::record_batch_estimated_size; use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, CACHE_HIT, CACHE_MISS}; use crate::read::Batch; use crate::read::range_cache::{RangeScanCacheKey, RangeScanCacheValue}; +use crate::read::read_columns::JsonTargetTypes; use crate::sst::file::{RegionFileId, RegionIndexId}; use crate::sst::parquet::PARQUET_METADATA_KEY; use crate::sst::parquet::read_columns::ParquetReadColumns; @@ -2025,6 +2026,11 @@ pub struct SelectorResultValue { pub result: SelectorResult, /// The read columns of rows. pub read_cols: ParquetReadColumns, + /// JSON2 target types used by flat-format reads. + /// + /// JSON2 projection is query-driven; the same parquet columns can produce + /// different cached batches under different type hints. + pub json_target_types: JsonTargetTypes, } impl SelectorResultValue { @@ -2033,6 +2039,7 @@ impl SelectorResultValue { SelectorResultValue { result: SelectorResult::PrimaryKey(result), read_cols, + json_target_types: Arc::default(), } } @@ -2040,21 +2047,26 @@ impl SelectorResultValue { pub fn new_flat( result: Vec, read_cols: ParquetReadColumns, + json_target_types: JsonTargetTypes, ) -> SelectorResultValue { SelectorResultValue { result: SelectorResult::Flat(result), read_cols, + json_target_types, } } /// Returns memory used by the value (estimated). fn estimated_size(&self) -> usize { - match &self.result { + let result_size: usize = match &self.result { SelectorResult::PrimaryKey(batches) => { batches.iter().map(|batch| batch.memory_size()).sum() } SelectorResult::Flat(batches) => batches.iter().map(record_batch_estimated_size).sum(), - } + }; + result_size + + self.json_target_types.len() + * (mem::size_of::() + mem::size_of::()) } } @@ -2752,7 +2764,7 @@ mod tests { region_id: RegionId::new(1, 1), row_groups: vec![(FileId::random(), 0)], scan: ScanRequestFingerprintBuilder { - read_columns: ReadColumns::from_deduped_column_ids(std::iter::empty()), + read_columns: ReadColumns::new(std::iter::empty()), read_column_types: vec![], filters: vec!["tag_0 = 1".to_string()], time_filters: vec![], diff --git a/src/mito2/src/compaction/reader.rs b/src/mito2/src/compaction/reader.rs index 89b70f89c0e..f862d11e1a9 100644 --- a/src/mito2/src/compaction/reader.rs +++ b/src/mito2/src/compaction/reader.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; use common_time::Timestamp; @@ -118,15 +118,28 @@ impl CompactionSstReaderBuilder<'_> { }; let projection = (0..self.metadata.column_metadatas.len()).collect(); - let read_columns = ReadColumns::from_deduped_column_ids( - self.metadata.column_metadatas.iter().map(|x| x.column_id), - ); - let mapper = FlatProjectionMapper::new_with_read_columns( - &self.metadata, - projection, - read_columns, - json_type_hint.as_ref(), - )?; + let read_column_ids = self + .metadata + .column_metadatas + .iter() + .map(|x| x.column_id) + .collect::>(); + let json_target_types = json_type_hint + .as_ref() + .map(|hint| { + hint.iter() + .filter_map(|(col_name, json_type)| { + self.metadata + .column_by_name(col_name) + .map(|col| (col.column_id, json_type.clone())) + }) + .collect::>() + }) + .unwrap_or_default(); + let read_columns = + ReadColumns::new(read_column_ids).with_json_target_types(json_target_types); + let mapper = + FlatProjectionMapper::new_with_read_columns(&self.metadata, projection, read_columns)?; let mut scan_input = ScanInput::new(self.sst_layer, mapper) .with_files(self.inputs.to_vec()) diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index b9eed37cb30..4baa707f4e7 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -38,7 +38,7 @@ use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution}; use crate::config::MitoConfig; use crate::error::Error; -use crate::read::read_columns::{ReadColumn, ReadColumns}; +use crate::read::read_columns::ReadColumns; use crate::read::scan_region::Scanner; use crate::test_util; use crate::test_util::{CreateRequestBuilder, TestEnv}; @@ -108,18 +108,19 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul }; assert_eq!( seq_scan.input().read_cols, - ReadColumns::from_deduped_column_ids([1, 0]) + ReadColumns::new([1, 0]) + .with_json_target_types(BTreeMap::from([(1, JsonNativeType::Variant)])) ); let stream = scanner.scan().await?; let batches = RecordBatches::try_collect(stream).await?; let expected = r#" -+------------------------------------------+-------+ -| field_0 | tag_0 | -+------------------------------------------+-------+ -| {a: {x: 10, y: ignored-a}, b: ignored-b} | tag-1 | -| {a: {x: 20, y: ignored-c}, b: ignored-d} | tag-2 | -+------------------------------------------+-------+ ++------------------------------------------------+-------+ +| field_0 | tag_0 | ++------------------------------------------------+-------+ +| {"a":{"x":10,"y":"ignored-a"},"b":"ignored-b"} | tag-1 | +| {"a":{"x":20,"y":"ignored-c"},"b":"ignored-d"} | tag-2 | ++------------------------------------------------+-------+ "#; assert_eq!(batches.pretty_print()?, expected.trim()); @@ -149,19 +150,8 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul // whole JSON2 struct. tag_0 is still read as a normal root column. assert_eq!( seq_scan.input().read_cols, - ReadColumns { - cols: vec![ - ReadColumn::new( - 1, - vec![vec![ - "field_0".to_string(), - "a".to_string(), - "x".to_string() - ]] - ), - ReadColumn::new(0, vec![]), - ] - } + ReadColumns::new([1, 0]) + .with_json_target_types(BTreeMap::from([(1, json_type_hint["field_0"].clone())])) ); // The scanner should still return a valid RecordBatch in the requested logical projection. @@ -189,10 +179,7 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul let Scanner::Seq(seq_scan) = &scanner else { unreachable!(); }; - assert_eq!( - seq_scan.input().read_cols, - ReadColumns::from_deduped_column_ids([0]) - ); + assert_eq!(seq_scan.input().read_cols, ReadColumns::new([0])); let stream = scanner.scan().await?; let batches = RecordBatches::try_collect(stream).await?; @@ -205,6 +192,19 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul +-------+ "#; assert_eq!(batches.pretty_print()?, expected.trim()); + + // JSON type hints are derived from the whole query plan, so they may include + // columns unrelated to this scan. Mito2 should ignore hints for unknown + // columns (`other_json`) instead of failing. + let request = ScanRequest { + projection: Some(vec![0]), + json_type_hint: HashMap::from([("other_json".to_string(), JsonNativeType::String)]), + ..Default::default() + }; + let scanner = engine.scanner(region_id, request).await?; + let stream = scanner.scan().await?; + let batches = RecordBatches::try_collect(stream).await?; + assert_eq!(batches.pretty_print()?, expected.trim()); Ok(()) } diff --git a/src/mito2/src/memtable/bulk/context.rs b/src/mito2/src/memtable/bulk/context.rs index af38ef66c14..994d3a027d4 100644 --- a/src/mito2/src/memtable/bulk/context.rs +++ b/src/mito2/src/memtable/bulk/context.rs @@ -72,9 +72,9 @@ impl BulkIterContext { let codec = build_primary_key_codec(®ion_metadata); let read_cols = if let Some(col_ids) = projection { - ReadColumns::from_deduped_column_ids(col_ids.iter().copied()) + ReadColumns::new(col_ids.iter().copied()) } else { - ReadColumns::from_deduped_column_ids( + ReadColumns::new( region_metadata .column_metadatas .iter() diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index af0bdc58c53..3fa625694c6 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -25,11 +25,9 @@ use datatypes::arrow::compute::{TakeOptions, take}; use datatypes::arrow::datatypes::{FieldRef, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; -use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::DataType; use datatypes::value::Value; use datatypes::vectors::VectorRef; -use datatypes::vectors::json::array::JsonArray; use mito_codec::row_converter::{ CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec, build_primary_key_codec_with_fields, @@ -40,8 +38,8 @@ use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::ColumnId; use crate::error::{ - CompatReaderSnafu, ComputeArrowSnafu, ConvertValueSnafu, CreateDefaultSnafu, DecodeSnafu, - EncodeSnafu, NewRecordBatchSnafu, Result, UnexpectedSnafu, UnsupportedOperationSnafu, + CompatReaderSnafu, ComputeArrowSnafu, CreateDefaultSnafu, DecodeSnafu, EncodeSnafu, + NewRecordBatchSnafu, Result, UnexpectedSnafu, UnsupportedOperationSnafu, }; use crate::read::flat_projection::{FlatProjectionMapper, flat_projected_columns}; use crate::sst::parquet::flat_format::{FlatReadFormat, primary_key_column_index}; @@ -99,20 +97,12 @@ impl FlatCompatBatch { let actual = read_format.metadata(); let format_projection = read_format.format_projection(); let mut actual_schema = flat_projected_columns(actual, format_projection); - if read_format - .arrow_schema() - .fields() - .iter() - .any(is_json2_extension_type) - { - for field in read_format.arrow_schema().fields() { - if is_json2_extension_type(field) - && let Some(column_id) = - actual.column_by_name(field.name()).map(|x| x.column_id) - && let Some(i) = actual_schema.iter().position(|x| x.0 == column_id) - { - actual_schema[i].1 = ConcreteDataType::from_arrow_type(field.data_type()); - } + for (column_id, target_type) in read_format.json_target_types().iter() { + if let Some(i) = actual_schema + .iter() + .position(|(actual_column_id, _)| actual_column_id == column_id) + { + actual_schema[i].1 = ConcreteDataType::json2(target_type.clone()); } } @@ -270,16 +260,9 @@ impl FlatCompatBatch { let old_column = batch.column(*pos); if let Some(ty) = cast_type { - let casted = if let Some(json_type) = ty.as_json() - && json_type.is_json2() - { - JsonArray::from(old_column) - .project_to(&json_type.as_arrow_type()) - .context(ConvertValueSnafu)? - } else { + let casted = datatypes::arrow::compute::cast(old_column, &ty.as_arrow_type()) - .context(ComputeArrowSnafu)? - }; + .context(ComputeArrowSnafu)?; Ok(casted) } else { Ok(old_column.clone()) @@ -777,7 +760,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([0, 1, 2, 3]), + ReadColumns::new([0, 1, 2, 3]), None, "test", false, @@ -864,13 +847,12 @@ mod tests { let mapper = FlatProjectionMapper::new_with_read_columns( &expected_metadata, vec![1, 2], - ReadColumns::from_deduped_column_ids([1, 2, 3]), - None, + ReadColumns::new([1, 2, 3]), ) .unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([1, 2, 3]), + ReadColumns::new([1, 2, 3]), None, "test", false, @@ -959,7 +941,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([0, 1, 2, 3]), + ReadColumns::new([0, 1, 2, 3]), None, "test", false, @@ -1041,7 +1023,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata, - ReadColumns::from_deduped_column_ids([0, 1, 2]), + ReadColumns::new([0, 1, 2]), None, "test", false, @@ -1109,7 +1091,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([0, 2, 3]), + ReadColumns::new([0, 2, 3]), None, "test", true, diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index fbd6088f670..e631cf45db4 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -14,7 +14,6 @@ //! Utilities for projection on flat format. -use std::collections::HashMap; use std::sync::Arc; use api::v1::SemanticType; @@ -27,7 +26,7 @@ use datatypes::arrow::array::Array; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field}; use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::{ConcreteDataType, DataType}; -use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; +use datatypes::schema::{Schema, SchemaRef}; use datatypes::types::JsonType; use datatypes::types::json_type::JsonNativeType; use datatypes::value::Value; @@ -86,8 +85,8 @@ impl FlatProjectionMapper { ) -> Result { let projection: Vec<_> = projection.into_iter().collect(); let read_column_ids = read_column_ids_from_projection(metadata, &projection)?; - let read_cols = ReadColumns::from_deduped_column_ids(read_column_ids); - Self::new_with_read_columns(metadata, projection, read_cols, None) + let read_cols = ReadColumns::new(read_column_ids); + Self::new_with_read_columns(metadata, projection, read_cols) } /// Returns a new mapper with output projection and explicit read columns. @@ -95,7 +94,6 @@ impl FlatProjectionMapper { metadata: &RegionMetadataRef, projection: Vec, read_cols: ReadColumns, - json_type_hint: Option<&HashMap>, ) -> Result { // If the original projection is empty. let is_empty_projection = projection.is_empty(); @@ -115,7 +113,11 @@ impl FlatProjectionMapper { output_col_ids.push(col.column_id); let mut schema = col.column_schema.clone(); - maybe_concretize_json2_datatype(&mut schema, json_type_hint); + if let Some(data_type) = + json2_read_datatype(col.column_id, &schema.data_type, &read_cols) + { + schema.data_type = data_type; + } col_schemas.push(schema); } @@ -124,6 +126,7 @@ impl FlatProjectionMapper { // TODO(yingwen): Support different flat schema options. let format_projection = FormatProjection::compute_format_projection( + metadata, &id_to_index, // All columns with internal columns. metadata.column_metadatas.len() + 3, @@ -133,21 +136,8 @@ impl FlatProjectionMapper { let mut batch_schema = flat_projected_columns(metadata, &format_projection); for (column_id, data_type) in batch_schema.iter_mut() { - if let Some(json_type) = data_type.as_json() - && json_type.is_json2() - { - if let Some(concretized) = metadata - .column_by_id(*column_id) - .and_then(|metadata| { - json_type_hint.and_then(|x| x.get(&metadata.column_schema.name).cloned()) - }) - .map(ConcreteDataType::json2) - { - *data_type = concretized; - } else if is_empty_json2_type(json_type) { - // see `merge_scan::maybe_amend_json2_field` - *data_type = ConcreteDataType::json2(JsonNativeType::Variant); - } + if let Some(updated) = json2_read_datatype(*column_id, data_type, &read_cols) { + *data_type = updated; } } @@ -401,24 +391,25 @@ impl FlatProjectionMapper { } } -fn maybe_concretize_json2_datatype( - schema: &mut ColumnSchema, - json_type_hint: Option<&HashMap>, -) { - if let Some(json_type) = schema.data_type.as_json() - && json_type.is_json2() - { - if let Some(concretized) = json_type_hint - .and_then(|x| x.get(&schema.name)) - .cloned() - .map(ConcreteDataType::json2) - { - schema.data_type = concretized; - } else if is_empty_json2_type(json_type) { - // see `merge_scan::maybe_amend_json2_field` - schema.data_type = ConcreteDataType::json2(JsonNativeType::Variant); - } +fn json2_read_datatype( + column_id: ColumnId, + data_type: &ConcreteDataType, + read_cols: &ReadColumns, +) -> Option { + let json_type = data_type.as_json()?; + if !json_type.is_json2() { + return None; } + + if let Some(concretized) = read_cols.json_target_type(column_id).cloned() { + return Some(ConcreteDataType::json2(concretized)); + } + + if is_empty_json2_type(json_type) { + return Some(ConcreteDataType::json2(JsonNativeType::Variant)); + } + + None } fn is_empty_json2_type(json_type: &JsonType) -> bool { @@ -547,9 +538,8 @@ impl CompactionProjectionMapper { .collect::>(); let read_col_ids = metadata.column_metadatas.iter().map(|col| col.column_id); - let read_cols = ReadColumns::from_deduped_column_ids(read_col_ids); - let mapper = - FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols, None)?; + let read_cols = ReadColumns::new(read_col_ids); + let mapper = FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols)?; let assembler = DfBatchAssembler::new(mapper.output_schema()); Ok(Self { mapper, assembler }) @@ -617,7 +607,6 @@ impl DfBatchAssembler { #[cfg(test)] mod tests { use datatypes::schema::ColumnSchema; - use datatypes::types::json_type::JsonObjectType; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::storage::RegionId; @@ -687,18 +676,10 @@ mod tests { #[test] fn test_json_type_hint_does_not_concretize_legacy_json() { let metadata = metadata_with_legacy_json(); - let hint = HashMap::from([( - "j".to_string(), - JsonNativeType::Object(JsonObjectType::from([( - "a".to_string(), - JsonNativeType::i64(), - )])), - )]); let mapper = FlatProjectionMapper::new_with_read_columns( &metadata, vec![0, 1], - ReadColumns::from_deduped_column_ids([0, 1]), - Some(&hint), + ReadColumns::new([0, 1]), ) .unwrap(); diff --git a/src/mito2/src/read/last_row.rs b/src/mito2/src/read/last_row.rs index 6e1a62d806e..c2316525f70 100644 --- a/src/mito2/src/read/last_row.rs +++ b/src/mito2/src/read/last_row.rs @@ -29,6 +29,7 @@ use crate::cache::{ selector_result_cache_hit, selector_result_cache_miss, }; use crate::error::{ComputeArrowSnafu, Result}; +use crate::read::read_columns::JsonTargetTypes; use crate::read::{ Batch, BatchReader, BoxedBatchReader, BoxedRecordBatchStream, timestamp_array_to_i64_slice, }; @@ -137,6 +138,7 @@ impl FlatRowGroupLastRowCachedReader { row_group_idx: usize, cache_strategy: CacheStrategy, read_cols: &ParquetReadColumns, + json_target_types: JsonTargetTypes, reader: FlatRowGroupReader, ) -> Self { let key = SelectorResultKey { @@ -148,13 +150,14 @@ impl FlatRowGroupLastRowCachedReader { if let Some(value) = cache_strategy.get_selector_result(&key) { let is_flat = matches!(&value.result, SelectorResult::Flat(_)); let schema_matches = value.read_cols == *read_cols; - if is_flat && schema_matches { + let json_target_types_matches = value.json_target_types == json_target_types; + if is_flat && schema_matches && json_target_types_matches { Self::new_hit(value) } else { - Self::new_miss(key, read_cols, reader, cache_strategy) + Self::new_miss(key, read_cols, json_target_types, reader, cache_strategy) } } else { - Self::new_miss(key, read_cols, reader, cache_strategy) + Self::new_miss(key, read_cols, json_target_types, reader, cache_strategy) } } @@ -174,6 +177,7 @@ impl FlatRowGroupLastRowCachedReader { fn new_miss( key: SelectorResultKey, read_cols: &ParquetReadColumns, + json_target_types: JsonTargetTypes, reader: FlatRowGroupReader, cache_strategy: CacheStrategy, ) -> Self { @@ -181,6 +185,7 @@ impl FlatRowGroupLastRowCachedReader { Self::Miss(FlatRowGroupLastRowReader::new( key, read_cols.clone(), + json_target_types, reader, cache_strategy, )) @@ -260,6 +265,7 @@ pub(crate) struct FlatRowGroupLastRowReader { yielded_batches: Vec, cache_strategy: CacheStrategy, read_cols: ParquetReadColumns, + json_target_types: JsonTargetTypes, /// Accumulates small selector-output batches before concatenating. pending: BatchBuffer, } @@ -268,6 +274,7 @@ impl FlatRowGroupLastRowReader { fn new( key: SelectorResultKey, read_cols: ParquetReadColumns, + json_target_types: JsonTargetTypes, reader: FlatRowGroupReader, cache_strategy: CacheStrategy, ) -> Self { @@ -278,6 +285,7 @@ impl FlatRowGroupLastRowReader { yielded_batches: vec![], cache_strategy, read_cols, + json_target_types, pending: BatchBuffer::new(), } } @@ -326,6 +334,7 @@ impl FlatRowGroupLastRowReader { let value = Arc::new(SelectorResultValue::new_flat( batches, self.read_cols.clone(), + self.json_target_types.clone(), )); self.cache_strategy.put_selector_result(self.key, value); } diff --git a/src/mito2/src/read/projection.rs b/src/mito2/src/read/projection.rs index 27fe76345b6..c6d4baee3ac 100644 --- a/src/mito2/src/read/projection.rs +++ b/src/mito2/src/read/projection.rs @@ -295,8 +295,7 @@ mod tests { let mapper = FlatProjectionMapper::new_with_read_columns( &metadata, vec![4, 1], - ReadColumns::from_deduped_column_ids([4, 1, 3]), - None, + ReadColumns::new([4, 1, 3]), ) .unwrap(); assert_eq!(&[4, 1, 3], mapper.read_columns().column_ids().as_slice()); diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index 036a574c1f5..ee0e421f86e 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -852,7 +852,7 @@ pub fn bench_cache_flat_range_stream( let cache_strategy = CacheStrategy::EnableAll(cache_manager); let fingerprint = ScanRequestFingerprintBuilder { - read_columns: ReadColumns::from_deduped_column_ids(std::iter::empty()), + read_columns: ReadColumns::new(std::iter::empty()), read_column_types: vec![], filters: vec![], time_filters: vec![], @@ -916,7 +916,7 @@ mod tests { filter_deleted: bool, partition_expr_version: u64, ) -> ScanRequestFingerprint { - let read_columns = ReadColumns::from_deduped_column_ids([1, 2]); + let read_columns = ReadColumns::new([1, 2]); ScanRequestFingerprintBuilder { read_columns, read_column_types: vec![None, None], diff --git a/src/mito2/src/read/read_columns.rs b/src/mito2/src/read/read_columns.rs index 21699b8c439..42a9f71ddcb 100644 --- a/src/mito2/src/read/read_columns.rs +++ b/src/mito2/src/read/read_columns.rs @@ -12,119 +12,72 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::BTreeMap; use std::mem; +use std::sync::Arc; -use store_api::storage::{ColumnId, NestedPath}; +use datatypes::types::json_type::JsonNativeType; +use store_api::storage::ColumnId; + +pub(crate) type JsonTargetTypes = Arc>; /// Logical columns to read from a region. /// -/// Read columns describe which logical columns and nested fields should be read -/// from storage. Each read column is identified by its [`ColumnId`], -/// which represents the root column in the storage schema. -/// -/// Nested fields under the column are specified by [`NestedPath`] entries. -/// Each path includes the root column name as its first element. -/// -/// For example, assume column id `9` corresponds to a root column named `j` -/// with nested fields: -/// -/// ```text -/// j -/// ├── a -/// └── b -/// └── c -/// ``` -/// -/// The following SQL: -/// -/// SELECT j.a, j.b.c FROM t -/// -/// may produce read columns like: -/// -/// ```text -/// ReadColumn { -/// column_id: 9, -/// nested_paths: [ -/// ["j", "a"], -/// ["j", "b", "c"], -/// ] -/// } -/// ``` -/// -/// If `nested_paths` is empty, the whole column will be read. +/// Read columns describe which logical root columns should be read from storage. +/// JSON2 columns can carry query-time target types that are later translated to +/// physical nested parquet paths by the parquet reader. #[derive(Debug, Default, Clone, PartialEq, Eq, Hash)] pub struct ReadColumns { - pub cols: Vec, + pub col_ids: Vec, + json_target_types: JsonTargetTypes, } impl ReadColumns { - pub fn from_deduped_column_ids(column_ids: I) -> Self + /// Creates read columns from logical column ids. + /// + /// This preserves the input order and duplicate ids. + pub fn new(col_ids: I) -> Self where I: IntoIterator, { - let cols = column_ids - .into_iter() - .map(|col_id| ReadColumn::new(col_id, vec![])) - .collect(); - ReadColumns { cols } + Self { + col_ids: col_ids.into_iter().collect(), + json_target_types: Arc::default(), + } + } + + pub fn with_json_target_types( + mut self, + json_target_types: BTreeMap, + ) -> Self { + self.json_target_types = Arc::new(json_target_types); + self } pub fn is_empty(&self) -> bool { - self.cols.is_empty() + self.col_ids.is_empty() } pub fn column_ids_iter(&self) -> impl Iterator + '_ { - self.cols.iter().map(|column| column.column_id) + self.col_ids.iter().copied() } pub fn column_ids(&self) -> Vec { self.column_ids_iter().collect() } - pub fn columns(&self) -> &[ReadColumn] { - &self.cols + pub fn json_target_types(&self) -> &JsonTargetTypes { + &self.json_target_types + } + + pub fn json_target_type(&self, column_id: ColumnId) -> Option<&JsonNativeType> { + self.json_target_types.get(&column_id) } pub fn estimated_size(&self) -> usize { - self.cols.capacity() * mem::size_of::() - + self - .cols - .iter() - .map(ReadColumn::estimated_size) - .sum::() - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct ReadColumn { - pub column_id: ColumnId, - /// Nested field paths under this column. - /// Empty means reading the whole column. - pub nested_paths: Vec, -} - -impl ReadColumn { - pub fn new(column_id: ColumnId, nested_paths: Vec) -> Self { - Self { - column_id, - nested_paths, - } - } - - pub fn nested_paths(&self) -> &[NestedPath] { - &self.nested_paths - } - - pub fn estimated_size(&self) -> usize { - mem::size_of::() - + self.nested_paths.capacity() * mem::size_of::() - + self - .nested_paths - .iter() - .map(|path| { - path.capacity() * mem::size_of::() - + path.iter().map(|node| node.capacity()).sum::() - }) - .sum::() + self.col_ids.capacity() * mem::size_of::() + + self.col_ids.len() * mem::size_of::() + + self.json_target_types.len() + * (mem::size_of::() + mem::size_of::()) } } diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 250f0d7fc27..c0340464ecc 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -14,7 +14,7 @@ //! Scans a region according to the scan request. -use std::collections::HashSet; +use std::collections::{BTreeMap, HashSet}; use std::fmt; use std::num::NonZeroU64; use std::sync::Arc; @@ -35,18 +35,18 @@ use datafusion_expr::Expr; use datafusion_expr::utils::expr_to_columns; use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use datatypes::extension::json::is_json2_extension_type; +use datatypes::prelude::ConcreteDataType; use datatypes::types::json_type::JsonNativeType; use datatypes::value::timestamp_to_scalar_value; use futures::StreamExt; -use itertools::Itertools; use partition::expr::PartitionExpr; use smallvec::SmallVec; -use snafu::{OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::region_engine::{PartitionRange, RegionScannerRef}; use store_api::storage::{ - ColumnId, NestedPath, RegionId, ScanRequest, SequenceNumber, SequenceRange, - TimeSeriesDistribution, TimeSeriesRowSelector, + ColumnId, RegionId, ScanRequest, SequenceNumber, SequenceRange, TimeSeriesDistribution, + TimeSeriesRowSelector, }; use table::predicate::{Predicate, build_time_range_predicate, extract_time_range_from_expr}; use tokio::sync::{Semaphore, mpsc}; @@ -64,7 +64,7 @@ use crate::read::compat::{self, FlatCompatBatch}; use crate::read::flat_projection::FlatProjectionMapper; use crate::read::range::{FileRangeBuilder, MemRangeBuilder, RangeMeta, RowGroupIndex}; use crate::read::range_cache::{ScanRequestFingerprint, implied_time_range_from_exprs}; -use crate::read::read_columns::{ReadColumn, ReadColumns}; +use crate::read::read_columns::ReadColumns; use crate::read::seq_scan::SeqScan; use crate::read::series_scan::SeriesScan; use crate::read::stream::ScanBatchStream; @@ -420,23 +420,7 @@ impl ScanRegion { let read_col_ids = self.build_read_col_ids(self.request.projection.as_deref(), &predicate)?; - - // Narrow JSON2 columns to avoid reading unnecessary nested fields. - // - // `read_col_ids` selects the root columns required by the projection and predicates, - // while nested projection is currently only applied to JSON2 columns, whose type hints - // further narrow them to the requested nested fields. - let has_structured_json = metadata - .schema - .arrow_schema() - .fields() - .iter() - .any(is_json2_extension_type); - let read_cols = if has_structured_json { - self.read_columns_with_json_type_hint(&read_col_ids) - } else { - ReadColumns::from_deduped_column_ids(read_col_ids.iter().copied()) - }; + let read_cols = self.build_read_columns(&read_col_ids)?; // The mapper always computes projected column ids as the schema of SSTs may change. let projection = self @@ -444,23 +428,7 @@ impl ScanRegion { .projection .clone() .unwrap_or_else(|| (0..metadata.column_metadatas.len()).collect()); - let json_type_hint = has_structured_json - .then_some(&self.request.json_type_hint) - .inspect(|json_type_hint| { - debug!( - "Concretized JSON type: {{{}}}", - json_type_hint - .iter() - .map(|(k, v)| format!("{}: {}", k, v)) - .join(", ") - ); - }); - let mapper = FlatProjectionMapper::new_with_read_columns( - metadata, - projection, - read_cols, - json_type_hint, - )?; + let mapper = FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols)?; let mapper = if self.request.preserve_pk_dictionary_encoding { mapper.with_pk_dictionary_encoding() } else { @@ -609,7 +577,8 @@ impl ScanRegion { Ok(input) } - /// Builds the ordered root column ids required by the pushed-down projection and predicate. + /// Builds the deduplicated root column ids required by the projection and + /// predicate. fn build_read_col_ids( &self, projection: Option<&[usize]>, @@ -688,27 +657,50 @@ impl ScanRegion { Ok(read_col_ids) } - /// Builds read columns with nested paths derived from JSON type hints. - fn read_columns_with_json_type_hint(&self, col_ids: &[ColumnId]) -> ReadColumns { - let cols = col_ids + /// Builds logical read columns and attaches JSON2 target types when needed. + /// + /// The behavior about JSON2 is as follows: + /// - JSON2 columns without hints use Variant to read the whole column. + /// - Hints targeting non-JSON2 read columns are rejected. + fn build_read_columns(&self, col_ids: &[ColumnId]) -> Result { + let metadata = &self.version.metadata; + let json_type_hint = &self.request.json_type_hint; + + let has_json2 = metadata + .schema + .arrow_schema() + .fields() .iter() - .map(|&col_id| { - let nested_paths = self - .version - .metadata - .column_by_id(col_id) - .and_then(|column| { - let col_name = &column.column_schema.name; - self.request - .json_type_hint - .get(col_name) - .map(|json_type| json_nested_paths(col_name, json_type)) - }) - .unwrap_or_default(); - ReadColumn::new(col_id, nested_paths) - }) - .collect(); - ReadColumns { cols } + .any(is_json2_extension_type); + + if !has_json2 && json_type_hint.is_empty() { + return Ok(ReadColumns::new(col_ids.iter().copied())); + } + + let mut json_target_types = BTreeMap::new(); + for &col_id in col_ids { + let Some(col) = metadata.column_by_id(col_id) else { + continue; + }; + let col_name = &col.column_schema.name; + let hint = json_type_hint.get(col_name); + if !col.column_schema.data_type.is_json2() { + ensure!( + hint.is_none(), + InvalidRequestSnafu { + region_id: metadata.region_id, + reason: format!( + "JSON type hint targets non-JSON2 column {} (id: {}, type: {})", + col_name, col_id, col.column_schema.data_type + ), + } + ); + continue; + } + let target_type = hint.cloned().unwrap_or(JsonNativeType::Variant); + json_target_types.insert(col_id, target_type); + } + Ok(ReadColumns::new(col_ids.iter().copied()).with_json_target_types(json_target_types)) } fn region_id(&self) -> RegionId { @@ -1583,30 +1575,6 @@ fn pre_filter_mode(append_mode: bool, merge_mode: MergeMode) -> PreFilterMode { } } -fn json_nested_paths(column_name: &str, json_type: &JsonNativeType) -> Vec { - let mut paths = Vec::new(); - let mut current = vec![column_name.to_string()]; - collect_json_nested_paths(json_type, &mut current, &mut paths); - paths -} - -fn collect_json_nested_paths( - json_type: &JsonNativeType, - current: &mut NestedPath, - paths: &mut Vec, -) { - match json_type { - JsonNativeType::Object(fields) if !fields.is_empty() => { - for (field, child) in fields { - current.push(field.clone()); - collect_json_nested_paths(child, current, paths); - current.pop(); - } - } - _ => paths.push(current.clone()), - } -} - /// Output of [build_scan_fingerprint]: the cache fingerprint plus the derived /// implied time range used to decide whether the cache key can drop the time /// predicates for a given partition (see `build_range_cache_key`). @@ -1703,9 +1671,15 @@ pub(crate) fn build_scan_fingerprint(input: &ScanInput) -> Option Result<()> { - let hint = JsonNativeType::Object(JsonObjectType::from([ - ("a".to_string(), JsonNativeType::i64()), - ( - "b".to_string(), - JsonNativeType::Object(JsonObjectType::from([( - "c".to_string(), - JsonNativeType::String, - )])), - ), - ])); - - fn nested_path(parts: &[&str]) -> NestedPath { - parts.iter().map(|part| part.to_string()).collect() - } - - assert_eq!( - json_nested_paths("j", &hint), - vec![nested_path(&["j", "a"]), nested_path(&["j", "b", "c"])] - ); - Ok(()) - } - #[tokio::test] async fn test_build_scan_fingerprint_for_eligible_scan() { let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); @@ -2370,6 +2321,158 @@ mod tests { assert_ne!(0, metadata.partition_expr_version); } + #[tokio::test] + async fn test_build_scan_fingerprint_uses_json_target_types() { + let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "k0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "j".to_string(), + ConcreteDataType::json2(JsonNativeType::Variant), + true, + ), + semantic_type: SemanticType::Field, + column_id: 2, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + + let make_input = |target_type| async { + let env = SchedulerEnv::new().await; + let read_cols = ReadColumns::new([0, 1, 2]) + .with_json_target_types(BTreeMap::from([(2, target_type)])); + let mapper = + FlatProjectionMapper::new_with_read_columns(&metadata, vec![0, 1, 2], read_cols) + .unwrap(); + let predicate = + PredicateGroup::new(metadata.as_ref(), &[col("k0").eq(lit("foo"))]).unwrap(); + let file = FileHandle::new( + FileMeta::default(), + Arc::new(crate::sst::file_purger::NoopFilePurger), + ); + ScanInput::new(env.access_layer.clone(), mapper) + .with_predicate(predicate) + .with_cache(CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .range_result_cache_size(1024) + .build(), + ))) + .with_files(vec![file]) + }; + + let int_target = JsonNativeType::i64(); + let string_target = JsonNativeType::String; + let int_fingerprint = build_scan_fingerprint(&make_input(int_target.clone()).await) + .unwrap() + .fingerprint; + let string_fingerprint = build_scan_fingerprint(&make_input(string_target).await) + .unwrap() + .fingerprint; + + assert_ne!(int_fingerprint, string_fingerprint); + assert_eq!( + Some(&Some(ConcreteDataType::json2(int_target))), + int_fingerprint.read_column_types().get(2) + ); + } + + #[tokio::test] + async fn test_scan_input_rejects_json_type_hint_for_non_json2_column() { + let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "k0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "j".to_string(), + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::from([( + "a".to_string(), + JsonNativeType::i64(), + )]))), + true, + ), + semantic_type: SemanticType::Field, + column_id: 2, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "v0".to_string(), + ConcreteDataType::int64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 3, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + metadata.clone(), + Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), + 0, + None, + )); + let version = Arc::new( + crate::region::version::VersionBuilder::new(metadata.clone(), mutable).build(), + ); + let env = SchedulerEnv::new().await; + let request = ScanRequest { + projection: Some(vec![0, 1, 2, 3]), + json_type_hint: std::collections::HashMap::from([( + "v0".to_string(), + JsonNativeType::i64(), + )]), + ..Default::default() + }; + + let err = ScanRegion::new( + version, + env.access_layer.clone(), + request, + CacheStrategy::Disabled, + ) + .scan_input() + .await; + let Err(err) = err else { + panic!("scan input should reject JSON type hint for non-JSON2 column"); + }; + + assert!(err.to_string().contains("non-JSON2 column v0")); + } + #[test] fn test_update_dyn_filters_with_empty_base_predicates() { let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index 2b9e800fb23..a3a94ef25df 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -235,6 +235,7 @@ impl FileRange { self.row_group_idx, cache_strategy, self.context.read_format().parquet_read_columns(), + self.context.read_format().json_target_types().clone(), flat_row_group_reader, ); FlatPruneReader::new_with_last_row_reader(self.context.clone(), reader, skip_fields) @@ -827,9 +828,7 @@ mod tests { let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -895,9 +894,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index 237b7c0f2c6..e381ea76f49 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -51,12 +51,11 @@ use crate::error::{ ComputeArrowSnafu, DecodeSnafu, InvalidParquetSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result, }; -use crate::read::read_columns::ReadColumns; +use crate::read::read_columns::{JsonTargetTypes, ReadColumns}; use crate::sst::parquet::format::{ FIXED_POS_COLUMN_NUM, FormatProjection, INTERNAL_COLUMN_NUM, PrimaryKeyArray, PrimaryKeyReadFormat, StatValues, column_null_counts, column_values, }; -use crate::sst::parquet::json_align::align_schema_by_nested_paths; use crate::sst::parquet::read_columns::ParquetReadColumns; use crate::sst::{ FlatSchemaOptions, flat_sst_arrow_schema_column_num, tag_maybe_to_dictionary_field, @@ -181,6 +180,8 @@ pub(crate) fn field_column_start(metadata: &RegionMetadata, num_columns: usize) pub struct FlatReadFormat { /// Sequence number to override the sequence read from the SST. override_sequence: Option, + /// Logical columns requested by this read. + read_cols: ReadColumns, /// Parquet format adapter. parquet_adapter: ParquetAdapter, /// Output schema to wrap binary `__primary_key` back to a dictionary; `None` disables wrapping. @@ -210,22 +211,25 @@ impl FlatReadFormat { // Only skip auto convert when the primary key encoding is sparse. ParquetAdapter::PrimaryKeyToFlat(ParquetPrimaryKeyToFlat::new( metadata, - read_cols, + read_cols.clone(), skip_auto_convert, )) } else { ParquetAdapter::PrimaryKeyToFlat(ParquetPrimaryKeyToFlat::new( - metadata, read_cols, false, + metadata, + read_cols.clone(), + false, )) } } else { let file_schema = file_schema .unwrap_or_else(|| to_flat_sst_arrow_schema(&metadata, &Default::default())); - ParquetAdapter::Flat(ParquetFlat::new(metadata, read_cols, file_schema)) + ParquetAdapter::Flat(ParquetFlat::new(metadata, read_cols.clone(), file_schema)) }; Ok(FlatReadFormat { override_sequence: None, + read_cols, parquet_adapter, pk_dict_wrap_schema: None, }) @@ -294,22 +298,51 @@ impl FlatReadFormat { } } - /// Gets the projected output schema produced by parquet reading. + /// Gets the projected output schema expected by the scan. pub(crate) fn output_arrow_schema(&self) -> Result { - let read_columns = self.parquet_read_columns(); - let projection = read_columns.root_indices(); + let projection = self.parquet_read_columns().root_indices(); let mut schema = self .arrow_schema() .project(projection) .context(ComputeArrowSnafu)?; - if read_columns.has_nested() { - debug_assert_eq!(schema.fields().len(), read_columns.columns().len()); - let nested_paths = read_columns.columns().iter().map(|x| x.nested_paths()); - align_schema_by_nested_paths(&mut schema, nested_paths); + let mut fields = schema.fields().iter().cloned().collect::>(); + for (column_id, target_type) in self.json_target_types().iter() { + let Some(index) = self.parquet_projected_index_by_id(*column_id) else { + continue; + }; + let Some(field) = schema.fields().get(index) else { + continue; + }; + fields[index] = Arc::new( + field + .as_ref() + .clone() + .with_data_type(ConcreteDataType::json2(target_type.clone()).as_arrow_type()), + ); } + schema.fields = fields.into(); Ok(Arc::new(schema)) } + /// Index of a column in the projected schema produced directly by parquet + /// reading, before any primary-key-to-flat conversion. + fn parquet_projected_index_by_id(&self, column_id: ColumnId) -> Option { + match &self.parquet_adapter { + ParquetAdapter::Flat(p) => p + .format_projection + .column_id_to_projected_index + .get(&column_id) + .copied(), + // `format_projection` addresses the post-conversion flat batch here. + // This helper needs the raw primary-key projection used by parquet reading. + ParquetAdapter::PrimaryKeyToFlat(p) => p + .format + .field_id_to_projected_index() + .get(&column_id) + .copied(), + } + } + /// Gets the metadata of the SST. pub(crate) fn metadata(&self) -> &RegionMetadataRef { match &self.parquet_adapter { @@ -326,6 +359,11 @@ impl FlatReadFormat { } } + /// Gets JSON2 target types keyed by column id. + pub(crate) fn json_target_types(&self) -> &JsonTargetTypes { + self.read_cols.json_target_types() + } + /// Gets the projection in the flat format. /// /// When `skip_auto_convert` is enabled (primary-key format read), this returns the @@ -491,6 +529,7 @@ impl ParquetPrimaryKeyToFlat { } else { // Computes the format projection for the new format. let format_projection = FormatProjection::compute_format_projection( + &metadata, &id_to_index, sst_column_num, read_cols.clone(), @@ -540,8 +579,12 @@ impl ParquetFlat { let id_to_index = sst_column_id_indices(&metadata); let sst_column_num = flat_sst_arrow_schema_column_num(&metadata, &FlatSchemaOptions::default()); - let format_projection = - FormatProjection::compute_format_projection(&id_to_index, sst_column_num, read_cols); + let format_projection = FormatProjection::compute_format_projection( + &metadata, + &id_to_index, + sst_column_num, + read_cols, + ); Self { metadata, @@ -841,9 +884,7 @@ impl FlatReadFormat { pub fn new_with_all_columns(metadata: RegionMetadataRef) -> FlatReadFormat { Self::new( Arc::clone(&metadata), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, @@ -965,7 +1006,7 @@ mod tests { .collect(); let mut read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids), + ReadColumns::new(column_ids), None, "test", false, @@ -1045,7 +1086,7 @@ mod tests { let metadata = Arc::new(build_metadata(1, 2, PrimaryKeyEncoding::Dense)); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids([0_u32, 2_u32]), + ReadColumns::new([0_u32, 2_u32]), None, "test", false, diff --git a/src/mito2/src/sst/parquet/format.rs b/src/mito2/src/sst/parquet/format.rs index a6338c1a1c9..f3f0a34a682 100644 --- a/src/mito2/src/sst/parquet/format.rs +++ b/src/mito2/src/sst/parquet/format.rs @@ -39,6 +39,7 @@ use datatypes::arrow::array::{ use datatypes::arrow::datatypes::{SchemaRef, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::DataType; +use datatypes::types::json_type::JsonNativeType; use datatypes::vectors::Helper; use mito_codec::row_converter::{ CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec_with_fields, @@ -47,12 +48,12 @@ use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; use parquet::file::statistics::Statistics; use snafu::{OptionExt, ResultExt, ensure}; use store_api::metadata::{ColumnMetadata, RegionMetadataRef}; -use store_api::storage::{ColumnId, SequenceNumber}; +use store_api::storage::{ColumnId, NestedPath, SequenceNumber}; use crate::error::{ ConvertVectorSnafu, DecodeSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result, }; -use crate::read::read_columns::ReadColumns; +use crate::read::read_columns::{JsonTargetTypes, ReadColumns}; use crate::read::{Batch, BatchBuilder, BatchColumn}; use crate::sst::file::{FileMeta, FileTimeRange}; use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns}; @@ -233,6 +234,7 @@ impl PrimaryKeyReadFormat { let arrow_schema = to_sst_arrow_schema(&metadata); let format_projection = FormatProjection::compute_format_projection( + &metadata, &field_id_to_index, arrow_schema.fields.len(), read_cols, @@ -592,18 +594,21 @@ impl FormatProjection { /// /// `id_to_index` is a mapping from column id to the index of the column in the SST. pub(crate) fn compute_format_projection( + metadata: &RegionMetadataRef, id_to_index: &HashMap, sst_column_num: usize, cols: ReadColumns, ) -> Self { + let json_target_types = cols.json_target_types().clone(); let mut projected_columns: Vec<_> = cols - .cols + .col_ids .into_iter() - .filter_map(|col| { - id_to_index - .get(&col.column_id) - .copied() - .map(|index_of_sst| (col.column_id, index_of_sst, col.nested_paths)) + .filter_map(|col_id| { + id_to_index.get(&col_id).copied().map(|index_of_sst| { + let nested_paths = + json_target_nested_paths(metadata, &json_target_types, col_id); + (col_id, index_of_sst, nested_paths) + }) }) .collect(); // Sorts columns by their indices in the SST. SST uses a bitmap for projection. @@ -681,6 +686,45 @@ impl FormatProjection { } } +fn json_target_nested_paths( + metadata: &RegionMetadataRef, + json_target_types: &JsonTargetTypes, + column_id: ColumnId, +) -> Vec { + let Some(target_type) = json_target_types.get(&column_id) else { + return Vec::new(); + }; + let Some(column) = metadata.column_by_id(column_id) else { + return Vec::new(); + }; + + json_nested_paths(&column.column_schema.name, target_type) +} + +fn json_nested_paths(column_name: &str, json_type: &JsonNativeType) -> Vec { + let mut paths = Vec::new(); + let mut current = vec![column_name.to_string()]; + collect_json_nested_paths(json_type, &mut current, &mut paths); + paths +} + +fn collect_json_nested_paths( + json_type: &JsonNativeType, + current: &mut NestedPath, + paths: &mut Vec, +) { + match json_type { + JsonNativeType::Object(fields) if !fields.is_empty() => { + for (field, child) in fields { + current.push(field.clone()); + collect_json_nested_paths(child, current, paths); + current.pop(); + } + } + _ => paths.push(current.clone()), + } +} + /// Values of column statistics of the SST. /// /// It also distinguishes the case that a column is not found and @@ -710,9 +754,7 @@ impl PrimaryKeyReadFormat { pub fn new_with_all_columns(metadata: RegionMetadataRef) -> PrimaryKeyReadFormat { Self::new( Arc::clone(&metadata), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), ) } } @@ -818,6 +860,7 @@ pub(crate) fn need_override_sequence(parquet_meta: &ParquetMetaData) -> bool { #[cfg(test)] mod tests { + use std::collections::BTreeMap; use std::sync::Arc; use api::v1::OpType; @@ -840,7 +883,6 @@ mod tests { use super::*; use crate::error::InvalidMetadataSnafu; - use crate::read::read_columns::ReadColumn; use crate::sst::parquet::flat_format::{ FlatReadFormat, FlatWriteFormat, sequence_column_index, sst_column_id_indices, }; @@ -986,29 +1028,25 @@ mod tests { fn test_projection_indices() { let metadata = build_test_region_metadata(); // Only read tag1 - let read_format = - PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([3])); + let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([3])); assert_eq!( &[2, 3, 4, 5], read_format.parquet_read_columns().root_indices() ); // Only read field1 - let read_format = - PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([4])); + let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([4])); assert_eq!( &[0, 2, 3, 4, 5], read_format.parquet_read_columns().root_indices() ); // Only read ts - let read_format = - PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([5])); + let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([5])); assert_eq!( &[2, 3, 4, 5], read_format.parquet_read_columns().root_indices() ); // Read field0, tag0, ts - let read_format = - PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids([2, 1, 5])); + let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new([2, 1, 5])); assert_eq!( &[1, 2, 3, 4, 5], read_format.parquet_read_columns().root_indices() @@ -1049,14 +1087,16 @@ mod tests { let metadata = Arc::new(builder.build().context(InvalidMetadataSnafu)?); let column_id_to_parquet_index = sst_column_id_indices(&metadata); let projection = FormatProjection::compute_format_projection( + &metadata, &column_id_to_parquet_index, metadata.column_metadatas.len() + FIXED_POS_COLUMN_NUM, - ReadColumns { - cols: vec![ReadColumn::new( - 4, - vec![vec!["j".to_string(), "a".to_string()]], - )], - }, + ReadColumns::new([4]).with_json_target_types(BTreeMap::from([( + 4, + JsonNativeType::Object(JsonObjectType::from([( + "a".to_string(), + JsonNativeType::i64(), + )])), + )])), ); let columns = projection.parquet_read_cols.columns(); @@ -1111,8 +1151,7 @@ mod tests { .iter() .map(|col| col.column_id) .collect(); - let read_format = - PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids(column_ids)); + let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new(column_ids)); assert_eq!(arrow_schema, *read_format.arrow_schema()); let record_batch = RecordBatch::new_empty(arrow_schema); @@ -1131,8 +1170,7 @@ mod tests { .iter() .map(|col| col.column_id) .collect(); - let read_format = - PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids(column_ids)); + let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new(column_ids)); let columns: Vec = vec![ Arc::new(Int64Array::from(vec![1, 1, 10, 10])), // field1 @@ -1160,9 +1198,7 @@ mod tests { let metadata = build_test_region_metadata(); let read_format = PrimaryKeyReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), ); let columns: Vec = vec![ @@ -1330,56 +1366,36 @@ mod tests { // The projection includes all "fixed position" columns: ts(4), __primary_key(5), __sequence(6), __op_type(7) // Only read tag1 (column_id=3, index=1) + fixed columns - let read_format = FlatReadFormat::new( - metadata.clone(), - ReadColumns::from_deduped_column_ids([3]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata.clone(), ReadColumns::new([3]), None, "test", false) + .unwrap(); assert_eq!( &[1, 4, 5, 6, 7], read_format.parquet_read_columns().root_indices() ); // Only read field1 (column_id=4, index=2) + fixed columns - let read_format = FlatReadFormat::new( - metadata.clone(), - ReadColumns::from_deduped_column_ids([4]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata.clone(), ReadColumns::new([4]), None, "test", false) + .unwrap(); assert_eq!( &[2, 4, 5, 6, 7], read_format.parquet_read_columns().root_indices() ); // Only read ts (column_id=5, index=4) + fixed columns (ts is already included in fixed) - let read_format = FlatReadFormat::new( - metadata.clone(), - ReadColumns::from_deduped_column_ids([5]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata.clone(), ReadColumns::new([5]), None, "test", false) + .unwrap(); assert_eq!( &[4, 5, 6, 7], read_format.parquet_read_columns().root_indices() ); // Read field0(column_id=2, index=3), tag0(column_id=1, index=0), ts(column_id=5, index=4) + fixed columns - let read_format = FlatReadFormat::new( - metadata, - ReadColumns::from_deduped_column_ids([2, 1, 5]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata, ReadColumns::new([2, 1, 5]), None, "test", false) + .unwrap(); assert_eq!( &[0, 3, 4, 5, 6, 7], read_format.parquet_read_columns().root_indices() @@ -1391,7 +1407,7 @@ mod tests { let metadata = build_test_region_metadata(); let mut format = FlatReadFormat::new( metadata, - ReadColumns::from_deduped_column_ids(std::iter::once(1)), // Just read tag0 + ReadColumns::new(std::iter::once(1)), // Just read tag0 Some(build_test_flat_sst_schema()), "test", false, @@ -1608,7 +1624,7 @@ mod tests { .collect(); let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids), + ReadColumns::new(column_ids), Some(build_test_arrow_schema()), "test", false, @@ -1674,7 +1690,7 @@ mod tests { .collect(); let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids.clone()), + ReadColumns::new(column_ids.clone()), None, "test", false, @@ -1744,7 +1760,7 @@ mod tests { let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids), + ReadColumns::new(column_ids), None, "test", true, diff --git a/src/mito2/src/sst/parquet/json_align/mod.rs b/src/mito2/src/sst/parquet/json_align/mod.rs index 89b993d3d09..54c97bdbb89 100644 --- a/src/mito2/src/sst/parquet/json_align/mod.rs +++ b/src/mito2/src/sst/parquet/json_align/mod.rs @@ -15,10 +15,8 @@ use datatypes::arrow::record_batch::RecordBatch; use futures::stream::BoxStream; -mod schema; mod stream; -pub(crate) use schema::align_schema_by_nested_paths; pub(crate) use stream::NestedSchemaAligner; use crate::error::Result; diff --git a/src/mito2/src/sst/parquet/json_align/schema.rs b/src/mito2/src/sst/parquet/json_align/schema.rs deleted file mode 100644 index 86bd4d4d0d7..00000000000 --- a/src/mito2/src/sst/parquet/json_align/schema.rs +++ /dev/null @@ -1,250 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::collections::BTreeMap; -use std::sync::Arc; - -use arrow_schema::{DataType as ArrowDataType, FieldRef}; -use datatypes::arrow::datatypes::Schema; -use datatypes::extension::json::is_json2_extension_type; -use store_api::storage::NestedPath; - -/// Aligns nested struct fields according to the requested nested paths. -/// -/// For each root field: -/// - An empty path list keeps the whole field unchanged. -/// - Non-JSON root fields ignore nested paths and keep the whole field unchanged. -/// - JSON2 root fields are rebuilt from `nested_paths`. -/// - Existing schema fields are preserved only when they are the requested leaf. -/// - Requested paths missing from the schema are synthesized with JSONB (`Binary`) -/// leaves. -/// -/// For example, if the schema has `j: struct>` -/// and `nested_paths` requests `j.a.x` and `j.a.z`, the result is -/// `j: struct>`. -pub(crate) fn align_schema_by_nested_paths<'a, I>(schema: &mut Schema, nested_paths: I) -where - I: IntoIterator, -{ - let fields = schema - .fields - .into_iter() - .zip(nested_paths) - .map(|(field, paths)| { - if !paths.is_empty() && is_json2_extension_type(field) { - let child_paths = paths - .iter() - .map(|path| { - if path.first().is_some_and(|root| root == field.name()) { - &path[1..] - } else { - path - } - }) - .collect::>(); - rebuild_field_by_nested_paths(field, &child_paths) - } else { - field.clone() - } - }) - .collect::>(); - schema.fields = fields.into() -} - -fn rebuild_field_by_nested_paths(field: &FieldRef, nested_paths: &[&[String]]) -> FieldRef { - if nested_paths.iter().any(|path| path.is_empty()) { - return field.clone(); - }; - - let fields = group_child_paths(nested_paths) - .into_iter() - .map(|(name, paths)| { - let existing = find_struct_child(field, &name); - build_field_from_nested_paths(&name, existing, &paths) - }) - .collect::>(); - - Arc::new( - field - .as_ref() - .clone() - .with_data_type(ArrowDataType::Struct(fields.into())), - ) -} - -fn group_child_paths<'a>(nested_paths: &[&'a [String]]) -> BTreeMap> { - let mut child_paths = BTreeMap::>::new(); - for path in nested_paths { - let Some((name, remaining)) = path.split_first() else { - continue; - }; - child_paths.entry(name.clone()).or_default().push(remaining); - } - child_paths -} - -fn find_struct_child<'a>(field: &'a FieldRef, name: &str) -> Option<&'a FieldRef> { - let ArrowDataType::Struct(fields) = field.data_type() else { - return None; - }; - fields.iter().find(|field| field.name() == name) -} - -fn build_field_from_nested_paths( - name: &str, - existing: Option<&FieldRef>, - nested_paths: &[&[String]], -) -> FieldRef { - if nested_paths.iter().any(|path| path.is_empty()) { - return existing.cloned().unwrap_or_else(|| new_jsonb_field(name)); - } - - let fields = group_child_paths(nested_paths) - .into_iter() - .map(|(name, paths)| { - let existing_child = existing.and_then(|field| find_struct_child(field, &name)); - build_field_from_nested_paths(&name, existing_child, &paths) - }) - .collect::>(); - - let field = existing - .map(|field| field.as_ref().clone()) - .unwrap_or_else(|| arrow_schema::Field::new(name, ArrowDataType::Binary, true)); - Arc::new(field.with_data_type(ArrowDataType::Struct(fields.into()))) -} - -fn new_jsonb_field(name: &str) -> FieldRef { - Arc::new(arrow_schema::Field::new(name, ArrowDataType::Binary, true)) -} - -#[cfg(test)] -mod tests { - use arrow_schema::Field; - use datatypes::extension::json::Json2ExtensionType; - - use super::*; - - #[test] - fn test_align_schema_by_nested_paths() { - fn new_field(name: &str, data_type: ArrowDataType) -> FieldRef { - Arc::new(Field::new(name, data_type, true)) - } - - fn struct_field(name: &str, fields: impl IntoIterator) -> FieldRef { - new_field(name, ArrowDataType::Struct(fields.into_iter().collect())) - } - - fn json_struct_field(name: &str, fields: impl IntoIterator) -> FieldRef { - Arc::new( - Field::new( - name, - ArrowDataType::Struct(fields.into_iter().collect()), - true, - ) - .with_extension_type(Json2ExtensionType::default()), - ) - } - - let mut schema = Schema::new([ - json_struct_field( - "j", - [ - struct_field( - "a", - [ - new_field("x", ArrowDataType::Int64), - new_field("y", ArrowDataType::Utf8), - struct_field( - "z", - [ - new_field("q", ArrowDataType::Boolean), - new_field("r", ArrowDataType::Float64), - ], - ), - ], - ), - new_field("b", ArrowDataType::Utf8), - struct_field( - "c", - vec![ - new_field("d", ArrowDataType::Int64), - new_field("e", ArrowDataType::Utf8), - ], - ), - ], - ), - new_field("tag", ArrowDataType::Utf8), - struct_field( - "k", - [ - new_field("k_0", ArrowDataType::Int64), - new_field("k_1", ArrowDataType::Utf8), - ], - ), - ]); - - let nested_paths = [ - vec![ - ["j", "a", "x"].iter().map(|x| x.to_string()).collect(), - ["j", "a", "z", "q"].iter().map(|x| x.to_string()).collect(), - ["j", "a", "m"].iter().map(|x| x.to_string()).collect(), - ["j", "b", "x"].iter().map(|x| x.to_string()).collect(), - ["j", "c"].iter().map(|x| x.to_string()).collect(), - ["j", "d", "e"].iter().map(|x| x.to_string()).collect(), - ], - vec![["tag", "ignored"].iter().map(|x| x.to_string()).collect()], - vec![], - ]; - - align_schema_by_nested_paths( - &mut schema, - nested_paths.iter().map(|paths| paths.as_slice()), - ); - - let expected = Schema::new([ - json_struct_field( - "j", - [ - struct_field( - "a", - [ - new_field("m", ArrowDataType::Binary), - new_field("x", ArrowDataType::Int64), - struct_field("z", vec![new_field("q", ArrowDataType::Boolean)]), - ], - ), - struct_field("b", [new_field("x", ArrowDataType::Binary)]), - struct_field( - "c", - [ - new_field("d", ArrowDataType::Int64), - new_field("e", ArrowDataType::Utf8), - ], - ), - struct_field("d", [new_field("e", ArrowDataType::Binary)]), - ], - ), - new_field("tag", ArrowDataType::Utf8), - struct_field( - "k", - [ - new_field("k_0", ArrowDataType::Int64), - new_field("k_1", ArrowDataType::Utf8), - ], - ), - ]); - - assert_eq!(schema, expected); - } -} diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index 5b08a0e4957..fcaaae4a54d 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -1457,9 +1457,7 @@ mod tests { Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense)); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, @@ -1495,9 +1493,7 @@ mod tests { )); let legacy_read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "memtable", false, @@ -1524,9 +1520,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let raw_pk_read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "memtable", true, @@ -1560,9 +1554,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let full_read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -1594,7 +1586,7 @@ mod tests { let ts = metric_metadata.time_index_column().column_id; let projected_read_format = FlatReadFormat::new( metric_metadata.clone(), - ReadColumns::from_deduped_column_ids([field_0, ts]), + ReadColumns::new([field_0, ts]), None, "test", true, @@ -1646,9 +1638,7 @@ mod tests { )); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, @@ -1686,9 +1676,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -1734,9 +1722,7 @@ mod tests { Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense)); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 97a4abfac02..4f890cb83e8 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -451,7 +451,7 @@ impl ParquetReaderBuilder { } else { let expected_meta = self.expected_metadata.as_ref().unwrap_or(®ion_meta); // Lists all column ids to read, we always use the expected metadata if possible. - ReadColumns::from_deduped_column_ids( + ReadColumns::new( expected_meta .column_metadatas .iter() @@ -2472,9 +2472,7 @@ mod tests { let metadata = Arc::new(sst_region_metadata()); let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -2679,7 +2677,7 @@ mod tests { let region_metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( region_metadata.clone(), - ReadColumns::from_deduped_column_ids( + ReadColumns::new( region_metadata .column_metadatas .iter() @@ -2946,9 +2944,7 @@ mod tests { let expected_metadata = expected_metadata_with_reused_tag_name(metadata.as_ref()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -2970,9 +2966,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, diff --git a/src/query/src/optimizer/json_type_concretize.rs b/src/query/src/optimizer/json_type_concretize.rs index 4210433eb3a..4625b619784 100644 --- a/src/query/src/optimizer/json_type_concretize.rs +++ b/src/query/src/optimizer/json_type_concretize.rs @@ -16,12 +16,14 @@ use std::collections::HashMap; use arrow_schema::DataType; use common_function::scalars::json::json_get::JsonGetWithType; -use datafusion::datasource::DefaultTableSource; +use datafusion::datasource::{DefaultTableSource, TableProvider}; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_common::{Result, plan_datafusion_err, plan_err}; use datafusion_expr::{Expr, LogicalPlan}; use datafusion_optimizer::{OptimizerConfig, OptimizerRule}; +use datatypes::extension::json::is_json2_extension_type; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; +use table::table::adapter::DfTableProviderAdapter; use crate::dummy_catalog::DummyTableProvider; @@ -56,22 +58,56 @@ impl OptimizerRule for JsonTypeConcretizeRule { return Ok(Transformed::no(plan)); }; - let Some(adapter) = source - .table_provider - .as_any() - .downcast_ref::() - else { - return Ok(Transformed::no(plan)); - }; - - adapter.with_json_type_hint(json_types.clone()); - Ok(Transformed::yes(plan)) + if apply_json_type_hint(source.table_provider.as_ref(), &json_types) { + Ok(Transformed::yes(plan)) + } else { + Ok(Transformed::no(plan)) + } } _ => Ok(Transformed::no(plan)), }) } } +// FIXME: `json_types` is keyed only by unqualified column name. In joins with +// same-named JSON2 columns, a hint deduced from one scan can be applied to +// another scan. Carry the originating relation/scan when deducing hints. +/// Applies JSON type hints to providers that can carry scan request hints. +/// +/// Returns `true` if at least one JSON2 hint is retained and written to the provider. +fn apply_json_type_hint( + provider: &dyn TableProvider, + json_types: &HashMap, +) -> bool { + let schema = provider.schema(); + let json_types = json_types + .iter() + .filter(|(column, _)| { + schema + .fields() + .iter() + .any(|field| field.name() == *column && is_json2_extension_type(field)) + }) + .map(|(column, json_type)| (column.clone(), json_type.clone())) + .collect::>(); + + if json_types.is_empty() { + return false; + } + + if let Some(adapter) = provider.as_any().downcast_ref::() { + adapter.with_json_type_hint(json_types); + return true; + } + + if let Some(adapter) = provider.as_any().downcast_ref::() { + adapter.with_json_type_hint(json_types); + return true; + } + + false +} + fn deduce_json_types(plan: &LogicalPlan) -> Result> { let mut json_types = HashMap::::new(); @@ -228,11 +264,11 @@ mod tests { #[test] fn test_json_type_concretize_rule_rewrite() -> Result<()> { let exprs = vec![ - json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab"), - json_get_expr(col("k0"), path_expr("a.c"), None)?.alias("ac"), - json_get_expr(col("k0"), path_expr("d"), Some(DataType::Boolean))?.alias("d"), + json_get_expr(col("j"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab"), + json_get_expr(col("j"), path_expr("a.c"), None)?.alias("ac"), + json_get_expr(col("j"), path_expr("d"), Some(DataType::Boolean))?.alias("d"), ]; - let (provider, plan) = build_plan(exprs)?; + let (provider, plan) = build_json2_plan(exprs)?; assert!( JsonTypeConcretizeRule @@ -253,17 +289,17 @@ mod tests { let request = provider.scan_request(); assert_eq!(1, request.json_type_hint.len()); - assert_eq!(Some(&expected), request.json_type_hint.get("k0")); + assert_eq!(Some(&expected), request.json_type_hint.get("j")); Ok(()) } #[test] fn test_json_type_concretize_rule_conflict_to_variant() -> Result<()> { let exprs = vec![ - json_get_expr(col("k0"), path_expr("a"), Some(DataType::Int64))?.alias("a_num"), - json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Boolean))?.alias("a_obj"), + json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?.alias("a_num"), + json_get_expr(col("j"), path_expr("a.b"), Some(DataType::Boolean))?.alias("a_obj"), ]; - let (provider, plan) = build_plan(exprs)?; + let (provider, plan) = build_json2_plan(exprs)?; assert!( JsonTypeConcretizeRule @@ -277,11 +313,26 @@ mod tests { )])); assert_eq!( Some(&expected), - provider.scan_request().json_type_hint.get("k0") + provider.scan_request().json_type_hint.get("j") ); Ok(()) } + #[test] + fn test_json_type_concretize_rule_ignores_non_json2_columns() -> Result<()> { + let exprs = + vec![json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab")]; + let (provider, plan) = build_plan(exprs)?; + + assert!( + !JsonTypeConcretizeRule + .rewrite(plan, &OptimizerContext::default())? + .transformed + ); + assert!(provider.scan_request().json_type_hint.is_empty()); + Ok(()) + } + #[test] fn test_json_type_concretize_rule_no_json_get() -> Result<()> { let (provider, plan) = build_plan(vec![col("k0"), col("v0")])?; diff --git a/src/table/src/table/adapter.rs b/src/table/src/table/adapter.rs index 0ae56ce1c31..ac71b448e94 100644 --- a/src/table/src/table/adapter.rs +++ b/src/table/src/table/adapter.rs @@ -28,6 +28,7 @@ use datafusion_expr::TableProviderFilterPushDown as DfTableProviderFilterPushDow use datafusion_expr::expr::Expr; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; +use datatypes::types::json_type::JsonNativeType; use store_api::storage::{ScanRequest, VectorSearchRequest}; use crate::table::{TableRef, TableType}; @@ -63,6 +64,10 @@ impl DfTableProviderAdapter { self.scan_req.lock().unwrap().vector_search = Some(hint); } + pub fn with_json_type_hint(&self, hint: std::collections::HashMap) { + self.scan_req.lock().unwrap().json_type_hint = hint; + } + pub fn get_vector_search_hint(&self) -> Option { self.scan_req.lock().unwrap().vector_search.clone() } diff --git a/tests/cases/standalone/common/types/json/json2_limit.result b/tests/cases/standalone/common/types/json/json2_limit.result index 3cb462a8a7e..035c4ed07bf 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.result +++ b/tests/cases/standalone/common/types/json/json2_limit.result @@ -77,6 +77,75 @@ drop table json2_whole_and_path_read; Affected Rows: 0 +create table json2_join_same_name_left ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +create table json2_join_same_name_right ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +insert into json2_join_same_name_left values + (1, 'a', '{"a": 1, "left_only": "kept"}'); + +Affected Rows: 1 + +insert into json2_join_same_name_right values + (1, 'a', '{"a": "right", "right_only": "should be kept"}'); + +Affected Rows: 1 + +admin flush_table('json2_join_same_name_left'); + ++------------------------------------------------+ +| ADMIN flush_table('json2_join_same_name_left') | ++------------------------------------------------+ +| 0 | ++------------------------------------------------+ + +admin flush_table('json2_join_same_name_right'); + ++-------------------------------------------------+ +| ADMIN flush_table('json2_join_same_name_right') | ++-------------------------------------------------+ +| 0 | ++-------------------------------------------------+ + +-- FIXME: This should return `right` and `1`. The current NULL values are caused +-- by JSON type hints losing the table qualifier in joins. +select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 +from json2_join_same_name_left l +join json2_join_same_name_right r +on l.k = r.k; + ++---------------------------------------------+---------------------------------------------------+ +| json_get(r.j,Utf8("a")) | arrow_cast(json_get(l.j,Utf8("a")),Utf8("Int64")) | ++---------------------------------------------+---------------------------------------------------+ +| {"a":"right","right_only":"should be kept"} | | ++---------------------------------------------+---------------------------------------------------+ + +drop table json2_join_same_name_left; + +Affected Rows: 0 + +drop table json2_join_same_name_right; + +Affected Rows: 0 + create table json2_without_append_mode ( ts timestamp time index, j json2 diff --git a/tests/cases/standalone/common/types/json/json2_limit.sql b/tests/cases/standalone/common/types/json/json2_limit.sql index 4917b9690e5..986ed58eb2a 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.sql +++ b/tests/cases/standalone/common/types/json/json2_limit.sql @@ -44,6 +44,45 @@ select j from json2_whole_and_path_read where j.a.b = 1; drop table json2_whole_and_path_read; +create table json2_join_same_name_left ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +create table json2_join_same_name_right ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +insert into json2_join_same_name_left values + (1, 'a', '{"a": 1, "left_only": "kept"}'); + +insert into json2_join_same_name_right values + (1, 'a', '{"a": "right", "right_only": "should be kept"}'); + +admin flush_table('json2_join_same_name_left'); + +admin flush_table('json2_join_same_name_right'); + +-- FIXME: This should return `right` and `1`. The current NULL values are caused +-- by JSON type hints losing the table qualifier in joins. +select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 +from json2_join_same_name_left l +join json2_join_same_name_right r +on l.k = r.k; + +drop table json2_join_same_name_left; + +drop table json2_join_same_name_right; + create table json2_without_append_mode ( ts timestamp time index, j json2 From 76924c2d36b8d6b3d2afb111536e7bda04d0bda8 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 14 Aug 2026 06:32:17 +0000 Subject: [PATCH 070/126] feat(mito2): introduce two-phase metric series scans (#8826) * feat(mito2): add two-phase series scan Signed-off-by: evenyag * docs: regenerate configuration reference Signed-off-by: evenyag * test(sqlness): update series scan explain results Signed-off-by: evenyag * test: update config API expectation Signed-off-by: evenyag * fix(mito2): bound two-phase series discovery Signed-off-by: evenyag * fix(mito2): avoid candidate distribution deadlock Signed-off-by: evenyag * chore(mito2): remove obsolete dead code allowances Signed-off-by: evenyag * fix(mito2): share series scan memory pool Signed-off-by: evenyag --------- Signed-off-by: evenyag --- config/config.md | 2 + config/datanode.example.toml | 3 + config/standalone.example.toml | 3 + src/metric-engine/src/engine/read.rs | 173 +++++++- src/mito2/src/config.rs | 12 + src/mito2/src/engine.rs | 32 ++ src/mito2/src/engine/scan_test.rs | 196 ++++++++- src/mito2/src/read/scan_region.rs | 35 +- src/mito2/src/read/series_candidate.rs | 13 +- src/mito2/src/read/series_reader.rs | 68 ++- src/mito2/src/read/series_scan.rs | 408 ++++++++++++++++-- tests-integration/tests/http.rs | 1 + .../explain/step_aggr_advance.result | 28 +- .../tql-explain-analyze/analyze.result | 14 +- .../promql/histogram_multi_partition.result | 4 +- .../promql/tsid_binary_join_regression.result | 56 +-- .../tsid_histogram_quantile_regression.result | 2 +- .../common/select/tql_filter.result | 6 +- .../standalone/common/tql/partition.result | 10 +- .../standalone/optimizer/windowed_sort.result | 2 +- .../tql-explain-analyze/analyze.result | 22 +- .../tql-explain-analyze/column_pruning.result | 4 +- .../tql-explain-analyze/tsid_column.result | 12 +- 23 files changed, 963 insertions(+), 143 deletions(-) diff --git a/config/config.md b/config/config.md index c7737a164b8..123eaf90a20 100644 --- a/config/config.md +++ b/config/config.md @@ -185,6 +185,7 @@ | `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). | | `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.

Setting this to "true" is a necessary but not sufficient condition for scheduling compaction after a region edit.
Other constraints, such as "min_compaction_interval", may still prevent compaction from being scheduled.
Setting this to "false", however, guarantees that compaction will not be scheduled after a region edit. | | `region_engine.mito.default_flat_format` | Bool | `true` | Whether to enable flat format as the default SST format. | +| `region_engine.mito.experimental_series_scan_v2` | Bool | `true` | Whether to enable the experimental two-phase mode for eligible metric series scans. | | `region_engine.mito.index` | -- | -- | The options for index in Mito engine. | | `region_engine.mito.index.aux_path` | String | `""` | Auxiliary directory path for the index in filesystem, used to store intermediate files for
creating the index and staging files for searching the index, defaults to `{data_home}/index_intermediate`.
The default name for this directory is `index_intermediate` for backward compatibility.

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | @@ -606,6 +607,7 @@ | `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). | | `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.

Setting this to "true" is a necessary but not sufficient condition for scheduling compaction after a region edit.
Other constraints, such as "min_compaction_interval", may still prevent compaction from being scheduled.
Setting this to "false", however, guarantees that compaction will not be scheduled after a region edit. | | `region_engine.mito.default_flat_format` | Bool | `true` | Whether to enable flat format as the default SST format. | +| `region_engine.mito.experimental_series_scan_v2` | Bool | `true` | Whether to enable the experimental two-phase mode for eligible metric series scans. | | `region_engine.mito.index` | -- | -- | The options for index in Mito engine. | | `region_engine.mito.index.aux_path` | String | `""` | Auxiliary directory path for the index in filesystem, used to store intermediate files for
creating the index and staging files for searching the index, defaults to `{data_home}/index_intermediate`.
The default name for this directory is `index_intermediate` for backward compatibility.

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 3426630d7b6..1f1377d91cc 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -571,6 +571,9 @@ schedule_compaction_after_edit = true ## Whether to enable flat format as the default SST format. default_flat_format = true +## Whether to enable the experimental two-phase mode for eligible metric series scans. +experimental_series_scan_v2 = true + ## The options for index in Mito engine. [region_engine.mito.index] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index c8dff4944c1..068025e2a28 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -729,6 +729,9 @@ schedule_compaction_after_edit = true ## Whether to enable flat format as the default SST format. default_flat_format = true +## Whether to enable the experimental two-phase mode for eligible metric series scans. +experimental_series_scan_v2 = true + ## The options for index in Mito engine. [region_engine.mito.index] diff --git a/src/metric-engine/src/engine/read.rs b/src/metric-engine/src/engine/read.rs index 53dd086d225..ee086184c23 100644 --- a/src/metric-engine/src/engine/read.rs +++ b/src/metric-engine/src/engine/read.rs @@ -306,15 +306,72 @@ impl MetricEngineInner { #[cfg(test)] mod test { + use std::fmt; + + use api::v1::Rows; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; - use store_api::region_request::RegionRequest; + use datafusion::physical_plan::DisplayFormatType; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; + use futures_util::TryStreamExt; + use futures_util::future::try_join_all; + use mito2::config::MitoConfig; + use store_api::region_engine::{PrepareRequest, QueryScanContext}; + use store_api::region_request::{RegionFlushRequest, RegionPutRequest, RegionRequest}; + use store_api::storage::TimeSeriesDistribution; use super::*; + use crate::config::EngineConfig; use crate::test_util::{ - TestEnv, alter_logical_region_add_tag_columns, create_logical_region_request, + self, TestEnv, alter_logical_region_add_tag_columns, create_logical_region_request, }; + struct ScannerDisplay<'a>(&'a dyn store_api::region_engine::RegionScanner); + + impl fmt::Display for ScannerDisplay<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt_as(DisplayFormatType::Default, f) + } + } + + async fn count_scanner_rows(scanner: &mut RegionScannerRef, partitions: usize) -> usize { + let ranges = scanner + .properties() + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + let mut prepared = vec![Vec::new(); partitions]; + prepared[0] = ranges; + scanner + .prepare( + PrepareRequest::default() + .with_ranges(prepared) + .with_target_partitions(partitions), + ) + .unwrap(); + + let metrics = ExecutionPlanMetricsSet::default(); + let context = QueryScanContext::default(); + let streams = (0..partitions) + .map(|partition| { + scanner + .scan_partition(&context, &metrics, partition) + .unwrap() + }) + .collect::>(); + try_join_all(streams.into_iter().map(|stream| async move { + stream + .try_fold(0, |rows, batch| async move { Ok(rows + batch.num_rows()) }) + .await + })) + .await + .unwrap() + .into_iter() + .sum() + } + #[tokio::test] async fn test_invalid_logical_projection() { let env = TestEnv::new().await; @@ -406,4 +463,116 @@ mod test { &[11, 10, 9, 8, 0, 1, 4] ); } + + #[tokio::test] + async fn test_two_phase_series_scan_reads_metric_region() { + let env = TestEnv::with_mito_config( + "test_two_phase_series_scan", + MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }, + EngineConfig::default(), + ) + .await; + env.init_metric_region().await; + + let physical_region_id = env.default_physical_region_id(); + let logical_region_id = env.default_logical_region_id(); + let logical_region_id_2 = RegionId::new(1024, logical_region_id.region_number()); + env.metric() + .handle_request( + logical_region_id_2, + RegionRequest::Create(create_logical_region_request( + &["job"], + physical_region_id, + "test_metric_region_2", + )), + ) + .await + .unwrap(); + + let schema = test_util::row_schema_with_tags(&["job"]); + let put = |rows| { + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows: test_util::build_rows(1, rows), + }, + hint: None, + partition_expr_version: None, + }) + }; + env.metric() + .handle_request(logical_region_id, put(3)) + .await + .unwrap(); + env.metric() + .handle_request( + physical_region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + env.metric() + .handle_request(logical_region_id, put(2)) + .await + .unwrap(); + env.metric() + .handle_request(logical_region_id_2, put(4)) + .await + .unwrap(); + + let data_region_id = utils::to_data_region_id(physical_region_id); + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }; + let mut scanner = env + .mito() + .handle_query(data_region_id, request) + .await + .unwrap(); + + let rows = tokio::time::timeout( + std::time::Duration::from_secs(30), + count_scanner_rows(&mut scanner, 2), + ) + .await + .expect("two-phase series scan should not deadlock"); + assert_eq!(7, rows); + + let explain = ScannerDisplay(scanner.as_ref()).to_string(); + assert!(explain.contains("\"mode\":\"two_phase\""), "{explain}"); + } + + #[tokio::test] + async fn test_series_scan_v2_flag_disables_two_phase_mode() { + let env = TestEnv::with_mito_config( + "test_disable_two_phase_series_scan", + MitoConfig { + experimental_series_scan_v2: false, + ..Default::default() + }, + EngineConfig::default(), + ) + .await; + env.init_metric_region().await; + + let data_region_id = utils::to_data_region_id(env.default_physical_region_id()); + let scanner = env + .mito() + .handle_query( + data_region_id, + ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }, + ) + .await + .unwrap(); + + let explain = ScannerDisplay(scanner.as_ref()).to_string(); + assert!(explain.contains("\"mode\":\"legacy\""), "{explain}"); + } } diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index f7489709ee3..54445549a05 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -188,6 +188,9 @@ pub struct MitoConfig { /// When enabled, forces using BulkMemtable and BulkMemtableBuilder. pub default_flat_format: bool, + /// Whether to enable the experimental two-phase mode for eligible metric series scans. + pub experimental_series_scan_v2: bool, + pub gc: GcConfig, } @@ -239,6 +242,7 @@ impl Default for MitoConfig { min_compaction_interval: Duration::from_secs(0), schedule_compaction_after_edit: true, default_flat_format: true, + experimental_series_scan_v2: true, gc: GcConfig::default(), }; @@ -389,6 +393,14 @@ mod tests { assert_eq!(ReadableSize::mb(512), config.sst_meta_cache_size); assert_eq!(ReadableSize::mb(128), config.prefilter_result_cache_size); } + + #[test] + fn test_experimental_series_scan_v2_config() { + assert!(MitoConfig::default().experimental_series_scan_v2); + + let config: MitoConfig = toml::from_str("experimental_series_scan_v2 = false").unwrap(); + assert!(!config.experimental_series_scan_v2); + } } /// Index build mode. diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index a128531b535..ddd34669eb7 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -102,6 +102,7 @@ use common_recordbatch::{QueryMemoryTracker, SendableRecordBatchStream}; use common_stat::get_total_memory_bytes; use common_telemetry::{info, tracing, warn}; use common_wal::options::WalOptions; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; use futures::future::{join_all, try_join_all}; use futures::stream::{self, Stream, StreamExt}; use object_store::manager::ObjectStoreManagerRef; @@ -236,6 +237,7 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> { let wal_raw_entry_reader = Arc::new(LogStoreRawEntryReader::new(self.log_store)); let total_memory = get_total_memory_bytes().max(0) as u64; let scan_memory_limit = config.scan_memory_limit.resolve(total_memory) as usize; + let scan_memory_pool = new_scan_memory_pool(scan_memory_limit); let scan_memory_tracker = QueryMemoryTracker::builder(scan_memory_limit, config.scan_memory_on_exhausted) .on_update(|usage| { @@ -254,6 +256,7 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> { config, wal_raw_entry_reader, scan_memory_tracker, + scan_memory_pool, region_hook, #[cfg(feature = "enterprise")] extension_range_provider_factory: None, @@ -739,6 +742,8 @@ struct EngineInner { wal_raw_entry_reader: Arc, /// Memory tracker for table scans. scan_memory_tracker: QueryMemoryTracker, + /// Memory pool shared by internal scan operators across all queries. + scan_memory_pool: Arc, /// The region hook (if any) registered via plugins; exposed for the GC worker /// to fire [`RegionHook::on_region_gc`]. region_hook: Option, @@ -1097,6 +1102,8 @@ impl EngineInner { ) .with_query_stat_counters(region.region_stats.query_stat_counters()) .with_max_concurrent_scan_files(self.config.max_concurrent_scan_files) + .with_scan_memory_pool(self.scan_memory_pool.clone()) + .with_experimental_series_scan_v2(self.config.experimental_series_scan_v2) .with_ignore_inverted_index(self.config.inverted_index.apply_on_query.disabled()) .with_ignore_fulltext_index(self.config.fulltext_index.apply_on_query.disabled()) .with_ignore_bloom_filter(self.config.bloom_filter_index.apply_on_query.disabled()) @@ -1443,6 +1450,7 @@ impl MitoEngine { let wal_raw_entry_reader = Arc::new(LogStoreRawEntryReader::new(log_store.clone())); let total_memory = get_total_memory_bytes().max(0) as u64; let scan_memory_limit = config.scan_memory_limit.resolve(total_memory) as usize; + let scan_memory_pool = new_scan_memory_pool(scan_memory_limit); let scan_memory_tracker = QueryMemoryTracker::builder(scan_memory_limit, config.scan_memory_on_exhausted) .on_update(|usage| { @@ -1472,6 +1480,7 @@ impl MitoEngine { config, wal_raw_entry_reader, scan_memory_tracker, + scan_memory_pool, region_hook: None, #[cfg(feature = "enterprise")] extension_range_provider_factory: None, @@ -1485,10 +1494,20 @@ impl MitoEngine { } } +fn new_scan_memory_pool(scan_memory_limit: usize) -> Arc { + if scan_memory_limit == 0 { + Arc::new(UnboundedMemoryPool::default()) + } else { + Arc::new(GreedyMemoryPool::new(scan_memory_limit)) + } +} + #[cfg(test)] mod tests { use std::time::Duration; + use datafusion::execution::memory_pool::MemoryConsumer; + use super::*; use crate::sst::file::FileMeta; @@ -1574,4 +1593,17 @@ mod tests { }; assert!(!is_valid_region_edit(&edit)); } + + #[test] + fn test_scan_memory_pool_is_shared_across_consumers() { + let pool = new_scan_memory_pool(100); + let cloned_pool = pool.clone(); + let first = MemoryConsumer::new("first-scan").register(&pool); + let second = MemoryConsumer::new("second-scan").register(&cloned_pool); + + first.try_grow(60).unwrap(); + assert!(second.try_grow(50).is_err()); + second.try_grow(40).unwrap(); + assert_eq!(100, pool.reserved()); + } } diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 4baa707f4e7..a53b77c8253 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -13,11 +13,12 @@ // limitations under the License. use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; use api::helper::encode_json_value; -use api::v1::Rows; use api::v1::helper::row; use api::v1::value::ValueData; +use api::v1::{ColumnDataType, Rows, SemanticType, WriteHint}; use common_base::readable_size::ReadableSize; use common_error::ext::{ErrorExt, WhateverResult}; use common_error::status_code::StatusCode; @@ -26,14 +27,18 @@ use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_common::ScalarValue; use datafusion_expr::{col, lit}; use datatypes::arrow::array::AsArray; -use datatypes::arrow::datatypes::{Float64Type, TimestampMillisecondType}; +use datatypes::arrow::datatypes::{Float64Type, TimestampMillisecondType, UInt64Type}; use datatypes::json::value::JsonValue; use datatypes::prelude::ConcreteDataType; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; use futures::TryStreamExt; +use futures::future::try_join_all; use serde_json::json; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metric_engine_consts::PRIMARY_KEY_ENCODING; use store_api::region_engine::{PrepareRequest, RegionEngine, RegionScanner}; -use store_api::region_request::RegionRequest; +use store_api::region_request::{RegionPutRequest, RegionRequest}; +use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution}; use crate::config::MitoConfig; @@ -41,6 +46,7 @@ use crate::error::Error; use crate::read::read_columns::ReadColumns; use crate::read::scan_region::Scanner; use crate::test_util; +use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; use crate::test_util::{CreateRequestBuilder, TestEnv}; #[tokio::test] @@ -753,6 +759,7 @@ async fn test_series_scan_with_format(flat_format: bool) { let Scanner::Series(mut scanner) = scanner else { panic!("Scanner should be series scan"); }; + assert_eq!("legacy", scanner.mode()); // 3 partition ranges for 3 time window. assert_eq!( 3, @@ -789,6 +796,189 @@ async fn test_series_scan_with_format(flat_format: bool) { assert_eq!(expected_rows, actual_rows); } +#[tokio::test] +async fn test_two_phase_series_scan() { + let mut env = TestEnv::with_prefix("test_two_phase_series_scan").await; + let engine = env + .create_engine(MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let mut request = CreateRequestBuilder::new().build(); + request.column_metadatas = metadata.column_metadatas.clone(); + request.primary_key = metadata.primary_key.clone(); + request + .options + .insert(PRIMARY_KEY_ENCODING.to_string(), "sparse".to_string()); + request + .options + .insert("memtable.type".to_string(), "bulk".to_string()); + request + .options + .insert("sst_format".to_string(), "flat".to_string()); + let full_row_schema = test_util::rows_schema(&request); + let mut encoded_primary_key_schema = full_row_schema[0].clone(); + encoded_primary_key_schema.column_name = PRIMARY_KEY_COLUMN_NAME.to_string(); + encoded_primary_key_schema.datatype = ColumnDataType::Binary.into(); + encoded_primary_key_schema.semantic_type = SemanticType::Tag.into(); + let row_schema = vec![ + encoded_primary_key_schema, + full_row_schema[5].clone(), + full_row_schema[4].clone(), + ]; + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let rows = |values: &[(u32, u64, &str, &str, u64, i64)]| Rows { + schema: row_schema.clone(), + rows: values + .iter() + .map(|(table_id, tsid, tag_0, tag_1, field, ts)| { + row(vec![ + ValueData::BinaryValue(new_sparse_primary_key( + &[*tag_0, *tag_1], + &metadata, + *table_id, + *tsid, + )), + ValueData::TimestampMillisecondValue(*ts), + ValueData::U64Value(*field), + ]) + }) + .collect(), + }; + let put = |rows| { + RegionRequest::Put(RegionPutRequest { + rows, + hint: Some(WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + }), + partition_expr_version: None, + }) + }; + + engine + .handle_request( + region_id, + put(rows(&[ + (10, 0, "a", "x", 10, 1000), + (10, u64::MAX, "b", "y", 20, 1000), + (20, 0, "c", "z", 30, 1000), + ])), + ) + .await + .unwrap(); + test_util::flush_region(&engine, region_id, None).await; + engine + .handle_request( + region_id, + put(rows(&[ + // Replaces the flushed row for this series and timestamp. + (10, 0, "a", "x", 11, 1000), + (10, 0, "a", "x", 12, 2000), + (10, u64::MAX, "b", "y", 21, 2000), + (20, u64::MAX, "d", "w", 40, 1000), + ])), + ) + .await + .unwrap(); + + let scanner = engine + .scanner( + region_id, + ScanRequest { + // Internal metric identifiers must still be read for candidate discovery. + projection: Some(vec![2, 4, 5]), + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }, + ) + .await + .unwrap(); + let Scanner::Series(mut scanner) = scanner else { + panic!("Scanner should be series scan"); + }; + assert_eq!("two_phase", scanner.mode()); + + let ranges = scanner + .properties() + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + scanner + .prepare( + PrepareRequest::default() + .with_ranges(vec![ranges, Vec::new(), Vec::new()]) + .with_target_partitions(3), + ) + .unwrap(); + + let metrics_set = ExecutionPlanMetricsSet::default(); + let context = store_api::region_engine::QueryScanContext::default(); + let partition_batches = try_join_all((0..3).map(|partition| { + let stream = scanner + .scan_partition(&context, &metrics_set, partition) + .unwrap(); + async move { stream.try_collect::>().await } + })) + .await + .unwrap(); + + let mut series_to_partition = BTreeMap::new(); + let mut actual_rows = Vec::new(); + for (partition, batches) in partition_batches.into_iter().enumerate() { + for batch in batches { + let tags = batch.column_by_name("tag_0").unwrap(); + let fields = batch + .column_by_name("field_0") + .unwrap() + .as_primitive::(); + let timestamps = batch + .column_by_name("ts") + .unwrap() + .as_primitive::(); + for row in 0..batch.num_rows() { + let tag = datatypes::arrow_array::string_array_value_at_index(tags, row) + .unwrap() + .to_string(); + let previous = series_to_partition.insert(tag.clone(), partition); + if let Some(previous) = previous { + assert_eq!(previous, partition, "series {tag} moved partitions"); + } + actual_rows.push((tag, fields.value(row), timestamps.value(row))); + } + } + } + actual_rows.sort(); + assert_eq!( + vec![ + ("a".to_string(), 11, 1000), + ("a".to_string(), 12, 2000), + ("b".to_string(), 20, 1000), + ("b".to_string(), 21, 2000), + ("c".to_string(), 30, 1000), + ("d".to_string(), 40, 1000), + ], + actual_rows + ); + assert_eq!(4, series_to_partition.len()); + assert_eq!(Some(&0), series_to_partition.get("a")); + assert_eq!(Some(&0), series_to_partition.get("c")); + assert_eq!(Some(&2), series_to_partition.get("b")); + assert_eq!(Some(&2), series_to_partition.get("d")); +} + /// Scans all partitions in round-robin fashion and returns rows sorted by (tag, ts). /// Also asserts that each series appears in only one partition. async fn collect_partition_rows_round_robin( diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index c0340464ecc..df199d4bac2 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -28,6 +28,7 @@ use common_recordbatch::filter::SimpleFilterEvaluator; use common_telemetry::tracing::Instrument; use common_telemetry::{debug, error, tracing, warn}; use common_time::range::TimestampRange; +use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; use datafusion::physical_plan::expressions::DynamicFilterPhysicalExpr; use datafusion_common::pruning::PruningStatistics; use datafusion_common::{Column, ScalarValue}; @@ -238,6 +239,10 @@ pub(crate) struct ScanRegion { cache_strategy: CacheStrategy, /// Maximum number of SST files to scan concurrently. max_concurrent_scan_files: usize, + /// Memory pool shared by internal scan operators across all queries. + scan_memory_pool: Arc, + /// Whether to enable the experimental two-phase metric series scan. + experimental_series_scan_v2: bool, /// Whether to ignore inverted index. ignore_inverted_index: bool, /// Whether to ignore fulltext index. @@ -269,6 +274,8 @@ impl ScanRegion { request, cache_strategy, max_concurrent_scan_files: DEFAULT_MAX_CONCURRENT_SCAN_FILES, + scan_memory_pool: Arc::new(UnboundedMemoryPool::default()), + experimental_series_scan_v2: false, ignore_inverted_index: false, ignore_fulltext_index: false, ignore_bloom_filter: false, @@ -297,6 +304,20 @@ impl ScanRegion { self } + /// Sets the memory pool shared by internal scan operators. + #[must_use] + pub(crate) fn with_scan_memory_pool(mut self, scan_memory_pool: Arc) -> Self { + self.scan_memory_pool = scan_memory_pool; + self + } + + /// Sets whether eligible metric series scans use the two-phase implementation. + #[must_use] + pub(crate) fn with_experimental_series_scan_v2(mut self, enabled: bool) -> Self { + self.experimental_series_scan_v2 = enabled; + self + } + /// Sets whether to ignore inverted index. #[must_use] pub(crate) fn with_ignore_inverted_index(mut self, ignore: bool) -> Self { @@ -387,8 +408,9 @@ impl ScanRegion { /// Scans by series. #[tracing::instrument(skip_all, fields(region_id = %self.region_id()))] pub(crate) async fn series_scan(self) -> Result { + let experimental_series_scan_v2 = self.experimental_series_scan_v2; let input = self.scan_input().await?; - Ok(SeriesScan::new(input)) + Ok(SeriesScan::new(input, experimental_series_scan_v2)) } /// Returns true if the region can use unordered scan for current request. @@ -541,6 +563,7 @@ impl ScanRegion { .with_bloom_filter_index_appliers(bloom_filter_appliers) .with_fulltext_index_appliers(fulltext_index_appliers) .with_max_concurrent_scan_files(self.max_concurrent_scan_files) + .with_scan_memory_pool(self.scan_memory_pool) .with_start_time(self.start_time) .with_append_mode(self.version.options.append_mode) .with_filter_deleted(self.filter_deleted) @@ -899,6 +922,8 @@ pub struct ScanInput { ignore_file_not_found: bool, /// Maximum number of SST files to scan concurrently. pub(crate) max_concurrent_scan_files: usize, + /// Memory pool shared by internal scan operators across all queries. + pub(crate) scan_memory_pool: Arc, /// Index appliers. inverted_index_appliers: [Option; 2], bloom_filter_index_appliers: [Option; 2], @@ -950,6 +975,7 @@ impl ScanInput { cache_strategy: CacheStrategy::Disabled, ignore_file_not_found: false, max_concurrent_scan_files: DEFAULT_MAX_CONCURRENT_SCAN_FILES, + scan_memory_pool: Arc::new(UnboundedMemoryPool::default()), inverted_index_appliers: [None, None], bloom_filter_index_appliers: [None, None], fulltext_index_appliers: [None, None], @@ -1037,6 +1063,13 @@ impl ScanInput { self } + /// Sets the memory pool shared by internal scan operators. + #[must_use] + pub(crate) fn with_scan_memory_pool(mut self, scan_memory_pool: Arc) -> Self { + self.scan_memory_pool = scan_memory_pool; + self + } + /// Sets inverted index appliers. #[must_use] pub(crate) fn with_inverted_index_appliers( diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index 03d76d2adf6..2933327a38e 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -71,7 +71,6 @@ pub(crate) struct MetricSeriesId { pub(crate) type MetricSeriesIdStream = BoxStream<'static, Result>>; /// Builds candidate metric series from the ranges assigned to a [`SeriesScan`](super::series_scan::SeriesScan). -#[allow(dead_code)] pub(crate) struct SeriesCandidateScanner { stream_ctx: Arc, partitions: Vec>, @@ -82,7 +81,6 @@ pub(crate) struct SeriesCandidateScanner { part_metrics: PartitionMetrics, } -#[allow(dead_code)] impl SeriesCandidateScanner { /// Creates a candidate-series scanner for native memtable and SST ranges. /// @@ -346,8 +344,7 @@ impl SeriesCandidateRangeBuilder { } } -pub(crate) fn validate_metric_metadata(stream_ctx: &StreamContext) -> Result<()> { - let metadata = stream_ctx.input.region_metadata(); +pub(crate) fn is_sparse_metric_metadata(metadata: &store_api::metadata::RegionMetadataRef) -> bool { let valid_prefix = metadata .primary_key .starts_with(&[ReservedColumnId::table_id(), ReservedColumnId::tsid()]); @@ -358,8 +355,14 @@ pub(crate) fn validate_metric_metadata(stream_ctx: &StreamContext) -> Result<()> table_id.column_schema.data_type == ConcreteDataType::uint32_datatype() && tsid.column_schema.data_type == ConcreteDataType::uint64_datatype() }); + + metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse && valid_prefix && valid_types +} + +pub(crate) fn validate_metric_metadata(stream_ctx: &StreamContext) -> Result<()> { + let metadata = stream_ctx.input.region_metadata(); ensure!( - metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse && valid_prefix && valid_types, + is_sparse_metric_metadata(metadata), InvalidRequestSnafu { region_id: metadata.region_id, reason: "candidate-series scan requires sparse (__table_id, __tsid) primary keys", diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index d16d13f2dac..0a3e008f27a 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -84,12 +84,16 @@ impl SeriesRange { pub(crate) struct AssignedSeriesBatch { range: SeriesRange, series: Vec, + enable_range_cache: bool, } -#[allow(dead_code)] impl AssignedSeriesBatch { - fn new(range: SeriesRange, series: Vec) -> Self { - Self { range, series } + fn new(range: SeriesRange, series: Vec, enable_range_cache: bool) -> Self { + Self { + range, + series, + enable_range_cache, + } } pub(crate) fn range(&self) -> SeriesRange { @@ -99,23 +103,28 @@ impl AssignedSeriesBatch { pub(crate) fn series(&self) -> &[MetricSeriesId] { &self.series } + + pub(crate) fn enable_range_cache(&self) -> bool { + self.enable_range_cache + } } /// Collects candidate batches and assigns every TSID by its integer range. -#[allow(dead_code)] pub(crate) struct SeriesBatchCollector { assignments: Vec>, + num_series: usize, } -#[allow(dead_code)] impl SeriesBatchCollector { pub(crate) fn new(partitions: usize) -> Option { (partitions > 0).then(|| Self { assignments: (0..partitions).map(|_| Vec::new()).collect(), + num_series: 0, }) } pub(crate) fn push(&mut self, batch: Vec) { + self.num_series += batch.len(); let partitions = self.assignments.len(); for series in batch { let partition = SeriesRange::partition_for(series.tsid, partitions); @@ -123,13 +132,21 @@ impl SeriesBatchCollector { } } - pub(crate) fn finish(self) -> Vec { + pub(crate) fn len(&self) -> usize { + self.num_series + } + + pub(crate) fn finish(self, enable_range_cache: bool) -> Vec { let partitions = self.assignments.len(); self.assignments .into_iter() .enumerate() .map(|(partition, series)| { - AssignedSeriesBatch::new(SeriesRange::new(partition, partitions).unwrap(), series) + AssignedSeriesBatch::new( + SeriesRange::new(partition, partitions).unwrap(), + series, + enable_range_cache, + ) }) .collect() } @@ -140,14 +157,16 @@ impl SeriesBatchCollector { struct MetricSeriesFilter { range: SeriesRange, series: Arc>, + enable_range_cache: bool, } impl MetricSeriesFilter { fn new(assigned: &AssignedSeriesBatch) -> Self { - let series = assigned.series.iter().copied().collect(); + let series = assigned.series().iter().copied().collect(); Self { - range: assigned.range, + range: assigned.range(), series: Arc::new(series), + enable_range_cache: assigned.enable_range_cache(), } } @@ -232,7 +251,6 @@ fn filter_flat_stream_by_series( } /// Reads all collected metric series assigned to one partition. -#[allow(dead_code)] pub(crate) struct SeriesReader { stream_ctx: Arc, partition_ranges: Vec, @@ -244,7 +262,6 @@ pub(crate) struct SeriesReader { part_metrics: PartitionMetrics, } -#[allow(dead_code)] impl SeriesReader { /// Creates a reader for the series assigned to one data partition. /// @@ -354,7 +371,10 @@ async fn build_series_partition_range( partition_pruner: Arc, part_metrics: PartitionMetrics, ) -> Result<(BoxedRecordBatchStream, usize)> { - let cache_key = build_series_range_cache_key(&stream_ctx, &part_range, range); + let cache_key = filter + .enable_range_cache + .then(|| build_series_range_cache_key(&stream_ctx, &part_range, range)) + .flatten(); if let Some(key) = cache_key.as_ref() { if let Some(value) = stream_ctx.input.cache_strategy.get_range_result(key) { part_metrics.inc_range_cache_hit(); @@ -557,7 +577,7 @@ mod tests { ) -> AssignedSeriesBatch { let mut collector = SeriesBatchCollector::new(partitions).unwrap(); collector.push(series); - collector.finish().remove(partition) + collector.finish(true).remove(partition) } #[test] @@ -572,13 +592,13 @@ mod tests { ]; let mut first = SeriesBatchCollector::new(4).unwrap(); first.push(input.clone()); - let first = first.finish(); + let first = first.finish(true); let mut second = SeriesBatchCollector::new(4).unwrap(); for chunk in input.chunks(2) { second.push(chunk.to_vec()); } - let second = second.finish(); + let second = second.finish(true); assert_eq!( first @@ -601,6 +621,24 @@ mod tests { ); } + #[test] + fn collector_tracks_size_and_cache_policy() { + let mut collector = SeriesBatchCollector::new(2).unwrap(); + collector.push(vec![series(1, 1), series(1, u64::MAX)]); + collector.push(vec![series(2, 2)]); + assert_eq!(3, collector.len()); + + let assignments = collector.finish(false); + assert_eq!( + 3, + assignments + .iter() + .map(|batch| batch.series().len()) + .sum::() + ); + assert!(assignments.iter().all(|batch| !batch.enable_range_cache())); + } + #[test] fn series_ranges_cover_the_tsid_domain() { let ranges = (0..3) diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index adefae91173..1a1369830f3 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -15,6 +15,7 @@ //! Per-series scan implementation. use std::fmt; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -45,13 +46,15 @@ use crate::error::{ ScanSeriesSnafu, TooManyFilesToReadSnafu, }; use crate::read::ScannerMetrics; -use crate::read::pruner::{PartitionPruner, Pruner}; +use crate::read::pruner::{PartitionPruner, Pruner, PrunerOptions}; use crate::read::scan_region::{ScanInput, StreamContext}; use crate::read::scan_util::{ PartitionMetrics, PartitionMetricsList, SeriesDistributorMetrics, compute_average_batch_size, compute_parallel_channel_size, }; use crate::read::seq_scan::SeqScan; +use crate::read::series_candidate::{SeriesCandidateScanner, is_sparse_metric_metadata}; +use crate::read::series_reader::{AssignedSeriesBatch, SeriesBatchCollector, SeriesReader}; use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::format::PrimaryKeyArray; @@ -59,8 +62,57 @@ use crate::sst::parquet::format::PrimaryKeyArray; /// Timeout to send a batch to a sender. const SEND_TIMEOUT: Duration = Duration::from_micros(100); -/// List of receivers. -type ReceiverList = Vec>>>; +/// Maximum number of candidate series retained before distributing assignments. +const CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD: usize = 1_000_000; + +/// Legacy data receivers for output partitions. +type LegacyReceiverList = Vec>>>; + +/// Candidate assignment receivers for output partitions. +type CandidateReceiverList = Vec>; + +/// Candidate assignment senders for output partitions. +type CandidateSenderList = Vec>>>; + +struct CandidateReceiver { + receiver: mpsc::UnboundedReceiver>, + active_receivers: Arc, +} + +impl CandidateReceiver { + async fn recv(&mut self) -> Option> { + self.receiver.recv().await + } +} + +impl Drop for CandidateReceiver { + fn drop(&mut self) { + let previous = self.active_receivers.fetch_sub(1, Ordering::Relaxed); + debug_assert!(previous > 0); + } +} + +/// Input required by a partition-local series reader. +struct SeriesReaderInput { + assigned_series: AssignedSeriesBatch, + partition_pruner: Arc, + range_semaphore: Arc, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SeriesScanMode { + Legacy, + TwoPhase, +} + +impl SeriesScanMode { + fn as_str(self) -> &'static str { + match self { + Self::Legacy => "legacy", + Self::TwoPhase => "two_phase", + } + } +} /// Scans a region and returns sorted rows of a series in the same partition. /// @@ -68,14 +120,18 @@ type ReceiverList = Vec>>>; /// partition. /// Always returns the same series (primary key) to the same partition. pub struct SeriesScan { + /// Implementation used by this scan. + mode: SeriesScanMode, /// Properties of the scanner. properties: ScannerProperties, /// Context of streams. stream_ctx: Arc, /// Shared pruner for file range building. pruner: Arc, - /// Receivers of each partition. - receivers: Mutex, + /// Legacy data receivers for each partition. + legacy_receivers: Mutex, + /// Candidate assignment receivers for each partition. + candidate_receivers: Mutex, /// Metrics for each partition. /// The scanner only sets in query and keeps it empty during compaction. metrics_list: Arc, @@ -83,7 +139,12 @@ pub struct SeriesScan { impl SeriesScan { /// Creates a new [SeriesScan]. - pub(crate) fn new(input: ScanInput) -> Self { + pub(crate) fn new(input: ScanInput, experimental_series_scan_v2: bool) -> Self { + let mode = if experimental_series_scan_v2 && Self::supports_two_phase(&input) { + SeriesScanMode::TwoPhase + } else { + SeriesScanMode::Legacy + }; let mut properties = ScannerProperties::default() .with_append_mode(input.append_mode) .with_total_rows(input.total_rows()); @@ -95,17 +156,40 @@ impl SeriesScan { // Create the shared pruner with number of workers equal to CPU cores. let num_workers = common_stat::get_total_cpu_cores().max(1); - let pruner = Arc::new(Pruner::new(stream_ctx.clone(), num_workers)); + let pruner = match mode { + SeriesScanMode::Legacy => Arc::new(Pruner::new(stream_ctx.clone(), num_workers)), + SeriesScanMode::TwoPhase => Arc::new(Pruner::new_with_options( + stream_ctx.clone(), + num_workers, + PrunerOptions { + retain_builders: true, + enable_predicate_prefilter: false, + }, + )), + }; Self { + mode, properties, stream_ctx, pruner, - receivers: Mutex::new(Vec::new()), + legacy_receivers: Mutex::new(Vec::new()), + candidate_receivers: Mutex::new(Vec::new()), metrics_list: Arc::new(PartitionMetricsList::default()), } } + fn supports_two_phase(input: &ScanInput) -> bool { + if !is_sparse_metric_metadata(input.region_metadata()) { + return false; + } + #[cfg(feature = "enterprise")] + if !input.extension_ranges().is_empty() { + return false; + } + true + } + #[tracing::instrument( skip_all, fields( @@ -174,60 +258,128 @@ impl SeriesScan { } ); - self.maybe_start_distributor(metrics_set, &self.metrics_list, ctx.explain_verbose); + match self.mode { + SeriesScanMode::Legacy => self.scan_legacy_batch_in_partition( + partition, + part_metrics, + metrics_set, + ctx.explain_verbose, + ), + SeriesScanMode::TwoPhase => self.scan_two_phase_batch_in_partition( + partition, + part_metrics, + metrics_set, + ctx.explain_verbose, + ), + } + } - let mut receiver = self.take_receiver(partition)?; + fn scan_legacy_batch_in_partition( + &self, + partition: usize, + part_metrics: PartitionMetrics, + metrics_set: &ExecutionPlanMetricsSet, + explain_verbose: bool, + ) -> Result { + self.maybe_start_legacy_distributor(metrics_set, explain_verbose); + let mut receiver = self.legacy_receivers.lock().unwrap()[partition] + .take() + .context(ScanMultiTimesSnafu { partition })?; let stream = try_stream! { part_metrics.on_first_poll(); let mut fetch_start = Instant::now(); + let mut metrics = ScannerMetrics::default(); while let Some(series) = receiver.recv().await { let series = series?; - - let mut metrics = ScannerMetrics::default(); metrics.scan_cost += fetch_start.elapsed(); - fetch_start = Instant::now(); - metrics.num_batches += series.num_batches(); metrics.num_rows += series.num_rows(); let yield_start = Instant::now(); yield ScanBatch::Series(series); metrics.yield_cost += yield_start.elapsed(); - - part_metrics.merge_metrics(&metrics); + fetch_start = Instant::now(); } + part_metrics.merge_metrics(&metrics); part_metrics.on_finish(); }; Ok(Box::pin(stream)) } - /// Takes the receiver for the partition. - fn take_receiver(&self, partition: usize) -> Result>> { - let mut rx_list = self.receivers.lock().unwrap(); - rx_list[partition] + fn scan_two_phase_batch_in_partition( + &self, + partition: usize, + part_metrics: PartitionMetrics, + metrics_set: &ExecutionPlanMetricsSet, + explain_verbose: bool, + ) -> Result { + self.maybe_start_candidate_distributor(metrics_set, explain_verbose); + // Safety: `scan_batch_in_partition` validates the index, and the receiver + // list always contains one entry for each output partition. + let mut receiver = self.candidate_receivers.lock().unwrap()[partition] .take() - .context(ScanMultiTimesSnafu { partition }) + .context(ScanMultiTimesSnafu { partition })?; + let stream_ctx = self.stream_ctx.clone(); + let partition_ranges = self + .properties + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + let stream = try_stream! { + part_metrics.on_first_poll(); + + let mut fetch_start = Instant::now(); + let mut metrics = ScannerMetrics::default(); + while let Some(input) = receiver.recv().await { + let input = input?; + metrics.scan_cost += fetch_start.elapsed(); + + let build_start = Instant::now(); + let reader = SeriesReader::try_new( + stream_ctx.clone(), + partition_ranges.clone(), + input.assigned_series, + input.partition_pruner, + input.range_semaphore, + part_metrics.clone(), + )?; + let mut reader_stream = reader.build_stream().await?; + metrics.scan_cost += build_start.elapsed(); + fetch_start = Instant::now(); + while let Some(record_batch) = reader_stream.try_next().await? { + metrics.scan_cost += fetch_start.elapsed(); + metrics.num_batches += 1; + metrics.num_rows += record_batch.num_rows(); + + let yield_start = Instant::now(); + yield ScanBatch::RecordBatch(record_batch); + metrics.yield_cost += yield_start.elapsed(); + fetch_start = Instant::now(); + } + } + metrics.scan_cost += fetch_start.elapsed(); + + part_metrics.merge_metrics(&metrics); + part_metrics.on_finish(); + }; + Ok(Box::pin(stream)) } - /// Starts the distributor if the receiver list is empty. - #[tracing::instrument( - skip(self, metrics_set, metrics_list), - fields(region_id = %self.stream_ctx.input.mapper.metadata().region_id) - )] - fn maybe_start_distributor( + fn maybe_start_legacy_distributor( &self, metrics_set: &ExecutionPlanMetricsSet, - metrics_list: &Arc, explain_verbose: bool, ) { - let mut rx_list = self.receivers.lock().unwrap(); + let mut rx_list = self.legacy_receivers.lock().unwrap(); if !rx_list.is_empty() { return; } - let (senders, receivers) = new_channel_list(self.properties.num_partitions()); + let (senders, receivers) = new_legacy_channel_list(self.properties.num_partitions()); let mut distributor = SeriesDistributor { stream_ctx: self.stream_ctx.clone(), range_semaphore: Some(Arc::new(Semaphore::new(self.properties.num_partitions()))), @@ -236,7 +388,7 @@ impl SeriesScan { pruner: self.pruner.clone(), senders, metrics_set: metrics_set.clone(), - metrics_list: metrics_list.clone(), + metrics_list: self.metrics_list.clone(), explain_verbose, }; let region_id = distributor.stream_ctx.input.mapper.metadata().region_id; @@ -251,6 +403,41 @@ impl SeriesScan { *rx_list = receivers; } + fn maybe_start_candidate_distributor( + &self, + metrics_set: &ExecutionPlanMetricsSet, + explain_verbose: bool, + ) { + let mut rx_list = self.candidate_receivers.lock().unwrap(); + if !rx_list.is_empty() { + return; + } + + let (senders, receivers, active_receivers) = + new_candidate_channel_list(self.properties.num_partitions()); + let mut distributor = SeriesCandidateDistributor { + stream_ctx: self.stream_ctx.clone(), + range_semaphore: Arc::new(Semaphore::new(self.properties.num_partitions())), + partitions: self.properties.partitions.clone(), + pruner: self.pruner.clone(), + senders, + active_receivers, + metrics_set: metrics_set.clone(), + metrics_list: self.metrics_list.clone(), + explain_verbose, + }; + let region_id = distributor.stream_ctx.input.mapper.metadata().region_id; + let span = tracing::info_span!("SeriesScan::candidate_distributor", region_id = %region_id); + common_runtime::spawn_query( + async move { + distributor.execute().await; + } + .instrument(span), + ); + + *rx_list = receivers; + } + /// Scans the region and returns a stream. #[tracing::instrument( skip_all, @@ -319,7 +506,7 @@ impl SeriesScan { } } -fn new_channel_list(num_partitions: usize) -> (SenderList, ReceiverList) { +fn new_legacy_channel_list(num_partitions: usize) -> (SenderList, LegacyReceiverList) { let (senders, receivers): (Vec<_>, Vec<_>) = (0..num_partitions) .map(|_| { let (sender, receiver) = mpsc::channel(1); @@ -329,6 +516,26 @@ fn new_channel_list(num_partitions: usize) -> (SenderList, ReceiverList) { (SenderList::new(senders), receivers) } +fn new_candidate_channel_list( + num_partitions: usize, +) -> (CandidateSenderList, CandidateReceiverList, Arc) { + let active_receivers = Arc::new(AtomicUsize::new(num_partitions)); + let (senders, receivers) = (0..num_partitions) + .map(|_| { + // Partition streams can be consumed sequentially. A bounded channel for an unpolled + // partition could block the distributor and prevent it from sending more assignments + // to the partition currently being consumed, causing a deadlock. + let (sender, receiver) = mpsc::unbounded_channel(); + let receiver = CandidateReceiver { + receiver, + active_receivers: active_receivers.clone(), + }; + (Some(sender), Some(receiver)) + }) + .unzip(); + (senders, receivers, active_receivers) +} + impl RegionScanner for SeriesScan { fn name(&self) -> &str { "SeriesScan" @@ -402,19 +609,24 @@ impl DisplayAs for SeriesScan { )?; match t { DisplayFormatType::Default | DisplayFormatType::TreeRender => { - self.stream_ctx.format_for_explain(false, f) + self.stream_ctx.format_for_explain(false, f)?; } DisplayFormatType::Verbose => { self.stream_ctx.format_for_explain(true, f)?; - self.metrics_list.format_verbose_metrics(f) } } + write!(f, ", \"mode\":\"{}\"", self.mode.as_str())?; + if matches!(t, DisplayFormatType::Verbose) { + self.metrics_list.format_verbose_metrics(f)?; + } + Ok(()) } } impl fmt::Debug for SeriesScan { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("SeriesScan") + .field("mode", &self.mode) .field("num_ranges", &self.stream_ctx.ranges.len()) .finish() } @@ -426,6 +638,11 @@ impl SeriesScan { pub(crate) fn input(&self) -> &ScanInput { &self.stream_ctx.input } + + /// Returns the scan mode for tests. + pub(crate) fn mode(&self) -> &'static str { + self.mode.as_str() + } } /// The distributor scans series and distributes them to different partitions. @@ -463,9 +680,7 @@ impl SeriesDistributor { fields(region_id = %self.stream_ctx.input.mapper.metadata().region_id) )] async fn execute(&mut self) { - let result = self.scan_partitions_flat().await; - - if let Err(e) = result { + if let Err(e) = self.scan_partitions_flat().await { self.senders.send_error(e).await; } } @@ -605,6 +820,113 @@ impl SeriesDistributor { } } +/// Discovers metric series and sends one reader assignment to each partition. +struct SeriesCandidateDistributor { + stream_ctx: Arc, + range_semaphore: Arc, + partitions: Vec>, + pruner: Arc, + senders: CandidateSenderList, + active_receivers: Arc, + metrics_set: ExecutionPlanMetricsSet, + metrics_list: Arc, + explain_verbose: bool, +} + +impl SeriesCandidateDistributor { + async fn execute(&mut self) { + if let Err(e) = self.distribute().await { + self.send_error(e); + } + } + + async fn distribute(&mut self) -> Result<()> { + let part_metrics = new_partition_metrics( + &self.stream_ctx, + self.explain_verbose, + &self.metrics_set, + self.partitions.len(), + &self.metrics_list, + ); + part_metrics.on_first_poll(); + + let candidate_scanner = SeriesCandidateScanner::try_new( + self.stream_ctx.clone(), + self.partitions.clone(), + self.pruner.clone(), + self.range_semaphore.clone(), + self.stream_ctx.input.scan_memory_pool.clone(), + self.metrics_set.clone(), + part_metrics.clone(), + )?; + let partition_pruner = candidate_scanner.partition_pruner(); + let mut candidates = candidate_scanner.build_stream().await?; + let mut collector = + SeriesBatchCollector::new(self.partitions.len()).context(InvalidSenderSnafu)?; + let mut chunked = false; + while let Some(batch) = candidates.try_next().await? { + if !self.should_fetch_candidates() { + part_metrics.on_finish(); + return Ok(()); + } + collector.push(batch); + if collector.len() >= CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD { + chunked = true; + self.send_assignments(collector.finish(false), &partition_pruner); + if !self.should_fetch_candidates() { + part_metrics.on_finish(); + return Ok(()); + } + collector = + SeriesBatchCollector::new(self.partitions.len()).context(InvalidSenderSnafu)?; + } + } + + if collector.len() > 0 { + self.send_assignments(collector.finish(!chunked), &partition_pruner); + } + part_metrics.on_finish(); + Ok(()) + } + + fn send_assignments( + &mut self, + assignments: Vec, + partition_pruner: &Arc, + ) { + for (partition, assigned_series) in assignments.into_iter().enumerate() { + if assigned_series.series().is_empty() { + continue; + } + let Some(sender) = self.senders[partition].as_ref() else { + continue; + }; + let sent = sender + .send(Ok(SeriesReaderInput { + assigned_series, + partition_pruner: partition_pruner.clone(), + range_semaphore: self.range_semaphore.clone(), + })) + .is_ok(); + if !sent { + self.senders[partition] = None; + } + } + } + + fn should_fetch_candidates(&self) -> bool { + self.active_receivers.load(Ordering::Relaxed) > 0 + } + + fn send_error(&mut self, error: Error) { + let error = Arc::new(error); + for sender in self.senders.iter_mut().filter_map(Option::take) { + let result = Err(error.clone()).context(ScanSeriesSnafu); + let _ = sender.send(result); + } + } +} + /// Batches of the same series. #[derive(Debug)] pub enum SeriesBatch { @@ -893,6 +1215,18 @@ mod tests { use super::*; + #[test] + fn candidate_distributor_stops_after_all_receivers_close() { + let (_senders, mut receivers, active_receivers) = new_candidate_channel_list(2); + assert_eq!(2, active_receivers.load(Ordering::Relaxed)); + + drop(receivers[0].take()); + assert_eq!(1, active_receivers.load(Ordering::Relaxed)); + + drop(receivers[1].take()); + assert_eq!(0, active_receivers.load(Ordering::Relaxed)); + } + fn new_test_record_batch( primary_keys: &[&[u8]], timestamps: &[i64], diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index b2c8fe394e2..530b581a2ca 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2258,6 +2258,7 @@ scan_memory_on_exhausted = "fail" min_compaction_interval = "0s" schedule_compaction_after_edit = true default_flat_format = true +experimental_series_scan_v2 = true [region_engine.mito.index] aux_path = "" diff --git a/tests/cases/distributed/explain/step_aggr_advance.result b/tests/cases/distributed/explain/step_aggr_advance.result index 9fc8dfe5a1c..59f2946e8bc 100644 --- a/tests/cases/distributed/explain/step_aggr_advance.result +++ b/tests/cases/distributed/explain/step_aggr_advance.result @@ -67,7 +67,7 @@ tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -79,7 +79,7 @@ tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -144,7 +144,7 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[__sum_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -154,7 +154,7 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -219,7 +219,7 @@ tql analyze (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[__avg_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -229,7 +229,7 @@ tql analyze (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -284,7 +284,7 @@ tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -294,7 +294,7 @@ tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -359,7 +359,7 @@ tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -369,7 +369,7 @@ tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -429,7 +429,7 @@ tql analyze sum(aggr_optimize_not); |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, greptime_value@5 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[__sum_state(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -437,7 +437,7 @@ tql analyze sum(aggr_optimize_not); |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, greptime_value@5 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -563,7 +563,7 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp, sum(prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000)))@4 as sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED |_|_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED @@ -576,7 +576,7 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/distributed/tql-explain-analyze/analyze.result b/tests/cases/distributed/tql-explain-analyze/analyze.result index bdd7c4811e1..059cf55df64 100644 --- a/tests/cases/distributed/tql-explain-analyze/analyze.result +++ b/tests/cases/distributed/tql-explain-analyze/analyze.result @@ -26,7 +26,7 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -51,7 +51,7 @@ TQL ANALYZE (0, 10, '1s', '2s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[2000], interval=[1000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -75,7 +75,7 @@ TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -134,12 +134,12 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -166,7 +166,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[j@0 as j, prom_rate(j_range,i,test.j,Int64(10000))@1 as prom_rate(j_range,i,j,Int64(10000)), k@2 as k, l@3 as l] REDACTED |_|_|_FilterExec: prom_rate(j_range,i,test.j,Int64(10000))@1 IS NOT NULL REDACTED @@ -175,7 +175,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/histogram_multi_partition.result b/tests/cases/standalone/common/promql/histogram_multi_partition.result index 3d5490f38a1..93168ce449f 100644 --- a/tests/cases/standalone/common/promql/histogram_multi_partition.result +++ b/tests/cases/standalone/common/promql/histogram_multi_partition.result @@ -56,7 +56,7 @@ tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bu |_|_|_ProjectionExec: expr=[ts@0 as ts, le@1 as le, val@3 as val] REDACTED |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[10000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard", "le"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[le@0 as le, ts@1 as ts], aggr=[__sum_state(histogram_gap_bucket.val)] REDACTED |_|_|_RepartitionExec: partitioning=Hash([le@0, ts@1],REDACTED @@ -64,7 +64,7 @@ tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bu |_|_|_ProjectionExec: expr=[ts@0 as ts, le@1 as le, val@3 as val] REDACTED |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[10000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard", "le"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result index 72a6dc9c5b4..a3983260e33 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result @@ -127,13 +127,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -167,13 +167,13 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left + tsid_binary_join_right) / tsid |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -210,19 +210,19 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left + tsid_binary_join_right) * (ts |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -255,13 +255,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -294,13 +294,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_ |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -334,13 +334,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -373,13 +373,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -418,19 +418,19 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -473,25 +473,25 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -529,7 +529,7 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[ts@2 as ts, __tsid@3 as __tsid, greptime_value@4 as greptime_value, host@5 as host, job@6 as job, CASE WHEN __common_expr_1@0 IS NOT NULL THEN __common_expr_1@0 ELSE_END as __promql_or_match_0, CASE WHEN __common_expr_2@1 IS NOT NULL THEN __common_expr_2@1 ELSE_END as __promql_or_match_1] REDACTED |_|_|_ProjectionExec: expr=[CAST(host@1 AS Utf8) as __common_expr_1, CAST(job@2 AS Utf8) as __common_expr_2, ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job] REDACTED @@ -537,13 +537,13 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -582,19 +582,19 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result index ba2e97bcc2d..af5f3a9a4f8 100644 --- a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result +++ b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result @@ -118,7 +118,7 @@ TQL ANALYZE (0, 10, '5s') histogram_quantile(0.5, tsid_no_aggr_histogram_bucket) |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, le@4 as le, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ diff --git a/tests/cases/standalone/common/select/tql_filter.result b/tests/cases/standalone/common/select/tql_filter.result index 345671376f7..b4fbfdb5862 100644 --- a/tests/cases/standalone/common/select/tql_filter.result +++ b/tests/cases/standalone/common/select/tql_filter.result @@ -24,7 +24,7 @@ tql analyze (1, 3, '1s') t1{ a = "a" }; | 1_| 0_|_PromInstantManipulateExec: range=[1000..3000], lookback=[300000], interval=[1000], time index=[b] REDACTED |_|_|_PromSeriesDivideExec: tags=["a"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -47,7 +47,7 @@ tql analyze (1, 3, '1s') t1{ a =~ ".*" }; | 1_| 0_|_PromInstantManipulateExec: range=[1000..3000], lookback=[300000], interval=[1000], time index=[b] REDACTED |_|_|_PromSeriesDivideExec: tags=["a"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 6_| +-+-+-+ @@ -70,7 +70,7 @@ tql analyze (1, 3, '1s') t1{ a =~ "a.*" }; | 1_| 0_|_PromInstantManipulateExec: range=[1000..3000], lookback=[300000], interval=[1000], time index=[b] REDACTED |_|_|_PromSeriesDivideExec: tags=["a"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ diff --git a/tests/cases/standalone/common/tql/partition.result b/tests/cases/standalone/common/tql/partition.result index f1072edbb32..e32d13043bd 100644 --- a/tests/cases/standalone/common/tql/partition.result +++ b/tests/cases/standalone/common/tql/partition.result @@ -32,7 +32,7 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -82,7 +82,7 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[k@0 as k, j@1 as j], aggr=[__avg_state(prom_irate(j_range,i))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -92,7 +92,7 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -141,10 +141,10 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/windowed_sort.result b/tests/cases/standalone/optimizer/windowed_sort.result index 544563ad1a7..3be0dd1c3e4 100644 --- a/tests/cases/standalone/optimizer/windowed_sort.result +++ b/tests/cases/standalone/optimizer/windowed_sort.result @@ -290,7 +290,7 @@ TQL ANALYZE('2023-10-01 00:00:00.000'::TIMESTAMP, '2023-10-01 00:00:05.000'::TIM |_|_|_ProjectionExec: expr=[greptime_timestamp@0 as greptime_timestamp, greptime_value@1 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[1696118400000..1696118405000], lookback=[300000], interval=[1000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["instance", "job"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ diff --git a/tests/cases/standalone/tql-explain-analyze/analyze.result b/tests/cases/standalone/tql-explain-analyze/analyze.result index 203debf90bf..954020ea4cf 100644 --- a/tests/cases/standalone/tql-explain-analyze/analyze.result +++ b/tests/cases/standalone/tql-explain-analyze/analyze.result @@ -26,7 +26,7 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -51,7 +51,7 @@ TQL ANALYZE (0, 10, '1s', '2s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[2000], interval=[1000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -75,7 +75,7 @@ TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -134,12 +134,12 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -165,7 +165,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_FilterExec: prom_rate(j_range,i,j,Int64(10000))@1 IS NOT NULL REDACTED |_|_|_ProjectionExec: expr=[j@1 as j, prom_rate(j_range@4, i@0, j@1, 10000) as prom_rate(j_range,i,j,Int64(10000)), k@2 as k, l@3 as l] REDACTED @@ -173,7 +173,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -238,12 +238,12 @@ TQL ANALYZE FORMAT TEXT (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -297,7 +297,7 @@ TQL ANALYZE sum(test2); |_|_|_ProjectionExec: expr=[greptime_timestamp@0 as greptime_timestamp, greptime_value@1 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[__sum_state(test2.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -305,7 +305,7 @@ TQL ANALYZE sum(test2); |_|_|_ProjectionExec: expr=[greptime_timestamp@0 as greptime_timestamp, greptime_value@1 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/tql-explain-analyze/column_pruning.result b/tests/cases/standalone/tql-explain-analyze/column_pruning.result index 4bce77d7832..cecb72fcee1 100644 --- a/tests/cases/standalone/tql-explain-analyze/column_pruning.result +++ b/tests/cases/standalone/tql-explain-analyze/column_pruning.result @@ -45,7 +45,7 @@ TQL ANALYZE (0, 10, '5s') sum(promql_column_pruning); |_|_|_ProjectionExec: expr=[ts@0 as ts, greptime_value@4 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["job", "instance", "region"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -75,7 +75,7 @@ TQL ANALYZE (0, 10, '5s') sum(rate(promql_column_pruning[5s])); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[5000], eval range=[5000], time index=[ts] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["job", "instance", "region"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/tql-explain-analyze/tsid_column.result b/tests/cases/standalone/tql-explain-analyze/tsid_column.result index 8d27a55ab38..cf2deb76647 100644 --- a/tests/cases/standalone/tql-explain-analyze/tsid_column.result +++ b/tests/cases/standalone/tql-explain-analyze/tsid_column.result @@ -54,7 +54,7 @@ TQL ANALYZE (0, 10, '5s') sum(tsid_metric); |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -82,7 +82,7 @@ TQL ANALYZE (0, 10, '5s') sum by (job, instance) (tsid_metric); |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, instance@3 as instance, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 6_| +-+-+-+ @@ -120,7 +120,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -133,7 +133,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -173,7 +173,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -186,7 +186,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ From c3cd186997538e2ed665ed32ce1755baf62bd857 Mon Sep 17 00:00:00 2001 From: Morax <100508620+fzlzjerry@users.noreply.github.com> Date: Fri, 14 Aug 2026 07:04:51 +0000 Subject: [PATCH 071/126] fix(query): respect query timezone in timestamp casts (#8859) * fix(query): respect timezone in insert values Signed-off-by: Morax * fix(query): normalize timestamp casts with query timezone Signed-off-by: Morax * fix(query): scope timestamp conversion to insert assignments Signed-off-by: Morax * style(query): simplify Arc usage Signed-off-by: Morax --------- Signed-off-by: Morax --- src/query/src/optimizer/type_conversion.rs | 284 +++++++++++++++--- src/query/src/planner.rs | 70 ++++- .../insert/insert_default_timezone.result | 39 +++ .../common/insert/insert_default_timezone.sql | 18 ++ 4 files changed, 370 insertions(+), 41 deletions(-) diff --git a/src/query/src/optimizer/type_conversion.rs b/src/query/src/optimizer/type_conversion.rs index 1702aacce83..366a7d8a31d 100644 --- a/src/query/src/optimizer/type_conversion.rs +++ b/src/query/src/optimizer/type_conversion.rs @@ -12,14 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use common_time::Timezone; use common_time::timestamp::{TimeUnit, Timestamp}; use datafusion::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; use datafusion_common::{DFSchemaRef, DataFusionError, Result, ScalarValue}; -use datafusion_expr::expr::InList; +use datafusion_expr::expr::{Cast, InList}; use datafusion_expr::{ - Between, BinaryExpr, Expr, ExprSchemable, Filter, LogicalPlan, Operator, TableScan, + Between, BinaryExpr, Expr, ExprSchemable, Filter, LogicalPlan, Operator, TableScan, WriteOp, }; use datatypes::arrow::compute; use datatypes::arrow::datatypes::DataType; @@ -32,7 +34,7 @@ use crate::plan::ExtractExpr; /// TypeConversionRule converts some literal values in logical plan to other types according /// to data type of corresponding columns. /// Specifically: -/// - string literal of timestamp is converted to `Expr::Literal(ScalarValue::TimestampMillis)` +/// - string literal of timestamp is converted to the target timestamp type /// - string literal of boolean is converted to `Expr::Literal(ScalarValue::Boolean)` pub struct TypeConversionRule; @@ -45,10 +47,8 @@ impl ExtensionAnalyzerRule for TypeConversionRule { ) -> Result { plan.transform_up_with_subqueries(|plan| match plan { LogicalPlan::Filter(filter) => { - let mut converter = TypeConverter { - schema: filter.input.schema().clone(), - query_ctx: ctx.query_ctx(), - }; + let mut converter = + TypeConverter::new(filter.input.schema().clone(), ctx.query_ctx()); let rewritten = filter.predicate.clone().rewrite(&mut converter)?.data; Ok(Transformed::yes(LogicalPlan::Filter(Filter::try_new( rewritten, @@ -63,10 +63,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule { filters, fetch, }) => { - let mut converter = TypeConverter { - schema: projected_schema.clone(), - query_ctx: ctx.query_ctx(), - }; + let mut converter = TypeConverter::new(projected_schema.clone(), ctx.query_ctx()); let rewrite_filters = filters .into_iter() .map(|e| e.rewrite(&mut converter).map(|x| x.data)) @@ -80,8 +77,18 @@ impl ExtensionAnalyzerRule for TypeConversionRule { fetch, }))) } - LogicalPlan::Projection { .. } - | LogicalPlan::Window { .. } + LogicalPlan::Projection { .. } => { + let mut converter = TypeConverter::new(plan.schema().clone(), ctx.query_ctx()); + let inputs = plan.inputs().into_iter().cloned().collect::>(); + let expr = plan + .expressions_consider_join() + .into_iter() + .map(|e| e.rewrite(&mut converter).map(|x| x.data)) + .collect::>>()?; + + plan.with_new_exprs(expr, inputs).map(Transformed::yes) + } + LogicalPlan::Window { .. } | LogicalPlan::Aggregate { .. } | LogicalPlan::Repartition { .. } | LogicalPlan::Extension { .. } @@ -89,10 +96,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule { | LogicalPlan::Union { .. } | LogicalPlan::Values { .. } | LogicalPlan::Analyze { .. } => { - let mut converter = TypeConverter { - schema: plan.schema().clone(), - query_ctx: ctx.query_ctx(), - }; + let mut converter = TypeConverter::new(plan.schema().clone(), ctx.query_ctx()); let inputs = plan.inputs().into_iter().cloned().collect::>(); let expr = plan .expressions_consider_join() @@ -107,10 +111,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule { let Ok(schema) = join.left.schema().join(join.right.schema()) else { return Ok(Transformed::no(LogicalPlan::Join(join))); }; - let mut converter = TypeConverter { - schema: std::sync::Arc::new(schema), - query_ctx: ctx.query_ctx(), - }; + let mut converter = TypeConverter::new(Arc::new(schema), ctx.query_ctx()); let plan = LogicalPlan::Join(join); let inputs = plan.inputs().into_iter().cloned().collect::>(); let expr = plan @@ -122,6 +123,14 @@ impl ExtensionAnalyzerRule for TypeConversionRule { plan.with_new_exprs(expr, inputs).map(Transformed::yes) } + LogicalPlan::Dml(mut dml) if matches!(dml.op, WriteOp::Insert(_)) => { + dml.input = Arc::new(rewrite_insert_assignments( + dml.input.as_ref().clone(), + ctx.query_ctx(), + )?); + Ok(Transformed::yes(LogicalPlan::Dml(dml))) + } + LogicalPlan::Distinct { .. } | LogicalPlan::Limit { .. } | LogicalPlan::Subquery { .. } @@ -146,6 +155,10 @@ struct TypeConverter { } impl TypeConverter { + fn new(schema: DFSchemaRef, query_ctx: QueryContextRef) -> Self { + Self { query_ctx, schema } + } + fn column_type(&self, expr: &Expr) -> Option { if let Expr::Column(_) = expr && let Ok(v) = expr.get_type(&self.schema) @@ -162,7 +175,7 @@ impl TypeConverter { ) -> Result { match (target_type, value) { (DataType::Timestamp(_, _), ScalarValue::Utf8(Some(v))) => { - string_to_timestamp_ms(v, Some(&self.query_ctx.timezone())) + parse_string_to_timestamp(v, Some(&self.query_ctx.timezone())) } (DataType::Boolean, ScalarValue::Utf8(Some(v))) => match v.to_lowercase().as_str() { "true" => Ok(ScalarValue::Boolean(Some(true))), @@ -294,6 +307,120 @@ impl TreeNodeRewriter for TypeConverter { } } +/// Rewrites timestamp assignment casts at the INSERT boundary. The source query +/// is intentionally left untouched so explicit casts keep DataFusion semantics. +fn rewrite_insert_assignments( + plan: LogicalPlan, + query_ctx: QueryContextRef, +) -> Result { + let LogicalPlan::Projection(projection) = plan else { + return Ok(plan); + }; + + // VALUES type coercion is performed inside the Values node. Only rewrite + // the direct source of the assignment projection, not nested query plans. + let input = rewrite_insert_values(projection.input.as_ref().clone(), query_ctx.clone())?; + let input_literals = projected_literals(&input); + let mut converter = InsertAssignmentConverter { + query_ctx, + input_literals, + }; + let plan = LogicalPlan::Projection(projection); + let expr = plan + .expressions_consider_join() + .into_iter() + .map(|expr| expr.rewrite(&mut converter).map(|x| x.data)) + .collect::>>()?; + + plan.with_new_exprs(expr, vec![input]) +} + +fn rewrite_insert_values(plan: LogicalPlan, query_ctx: QueryContextRef) -> Result { + let LogicalPlan::Values(mut values) = plan else { + return Ok(plan); + }; + let mut converter = InsertAssignmentConverter { + query_ctx, + input_literals: None, + }; + values.values = values + .values + .into_iter() + .map(|row| { + row.into_iter() + .map(|expr| expr.rewrite(&mut converter).map(|x| x.data)) + .collect::>>() + }) + .collect::>>()?; + + Ok(LogicalPlan::Values(values)) +} + +fn projected_literals(plan: &LogicalPlan) -> Option<(DFSchemaRef, Vec>)> { + let LogicalPlan::Projection(projection) = plan else { + return None; + }; + let literals = projection.expr.iter().map(projected_literal).collect(); + Some((projection.schema.clone(), literals)) +} + +fn projected_literal(expr: &Expr) -> Option { + match expr { + Expr::Literal(value, _) => Some(value.clone()), + Expr::Alias(alias) => projected_literal(&alias.expr), + _ => None, + } +} + +struct InsertAssignmentConverter { + query_ctx: QueryContextRef, + input_literals: Option<(DFSchemaRef, Vec>)>, +} + +impl InsertAssignmentConverter { + fn assignment_literal(&self, expr: &Expr) -> Option { + let Expr::Column(column) = expr else { + return expr.as_literal().cloned(); + }; + let (schema, literals) = self.input_literals.as_ref()?; + let index = schema.maybe_index_of_column(column)?; + literals.get(index)?.clone() + } +} + +impl TreeNodeRewriter for InsertAssignmentConverter { + type Node = Expr; + + fn f_up(&mut self, expr: Expr) -> Result> { + let Expr::Cast(Cast { + expr: inner, + data_type: target_type, + }) = expr + else { + return Ok(Transformed::no(expr)); + }; + let original = || { + Expr::Cast(Cast { + expr: inner.clone(), + data_type: target_type.clone(), + }) + }; + if !matches!(&target_type, DataType::Timestamp(_, _)) { + return Ok(Transformed::no(original())); + } + let Some(ScalarValue::Utf8(Some(value))) = self.assignment_literal(inner.as_ref()) else { + return Ok(Transformed::no(original())); + }; + + match cast_string_to_timestamp(&value, &target_type, Some(&self.query_ctx.timezone())) { + Ok(value) if !value.is_null() => Ok(Transformed::yes(Expr::Literal(value, None))), + // Preserve DataFusion's existing parser for unsupported timestamp + // syntax and out-of-range values. + Ok(_) | Err(_) => Ok(Transformed::no(original())), + } + } +} + fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr { let timestamp = match unit { TimeUnit::Second => val * 1_000, @@ -308,18 +435,34 @@ fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr { ) } -fn string_to_timestamp_ms(string: &str, timezone: Option<&Timezone>) -> Result { +fn cast_string_to_timestamp( + string: &str, + target_type: &DataType, + timezone: Option<&Timezone>, +) -> Result { + let parsed = parse_string_to_timestamp(string, timezone)?; + cast_timestamp(parsed, target_type) +} + +fn parse_string_to_timestamp(string: &str, timezone: Option<&Timezone>) -> Result { let ts = Timestamp::from_str(string, timezone) .map_err(|e| DataFusionError::External(Box::new(e)))?; let value = Some(ts.value()); - let scalar = match ts.unit() { + Ok(match ts.unit() { TimeUnit::Second => ScalarValue::TimestampSecond(value, None), TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(value, None), TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(value, None), TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(value, None), - }; - Ok(scalar) + }) +} + +fn cast_timestamp(parsed: ScalarValue, target_type: &DataType) -> Result { + let parsed = parsed.to_array()?; + let casted = compute::cast(&parsed, target_type) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + + ScalarValue::try_from_array(&casted, 0) } #[cfg(test)] @@ -337,19 +480,23 @@ mod tests { use super::*; #[test] - fn test_string_to_timestamp_ms() { + fn test_cast_string_to_timestamp() { + use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; + + let target_type = DataType::Timestamp(ArrowTimeUnit::Second, None); assert_eq!( - string_to_timestamp_ms("2022-02-02 19:00:00+08:00", None).unwrap(), + cast_string_to_timestamp("2022-02-02 19:00:00+08:00", &target_type, None).unwrap(), ScalarValue::TimestampSecond(Some(1643799600), None) ); assert_eq!( - string_to_timestamp_ms("2009-02-13 23:31:30Z", None).unwrap(), + cast_string_to_timestamp("2009-02-13 23:31:30Z", &target_type, None).unwrap(), ScalarValue::TimestampSecond(Some(1234567890), None) ); assert_eq!( - string_to_timestamp_ms( + cast_string_to_timestamp( "2009-02-13 23:31:30", + &target_type, Some(&Timezone::from_tz_string("Asia/Shanghai").unwrap()) ) .unwrap(), @@ -357,13 +504,24 @@ mod tests { ); assert_eq!( - string_to_timestamp_ms( + cast_string_to_timestamp( "2009-02-13 23:31:30", + &target_type, Some(&Timezone::from_tz_string("-8:00").unwrap()) ) .unwrap(), ScalarValue::TimestampSecond(Some(1234567890 + 8 * 3600), None) ); + + assert_eq!( + cast_string_to_timestamp( + "2009-02-13 23:31:30.123456789Z", + &DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), + None, + ) + .unwrap(), + ScalarValue::TimestampNanosecond(Some(1_234_567_890_123_456_789), None) + ); } #[test] @@ -421,14 +579,11 @@ mod tests { ) .unwrap(), ); - let mut converter = TypeConverter { - schema, - query_ctx: QueryContext::arc(), - }; + let mut converter = TypeConverter::new(schema, QueryContext::arc()); assert_eq!( Expr::Column(Column::from_name("ts")).gt(ScalarValue::TimestampSecond( - Some(1599514949), + Some(1_599_514_949), None ) .lit()), @@ -439,6 +594,58 @@ mod tests { ); } + #[test] + fn test_insert_assignment_uses_query_timezone_and_target_precision() { + use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; + + let query_ctx = QueryContext::arc(); + query_ctx.set_timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()); + let mut converter = InsertAssignmentConverter { + query_ctx, + input_literals: None, + }; + + let expr = Expr::Cast(Cast { + expr: Box::new("2009-02-13 23:31:30.123456789".lit()), + data_type: DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), + }); + assert_eq!( + converter.f_up(expr).unwrap().data, + ScalarValue::TimestampNanosecond(Some(1_234_539_090_123_456_789), None).lit() + ); + } + + #[test] + fn test_insert_assignment_falls_back_for_unsupported_literal() { + use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; + + let mut converter = InsertAssignmentConverter { + query_ctx: QueryContext::arc(), + input_literals: None, + }; + for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] { + let expr = Expr::Cast(Cast { + expr: Box::new(literal.lit()), + data_type: DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), + }); + + assert_eq!(converter.f_up(expr.clone()).unwrap().data, expr); + } + } + + #[test] + fn test_type_converter_leaves_explicit_timestamp_cast_unchanged() { + use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; + + let mut converter = TypeConverter::new(Arc::new(DFSchema::empty()), QueryContext::arc()); + let expr = Expr::Cast(Cast { + expr: Box::new("2009-02-13 23:31:30".lit()), + data_type: DataType::Timestamp(ArrowTimeUnit::Millisecond, None), + }); + + assert_eq!(converter.f_up(expr.clone()).unwrap().data, expr); + } + #[test] fn test_convert_bool() { let col_name = "is_valid"; @@ -452,10 +659,7 @@ mod tests { ) .unwrap(), ); - let mut converter = TypeConverter { - schema, - query_ctx: QueryContext::arc(), - }; + let mut converter = TypeConverter::new(schema, QueryContext::arc()); assert_eq!( Expr::Column(Column::from_name(col_name)).eq(true.lit()), diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index cd0366cfac4..c941839c0ad 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -664,9 +664,10 @@ mod tests { use catalog::RegisterTableRequest; use catalog::memory::MemoryCatalogManager; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_time::Timezone; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; - use session::context::QueryContext; + use session::context::{QueryContext, QueryContextBuilder}; use store_api::metric_engine_consts::{ DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, @@ -697,6 +698,36 @@ mod tests { crate::tests::new_query_engine_with_table(table) } + async fn create_timestamp_test_engine() -> QueryEngineRef { + let columns = vec![ + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + "st", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![]) + .value_indices(vec![0, 1]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::new("timestamps", table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + crate::tests::new_query_engine_with_table(table) + } + fn create_promql_test_engine() -> QueryEngineRef { let catalog_manager = MemoryCatalogManager::with_default_setup(); let physical_table_name = "phy"; @@ -829,6 +860,43 @@ mod tests { engine.planner().plan(&stmt, query_ctx).await.unwrap() } + #[tokio::test] + async fn test_insert_timestamp_literals_use_query_timezone() { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) + .build(), + ); + let engine = create_timestamp_test_engine().await; + + for (sql, expected_timestamp) in [ + ( + "INSERT INTO timestamps (ts, st) \ + VALUES ('2026-08-02 12:00:00.001', now())", + 1_785_643_200_001_i64, + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-03 12:00:00.001', now()", + 1_785_729_600_001_i64, + ), + ] { + let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap() + .display_indent() + .to_string(); + + assert!( + plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")), + "{plan}" + ); + } + } + #[tokio::test] async fn test_extract_placeholder_cast_types_multiple() { let plan = parse_sql_to_plan( diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.result b/tests/cases/standalone/common/insert/insert_default_timezone.result index 9b12e4a15d0..da8b1261e7a 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.result +++ b/tests/cases/standalone/common/insert/insert_default_timezone.result @@ -52,6 +52,41 @@ SELECT * FROM test1 ORDER BY j; | 4 | 2025-02-01T00:01:01 | +---+---------------------+ +CREATE TABLE test3 (ts TIMESTAMP TIME INDEX, st TIMESTAMP, ts_ns TIMESTAMP(9)); + +Affected Rows: 0 + +INSERT INTO test3 (ts) VALUES ('2026-08-01 12:00:00.001'); + +Affected Rows: 1 + +INSERT INTO test3 (ts, st) VALUES ('2026-08-02 12:00:00.001', now()); + +Affected Rows: 1 + +INSERT INTO test3 (ts, st) SELECT '2026-08-03 12:00:00.001', now(); + +Affected Rows: 1 + +INSERT INTO test3 (ts, st, ts_ns) VALUES ( + '2026-08-04 12:00:00.001', + now(), + '2026-08-04 12:00:00.123456789' +); + +Affected Rows: 1 + +SELECT ts, ts_ns FROM test3 ORDER BY ts; + ++-------------------------+-------------------------------+ +| ts | ts_ns | ++-------------------------+-------------------------------+ +| 2026-08-01T04:00:00.001 | | +| 2026-08-02T04:00:00.001 | | +| 2026-08-03T04:00:00.001 | | +| 2026-08-04T04:00:00.001 | 2026-08-04T04:00:00.123456789 | ++-------------------------+-------------------------------+ + SET time_zone = 'UTC'; Affected Rows: 0 @@ -64,3 +99,7 @@ DROP TABLE test2; Affected Rows: 0 +DROP TABLE test3; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.sql b/tests/cases/standalone/common/insert/insert_default_timezone.sql index 2473cdcee49..fba8c134fa0 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.sql +++ b/tests/cases/standalone/common/insert/insert_default_timezone.sql @@ -16,8 +16,26 @@ SELECT * FROM test2 ORDER BY j; SELECT * FROM test1 ORDER BY j; +CREATE TABLE test3 (ts TIMESTAMP TIME INDEX, st TIMESTAMP, ts_ns TIMESTAMP(9)); + +INSERT INTO test3 (ts) VALUES ('2026-08-01 12:00:00.001'); + +INSERT INTO test3 (ts, st) VALUES ('2026-08-02 12:00:00.001', now()); + +INSERT INTO test3 (ts, st) SELECT '2026-08-03 12:00:00.001', now(); + +INSERT INTO test3 (ts, st, ts_ns) VALUES ( + '2026-08-04 12:00:00.001', + now(), + '2026-08-04 12:00:00.123456789' +); + +SELECT ts, ts_ns FROM test3 ORDER BY ts; + SET time_zone = 'UTC'; DROP TABLE test1; DROP TABLE test2; + +DROP TABLE test3; From 4dd92c774e3cceca0b3e505f23334f4d55c18444 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 14 Aug 2026 08:24:50 +0000 Subject: [PATCH 072/126] feat: add json_object function and use it in the entity-graph derivation (#8870) * feat: add json_object scalar function Builds a JSONB object from interleaved (key, value, ...) arguments, like MySQL's JSON_OBJECT. Values are written into the binary directly, so JSON-hostile characters (quotes, backslashes, control characters) need no text-level escaping. Keys must be non-NULL strings; values may be strings, numbers, booleans, or NULL (JSON null). Signed-off-by: Dennis Zhuang * fix: build entity-graph JSON objects with json_object The derivation assembled entity_id_attrs and descriptive by concatenating a JSON text and parsing it, escaping only backslash and double quote in runtime values. A label containing a control character (e.g. a newline) produced unparseable text and failed the whole semantic_entities scan instead of one attribute. json_object assembles the JSONB binary directly from the value columns, so no text escaping is involved; NULL-to-'' stays at the call site. Signed-off-by: Dennis Zhuang * chore: trim comments and fold duplicate test coverage Signed-off-by: Dennis Zhuang * fix: json_object() returns an empty object; narrow values to integers and floats MySQL's JSON_OBJECT allows an empty pair list, so the signature accepts zero arguments and the row count falls back to number_rows. Decimals stay rejected instead of casting to Float64: JSONB numbers (i64/u64/f64) cannot represent them exactly and a silent precision loss is worse than an explicit cast. Signed-off-by: Dennis Zhuang * chore: document key-to-string conversion and align test naming Keys follow MySQL JSON_OBJECT: any castable type is converted to string. Rustdoc and the cast-failure message now say so, with a numeric-key test. Test names take the module-conventional test_ prefix. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/function/src/scalars/json.rs | 2 + .../function/src/scalars/json/json_object.rs | 281 ++++++++++++++++++ src/operator/src/statement/semantic_graph.rs | 66 ++-- .../common/function/json/json.result | 70 +++++ .../standalone/common/function/json/json.sql | 25 ++ 5 files changed, 406 insertions(+), 38 deletions(-) create mode 100644 src/common/function/src/scalars/json/json_object.rs diff --git a/src/common/function/src/scalars/json.rs b/src/common/function/src/scalars/json.rs index 8a06c972c9d..d0e7f412ac2 100644 --- a/src/common/function/src/scalars/json.rs +++ b/src/common/function/src/scalars/json.rs @@ -15,6 +15,7 @@ pub mod json_get; mod json_get_rewriter; mod json_is; +mod json_object; mod json_object_keys; mod json_path_exists; mod json_path_match; @@ -54,6 +55,7 @@ impl JsonFunction { registry.register_scalar(JsonIsArray::default()); registry.register_scalar(JsonIsObject::default()); + registry.register_scalar(json_object::JsonObjectFunction::default()); registry.register_scalar(json_object_keys::JsonObjectKeysFunction::default()); registry.register_scalar(json_path_exists::JsonPathExistsFunction::default()); registry.register_scalar(json_path_match::JsonPathMatchFunction::default()); diff --git a/src/common/function/src/scalars/json/json_object.rs b/src/common/function/src/scalars/json/json_object.rs new file mode 100644 index 00000000000..a9602c3abc5 --- /dev/null +++ b/src/common/function/src/scalars/json/json_object.rs @@ -0,0 +1,281 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt::{self, Display}; +use std::sync::Arc; + +use datafusion_common::DataFusionError; +use datafusion_common::arrow::array::{Array, ArrayRef, AsArray, BinaryViewBuilder}; +use datafusion_common::arrow::compute; +use datafusion_common::arrow::datatypes::{DataType, Float64Type, Int64Type, UInt64Type}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, TypeSignature, Volatility}; + +use crate::function::Function; + +const NAME: &str = "json_object"; + +/// Builds a `JSONB` object from interleaved `(key, value, key, value, ...)` +/// arguments, like MySQL's `JSON_OBJECT`; called with no arguments it returns +/// `{}`. Values are written into the binary directly, so they need no JSON +/// text escaping. Keys must be non-NULL and are converted to strings (so a +/// numeric key like `1` becomes `"1"`, as in MySQL); values may be strings, +/// integers, floats, booleans, or NULL (rendered as JSON null). Other types — +/// including decimals, which JSONB numbers cannot represent exactly — are +/// rejected; cast them explicitly. A duplicate key keeps the last value. +#[derive(Clone, Debug)] +pub(crate) struct JsonObjectFunction { + signature: Signature, +} + +impl Default for JsonObjectFunction { + fn default() -> Self { + Self { + signature: Signature::one_of( + vec![TypeSignature::Nullary, TypeSignature::VariadicAny], + Volatility::Immutable, + ), + } + } +} + +/// A value column normalized to the canonical arrow type its JSON rendering +/// reads from. +enum ValueColumn { + Null, + Bool(ArrayRef), + Int(ArrayRef), + UInt(ArrayRef), + Float(ArrayRef), + String(ArrayRef), +} + +impl ValueColumn { + fn try_new(array: &ArrayRef) -> datafusion_common::Result { + let normalized = match array.data_type() { + DataType::Null => return Ok(ValueColumn::Null), + DataType::Boolean => ValueColumn::Bool(array.clone()), + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => { + ValueColumn::Int(compute::cast(array, &DataType::Int64)?) + } + DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => { + ValueColumn::UInt(compute::cast(array, &DataType::UInt64)?) + } + DataType::Float16 | DataType::Float32 | DataType::Float64 => { + ValueColumn::Float(compute::cast(array, &DataType::Float64)?) + } + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { + ValueColumn::String(compute::cast(array, &DataType::Utf8View)?) + } + other => { + return Err(DataFusionError::Execution(format!( + "{NAME} does not support values of type {other}; cast the value to a string" + ))); + } + }; + Ok(normalized) + } + + fn value(&self, row: usize) -> jsonb::Value<'_> { + let array = match self { + ValueColumn::Null => return jsonb::Value::Null, + ValueColumn::Bool(array) + | ValueColumn::Int(array) + | ValueColumn::UInt(array) + | ValueColumn::Float(array) + | ValueColumn::String(array) => array, + }; + if !array.is_valid(row) { + return jsonb::Value::Null; + } + match self { + ValueColumn::Null => unreachable!(), + ValueColumn::Bool(array) => array.as_boolean().value(row).into(), + ValueColumn::Int(array) => array.as_primitive::().value(row).into(), + ValueColumn::UInt(array) => array.as_primitive::().value(row).into(), + ValueColumn::Float(array) => array.as_primitive::().value(row).into(), + ValueColumn::String(array) => array.as_string_view().value(row).into(), + } + } +} + +impl Function for JsonObjectFunction { + fn name(&self) -> &str { + NAME + } + + fn return_type(&self, _: &[DataType]) -> datafusion_common::Result { + Ok(DataType::BinaryView) + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + if arrays.len() % 2 != 0 { + return Err(DataFusionError::Execution(format!( + "{NAME} expects (key, value) argument pairs, got {} arguments", + arrays.len() + ))); + } + let pairs = arrays + .chunks(2) + .map(|pair| { + let keys = compute::cast(&pair[0], &DataType::Utf8View).map_err(|_| { + DataFusionError::Execution(format!( + "{NAME} cannot convert keys of type {} to string", + pair[0].data_type() + )) + })?; + if keys.null_count() > 0 { + return Err(DataFusionError::Execution(format!( + "{NAME} does not allow NULL keys" + ))); + } + Ok((keys, ValueColumn::try_new(&pair[1])?)) + }) + .collect::>>()?; + + let rows = arrays.first().map_or(args.number_rows, |a| a.len()); + let mut builder = BinaryViewBuilder::with_capacity(rows); + let mut buf = Vec::new(); + for row in 0..rows { + let mut object = jsonb::Object::new(); + for (keys, values) in &pairs { + object.insert( + keys.as_string_view().value(row).to_string(), + values.value(row), + ); + } + buf.clear(); + jsonb::Value::Object(object).write_to_vec(&mut buf); + builder.append_value(&buf); + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + } +} + +impl Display for JsonObjectFunction { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "JSON_OBJECT") + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_schema::Field; + use datafusion_common::arrow::array::{ + Int64Array, NullArray, StringArray, TimestampMillisecondArray, + }; + + use super::*; + + fn invoke(args: Vec, rows: usize) -> datafusion_common::Result> { + let function = JsonObjectFunction::default(); + let result = function.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: rows, + return_field: Arc::new(Field::new("x", DataType::BinaryView, true)), + config_options: Arc::new(Default::default()), + })?; + let array = result.to_array(rows)?; + let array = array.as_binary_view(); + Ok((0..array.len()) + .map(|i| jsonb::from_slice(array.value(i)).unwrap().to_string()) + .collect()) + } + + fn key(name: &str) -> ColumnarValue { + ColumnarValue::Scalar(datafusion_common::ScalarValue::Utf8(Some(name.to_string()))) + } + + #[test] + fn test_builds_objects_from_mixed_types_without_escaping() { + let texts = invoke( + vec![ + key("host"), + ColumnarValue::Array(Arc::new(StringArray::from(vec![ + Some("we\"ird\\\nhost"), + None, + ]))), + key("pid"), + ColumnarValue::Array(Arc::new(Int64Array::from(vec![42, 7]))), + key("up"), + ColumnarValue::Scalar(datafusion_common::ScalarValue::Boolean(Some(true))), + key("v"), + ColumnarValue::Array(Arc::new(NullArray::new(2))), + ], + 2, + ) + .unwrap(); + assert_eq!( + texts, + vec![ + r#"{"host":"we\"ird\\\nhost","pid":42,"up":true,"v":null}"#, + r#"{"host":null,"pid":7,"up":true,"v":null}"#, + ] + ); + } + + #[test] + fn test_empty_call_builds_one_empty_object_per_row() { + assert_eq!(invoke(vec![], 3).unwrap(), vec!["{}", "{}", "{}"]); + } + + #[test] + fn test_numeric_key_converts_to_string() { + let texts = invoke( + vec![ + ColumnarValue::Scalar(datafusion_common::ScalarValue::Int64(Some(7))), + key("v"), + ], + 1, + ) + .unwrap(); + assert_eq!(texts, vec![r#"{"7":"v"}"#]); + } + + #[test] + fn test_rejects_odd_arguments_null_keys_and_unsupported_values() { + let err = invoke(vec![key("a")], 1).unwrap_err(); + assert!(err.to_string().contains("(key, value) argument pairs")); + + let err = invoke( + vec![ + ColumnarValue::Array(Arc::new(StringArray::from(vec![None::<&str>]))), + key("v"), + ], + 1, + ) + .unwrap_err(); + assert!(err.to_string().contains("NULL keys")); + + let err = invoke( + vec![ + key("ts"), + ColumnarValue::Array(Arc::new(TimestampMillisecondArray::from(vec![1_000]))), + ], + 1, + ) + .unwrap_err(); + assert!(err.to_string().contains("does not support values")); + } +} diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index 777c6f677ea..7014a94f8ac 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -465,43 +465,34 @@ fn null_json() -> Expr { lit(ScalarValue::Binary(None)) } -/// Renders a compile-time-known string as JSON text (quoted, fully escaped — -/// including control characters, unlike the runtime value escaping). +/// Renders a compile-time-known string as JSON text (quoted, fully escaped). fn json_quote(value: &str) -> String { serde_json::Value::from(value).to_string() } -/// Wraps a column reference in `replace` calls so its runtime value is -/// JSON-escaped (`\` then `"`); NULL becomes `''` so one NULL column does not -/// invalidate the whole JSON text (descriptive columns are nullable). -fn json_escaped_value_expr(column: &str) -> Expr { - let escaped_backslash = - string_fns::replace().call(vec![cast_string_or_empty(column), lit("\\"), lit("\\\\")]); - string_fns::replace().call(vec![escaped_backslash, lit("\""), lit("\\\"")]) -} +/// The `json_object` UDF, resolved like [`PARSE_JSON_UDF`]. It assembles the +/// JSONB binary directly from the value columns, so runtime values need no +/// JSON text escaping. +static JSON_OBJECT_UDF: LazyLock> = LazyLock::new(|| { + Arc::new( + FUNCTION_REGISTRY + .get_function("json_object") + .expect("json_object must be registered") + .provide(FunctionContext::default()), + ) +}); -/// Builds a JSONB object from `columns` by concatenating a JSON text and parsing -/// it — GreptimeDB has no struct→json function. Keys are JSON-escaped in Rust; -/// values are JSON-escaped at runtime via [`json_escaped_value_expr`]. -/// -/// TODO(entity-graph): replace the text round-trip with a UDF that assembles -/// JSONB directly from the value columns (`jsonb::ObjectBuilder`, keys baked -/// in), dropping the escaping helpers and the per-row parse cost. +/// Builds a JSONB object with one entry per column: key = the column name, +/// value = the column rendered as a string, NULL coalesced to `""` so one NULL +/// column does not null the entry (descriptive columns are nullable). Keys come +/// out sorted — JSONB objects are key-ordered regardless of input order. fn json_object_expr(columns: &[String]) -> Expr { - if columns.is_empty() { - return parse_json_expr(lit("{}")); + let mut args = Vec::with_capacity(columns.len() * 2); + for column in columns { + args.push(lit(column.as_str())); + args.push(cast_string_or_empty(column)); } - let mut parts = vec![lit("{")]; - for (i, column) in columns.iter().enumerate() { - if i > 0 { - parts.push(lit(",")); - } - parts.push(lit(format!("{}:\"", json_quote(column)))); - parts.push(json_escaped_value_expr(column)); - parts.push(lit("\"")); - } - parts.push(lit("}")); - parse_json_expr(concat_expr(parts)) + cast(JSON_OBJECT_UDF.call(args), DataType::Binary) } /// Renders pre-sorted columns as a `k=v,k=v` concatenation. `nullable` @@ -599,9 +590,7 @@ fn registry_source( let entity_id_attrs = if decl.id_columns.len() == 1 { null_json() } else { - let mut cols = decl.id_columns.clone(); - cols.sort(); - json_object_expr(&cols) + json_object_expr(&decl.id_columns) }; let scope = match decl.scope_columns.as_slice() { @@ -799,7 +788,7 @@ mod tests { Arc::new(StringArray::from(vec!["cart", "cart", "cart"])), Arc::new(Int64Array::from(vec![42, 42, 42])), Arc::new(StringArray::from(vec![ - Some(r#"we"ird\host"#), + Some("we\"ird\\\nhost"), None, Some("h2"), ])), @@ -912,7 +901,8 @@ mod tests { rows.sort(); // Composite id -> sorted `k=v,k=v` plus a JSON object of the id columns; - // descriptive JSON escapes `\` and `"` in runtime values, NULL -> "". + // descriptive JSON keeps `\`, `"` and control characters intact in + // runtime values, NULL -> "". assert_eq!( rows, vec![ @@ -929,7 +919,7 @@ mod tests { ( "pid=42,service_name=cart".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), - Some(r#"{"host":"we\"ird\\host"}"#.to_string()), + Some(r#"{"host":"we\"ird\\\nhost"}"#.to_string()), ), ] ); @@ -955,7 +945,7 @@ mod tests { let batches = collect(&ctx, plan).await; let mut scopes: Vec = batches.iter().flat_map(|b| strings(b, 7)).collect(); scopes.sort(); - assert_eq!(scopes, vec!["", "h2", r#"we"ird\host"#]); + assert_eq!(scopes, vec!["", "h2", "we\"ird\\\nhost"]); // Multiple scope columns: sorted `k=v,k=v`. let mut multi = decl("service", &["service_name"]); @@ -978,7 +968,7 @@ mod tests { vec![ "host=,pid=42", "host=h2,pid=42", - r#"host=we"ird\host,pid=42"# + "host=we\"ird\\\nhost,pid=42" ] ); } diff --git a/tests/cases/standalone/common/function/json/json.result b/tests/cases/standalone/common/function/json/json.result index ab32c9c0132..c035db72e4f 100644 --- a/tests/cases/standalone/common/function/json/json.result +++ b/tests/cases/standalone/common/function/json/json.result @@ -145,3 +145,73 @@ SELECT json_path_match(parse_json('null'), '$.a == 1'); | | +------------------------------------------------------------+ +--- json_object --- +SELECT json_to_string(json_object('a', 1, 'b', 'text', 'c', true, 'd', 1.5)); + ++-----------------------------------------------------------------------------------------------------------------------+ +| json_to_string(json_object(Utf8("a"),Int64(1),Utf8("b"),Utf8("text"),Utf8("c"),Boolean(true),Utf8("d"),Float64(1.5))) | ++-----------------------------------------------------------------------------------------------------------------------+ +| {"a":1,"b":"text","c":true,"d":1.5} | ++-----------------------------------------------------------------------------------------------------------------------+ + +SELECT json_to_string(json_object()); + ++-------------------------------+ +| json_to_string(json_object()) | ++-------------------------------+ +| {} | ++-------------------------------+ + +SELECT json_to_string(json_object('nul', NULL)); + ++-----------------------------------------------+ +| json_to_string(json_object(Utf8("nul"),NULL)) | ++-----------------------------------------------+ +| {"nul":null} | ++-----------------------------------------------+ + +SELECT json_to_string(json_object('nl', concat('line1', chr(10), 'line2'), 'q', 'quote"back\slash')); + ++-------------------------------------------------------------------------------------------------------------------------------+ +| json_to_string(json_object(Utf8("nl"),concat(Utf8("line1"),chr(Int64(10)),Utf8("line2")),Utf8("q"),Utf8("quote"back\slash"))) | ++-------------------------------------------------------------------------------------------------------------------------------+ +| {"nl":"line1\nline2","q":"quote\"back\\slash"} | ++-------------------------------------------------------------------------------------------------------------------------------+ + +SELECT json_object('a'); + +Error: 3001(EngineExecuteQuery), Execution error: json_object expects (key, value) argument pairs, got 1 arguments + +SELECT json_object(NULL, 1); + +Error: 3001(EngineExecuteQuery), Execution error: json_object does not allow NULL keys + +SELECT json_object('ts', to_timestamp(0)); + +Error: 3001(EngineExecuteQuery), Execution error: json_object does not support values of type Timestamp(ns); cast the value to a string + +SELECT json_object('price', CAST('12.34' AS DECIMAL(10, 2))); + +Error: 3001(EngineExecuteQuery), Execution error: json_object does not support values of type Decimal128(10, 2); cast the value to a string + +CREATE TABLE json_object_src (ts TIMESTAMP TIME INDEX, host STRING, pid BIGINT); + +Affected Rows: 0 + +INSERT INTO json_object_src VALUES (0, 'h1', 42), (1, NULL, 7); + +Affected Rows: 2 + +SELECT json_to_string(json_object('host', host, 'pid', pid)) FROM json_object_src ORDER BY ts; + ++------------------------------------------------------------------------------------------------+ +| json_to_string(json_object(Utf8("host"),json_object_src.host,Utf8("pid"),json_object_src.pid)) | ++------------------------------------------------------------------------------------------------+ +| {"host":"h1","pid":42} | +| {"host":null,"pid":7} | ++------------------------------------------------------------------------------------------------+ + +DROP TABLE json_object_src; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/json/json.sql b/tests/cases/standalone/common/function/json/json.sql index e09f206df9f..16517b54f69 100644 --- a/tests/cases/standalone/common/function/json/json.sql +++ b/tests/cases/standalone/common/function/json/json.sql @@ -37,3 +37,28 @@ SELECT json_path_match(parse_json('{"a":1,"b":[1,2,3]}'), '$.b[1 to last] >= 2') SELECT json_path_match(parse_json('{"a":1,"b":[1,2,3]}'), 'null'); SELECT json_path_match(parse_json('null'), '$.a == 1'); + +--- json_object --- +SELECT json_to_string(json_object('a', 1, 'b', 'text', 'c', true, 'd', 1.5)); + +SELECT json_to_string(json_object()); + +SELECT json_to_string(json_object('nul', NULL)); + +SELECT json_to_string(json_object('nl', concat('line1', chr(10), 'line2'), 'q', 'quote"back\slash')); + +SELECT json_object('a'); + +SELECT json_object(NULL, 1); + +SELECT json_object('ts', to_timestamp(0)); + +SELECT json_object('price', CAST('12.34' AS DECIMAL(10, 2))); + +CREATE TABLE json_object_src (ts TIMESTAMP TIME INDEX, host STRING, pid BIGINT); + +INSERT INTO json_object_src VALUES (0, 'h1', 42), (1, NULL, 7); + +SELECT json_to_string(json_object('host', host, 'pid', pid)) FROM json_object_src ORDER BY ts; + +DROP TABLE json_object_src; From 2eb5f593f8dde448a81c8a3bfaac9867bc69b8aa Mon Sep 17 00:00:00 2001 From: sun Date: Fri, 14 Aug 2026 08:45:02 +0000 Subject: [PATCH 073/126] feat: update dashboard to v0.13.12 (#8882) --- src/servers/dashboard/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index dbbe182110e..3d431cfd5f0 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.13.11 +v0.13.12 From ed4271af403e6fa8898287ddd735c1c984644af7 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 14 Aug 2026 09:36:01 +0000 Subject: [PATCH 074/126] feat(procedure): record event actor (#8849) * feat(event): record procedure actor Signed-off-by: WenyXu * test: handle streamed region migration output Signed-off-by: WenyXu * test: cover procedure actors across SQL protocols Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/common/event-recorder/src/event_table.rs | 283 ------------------ .../meta/src/ddl/tests/event/database.rs | 4 +- src/common/meta/src/ddl/tests/event/flow.rs | 6 +- src/common/meta/src/ddl/tests/event/view.rs | 4 +- src/common/procedure/src/event.rs | 65 +++- src/meta-client/src/client/procedure.rs | 54 +--- src/meta-srv/src/event/gc.rs | 6 +- src/meta-srv/src/event/repartition.rs | 6 +- src/meta-srv/src/service/procedure.rs | 76 +---- tests-integration/src/test_util.rs | 55 +++- tests-integration/tests/database_ddl_event.rs | 60 +++- .../tests/event_recorder_test_util.rs | 57 ++++ tests-integration/tests/flow_ddl_event.rs | 48 +-- tests-integration/tests/gc_event.rs | 53 ++-- tests-integration/tests/http.rs | 25 ++ tests-integration/tests/region_migration.rs | 80 ++++- tests-integration/tests/repartition_event.rs | 43 +-- tests-integration/tests/sql.rs | 48 ++- tests-integration/tests/table_ddl_event.rs | 30 +- tests-integration/tests/view_ddl_event.rs | 51 ++-- tests-integration/tests/wal_prune_event.rs | 1 + 21 files changed, 504 insertions(+), 551 deletions(-) diff --git a/src/common/event-recorder/src/event_table.rs b/src/common/event-recorder/src/event_table.rs index 1ea4369a00a..25c16f433d1 100644 --- a/src/common/event-recorder/src/event_table.rs +++ b/src/common/event-recorder/src/event_table.rs @@ -310,286 +310,3 @@ pub fn nullable_json(value: Option<&serde_json::Value>) -> Value { pub fn jsonb_value(value: &serde_json::Value) -> Value { ValueData::BinaryValue(jsonb::Value::from(value).to_vec()).into() } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn base_schema_preserves_names_types_semantics_extensions_and_order() { - assert_eq!( - base_column_schemas(), - vec![ - ColumnSchema { - column_name: "type".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Tag.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "payload".to_string(), - datatype: ColumnDataType::Binary.into(), - semantic_type: SemanticType::Field.into(), - datatype_extension: Some(ColumnDataTypeExtension { - type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), - }), - ..Default::default() - }, - ColumnSchema { - column_name: "timestamp".to_string(), - datatype: ColumnDataType::TimestampNanosecond.into(), - semantic_type: SemanticType::Timestamp.into(), - ..Default::default() - }, - ] - ); - } - - #[test] - fn procedure_envelope_schema_preserves_names_types_semantics_and_order() { - assert_eq!( - procedure_event_column_schemas(), - vec![ - ColumnSchema { - column_name: "procedure_id".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "procedure_state".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "procedure_error".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "procedure_trigger".to_string(), - datatype: ColumnDataType::Binary.into(), - semantic_type: SemanticType::Field.into(), - datatype_extension: Some(ColumnDataTypeExtension { - type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), - }), - ..Default::default() - }, - ] - ); - } - - #[test] - fn shared_dimension_schema_preserves_names_types_semantics_and_order() { - assert_eq!( - column_schemas([&CATALOG_NAME_COLUMN, &SCHEMA_NAME_COLUMN]), - ["catalog_name", "schema_name"].map(|column_name| ColumnSchema { - column_name: column_name.to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }) - ); - } - - #[test] - fn flow_dimension_schema_preserves_names_types_semantics_and_order() { - assert_eq!( - column_schemas([&FLOW_NAME_COLUMN, &FLOW_ID_COLUMN]), - [ - ("flow_name", ColumnDataType::String), - ("flow_id", ColumnDataType::Uint32), - ] - .map(|(column_name, datatype)| ColumnSchema { - column_name: column_name.to_string(), - datatype: datatype.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }) - ); - } - - #[test] - fn admin_function_schema_preserves_names_types_semantics_extensions_and_order() { - assert_eq!( - column_schemas([ - &ACTOR_COLUMN, - &ADMIN_FUNCTION_NAME_COLUMN, - &ADMIN_FUNCTION_STATUS_COLUMN, - &ADMIN_FUNCTION_OUTPUT_COLUMN, - ]), - vec![ - ColumnSchema { - column_name: "actor".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "admin_function_name".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "admin_function_status".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "admin_function_output".to_string(), - datatype: ColumnDataType::Binary.into(), - semantic_type: SemanticType::Field.into(), - datatype_extension: Some(ColumnDataTypeExtension { - type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), - }), - ..Default::default() - }, - ] - ); - } - - #[test] - fn table_dimension_schema_preserves_names_types_semantics_and_order() { - assert_eq!( - column_schemas([ - &TABLE_NAME_COLUMN, - &TABLE_ID_COLUMN, - &PHYSICAL_TABLE_ID_COLUMN, - ]), - [ - ("table_name", ColumnDataType::String), - ("table_id", ColumnDataType::Uint32), - ("physical_table_id", ColumnDataType::Uint32), - ] - .map(|(column_name, datatype)| ColumnSchema { - column_name: column_name.to_string(), - datatype: datatype.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }) - ); - } - - #[test] - fn view_dimension_schema_preserves_names_types_semantics_and_order() { - assert_eq!( - column_schemas([&VIEW_NAME_COLUMN, &VIEW_ID_COLUMN]), - [ - ("view_name", ColumnDataType::String), - ("view_id", ColumnDataType::Uint32), - ] - .map(|(column_name, datatype)| ColumnSchema { - column_name: column_name.to_string(), - datatype: datatype.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }) - ); - } - - #[test] - fn region_dimension_schema_preserves_names_types_semantics_and_order() { - assert_field_columns( - [ - ®ION_ID_COLUMN, - ®ION_NUMBER_COLUMN, - ®ION_MIGRATION_TRIGGER_REASON_COLUMN, - ®ION_MIGRATION_SRC_NODE_ID_COLUMN, - ®ION_MIGRATION_SRC_PEER_ADDR_COLUMN, - ®ION_MIGRATION_DST_NODE_ID_COLUMN, - ®ION_MIGRATION_DST_PEER_ADDR_COLUMN, - ], - [ - ("region_id", ColumnDataType::Uint64), - ("region_number", ColumnDataType::Uint32), - ("region_migration_trigger_reason", ColumnDataType::String), - ("region_migration_src_node_id", ColumnDataType::Uint64), - ("region_migration_src_peer_addr", ColumnDataType::String), - ("region_migration_dst_node_id", ColumnDataType::Uint64), - ("region_migration_dst_peer_addr", ColumnDataType::String), - ], - ); - } - - #[test] - fn repartition_dimension_schema_preserves_names_types_semantics_and_order() { - assert_field_columns( - [ - &TABLE_NAME_COLUMN, - &TABLE_ID_COLUMN, - &PARENT_PROCEDURE_ID_COLUMN, - &REPARTITION_GROUP_ID_COLUMN, - &SOURCE_REGION_ID_COLUMN, - &SOURCE_REGION_NUMBER_COLUMN, - &SOURCE_PARTITION_EXPR_COLUMN, - &TARGET_REGION_ID_COLUMN, - &TARGET_REGION_NUMBER_COLUMN, - &TARGET_PARTITION_EXPR_COLUMN, - ], - [ - ("table_name", ColumnDataType::String), - ("table_id", ColumnDataType::Uint32), - ("parent_procedure_id", ColumnDataType::String), - ("repartition_group_id", ColumnDataType::String), - ("source_region_id", ColumnDataType::Uint64), - ("source_region_number", ColumnDataType::Uint32), - ("source_partition_expr", ColumnDataType::String), - ("target_region_id", ColumnDataType::Uint64), - ("target_region_number", ColumnDataType::Uint32), - ("target_partition_expr", ColumnDataType::String), - ], - ); - } - - #[test] - fn wal_prune_schema_preserves_names_types_semantics_and_order() { - assert_field_columns( - [ - &TOPIC_NAME_COLUMN, - &PRUNABLE_ENTRY_ID_COLUMN, - &LATEST_OFFSET_COLUMN, - ], - [ - ("topic_name", ColumnDataType::String), - ("prunable_entry_id", ColumnDataType::Uint64), - ("latest_offset", ColumnDataType::Uint64), - ], - ); - } - - #[test] - fn nullable_values_preserve_types_and_nulls() { - assert_eq!( - nullable_string(Some("catalog")), - Value { - value_data: Some(ValueData::StringValue("catalog".to_string())) - } - ); - assert_eq!(nullable_value(None), Value { value_data: None }); - assert_eq!( - nullable_value(Some(ValueData::BoolValue(true))), - Value { - value_data: Some(ValueData::BoolValue(true)) - } - ); - } - - fn assert_field_columns( - columns: [&EventTableColumn; N], - expected: [(&str, ColumnDataType); N], - ) { - assert_eq!( - column_schemas(columns), - expected.map(|(column_name, datatype)| ColumnSchema { - column_name: column_name.to_string(), - datatype: datatype.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }) - ); - } -} diff --git a/src/common/meta/src/ddl/tests/event/database.rs b/src/common/meta/src/ddl/tests/event/database.rs index 8649982ed15..a617a1f5f6a 100644 --- a/src/common/meta/src/ddl/tests/event/database.rs +++ b/src/common/meta/src/ddl/tests/event/database.rs @@ -20,7 +20,7 @@ use api::v1::value::ValueData; use api::v1::{AlterDatabaseExpr, Row, SetDatabaseOptions as PbSetDatabaseOptions, Value}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, + ACTOR_COLUMN, CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN as EVENT_TABLE_PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN as EVENT_TABLE_PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN as EVENT_TABLE_PROCEDURE_STATE_COLUMN, @@ -303,6 +303,7 @@ fn assert_procedure_event_contract( EVENT_TABLE_PROCEDURE_TRIGGER_COLUMN.column_schema(), EVENT_TABLE_CATALOG_NAME_COLUMN.column_schema(), EVENT_TABLE_SCHEMA_NAME_COLUMN.column_schema(), + ACTOR_COLUMN.column_schema(), EVENT_CONTEXT_COLUMN.column_schema(), ], &[Row { @@ -326,6 +327,7 @@ fn assert_procedure_event_contract( value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, Value { value_data: None }, + Value { value_data: None }, ], }], ); diff --git a/src/common/meta/src/ddl/tests/event/flow.rs b/src/common/meta/src/ddl/tests/event/flow.rs index cec0162adf8..297d6a6d9e2 100644 --- a/src/common/meta/src/ddl/tests/event/flow.rs +++ b/src/common/meta/src/ddl/tests/event/flow.rs @@ -19,8 +19,8 @@ use api::v1::{ColumnSchema, Row, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, jsonb_value, - procedure_event_column_schemas, + ACTOR_COLUMN, CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, + jsonb_value, procedure_event_column_schemas, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; @@ -241,6 +241,7 @@ fn assert_procedure_event_contract( ) { let mut schema = procedure_event_column_schemas(); schema.extend(flow_schema()); + schema.push(ACTOR_COLUMN.column_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); assert_event_contract( event, @@ -260,6 +261,7 @@ fn assert_procedure_event_contract( .map(Into::into) .unwrap_or(Value { value_data: None }), Value { value_data: None }, + Value { value_data: None }, ], }], ); diff --git a/src/common/meta/src/ddl/tests/event/view.rs b/src/common/meta/src/ddl/tests/event/view.rs index 6ebb62cda82..85ecbae4ac8 100644 --- a/src/common/meta/src/ddl/tests/event/view.rs +++ b/src/common/meta/src/ddl/tests/event/view.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use api::v1::value::ValueData; use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, + ACTOR_COLUMN, CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, VIEW_NAME_COLUMN, jsonb_value, procedure_event_column_schemas, }; use common_event_recorder::testing::assert_event_contract; @@ -335,6 +335,7 @@ fn assert_procedure_event_contract( ) { let mut schema = procedure_event_column_schemas(); schema.extend(view_schema()); + schema.push(ACTOR_COLUMN.column_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); let mut values = vec![ @@ -345,6 +346,7 @@ fn assert_procedure_event_contract( ]; values.extend(locator.values()); values.push(Value { value_data: None }); + values.push(Value { value_data: None }); assert_event_contract(event, event_type, &schema, &[Row { values }]); } diff --git a/src/common/procedure/src/event.rs b/src/common/procedure/src/event.rs index a4a6c290bf3..8ca3a6c67b3 100644 --- a/src/common/procedure/src/event.rs +++ b/src/common/procedure/src/event.rs @@ -19,8 +19,9 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, - PROCEDURE_TRIGGER_COLUMN, jsonb_value, nullable_json, procedure_event_column_schemas, + ACTOR_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, + PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, nullable_json, nullable_string, + procedure_event_column_schemas, }; use common_time::timestamp::{TimeUnit, Timestamp}; use snafu::ResultExt; @@ -100,6 +101,7 @@ impl Event for ProcedureEvent { let mut schema = procedure_event_column_schemas(); let mut internal_schema = self.internal_event.extra_schema(); schema.append(&mut internal_schema); + schema.push(ACTOR_COLUMN.column_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); schema } @@ -124,9 +126,10 @@ impl Event for ProcedureEvent { .map(serde_json::to_value) .transpose() .context(SerializeEventSnafu)?; + let actor = self.context.actor.as_deref(); for internal_event_extra_row in internal_event_extra_rows.iter_mut() { - let mut values = Vec::with_capacity(5 + internal_event_extra_row.values.len()); + let mut values = Vec::with_capacity(6 + internal_event_extra_row.values.len()); values.extend([ ValueData::StringValue(procedure_id.clone()).into(), ValueData::StringValue(state.clone()).into(), @@ -134,6 +137,7 @@ impl Event for ProcedureEvent { jsonb_value(&trigger), ]); values.append(&mut internal_event_extra_row.values); + values.push(nullable_string(actor)); values.push(nullable_json(event_context.as_ref())); rows.push(Row { values }); } @@ -155,7 +159,7 @@ mod tests { use common_error::mock::MockError; use common_error::status_code::StatusCode; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, + ACTOR_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, }; use common_event_recorder::{Event, PersistentEventContext, TriggerReason}; use serde_json::json; @@ -223,6 +227,7 @@ mod tests { jsonb_value(&json!({"type": "Submitted"})), ValueData::StringValue("test_event1".to_string()).into(), Value { value_data: None }, + Value { value_data: None }, ], }, Row { @@ -233,6 +238,7 @@ mod tests { jsonb_value(&json!({"type": "Submitted"})), ValueData::StringValue("test_event2".to_string()).into(), Value { value_data: None }, + Value { value_data: None }, ], }, Row { @@ -243,6 +249,7 @@ mod tests { jsonb_value(&json!({"type": "Submitted"})), Value { value_data: None }, Value { value_data: None }, + Value { value_data: None }, ], }, ] @@ -313,16 +320,18 @@ mod tests { semantic_type: SemanticType::Field.into(), ..Default::default() }, + ACTOR_COLUMN.column_schema(), EVENT_CONTEXT_COLUMN.column_schema(), ] ); } #[test] - fn procedure_event_only_records_context_on_submission() { - let context = ProcedureContext::from_event_context(PersistentEventContext::new( + fn procedure_event_records_actor_for_all_lifecycle_triggers() { + let mut context = ProcedureContext::from_event_context(PersistentEventContext::new( TriggerReason::AutoRepartition, )); + context.actor = Some("alice".to_string()); let submitted = ProcedureEvent::new_with_context( ProcedureId::random(), Box::new(TestEvent {}), @@ -330,22 +339,46 @@ mod tests { EventTrigger::Submitted, context.clone(), ); - let succeeded = ProcedureEvent::new_with_context( - ProcedureId::random(), - Box::new(TestEvent {}), - ProcedureState::Done { output: None }, - EventTrigger::Succeeded, - context, - ); assert_eq!( submitted.extra_rows().unwrap()[0].values[5], - jsonb_value(&json!({"reason": "auto_repartition"})) + ValueData::StringValue("alice".to_string()).into() ); assert_eq!( - succeeded.extra_rows().unwrap()[0].values[5], - Value { value_data: None } + submitted.extra_rows().unwrap()[0].values[6], + jsonb_value(&json!({"reason": "auto_repartition"})) ); + + let child_id = ProcedureId::random(); + for trigger in [ + EventTrigger::Recovered, + EventTrigger::ChildSubmitted { + procedure_id: child_id, + outcome: ChildSubmissionOutcome::Accepted, + }, + EventTrigger::Retrying { + phase: RetryPhase::Execute, + attempt: 1, + }, + EventTrigger::RollingBack, + EventTrigger::Succeeded, + EventTrigger::Failed, + EventTrigger::Poisoned, + ] { + let event = ProcedureEvent::new_with_context( + ProcedureId::random(), + Box::new(TestEvent {}), + ProcedureState::Running, + trigger, + context.clone(), + ); + let values = &event.extra_rows().unwrap()[0].values; + assert_eq!( + values[5], + ValueData::StringValue("alice".to_string()).into() + ); + assert_eq!(values[6], Value { value_data: None }); + } } #[test] diff --git a/src/meta-client/src/client/procedure.rs b/src/meta-client/src/client/procedure.rs index ccf2c28188f..9c32b4fc880 100644 --- a/src/meta-client/src/client/procedure.rs +++ b/src/meta-client/src/client/procedure.rs @@ -63,13 +63,9 @@ pub(crate) fn procedure_event_context(context: &ExecutorContext) -> Option Option { - context - .actor - .as_deref() - .filter(|username| !username.is_empty()) - .map(|username| ProcedureActor { - username: username.to_string(), - }) + context.actor.as_ref().map(|username| ProcedureActor { + username: username.clone(), + }) } #[derive(Clone, Debug)] @@ -531,7 +527,7 @@ mod tests { use tonic::{Request, Response, Status}; use crate::client::MetaClientBuilder; - use crate::client::procedure::{gc_timeout_secs, procedure_actor, procedure_event_context}; + use crate::client::procedure::{gc_timeout_secs, procedure_event_context}; #[test] fn test_gc_timeout_secs() { @@ -578,21 +574,6 @@ mod tests { ); } - #[test] - fn test_procedure_actor() { - let context = ExecutorContext { - actor: Some(String::new()), - ..Default::default() - }; - assert!(procedure_actor(&context).is_none()); - - let context = ExecutorContext { - actor: Some("alice".to_string()), - ..Default::default() - }; - assert_eq!(procedure_actor(&context).unwrap().username, "alice"); - } - #[derive(Clone)] struct MockHeartbeat { leader_addr: String, @@ -728,18 +709,18 @@ mod tests { username: "alice".to_string(), created_at_ns: 42, }; - let executor_context = |actor: String| ExecutorContext { + let executor_context = ExecutorContext { query_context: Some(QueryContext { channel: Channel::Postgres as u8, ..Default::default() }), - actor: Some(actor), + actor: Some("effective-user".to_string()), event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), ..Default::default() }; ProcedureExecutor::submit_ddl_task( &client, - executor_context("effective-user".to_string()), + executor_context, SubmitDdlTaskRequest::new(DdlTask::new_create_database( "greptime".to_string(), "metrics".to_string(), @@ -772,27 +753,6 @@ mod tests { let extensions = &request.query_context.unwrap().extensions; assert_eq!(extensions[CREATE_DATABASE_CREATOR_EXTENSION_KEY], encoded); - ProcedureExecutor::submit_ddl_task( - &client, - executor_context(String::new()), - SubmitDdlTaskRequest::new(DdlTask::new_drop_database( - "greptime".to_string(), - "metrics".to_string(), - false, - )), - ) - .await - .unwrap(); - assert!( - request_rx - .recv() - .await - .unwrap() - .into_inner() - .actor - .is_none() - ); - server_handle.abort(); } diff --git a/src/meta-srv/src/event/gc.rs b/src/meta-srv/src/event/gc.rs index 3575ac20c2f..7f679f1785b 100644 --- a/src/meta-srv/src/event/gc.rs +++ b/src/meta-srv/src/event/gc.rs @@ -212,8 +212,8 @@ mod tests { use api::v1::ColumnSchema; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, - PROCEDURE_TRIGGER_COLUMN, jsonb_value, + ACTOR_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, + PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; use common_event_recorder::{EventTypeFilter, PersistentEventContext, TriggerReason}; @@ -536,6 +536,7 @@ mod tests { ); let mut event_schema = procedure_schema(); event_schema.extend(schema()); + event_schema.push(ACTOR_COLUMN.column_schema()); event_schema.push(EVENT_CONTEXT_COLUMN.column_schema()); let mut values = vec![ ValueData::StringValue(procedure_id.to_string()).into(), @@ -554,6 +555,7 @@ mod tests { .values, ); values.push(Value { value_data: None }); + values.push(Value { value_data: None }); assert_event_contract( &event, diff --git a/src/meta-srv/src/event/repartition.rs b/src/meta-srv/src/event/repartition.rs index 1faf3f18212..c1e18fd565e 100644 --- a/src/meta-srv/src/event/repartition.rs +++ b/src/meta-srv/src/event/repartition.rs @@ -364,8 +364,8 @@ mod tests { use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, - PROCEDURE_TRIGGER_COLUMN, jsonb_value, + ACTOR_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, + PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; @@ -668,6 +668,7 @@ mod tests { ); let mut schema = procedure_schema(); schema.extend(parent_schema()); + schema.push(ACTOR_COLUMN.column_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); assert_event_contract( @@ -685,6 +686,7 @@ mod tests { ValueData::StringValue("repartition_events".to_string()).into(), ValueData::U32Value(1024).into(), Value { value_data: None }, + Value { value_data: None }, ], }], ); diff --git a/src/meta-srv/src/service/procedure.rs b/src/meta-srv/src/service/procedure.rs index 9e226bc18fe..df27e2cb32b 100644 --- a/src/meta-srv/src/service/procedure.rs +++ b/src/meta-srv/src/service/procedure.rs @@ -55,28 +55,14 @@ struct ProcedureSubmission { event_context: Option, } -impl TryFrom<(Option, Option)> for ProcedureSubmission { - type Error = Status; - - fn try_from( +impl From<(Option, Option)> for ProcedureSubmission { + fn from( (actor, event_context): (Option, Option), - ) -> Result { - let actor = actor - .map(|actor| { - if actor.username.is_empty() { - Err(Status::invalid_argument( - "Procedure actor must not be empty", - )) - } else { - Ok(actor.username) - } - }) - .transpose()?; - - Ok(Self { - actor, + ) -> Self { + Self { + actor: actor.map(|actor| actor.username), event_context, - }) + } } } @@ -139,7 +125,7 @@ impl procedure_service_server::ProcedureService for Metasrv { let ProcedureSubmission { actor, event_context, - } = ProcedureSubmission::try_from((actor, event_context))?; + } = ProcedureSubmission::from((actor, event_context)); let mut query_context = query_context .context(error::MissingRequiredParameterSnafu { param: "query_context", @@ -191,7 +177,7 @@ impl procedure_service_server::ProcedureService for Metasrv { let _header = header.context(error::MissingRequestHeaderSnafu)?; let procedure_context = - ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); + ProcedureContext::from(ProcedureSubmission::from((actor, event_context))); let from_peer = self .lookup_datanode_peer(from_peer) .await? @@ -329,7 +315,7 @@ impl procedure_service_server::ProcedureService for Metasrv { let _header = header.context(error::MissingRequestHeaderSnafu)?; let procedure_context = - ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); + ProcedureContext::from(ProcedureSubmission::from((actor, event_context))); let response = self .handle_gc_regions( @@ -361,7 +347,7 @@ impl procedure_service_server::ProcedureService for Metasrv { let _header = header.context(error::MissingRequestHeaderSnafu)?; let procedure_context = - ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); + ProcedureContext::from(ProcedureSubmission::from((actor, event_context))); let response = self .handle_gc_table( @@ -593,15 +579,14 @@ fn gc_response_to_table_pb(resp: GcResponse) -> GcTableResponse { mod tests { use std::time::Duration; - use api::v1::meta::{ProcedureActor, ProcedureEventContext, Role}; + use api::v1::meta::Role; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, CreatorGrantIntent, DdlTask, QueryContext, }; - use common_procedure::ProcedureContext; use tonic::metadata::{MetadataMap, MetadataValue}; - use super::{Metasrv, ProcedureSubmission, restore_create_database_creator}; + use super::{Metasrv, restore_create_database_creator}; fn create_database_task() -> DdlTask { DdlTask::new_create_database( @@ -738,41 +723,4 @@ mod tests { assert_eq!(task.creator, Some(creator)); assert!(trusted.extensions.is_empty()); } - - #[test] - fn test_procedure_submission_rejects_empty_actor() { - let submission = ProcedureSubmission::try_from(( - Some(ProcedureActor { - username: "alice".to_string(), - }), - Some(ProcedureEventContext { - reason: "manual".to_string(), - protocol: "postgres".to_string(), - extensions: Default::default(), - }), - )) - .unwrap(); - assert_eq!(submission.actor.as_deref(), Some("alice")); - assert_eq!( - submission - .event_context - .as_ref() - .map(|context| context.protocol.as_str()), - Some("postgres") - ); - - assert!( - ProcedureSubmission::try_from(( - Some(ProcedureActor { - username: String::new(), - }), - None, - )) - .is_err() - ); - - let procedure_context = - ProcedureContext::from(ProcedureSubmission::try_from((None, None)).unwrap()); - assert_eq!(procedure_context.event_context, None); - } } diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index d124c5c6cdb..30b9a07348d 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -18,9 +18,13 @@ use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; -use auth::{DefaultPermissionChecker, PermissionCheckerRef, UserProviderRef}; +use api::v1::Basic; +use auth::{ + DefaultPermissionChecker, PermissionCheckerRef, UserProviderRef, user_provider_from_option, +}; use axum::Router; use catalog::kvbackend::KvBackendCatalogManager; +use client::{Client, Database}; use common_base::Plugins; use common_config::Configurable; use common_meta::key::TableMetadataManager; @@ -732,14 +736,57 @@ async fn setup_grpc_server_for_instance( grpc_config: Option, memory_limiter: Option, ) -> (GreptimeDbStandalone, Arc) { + let grpc_server = setup_grpc_server_for_frontend_instance_with( + instance.fe_instance().clone(), + user_provider, + grpc_config, + memory_limiter, + ) + .await; + (instance, grpc_server) +} + +/// Builds and starts an authenticated gRPC server on an existing frontend instance. +pub async fn setup_grpc_server_for_frontend_instance( + instance: Arc, + user_provider: Option, +) -> Arc { + setup_grpc_server_for_frontend_instance_with(instance, user_provider, None, None).await +} + +/// Starts a gRPC server backed by `instance` and returns a client authenticated +/// as the supplied test user. The server handle keeps the listener alive. +pub async fn setup_authenticated_grpc_database( + instance: Arc, + username: &str, + password: &str, +) -> (Database, Arc) { + let user_provider = + user_provider_from_option(&format!("static_user_provider:cmd:{username}={password}")) + .unwrap(); + let grpc_server = setup_grpc_server_for_frontend_instance(instance, Some(user_provider)).await; + let grpc_addr = grpc_server.bind_addr().unwrap().to_string(); + let mut database = + Database::new_with_dbname("greptime-public", Client::with_urls(vec![grpc_addr])); + database.set_auth(api::v1::auth_header::AuthScheme::Basic(Basic { + username: username.to_string(), + password: password.to_string(), + })); + (database, grpc_server) +} + +async fn setup_grpc_server_for_frontend_instance_with( + fe_instance_ref: Arc, + user_provider: Option, + grpc_config: Option, + memory_limiter: Option, +) -> Arc { let runtime: Runtime = RuntimeBuilder::default() .worker_threads(2) .thread_name("grpc-handlers") .build() .unwrap(); - let fe_instance_ref = instance.fe_instance().clone(); - let greptime_request_handler = GreptimeRequestHandler::new( fe_instance_ref.clone(), user_provider.clone(), @@ -769,7 +816,7 @@ async fn setup_grpc_server_for_instance( let fe_grpc_addr = "127.0.0.1:0".parse::().unwrap(); grpc_server.start(fe_grpc_addr).await.unwrap(); - (instance, Arc::new(grpc_server)) + Arc::new(grpc_server) } pub async fn setup_mysql_server( diff --git a/tests-integration/tests/database_ddl_event.rs b/tests-integration/tests/database_ddl_event.rs index 6cbc65a2673..4dfda4f32b9 100644 --- a/tests-integration/tests/database_ddl_event.rs +++ b/tests-integration/tests/database_ddl_event.rs @@ -14,17 +14,21 @@ use std::sync::Arc; -use client::OutputData; +use client::{Database, OutputData}; use common_test_util::temp_dir::create_temp_dir; -use servers::query_handler::sql::SqlQueryHandler; -use session::context::QueryContext; use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; -use tests_integration::test_util::{StorageType, get_test_store_config}; +use tests_integration::test_util::{ + StorageType, get_test_store_config, setup_authenticated_grpc_database, +}; -use crate::event_recorder_test_util::{assert_single_event, find_eventually_string}; +use crate::event_recorder_test_util::{ + assert_procedure_actor, assert_single_event, find_eventually_string, +}; const DATABASE_NAME: &str = "database_ddl_events"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[tokio::test(flavor = "multi_thread")] async fn test_database_ddl_events() { @@ -44,8 +48,14 @@ async fn test_database_ddl_events() { .build(true) .await; let instance = cluster.fe_instance(); + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; - execute_database_ddl(instance).await; + execute_database_ddl(&database).await; assert_database_events(instance).await; } @@ -57,23 +67,25 @@ async fn test_standalone_database_ddl_events() { .build() .await; let instance = standalone.fe_instance(); + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; - execute_database_ddl(instance).await; + execute_database_ddl(&database).await; assert_database_events(instance).await; } -async fn execute_database_ddl(instance: &Arc) { +async fn execute_database_ddl(database: &Database) { for sql in [ format!("CREATE DATABASE IF NOT EXISTS {DATABASE_NAME} WITH (ttl = '1h')"), format!("ALTER DATABASE {DATABASE_NAME} SET 'ttl' = '2h'"), format!("DROP DATABASE IF EXISTS {DATABASE_NAME}"), ] { - let output = instance - .do_query(&sql, QueryContext::arc()) - .await - .remove(0) - .unwrap(); + let output = database.sql(sql).await.unwrap(); assert!(matches!(output.data, OutputData::AffectedRows(_))); } } @@ -111,10 +123,27 @@ async fn assert_database_event( event_type: &str, submitted_payload_predicate: &str, ) { + let procedure_id = find_eventually_string( + instance, + &format!( + r#"SELECT procedure_id +FROM greptime_private.events +WHERE type = '{event_type}' + AND schema_name = '{DATABASE_NAME}' + AND json_path_match(procedure_trigger, '$.type == "Submitted"') +ORDER BY timestamp DESC +LIMIT 1"# + ), + "procedure_id", + ) + .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; + let submitted = format!( r#"SELECT count(*) AS event_count FROM greptime_private.events WHERE type = '{event_type}' + AND procedure_id = '{procedure_id}' AND procedure_state = 'Running' AND json_path_match(procedure_trigger, '$.type == "Submitted"') AND catalog_name = 'greptime' @@ -125,17 +154,18 @@ WHERE type = '{event_type}' let actual = find_eventually_string( instance, &format!( - "SELECT json_to_string(event_context) AS event_context FROM greptime_private.events WHERE type = '{event_type}' AND procedure_state = 'Running' AND json_path_match(procedure_trigger, '$.type == \"Submitted\"') ORDER BY timestamp DESC LIMIT 1" + "SELECT json_to_string(event_context) AS event_context FROM greptime_private.events WHERE procedure_id = '{procedure_id}' AND json_path_match(procedure_trigger, '$.type == \"Submitted\"')" ), "event_context", ) .await; - assert_eq!(r#"{"reason":"manual"}"#, actual); + assert_eq!(r#"{"protocol":"grpc","reason":"manual"}"#, actual); let lifecycle = format!( r#"SELECT count(*) AS event_count FROM greptime_private.events WHERE type = '{event_type}' + AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') AND catalog_name = 'greptime' diff --git a/tests-integration/tests/event_recorder_test_util.rs b/tests-integration/tests/event_recorder_test_util.rs index 0441cc08d3b..38bb18861b3 100644 --- a/tests-integration/tests/event_recorder_test_util.rs +++ b/tests-integration/tests/event_recorder_test_util.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::future::Future; use std::sync::Arc; use std::time::Duration; @@ -26,6 +27,62 @@ use session::context::QueryContext; const MAX_ATTEMPTS: usize = 60; const POLL_INTERVAL: Duration = Duration::from_millis(250); +/// Asserts through the supplied SQL executor that a table procedure records +/// the expected actor for both submission and completion events. +pub(crate) async fn assert_procedure_actor_by_table( + event_type: &str, + table_name: &str, + actor: &str, + mut execute: F, +) where + F: FnMut(String) -> Fut, + Fut: Future, +{ + let query = format!( + "SELECT count(DISTINCT json_get_string(procedure_trigger, 'type')) = 2 \ + AND count(*) = count(CASE WHEN actor = '{actor}' THEN 1 END) AS actor_matches \ + FROM greptime_private.events \ + WHERE type = '{event_type}' AND table_name = '{table_name}' \ + AND json_get_string(procedure_trigger, 'type') IN ('Submitted', 'Succeeded')" + ); + + for _ in 0..MAX_ATTEMPTS { + if execute(query.clone()).await { + return; + } + tokio::time::sleep(POLL_INTERVAL).await; + } + + panic!("timed out waiting for procedure actor: {query}"); +} + +/// Asserts that submission and completion events record the expected actor. +pub(crate) async fn assert_procedure_actor( + instance: &Arc, + procedure_id: &str, + actor: Option<&str>, +) { + let actor_predicate = actor.map_or_else( + || "count(actor) = 0".to_string(), + |actor| format!("count(*) = count(CASE WHEN actor = '{actor}' THEN 1 END)"), + ); + for trigger in ["Submitted", "Succeeded"] { + assert_eventually_eq( + instance, + &format!( + "SELECT count(*) > 0 AND {actor_predicate} AS actor_matches FROM greptime_private.events WHERE procedure_id = '{procedure_id}' AND json_get_string(procedure_trigger, 'type') = '{trigger}'" + ), + "\ ++---------------+ +| actor_matches | ++---------------+ +| true | ++---------------+", + ) + .await; + } +} + /// Asserts that the query eventually returns exactly one event. pub(crate) async fn assert_single_event(instance: &Arc, query: &str) { for _ in 0..MAX_ATTEMPTS { diff --git a/tests-integration/tests/flow_ddl_event.rs b/tests-integration/tests/flow_ddl_event.rs index c5fc9be94c9..122c36c84f4 100644 --- a/tests-integration/tests/flow_ddl_event.rs +++ b/tests-integration/tests/flow_ddl_event.rs @@ -14,15 +14,18 @@ use std::sync::Arc; -use servers::query_handler::sql::SqlQueryHandler; -use session::context::QueryContext; use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::setup_authenticated_grpc_database; use uuid::Uuid; -use crate::event_recorder_test_util::{assert_single_event, find_eventually_string}; +use crate::event_recorder_test_util::{ + assert_procedure_actor, assert_single_event, find_eventually_string, +}; const CREATE_FLOW_EVENT_TYPE: &str = "create_flow"; const DROP_FLOW_EVENT_TYPE: &str = "drop_flow"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flow_ddl_events() { @@ -35,39 +38,39 @@ async fn test_standalone_flow_ddl_events() { let missing_source = format!("flow_ddl_event_missing_source_{suffix}"); let sink = format!("flow_ddl_event_sink_{suffix}"); let flow = format!("flow_ddl_event_{suffix}"); + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; - instance - .do_query( - &format!("CREATE TABLE {sink} (val STRING, ts TIMESTAMP TIME INDEX)"), - QueryContext::arc(), - ) + database + .sql(format!( + "CREATE TABLE {sink} (val STRING, ts TIMESTAMP TIME INDEX)" + )) .await - .remove(0) .unwrap(); - instance - .do_query( - &format!( - "CREATE FLOW IF NOT EXISTS {flow} SINK TO {sink} EVAL INTERVAL '10s' \ - WITH (defer_on_missing_source = true) AS SELECT val, ts FROM {missing_source}" - ), - QueryContext::arc(), - ) + database + .sql(format!( + "CREATE FLOW IF NOT EXISTS {flow} SINK TO {sink} EVAL INTERVAL '10s' \ + WITH (defer_on_missing_source = true) AS SELECT val, ts FROM {missing_source}" + )) .await - .remove(0) .unwrap(); assert_create_events(instance, &flow).await; - instance - .do_query(&format!("DROP FLOW IF EXISTS {flow}"), QueryContext::arc()) + database + .sql(format!("DROP FLOW IF EXISTS {flow}")) .await - .remove(0) .unwrap(); assert_drop_events(instance, &flow).await; } async fn assert_create_events(instance: &Arc, flow: &str) { let procedure_id = find_submitted_procedure_id(instance, CREATE_FLOW_EVENT_TYPE, flow).await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; assert_single_event( instance, &format!( @@ -120,11 +123,12 @@ async fn assert_event_context( "event_context", ) .await; - assert_eq!(r#"{"reason":"manual"}"#, actual); + assert_eq!(r#"{"protocol":"grpc","reason":"manual"}"#, actual); } async fn assert_drop_events(instance: &Arc, flow: &str) { let procedure_id = find_submitted_procedure_id(instance, DROP_FLOW_EVENT_TYPE, flow).await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; assert_single_event( instance, &format!( diff --git a/tests-integration/tests/gc_event.rs b/tests-integration/tests/gc_event.rs index 54a0787d12a..4721a442a80 100644 --- a/tests-integration/tests/gc_event.rs +++ b/tests-integration/tests/gc_event.rs @@ -15,19 +15,21 @@ use std::sync::Arc; use std::time::Duration; -use common_event_recorder::{PersistentEventContext, TriggerReason}; -use common_procedure::{ProcedureContext, ProcedureId, ProcedureWithId, watcher}; use common_test_util::temp_dir::create_temp_dir; -use meta_srv::gc::{BatchGcProcedure, GcSchedulerOptions}; +use meta_srv::gc::GcSchedulerOptions; use mito2::gc::GcConfig; use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; use tests_integration::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; -use tests_integration::test_util::{StorageType, get_test_store_config}; +use tests_integration::test_util::{ + StorageType, get_test_store_config, setup_authenticated_grpc_database, +}; -use crate::event_recorder_test_util::find_eventually_string; +use crate::event_recorder_test_util::{assert_procedure_actor, find_eventually_string}; const TABLE_NAME: &str = "batch_gc_events"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; const MAX_ATTEMPTS: usize = 60; const POLL_INTERVAL: Duration = Duration::from_millis(250); @@ -121,31 +123,28 @@ async fn test_batch_gc_event() { .collect::>(); assert_eq!(regions.len(), 1); - let procedure = BatchGcProcedure::new( - cluster.metasrv.mailbox().clone(), - cluster.metasrv.table_metadata_manager().clone(), - cluster.metasrv.options().grpc.server_addr.clone(), - regions, - false, - Duration::from_secs(10), - Default::default(), - ); - let procedure_id = ProcedureId::parse_str("00000000-0000-0000-0000-00000000bac0").unwrap(); - let mut watcher = cluster - .metasrv - .procedure_manager() - .submit(ProcedureWithId { - id: procedure_id, - procedure: Box::new(procedure), - context: ProcedureContext::from_event_context(PersistentEventContext::new( - TriggerReason::Manual, - )), - }) + let (actor_db, _actor_grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; + actor_db + .sql(format!("ADMIN GC_REGIONS({})", regions[0].as_u64())) .await .unwrap(); - watcher::wait(&mut watcher).await.unwrap(); assert_sst_count(&cluster, 1).await; + let procedure_id = find_eventually_string( + instance, + &format!( + "SELECT procedure_id FROM greptime_private.events WHERE type = 'batch_gc' AND actor = '{PROCEDURE_ACTOR}' AND json_get_string(procedure_trigger, 'type') = 'Submitted'" + ), + "procedure_id", + ) + .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; + let submitted = format!( r#"SELECT json_to_string(event_context) AS event_context FROM greptime_private.events @@ -155,7 +154,7 @@ WHERE type = 'batch_gc' AND json_get_string(procedure_trigger, 'type') = 'Submitted'"#, ); assert_eq!( - r#"{"reason":"manual"}"#, + r#"{"protocol":"grpc","reason":"manual"}"#, find_eventually_string(instance, &submitted, "event_context").await ); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 530b581a2ca..079439bf78c 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -84,6 +84,8 @@ use tests_integration::test_util::{ use urlencoding::encode; use yaml_rust::YamlLoader; +use crate::event_recorder_test_util::assert_procedure_actor_by_table; + #[macro_export] macro_rules! http_test { ($service:ident, $($(#[$meta:meta])* $test:ident),*,) => { @@ -258,6 +260,29 @@ pub async fn test_http_auth(store_type: StorageType) { .send() .await; assert_eq!(res.status(), StatusCode::OK); + let actor_query_client = &client; + assert_procedure_actor_by_table( + "create_table", + "auth_test", + "writeonly_user", + |query| async move { + let res = actor_query_client + .get(&format!("/v1/sql?db=public&sql={}", encode(&query))) + .header("Authorization", basic_auth("greptime_user", "greptime_pwd")) + .send() + .await; + let Ok(body) = serde_json::from_str::(&res.text().await) else { + return false; + }; + matches!( + body.output().first(), + Some(GreptimeQueryOutput::Records(records)) + if records.rows().first().and_then(|row| row.first()).and_then(Value::as_bool) + == Some(true) + ) + }, + ) + .await; let res = client .get("/v1/sql?db=public&sql=insert into auth_test values(1);") .header( diff --git a/tests-integration/tests/region_migration.rs b/tests-integration/tests/region_migration.rs index d67180249c9..17a58de2eed 100644 --- a/tests-integration/tests/region_migration.rs +++ b/tests-integration/tests/region_migration.rs @@ -15,11 +15,11 @@ use std::sync::Arc; use std::time::Duration; -use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, OutputData}; +use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputData}; use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME; use common_event_recorder::event_table::{ - REGION_ID_COLUMN, REGION_MIGRATION_DST_NODE_ID_COLUMN, REGION_MIGRATION_SRC_NODE_ID_COLUMN, - REGION_MIGRATION_TRIGGER_REASON_COLUMN, + ACTOR_COLUMN, REGION_ID_COLUMN, REGION_MIGRATION_DST_NODE_ID_COLUMN, + REGION_MIGRATION_SRC_NODE_ID_COLUMN, REGION_MIGRATION_TRIGGER_REASON_COLUMN, }; use common_event_recorder::{ DEFAULT_EVENTS_TABLE_NAME, DEFAULT_FLUSH_INTERVAL_SECONDS, EVENTS_TABLE_TIMESTAMP_COLUMN_NAME, @@ -58,10 +58,14 @@ use session::context::{QueryContext, QueryContextRef}; use store_api::storage::RegionId; use table::metadata::TableId; use tests_integration::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; -use tests_integration::test_util::{PEER_PLACEHOLDER_ADDR, StorageType, get_test_store_config}; +use tests_integration::test_util::{ + PEER_PLACEHOLDER_ADDR, StorageType, get_test_store_config, setup_authenticated_grpc_database, +}; use uuid::Uuid; const TEST_TABLE_NAME: &str = "migration_target"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[macro_export] macro_rules! region_migration_test { @@ -263,6 +267,7 @@ pub async fn test_region_migration(store_type: StorageType, endpoints: Vec String { + let output = database + .sql(format!( + "admin migrate_region({region_id}, {from_peer_id}, {to_peer_id})" + )) + .await + .unwrap(); + let recordbatches = match output.data { + OutputData::RecordBatches(recordbatches) => recordbatches, + OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(), + OutputData::AffectedRows(_) => unreachable!(), + }; + + let record_batch = &recordbatches.take()[0]; + record_batch + .column(0) + .as_string::() + .value(0) + .to_string() +} + /// Query procedure state by SQL. async fn query_procedure_by_sql(instance: &Arc, pid: &str) -> String { let OutputData::RecordBatches(recordbatches) = run_sql( @@ -1331,6 +1373,7 @@ async fn run_sql( } enum RegionMigrationEvents { + Actor, ProcedureId, Timestamp, ProcedureState, @@ -1349,6 +1392,7 @@ impl Iden for RegionMigrationEvents { s, "{}", match self { + Self::Actor => ACTOR_COLUMN.name(), Self::ProcedureId => EVENTS_TABLE_PROCEDURE_ID_COLUMN_NAME, Self::Timestamp => EVENTS_TABLE_TIMESTAMP_COLUMN_NAME, Self::ProcedureState => EVENTS_TABLE_PROCEDURE_STATE_COLUMN_NAME, @@ -1371,6 +1415,7 @@ async fn check_region_migration_events_system_table( region_id: u64, from_peer_id: u64, to_peer_id: u64, + actor: Option<&str>, ) { // Sleep enough time to ensure the event is recorded. tokio::time::sleep(DEFAULT_FLUSH_INTERVAL_SECONDS * 2).await; @@ -1386,7 +1431,8 @@ async fn check_region_migration_events_system_table( // region_migration_src_node_id = ${from_peer_id} AND // region_migration_dst_node_id = ${to_peer_id} // ORDER BY timestamp ASC - let query = Query::select() + let mut query = Query::select(); + query .column(RegionMigrationEvents::RegionMigrationTriggerReason) .column(RegionMigrationEvents::ProcedureState) .expr_as( @@ -1402,7 +1448,13 @@ async fn check_region_migration_events_system_table( .and_where(Expr::col(RegionMigrationEvents::ProcedureId).eq(procedure_id)) .and_where(Expr::col(RegionMigrationEvents::RegionId).eq(region_id)) .and_where(Expr::col(RegionMigrationEvents::SrcNodeId).eq(from_peer_id)) - .and_where(Expr::col(RegionMigrationEvents::DstNodeId).eq(to_peer_id)) + .and_where(Expr::col(RegionMigrationEvents::DstNodeId).eq(to_peer_id)); + if let Some(actor) = actor { + query.and_where(Expr::col(RegionMigrationEvents::Actor).eq(actor)); + } else { + query.and_where(Expr::col(RegionMigrationEvents::Actor).is_null()); + } + let query = query .order_by(RegionMigrationEvents::Timestamp, Order::Asc) .to_string(PostgresQueryBuilder); @@ -1411,12 +1463,22 @@ async fn check_region_migration_events_system_table( .await .remove(0); - let expected = "\ + let expected = if actor == Some(PROCEDURE_ACTOR) { + "\ ++---------------------------------+-----------------+-------------------+---------------------------------------+ +| region_migration_trigger_reason | procedure_state | procedure_trigger | event_context | ++---------------------------------+-----------------+-------------------+---------------------------------------+ +| Manual | Running | Submitted | {\"protocol\":\"grpc\",\"reason\":\"manual\"} | +| Manual | Done | Succeeded | | ++---------------------------------+-----------------+-------------------+---------------------------------------+" + } else { + "\ +---------------------------------+-----------------+-------------------+---------------------+ | region_migration_trigger_reason | procedure_state | procedure_trigger | event_context | +---------------------------------+-----------------+-------------------+---------------------+ | Manual | Running | Submitted | {\"reason\":\"manual\"} | | Manual | Done | Succeeded | | -+---------------------------------+-----------------+-------------------+---------------------+"; ++---------------------------------+-----------------+-------------------+---------------------+" + }; check_output_stream(result.unwrap().data, expected).await; } diff --git a/tests-integration/tests/repartition_event.rs b/tests-integration/tests/repartition_event.rs index 65503a4a8ee..906eb3cfee3 100644 --- a/tests-integration/tests/repartition_event.rs +++ b/tests-integration/tests/repartition_event.rs @@ -15,17 +15,22 @@ use std::sync::Arc; use std::time::Duration; +use client::Database; use common_test_util::temp_dir::create_temp_dir; use meta_srv::gc::GcSchedulerOptions; use mito2::gc::GcConfig; -use servers::query_handler::sql::SqlQueryHandler; -use session::context::QueryContext; use tests_integration::cluster::GreptimeDbClusterBuilder; -use tests_integration::test_util::{StorageType, get_test_store_config}; +use tests_integration::test_util::{ + StorageType, get_test_store_config, setup_authenticated_grpc_database, +}; -use crate::event_recorder_test_util::{assert_eventually_eq, find_eventually_string}; +use crate::event_recorder_test_util::{ + assert_eventually_eq, assert_procedure_actor, find_eventually_string, +}; const TABLE_NAME: &str = "repartition_events"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[tokio::test(flavor = "multi_thread")] async fn test_repartition_event() { @@ -56,11 +61,17 @@ async fn test_repartition_event() { .build(true) .await; let instance = cluster.fe_instance(); + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; - execute_partition(instance).await; + execute_partition(&database).await; let partition_procedure_id = assert_repartition_event(instance).await; assert_repartition_group_event(instance, &partition_procedure_id).await; - execute_merge(instance).await; + execute_merge(&database).await; let merge_procedure_id = find_submitted_procedure_id( instance, "repartition", @@ -70,10 +81,11 @@ async fn test_repartition_event() { ), ) .await; + assert_procedure_actor(instance, &merge_procedure_id, Some(PROCEDURE_ACTOR)).await; assert_merge_repartition_group_event(instance, &merge_procedure_id).await; } -async fn execute_partition(instance: &Arc) { +async fn execute_partition(database: &Database) { for sql in [ format!( "CREATE TABLE {TABLE_NAME} (host INT, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host)) ENGINE = mito" @@ -83,21 +95,13 @@ async fn execute_partition(instance: &Arc) { ), format!("ALTER TABLE {TABLE_NAME} PARTITION ON COLUMNS (host) (host < 10, host >= 10)"), ] { - instance - .do_query(&sql, QueryContext::arc()) - .await - .remove(0) - .unwrap(); + database.sql(sql).await.unwrap(); } } -async fn execute_merge(instance: &Arc) { +async fn execute_merge(database: &Database) { let sql = format!("ALTER TABLE {TABLE_NAME} MERGE PARTITION (host < 10, host >= 10)"); - instance - .do_query(&sql, QueryContext::arc()) - .await - .remove(0) - .unwrap(); + database.sql(sql).await.unwrap(); } async fn assert_repartition_event(instance: &Arc) -> String { @@ -110,6 +114,7 @@ async fn assert_repartition_event(instance: &Arc) ), ) .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; let query = format!( r#"SELECT type, catalog_name, schema_name, table_name, table_id, json_path_match(payload, '$.version == 2') AS payload_version, @@ -168,6 +173,7 @@ async fn assert_repartition_group_event( ), ) .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; let query = format!( r#"SELECT type, catalog_name, schema_name, table_name, table_id, parent_procedure_id = '{parent_procedure_id}' AS matches_parent_procedure_id, @@ -235,6 +241,7 @@ async fn assert_merge_repartition_group_event( ), ) .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; let query = format!( r#"SELECT type, catalog_name, schema_name, table_name, table_id, parent_procedure_id = '{parent_procedure_id}' AS matches_parent_procedure_id, diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 58df5424d39..e46a8f58d1b 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -35,6 +35,8 @@ use tests_integration::test_util::{ }; use tokio_postgres::{Client, NoTls, SimpleQueryMessage}; +use crate::event_recorder_test_util::assert_procedure_actor_by_table; + #[macro_export] macro_rules! sql_test { ($service:ident, $($(#[$meta:meta])* $test:ident),*,) => { @@ -148,6 +150,7 @@ pub async fn test_mysql_auth(store_type: StorageType) { .await; assert!(conn_re.is_ok()); + let read_pool = conn_re.unwrap(); // 4. readonly user let conn_re = MySqlPoolOptions::new() @@ -181,6 +184,19 @@ pub async fn test_mysql_auth(store_type: StorageType) { .execute("CREATE TABLE test (ts timestamp time index)") .await .unwrap(); + let actor_query_pool = &read_pool; + assert_procedure_actor_by_table( + "create_table", + "test", + "writeonly_user", + |query| async move { + sqlx::query_scalar::<_, bool>(&query) + .fetch_one(actor_query_pool) + .await + .unwrap_or(false) + }, + ) + .await; let err = pool.execute("SHOW TABLES").await.unwrap_err(); assert!( err.to_string() @@ -466,8 +482,10 @@ pub async fn test_mysql_timezone(store_type: StorageType) { } pub async fn test_postgres_auth(store_type: StorageType) { - let user_provider = - user_provider_from_option("static_user_provider:cmd:greptime_user=greptime_pwd").unwrap(); + let user_provider = user_provider_from_option( + "static_user_provider:cmd:greptime_user=greptime_pwd,writeonly_user:wo=writeonly_pwd", + ) + .unwrap(); let (mut guard, fe_pg_server) = setup_pg_server_with_user_provider(store_type, "sql_crud", Some(user_provider)).await; @@ -518,6 +536,32 @@ pub async fn test_postgres_auth(store_type: StorageType) { .await; assert!(conn_re.is_ok()); + let read_pool = conn_re.unwrap(); + + let write_pool = PgPoolOptions::new() + .max_connections(2) + .connect(&format!( + "postgres://writeonly_user:writeonly_pwd@{addr}/public" + )) + .await + .unwrap(); + write_pool + .execute("CREATE TABLE postgres_auth_actor (ts timestamp time index)") + .await + .unwrap(); + let actor_query_pool = &read_pool; + assert_procedure_actor_by_table( + "create_table", + "postgres_auth_actor", + "writeonly_user", + |query| async move { + sqlx::query_scalar::<_, bool>(&query) + .fetch_one(actor_query_pool) + .await + .unwrap_or(false) + }, + ) + .await; let _ = fe_pg_server.shutdown().await; guard.remove_all().await; diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index afd41b844ac..8bad66dae3c 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -41,9 +41,11 @@ use servers::query_handler::grpc::GrpcQueryHandler; use servers::query_handler::sql::SqlQueryHandler; use session::context::{Channel, QueryContext}; use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::test_util::setup_authenticated_grpc_database; use crate::event_recorder_test_util::{ - assert_eventually_eq, assert_single_event, find_eventually_string, find_eventually_u32, + assert_eventually_eq, assert_procedure_actor, assert_single_event, find_eventually_string, + find_eventually_u32, }; const EVENTS_TABLE: &str = "greptime_private.events"; @@ -53,6 +55,8 @@ const LOGICAL_TABLE: &str = "table_ddl_events_logical"; const AUTO_TABLE: &str = "table_ddl_events_auto"; const AUTO_INFLUX_TABLE: &str = "table_ddl_events_auto_influx"; const GRPC_TABLE: &str = "table_ddl_events_grpc"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[tokio::test(flavor = "multi_thread")] async fn test_table_ddl_procedure_events() { @@ -150,6 +154,7 @@ async fn test_table_ddl_procedure_events() { .unwrap(); let grpc_create_procedure_id = submitted_procedure_id(&frontend, "create_table", GRPC_TABLE).await; + assert_procedure_actor(&frontend, &grpc_create_procedure_id, Some("greptime")).await; assert_event_context( &frontend, "create_table", @@ -200,26 +205,27 @@ async fn test_table_ddl_procedure_events() { // Act / Assert: Create Table retains its rich submitted row and records the // created table ID when it completes. - run_sql_with_context( - &frontend, - &format!( - "CREATE TABLE {TABLE} (host STRING PRIMARY KEY, ts TIMESTAMP TIME INDEX, val DOUBLE)" - ), - Arc::new(QueryContext::with_channel( - "greptime", - "public", - Channel::HttpSql, - )), + let (actor_db, _actor_grpc_server) = setup_authenticated_grpc_database( + frontend.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, ) .await; + actor_db + .sql(format!( + "CREATE TABLE {TABLE} (host STRING PRIMARY KEY, ts TIMESTAMP TIME INDEX, val DOUBLE)" + )) + .await + .unwrap(); let table_id = find_table_id(&frontend, TABLE).await; let create_table_procedure_id = submitted_procedure_id(&frontend, "create_table", TABLE).await; + assert_procedure_actor(&frontend, &create_table_procedure_id, Some(PROCEDURE_ACTOR)).await; assert_event_context( &frontend, "create_table", &create_table_procedure_id, "manual", - Some("httpsql"), + Some("grpc"), ) .await; assert_named_submitted_event( diff --git a/tests-integration/tests/view_ddl_event.rs b/tests-integration/tests/view_ddl_event.rs index fce33d09bc8..b21838995c3 100644 --- a/tests-integration/tests/view_ddl_event.rs +++ b/tests-integration/tests/view_ddl_event.rs @@ -16,17 +16,21 @@ use std::sync::Arc; use common_test_util::temp_dir::create_temp_dir; use common_wal::config::DatanodeWalConfig; -use servers::query_handler::sql::SqlQueryHandler; -use session::context::QueryContext; use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; -use tests_integration::test_util::{StorageType, get_test_store_config}; +use tests_integration::test_util::{ + StorageType, get_test_store_config, setup_authenticated_grpc_database, +}; use uuid::Uuid; -use crate::event_recorder_test_util::{assert_single_event, find_eventually_string}; +use crate::event_recorder_test_util::{ + assert_procedure_actor, assert_single_event, find_eventually_string, +}; const CREATE_VIEW_EVENT_TYPE: &str = "create_view"; const DROP_VIEW_EVENT_TYPE: &str = "drop_view"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[tokio::test(flavor = "multi_thread")] async fn test_view_ddl_events() { @@ -72,37 +76,35 @@ async fn execute_view_ddl( source_table: &str, view: &str, ) { - instance - .do_query( - &format!( - "CREATE TABLE {source_table} (host STRING PRIMARY KEY, amount DOUBLE, ts TIMESTAMP TIME INDEX)" - ), - QueryContext::arc(), - ) + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; + + database + .sql(format!( + "CREATE TABLE {source_table} (host STRING PRIMARY KEY, amount DOUBLE, ts TIMESTAMP TIME INDEX)" + )) .await - .remove(0) .unwrap(); - instance - .do_query( - &format!("CREATE VIEW {view} AS SELECT amount FROM {source_table}"), - QueryContext::arc(), - ) + database + .sql(format!( + "CREATE VIEW {view} AS SELECT amount FROM {source_table}" + )) .await - .remove(0) .unwrap(); assert_create_events(instance, view).await; - instance - .do_query(&format!("DROP VIEW {view}"), QueryContext::arc()) - .await - .remove(0) - .unwrap(); + database.sql(format!("DROP VIEW {view}")).await.unwrap(); assert_drop_events(instance, view).await; } async fn assert_create_events(instance: &Arc, view: &str) { let procedure_id = find_submitted_procedure_id(instance, CREATE_VIEW_EVENT_TYPE, view).await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; assert_single_event( instance, &format!( @@ -157,11 +159,12 @@ async fn assert_event_context( "event_context", ) .await; - assert_eq!(r#"{"reason":"manual"}"#, actual); + assert_eq!(r#"{"protocol":"grpc","reason":"manual"}"#, actual); } async fn assert_drop_events(instance: &Arc, view: &str) { let procedure_id = find_submitted_procedure_id(instance, DROP_VIEW_EVENT_TYPE, view).await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; assert_single_event( instance, &format!( diff --git a/tests-integration/tests/wal_prune_event.rs b/tests-integration/tests/wal_prune_event.rs index 05ceb1b92d9..db86ade4d9d 100644 --- a/tests-integration/tests/wal_prune_event.rs +++ b/tests-integration/tests/wal_prune_event.rs @@ -100,6 +100,7 @@ WHERE type = 'wal_prune' AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') + AND actor IS NULL AND topic_name = '{topic_name}' AND prunable_entry_id = {pruned_entry_id} AND latest_offset = 3 From 072810159a6c19b153bcd59f9abbbc3fe1f39462 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Fri, 14 Aug 2026 09:40:50 +0000 Subject: [PATCH 075/126] chore: add v2 version label to prom metrics (#8885) Signed-off-by: shuiyisong --- src/servers/src/http/prom_store.rs | 28 ++++++++++++------------ src/servers/src/metrics.rs | 23 ++++++++----------- src/servers/src/prom_remote_write/mod.rs | 9 +++++++- src/servers/src/prom_remote_write/v2.rs | 10 ++++++--- src/servers/src/prom_store.rs | 5 ++++- 5 files changed, 42 insertions(+), 33 deletions(-) diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index ed960264b0d..1f0d26789eb 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -47,9 +47,11 @@ use crate::http::header::{ }; use crate::pending_rows_batcher::PendingRowsBatcher; use crate::prom_remote_write::decode::PromSeriesProcessor; -use crate::prom_remote_write::decode_remote_write_request; use crate::prom_remote_write::v2::decode_remote_write_v2; use crate::prom_remote_write::validation::PromValidationMode; +use crate::prom_remote_write::{ + REMOTE_WRITE_V1_VERSION, REMOTE_WRITE_V2_VERSION, decode_remote_write_request, +}; use crate::prom_store::snappy_decompress; use crate::query_handler::{PipelineHandlerRef, PromStoreProtocolHandlerRef, PromStoreResponse}; @@ -57,8 +59,6 @@ pub const PHYSICAL_TABLE_PARAM: &str = "physical_table"; pub const DEFAULT_ENCODING: &str = "snappy"; pub const VM_ENCODING: &str = "zstd"; pub const VM_PROTO_VERSION: &str = "1"; -const REMOTE_WRITE_V1_VERSION: &str = "1.0"; -const REMOTE_WRITE_V2_VERSION: &str = "2.0"; const REMOTE_WRITE_V1_PROTO: &str = "prometheus.WriteRequest"; const REMOTE_WRITE_V2_PROTO: &str = "io.prometheus.write.v2.Request"; const CONTENT_TYPE_PROTO_PARAM: &str = "proto"; @@ -194,11 +194,11 @@ async fn remote_write_v1( { Ok(outcome) => outcome, Err(error) => { - record_remote_write_samples(&db, error.rows_written); + record_remote_write_samples(&db, REMOTE_WRITE_V1_VERSION, error.rows_written); return Err(error.error); } }; - record_remote_write_samples(&db, outcome.rows_written); + record_remote_write_samples(&db, REMOTE_WRITE_V1_VERSION, outcome.rows_written); Ok(( StatusCode::NO_CONTENT, @@ -264,8 +264,8 @@ async fn remote_write_v2( { Ok(outcome) => outcome, Err(error) => { - record_remote_write_samples(&db, error.samples_written); - record_remote_write_histograms(&db, error.histograms_written); + record_remote_write_samples(&db, REMOTE_WRITE_V2_VERSION, error.samples_written); + record_remote_write_histograms(&db, REMOTE_WRITE_V2_VERSION, error.histograms_written); return Ok(remote_write_v2_error_response( error.error, error.samples_written, @@ -276,8 +276,8 @@ async fn remote_write_v2( }; debug_assert_eq!(outcome.samples_written, sample_count); debug_assert_eq!(outcome.histograms_written, histogram_count); - record_remote_write_samples(&db, outcome.samples_written); - record_remote_write_histograms(&db, outcome.histograms_written); + record_remote_write_samples(&db, REMOTE_WRITE_V2_VERSION, outcome.samples_written); + record_remote_write_histograms(&db, REMOTE_WRITE_V2_VERSION, outcome.histograms_written); let mut headers = write_cost_header_map(outcome.write_cost); append_remote_write_v2_written_headers( @@ -322,7 +322,7 @@ fn prepare_remote_write_context( query_ctx.set_extension(SEMANTIC_SOURCE_VERSION, remote_write_version); query_ctx.set_extension(SEMANTIC_METRIC_METADATA_QUALITY, METADATA_QUALITY_INFERRED); let timer = crate::metrics::METRIC_HTTP_PROM_STORE_WRITE_ELAPSED - .with_label_values(&[db.as_str()]) + .with_label_values(&[db.as_str(), remote_write_version]) .start_timer(); (db, query_ctx, timer) @@ -608,21 +608,21 @@ fn incomplete_prom_write_error() -> error::Error { .build() } -fn record_remote_write_samples(db: &str, rows: u64) { +fn record_remote_write_samples(db: &str, version: &str, rows: u64) { if rows == 0 { return; } crate::metrics::PROM_STORE_REMOTE_WRITE_SAMPLES - .with_label_values(&[db]) + .with_label_values(&[db, version]) .inc_by(rows); } -fn record_remote_write_histograms(db: &str, rows: u64) { +fn record_remote_write_histograms(db: &str, version: &str, rows: u64) { if rows == 0 { return; } crate::metrics::PROM_STORE_REMOTE_WRITE_HISTOGRAMS - .with_label_values(&[db]) + .with_label_values(&[db, version]) .inc_by(rows); } diff --git a/src/servers/src/metrics.rs b/src/servers/src/metrics.rs index 41a16bfea23..0d70540be31 100644 --- a/src/servers/src/metrics.rs +++ b/src/servers/src/metrics.rs @@ -44,6 +44,7 @@ pub(crate) const METRIC_POSTGRES_EXTENDED_QUERY: &str = "extended"; pub(crate) const METRIC_METHOD_LABEL: &str = "method"; pub(crate) const METRIC_PATH_LABEL: &str = "path"; pub(crate) const METRIC_RESULT_LABEL: &str = "result"; +pub(crate) const METRIC_VERSION_LABEL: &str = "version"; pub(crate) const METRIC_SUCCESS_VALUE: &str = "success"; pub(crate) const METRIC_FAILURE_VALUE: &str = "failure"; @@ -100,39 +101,33 @@ lazy_static! { vec![0.005, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0, 10.0, 60.0, 300.0] ) .unwrap(); - /// Http prometheus write duration per database. + /// Http prometheus write duration per database and remote write protocol version. pub static ref METRIC_HTTP_PROM_STORE_WRITE_ELAPSED: HistogramVec = register_histogram_vec!( "greptime_servers_http_prometheus_write_elapsed", "servers http prometheus write elapsed", - &[METRIC_DB_LABEL], + &[METRIC_DB_LABEL, METRIC_VERSION_LABEL], vec![0.005, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0, 10.0, 60.0, 300.0] ) .unwrap(); - /// Prometheus remote write codec duration. + /// Prometheus remote write codec duration per protocol version. pub static ref METRIC_HTTP_PROM_STORE_CODEC_ELAPSED: HistogramVec = register_histogram_vec!( "greptime_servers_http_prometheus_codec_elapsed", "servers http prometheus request codec duration", - &["type"], + &["type", METRIC_VERSION_LABEL], ) .unwrap(); - /// Decode duration of prometheus write request. - pub static ref METRIC_HTTP_PROM_STORE_DECODE_ELAPSED: Histogram = METRIC_HTTP_PROM_STORE_CODEC_ELAPSED - .with_label_values(&["decode"]); - /// Duration to convert prometheus write request to gRPC request. - pub static ref METRIC_HTTP_PROM_STORE_CONVERT_ELAPSED: Histogram = METRIC_HTTP_PROM_STORE_CODEC_ELAPSED - .with_label_values(&["convert"]); - /// The samples count of Prometheus remote write. + /// The samples count of Prometheus remote write per protocol version. pub static ref PROM_STORE_REMOTE_WRITE_SAMPLES: IntCounterVec = register_int_counter_vec!( "greptime_servers_prometheus_remote_write_samples", "frontend prometheus remote write samples", - &[METRIC_DB_LABEL] + &[METRIC_DB_LABEL, METRIC_VERSION_LABEL] ) .unwrap(); - /// The native histograms count of Prometheus remote write. + /// The native histograms count of Prometheus remote write per protocol version. pub static ref PROM_STORE_REMOTE_WRITE_HISTOGRAMS: IntCounterVec = register_int_counter_vec!( "greptime_servers_prometheus_remote_write_histograms", "frontend prometheus remote write native histograms", - &[METRIC_DB_LABEL] + &[METRIC_DB_LABEL, METRIC_VERSION_LABEL] ) .unwrap(); pub static ref PENDING_BATCHES: IntGauge = register_int_gauge!( diff --git a/src/servers/src/prom_remote_write/mod.rs b/src/servers/src/prom_remote_write/mod.rs index 1c6b972e3db..73f968b6084 100644 --- a/src/servers/src/prom_remote_write/mod.rs +++ b/src/servers/src/prom_remote_write/mod.rs @@ -37,6 +37,11 @@ use crate::prom_remote_write::row_builder::TablesBuilder; use crate::prom_remote_write::validation::PromValidationMode; use crate::prom_store::{snappy_decompress, zstd_decompress}; +/// Prometheus remote write protocol versions, also used as the `version` label +/// of the remote write metrics. +pub const REMOTE_WRITE_V1_VERSION: &str = "1.0"; +pub const REMOTE_WRITE_V2_VERSION: &str = "2.0"; + lazy_static! { static ref PROM_WRITE_REQUEST_POOL: Pool> = Pool::new(256, PromWriteRequest::default); @@ -56,7 +61,9 @@ pub fn decode_remote_write_request( prom_validation_mode: PromValidationMode, processor: &mut PromSeriesProcessor, ) -> crate::error::Result> { - let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_DECODE_ELAPSED.start_timer(); + let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED + .with_label_values(&["decode", REMOTE_WRITE_V1_VERSION]) + .start_timer(); // due to vmagent's limitation, there is a chance that vmagent is // sending content type wrong so we have to apply a fallback with decoding diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 31da7ac1f35..1bcec98d5f5 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -43,8 +43,8 @@ use table::requests::{ use crate::error::{self, Result}; use crate::prom_remote_write::row_builder::PromCtx; -use crate::prom_remote_write::try_decompress; use crate::prom_remote_write::validation::validate_label_name; +use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, try_decompress}; #[allow(deprecated)] use crate::prom_store::{ DATABASE_LABEL, DATABASE_LABEL_ALT, METRIC_NAME_LABEL, PHYSICAL_TABLE_LABEL, @@ -128,7 +128,9 @@ pub(crate) fn decode_remote_write_v2( body: Bytes, native_histograms_enabled: bool, ) -> Result { - let decode_timer = crate::metrics::METRIC_HTTP_PROM_STORE_DECODE_ELAPSED.start_timer(); + let decode_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED + .with_label_values(&["decode", REMOTE_WRITE_V2_VERSION]) + .start_timer(); // Match the v1 decoder's VictoriaMetrics fallback: some clients may send a // mismatched content-encoding header, so try the other compression on failure. @@ -140,7 +142,9 @@ pub(crate) fn decode_remote_write_v2( let request = BorrowedRequest::decode(&buf).context(error::DecodePromRemoteRequestSnafu)?; drop(decode_timer); - let _convert_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CONVERT_ELAPSED.start_timer(); + let _convert_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED + .with_label_values(&["convert", REMOTE_WRITE_V2_VERSION]) + .start_timer(); convert_remote_write_v2(request, native_histograms_enabled) } diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 33707b29358..c4f2f929cae 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -39,6 +39,7 @@ use snafu::{OptionExt, ResultExt, ensure}; use snap::raw::{Decoder, Encoder}; use crate::error::{self, Result}; +use crate::prom_remote_write::REMOTE_WRITE_V1_VERSION; use crate::row_writer::{self, MultiTableData}; pub const METRIC_NAME_LABEL: &str = "__name__"; @@ -458,7 +459,9 @@ fn recordbatch_to_timeseries( } pub fn to_grpc_row_insert_requests(request: &WriteRequest) -> Result<(RowInsertRequests, usize)> { - let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CONVERT_ELAPSED.start_timer(); + let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED + .with_label_values(&["convert", REMOTE_WRITE_V1_VERSION]) + .start_timer(); let mut multi_table_data = MultiTableData::new(); From 0ef521d28736644b15e0eaaff80379bac49d56c7 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 14 Aug 2026 10:17:43 +0000 Subject: [PATCH 076/126] docs: refresh README product copy (#8886) * docs: refresh README product copy Align the README with the current product messaging on the website: - Drop the Observability 2.0 / wide events framing and the unqualified cost multiples from the headline and feature table. - Remove the claim that the open-source distribution scales reads with horizontal replicas; read replicas are Enterprise. - Replace the competitor comparison table with a per-protocol compatibility matrix, since query-side coverage is narrower than ingestion (no LogQL, partial QueryDSL in Enterprise only). - Add a section stating what is in the Apache-2.0 build and what Enterprise adds. - Fix the canary badge caption, add gRPC to the ingest list. Signed-off-by: Dennis Zhuang * docs: link the Enterprise overview instead of a full feature list Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- README.md | 85 +++++++++++++++++++++++++++++++------------------------ 1 file changed, 48 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index 21980b1a22b..e95f7849e4b 100644 --- a/README.md +++ b/README.md @@ -6,10 +6,11 @@

-

One database for metrics, logs, and traces
-replacing Prometheus, Loki, and Elasticsearch

+

Metrics, logs, and traces.
+One engine, on your infrastructure.

-> The unified OpenTelemetry backend — with SQL + PromQL on object storage. +> A columnar database for metrics, logs, and traces on object storage. +> Apache-2.0 licensed core.

@@ -28,7 +29,7 @@ replacing Prometheus, Loki, and Elasticsearch

Nightly -stable for production  ·  latest includes pre-releases  ·  nightly is a weekly snapshot of main +stable for production  ·  canary includes pre-releases  ·  nightly is a weekly snapshot of main Docker Pulls @@ -57,9 +58,11 @@ replacing Prometheus, Loki, and Elasticsearch
- [Introduction](#introduction) +- [Why You Might Use It](#why-you-might-use-it) - [Overview](#overview) -- [Features](#features) -- [How GreptimeDB Compares](#how-greptimedb-compares) +- [What's Supported](#whats-supported) +- [Compatibility and Migration](#compatibility-and-migration) +- [Limitations and Edition Boundary](#limitations-and-edition-boundary) - [Architecture](#architecture) - [Try GreptimeDB](#try-greptimedb) - [Getting Started](#getting-started) @@ -74,9 +77,20 @@ replacing Prometheus, Loki, and Elasticsearch ## Introduction -**GreptimeDB** is an open-source observability database built for [Observability 2.0](https://docs.greptime.com/user-guide/concepts/observability-2/) — treating metrics, logs, and traces as one unified data model (wide events) instead of three separate pillars. +**GreptimeDB** is an open-source observability database. Metrics, logs, and traces run on one columnar engine over object storage and share one [table model](https://docs.greptime.com/user-guide/concepts/data-model/): tags, timestamp, and fields. When signals carry common identifiers such as service, host, or trace ID, you can correlate them in SQL without moving data between databases. -Use it as the single OpenTelemetry backend — replacing Prometheus, Loki, and Elasticsearch with one database built on object storage. Query with SQL and PromQL, scale without pain, cut costs up to 50×. +Ingest through OpenTelemetry, Prometheus Remote Write, Loki Push, or Elasticsearch Bulk. Use SQL across observability data and PromQL for metrics. Migrate ingestion one signal at a time without rebuilding your collectors. + +## Why You Might Use It + +- You run Prometheus plus Loki or Elasticsearch and want one backend instead of three +- You have outgrown Prometheus on cardinality or retention and don't want the Thanos/Mimir operational surface +- You need long retention on object storage without a separate analytics stack +- You want to query telemetry with SQL, not only a domain query language +- You are storing GenAI or agent telemetry ([OTel GenAI conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/)) alongside infrastructure signals +- You need the same engine and semantics on resource-constrained devices + +Learn more in [Why GreptimeDB](https://docs.greptime.com/user-guide/concepts/why-greptimedb). ## Overview @@ -88,42 +102,38 @@ A quick overview of what GreptimeDB ingests, how it connects to other systems, a

-## Features +## What's Supported -| Feature | Description | -|---------|-------------| -| **Observability 2.0 native** | Logs, metrics, and traces in one engine with [SQL + PromQL](https://docs.greptime.com/user-guide/query-data/overview/). Native [OpenTelemetry](https://docs.greptime.com/user-guide/ingest-data/for-observability/opentelemetry/), [Prometheus remote write](https://docs.greptime.com/user-guide/ingest-data/for-observability/prometheus/), and [Jaeger](https://docs.greptime.com/user-guide/query-data/jaeger/). Migrate one signal at a time, or use as a single backend. | -| **Elastic compute-storage separation** | Scale reads independently with horizontal replicas. Serve high-concurrency workloads from dashboards, alerting, and AI agents — without resharding or data migration. | -| **Sub-second on PB–EB-scale data** | Columnar engine with [fulltext, inverted, and skipping indexes](https://docs.greptime.com/user-guide/manage-data/data-index). Written in Rust. Designed for high-concurrency point queries, not just analytical scans. | -| **50× lower cost** | Object storage (S3, GCS, Azure Blob) as [primary storage](https://docs.greptime.com/user-guide/deployments-administration/configuration/#storage-options), with a tiered cache (memory + local disk) to keep writes and queries fast. | +| | | +|---|---| +| **Ingest** | [OpenTelemetry (OTLP)](https://docs.greptime.com/user-guide/ingest-data/for-observability/opentelemetry/), [Prometheus Remote Write](https://docs.greptime.com/user-guide/ingest-data/for-observability/prometheus/), [Loki Push](https://docs.greptime.com/user-guide/ingest-data/for-observability/loki), [Elasticsearch Bulk](https://docs.greptime.com/user-guide/ingest-data/for-observability/elasticsearch), [InfluxDB line protocol](https://docs.greptime.com/user-guide/protocols/influxdb-line-protocol), [gRPC](https://docs.greptime.com/user-guide/protocols/grpc/) | +| **Query** | [SQL](https://docs.greptime.com/user-guide/query-data/overview/), [PromQL](https://docs.greptime.com/user-guide/query-data/promql/), [Jaeger-compatible trace queries](https://docs.greptime.com/user-guide/query-data/jaeger/), MySQL and PostgreSQL wire protocols | +| **Storage** | S3, GCS, Azure Blob and S3-compatible endpoints as [primary storage](https://docs.greptime.com/user-guide/deployments-administration/configuration/#storage-options), with memory and local-disk caches | +| **Built in** | Retention policies, downsampling, [continuous aggregation](https://docs.greptime.com/user-guide/flow-computation/overview), explicit table partitioning, and [inverted / skipping / fulltext indexes](https://docs.greptime.com/user-guide/manage-data/data-index) | -**Perfect for:** - * Replacing Prometheus + Loki + Elasticsearch with a single observability backend - * Scaling past Prometheus — high cardinality, long-term storage, no Thanos/Mimir overhead - * AI/agent workloads — store GenAI telemetry ([OTel GenAI conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/)), and serve high-concurrency reads from SRE/developer agents via horizontal read replicas - * Cutting observability costs with object storage (up to 50× savings on traces, 30% on logs) - * Edge-to-cloud observability with unified APIs on resource-constrained devices +Compute and storage are disaggregated: object storage holds the data, while memory and local-disk caches keep recent and frequently queried data close to compute. -> **Why Observability 2.0?** Three separate databases for metrics, logs, and traces means three storage layers, three query languages, and three sets of dashboards. GreptimeDB stores all three as timestamped wide events in one columnar engine — JOIN across signals in SQL, run one stack instead of three, and ingest AI agent telemetry the same way. Read more: [Observability 2.0 and the Database for It](https://greptime.com/blogs/2025-04-25-greptimedb-observability2-new-database). +## Compatibility and Migration -Learn more in [Why GreptimeDB](https://docs.greptime.com/user-guide/concepts/why-greptimedb). +Compatibility is per protocol, and query-side coverage is narrower than ingestion. -## How GreptimeDB Compares - -| Capability | GreptimeDB | Prometheus / Thanos / Mimir | Grafana Loki | Elasticsearch | -|---|---|---|---|---| -| Data types | Metrics, logs, traces | Metrics only | Logs only | Logs, traces | -| Query language | SQL + PromQL | PromQL | LogQL | Query DSL | -| Storage | Native object storage (S3, etc.) | Local disk + object storage (Thanos/Mimir) | Object storage (chunks) | Local disk | -| Scaling | Compute-storage separation, stateless nodes | Federation / Thanos / Mimir — multi-component, ops heavy | Stateless + object storage | Shard-based, ops heavy | -| Cost efficiency | Up to 50× lower storage cost | High at scale | Moderate | High (inverted index overhead) | -| OpenTelemetry | Native (metrics + logs + traces) | Partial (metrics only) | Partial (logs only) | Via instrumentation | +| | Compatible | Not compatible | +|---|---|---| +| **Prometheus** | Remote Write ingestion; PromQL queries | Gaps are listed in [PromQL compatibility](https://docs.greptime.com/user-guide/query-data/promql/) | +| **Loki** | Push ingestion; dual-write through Grafana Alloy makes the [cutover gradual](https://greptime.com/blogs/2026-07-23-from-loki-to-greptimedb-dual-write-migration) | LogQL and the rest of the Loki query API | +| **Elasticsearch** | `_bulk` ingestion in the open-source core; [QueryDSL](https://docs.greptime.com/enterprise/elasticsearch-compatible/query/) partially, in Enterprise | Most other Elasticsearch APIs | **Benchmarks:** * [GreptimeDB tops JSONBench's billion-record cold run test](https://greptime.com/blogs/2025-03-18-jsonbench-greptimedb-performance) * [TSBS Benchmark](https://github.com/GreptimeTeam/greptimedb/tree/main/docs/benchmarks/tsbs) * [More benchmark reports](https://docs.greptime.com/user-guide/concepts/features-that-you-concern#how-is-greptimedbs-performance-compared-to-other-solutions) +## Limitations and Edition Boundary + +Cluster deployment, object storage, the Flow engine, and every ingestion protocol listed above are in the Apache-2.0 build. Repartitioning, region migration, and index creation are manual operations there. + +Read replicas, workload isolation, and automated repartitioning are **GreptimeDB Enterprise** features, along with enterprise security and governance. The [Enterprise overview](https://docs.greptime.com/enterprise/overview/) has the current list, and [pricing](https://www.greptime.com/pricing#differences) has the edition comparison. + ## Architecture GreptimeDB can run in two modes: @@ -131,7 +141,7 @@ GreptimeDB can run in two modes: * **Distributed** — four components, each independently scalable: - **Frontend** — protocol entry (OTel, Prometheus, MySQL/PostgreSQL, gRPC, ingestion APIs for Elasticsearch/InfluxDB/Loki) and the distributed query engine. Stateless, scales horizontally. - **Datanode** — region engine with WAL, memtable, SST, cache, compaction, and indexes. Persists data to object storage. Elastic. - - **Metasrv** — metadata, routing, repartitioning, autopilot, and security. Backed by a pluggable KV layer (etcd or RDS). + - **Metasrv** — metadata, routing, repartitioning, and security. Backed by a pluggable KV layer (etcd or RDS). - **Flownode** (optional) — continuous flow computation (streaming and materialized views). For deeper coverage, see the [architecture doc](https://docs.greptime.com/contributor-guide/overview/#architecture) or [DeepWiki](https://deepwiki.com/GreptimeTeam/greptimedb/1-overview). @@ -209,7 +219,7 @@ See the [Contribution Guidelines](CONTRIBUTING.md) for the full developer workfl ## Project Status -GreptimeDB is generally available, with stable APIs and regular releases. It runs in production at scale — [OceanBase Cloud](https://greptime.com/blogs/2025-07-22-user-case-obcloud-log-management-greptimedb) operates 80+ GreptimeDB clusters managing 300 TB of logs, cutting log storage cost by 60% after migrating from Grafana Loki. See more in [case studies](https://greptime.com/blogs/?category=Use%20Case). +GreptimeDB is generally available, with stable APIs and regular releases. It runs in production at scale — [OceanBase Cloud](https://greptime.com/blogs/2025-07-22-user-case-obcloud-log-management-greptimedb) operates 80+ GreptimeDB clusters managing 300 TB of logs, cutting log storage cost by 60%+ after migrating from Grafana Loki. See more in [case studies](https://greptime.com/blogs/?category=Use%20Case). Read the [v1.0 highlights](https://greptime.com/blogs/2025-11-05-greptimedb-v1-highlights) and [2026 roadmap](https://greptime.com/blogs/2026-02-11-greptimedb-roadmap-2026), or browse the [version reference](https://docs.greptime.com/nightly/reference/about-greptimedb-version). @@ -241,8 +251,9 @@ that license carry an explicit Enterprise License header. ## Commercial Support -Running GreptimeDB in your organization? -We offer enterprise add-ons, services, training, and consulting. +Scaling observability on your infrastructure? +[GreptimeDB Enterprise](https://docs.greptime.com/enterprise/overview/) adds the operational, +security, and support layer for production deployments. [Contact us](https://greptime.com/contactus) for details. ## Contributing From cb30837cd56c8e6abbdbfba098728de3ea61990d Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 14 Aug 2026 12:06:27 +0000 Subject: [PATCH 077/126] feat: add incremental primary key index writer (#8788) * feat: initial implementation of the pk index writer Signed-off-by: evenyag * perf: optimize primary key index writer Signed-off-by: evenyag * chore: add todo Signed-off-by: evenyag * feat(mito2): track series count in pk index metrics Signed-off-by: evenyag * refactor(mito2): simplify pk index writer cleanup Signed-off-by: evenyag * fix(mito2): clean up aborted pk index writers Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/AGENTS.md | 1 + src/mito2/src/lib.rs | 1 + src/mito2/src/pk_index.rs | 19 + src/mito2/src/pk_index/writer.rs | 1172 ++++++++++++++++++++++++++++++ 4 files changed, 1193 insertions(+) create mode 100644 src/mito2/src/pk_index.rs create mode 100644 src/mito2/src/pk_index/writer.rs diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index 73fc517fa7e..af60d784f1c 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -26,6 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | +| `pk_index` | `src/mito2/src/pk_index/` | Incremental writer for aggregate primary-key index files | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | | `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits | | `cache` | `src/mito2/src/cache.rs` | Write/file/page caches | diff --git a/src/mito2/src/lib.rs b/src/mito2/src/lib.rs index a212cf69533..9dfa478060b 100644 --- a/src/mito2/src/lib.rs +++ b/src/mito2/src/lib.rs @@ -38,6 +38,7 @@ pub mod gc; pub mod manifest; pub mod memtable; mod metrics; +pub mod pk_index; pub mod read; pub mod region; mod region_write_ctx; diff --git a/src/mito2/src/pk_index.rs b/src/mito2/src/pk_index.rs new file mode 100644 index 00000000000..61318026d6c --- /dev/null +++ b/src/mito2/src/pk_index.rs @@ -0,0 +1,19 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Primary-key index writer. + +mod writer; + +pub use writer::{PkIndexWriter, PkIndexWriterMetrics, PkIndexWriterOptions, pk_columns_schema}; diff --git a/src/mito2/src/pk_index/writer.rs b/src/mito2/src/pk_index/writer.rs new file mode 100644 index 00000000000..f2407d97423 --- /dev/null +++ b/src/mito2/src/pk_index/writer.rs @@ -0,0 +1,1172 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::cmp::Ordering; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use bytes::Bytes; +use datatypes::arrow::array::{ + Array, ArrayRef, BinaryArray, DictionaryArray, Int64Array, StringArray, UInt32Array, + UInt64Array, +}; +use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, UInt32Type}; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::prelude::ConcreteDataType; +use datatypes::timestamp::timestamp_array_to_primitive; +use datatypes::value::Value; +use futures::future::BoxFuture; +use mito_codec::row_converter::{CompositeValues, PrimaryKeyCodec, build_primary_key_codec}; +use object_store::{ObjectStore, Writer}; +use parquet::arrow::AsyncArrowWriter; +use parquet::arrow::async_writer::AsyncFileWriter; +use parquet::basic::{Compression, Encoding, ZstdLevel}; +use parquet::errors::ParquetError; +use parquet::file::properties::WriterProperties; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::RegionMetadataRef; +use store_api::storage::ColumnId; +use store_api::storage::consts::ReservedColumnId; + +use crate::access_layer::TempFileCleaner; +use crate::error::{ + DecodeSnafu, InvalidMetaSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, OpenDalSnafu, + Result, UnexpectedSnafu, WriteParquetSnafu, +}; +use crate::sst::parquet::DEFAULT_ROW_GROUP_SIZE; +use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index}; +use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; + +const MIN_TS_COLUMN: &str = "__pk_min_ts"; +const MAX_TS_COLUMN: &str = "__pk_max_ts"; +const ROW_COUNT_COLUMN: &str = "__pk_row_count"; +const TABLE_ID_COLUMN: &str = "__table_id"; +const TSID_COLUMN: &str = "__tsid"; +const WRITE_BATCH_SIZE: usize = 1024; + +type ParquetWriter = AsyncArrowWriter; + +/// Bridges an OpenDAL [`Writer`] with Parquet's [`AsyncFileWriter`] and tracks +/// the number of bytes successfully submitted to the object store. +struct AsyncWriter { + inner: Writer, + output_bytes: u64, +} + +impl AsyncWriter { + fn new(inner: Writer) -> Self { + Self { + inner, + output_bytes: 0, + } + } + + fn output_bytes(&self) -> u64 { + self.output_bytes + } + + fn into_inner(self) -> Writer { + self.inner + } +} + +impl AsyncFileWriter for AsyncWriter { + fn write(&mut self, bytes: Bytes) -> BoxFuture<'_, parquet::errors::Result<()>> { + Box::pin(async move { + let len = bytes.len() as u64; + self.inner + .write(bytes) + .await + .map_err(|error| ParquetError::External(Box::new(error)))?; + self.output_bytes += len; + Ok(()) + }) + } + + fn complete(&mut self) -> BoxFuture<'_, parquet::errors::Result<()>> { + Box::pin(async move { + self.inner + .close() + .await + .map(|_| ()) + .map_err(|error| ParquetError::External(Box::new(error))) + }) + } +} + +/// Options for writing a primary-key index. +#[derive(Debug, Clone)] +pub struct PkIndexWriterOptions { + /// Maximum number of rows in a Parquet row group. + pub row_group_size: usize, +} + +impl Default for PkIndexWriterOptions { + fn default() -> Self { + Self { + row_group_size: DEFAULT_ROW_GROUP_SIZE, + } + } +} + +/// Metrics collected by a [`PkIndexWriter`]. +#[derive(Debug, Clone, Default)] +pub struct PkIndexWriterMetrics { + /// Number of input record batches passed to the writer. + pub input_batches: usize, + /// Number of logical rows passed to the writer. + pub input_rows: usize, + /// Number of unique time series (primary keys) aggregated by the writer. + pub num_series: usize, + /// Size of the committed index file. This remains zero for an aborted writer. + pub output_bytes: u64, + /// Time spent opening the object-store and Parquet writers. + pub open_elapsed: Duration, + /// Time spent validating, aggregating, and decoding primary keys. + pub aggregate_elapsed: Duration, + /// Time spent encoding and writing Parquet batches. + pub write_elapsed: Duration, + /// Time spent closing a completed file. + pub finish_elapsed: Duration, + /// Time spent removing an incomplete output. + pub cleanup_elapsed: Duration, + /// Whether this writer was explicitly aborted. + pub aborted: bool, +} + +impl PkIndexWriterMetrics { + /// Returns time spent by writer-owned work. + pub fn total_elapsed(&self) -> Duration { + self.open_elapsed + + self.aggregate_elapsed + + self.write_elapsed + + self.finish_elapsed + + self.cleanup_elapsed + } +} + +#[derive(Debug)] +struct PkIndexRow { + min_ts: i64, + max_ts: i64, + row_count: u64, + table_id: u32, + tsid: u64, + tags: Vec>, +} + +/// Incrementally aggregates sorted flat record batches and writes `pk_columns` Parquet. +pub struct PkIndexWriter { + codec: Arc, + tag_columns: Vec<(ColumnId, String)>, + schema: SchemaRef, + object_store: ObjectStore, + file_name: String, + writer: Option, + current_primary_key: Option>, + current_row: Option, + buffered_rows: Vec, + metrics: PkIndexWriterMetrics, + failed: bool, +} + +impl PkIndexWriter { + /// Creates a writer for `path` in `object_store`. + /// + /// The file name in `path` must be unique in the object store so aborting + /// the writer only removes temporary files that belong to this writer. + pub async fn try_new( + metadata: RegionMetadataRef, + object_store: ObjectStore, + path: &str, + options: PkIndexWriterOptions, + ) -> Result { + let open_start = Instant::now(); + ensure!( + options.row_group_size > 0, + InvalidMetaSnafu { + reason: "primary-key index row group size must be greater than zero", + } + ); + let schema = pk_columns_schema(&metadata)?; + let tag_columns = tag_columns(&metadata); + let file_name = path.rsplit('/').next().unwrap_or(path).to_string(); + let output = object_store + .writer_with(path) + .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .concurrent(DEFAULT_WRITE_CONCURRENCY) + .await + .context(OpenDalSnafu)?; + let properties = WriterProperties::builder() + .set_compression(Compression::ZSTD(ZstdLevel::default())) + .set_encoding(Encoding::PLAIN) + .set_max_row_group_row_count(Some(options.row_group_size)) + .set_column_index_truncate_length(None) + .set_statistics_truncate_length(None) + .build(); + let writer = + AsyncArrowWriter::try_new(AsyncWriter::new(output), schema.clone(), Some(properties)) + .context(WriteParquetSnafu)?; + let codec = build_primary_key_codec(&metadata); + + Ok(Self { + codec, + tag_columns, + schema, + object_store, + file_name, + writer: Some(writer), + current_primary_key: None, + current_row: None, + buffered_rows: Vec::with_capacity(WRITE_BATCH_SIZE), + metrics: PkIndexWriterMetrics { + open_elapsed: open_start.elapsed(), + ..Default::default() + }, + failed: false, + }) + } + + /// Returns the metrics collected so far. + pub fn metrics(&self) -> &PkIndexWriterMetrics { + &self.metrics + } + + /// Aggregates one sorted flat record batch. + /// + /// After this method returns an error, callers must call [`Self::abort`]. + pub async fn write(&mut self, batch: &RecordBatch) -> Result<()> { + ensure!( + !self.failed, + InvalidRecordBatchSnafu { + reason: "cannot write to a failed primary-key index writer", + } + ); + + self.metrics.input_batches += 1; + self.metrics.input_rows += batch.num_rows(); + let aggregate_start = Instant::now(); + let write_before = self.metrics.write_elapsed; + let result = self.write_inner(batch).await; + let write_cost = self.metrics.write_elapsed.saturating_sub(write_before); + self.metrics.aggregate_elapsed += aggregate_start.elapsed().saturating_sub(write_cost); + if result.is_err() { + self.failed = true; + } + result + } + + /// Finishes and commits the index file. + pub async fn finish(mut self) -> Result { + if self.failed { + let error = InvalidRecordBatchSnafu { + reason: "cannot finish a failed primary-key index writer", + } + .build(); + self.cleanup().await; + return Err(error); + } + + let result = self.finish_inner().await; + if result.is_err() { + self.cleanup().await; + } + result.map(|_| self.metrics) + } + + /// Aborts the writer and removes incomplete output files. + pub async fn abort(mut self) -> Result { + self.metrics.aborted = true; + self.cleanup().await; + Ok(self.metrics) + } + + async fn write_inner(&mut self, batch: &RecordBatch) -> Result<()> { + if batch.num_rows() == 0 { + return Ok(()); + } + ensure!( + batch.num_columns() >= 4, + InvalidRecordBatchSnafu { + reason: format!( + "primary-key index input has too few columns: {}", + batch.num_columns() + ), + } + ); + + let pk_idx = primary_key_column_index(batch.num_columns()); + let ts_idx = time_index_column_index(batch.num_columns()); + let primary_keys = batch.column(pk_idx); + let timestamps = timestamp_values(batch.column(ts_idx))?; + ensure!( + primary_keys.len() == batch.num_rows() && timestamps.len() == batch.num_rows(), + InvalidRecordBatchSnafu { + reason: "primary-key or timestamp array length does not match the batch", + } + ); + + if let Some(array) = primary_keys.as_any().downcast_ref::() { + ensure!( + array.null_count() == 0, + InvalidRecordBatchSnafu { + reason: "primary-key index input contains null primary keys", + } + ); + self.write_binary_primary_keys(array, timestamps.values()) + .await + } else if let Some(array) = primary_keys + .as_any() + .downcast_ref::>() + { + ensure!( + array.null_count() == 0, + InvalidRecordBatchSnafu { + reason: "primary-key index input contains null primary keys", + } + ); + self.write_dictionary_primary_keys(array, timestamps.values()) + .await + } else { + InvalidRecordBatchSnafu { + reason: format!( + "primary-key index requires Binary or Dictionary(UInt32, Binary) primary keys, got {:?}", + primary_keys.data_type() + ), + } + .fail() + } + } + + async fn write_binary_primary_keys( + &mut self, + primary_keys: &BinaryArray, + timestamps: &[i64], + ) -> Result<()> { + let mut start = 0; + while start < primary_keys.len() { + let primary_key = primary_keys.value(start); + let mut end = start + 1; + while end < primary_keys.len() && primary_keys.value(end) == primary_key { + end += 1; + } + + self.update_primary_key( + primary_key, + timestamps[start], + timestamps[end - 1], + (end - start) as u64, + ) + .await?; + start = end; + } + Ok(()) + } + + async fn write_dictionary_primary_keys( + &mut self, + primary_keys: &DictionaryArray, + timestamps: &[i64], + ) -> Result<()> { + let values = primary_keys + .values() + .as_any() + .downcast_ref::() + .context(InvalidRecordBatchSnafu { + reason: "primary-key dictionary values are not binary", + })?; + let keys = primary_keys.keys().values(); + let mut start = 0; + while start < keys.len() { + let key = keys[start]; + let mut end = start + 1; + while end < keys.len() && keys[end] == key { + end += 1; + } + + self.update_primary_key( + values.value(key as usize), + timestamps[start], + timestamps[end - 1], + (end - start) as u64, + ) + .await?; + start = end; + } + Ok(()) + } + + async fn update_primary_key( + &mut self, + primary_key: &[u8], + min_ts: i64, + max_ts: i64, + row_count: u64, + ) -> Result<()> { + if let Some(current) = self.current_primary_key.as_deref() { + match primary_key.cmp(current) { + Ordering::Less => { + return InvalidRecordBatchSnafu { + reason: "primary-key index input is not sorted by primary key", + } + .fail(); + } + Ordering::Equal => { + let row = self.current_row.as_mut().context(InvalidRecordBatchSnafu { + reason: "primary-key index aggregation state is incomplete", + })?; + row.min_ts = row.min_ts.min(min_ts); + row.max_ts = row.max_ts.max(max_ts); + row.row_count += row_count; + return Ok(()); + } + Ordering::Greater => self.finish_current_row().await?, + } + } + + let row = decode_primary_key( + self.codec.as_ref(), + primary_key, + min_ts, + max_ts, + row_count, + &self.tag_columns, + )?; + self.current_primary_key = Some(primary_key.to_vec()); + self.current_row = Some(row); + Ok(()) + } + + async fn finish_current_row(&mut self) -> Result<()> { + self.current_primary_key = None; + if let Some(row) = self.current_row.take() { + self.buffered_rows.push(row); + self.metrics.num_series += 1; + } + if self.buffered_rows.len() >= WRITE_BATCH_SIZE { + self.flush_rows().await?; + } + Ok(()) + } + + async fn flush_rows(&mut self) -> Result<()> { + if self.buffered_rows.is_empty() { + return Ok(()); + } + let batch = rows_to_batch(&self.schema, &self.buffered_rows)?; + let start = Instant::now(); + let result = self + .writer + .as_mut() + .context(UnexpectedSnafu { + reason: "primary-key index Parquet writer is closed", + })? + .write(&batch) + .await + .context(WriteParquetSnafu); + self.metrics.write_elapsed += start.elapsed(); + result?; + self.buffered_rows.clear(); + Ok(()) + } + + async fn finish_inner(&mut self) -> Result<()> { + let aggregate_start = Instant::now(); + let write_before = self.metrics.write_elapsed; + self.finish_current_row().await?; + self.flush_rows().await?; + let write_cost = self.metrics.write_elapsed.saturating_sub(write_before); + self.metrics.aggregate_elapsed += aggregate_start.elapsed().saturating_sub(write_cost); + + let finish_start = Instant::now(); + self.writer + .as_mut() + .context(UnexpectedSnafu { + reason: "primary-key index Parquet writer is closed", + })? + .finish() + .await + .context(WriteParquetSnafu)?; + let writer = self.writer.take().context(UnexpectedSnafu { + reason: "primary-key index Parquet writer is closed", + })?; + self.metrics.output_bytes = writer.into_inner().output_bytes(); + self.metrics.finish_elapsed += finish_start.elapsed(); + Ok(()) + } + + async fn cleanup(&mut self) { + let start = Instant::now(); + if let Some(writer) = self.writer.take() { + let mut writer = writer.into_inner().into_inner(); + if let Err(error) = writer.abort().await { + common_telemetry::warn!(error; "Failed to abort primary-key index writer"); + } + } + self.current_primary_key = None; + self.current_row = None; + self.buffered_rows.clear(); + + TempFileCleaner::clean_atomic_dir_files(&self.object_store, &[&self.file_name]).await; + self.metrics.output_bytes = 0; + self.metrics.cleanup_elapsed += start.elapsed(); + } +} + +/// Returns the Arrow schema of a `pk_columns` primary-key index. +pub fn pk_columns_schema(metadata: &RegionMetadataRef) -> Result { + validate_metadata(metadata)?; + let mut fields = vec![ + Field::new(MIN_TS_COLUMN, DataType::Int64, false), + Field::new(MAX_TS_COLUMN, DataType::Int64, false), + Field::new(ROW_COUNT_COLUMN, DataType::UInt64, false), + Field::new(TABLE_ID_COLUMN, DataType::UInt32, false), + Field::new(TSID_COLUMN, DataType::UInt64, false), + ]; + fields.extend( + tag_columns(metadata) + .into_iter() + .map(|(_, name)| Field::new(name, DataType::Utf8, true)), + ); + Ok(Arc::new(Schema::new(fields))) +} + +fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> { + for column in &metadata.column_metadatas { + ensure!( + !matches!( + column.column_schema.name.as_str(), + MIN_TS_COLUMN | MAX_TS_COLUMN | ROW_COUNT_COLUMN + ), + InvalidMetaSnafu { + reason: format!( + "primary-key index internal column name {} is already in use", + column.column_schema.name + ), + } + ); + } + ensure!( + metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse, + InvalidMetaSnafu { + reason: "primary-key index only supports sparse primary-key encoding", + } + ); + ensure!( + metadata + .primary_key + .starts_with(&[ReservedColumnId::table_id(), ReservedColumnId::tsid()]), + InvalidMetaSnafu { + reason: "primary-key index requires (__table_id, __tsid) as the primary-key prefix", + } + ); + let table_id = metadata + .column_by_id(ReservedColumnId::table_id()) + .context(InvalidMetaSnafu { + reason: "primary-key index metadata is missing __table_id", + })?; + let tsid = metadata + .column_by_id(ReservedColumnId::tsid()) + .context(InvalidMetaSnafu { + reason: "primary-key index metadata is missing __tsid", + })?; + ensure!( + table_id.column_schema.data_type == ConcreteDataType::uint32_datatype() + && tsid.column_schema.data_type == ConcreteDataType::uint64_datatype(), + InvalidMetaSnafu { + reason: "primary-key index requires UInt32 __table_id and UInt64 __tsid", + } + ); + for column in metadata.primary_key_columns() { + if is_reserved_column(column.column_id) { + continue; + } + ensure!( + column.column_schema.data_type == ConcreteDataType::string_datatype(), + InvalidMetaSnafu { + reason: format!( + "primary-key index requires string tag column {}, got {}", + column.column_schema.name, column.column_schema.data_type + ), + } + ); + } + Ok(()) +} + +fn tag_columns(metadata: &RegionMetadataRef) -> Vec<(ColumnId, String)> { + metadata + .primary_key_columns() + .filter(|column| !is_reserved_column(column.column_id)) + .map(|column| (column.column_id, column.column_schema.name.clone())) + .collect() +} + +fn is_reserved_column(column_id: ColumnId) -> bool { + column_id == ReservedColumnId::table_id() || column_id == ReservedColumnId::tsid() +} + +fn timestamp_values(array: &ArrayRef) -> Result { + let timestamps = if let Some(array) = array.as_any().downcast_ref::() { + array.clone() + } else { + timestamp_array_to_primitive(array) + .map(|(array, _)| array) + .with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "primary-key index requires an Int64 or timestamp time index, got {:?}", + array.data_type() + ), + })? + }; + ensure!( + timestamps.null_count() == 0, + InvalidRecordBatchSnafu { + reason: "primary-key index input contains null timestamps", + } + ); + Ok(timestamps) +} + +// TODO(yingwen): Bench and optimize the performance if this is costly. +fn decode_primary_key( + codec: &dyn PrimaryKeyCodec, + primary_key: &[u8], + min_ts: i64, + max_ts: i64, + row_count: u64, + tag_columns: &[(ColumnId, String)], +) -> Result { + let CompositeValues::Sparse(values) = codec.decode(primary_key).context(DecodeSnafu)? else { + return InvalidRecordBatchSnafu { + reason: "decoded primary key is not sparse", + } + .fail(); + }; + let table_id = match values.get(&ReservedColumnId::table_id()) { + Some(Value::UInt32(value)) => *value, + value => { + return InvalidRecordBatchSnafu { + reason: format!("missing or invalid sparse __table_id: {value:?}"), + } + .fail(); + } + }; + let tsid = match values.get(&ReservedColumnId::tsid()) { + Some(Value::UInt64(value)) => *value, + value => { + return InvalidRecordBatchSnafu { + reason: format!("missing or invalid sparse __tsid: {value:?}"), + } + .fail(); + } + }; + let tags = tag_columns + .iter() + .map(|(column_id, _)| match values.get(column_id) { + None | Some(Value::Null) => Ok(None), + Some(Value::String(value)) => Ok(Some(value.as_utf8().to_string())), + value => InvalidRecordBatchSnafu { + reason: format!( + "invalid sparse string tag value for column {column_id}: {value:?}" + ), + } + .fail(), + }) + .collect::>>()?; + Ok(PkIndexRow { + min_ts, + max_ts, + row_count, + table_id, + tsid, + tags, + }) +} + +fn rows_to_batch(schema: &SchemaRef, rows: &[PkIndexRow]) -> Result { + let mut arrays: Vec = vec![ + Arc::new(Int64Array::from_iter_values( + rows.iter().map(|row| row.min_ts), + )), + Arc::new(Int64Array::from_iter_values( + rows.iter().map(|row| row.max_ts), + )), + Arc::new(UInt64Array::from_iter_values( + rows.iter().map(|row| row.row_count), + )), + Arc::new(UInt32Array::from_iter_values( + rows.iter().map(|row| row.table_id), + )), + Arc::new(UInt64Array::from_iter_values( + rows.iter().map(|row| row.tsid), + )), + ]; + for tag_idx in 0..schema.fields().len() - 5 { + arrays.push(Arc::new(StringArray::from_iter( + rows.iter().map(|row| row.tags[tag_idx].as_deref()), + ))); + } + RecordBatch::try_new(schema.clone(), arrays).context(NewRecordBatchSnafu) +} + +#[cfg(test)] +mod tests { + use api::v1::SemanticType; + use datatypes::arrow::array::{ + BinaryDictionaryBuilder, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, UInt8Array, + }; + use datatypes::arrow::datatypes::{TimeUnit, UInt32Type}; + use datatypes::schema::ColumnSchema; + use mito_codec::row_converter::{PrimaryKeyCodec, SparsePrimaryKeyCodec}; + use object_store::ErrorKind; + use object_store::services::Memory; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use store_api::codec::PrimaryKeyEncoding; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + + use super::*; + use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; + + fn object_store() -> ObjectStore { + ObjectStore::new(Memory::default()).unwrap().finish() + } + + fn flat_schema(primary_key_type: DataType) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("__primary_key", primary_key_type, false), + Field::new("__sequence", DataType::UInt64, false), + Field::new("__op_type", DataType::UInt8, false), + ])) + } + + fn binary_batch(primary_keys: &[&[u8]], timestamps: &[i64]) -> RecordBatch { + RecordBatch::try_new( + flat_schema(DataType::Binary), + vec![ + Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())), + Arc::new(BinaryArray::from_iter_values(primary_keys.iter().copied())), + Arc::new(UInt64Array::from(vec![1; timestamps.len()])), + Arc::new(UInt8Array::from(vec![0; timestamps.len()])), + ], + ) + .unwrap() + } + + fn dictionary_batch(primary_keys: &[&[u8]], timestamps: &[i64]) -> RecordBatch { + let mut builder = BinaryDictionaryBuilder::::new(); + for primary_key in primary_keys { + builder.append(*primary_key).unwrap(); + } + RecordBatch::try_new( + flat_schema(DataType::Dictionary( + Box::new(DataType::UInt32), + Box::new(DataType::Binary), + )), + vec![ + Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())), + Arc::new(builder.finish()), + Arc::new(UInt64Array::from(vec![1; timestamps.len()])), + Arc::new(UInt8Array::from(vec![0; timestamps.len()])), + ], + ) + .unwrap() + } + + async fn read_index(store: &ObjectStore, path: &str) -> (u64, usize, Vec) { + let bytes = store.read(path).await.unwrap().to_bytes(); + let output_bytes = bytes.len() as u64; + let builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap(); + let row_groups = builder.metadata().num_row_groups(); + let batches = builder + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + (output_bytes, row_groups, batches) + } + + #[test] + fn test_pk_columns_schema() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let schema = pk_columns_schema(&metadata).unwrap(); + assert_eq!( + schema + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + [ + "__pk_min_ts", + "__pk_max_ts", + "__pk_row_count", + "__table_id", + "__tsid", + "tag_0", + "tag_1", + ] + ); + assert!(!schema.field(4).is_nullable()); + assert!(schema.field(5).is_nullable()); + + let dense = Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense)); + assert!(pk_columns_schema(&dense).is_err()); + } + + #[tokio::test] + async fn test_reject_pk_index_internal_column_names_before_opening_writer() { + for (index, name) in [MIN_TS_COLUMN, MAX_TS_COLUMN, ROW_COUNT_COLUMN] + .into_iter() + .enumerate() + { + let mut builder = RegionMetadataBuilder::from_existing( + sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse), + ); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new(name, ConcreteDataType::string_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 100 + index as u32, + }); + let metadata = Arc::new(builder.build().unwrap()); + let store = object_store(); + let path = format!("collision-{index}.parquet"); + let error = PkIndexWriter::try_new( + metadata, + store.clone(), + &path, + PkIndexWriterOptions::default(), + ) + .await + .err() + .unwrap(); + + assert!(error.to_string().contains(name), "{error}"); + assert_eq!( + store.stat(&path).await.unwrap_err().kind(), + ErrorKind::NotFound + ); + } + } + + #[test] + fn test_timestamp_values() { + let arrays: Vec = vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(TimestampSecondArray::from(vec![1, 2])), + Arc::new(TimestampMillisecondArray::from(vec![1, 2])), + Arc::new(TimestampMicrosecondArray::from(vec![1, 2])), + Arc::new(TimestampNanosecondArray::from(vec![1, 2])), + ]; + for array in arrays { + assert_eq!(timestamp_values(&array).unwrap().values().as_ref(), &[1, 2]); + } + + let nulls: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![Some(1), None])); + let error = timestamp_values(&nulls).unwrap_err(); + assert!(error.to_string().contains("null timestamps"), "{error}"); + + let unsupported: ArrayRef = Arc::new(UInt8Array::from(vec![1, 2])); + let error = timestamp_values(&unsupported).unwrap_err(); + assert!( + error.to_string().contains("requires an Int64 or timestamp"), + "{error}" + ); + } + + #[tokio::test] + async fn test_write_batches_and_metrics() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let primary_key_1 = new_sparse_primary_key(&["a", "x"], &metadata, 1, 10); + let primary_key_2 = new_sparse_primary_key(&["b", "y"], &metadata, 1, 20); + let store = object_store(); + let mut writer = PkIndexWriter::try_new( + metadata, + store.clone(), + "pk.parquet", + PkIndexWriterOptions { row_group_size: 2 }, + ) + .await + .unwrap(); + + writer + .write(&dictionary_batch( + &[ + primary_key_1.as_slice(), + primary_key_1.as_slice(), + primary_key_1.as_slice(), + primary_key_1.as_slice(), + ], + &[70, 80, 90, 100], + )) + .await + .unwrap(); + writer + .write(&binary_batch( + &[ + primary_key_1.as_slice(), + primary_key_1.as_slice(), + primary_key_1.as_slice(), + primary_key_2.as_slice(), + primary_key_2.as_slice(), + primary_key_2.as_slice(), + primary_key_2.as_slice(), + ], + &[110, 120, 130, 200, 210, 220, 230], + )) + .await + .unwrap(); + assert_eq!(writer.metrics().input_batches, 2); + assert_eq!(writer.metrics().input_rows, 11); + + let metrics = writer.finish().await.unwrap(); + assert_eq!(metrics.input_batches, 2); + assert_eq!(metrics.input_rows, 11); + assert_eq!(metrics.num_series, 2); + assert!(!metrics.aborted); + + let (output_bytes, row_groups, batches) = read_index(&store, "pk.parquet").await; + assert_eq!(metrics.output_bytes, output_bytes); + assert_eq!(row_groups, 1); + assert_eq!(batches.len(), 1); + let batch = &batches[0]; + assert_eq!(batch.num_rows(), 2); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(), + &Int64Array::from(vec![70, 200]) + ); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(), + &Int64Array::from(vec![130, 230]) + ); + assert_eq!( + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(), + &UInt64Array::from(vec![7, 4]) + ); + assert_eq!( + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap(), + &UInt32Array::from(vec![1, 1]) + ); + assert_eq!( + batch + .column(4) + .as_any() + .downcast_ref::() + .unwrap(), + &UInt64Array::from(vec![10, 20]) + ); + assert_eq!( + batch + .column(5) + .as_any() + .downcast_ref::() + .unwrap(), + &StringArray::from(vec![Some("a"), Some("b")]) + ); + } + + #[tokio::test] + async fn test_row_group_size_and_empty_file() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let store = object_store(); + let mut writer = PkIndexWriter::try_new( + metadata.clone(), + store.clone(), + "groups.parquet", + PkIndexWriterOptions { row_group_size: 2 }, + ) + .await + .unwrap(); + let keys = (0..5) + .map(|tsid| new_sparse_primary_key(&["a", "x"], &metadata, 1, tsid)) + .collect::>(); + let key_refs = keys.iter().map(Vec::as_slice).collect::>(); + writer + .write(&binary_batch(&key_refs, &[1, 2, 3, 4, 5])) + .await + .unwrap(); + writer.finish().await.unwrap(); + let (_, row_groups, _) = read_index(&store, "groups.parquet").await; + assert_eq!(row_groups, 3); + + let empty = PkIndexWriter::try_new( + metadata, + store.clone(), + "empty.parquet", + PkIndexWriterOptions::default(), + ) + .await + .unwrap() + .finish() + .await + .unwrap(); + assert_eq!(empty.input_rows, 0); + assert_eq!(empty.num_series, 0); + let (output_bytes, _, batches) = read_index(&store, "empty.parquet").await; + assert_eq!(empty.output_bytes, output_bytes); + assert!(batches.is_empty()); + } + + #[tokio::test] + async fn test_abort_and_out_of_order_input() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let primary_key_1 = new_sparse_primary_key(&["a", "x"], &metadata, 1, 10); + let primary_key_2 = new_sparse_primary_key(&["b", "y"], &metadata, 1, 20); + let store = object_store(); + let mut writer = PkIndexWriter::try_new( + metadata.clone(), + store.clone(), + "abort.parquet", + PkIndexWriterOptions { row_group_size: 1 }, + ) + .await + .unwrap(); + let error = writer + .write(&binary_batch( + &[primary_key_2.as_slice(), primary_key_1.as_slice()], + &[1, 2], + )) + .await + .unwrap_err(); + assert!(error.to_string().contains("not sorted"), "{error}"); + store + .write("abort.parquet", Bytes::from_static(b"existing")) + .await + .unwrap(); + let metrics = writer.abort().await.unwrap(); + assert!(metrics.aborted); + assert_eq!(metrics.output_bytes, 0); + assert_eq!( + store.read("abort.parquet").await.unwrap().to_bytes(), + Bytes::from_static(b"existing") + ); + + let mut writer = PkIndexWriter::try_new( + metadata, + store.clone(), + "dictionary-abort.parquet", + PkIndexWriterOptions { row_group_size: 1 }, + ) + .await + .unwrap(); + let error = writer + .write(&dictionary_batch( + &[ + primary_key_2.as_slice(), + primary_key_2.as_slice(), + primary_key_1.as_slice(), + ], + &[1, 2, 3], + )) + .await + .unwrap_err(); + assert!(error.to_string().contains("not sorted"), "{error}"); + writer.abort().await.unwrap(); + assert_eq!( + store + .stat("dictionary-abort.parquet") + .await + .unwrap_err() + .kind(), + ErrorKind::NotFound + ); + } + + #[tokio::test] + async fn test_nullable_tag_and_invalid_options() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let codec = SparsePrimaryKeyCodec::new(&metadata); + let mut primary_key = Vec::new(); + codec + .encode_value_refs( + &[ + ( + ReservedColumnId::table_id(), + datatypes::value::ValueRef::UInt32(1), + ), + ( + ReservedColumnId::tsid(), + datatypes::value::ValueRef::UInt64(10), + ), + (0, datatypes::value::ValueRef::String("a")), + ], + &mut primary_key, + ) + .unwrap(); + let store = object_store(); + let mut writer = PkIndexWriter::try_new( + metadata.clone(), + store.clone(), + "nullable.parquet", + PkIndexWriterOptions::default(), + ) + .await + .unwrap(); + writer + .write(&binary_batch(&[primary_key.as_slice()], &[1])) + .await + .unwrap(); + writer.finish().await.unwrap(); + let (_, _, batches) = read_index(&store, "nullable.parquet").await; + let tag_1 = batches[0] + .column(6) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(tag_1.is_null(0)); + + assert!( + PkIndexWriter::try_new( + metadata, + store, + "invalid.parquet", + PkIndexWriterOptions { row_group_size: 0 }, + ) + .await + .is_err() + ); + } +} From a8924bb95c43b562e4312af309d2a5e80c103185 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:20:04 +0000 Subject: [PATCH 078/126] refactor(udaf): replace uddsketch implementation (#8867) * refactor(function): replace uddsketch implementation Signed-off-by: Lei, HUANG * bench(function): compare uddsketch batch ingestion Signed-off-by: Lei, HUANG * perf(function): avoid copying non-null uddsketch batches Signed-off-by: Lei, HUANG * fix(function): decode legacy uddsketch states Signed-off-by: Lei, HUANG * fix(function): harden legacy uddsketch validation Signed-off-by: Lei, HUANG * format: taplo Signed-off-by: Lei, HUANG * test: add compatibility tests for uddsketch functions Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- Cargo.lock | 18 +- Cargo.toml | 2 + src/common/function/Cargo.toml | 8 +- src/common/function/benches/uddsketch.rs | 201 ++++++ .../src/aggrs/approximate/uddsketch.rs | 299 +++++---- src/common/function/src/lib.rs | 1 + .../function/src/scalars/uddsketch_calc.rs | 99 ++- src/common/function/src/uddsketch_compat.rs | 619 ++++++++++++++++++ .../cases/uddsketch_legacy_state/case.toml | 9 + .../cases/uddsketch_legacy_state/setup.sql | 31 + .../uddsketch_legacy_state/verify.result | 63 ++ .../cases/uddsketch_legacy_state/verify.sql | 33 + 12 files changed, 1235 insertions(+), 148 deletions(-) create mode 100644 src/common/function/benches/uddsketch.rs create mode 100644 src/common/function/src/uddsketch_compat.rs create mode 100644 tests/compatibility/cases/uddsketch_legacy_state/case.toml create mode 100644 tests/compatibility/cases/uddsketch_legacy_state/setup.sql create mode 100644 tests/compatibility/cases/uddsketch_legacy_state/verify.result create mode 100644 tests/compatibility/cases/uddsketch_legacy_state/verify.sql diff --git a/Cargo.lock b/Cargo.lock index ece381b7b84..0f473d4e030 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2477,6 +2477,7 @@ dependencies = [ "common-telemetry", "common-time", "common-version", + "criterion 0.7.0", "datafusion", "datafusion-common", "datafusion-expr", @@ -2511,7 +2512,8 @@ dependencies = [ "store-api", "table", "tokio", - "uddsketch", + "uddsketch-rs", + "vu128", "wkt", ] @@ -15823,11 +15825,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" [[package]] -name = "uddsketch" +name = "uddsketch-rs" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/timescaledb-toolkit.git?rev=84828fe8fb494a6a61412a3da96517fc80f7bb20#84828fe8fb494a6a61412a3da96517fc80f7bb20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f0495d3e02c0e9b074e5631de7af419b9f4fdbfec4b1ca1b3180c8b3a926cb5" dependencies = [ - "serde", + "thiserror 2.0.17", + "vu128", ] [[package]] @@ -16216,6 +16220,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "vu128" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b18da3bd753c6f4373511e5f025423986560dfe4a5e7d642cc9a0266847f9fdd" + [[package]] name = "walkdir" version = "2.5.0" diff --git a/Cargo.toml b/Cargo.toml index bfd6153bf5b..daee1bd4090 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -266,8 +266,10 @@ tracing-appender = "0.2" tracing-opentelemetry = "0.33.0" tracing-subscriber = { version = "0.3", features = ["env-filter", "json", "fmt"] } typetag = "0.2" +uddsketch-rs = "0.1.0" uuid = { version = "1.17", features = ["serde", "v4", "v7", "fast-rng"] } vrl = "0.33" +vu128 = "1.1" zstd = "0.13" # DO_NOT_REMOVE_THIS: END_OF_EXTERNAL_DEPENDENCIES diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml index 622f3ddfd23..a1ea31bc32d 100644 --- a/src/common/function/Cargo.toml +++ b/src/common/function/Cargo.toml @@ -65,12 +65,18 @@ snafu.workspace = true sql.workspace = true store-api.workspace = true table.workspace = true -uddsketch = { git = "https://github.com/GreptimeTeam/timescaledb-toolkit.git", rev = "84828fe8fb494a6a61412a3da96517fc80f7bb20" } +uddsketch-rs.workspace = true +vu128.workspace = true wkt = { version = "0.11", optional = true } [dev-dependencies] approx = "0.5" +criterion.workspace = true futures.workspace = true pretty_assertions.workspace = true serde = { version = "1.0", features = ["derive"] } tokio.workspace = true + +[[bench]] +name = "uddsketch" +harness = false diff --git a/src/common/function/benches/uddsketch.rs b/src/common/function/benches/uddsketch.rs new file mode 100644 index 00000000000..6c8ad8779c2 --- /dev/null +++ b/src/common/function/benches/uddsketch.rs @@ -0,0 +1,201 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Float64Array, Int64Array}; +use arrow::datatypes::{DataType, Field, Schema}; +use common_function::aggrs::approximate::uddsketch::UddSketchState; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion::common::ScalarValue; +use datafusion::logical_expr::Accumulator; +use datafusion::logical_expr::function::AccumulatorArgs; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::{Column, Literal}; + +const BATCH_SIZES: [usize; 4] = [128, 256, 1024, 2048]; +const BUCKET_SIZE: i64 = 128; +const ERROR_RATE: f64 = 0.01; + +struct AccumulatorFactory { + udf: datafusion::logical_expr::AggregateUDF, + schema: Schema, + exprs: Vec>, + expr_fields: Vec>, + return_field: Arc, +} + +impl AccumulatorFactory { + fn new() -> Self { + let udf = UddSketchState::state_udf_impl(); + let schema = Schema::new(vec![ + Field::new("bucket_size", DataType::Int64, false), + Field::new("error", DataType::Float64, false), + Field::new("value", DataType::Float64, false), + ]); + let exprs: Vec> = vec![ + Arc::new(Literal::new(ScalarValue::Int64(Some(BUCKET_SIZE)))), + Arc::new(Literal::new(ScalarValue::Float64(Some(ERROR_RATE)))), + Arc::new(Column::new("value", 2)), + ]; + let expr_fields = exprs + .iter() + .map(|expr| expr.return_field(&schema).unwrap()) + .collect::>(); + let return_type = udf + .return_type(&[DataType::Int64, DataType::Float64, DataType::Float64]) + .unwrap(); + + Self { + udf, + schema, + exprs, + expr_fields, + return_field: Arc::new(Field::new("uddsketch_state", return_type, true)), + } + } + + fn create(&self) -> Box { + self.udf + .accumulator(AccumulatorArgs { + return_field: Arc::clone(&self.return_field), + schema: &self.schema, + ignore_nulls: false, + order_bys: &[], + is_reversed: false, + name: "uddsketch_state", + is_distinct: false, + exprs: &self.exprs, + expr_fields: &self.expr_fields, + }) + .unwrap() + } +} + +fn input_arrays(batch_size: usize) -> Vec { + let bucket_sizes = Arc::new(Int64Array::from_value(BUCKET_SIZE, batch_size)) as ArrayRef; + let errors = Arc::new(Float64Array::from_value(ERROR_RATE, batch_size)) as ArrayRef; + let mut state = 0x9e37_79b9_7f4a_7c15_u64; + let values = (0..batch_size) + .map(|index| { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407) + .wrapping_add(index as u64); + let unit = (state >> 11) as f64 * (1.0 / (1_u64 << 53) as f64); + let magnitude = 10_f64.powf(-9.0 + 18.0 * unit); + if state & 1 == 0 { + magnitude + } else { + -magnitude + } + }) + .collect::>(); + let values = Arc::new(Float64Array::from(values)) as ArrayRef; + + vec![bucket_sizes, errors, values] +} + +fn validate(factory: &AccumulatorFactory, values: &[ArrayRef]) { + let mut accumulator = factory.create(); + accumulator.update_batch(values).unwrap(); + match accumulator.evaluate().unwrap() { + ScalarValue::Binary(Some(encoded)) => assert!(!encoded.is_empty()), + encoded => panic!("expected non-empty Binary, got {encoded:?}"), + } +} + +fn bench_uddsketch(c: &mut Criterion) { + let factory = AccumulatorFactory::new(); + let inputs = BATCH_SIZES + .into_iter() + .map(|batch_size| { + let values = input_arrays(batch_size); + validate(&factory, &values); + (batch_size, values) + }) + .collect::>(); + + let mut group = c.benchmark_group("uddsketch/ingest/fresh"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + b.iter(|| { + let mut accumulator = factory.create(); + accumulator.update_batch(black_box(values)).unwrap(); + black_box(accumulator); + }); + }, + ); + } + group.finish(); + + let mut group = c.benchmark_group("uddsketch/ingest/reused"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + let mut accumulator = factory.create(); + b.iter(|| { + accumulator.update_batch(black_box(values)).unwrap(); + black_box(&mut accumulator); + }); + }, + ); + } + group.finish(); + + let mut group = c.benchmark_group("uddsketch/ingest_evaluate/fresh"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + b.iter(|| { + let mut accumulator = factory.create(); + accumulator.update_batch(black_box(values)).unwrap(); + black_box(accumulator.evaluate().unwrap()); + }); + }, + ); + } + group.finish(); + + let mut group = c.benchmark_group("uddsketch/ingest_evaluate/reused"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + let mut accumulator = factory.create(); + b.iter(|| { + accumulator.update_batch(black_box(values)).unwrap(); + black_box(accumulator.evaluate().unwrap()); + }); + }, + ); + } + group.finish(); +} + +criterion_group!(benches, bench_uddsketch); +criterion_main!(benches); diff --git a/src/common/function/src/aggrs/approximate/uddsketch.rs b/src/common/function/src/aggrs/approximate/uddsketch.rs index 24b5c5455c0..f7d1558d13c 100644 --- a/src/common/function/src/aggrs/approximate/uddsketch.rs +++ b/src/common/function/src/aggrs/approximate/uddsketch.rs @@ -21,7 +21,6 @@ use std::sync::Arc; use common_query::prelude::*; -use common_telemetry::trace; use datafusion::common::cast::{as_binary_array, as_primitive_array}; use datafusion::common::not_impl_err; use datafusion::error::{DataFusionError, Result as DfResult}; @@ -29,27 +28,40 @@ use datafusion::logical_expr::function::AccumulatorArgs; use datafusion::logical_expr::{Accumulator as DfAccumulator, AggregateUDF, Volatility}; use datafusion::physical_plan::expressions::Literal; use datafusion::prelude::create_udaf; -use datatypes::arrow::array::ArrayRef; +use datatypes::arrow::array::{Array, ArrayRef}; use datatypes::arrow::datatypes::{DataType, Float64Type}; -use serde::{Deserialize, Serialize}; -use uddsketch::{SketchHashKey, UDDSketch}; +use uddsketch::{BatchWorkspace, UddSketch}; + +use crate::uddsketch_compat; pub const UDDSKETCH_STATE_NAME: &str = "uddsketch_state"; pub const UDDSKETCH_MERGE_NAME: &str = "uddsketch_merge"; -#[derive(Debug, Serialize, Deserialize)] +const MAX_BUCKETS: u32 = 1_000_000; + +#[derive(Debug)] pub struct UddSketchState { - uddsketch: UDDSketch, - error_rate: f64, + uddsketch: UddSketch, + workspace: BatchWorkspace, + values: Vec, } impl UddSketchState { - pub fn new(bucket_size: u64, error_rate: f64) -> Self { - Self { - uddsketch: UDDSketch::new(bucket_size, error_rate), - error_rate, + pub fn new(bucket_size: u32, error_rate: f64) -> DfResult { + if bucket_size > MAX_BUCKETS { + return Err(DataFusionError::Plan(format!( + "UDDSketch bucket size exceeds the maximum of {}", + MAX_BUCKETS + ))); } + let uddsketch = UddSketch::new(bucket_size, error_rate) + .map_err(|e| DataFusionError::Plan(e.to_string()))?; + Ok(Self { + uddsketch, + workspace: BatchWorkspace::default(), + values: Vec::new(), + }) } pub fn state_udf_impl() -> AggregateUDF { @@ -60,7 +72,7 @@ impl UddSketchState { Volatility::Immutable, Arc::new(|args| { let (bucket_size, error_rate) = downcast_accumulator_args(args)?; - Ok(Box::new(UddSketchState::new(bucket_size, error_rate))) + Ok(Box::new(UddSketchState::new(bucket_size, error_rate)?)) }), Arc::new(vec![DataType::Binary]), ) @@ -80,51 +92,44 @@ impl UddSketchState { Volatility::Immutable, Arc::new(|args| { let (bucket_size, error_rate) = downcast_accumulator_args(args)?; - Ok(Box::new(UddSketchState::new(bucket_size, error_rate))) + Ok(Box::new(UddSketchState::new(bucket_size, error_rate)?)) }), Arc::new(vec![DataType::Binary]), ) } - fn update(&mut self, value: f64) { - self.uddsketch.add_value(value); - } - fn merge(&mut self, raw: &[u8]) -> DfResult<()> { - if let Ok(uddsketch) = bincode::deserialize::(raw) { - if uddsketch.uddsketch.count() != 0 { - if self.uddsketch.max_allowed_buckets() != uddsketch.uddsketch.max_allowed_buckets() - || (self.error_rate - uddsketch.error_rate).abs() >= 1e-9 - { - return Err(DataFusionError::Plan(format!( - "Merging UDDSketch with different parameters: arguments={:?} vs actual input={:?}", - (self.uddsketch.max_allowed_buckets(), self.error_rate), - ( - uddsketch.uddsketch.max_allowed_buckets(), - uddsketch.error_rate - ) - ))); - } - self.uddsketch.merge_sketch(&uddsketch.uddsketch); - } - } else { - trace!("Warning: Failed to deserialize UDDSketch from {:?}", raw); - return Err(DataFusionError::Plan( - "Failed to deserialize UDDSketch from binary".to_string(), - )); + let uddsketch = uddsketch_compat::decode(raw).map_err(|e| { + common_telemetry::trace!("Failed to deserialize UDDSketch: {}", e); + DataFusionError::Plan("Failed to deserialize UDDSketch from binary".to_string()) + })?; + if uddsketch.count() == 0 { + return Ok(()); } - - Ok(()) + if self.uddsketch.max_buckets() != uddsketch.max_buckets() + || self.uddsketch.initial_error().to_bits() != uddsketch.initial_error().to_bits() + { + return Err(DataFusionError::Plan(format!( + "Merging UDDSketch with different parameters: arguments={:?} vs actual input={:?}", + (self.uddsketch.max_buckets(), self.uddsketch.initial_error()), + (uddsketch.max_buckets(), uddsketch.initial_error()) + ))); + } + self.uddsketch + .merge(&uddsketch) + .map_err(|e| DataFusionError::Plan(e.to_string())) } } -fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u64, f64)> { +fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u32, f64)> { let bucket_size = match args.exprs[0] .as_any() .downcast_ref::() .map(|lit| lit.value()) { - Some(ScalarValue::Int64(Some(value))) => *value as u64, + Some(ScalarValue::Int64(Some(value))) => u32::try_from(*value).map_err(|_| { + DataFusionError::Plan(format!("Invalid UDDSketch bucket size: {}", value)) + })?, _ => { return not_impl_err!( "{} not supported for bucket size: {}", @@ -158,9 +163,16 @@ impl DfAccumulator for UddSketchState { match array.data_type() { DataType::Float64 => { let f64_array = as_primitive_array::(array)?; - for v in f64_array.iter().flatten() { - self.update(v); - } + let values: &[f64] = if f64_array.null_count() == 0 { + f64_array.values().as_ref() + } else { + self.values.clear(); + self.values.extend(f64_array.iter().flatten()); + self.values.as_slice() + }; + self.uddsketch + .add_batch_with_workspace(values, &mut self.workspace) + .map_err(|e| DataFusionError::Execution(e.to_string()))?; } // meaning instantiate as `uddsketch_merge` DataType::Binary => self.merge_batch(std::slice::from_ref(array))?, @@ -176,35 +188,21 @@ impl DfAccumulator for UddSketchState { } fn evaluate(&mut self) -> DfResult { - Ok(ScalarValue::Binary(Some( - bincode::serialize(&self).map_err(|e| { - DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e)) - })?, - ))) + Ok(ScalarValue::Binary(Some(self.uddsketch.encode().map_err( + |e| DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e)), + )?))) } fn size(&self) -> usize { - // Base size of UDDSketch struct fields - let mut total_size = std::mem::size_of::() * 3 + // alpha, gamma, values_sum - std::mem::size_of::() + // compactions - std::mem::size_of::() * 2; // max_buckets, num_values - - // Size of buckets (SketchHashMap) - // Each bucket entry contains: - // - SketchHashKey (enum with i64/Zero/Invalid variants) - // - SketchHashEntry (count: u64, next: SketchHashKey) - let bucket_entry_size = std::mem::size_of::() + // key - std::mem::size_of::() + // count - std::mem::size_of::(); // next - - total_size += self.uddsketch.current_buckets_count() * bucket_entry_size; - - total_size + std::mem::size_of::() - std::mem::size_of::() + + self.uddsketch.allocated_size() + + self.workspace.allocated_size() + + self.values.capacity() * std::mem::size_of::() } fn state(&mut self) -> DfResult> { Ok(vec![ScalarValue::Binary(Some( - bincode::serialize(&self).map_err(|e| { + self.uddsketch.encode().map_err(|e| { DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e)) })?, ))]) @@ -224,20 +222,21 @@ impl DfAccumulator for UddSketchState { #[cfg(test)] mod tests { use datafusion::arrow::array::{BinaryArray, Float64Array}; + use uddsketch::UddSketchRef; use super::*; #[test] fn test_uddsketch_state_basic() { - let mut state = UddSketchState::new(10, 0.01); - state.update(1.0); - state.update(2.0); - state.update(3.0); + let mut state = UddSketchState::new(10, 0.01).unwrap(); + state.uddsketch.add(1.0).unwrap(); + state.uddsketch.add(2.0).unwrap(); + state.uddsketch.add(3.0).unwrap(); let result = state.evaluate().unwrap(); if let ScalarValue::Binary(Some(bytes)) = result { - let deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap(); - assert_eq!(deserialized.uddsketch.count(), 3); + let encoded = UddSketchRef::parse(&bytes).unwrap(); + assert_eq!(encoded.count(), 3); } else { panic!("Expected binary scalar value"); } @@ -245,39 +244,38 @@ mod tests { #[test] fn test_uddsketch_state_roundtrip() { - let mut state = UddSketchState::new(10, 0.01); - state.update(1.0); - state.update(2.0); + let mut state = UddSketchState::new(10, 0.01).unwrap(); + state.uddsketch.add(1.0).unwrap(); + state.uddsketch.add(2.0).unwrap(); // Serialize let serialized = state.evaluate().unwrap(); // Create new state and merge the serialized data - let mut new_state = UddSketchState::new(10, 0.01); + let mut new_state = UddSketchState::new(10, 0.01).unwrap(); if let ScalarValue::Binary(Some(bytes)) = &serialized { new_state.merge(bytes).unwrap(); - // Verify the merged state matches original by comparing deserialized values - let original_sketch: UddSketchState = bincode::deserialize(bytes).unwrap(); - let original_sketch = original_sketch.uddsketch; + let original_sketch = UddSketchRef::parse(bytes).unwrap(); let new_result = new_state.evaluate().unwrap(); if let ScalarValue::Binary(Some(new_bytes)) = new_result { - let new_sketch: UddSketchState = bincode::deserialize(&new_bytes).unwrap(); - let new_sketch = new_sketch.uddsketch; + let new_sketch = UddSketchRef::parse(&new_bytes).unwrap(); assert_eq!(original_sketch.count(), new_sketch.count()); assert_eq!(original_sketch.sum(), new_sketch.sum()); - assert_eq!(original_sketch.mean(), new_sketch.mean()); - assert_eq!(original_sketch.max_error(), new_sketch.max_error()); + assert_eq!( + original_sketch.max_error().unwrap(), + new_sketch.max_error().unwrap() + ); // Compare a few quantiles to ensure statistical equivalence for q in [0.1, 0.5, 0.9].iter() { + let original = original_sketch.quantile(*q).unwrap().unwrap(); + let merged = new_sketch.quantile(*q).unwrap().unwrap(); assert!( - (original_sketch.estimate_quantile(*q) - new_sketch.estimate_quantile(*q)) - .abs() - < 1e-10, + (original - merged).abs() < 1e-10, "Quantile {} mismatch: original={}, new={}", q, - original_sketch.estimate_quantile(*q), - new_sketch.estimate_quantile(*q) + original, + merged ); } } else { @@ -288,10 +286,41 @@ mod tests { } } + #[test] + fn test_uddsketch_state_merges_legacy_state() { + let mut state = UddSketchState::new(128, 0.01).unwrap(); + + state.merge(uddsketch_compat::LEGACY_STATE).unwrap(); + + let ScalarValue::Binary(Some(encoded)) = state.evaluate().unwrap() else { + panic!("Expected binary scalar value"); + }; + let sketch = UddSketchRef::parse(&encoded).unwrap(); + assert_eq!(sketch.count(), 4); + assert_eq!(sketch.sum(), 1.0); + assert_eq!(sketch.quantile(0.5).unwrap(), Some(0.9900000000000001)); + } + + #[test] + fn test_uddsketch_state_merges_compacted_legacy_state() { + let mut legacy_state = uddsketch_compat::COMPACTED_LEGACY_SKETCH.to_vec(); + legacy_state.extend_from_slice(&0.01_f64.to_le_bytes()); + let mut state = UddSketchState::new(7, 0.01).unwrap(); + + state.merge(&legacy_state).unwrap(); + + let ScalarValue::Binary(Some(encoded)) = state.evaluate().unwrap() else { + panic!("Expected binary scalar value"); + }; + let sketch = UddSketchRef::parse(&encoded).unwrap(); + assert_eq!(sketch.count(), 201); + assert_eq!(sketch.times_compacted(), 12); + } + #[test] fn test_uddsketch_state_batch_update() { - let mut state = UddSketchState::new(10, 0.01); - let values = vec![1.0f64, 2.0, 3.0]; + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let values = vec![Some(1.0f64), None, Some(2.0), Some(3.0)]; let array = Arc::new(Float64Array::from(values)) as ArrayRef; state @@ -300,25 +329,52 @@ mod tests { let result = state.evaluate().unwrap(); if let ScalarValue::Binary(Some(bytes)) = result { - let deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap(); - let deserialized = deserialized.uddsketch; - assert_eq!(deserialized.count(), 3); + let encoded = UddSketchRef::parse(&bytes).unwrap(); + assert_eq!(encoded.count(), 3); } else { panic!("Expected binary scalar value"); } } + #[test] + fn test_uddsketch_state_non_null_batch_avoids_values_buffer() { + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let array = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])) as ArrayRef; + + state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap(); + + assert_eq!(state.uddsketch.count(), 3); + assert_eq!(state.values.capacity(), 0); + } + + #[test] + fn test_uddsketch_state_non_null_sliced_batch() { + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let array = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0]); + let array = array.slice(1, 2); + let array = Arc::new(array) as ArrayRef; + + state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap(); + + assert_eq!(state.uddsketch.count(), 2); + assert_eq!(state.uddsketch.sum(), 5.0); + } + #[test] fn test_uddsketch_state_merge_batch() { - let mut state1 = UddSketchState::new(10, 0.01); - state1.update(1.0); + let mut state1 = UddSketchState::new(10, 0.01).unwrap(); + state1.uddsketch.add(1.0).unwrap(); let state1_binary = state1.evaluate().unwrap(); - let mut state2 = UddSketchState::new(10, 0.01); - state2.update(2.0); + let mut state2 = UddSketchState::new(10, 0.01).unwrap(); + state2.uddsketch.add(2.0).unwrap(); let state2_binary = state2.evaluate().unwrap(); - let mut merged_state = UddSketchState::new(10, 0.01); + let mut merged_state = UddSketchState::new(10, 0.01).unwrap(); if let (ScalarValue::Binary(Some(bytes1)), ScalarValue::Binary(Some(bytes2))) = (&state1_binary, &state2_binary) { @@ -330,9 +386,8 @@ mod tests { let result = merged_state.evaluate().unwrap(); if let ScalarValue::Binary(Some(bytes)) = result { - let deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap(); - let deserialized = deserialized.uddsketch; - assert_eq!(deserialized.count(), 2); + let encoded = UddSketchRef::parse(&bytes).unwrap(); + assert_eq!(encoded.count(), 2); } else { panic!("Expected binary scalar value"); } @@ -343,13 +398,14 @@ mod tests { #[test] fn test_uddsketch_state_size() { - let mut state = UddSketchState::new(10, 0.01); + let mut state = UddSketchState::new(10, 0.01).unwrap(); let initial_size = state.size(); // Add some values to create buckets - state.update(1.0); - state.update(2.0); - state.update(3.0); + let array = Arc::new(Float64Array::from_iter_values((0..64).map(f64::from))) as ArrayRef; + state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap(); let size_with_values = state.size(); assert!( @@ -358,14 +414,29 @@ mod tests { initial_size, size_with_values ); + } - // Verify size increases with more buckets - state.update(10.0); // This should create a new bucket - assert!( - state.size() > size_with_values, - "Size should increase after adding new bucket: prev={}, new={}", - size_with_values, - state.size() - ); + #[test] + fn test_uddsketch_state_rejects_invalid_config() { + assert!(UddSketchState::new(6, 0.01).is_err()); + assert!(UddSketchState::new(10, 1.0).is_err()); + + let mut maximum = UddSketchState::new(1_000_000, 0.01).unwrap(); + let ScalarValue::Binary(Some(encoded)) = maximum.evaluate().unwrap() else { + panic!("Expected binary scalar value"); + }; + UddSketchRef::parse(&encoded).unwrap(); + assert!(UddSketchState::new(1_000_001, 0.01).is_err()); + } + + #[test] + fn test_uddsketch_state_rejects_nan_batch() { + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let array = Arc::new(Float64Array::from(vec![1.0, f64::NAN])) as ArrayRef; + + let error = state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap_err(); + assert!(error.to_string().contains("NaN values are not supported")); } } diff --git a/src/common/function/src/lib.rs b/src/common/function/src/lib.rs index 7abd5953671..336f8d489ca 100644 --- a/src/common/function/src/lib.rs +++ b/src/common/function/src/lib.rs @@ -18,6 +18,7 @@ mod admin; mod flush_flow; mod macros; mod system; +mod uddsketch_compat; pub mod aggrs; pub mod function; diff --git a/src/common/function/src/scalars/uddsketch_calc.rs b/src/common/function/src/scalars/uddsketch_calc.rs index d5d91e7edb6..d28bb78d1be 100644 --- a/src/common/function/src/scalars/uddsketch_calc.rs +++ b/src/common/function/src/scalars/uddsketch_calc.rs @@ -22,10 +22,10 @@ use datafusion_common::DataFusionError; use datafusion_common::arrow::array::{Array, AsArray, Float64Builder}; use datafusion_common::arrow::datatypes::{DataType, Float64Type}; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility}; -use uddsketch::UDDSketch; use crate::function::{Function, extract_args}; use crate::function_registry::FunctionRegistry; +use crate::uddsketch_compat; const NAME: &str = "uddsketch_calc"; @@ -114,26 +114,19 @@ impl Function for UddSketchCalcFunction { let sketch_bytes = sketch_opt.unwrap(); let perc = perc_opt.unwrap(); - // Deserialize the UDDSketch from its bincode representation - let sketch: UDDSketch = match bincode::deserialize(sketch_bytes) { - Ok(s) => s, + let value = match uddsketch_compat::quantile(sketch_bytes, perc) { + Ok(value) => value, Err(e) => { - common_telemetry::trace!("Failed to deserialize UDDSketch: {}", e); + common_telemetry::trace!("Failed to parse UDDSketch: {}", e); builder.append_null(); continue; } }; - // Check if the sketch is empty, if so, return null - // This is important to avoid panics when calling estimate_quantile on an empty sketch - // In practice, this will happen if input is all null - if sketch.bucket_iter().count() == 0 { - builder.append_null(); - continue; + match value { + Some(value) => builder.append_value(value), + None => builder.append_null(), } - // Compute the estimated quantile from the sketch - let result = sketch.estimate_quantile(perc); - builder.append_value(result); } Ok(ColumnarValue::Array(Arc::new(builder.finish()))) @@ -146,6 +139,7 @@ mod tests { use arrow_schema::Field; use datafusion_common::arrow::array::{BinaryArray, Float64Array}; + use uddsketch::UddSketch; use super::*; @@ -159,24 +153,17 @@ mod tests { ); // Create a test sketch - let mut sketch = UDDSketch::new(128, 0.01); - sketch.add_value(10.0); - sketch.add_value(20.0); - sketch.add_value(30.0); - sketch.add_value(40.0); - sketch.add_value(50.0); - sketch.add_value(60.0); - sketch.add_value(70.0); - sketch.add_value(80.0); - sketch.add_value(90.0); - sketch.add_value(100.0); + let mut sketch = UddSketch::new(128, 0.01).unwrap(); + sketch + .add_batch(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0]) + .unwrap(); // Get expected values directly from the sketch - let expected_p50 = sketch.estimate_quantile(0.5); - let expected_p90 = sketch.estimate_quantile(0.9); - let expected_p95 = sketch.estimate_quantile(0.95); + let expected_p50 = sketch.quantile(0.5).unwrap().unwrap(); + let expected_p90 = sketch.quantile(0.9).unwrap().unwrap(); + let expected_p95 = sketch.quantile(0.95).unwrap().unwrap(); - let serialized = bincode::serialize(&sketch).unwrap(); + let serialized = sketch.encode().unwrap(); let percentiles = vec![0.5, 0.9, 0.95]; let args = vec![ @@ -207,6 +194,34 @@ mod tests { assert!((result.value(2) - expected_p95).abs() < 1e-10); } + #[test] + fn test_uddsketch_calc_function_reads_legacy_state() { + let function = UddSketchCalcFunction::default(); + let args = vec![ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![0.5]))), + ColumnarValue::Array(Arc::new(BinaryArray::from_iter_values(vec![ + uddsketch_compat::LEGACY_STATE, + ]))), + ]; + + let result = function + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: 1, + return_field: Arc::new(Field::new("x", DataType::Float64, false)), + config_options: Arc::new(Default::default()), + }) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + let result = result.as_primitive::(); + assert_eq!(result.len(), 1); + assert!(!result.is_null(0)); + assert_eq!(result.value(0), 0.9900000000000001); + } + #[test] fn test_uddsketch_calc_function_errors() { let function = UddSketchCalcFunction::default(); @@ -249,5 +264,31 @@ mod tests { let result = result.as_primitive::(); assert_eq!(result.len(), 1); assert!(result.is_null(0)); + + let empty = UddSketch::new(128, 0.01).unwrap().encode().unwrap(); + let mut populated = UddSketch::new(128, 0.01).unwrap(); + populated.add(1.0).unwrap(); + let populated = populated.encode().unwrap(); + let args = vec![ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![0.5, -0.1, f64::NAN]))), + ColumnarValue::Array(Arc::new(BinaryArray::from_iter_values(vec![ + empty, + populated.clone(), + populated, + ]))), + ]; + let result = function + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: 3, + return_field: Arc::new(Field::new("x", DataType::Float64, false)), + config_options: Arc::new(Default::default()), + }) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + assert_eq!(result.as_primitive::().null_count(), 3); } } diff --git a/src/common/function/src/uddsketch_compat.rs b/src/common/function/src/uddsketch_compat.rs new file mode 100644 index 00000000000..733b4ab043a --- /dev/null +++ b/src/common/function/src/uddsketch_compat.rs @@ -0,0 +1,619 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Compatibility decoder for UDDSketch states written before the canonical v1 format. + +use std::collections::{HashMap, HashSet}; + +use bincode::Options; +use serde::Deserialize; +use uddsketch::{UddSketch, UddSketchRef}; + +const MAX_BYTES: usize = 64 * 1024 * 1024; +const MAX_BUCKETS: usize = 1_000_000; +const HEADER_LEN: usize = 48; + +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq)] +enum LegacyBucketKey { + Negative(i64), + Zero, + Positive(i64), + Invalid, +} + +#[derive(Debug, Deserialize)] +struct LegacyBucket { + count: u64, + next: LegacyBucketKey, +} + +#[derive(Debug, Deserialize)] +struct LegacyBucketStore { + map: HashMap, + head: LegacyBucketKey, +} + +#[derive(Debug, Deserialize)] +struct LegacyUddSketch { + buckets: LegacyBucketStore, + alpha: f64, + gamma: f64, + compactions: u32, + max_buckets: u64, + count: u64, + sum: f64, +} + +#[derive(Debug, Deserialize)] +struct LegacyUddSketchState { + uddsketch: LegacyUddSketch, + initial_error: f64, +} + +pub(crate) fn decode(raw: &[u8]) -> Result { + match UddSketch::decode(raw) { + Ok(sketch) => Ok(sketch), + Err(current_error) => decode_legacy_state(raw).map_err(|legacy_error| { + format!( + "canonical decode failed: {current_error}; legacy decode failed: {legacy_error}" + ) + }), + } +} + +pub(crate) fn quantile(raw: &[u8], quantile: f64) -> Result, String> { + match UddSketchRef::parse(raw) { + Ok(sketch) => sketch.quantile(quantile).map_err(|error| error.to_string()), + Err(current_error) => decode_legacy_sketch(raw) + .and_then(|sketch| sketch.quantile(quantile)) + .map_err(|legacy_error| { + format!( + "canonical decode failed: {current_error}; legacy decode failed: {legacy_error}" + ) + }), + } +} + +fn validate_legacy_input(raw: &[u8]) -> Result<(), String> { + if raw.len() > MAX_BYTES { + return Err("input exceeds the legacy decode byte limit".to_string()); + } + let map_len = raw + .get(..8) + .and_then(|bytes| bytes.try_into().ok()) + .map(u64::from_le_bytes) + .ok_or_else(|| "legacy input is truncated before the bucket count".to_string())?; + if map_len > MAX_BUCKETS as u64 { + return Err("legacy populated bucket count exceeds decode limit".to_string()); + } + Ok(()) +} + +fn legacy_options() -> impl Options { + bincode::DefaultOptions::new() + .with_fixint_encoding() + .with_limit(MAX_BYTES as u64) + .reject_trailing_bytes() +} + +fn decode_legacy_state(raw: &[u8]) -> Result { + validate_legacy_input(raw)?; + let state = legacy_options() + .deserialize::(raw) + .map_err(|error| error.to_string())?; + let encoded = state.into_canonical()?; + UddSketch::decode(&encoded).map_err(|error| error.to_string()) +} + +fn decode_legacy_sketch(raw: &[u8]) -> Result { + validate_legacy_input(raw)?; + legacy_options() + .deserialize::(raw) + .map(|state| state.uddsketch) + .or_else(|state_error| { + legacy_options() + .deserialize::(raw) + .map_err(|sketch_error| { + format!( + "legacy state decode failed: {state_error}; legacy sketch decode failed: {sketch_error}" + ) + }) + }) +} + +impl LegacyUddSketchState { + fn into_canonical(self) -> Result, String> { + let sketch = self.uddsketch; + let max_buckets = u32::try_from(sketch.max_buckets) + .map_err(|_| "legacy maximum bucket count exceeds u32".to_string())?; + if !(7..=MAX_BUCKETS as u32).contains(&max_buckets) { + return Err("legacy maximum bucket count is outside supported limits".to_string()); + } + let compactions = u8::try_from(sketch.compactions) + .map_err(|_| "legacy compaction count exceeds u8".to_string())?; + if compactions > 63 { + return Err("legacy compaction count exceeds 63".to_string()); + } + + let (expected_alpha, expected_gamma) = mapping(self.initial_error, compactions)?; + if sketch.alpha.to_bits() != expected_alpha.to_bits() + || sketch.gamma.to_bits() != expected_gamma.to_bits() + { + return Err("legacy mapping metadata is inconsistent".to_string()); + } + + let buckets = sketch.buckets.ordered()?; + if buckets.len() > max_buckets as usize { + return Err("legacy populated bucket count exceeds its configured limit".to_string()); + } + let decoded_count = buckets.iter().try_fold(0_u64, |total, (_, count)| { + total + .checked_add(*count) + .ok_or_else(|| "legacy bucket count sum overflows u64".to_string()) + })?; + if decoded_count != sketch.count { + return Err("legacy bucket counts do not match the value count".to_string()); + } + if sketch.count == 0 && sketch.sum.to_bits() != 0.0_f64.to_bits() { + return Err("legacy empty sketch sum is not positive zero".to_string()); + } + + encode_canonical( + max_buckets, + self.initial_error, + compactions, + sketch.count, + sketch.sum, + &buckets, + ) + } +} + +impl LegacyUddSketch { + fn quantile(self, quantile: f64) -> Result, String> { + if !quantile.is_finite() || !(0.0..=1.0).contains(&quantile) { + return Err("invalid quantile".to_string()); + } + if self.compactions >= 64 { + return Err("legacy compaction count must be below 64".to_string()); + } + validate_current_mapping(self.alpha, self.gamma, self.compactions)?; + if self.max_buckets == 0 || self.max_buckets > MAX_BUCKETS as u64 { + return Err("legacy maximum bucket count is outside supported limits".to_string()); + } + + let count = self.count; + let alpha = self.alpha; + let gamma = self.gamma; + let buckets = self.buckets.ordered()?; + if buckets.len() > self.max_buckets as usize { + return Err("legacy populated bucket count exceeds its configured limit".to_string()); + } + for (key, _) in &buckets { + if !legacy_bucket_key_is_attainable(self.gamma, *key) { + return Err("legacy bucket index cannot represent an f64".to_string()); + } + } + let decoded_count = buckets.iter().try_fold(0_u64, |total, (_, count)| { + total + .checked_add(*count) + .ok_or_else(|| "legacy bucket count sum overflows u64".to_string()) + })?; + if decoded_count != count { + return Err("legacy bucket counts do not match the value count".to_string()); + } + if count == 0 { + return Ok(None); + } + + let target = if quantile == 1.0 { + count + } else { + ((count as f64 * quantile) as u64) + .saturating_add(1) + .min(count) + }; + let mut seen = 0_u64; + for (key, bucket_count) in buckets { + seen += bucket_count; + if seen >= target { + return Ok(Some(legacy_bucket_value(alpha, gamma, key)?)); + } + } + Err("legacy bucket counts do not cover the quantile rank".to_string()) + } +} + +fn validate_current_mapping( + mut alpha: f64, + mut gamma: f64, + compactions: u32, +) -> Result<(), String> { + if !alpha.is_finite() || !(0.0..=1.0).contains(&alpha) { + return Err("legacy current error is outside [0, 1]".to_string()); + } + if !gamma.is_finite() || gamma <= 1.0 { + return (alpha == 1.0 && gamma == f64::INFINITY && compactions >= 5) + .then_some(()) + .ok_or_else(|| "legacy gamma must be greater than one".to_string()); + } + if alpha == 1.0 { + return (compactions > 0 && 1.0 - 2.0 / (gamma + 1.0) == 1.0) + .then_some(()) + .ok_or_else(|| "legacy saturated mapping metadata is inconsistent".to_string()); + } + + for _ in 0..compactions { + alpha /= 1.0 + (1.0 - alpha * alpha).sqrt(); + gamma = gamma.sqrt(); + } + let expected_gamma = (1.0 + alpha) / (1.0 - alpha); + let relative_difference = (gamma - expected_gamma).abs() / expected_gamma; + if relative_difference > 1e-10 { + return Err("legacy mapping metadata is inconsistent".to_string()); + } + Ok(()) +} + +fn legacy_bucket_key_is_attainable(gamma: f64, key: LegacyBucketKey) -> bool { + let index = match key { + LegacyBucketKey::Zero => return true, + LegacyBucketKey::Negative(index) | LegacyBucketKey::Positive(index) => index, + LegacyBucketKey::Invalid => return false, + }; + if index == i64::MAX { + return true; + } + + let minimum = f64::from_bits(1).log(gamma).ceil() as i64; + let maximum = f64::MAX.log(gamma).ceil() as i64; + (minimum..=maximum).contains(&index) +} + +fn legacy_bucket_value(alpha: f64, gamma: f64, key: LegacyBucketKey) -> Result { + let magnitude = |index: i64| gamma.powf(index as f64 - 1.0) * (1.0 + alpha); + match key { + LegacyBucketKey::Negative(index) => Ok(-magnitude(index)), + LegacyBucketKey::Zero => Ok(0.0), + LegacyBucketKey::Positive(index) => Ok(magnitude(index)), + LegacyBucketKey::Invalid => Err("legacy bucket chain contains the end marker".to_string()), + } +} + +impl LegacyBucketStore { + fn ordered(self) -> Result, String> { + if self.map.len() > MAX_BUCKETS { + return Err("legacy populated bucket count exceeds decode limit".to_string()); + } + if self.map.is_empty() { + if self.head != LegacyBucketKey::Invalid { + return Err("legacy empty bucket store has a nonempty head".to_string()); + } + return Ok(Vec::new()); + } + + let mut buckets = Vec::with_capacity(self.map.len()); + let mut visited = HashSet::with_capacity(self.map.len()); + let mut key = self.head; + while key != LegacyBucketKey::Invalid { + if !visited.insert(key) { + return Err("legacy bucket chain contains a cycle".to_string()); + } + let bucket = self + .map + .get(&key) + .ok_or_else(|| "legacy bucket chain references a missing bucket".to_string())?; + if bucket.count == 0 { + return Err("legacy bucket has a zero count".to_string()); + } + buckets.push((key, bucket.count)); + key = bucket.next; + } + if buckets.len() != self.map.len() || self.map.contains_key(&LegacyBucketKey::Invalid) { + return Err("legacy bucket store contains unreachable buckets".to_string()); + } + if !buckets.windows(2).all(|pair| key_lt(pair[0].0, pair[1].0)) { + return Err("legacy bucket chain is not strictly ordered".to_string()); + } + Ok(buckets) + } +} + +fn mapping(initial_error: f64, compactions: u8) -> Result<(f64, f64), String> { + if !initial_error.is_finite() || !(1e-12..1.0).contains(&initial_error) { + return Err("legacy initial error is outside [1e-12, 1)".to_string()); + } + let mut alpha = initial_error; + let mut gamma = (1.0 + initial_error) / (1.0 - initial_error); + for _ in 0..compactions { + gamma *= gamma; + alpha = 2.0 * alpha / (1.0 + alpha.powi(2)); + } + Ok((alpha, gamma)) +} + +fn encode_canonical( + max_buckets: u32, + initial_error: f64, + compactions: u8, + count: u64, + sum: f64, + buckets: &[(LegacyBucketKey, u64)], +) -> Result, String> { + let negative = buckets + .iter() + .filter_map(|(key, count)| match key { + LegacyBucketKey::Negative(index) => Some((*index, *count)), + _ => None, + }) + .collect::>(); + let zero_count = buckets + .iter() + .find_map(|(key, count)| (*key == LegacyBucketKey::Zero).then_some(*count)) + .unwrap_or(0); + let positive = buckets + .iter() + .filter_map(|(key, count)| match key { + LegacyBucketKey::Positive(index) => Some((*index, *count)), + _ => None, + }) + .collect::>(); + + let mut encoded = vec![0; HEADER_LEN]; + put_varint(&mut encoded, negative.len() as u64); + put_varint(&mut encoded, zero_count); + put_varint(&mut encoded, positive.len() as u64); + encode_section(&mut encoded, &negative, true)?; + encode_section(&mut encoded, &positive, false)?; + + let payload_len = u32::try_from(encoded.len() - HEADER_LEN) + .map_err(|_| "legacy canonical payload exceeds u32".to_string())?; + encoded[0..4].copy_from_slice(b"UDDS"); + encoded[4] = 1; + encoded[6] = compactions; + encoded[8..12].copy_from_slice(&max_buckets.to_le_bytes()); + encoded[12..16].copy_from_slice(&(buckets.len() as u32).to_le_bytes()); + encoded[16..24].copy_from_slice(&initial_error.to_bits().to_le_bytes()); + encoded[24..32].copy_from_slice(&count.to_le_bytes()); + encoded[32..40].copy_from_slice(&sum.to_bits().to_le_bytes()); + encoded[40..44].copy_from_slice(&payload_len.to_le_bytes()); + Ok(encoded) +} + +fn encode_section( + output: &mut Vec, + buckets: &[(i64, u64)], + descending: bool, +) -> Result<(), String> { + let mut previous = None; + for &(index, count) in buckets { + let encoded_index = match previous { + None => zigzag(index), + Some((previous_index, _)) => { + let delta = if descending { + previous_index as i128 - index as i128 + } else { + index as i128 - previous_index as i128 + }; + u64::try_from(delta) + .ok() + .filter(|delta| *delta != 0) + .ok_or_else(|| "legacy bucket indices are not strictly ordered".to_string())? + } + }; + let encoded_count = match previous { + None => count, + Some((_, previous_count)) => zigzag(count.wrapping_sub(previous_count) as i64), + }; + put_varint(output, encoded_index); + put_varint(output, encoded_count); + previous = Some((index, count)); + } + Ok(()) +} + +fn put_varint(output: &mut Vec, value: u64) { + let mut buffer = [0; 9]; + let len = vu128::encode_u64(&mut buffer, value); + output.extend_from_slice(&buffer[..len]); +} + +const fn zigzag(value: i64) -> u64 { + (value.wrapping_shl(1) ^ (value >> 63)) as u64 +} + +fn key_lt(left: LegacyBucketKey, right: LegacyBucketKey) -> bool { + use LegacyBucketKey::*; + match (left, right) { + (Negative(left), Negative(right)) => left > right, + (Negative(_), Zero | Positive(_)) | (Zero, Positive(_)) => true, + (Positive(left), Positive(right)) => left < right, + _ => false, + } +} + +#[cfg(test)] +pub(crate) const LEGACY_STATE: &[u8] = &[ + 4, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 116, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 116, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 116, + 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 116, 0, 0, 0, 0, 0, 0, 0, + 123, 20, 174, 71, 225, 122, 132, 63, 253, 74, 129, 90, 191, 82, 240, 63, 0, 0, 0, 0, 128, 0, 0, + 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 240, 63, 123, 20, 174, 71, 225, 122, + 132, 63, +]; + +#[cfg(test)] +pub(crate) const COMPACTED_LEGACY_SKETCH: &[u8] = &[ + 6, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 255, 255, 255, 255, 255, 255, 255, 255, 36, 0, 0, 0, 0, 0, + 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 36, 0, 0, 0, 0, + 0, 0, 0, 2, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 254, 255, 255, 255, 255, 255, 255, + 255, 29, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 255, 255, 255, 255, 255, 255, 255, 255, 2, 0, 0, 0, + 2, 0, 0, 0, 0, 0, 0, 0, 36, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, + 0, 3, 0, 0, 0, 0, 0, 0, 0, 29, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 2, 0, 0, 0, 1, 0, 0, 0, 0, 0, + 0, 0, 35, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 254, 255, 255, + 255, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0, 240, 63, 169, 137, 186, 120, 1, 63, 82, 71, 12, 0, + 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 201, 0, 0, 0, 0, 0, 0, 0, 112, 103, 108, 212, 220, 81, 180, 84, +]; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn current_format_is_decoded_without_legacy_fallback() { + let mut sketch = UddSketch::new(128, 0.01).unwrap(); + sketch.add(42.0).unwrap(); + let expected = sketch.quantile(0.5).unwrap(); + let encoded = sketch.encode().unwrap(); + + assert!(decode_legacy_state(&encoded).is_err()); + assert_eq!(decode(&encoded).unwrap(), sketch); + assert_eq!(quantile(&encoded, 0.5).unwrap(), expected); + } + + #[test] + fn legacy_decoder_rejects_oversized_bucket_count_before_deserializing() { + let mut encoded = LEGACY_STATE.to_vec(); + encoded[..8].copy_from_slice(&((MAX_BUCKETS as u64) + 1).to_le_bytes()); + + assert_eq!( + decode_legacy_state(&encoded).unwrap_err(), + "legacy populated bucket count exceeds decode limit" + ); + } + + #[test] + fn legacy_decoder_reads_bare_sketch_from_scalar_callers() { + let bare_sketch = &LEGACY_STATE[..LEGACY_STATE.len() - std::mem::size_of::()]; + + assert_eq!( + quantile(bare_sketch, 0.5).unwrap(), + Some(0.9900000000000001) + ); + } + + #[test] + fn legacy_decoder_reads_compacted_bare_sketch() { + assert!(decode_legacy_state(COMPACTED_LEGACY_SKETCH).is_err()); + assert!(quantile(COMPACTED_LEGACY_SKETCH, 0.5).unwrap().is_some()); + } + + #[test] + fn legacy_quantile_rejects_invalid_mapping() { + let bare_len = LEGACY_STATE.len() - std::mem::size_of::(); + let mut invalid = LEGACY_STATE[..bare_len].to_vec(); + let alpha_offset = bare_len - 44; + invalid[alpha_offset..alpha_offset + 8].copy_from_slice(&f64::NAN.to_le_bytes()); + + assert!(quantile(&invalid, 0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.gamma = 2.0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1.0; + sketch.gamma = 2.0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1e-12; + sketch.gamma = 1.0 + 1e-9; + sketch.compactions = 0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1.0; + sketch.gamma = 1e20; + sketch.compactions = 0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1.0; + sketch.gamma = f64::INFINITY; + sketch.compactions = 1; + assert!(validate_current_mapping(sketch.alpha, sketch.gamma, sketch.compactions).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 0.99999999; + sketch.gamma = 2.0; + sketch.compactions = 0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + let LegacyBucketKey::Negative(index) = sketch.buckets.head else { + panic!("Expected negative head bucket"); + }; + let bucket = sketch + .buckets + .map + .remove(&LegacyBucketKey::Negative(index)) + .unwrap(); + sketch + .buckets + .map + .insert(LegacyBucketKey::Negative(i64::MIN), bucket); + sketch.buckets.head = LegacyBucketKey::Negative(i64::MIN); + assert!(sketch.quantile(0.5).is_err()); + } + + #[test] + fn legacy_quantile_accepts_saturated_mapping() { + let sketch = LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([( + LegacyBucketKey::Positive(0), + LegacyBucket { + count: 1, + next: LegacyBucketKey::Invalid, + }, + )]), + head: LegacyBucketKey::Positive(0), + }, + alpha: 1.0, + gamma: f64::INFINITY, + compactions: 63, + max_buckets: 7, + count: 1, + sum: f64::INFINITY, + }; + + assert_eq!(sketch.quantile(0.5).unwrap(), Some(0.0)); + } + + #[test] + fn legacy_quantile_handles_maximum_count_at_one() { + let sketch = LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([( + LegacyBucketKey::Zero, + LegacyBucket { + count: u64::MAX, + next: LegacyBucketKey::Invalid, + }, + )]), + head: LegacyBucketKey::Zero, + }, + alpha: 0.01, + gamma: (1.0 + 0.01) / (1.0 - 0.01), + compactions: 0, + max_buckets: 128, + count: u64::MAX, + sum: 0.0, + }; + + assert_eq!(sketch.quantile(1.0).unwrap(), Some(0.0)); + } +} diff --git a/tests/compatibility/cases/uddsketch_legacy_state/case.toml b/tests/compatibility/cases/uddsketch_legacy_state/case.toml new file mode 100644 index 00000000000..b22670037d1 --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/case.toml @@ -0,0 +1,9 @@ +name = "uddsketch_legacy_state" +reason = "Verify UDDSketch state bytes persisted by old binaries (wrapped legacy bincode UddSketchState) are read directly by uddsketch_calc and uddsketch_merge on the new binary, merged into canonical v1, and merged together with newly generated canonical states." +introduced_by = "PR #8867" +topologies = ["distributed", "standalone"] +from_range = [">=v1.0.0"] +to_range = ["*"] +features = ["table", "query", "udf", "aggregate"] +owner = "query" +namespace = "uddsketch_legacy_state" diff --git a/tests/compatibility/cases/uddsketch_legacy_state/setup.sql b/tests/compatibility/cases/uddsketch_legacy_state/setup.sql new file mode 100644 index 00000000000..abe681a7b98 --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/setup.sql @@ -0,0 +1,31 @@ +CREATE TABLE uddsketch_values ( + seq_id INT PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO uddsketch_values (seq_id, val) VALUES + (1, 10.0), + (2, 20.0), + (3, 30.0), + (4, 40.0), + (5, 50.0), + (6, 60.0), + (7, 70.0), + (8, 80.0), + (9, 90.0), + (10, 100.0); + +CREATE TABLE uddsketch_states ( + state BINARY, + grp INT PRIMARY KEY, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO uddsketch_states (state, grp) +SELECT uddsketch_state(128, 0.01, val), seq_id / 5 * 5 AS grp +FROM uddsketch_values +GROUP BY grp; + +ADMIN FLUSH_TABLE('uddsketch_values'); +ADMIN FLUSH_TABLE('uddsketch_states'); diff --git a/tests/compatibility/cases/uddsketch_legacy_state/verify.result b/tests/compatibility/cases/uddsketch_legacy_state/verify.result new file mode 100644 index 00000000000..b4c8a3d1566 --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/verify.result @@ -0,0 +1,63 @@ +-- (1) direct uddsketch_calc of each persisted legacy row +SELECT grp, uddsketch_calc(0.5, state) AS p50 +FROM uddsketch_states +ORDER BY grp; + ++-----+--------------------+ +| grp | p50 | ++-----+--------------------+ +| 0 | 30.267171338721894 | +| 5 | 70.11183939140265 | +| 10 | 100.49456770856492 | ++-----+--------------------+ + +-- (2) uddsketch_merge over the persisted legacy rows, then calc. +-- Succeeds only if every legacy row decodes to the exact declared +-- parameters (128, 0.01); the merge parameter check is the guard. +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_states; + ++--------------------+ +| p50 | ++--------------------+ +| 59.745049810145126 | ++--------------------+ + +-- (3) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows +CREATE TABLE uddsketch_new_states ( + state BINARY, + grp INT PRIMARY KEY, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +Affected Rows: 0 + +INSERT INTO uddsketch_new_states (state, grp) +SELECT uddsketch_state(128, 0.01, val), seq_id / 5 * 5 AS grp +FROM uddsketch_values +GROUP BY grp; + +Affected Rows: 3 + +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM ( + SELECT state FROM uddsketch_states + UNION ALL + SELECT state FROM uddsketch_new_states +) AS all_states; + ++--------------------+ +| p50 | ++--------------------+ +| 59.745049810145126 | ++--------------------+ + +-- (4) canonical-only control +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_new_states; + ++--------------------+ +| p50 | ++--------------------+ +| 59.745049810145126 | ++--------------------+ diff --git a/tests/compatibility/cases/uddsketch_legacy_state/verify.sql b/tests/compatibility/cases/uddsketch_legacy_state/verify.sql new file mode 100644 index 00000000000..8cf84cfbbc5 --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/verify.sql @@ -0,0 +1,33 @@ +-- (1) direct uddsketch_calc of each persisted legacy row +SELECT grp, uddsketch_calc(0.5, state) AS p50 +FROM uddsketch_states +ORDER BY grp; + +-- (2) uddsketch_merge over the persisted legacy rows, then calc. +-- Succeeds only if every legacy row decodes to the exact declared +-- parameters (128, 0.01); the merge parameter check is the guard. +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_states; + +-- (3) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows +CREATE TABLE uddsketch_new_states ( + state BINARY, + grp INT PRIMARY KEY, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO uddsketch_new_states (state, grp) +SELECT uddsketch_state(128, 0.01, val), seq_id / 5 * 5 AS grp +FROM uddsketch_values +GROUP BY grp; + +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM ( + SELECT state FROM uddsketch_states + UNION ALL + SELECT state FROM uddsketch_new_states +) AS all_states; + +-- (4) canonical-only control +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_new_states; From 4126cf99b6cd1d58d74b26e92490039198ecac6e Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 17 Aug 2026 07:21:29 +0000 Subject: [PATCH 079/126] refactor(mito2): rename pk index to series index (#8893) Signed-off-by: evenyag --- src/mito2/AGENTS.md | 2 +- src/mito2/src/lib.rs | 2 +- .../src/{pk_index.rs => series_index.rs} | 6 +- .../src/{pk_index => series_index}/writer.rs | 148 +++++++++--------- 4 files changed, 80 insertions(+), 78 deletions(-) rename src/mito2/src/{pk_index.rs => series_index.rs} (80%) rename src/mito2/src/{pk_index => series_index}/writer.rs (89%) diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index af60d784f1c..e68ca586096 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -26,7 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | -| `pk_index` | `src/mito2/src/pk_index/` | Incremental writer for aggregate primary-key index files | +| `series_index` | `src/mito2/src/series_index/` | Incremental writer for aggregate series index files | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | | `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits | | `cache` | `src/mito2/src/cache.rs` | Write/file/page caches | diff --git a/src/mito2/src/lib.rs b/src/mito2/src/lib.rs index 9dfa478060b..4f4409afd8a 100644 --- a/src/mito2/src/lib.rs +++ b/src/mito2/src/lib.rs @@ -38,13 +38,13 @@ pub mod gc; pub mod manifest; pub mod memtable; mod metrics; -pub mod pk_index; pub mod read; pub mod region; mod region_write_ctx; pub mod remap_manifest; pub mod request; pub mod schedule; +pub mod series_index; pub mod sst; mod time_provider; pub mod wal; diff --git a/src/mito2/src/pk_index.rs b/src/mito2/src/series_index.rs similarity index 80% rename from src/mito2/src/pk_index.rs rename to src/mito2/src/series_index.rs index 61318026d6c..21ced00cf64 100644 --- a/src/mito2/src/pk_index.rs +++ b/src/mito2/src/series_index.rs @@ -12,8 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Primary-key index writer. +//! Series index writer. mod writer; -pub use writer::{PkIndexWriter, PkIndexWriterMetrics, PkIndexWriterOptions, pk_columns_schema}; +pub use writer::{ + SeriesIndexWriter, SeriesIndexWriterMetrics, SeriesIndexWriterOptions, series_index_schema, +}; diff --git a/src/mito2/src/pk_index/writer.rs b/src/mito2/src/series_index/writer.rs similarity index 89% rename from src/mito2/src/pk_index/writer.rs rename to src/mito2/src/series_index/writer.rs index f2407d97423..ce5b7405f21 100644 --- a/src/mito2/src/pk_index/writer.rs +++ b/src/mito2/src/series_index/writer.rs @@ -49,9 +49,9 @@ use crate::sst::parquet::DEFAULT_ROW_GROUP_SIZE; use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index}; use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; -const MIN_TS_COLUMN: &str = "__pk_min_ts"; -const MAX_TS_COLUMN: &str = "__pk_max_ts"; -const ROW_COUNT_COLUMN: &str = "__pk_row_count"; +const MIN_TS_COLUMN: &str = "__series_min_ts"; +const MAX_TS_COLUMN: &str = "__series_max_ts"; +const ROW_COUNT_COLUMN: &str = "__series_row_count"; const TABLE_ID_COLUMN: &str = "__table_id"; const TSID_COLUMN: &str = "__tsid"; const WRITE_BATCH_SIZE: usize = 1024; @@ -106,14 +106,14 @@ impl AsyncFileWriter for AsyncWriter { } } -/// Options for writing a primary-key index. +/// Options for writing a series index. #[derive(Debug, Clone)] -pub struct PkIndexWriterOptions { +pub struct SeriesIndexWriterOptions { /// Maximum number of rows in a Parquet row group. pub row_group_size: usize, } -impl Default for PkIndexWriterOptions { +impl Default for SeriesIndexWriterOptions { fn default() -> Self { Self { row_group_size: DEFAULT_ROW_GROUP_SIZE, @@ -121,9 +121,9 @@ impl Default for PkIndexWriterOptions { } } -/// Metrics collected by a [`PkIndexWriter`]. +/// Metrics collected by a [`SeriesIndexWriter`]. #[derive(Debug, Clone, Default)] -pub struct PkIndexWriterMetrics { +pub struct SeriesIndexWriterMetrics { /// Number of input record batches passed to the writer. pub input_batches: usize, /// Number of logical rows passed to the writer. @@ -146,7 +146,7 @@ pub struct PkIndexWriterMetrics { pub aborted: bool, } -impl PkIndexWriterMetrics { +impl SeriesIndexWriterMetrics { /// Returns time spent by writer-owned work. pub fn total_elapsed(&self) -> Duration { self.open_elapsed @@ -158,7 +158,7 @@ impl PkIndexWriterMetrics { } #[derive(Debug)] -struct PkIndexRow { +struct SeriesIndexRow { min_ts: i64, max_ts: i64, row_count: u64, @@ -167,8 +167,8 @@ struct PkIndexRow { tags: Vec>, } -/// Incrementally aggregates sorted flat record batches and writes `pk_columns` Parquet. -pub struct PkIndexWriter { +/// Incrementally aggregates sorted flat record batches into series-index Parquet files. +pub struct SeriesIndexWriter { codec: Arc, tag_columns: Vec<(ColumnId, String)>, schema: SchemaRef, @@ -176,13 +176,13 @@ pub struct PkIndexWriter { file_name: String, writer: Option, current_primary_key: Option>, - current_row: Option, - buffered_rows: Vec, - metrics: PkIndexWriterMetrics, + current_row: Option, + buffered_rows: Vec, + metrics: SeriesIndexWriterMetrics, failed: bool, } -impl PkIndexWriter { +impl SeriesIndexWriter { /// Creates a writer for `path` in `object_store`. /// /// The file name in `path` must be unique in the object store so aborting @@ -191,16 +191,16 @@ impl PkIndexWriter { metadata: RegionMetadataRef, object_store: ObjectStore, path: &str, - options: PkIndexWriterOptions, + options: SeriesIndexWriterOptions, ) -> Result { let open_start = Instant::now(); ensure!( options.row_group_size > 0, InvalidMetaSnafu { - reason: "primary-key index row group size must be greater than zero", + reason: "series index row group size must be greater than zero", } ); - let schema = pk_columns_schema(&metadata)?; + let schema = series_index_schema(&metadata)?; let tag_columns = tag_columns(&metadata); let file_name = path.rsplit('/').next().unwrap_or(path).to_string(); let output = object_store @@ -231,7 +231,7 @@ impl PkIndexWriter { current_primary_key: None, current_row: None, buffered_rows: Vec::with_capacity(WRITE_BATCH_SIZE), - metrics: PkIndexWriterMetrics { + metrics: SeriesIndexWriterMetrics { open_elapsed: open_start.elapsed(), ..Default::default() }, @@ -240,7 +240,7 @@ impl PkIndexWriter { } /// Returns the metrics collected so far. - pub fn metrics(&self) -> &PkIndexWriterMetrics { + pub fn metrics(&self) -> &SeriesIndexWriterMetrics { &self.metrics } @@ -251,7 +251,7 @@ impl PkIndexWriter { ensure!( !self.failed, InvalidRecordBatchSnafu { - reason: "cannot write to a failed primary-key index writer", + reason: "cannot write to a failed series index writer", } ); @@ -269,10 +269,10 @@ impl PkIndexWriter { } /// Finishes and commits the index file. - pub async fn finish(mut self) -> Result { + pub async fn finish(mut self) -> Result { if self.failed { let error = InvalidRecordBatchSnafu { - reason: "cannot finish a failed primary-key index writer", + reason: "cannot finish a failed series index writer", } .build(); self.cleanup().await; @@ -287,7 +287,7 @@ impl PkIndexWriter { } /// Aborts the writer and removes incomplete output files. - pub async fn abort(mut self) -> Result { + pub async fn abort(mut self) -> Result { self.metrics.aborted = true; self.cleanup().await; Ok(self.metrics) @@ -301,7 +301,7 @@ impl PkIndexWriter { batch.num_columns() >= 4, InvalidRecordBatchSnafu { reason: format!( - "primary-key index input has too few columns: {}", + "series index input has too few columns: {}", batch.num_columns() ), } @@ -322,7 +322,7 @@ impl PkIndexWriter { ensure!( array.null_count() == 0, InvalidRecordBatchSnafu { - reason: "primary-key index input contains null primary keys", + reason: "series index input contains null primary keys", } ); self.write_binary_primary_keys(array, timestamps.values()) @@ -334,7 +334,7 @@ impl PkIndexWriter { ensure!( array.null_count() == 0, InvalidRecordBatchSnafu { - reason: "primary-key index input contains null primary keys", + reason: "series index input contains null primary keys", } ); self.write_dictionary_primary_keys(array, timestamps.values()) @@ -342,7 +342,7 @@ impl PkIndexWriter { } else { InvalidRecordBatchSnafu { reason: format!( - "primary-key index requires Binary or Dictionary(UInt32, Binary) primary keys, got {:?}", + "series index requires Binary or Dictionary(UInt32, Binary) primary keys, got {:?}", primary_keys.data_type() ), } @@ -419,13 +419,13 @@ impl PkIndexWriter { match primary_key.cmp(current) { Ordering::Less => { return InvalidRecordBatchSnafu { - reason: "primary-key index input is not sorted by primary key", + reason: "series index input is not sorted by primary key", } .fail(); } Ordering::Equal => { let row = self.current_row.as_mut().context(InvalidRecordBatchSnafu { - reason: "primary-key index aggregation state is incomplete", + reason: "series index aggregation state is incomplete", })?; row.min_ts = row.min_ts.min(min_ts); row.max_ts = row.max_ts.max(max_ts); @@ -471,7 +471,7 @@ impl PkIndexWriter { .writer .as_mut() .context(UnexpectedSnafu { - reason: "primary-key index Parquet writer is closed", + reason: "series index Parquet writer is closed", })? .write(&batch) .await @@ -494,13 +494,13 @@ impl PkIndexWriter { self.writer .as_mut() .context(UnexpectedSnafu { - reason: "primary-key index Parquet writer is closed", + reason: "series index Parquet writer is closed", })? .finish() .await .context(WriteParquetSnafu)?; let writer = self.writer.take().context(UnexpectedSnafu { - reason: "primary-key index Parquet writer is closed", + reason: "series index Parquet writer is closed", })?; self.metrics.output_bytes = writer.into_inner().output_bytes(); self.metrics.finish_elapsed += finish_start.elapsed(); @@ -512,7 +512,7 @@ impl PkIndexWriter { if let Some(writer) = self.writer.take() { let mut writer = writer.into_inner().into_inner(); if let Err(error) = writer.abort().await { - common_telemetry::warn!(error; "Failed to abort primary-key index writer"); + common_telemetry::warn!(error; "Failed to abort series index writer"); } } self.current_primary_key = None; @@ -525,8 +525,8 @@ impl PkIndexWriter { } } -/// Returns the Arrow schema of a `pk_columns` primary-key index. -pub fn pk_columns_schema(metadata: &RegionMetadataRef) -> Result { +/// Returns the Arrow schema of a series index. +pub fn series_index_schema(metadata: &RegionMetadataRef) -> Result { validate_metadata(metadata)?; let mut fields = vec![ Field::new(MIN_TS_COLUMN, DataType::Int64, false), @@ -552,7 +552,7 @@ fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> { ), InvalidMetaSnafu { reason: format!( - "primary-key index internal column name {} is already in use", + "series index internal column name {} is already in use", column.column_schema.name ), } @@ -561,7 +561,7 @@ fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> { ensure!( metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse, InvalidMetaSnafu { - reason: "primary-key index only supports sparse primary-key encoding", + reason: "series index only supports sparse primary-key encoding", } ); ensure!( @@ -569,24 +569,24 @@ fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> { .primary_key .starts_with(&[ReservedColumnId::table_id(), ReservedColumnId::tsid()]), InvalidMetaSnafu { - reason: "primary-key index requires (__table_id, __tsid) as the primary-key prefix", + reason: "series index requires (__table_id, __tsid) as the primary-key prefix", } ); let table_id = metadata .column_by_id(ReservedColumnId::table_id()) .context(InvalidMetaSnafu { - reason: "primary-key index metadata is missing __table_id", + reason: "series index metadata is missing __table_id", })?; let tsid = metadata .column_by_id(ReservedColumnId::tsid()) .context(InvalidMetaSnafu { - reason: "primary-key index metadata is missing __tsid", + reason: "series index metadata is missing __tsid", })?; ensure!( table_id.column_schema.data_type == ConcreteDataType::uint32_datatype() && tsid.column_schema.data_type == ConcreteDataType::uint64_datatype(), InvalidMetaSnafu { - reason: "primary-key index requires UInt32 __table_id and UInt64 __tsid", + reason: "series index requires UInt32 __table_id and UInt64 __tsid", } ); for column in metadata.primary_key_columns() { @@ -597,7 +597,7 @@ fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> { column.column_schema.data_type == ConcreteDataType::string_datatype(), InvalidMetaSnafu { reason: format!( - "primary-key index requires string tag column {}, got {}", + "series index requires string tag column {}, got {}", column.column_schema.name, column.column_schema.data_type ), } @@ -626,7 +626,7 @@ fn timestamp_values(array: &ArrayRef) -> Result { .map(|(array, _)| array) .with_context(|| InvalidRecordBatchSnafu { reason: format!( - "primary-key index requires an Int64 or timestamp time index, got {:?}", + "series index requires an Int64 or timestamp time index, got {:?}", array.data_type() ), })? @@ -634,7 +634,7 @@ fn timestamp_values(array: &ArrayRef) -> Result { ensure!( timestamps.null_count() == 0, InvalidRecordBatchSnafu { - reason: "primary-key index input contains null timestamps", + reason: "series index input contains null timestamps", } ); Ok(timestamps) @@ -648,7 +648,7 @@ fn decode_primary_key( max_ts: i64, row_count: u64, tag_columns: &[(ColumnId, String)], -) -> Result { +) -> Result { let CompositeValues::Sparse(values) = codec.decode(primary_key).context(DecodeSnafu)? else { return InvalidRecordBatchSnafu { reason: "decoded primary key is not sparse", @@ -686,7 +686,7 @@ fn decode_primary_key( .fail(), }) .collect::>>()?; - Ok(PkIndexRow { + Ok(SeriesIndexRow { min_ts, max_ts, row_count, @@ -696,7 +696,7 @@ fn decode_primary_key( }) } -fn rows_to_batch(schema: &SchemaRef, rows: &[PkIndexRow]) -> Result { +fn rows_to_batch(schema: &SchemaRef, rows: &[SeriesIndexRow]) -> Result { let mut arrays: Vec = vec![ Arc::new(Int64Array::from_iter_values( rows.iter().map(|row| row.min_ts), @@ -805,11 +805,11 @@ mod tests { } #[test] - fn test_pk_columns_schema() { + fn test_series_index_schema() { let metadata = Arc::new(sst_region_metadata_with_encoding( PrimaryKeyEncoding::Sparse, )); - let schema = pk_columns_schema(&metadata).unwrap(); + let schema = series_index_schema(&metadata).unwrap(); assert_eq!( schema .fields() @@ -817,9 +817,9 @@ mod tests { .map(|field| field.name().as_str()) .collect::>(), [ - "__pk_min_ts", - "__pk_max_ts", - "__pk_row_count", + "__series_min_ts", + "__series_max_ts", + "__series_row_count", "__table_id", "__tsid", "tag_0", @@ -830,11 +830,11 @@ mod tests { assert!(schema.field(5).is_nullable()); let dense = Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense)); - assert!(pk_columns_schema(&dense).is_err()); + assert!(series_index_schema(&dense).is_err()); } #[tokio::test] - async fn test_reject_pk_index_internal_column_names_before_opening_writer() { + async fn test_reject_series_index_internal_column_names_before_opening_writer() { for (index, name) in [MIN_TS_COLUMN, MAX_TS_COLUMN, ROW_COUNT_COLUMN] .into_iter() .enumerate() @@ -850,11 +850,11 @@ mod tests { let metadata = Arc::new(builder.build().unwrap()); let store = object_store(); let path = format!("collision-{index}.parquet"); - let error = PkIndexWriter::try_new( + let error = SeriesIndexWriter::try_new( metadata, store.clone(), &path, - PkIndexWriterOptions::default(), + SeriesIndexWriterOptions::default(), ) .await .err() @@ -901,11 +901,11 @@ mod tests { let primary_key_1 = new_sparse_primary_key(&["a", "x"], &metadata, 1, 10); let primary_key_2 = new_sparse_primary_key(&["b", "y"], &metadata, 1, 20); let store = object_store(); - let mut writer = PkIndexWriter::try_new( + let mut writer = SeriesIndexWriter::try_new( metadata, store.clone(), - "pk.parquet", - PkIndexWriterOptions { row_group_size: 2 }, + "series.parquet", + SeriesIndexWriterOptions { row_group_size: 2 }, ) .await .unwrap(); @@ -946,7 +946,7 @@ mod tests { assert_eq!(metrics.num_series, 2); assert!(!metrics.aborted); - let (output_bytes, row_groups, batches) = read_index(&store, "pk.parquet").await; + let (output_bytes, row_groups, batches) = read_index(&store, "series.parquet").await; assert_eq!(metrics.output_bytes, output_bytes); assert_eq!(row_groups, 1); assert_eq!(batches.len(), 1); @@ -1008,11 +1008,11 @@ mod tests { PrimaryKeyEncoding::Sparse, )); let store = object_store(); - let mut writer = PkIndexWriter::try_new( + let mut writer = SeriesIndexWriter::try_new( metadata.clone(), store.clone(), "groups.parquet", - PkIndexWriterOptions { row_group_size: 2 }, + SeriesIndexWriterOptions { row_group_size: 2 }, ) .await .unwrap(); @@ -1028,11 +1028,11 @@ mod tests { let (_, row_groups, _) = read_index(&store, "groups.parquet").await; assert_eq!(row_groups, 3); - let empty = PkIndexWriter::try_new( + let empty = SeriesIndexWriter::try_new( metadata, store.clone(), "empty.parquet", - PkIndexWriterOptions::default(), + SeriesIndexWriterOptions::default(), ) .await .unwrap() @@ -1054,11 +1054,11 @@ mod tests { let primary_key_1 = new_sparse_primary_key(&["a", "x"], &metadata, 1, 10); let primary_key_2 = new_sparse_primary_key(&["b", "y"], &metadata, 1, 20); let store = object_store(); - let mut writer = PkIndexWriter::try_new( + let mut writer = SeriesIndexWriter::try_new( metadata.clone(), store.clone(), "abort.parquet", - PkIndexWriterOptions { row_group_size: 1 }, + SeriesIndexWriterOptions { row_group_size: 1 }, ) .await .unwrap(); @@ -1082,11 +1082,11 @@ mod tests { Bytes::from_static(b"existing") ); - let mut writer = PkIndexWriter::try_new( + let mut writer = SeriesIndexWriter::try_new( metadata, store.clone(), "dictionary-abort.parquet", - PkIndexWriterOptions { row_group_size: 1 }, + SeriesIndexWriterOptions { row_group_size: 1 }, ) .await .unwrap(); @@ -1137,11 +1137,11 @@ mod tests { ) .unwrap(); let store = object_store(); - let mut writer = PkIndexWriter::try_new( + let mut writer = SeriesIndexWriter::try_new( metadata.clone(), store.clone(), "nullable.parquet", - PkIndexWriterOptions::default(), + SeriesIndexWriterOptions::default(), ) .await .unwrap(); @@ -1159,11 +1159,11 @@ mod tests { assert!(tag_1.is_null(0)); assert!( - PkIndexWriter::try_new( + SeriesIndexWriter::try_new( metadata, store, "invalid.parquet", - PkIndexWriterOptions { row_group_size: 0 }, + SeriesIndexWriterOptions { row_group_size: 0 }, ) .await .is_err() From 09c0b23a23450d0aa3bb2d89938ad0f3f5be253c Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 17 Aug 2026 09:53:59 +0000 Subject: [PATCH 080/126] feat: manage semantic table options via ALTER TABLE SET/UNSET (#8880) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(meta): actually acquire logical table locks in alter-logical-tables procedure The procedure listed its logical table locks from table_info_values, which is only filled during Prepare, while procedure lock keys are fixed at submission — so the logical locks were never acquired. Today every writer of a logical table's info is serialized by the physical table lock, which hides the problem; a metadata-only alter procedure targeting a single logical table would race it. Resolve the logical table ids at submission, persist them in the procedure state (serde(default): state dumped by older versions keeps the previous behavior), lock physical + logical tables, and re-check the resolved ids against the locked set at Prepare so a table dropped and recreated after submission cannot be mutated without a lock. Signed-off-by: Dennis Zhuang * feat: manage semantic table options via ALTER TABLE SET/UNSET CREATE TABLE accepts greptime.semantic.* options, but ALTER TABLE SET routed every option through SetRegionOption, whose closed match rejects them — tables auto-created by ingestion could never receive semantic declarations after the fact. Semantic options are pure metadata markers no region consumes, so they now take a metadata-only alter, following the repartition-hint precedent: - New AlterKind::SetAnnotations/UnsetAnnotations carrying an AnnotationFamily (currently only Semantic), so future marker-style option families reuse the same machinery. The converter classifies a SET/UNSET batch by key prefix and rejects batches that mix annotation keys with regular options. - The procedure reuses the MetadataOnly flow: no region dispatch, table-info update plus cache invalidation only. - Validation lives in the table-meta mutation layer, so it runs at frontend verification and again in the procedure's prepare step under the table lock: SET is strict (known key, value domain, entity columns exist and render as strings); UNSET is lenient inside the namespace so stale keys can be cleaned up. ModifyColumnTypes re-checks columns referenced by entity declarations at the same layer, closing a verify-then-execute race. - Logical metric tables are supported: an annotation alter submits a regular alter-table task locking only the logical table, and the DDL manager's physical-route guard admits it. - create_table_info re-checks semantic value domains for gRPC-built expressions that bypass the SQL parser. Signed-off-by: Dennis Zhuang * refactor(table): centralize annotation option classification and validation Address review feedback on the AnnotationFamily abstraction: with only one variant that every consumer immediately destructured, the generality was fake. Make it real and exhaustive instead: - AnnotationFamily gains RepartitionHint: repartition.column.hint is the same kind of marker option (pure metadata, no region consumes it) and previously had a hand-rolled special case in the converter, the metadata-only classifier, and a dedicated AlterKind pair — all deleted, one classification API remains. Per-family logical-table eligibility (allows_logical_tables) replaces the hard-coded Semantic check in the DDL manager guard. - One validation core in the table crate (check_annotation) serves both DDL entry points. CREATE and ALTER previously duplicated the rules; each keeps its existing error variants, status codes and messages via thin adapters over a typed error (ALTER missing column stays 4002 TableColumnNotFound, CREATE stays InvalidArguments). - The batch classifier returns Result instead of swallowing the mixed-batch error: a mixed SET on a logical table now reports the actual problem instead of UnexpectedLogicalRouteTable, and the flow classifiers propagate instead of guessing. The converter also moves its owned payloads instead of cloning them. Signed-off-by: Dennis Zhuang * test(meta): cover logical-table annotation alter routing The route-guard branch admitting metadata-only annotation alters on logical tables was only exercised end to end by sqlness. Pin it at the DDL manager level: a semantic SET on a logical table succeeds, updates only the logical table's metadata and dispatches nothing to datanodes; a mixed batch reports its own error instead of the route guard's; the repartition hint stays rejected on logical routes. Signed-off-by: Dennis Zhuang * fix(table): keep entity guard on ADD COLUMN and report missing columns first Review follow-ups: the old verify_alter loop scanned the post-alter schema, so it also caught DROP COLUMN followed by re-adding the declared column with a non-string type — the mutation-layer move only kept the MODIFY path. Guard add_columns the same way (this also covers ingestion auto-alter). And run the MODIFY drift check after the existence lookup, so altering a dropped-but-still-declared column reports ColumnNotExists (4002) like every other MODIFY on a missing column. Signed-off-by: Dennis Zhuang * style(grpc-expr): drop a test comment restating the classifier doc Signed-off-by: Dennis Zhuang * refactor(table): rename annotation validation helpers per review check_annotation* validated and normalized; align the names with the validate_and_normalize_* convention nearby, and spell out AnnotationContext (Cx is not used in this repo). Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/grpc-expr/src/alter.rs | 256 +++++++-- src/common/grpc-expr/src/lib.rs | 2 +- src/common/meta/AGENTS.md | 5 + .../meta/src/ddl/alter_logical_tables.rs | 57 +- src/common/meta/src/ddl/alter_table.rs | 42 +- .../meta/src/ddl/alter_table/executor.rs | 4 +- .../src/ddl/tests/alter_logical_tables.rs | 103 +++- src/common/meta/src/ddl/tests/alter_table.rs | 59 +- src/common/meta/src/ddl/tests/event/table.rs | 1 + src/common/meta/src/ddl_manager.rs | 162 +++++- src/operator/src/statement/ddl.rs | 209 +++---- src/table/src/metadata.rs | 510 ++++++++++++++++-- src/table/src/requests.rs | 217 +++++++- src/table/src/requests/semantic.rs | 26 +- .../information_schema/table_semantics.result | 118 ++++ .../information_schema/table_semantics.sql | 62 +++ .../common/system/semantic_graph.result | 55 ++ .../common/system/semantic_graph.sql | 29 + 18 files changed, 1630 insertions(+), 287 deletions(-) diff --git a/src/common/grpc-expr/src/alter.rs b/src/common/grpc-expr/src/alter.rs index 46de85da4c7..85de11d5cdb 100644 --- a/src/common/grpc-expr/src/alter.rs +++ b/src/common/grpc-expr/src/alter.rs @@ -29,8 +29,8 @@ use snafu::{OptionExt, ResultExt, ensure}; use store_api::region_request::{SetRegionOption, UnsetRegionOption}; use table::metadata::{TableId, TableMeta}; use table::requests::{ - AddColumnRequest, AlterKind, AlterTableRequest, ModifyColumnTypeRequest, - REPARTITION_COLUMN_HINT_KEY, SetDefaultRequest, SetIndexOption, UnsetIndexOption, + AddColumnRequest, AlterKind, AlterTableRequest, AnnotationFamily, ModifyColumnTypeRequest, + SetDefaultRequest, SetIndexOption, UnsetIndexOption, }; use crate::error::{ @@ -44,6 +44,57 @@ use crate::error::{ const LOCATION_TYPE_FIRST: i32 = LocationType::First as i32; const LOCATION_TYPE_AFTER: i32 = LocationType::After as i32; +/// Classifies a SET/UNSET key batch: `Ok(Some(family))` when every key belongs +/// to the same annotation family, `Ok(None)` when none does, and an error on a +/// mixed batch — annotation alters skip region dispatch, so they cannot share +/// a statement with options that regions must see. +fn annotation_family_of_keys<'a>( + mut keys: impl Iterator, +) -> Result> { + let Some(first) = keys.next() else { + return Ok(None); + }; + let family = AnnotationFamily::of_key(first); + let mut count = 1usize; + for key in keys { + count += 1; + let this = AnnotationFamily::of_key(key); + if this != family { + return error::InvalidTableOptionRequestSnafu { + err_msg: family.or(this).unwrap().mixed_batch_error(), + } + .fail(); + } + } + if let Some(family) = family + && family.requires_single_key() + && count > 1 + { + return error::InvalidTableOptionRequestSnafu { + err_msg: family.mixed_batch_error(), + } + .fail(); + } + Ok(family) +} + +/// Returns the annotation family when `kind` is a SET/UNSET whose keys all +/// belong to one family — the alters that only rewrite table metadata and skip +/// region dispatch. A mixed batch is an error; never interpret it as "not an +/// annotation alter", or the batch falls through to a path that reports a +/// misleading error (or dispatches to regions). +pub fn annotation_alter_family(kind: &Kind) -> Result> { + match kind { + Kind::SetTableOptions(api::v1::SetTableOptions { table_options }) => { + annotation_family_of_keys(table_options.iter().map(|option| option.key.as_str())) + } + Kind::UnsetTableOptions(api::v1::UnsetTableOptions { keys }) => { + annotation_family_of_keys(keys.iter().map(|key| key.as_str())) + } + _ => Ok(None), + } +} + fn set_index_option_from_proto(set_index: api::v1::SetIndex) -> Result { let options = set_index.options.context(MissingAlterIndexOptionSnafu)?; Ok(match options { @@ -164,20 +215,15 @@ pub fn alter_expr_to_request( AlterKind::RenameTable { new_table_name } } Kind::SetTableOptions(api::v1::SetTableOptions { table_options }) => { - let repartition_column_hint = table_options - .iter() - .find(|option| option.key == REPARTITION_COLUMN_HINT_KEY); - if let Some(option) = repartition_column_hint { - ensure!( - table_options.len() == 1, - error::InvalidTableOptionRequestSnafu { - err_msg: format!( - "{REPARTITION_COLUMN_HINT_KEY} must be altered separately" - ), - } - ); - AlterKind::SetRepartitionColumnHint { - column_name: option.value.clone(), + if let Some(family) = annotation_family_of_keys( + table_options.iter().map(|option| option.key.as_str()), + )? { + AlterKind::SetAnnotations { + family, + options: table_options + .into_iter() + .map(|option| (option.key, option.value)) + .collect(), } } else { AlterKind::SetTableOptions { @@ -190,19 +236,8 @@ pub fn alter_expr_to_request( } } Kind::UnsetTableOptions(api::v1::UnsetTableOptions { keys }) => { - let unset_repartition_column_hint = keys - .iter() - .any(|key| key.as_str() == REPARTITION_COLUMN_HINT_KEY); - if unset_repartition_column_hint { - ensure!( - keys.len() == 1, - error::InvalidTableOptionRequestSnafu { - err_msg: format!( - "{REPARTITION_COLUMN_HINT_KEY} must be altered separately" - ), - } - ); - AlterKind::UnsetRepartitionColumnHint + if let Some(family) = annotation_family_of_keys(keys.iter().map(|key| key.as_str()))? { + AlterKind::UnsetAnnotations { family, keys } } else { AlterKind::UnsetTableOptions { keys: keys @@ -359,6 +394,7 @@ mod tests { Option as PbOption, SemanticType, SetTableOptions, UnsetTableOptions, }; use datatypes::prelude::ConcreteDataType; + use table::requests::REPARTITION_COLUMN_HINT_KEY; use super::*; @@ -563,8 +599,12 @@ mod tests { let alter_request = alter_expr_to_request(1, expr, None).unwrap(); match alter_request.alter_kind { - AlterKind::SetRepartitionColumnHint { column_name } => { - assert_eq!("host", column_name); + AlterKind::SetAnnotations { family, options } => { + assert_eq!(AnnotationFamily::RepartitionHint, family); + assert_eq!( + vec![(REPARTITION_COLUMN_HINT_KEY.to_string(), "host".to_string())], + options + ); } _ => unreachable!(), } @@ -595,6 +635,30 @@ mod tests { err.to_string() .contains("repartition.column.hint must be altered separately") ); + + // Duplicate hint entries are not a meaningful batch either. + let dup = AlterTableExpr { + catalog_name: "test_catalog".to_string(), + schema_name: "test_schema".to_string(), + table_name: "monitor".to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![ + PbOption { + key: REPARTITION_COLUMN_HINT_KEY.to_string(), + value: "host".to_string(), + }, + PbOption { + key: REPARTITION_COLUMN_HINT_KEY.to_string(), + value: "region".to_string(), + }, + ], + })), + }; + let err = alter_expr_to_request(1, dup, None).unwrap_err(); + assert!( + err.to_string() + .contains("repartition.column.hint must be altered separately") + ); } #[test] @@ -608,10 +672,138 @@ mod tests { })), }; + let alter_request = alter_expr_to_request(1, expr, None).unwrap(); + match alter_request.alter_kind { + AlterKind::UnsetAnnotations { family, keys } => { + assert_eq!(AnnotationFamily::RepartitionHint, family); + assert_eq!(vec![REPARTITION_COLUMN_HINT_KEY.to_string()], keys); + } + _ => unreachable!(), + } + } + + #[test] + fn test_semantic_options_classified_as_annotations() { + let expr = AlterTableExpr { + catalog_name: "test_catalog".to_string(), + schema_name: "test_schema".to_string(), + table_name: "monitor".to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![ + PbOption { + key: "greptime.semantic.signal_type".to_string(), + value: "trace".to_string(), + }, + PbOption { + key: "greptime.semantic.entity.host.id".to_string(), + value: "host".to_string(), + }, + ], + })), + }; + + let alter_request = alter_expr_to_request(1, expr, None).unwrap(); + let AlterKind::SetAnnotations { family, options } = alter_request.alter_kind else { + panic!( + "expected SetAnnotations, got {:?}", + alter_request.alter_kind + ); + }; + assert_eq!(family, AnnotationFamily::Semantic); + assert_eq!(options.len(), 2); + + let expr = AlterTableExpr { + catalog_name: "test_catalog".to_string(), + schema_name: "test_schema".to_string(), + table_name: "monitor".to_string(), + kind: Some(Kind::UnsetTableOptions(UnsetTableOptions { + keys: vec!["greptime.semantic.signal_type".to_string()], + })), + }; let alter_request = alter_expr_to_request(1, expr, None).unwrap(); assert!(matches!( alter_request.alter_kind, - AlterKind::UnsetRepartitionColumnHint + AlterKind::UnsetAnnotations { + family: AnnotationFamily::Semantic, + .. + } )); } + + #[test] + fn test_semantic_options_reject_mixed_batch() { + let mixed_set = AlterTableExpr { + catalog_name: "test_catalog".to_string(), + schema_name: "test_schema".to_string(), + table_name: "monitor".to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![ + PbOption { + key: "greptime.semantic.signal_type".to_string(), + value: "trace".to_string(), + }, + PbOption { + key: table::requests::TTL_KEY.to_string(), + value: "7d".to_string(), + }, + ], + })), + }; + let err = alter_expr_to_request(1, mixed_set, None).unwrap_err(); + assert!(err.to_string().contains("altered separately"), "{err}"); + + let mixed_unset = AlterTableExpr { + catalog_name: "test_catalog".to_string(), + schema_name: "test_schema".to_string(), + table_name: "monitor".to_string(), + kind: Some(Kind::UnsetTableOptions(UnsetTableOptions { + keys: vec![ + "ttl".to_string(), + "greptime.semantic.signal_type".to_string(), + ], + })), + }; + let err = alter_expr_to_request(1, mixed_unset, None).unwrap_err(); + assert!(err.to_string().contains("altered separately"), "{err}"); + } + + #[test] + fn test_annotation_alter_family() { + let mixed = Kind::SetTableOptions(SetTableOptions { + table_options: vec![ + PbOption { + key: "greptime.semantic.signal_type".to_string(), + value: "trace".to_string(), + }, + PbOption { + key: "ttl".to_string(), + value: "7d".to_string(), + }, + ], + }); + let err = annotation_alter_family(&mixed).unwrap_err(); + assert!( + err.to_string().contains("must be altered separately"), + "{err}" + ); + + // Two annotation families cannot share a batch either. + let cross = Kind::SetTableOptions(SetTableOptions { + table_options: vec![ + PbOption { + key: "greptime.semantic.signal_type".to_string(), + value: "trace".to_string(), + }, + PbOption { + key: REPARTITION_COLUMN_HINT_KEY.to_string(), + value: "host".to_string(), + }, + ], + }); + let err = annotation_alter_family(&cross).unwrap_err(); + assert!( + err.to_string().contains("must be altered separately"), + "{err}" + ); + } } diff --git a/src/common/grpc-expr/src/lib.rs b/src/common/grpc-expr/src/lib.rs index 0941cb650d0..9b9994a8f02 100644 --- a/src/common/grpc-expr/src/lib.rs +++ b/src/common/grpc-expr/src/lib.rs @@ -16,4 +16,4 @@ mod alter; pub mod error; pub mod util; -pub use alter::{alter_expr_to_request, create_table_schema}; +pub use alter::{alter_expr_to_request, annotation_alter_family, create_table_schema}; diff --git a/src/common/meta/AGENTS.md b/src/common/meta/AGENTS.md index fd5da08a933..3684f8f9db3 100644 --- a/src/common/meta/AGENTS.md +++ b/src/common/meta/AGENTS.md @@ -26,6 +26,11 @@ and durable DDL procedures. Metasrv process and service wiring live in invalidation. - Procedure changes must preserve persisted state and `TYPE_NAME`; new procedures need loader registration in `src/common/meta/src/ddl_manager.rs`. +- Marker-style table options that no region consumes (`greptime.semantic.*`, + `repartition.column.hint`) go through `AnnotationFamily` in + `src/table/src/requests.rs` — classification, validation, logical-table + eligibility, and the metadata-only alter flow come with it. Do not add + per-option special cases to the alter conversion or DDL routing. - `KvBackend` behavior changes should extend the shared tests in `src/common/meta/src/kv_backend/test.rs`. diff --git a/src/common/meta/src/ddl/alter_logical_tables.rs b/src/common/meta/src/ddl/alter_logical_tables.rs index 9207f910577..ccf130cec8c 100644 --- a/src/common/meta/src/ddl/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/alter_logical_tables.rs @@ -24,7 +24,7 @@ use common_procedure::{Context, EventContext, EventTrigger, LockKey, Procedure, use common_telemetry::{debug, error, info, warn}; pub use executor::make_alter_region_request; use serde::{Deserialize, Serialize}; -use snafu::ResultExt; +use snafu::{ResultExt, ensure}; use store_api::metadata::ColumnMetadata; use store_api::metric_engine_consts::ALTER_PHYSICAL_EXTENSION_KEY; use strum::AsRefStr; @@ -40,7 +40,7 @@ use crate::ddl::event::table::{ TableDdlEvent, TableDdlEventType, TableDdlLocator, alter_table_kind_name, }; use crate::ddl::utils::{extract_column_metadatas, map_to_procedure_error, sync_follower_regions}; -use crate::error::Result; +use crate::error::{self, Result}; use crate::instruction::CacheIdent; use crate::key::DeserializedValueWithBytes; use crate::key::table_info::TableInfoValue; @@ -88,6 +88,7 @@ impl AlterLogicalTablesProcedure { pub fn new( tasks: Vec, physical_table_id: TableId, + logical_table_ids: Vec, context: DdlContext, ) -> Self { Self { @@ -97,6 +98,7 @@ impl AlterLogicalTablesProcedure { tasks, table_info_values: vec![], physical_table_id, + logical_table_ids, physical_table_info: None, physical_columns: vec![], table_cache_keys_to_invalidate: vec![], @@ -147,6 +149,26 @@ impl AlterLogicalTablesProcedure { ); } + // Locks are fixed at submission. A logical table that resolves to an + // id outside the locked set here was dropped and recreated in between, + // so this procedure holds no lock for it. Procedures restored from + // pre-upgrade state have no recorded ids and keep the old behavior. + if !self.data.logical_table_ids.is_empty() { + for value in &table_info_values { + let table_id = value.get_inner_ref().table_info.ident.table_id; + ensure!( + self.data.logical_table_ids.contains(&table_id), + error::UnexpectedSnafu { + err_msg: format!( + "logical table {} (id {table_id}) is not covered by the \ + procedure locks; retry the statement", + value.get_inner_ref().table_info.name + ), + } + ); + } + } + // Updates the procedure state. retain_unskipped(&mut self.data.tasks, &skip_alter); self.data.physical_table_info = Some(physical_table_info); @@ -305,17 +327,29 @@ impl Procedure for AlterLogicalTablesProcedure { // CatalogLock, SchemaLock, // TableLock // TableNameLock(s) - let mut lock_key = Vec::with_capacity(2 + 1 + self.data.tasks.len()); + let mut lock_key = Vec::with_capacity(2 + 1 + self.data.logical_table_ids.len()); let table_ref = self.data.tasks[0].table_ref(); lock_key.push(CatalogLock::Read(table_ref.catalog).into()); lock_key.push(SchemaLock::read(table_ref.catalog, table_ref.schema).into()); lock_key.push(TableLock::Write(self.data.physical_table_id).into()); - lock_key.extend( - self.data - .table_info_values - .iter() - .map(|table| TableLock::Write(table.table_info.ident.table_id).into()), - ); + if self.data.logical_table_ids.is_empty() { + // Pre-upgrade procedure state has no `logical_table_ids`, but a + // dump taken after `Prepare` still carries the resolved table + // snapshots — recover the logical locks from them. + lock_key.extend( + self.data + .table_info_values + .iter() + .map(|table| TableLock::Write(table.table_info.ident.table_id).into()), + ); + } else { + lock_key.extend( + self.data + .logical_table_ids + .iter() + .map(|table_id| TableLock::Write(*table_id).into()), + ); + } LockKey::new(lock_key) } @@ -366,6 +400,11 @@ pub struct AlterTablesData { table_info_values: Vec>, /// Physical table info physical_table_id: TableId, + /// Logical table ids resolved at submission time, so `lock_key` can name + /// them before `Prepare` runs (procedure locks are fixed at submission). + /// Empty when restored from pre-upgrade procedure state. + #[serde(default)] + logical_table_ids: Vec, physical_table_info: Option>, physical_columns: Vec, table_cache_keys_to_invalidate: Vec, diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index 63342a7bbb6..8e24f3ae2ab 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -21,7 +21,7 @@ use std::vec; use api::region::RegionResponse; use api::v1::alter_table_expr::Kind; -use api::v1::{RenameTable, SetTableOptions, UnsetTableOptions}; +use api::v1::{RenameTable, SetTableOptions}; use async_trait::async_trait; use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE}; use common_error::ext::BoxedError; @@ -38,7 +38,7 @@ use store_api::metric_engine_consts::TABLE_COLUMN_METADATA_EXTENSION_KEY; use store_api::storage::RegionId; use strum::AsRefStr; use table::metadata::{TableId, TableInfo}; -use table::requests::{REPARTITION_COLUMN_HINT_KEY, SKIP_WAL_KEY}; +use table::requests::SKIP_WAL_KEY; use table::table_reference::TableReference; use crate::ddl::DdlContext; @@ -51,7 +51,8 @@ use crate::ddl::utils::{ sync_follower_regions, }; use crate::error::{ - AbortProcedureSnafu, NoLeaderSnafu, PutPoisonSnafu, Result, RetryLaterSnafu, UnsupportedSnafu, + AbortProcedureSnafu, ConvertAlterTableRequestSnafu, NoLeaderSnafu, PutPoisonSnafu, Result, + RetryLaterSnafu, UnsupportedSnafu, }; use crate::key::table_info::TableInfoValue; use crate::key::{DeserializedValueWithBytes, RegionDistribution}; @@ -201,7 +202,7 @@ impl AlterTableProcedure { self.data.region_distribution = Some(region_distribution(&physical_table_route.region_routes)); } - self.data.state = self.data.flow().after_prepare(); + self.data.state = self.data.flow()?.after_prepare(); Ok(Status::executing(true)) } @@ -249,7 +250,7 @@ impl AlterTableProcedure { ensure!(!leaders.is_empty(), NoLeaderSnafu { table_id }); // Puts the poison before submitting alter region requests to datanodes. self.put_poison(ctx_provider, procedure_id).await?; - let flow = self.data.flow(); + let flow = self.data.flow()?; if flow == AlterTableFlow::MetadataFirst { let results = self .executor @@ -340,7 +341,7 @@ impl AlterTableProcedure { "altering table result doesn't contains extension key `{TABLE_COLUMN_METADATA_EXTENSION_KEY}`,leaving the table's column metadata unchanged" ); } - self.data.state = self.data.flow().after_regions(); + self.data.state = self.data.flow()?.after_regions(); Ok(()) } @@ -372,7 +373,7 @@ impl AlterTableProcedure { let table_info_value = self.data.table_info_value.as_ref().unwrap(); // Safety: Checked in `AlterTableProcedure::new`. let alter_kind = self.data.task.alter_table.kind.as_ref().unwrap(); - let flow = self.data.flow(); + let flow = self.data.flow()?; let metadata_only_alter = flow == AlterTableFlow::MetadataOnly; // Gets the table info from the cache or builds it. @@ -469,17 +470,12 @@ pub(crate) fn only_enables_skip_wal(alter_kind: &Kind) -> bool { && table_options[0].value == "true" } -fn is_metadata_only_alter(alter_kind: &Kind) -> bool { - match alter_kind { - Kind::RenameTable { .. } => true, - Kind::SetTableOptions(SetTableOptions { table_options }) => { - table_options.len() == 1 && table_options[0].key.as_str() == REPARTITION_COLUMN_HINT_KEY - } - Kind::UnsetTableOptions(UnsetTableOptions { keys }) => { - keys.len() == 1 && keys[0].as_str() == REPARTITION_COLUMN_HINT_KEY - } - _ => false, - } +fn is_metadata_only_alter(alter_kind: &Kind) -> Result { + // A mixed annotation batch is an error, never "not metadata-only": falling + // through to the region-first flow would dispatch it to regions. + let family = common_grpc_expr::annotation_alter_family(alter_kind) + .context(ConvertAlterTableRequestSnafu)?; + Ok(family.is_some() || matches!(alter_kind, Kind::RenameTable { .. })) } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -490,14 +486,14 @@ enum AlterTableFlow { } impl AlterTableFlow { - fn from_kind(kind: &Kind) -> Self { - if only_enables_skip_wal(kind) { + fn from_kind(kind: &Kind) -> Result { + Ok(if only_enables_skip_wal(kind) { Self::MetadataFirst - } else if is_metadata_only_alter(kind) { + } else if is_metadata_only_alter(kind)? { Self::MetadataOnly } else { Self::RegionFirst - } + }) } fn after_prepare(self) -> AlterTableState { @@ -647,7 +643,7 @@ impl AlterTableData { .map(|value| &value.table_info) } - fn flow(&self) -> AlterTableFlow { + fn flow(&self) -> Result { // Safety: Checked in `AlterTableProcedure::new`. AlterTableFlow::from_kind(self.task.alter_table.kind.as_ref().unwrap()) } diff --git a/src/common/meta/src/ddl/alter_table/executor.rs b/src/common/meta/src/ddl/alter_table/executor.rs index fbacdd1b2f5..65fa45d0a8f 100644 --- a/src/common/meta/src/ddl/alter_table/executor.rs +++ b/src/common/meta/src/ddl/alter_table/executor.rs @@ -335,8 +335,8 @@ fn build_new_table_info( | AlterKind::ModifyColumnTypes { .. } | AlterKind::SetTableOptions { .. } | AlterKind::UnsetTableOptions { .. } - | AlterKind::SetRepartitionColumnHint { .. } - | AlterKind::UnsetRepartitionColumnHint + | AlterKind::SetAnnotations { .. } + | AlterKind::UnsetAnnotations { .. } | AlterKind::SetIndexes { .. } | AlterKind::UnsetIndexes { .. } | AlterKind::DropDefaults { .. } diff --git a/src/common/meta/src/ddl/tests/alter_logical_tables.rs b/src/common/meta/src/ddl/tests/alter_logical_tables.rs index 351c0fd96ec..619e02bb9b9 100644 --- a/src/common/meta/src/ddl/tests/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/tests/alter_logical_tables.rs @@ -43,6 +43,7 @@ use crate::error::Error::{AlterLogicalTablesInvalidArguments, TableNotFound}; use crate::error::Result; use crate::key::table_name::TableNameKey; use crate::key::table_route::{PhysicalTableRouteValue, TableRouteValue}; +use crate::lock_key::TableLock; use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -162,7 +163,8 @@ async fn test_on_prepare_check_schema() { ), ]; let physical_table_id = 1024u32; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, physical_table_id, vec![], ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -177,7 +179,8 @@ async fn test_on_prepare_check_alter_kind() { "new_table1", )]; let physical_table_id = 1024u32; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, physical_table_id, vec![], ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -197,7 +200,7 @@ async fn test_on_prepare_different_physical_table() { make_alter_logical_table_add_column_task(None, "table2", vec!["column2".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy1_id, ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy1_id, vec![], ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -218,7 +221,7 @@ async fn test_on_prepare_logical_table_not_exists() { make_alter_logical_table_add_column_task(None, "table2", vec!["column2".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, TableNotFound { .. }); } @@ -241,7 +244,7 @@ async fn test_on_prepare() { make_alter_logical_table_add_column_task(None, "table3", vec!["column3".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context); let result = procedure.on_prepare().await; assert_matches!( result, @@ -277,7 +280,8 @@ async fn test_on_update_metadata() { make_alter_logical_table_add_column_task(None, "table3", vec!["new_col".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -362,7 +366,8 @@ async fn test_on_part_duplicate_alter_request() { make_alter_logical_table_add_column_task(None, "table2", vec!["col_0".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -448,7 +453,8 @@ async fn test_on_part_duplicate_alter_request() { ), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -620,7 +626,7 @@ async fn test_on_submit_alter_region_request() { make_alter_logical_table_add_column_task(None, "table2", vec!["mew_col".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context); procedure.on_prepare().await.unwrap(); procedure.on_submit_alter_region_requests().await.unwrap(); let mut results = Vec::new(); @@ -646,3 +652,82 @@ async fn test_on_submit_alter_region_request() { }) ); } + +#[test] +fn test_lock_key_covers_submission_resolved_logical_tables() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let ddl_context = new_ddl_context(node_manager); + let tasks = vec![make_alter_logical_table_add_column_task( + None, + "table1", + vec!["column1".to_string()], + )]; + let procedure = AlterLogicalTablesProcedure::new(tasks, 1024, vec![1025, 1026], ddl_context); + let keys = procedure.lock_key(); + let keys = keys.keys_to_lock().collect::>(); + for table_id in [1024u32, 1025, 1026] { + let expected: common_procedure::StringKey = TableLock::Write(table_id).into(); + assert!(keys.contains(&&expected), "missing lock for {table_id}"); + } +} + +#[tokio::test] +async fn test_from_json_without_logical_table_ids() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let ddl_context = new_ddl_context(node_manager); + let phy_id = create_physical_table(&ddl_context, "phy").await; + let logical_id = create_logical_table(ddl_context.clone(), phy_id, "table1").await; + let tasks = vec![make_alter_logical_table_add_column_task( + None, + "table1", + vec!["column1".to_string()], + )]; + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, vec![logical_id], ddl_context.clone()); + procedure.on_prepare().await.unwrap(); + let json = procedure.dump().unwrap(); + + // Pre-upgrade procedure state carries no `logical_table_ids`, but a + // post-Prepare dump still holds the resolved table snapshots. + let mut value: serde_json::Value = serde_json::from_str(&json).unwrap(); + value + .as_object_mut() + .unwrap() + .remove("logical_table_ids") + .unwrap(); + let restored = AlterLogicalTablesProcedure::from_json(&value.to_string(), ddl_context).unwrap(); + + let keys = restored.lock_key(); + let keys = keys.keys_to_lock().collect::>(); + for table_id in [phy_id, logical_id] { + let expected: common_procedure::StringKey = TableLock::Write(table_id).into(); + assert!(keys.contains(&&expected), "missing lock for {table_id}"); + } +} + +#[tokio::test] +async fn test_on_prepare_rejects_table_outside_locked_set() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let ddl_context = new_ddl_context(node_manager); + + let phy_id = create_physical_table(&ddl_context, "phy").await; + create_logical_table(ddl_context.clone(), phy_id, "table1").await; + + let tasks = vec![make_alter_logical_table_add_column_task( + None, + "table1", + vec!["column1".to_string()], + )]; + + // The table resolved at prepare time is not in the locked set, as if it + // had been dropped and recreated after submission. + let stale_locked_id = phy_id + 1000; + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, vec![stale_locked_id], ddl_context); + let err = procedure.on_prepare().await.unwrap_err(); + assert!( + err.to_string() + .contains("not covered by the procedure locks"), + "{err}" + ); +} diff --git a/src/common/meta/src/ddl/tests/alter_table.rs b/src/common/meta/src/ddl/tests/alter_table.rs index b35666e52e1..30910eb0ba5 100644 --- a/src/common/meta/src/ddl/tests/alter_table.rs +++ b/src/common/meta/src/ddl/tests/alter_table.rs @@ -29,7 +29,7 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_procedure::store::poison_store::PoisonStore; use common_procedure::{Procedure, ProcedureId, Status}; -use common_procedure_test::MockContextProvider; +use common_procedure_test::{MockContextProvider, execute_procedure_until_done}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; use store_api::metadata::ColumnMetadata; @@ -1121,3 +1121,60 @@ async fn test_on_submit_alter_request_with_exist_poison() { .unwrap_err(); assert_matches!(err, Error::PutPoison { .. }); } + +#[tokio::test] +async fn test_semantic_annotation_alter_is_metadata_only() { + let (tx, mut rx) = mpsc::channel(8); + let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); + let ddl_context = new_ddl_context(node_manager); + let table_id = 1024; + let table_name = "foo"; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + prepare_table_route(table_id), + HashMap::new(), + ) + .await + .unwrap(); + + let alter_table_task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![api::v1::Option { + key: "greptime.semantic.signal_type".to_string(), + value: "metric".to_string(), + }], + })), + }, + }; + let mut procedure = + AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); + execute_procedure_until_done(&mut procedure).await; + + // Metadata-only: no region request reaches any datanode. + rx.try_recv().unwrap_err(); + + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert_eq!( + table_info + .meta + .options + .extra_options + .get("greptime.semantic.signal_type"), + Some(&"metric".to_string()) + ); +} diff --git a/src/common/meta/src/ddl/tests/event/table.rs b/src/common/meta/src/ddl/tests/event/table.rs index 24182fd19af..92fa10ee9e0 100644 --- a/src/common/meta/src/ddl/tests/event/table.rs +++ b/src/common/meta/src/ddl/tests/event/table.rs @@ -391,6 +391,7 @@ fn procedure_cases() -> Vec { ), ], 43, + vec![], test_context(), ); let drop_table = DropTableProcedure::new( diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 7ba408ca2d9..d88f08918e9 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -53,7 +53,7 @@ use crate::ddl::truncate_table::TruncateTableProcedure; use crate::ddl::undrop_table::UndropTableProcedure; use crate::ddl::{DdlContext, utils}; use crate::error::{ - self, CreateRepartitionProcedureSnafu, EmptyDdlTasksSnafu, + self, ConvertAlterTableRequestSnafu, CreateRepartitionProcedureSnafu, EmptyDdlTasksSnafu, PersistRepartitionGcRequirementSnafu, ProcedureOutputSnafu, RegisterProcedureLoaderSnafu, RegisterRepartitionProcedureLoaderSnafu, Result, SubmitProcedureSnafu, TableInfoNotFoundSnafu, TableNotFoundSnafu, TableRouteNotFoundSnafu, UnexpectedLogicalRouteTableSnafu, @@ -554,8 +554,26 @@ impl DdlManager { ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = - AlterLogicalTablesProcedure::new(alter_table_tasks, physical_table_id, context); + // Resolve the logical table ids up front: procedure locks are fixed + // at submission, so `lock_key` cannot derive them during `Prepare`. + let logical_table_ids = { + let table_refs = alter_table_tasks + .iter() + .map(|task| task.table_ref()) + .collect::>(); + utils::table_id::get_all_table_ids_by_names( + self.table_metadata_manager().table_name_manager(), + &table_refs, + ) + .await? + }; + + let procedure = AlterLogicalTablesProcedure::new( + alter_table_tasks, + physical_table_id, + logical_table_ids, + context, + ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); @@ -1053,8 +1071,17 @@ async fn handle_alter_table_task( .get(table_id) .await? .context(TableRouteNotFoundSnafu { table_id })?; + // Classify before the route guard: a mixed annotation batch must surface + // its own error here, not a misleading "non-physical route" one. Families + // that only rewrite the logical table's metadata may target logical tables. + let annotation_family = match alter_table_task.alter_table.kind.as_ref() { + Some(kind) => common_grpc_expr::annotation_alter_family(kind) + .context(ConvertAlterTableRequestSnafu)?, + None => None, + }; ensure!( - table_route_value.is_physical(), + table_route_value.is_physical() + || annotation_family.is_some_and(|family| family.allows_logical_tables()), UnexpectedLogicalRouteTableSnafu { err_msg: format!("{:?} is a non-physical TableRouteValue.", table_ref), } @@ -1499,7 +1526,9 @@ mod tests { use common_event_recorder::{PersistentEventContext, ProcedureEventInput, TriggerReason}; use common_procedure::local::LocalManager; use common_procedure::test_util::InMemoryPoisonStore; - use common_procedure::{BoxedProcedure, ProcedureContext, ProcedureManagerRef}; + use common_procedure::{ + BoxedProcedure, ProcedureContext, ProcedureManager, ProcedureManagerRef, + }; use store_api::storage::TableId; use table::table_name::TableName; @@ -1900,4 +1929,127 @@ mod tests { assert!(matches!(err, crate::error::Error::Unsupported { .. })); } } + + async fn ddl_manager_with_context(ddl_context: DdlContext) -> DdlManager { + let kv_backend = Arc::new(MemoryKvBackend::new()); + let state_store = Arc::new(KvStateStore::new(kv_backend.clone())); + let poison_manager = Arc::new(InMemoryPoisonStore::default()); + let procedure_manager = Arc::new(LocalManager::new( + Default::default(), + state_store, + poison_manager, + None, + None, + )); + procedure_manager.start().await.unwrap(); + let ddl_manager = DdlManager::new( + ddl_context, + procedure_manager, + Arc::new(DummyRepartitionProcedureFactory), + ); + ddl_manager.register_loaders().unwrap(); + ddl_manager + } + + fn set_options_expr(table_name: &str, options: &[(&str, &str)]) -> api::v1::AlterTableExpr { + api::v1::AlterTableExpr { + catalog_name: common_catalog::consts::DEFAULT_CATALOG_NAME.to_string(), + schema_name: common_catalog::consts::DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(api::v1::alter_table_expr::Kind::SetTableOptions( + api::v1::SetTableOptions { + table_options: options + .iter() + .map(|(key, value)| api::v1::Option { + key: key.to_string(), + value: value.to_string(), + }) + .collect(), + }, + )), + } + } + + #[tokio::test] + async fn test_logical_table_annotation_alter_routing() { + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let node_manager = Arc::new(crate::test_util::MockDatanodeManager::new( + crate::ddl::test_util::datanode_handler::DatanodeWatcher::new(tx), + )); + let ddl_context = crate::test_util::new_ddl_context(node_manager); + let phy_id = crate::ddl::test_util::create_physical_table(&ddl_context, "phy").await; + let logical_id = + crate::ddl::test_util::create_logical_table(ddl_context.clone(), phy_id, "logical") + .await; + let ddl_manager = ddl_manager_with_context(ddl_context.clone()).await; + + // A semantic alter on a logical table passes the route guard, updates + // only the logical table's metadata, and dispatches nothing. + ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr( + "logical", + &[("greptime.semantic.signal_type", "metric")], + ))), + ) + .await + .unwrap(); + rx.try_recv().unwrap_err(); + let table_info = ddl_manager + .table_metadata_manager() + .table_info_manager() + .get(logical_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert_eq!( + table_info + .meta + .options + .extra_options + .get("greptime.semantic.signal_type"), + Some(&"metric".to_string()) + ); + + // A mixed batch fails with its own error, not the route guard's. + let err = ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr( + "logical", + &[("greptime.semantic.source", "prometheus"), ("ttl", "7d")], + ))), + ) + .await + .unwrap_err(); + let msg = common_error::ext::ErrorExt::output_msg(&err); + assert!(msg.contains("must be altered separately"), "{msg}"); + + // The repartition hint drives physical repartitioning; on a logical + // route it stays rejected by the guard. + let err = ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr( + "logical", + &[("repartition.column.hint", "host")], + ))), + ) + .await + .unwrap_err(); + let msg = common_error::ext::ErrorExt::output_msg(&err); + assert!(msg.contains("non-physical TableRouteValue"), "{msg}"); + } } diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 3c10186ca4e..f89f5de73a7 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -98,8 +98,8 @@ use table::TableRef; use table::dist_table::DistTable; use table::metadata::{self, TableId, TableInfo, TableMeta, TableType}; use table::requests::{ - AlterKind, AlterTableRequest, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, REPARTITION_COLUMN_HINT_KEY, - TableOptions, parse_entity_columns, parse_entity_option_key, + AlterKind, AlterTableRequest, AnnotationContext, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, + TableOptions, validate_and_normalize_annotation_options, }; use table::table_name::TableName; use table::table_reference::TableReference; @@ -1881,8 +1881,20 @@ impl StatementExecutor { (req, invalidate_keys) } else { - // This is logical table - let req = SubmitDdlTaskRequest::new(DdlTask::new_alter_logical_tables(vec![expr])); + // This is logical table. Annotation alters only rewrite its own + // metadata; `AlterLogicalTablesProcedure` only handles column adds. + let annotation_alter = match expr.kind.as_ref() { + Some(kind) => common_grpc_expr::annotation_alter_family(kind) + .context(AlterExprToRequestSnafu)? + .is_some_and(|family| family.allows_logical_tables()), + None => false, + }; + let task = if annotation_alter { + DdlTask::new_alter_table(expr) + } else { + DdlTask::new_alter_logical_tables(vec![expr]) + }; + let req = SubmitDdlTaskRequest::new(task); let mut invalidate_keys = vec![ CacheIdent::TableId(physical_table_id), @@ -2338,19 +2350,6 @@ pub fn verify_alter( .context(error::BuildTableMetaSnafu { table_name })?; validate_json2_columns_append_mode(&new_meta.schema, &new_meta.options)?; - // A column referenced by an entity declaration may be dropped (the - // read-time derivation skips the stale declaration), but must not change - // to a type without a stable string form. - for (key, value) in &new_meta.options.extra_options { - if parse_entity_option_key(key).is_none() { - continue; - } - for col in &parse_entity_columns(value) { - if let Some(column) = new_meta.schema.column_schema_by_name(col) { - ensure_entity_column_renders_as_string(key, col, &column.data_type)?; - } - } - } Ok(true) } @@ -2402,14 +2401,7 @@ pub fn create_table_info( validate_json2_columns_append_mode(&schema, &table_options)?; - validate_repartition_column_hint( - &mut table_options, - &column_name_to_index_map, - &partition_key_indices, - &create_table.time_index, - )?; - - validate_entity_semantic_options(&table_options, &schema)?; + validate_and_normalize_annotations(&mut table_options, &schema, &partition_key_indices)?; let meta = TableMeta { schema, @@ -2469,61 +2461,6 @@ fn validate_json2_columns_append_mode(schema: &Schema, table_options: &TableOpti Ok(()) } -fn validate_repartition_column_hint( - table_options: &mut TableOptions, - column_name_to_index_map: &HashMap, - partition_key_indices: &[usize], - time_index: &str, -) -> Result<()> { - let Some(column_name) = table_options - .extra_options - .get(REPARTITION_COLUMN_HINT_KEY) - .map(|value| value.trim().to_string()) - else { - return Ok(()); - }; - - ensure!( - !column_name.is_empty(), - InvalidPartitionRuleSnafu { - reason: format!("{REPARTITION_COLUMN_HINT_KEY} expects exactly one column name"), - } - ); - - ensure!( - !column_name.contains(','), - InvalidPartitionRuleSnafu { - reason: format!("{REPARTITION_COLUMN_HINT_KEY} expects exactly one column name"), - } - ); - - ensure!( - partition_key_indices.is_empty(), - InvalidPartitionRuleSnafu { - reason: format!( - "cannot set {REPARTITION_COLUMN_HINT_KEY} on a table with partition metadata" - ), - } - ); - - column_name_to_index_map - .get(&column_name) - .context(ColumnNotFoundSnafu { msg: &column_name })?; - - ensure!( - column_name != time_index, - InvalidPartitionRuleSnafu { - reason: format!("cannot set {REPARTITION_COLUMN_HINT_KEY} to the time index column"), - } - ); - - table_options - .extra_options - .insert(REPARTITION_COLUMN_HINT_KEY.to_string(), column_name); - - Ok(()) -} - /// Rejects DDL against read-only schemas and computed entity-graph tables. fn ensure_table_writable(schema: &str, table: &str) -> Result<()> { ensure!( @@ -2557,57 +2494,34 @@ fn ensure_table_definition_writable(schema: &str, table: &str) -> Result<()> { Ok(()) } -/// Whether `CAST(column AS Utf8)` yields a stable entity-id string — the -/// binary-backed and nested types do not, and without the DDL check the -/// failure would surface only when the graph is scanned. -fn has_stable_string_form(data_type: &datatypes::prelude::ConcreteDataType) -> bool { - use datatypes::prelude::ConcreteDataType; - !matches!( - data_type, - ConcreteDataType::Binary(_) - | ConcreteDataType::Json(_) - | ConcreteDataType::Vector(_) - | ConcreteDataType::List(_) - | ConcreteDataType::Struct(_) - | ConcreteDataType::Dictionary(_) - | ConcreteDataType::Null(_) - ) -} - -fn ensure_entity_column_renders_as_string( - key: &str, - col: &str, - data_type: &datatypes::prelude::ConcreteDataType, +/// CREATE-side annotation validation: one rule source in the table crate, +/// mapped onto this crate's existing error variants so client-visible codes +/// and messages stay put. +fn validate_and_normalize_annotations( + options: &mut TableOptions, + schema: &Schema, + partition_key_indices: &[usize], ) -> Result<()> { - ensure!( - has_stable_string_form(data_type), - InvalidSqlSnafu { - err_msg: format!( - "entity column `{col}` (option `{key}`) has type `{data_type}`, which cannot \ - render as a string", - ), + use table::requests::AnnotationValidationError as CheckError; + let cx = AnnotationContext { + schema, + partition_key_indices, + }; + validate_and_normalize_annotation_options(options, &cx).map_err(|e| match e { + CheckError::ColumnNotFound { column } => ColumnNotFoundSnafu { msg: column }.build(), + e @ (CheckError::UnknownKey { .. } + | CheckError::InvalidValue { .. } + | CheckError::ColumnNotStringForm { .. }) => InvalidSqlSnafu { + err_msg: e.to_string(), } - ); - Ok(()) -} - -/// Validates `greptime.semantic.entity..{id|descriptive|scope}` options -/// against the table schema: every named column must exist (tag or field) and -/// have a stable string form — the derivation renders id, scope and -/// descriptive values as strings. -fn validate_entity_semantic_options(table_options: &TableOptions, schema: &Schema) -> Result<()> { - for (key, value) in &table_options.extra_options { - if parse_entity_option_key(key).is_none() { - continue; - }; - for col in &parse_entity_columns(value) { - let column = schema - .column_schema_by_name(col) - .context(ColumnNotFoundSnafu { msg: col })?; - ensure_entity_column_renders_as_string(key, col, &column.data_type)?; + .build(), + e @ (CheckError::NotSingleColumn + | CheckError::PartitionMetadataConflict + | CheckError::TimeIndexConflict) => InvalidPartitionRuleSnafu { + reason: e.to_string(), } - } - Ok(()) + .build(), + }) } fn find_partition_columns(partitions: &Option) -> Result> { @@ -2904,6 +2818,7 @@ mod test { use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; use sqlparser::parser::Parser; + use table::requests::REPARTITION_COLUMN_HINT_KEY; use super::*; use crate::expr_helper; @@ -3133,7 +3048,7 @@ mod test { } #[test] - fn test_validate_entity_semantic_options() { + fn test_validate_and_normalize_annotations() { let schema = Schema::new(vec![ ColumnSchema::new("service_name", ConcreteDataType::string_datatype(), true), ColumnSchema::new("host_id", ConcreteDataType::string_datatype(), true), @@ -3147,6 +3062,10 @@ mod test { .collect(), ..Default::default() }; + let check = |pairs: &[(&str, &str)]| { + let mut options = opts(pairs); + validate_and_normalize_annotations(&mut options, &schema, &[]).map(|()| options) + }; // Any existing column with a string form may be an id, tag or field. for pairs in [ @@ -3156,22 +3075,27 @@ mod test { )][..], &[("greptime.semantic.entity.service.id", "value")][..], ] { - assert!(validate_entity_semantic_options(&opts(pairs), &schema).is_ok()); + assert!(check(pairs).is_ok()); } - let missing = validate_entity_semantic_options( - &opts(&[("greptime.semantic.entity.service.id", "nope")]), - &schema, - ) - .unwrap_err(); + // Missing columns keep this crate's error variant and status code. + let missing = check(&[("greptime.semantic.entity.service.id", "nope")]).unwrap_err(); assert!(matches!(missing, error::Error::ColumnNotFound { .. })); + assert_eq!( + common_error::status_code::StatusCode::InvalidArguments, + common_error::ext::ErrorExt::status_code(&missing) + ); - let binary_id = validate_entity_semantic_options( - &opts(&[("greptime.semantic.entity.service.id", "payload")]), - &schema, - ) - .unwrap_err(); + let binary_id = check(&[("greptime.semantic.entity.service.id", "payload")]).unwrap_err(); assert!(matches!(binary_id, error::Error::InvalidSql { .. })); + + // The SQL parser checks keys and value domains, but gRPC expressions + // bypass it. + let bad_value = check(&[("greptime.semantic.signal_type", "garbage")]).unwrap_err(); + assert!(matches!(bad_value, error::Error::InvalidSql { .. })); + + let unknown_key = check(&[("greptime.semantic.nonsense", "x")]).unwrap_err(); + assert!(matches!(unknown_key, error::Error::InvalidSql { .. })); } #[test] @@ -3559,9 +3483,10 @@ WITH ('repartition.column.hint' = ' host ')", }], })); let err = verify_alter(1, table_info.clone(), modify).unwrap_err(); + let msg = common_error::ext::ErrorExt::output_msg(&err); assert!( - err.to_string().contains("cannot render as a string"), - "{err}" + msg.contains("must keep a type that renders as a string"), + "{msg}" ); // Dropping a declared column stays allowed: the derivation skips the diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index 7b8958e22ce..4f7ed1c057d 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -37,8 +37,10 @@ use store_api::storage::{ColumnDescriptor, ColumnDescriptorBuilder, ColumnId}; use crate::error::{self, Result}; use crate::requests::{ - AddColumnRequest, AlterKind, ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, - SetDefaultRequest, SetIndexOption, TableOptions, UnsetIndexOption, + AddColumnRequest, AlterKind, AnnotationContext, AnnotationFamily, AnnotationValidationError, + ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, SetDefaultRequest, SetIndexOption, + TableOptions, UnsetIndexOption, has_stable_string_form, parse_entity_columns, + parse_entity_option_key, validate_and_normalize_annotation, }; use crate::table_reference::TableReference; @@ -331,10 +333,12 @@ impl TableMeta { AlterKind::RenameTable { .. } => Ok(self.new_meta_builder()), AlterKind::SetTableOptions { options } => self.set_table_options(options), AlterKind::UnsetTableOptions { keys } => self.unset_table_options(keys), - AlterKind::SetRepartitionColumnHint { column_name } => { - self.set_repartition_column_hint(table_name, column_name) + AlterKind::SetAnnotations { family, options } => { + self.set_annotations(table_name, *family, options) + } + AlterKind::UnsetAnnotations { family, keys } => { + self.unset_annotations(table_name, *family, keys) } - AlterKind::UnsetRepartitionColumnHint => self.unset_repartition_column_hint(), AlterKind::SetIndexes { options } => self.set_indexes(table_name, options), AlterKind::UnsetIndexes { options } => self.unset_indexes(table_name, options), AlterKind::DropDefaults { names } => self.drop_defaults(table_name, names), @@ -422,67 +426,92 @@ impl TableMeta { self.set_table_options(&requests) } - fn set_repartition_column_hint( + /// Applies an annotation SET. Validation lives here, on the mutation + /// path, so it runs both at frontend verification and again inside the + /// alter procedure's prepare step — under the table lock, against fresh + /// metadata. + fn set_annotations( &self, table_name: &str, - column_name: &str, + family: AnnotationFamily, + options: &[(String, String)], ) -> Result { - let column_name = column_name.trim(); ensure!( - !column_name.is_empty() && !column_name.contains(','), + !family.requires_single_key() || options.len() == 1, error::InvalidAlterRequestSnafu { table: table_name, - err: format!("{REPARTITION_COLUMN_HINT_KEY} expects exactly one column name"), + err: family.mixed_batch_error(), } ); - - ensure!( - self.partition_key_indices.is_empty(), - error::InvalidAlterRequestSnafu { - table: table_name, - err: format!( - "cannot set {REPARTITION_COLUMN_HINT_KEY} on a table with partition metadata" - ), - } - ); - - let column_index = self - .schema - .column_index_by_name(column_name) - .with_context(|| error::ColumnNotExistsSnafu { - column_name, - table_name, - })?; - - if let Some(time_index) = self.schema.timestamp_index() { + let cx = AnnotationContext { + schema: &self.schema, + partition_key_indices: &self.partition_key_indices, + }; + let mut new_options = self.options.clone(); + for (key, value) in options { ensure!( - column_index != time_index, + AnnotationFamily::of_key(key) == Some(family), error::InvalidAlterRequestSnafu { table: table_name, err: format!( - "cannot set {REPARTITION_COLUMN_HINT_KEY} to the time index column" + "`{key}` is outside the `{}` annotation namespace", + family.namespace() ), } ); + let checked = validate_and_normalize_annotation(family, &cx, key, value).map_err( + |e| match e { + AnnotationValidationError::ColumnNotFound { column } => { + error::ColumnNotExistsSnafu { + column_name: column, + table_name, + } + .build() + } + other => error::InvalidAlterRequestSnafu { + table: table_name, + err: other.to_string(), + } + .build(), + }, + )?; + new_options.extra_options.insert(key.clone(), checked); } - - let mut new_options = self.options.clone(); - new_options.extra_options.insert( - REPARTITION_COLUMN_HINT_KEY.to_string(), - column_name.to_string(), - ); - let mut builder = self.new_meta_builder(); builder.options(new_options); Ok(builder) } - fn unset_repartition_column_hint(&self) -> Result { + /// Applies an annotation UNSET. Deliberately lenient inside the family's + /// namespace: keys this version does not recognise may still be removed, + /// so options left behind by other versions can be cleaned up. + fn unset_annotations( + &self, + table_name: &str, + family: AnnotationFamily, + keys: &[String], + ) -> Result { + ensure!( + !family.requires_single_key() || keys.len() == 1, + error::InvalidAlterRequestSnafu { + table: table_name, + err: family.mixed_batch_error(), + } + ); let mut new_options = self.options.clone(); - new_options - .extra_options - .remove(REPARTITION_COLUMN_HINT_KEY); - + for key in keys { + ensure!( + AnnotationFamily::of_key(key) == Some(family), + error::InvalidAlterRequestSnafu { + table: table_name, + err: format!( + "`{key}` is outside the `{}` annotation namespace", + family.namespace() + ), + } + ); + new_options.extra_options.remove(key); + } let mut builder = self.new_meta_builder(); builder.options(new_options); Ok(builder) @@ -773,6 +802,24 @@ impl TableMeta { }, ); + // A dropped column may leave a stale entity declaration behind; + // re-adding it must not hand the declaration a type without a + // stable string form. + if !has_stable_string_form(&col_to_add.column_schema.data_type) + && let Some(key) = + entity_option_referencing(&self.options, &col_to_add.column_schema.name) + { + return error::InvalidAlterRequestSnafu { + table: table_name, + err: format!( + "column `{}` is referenced by entity option `{key}` and must \ + keep a type that renders as a string, got `{}`", + col_to_add.column_schema.name, col_to_add.column_schema.data_type + ), + } + .fail(); + } + new_columns.push(col_to_add.clone()); } } @@ -966,6 +1013,25 @@ impl TableMeta { table_name, })?; + // A column referenced by an entity declaration may be dropped (the + // read-time derivation skips the stale declaration), but must not + // change to a type without a stable string form. Checked after the + // existence lookup so a missing column keeps reporting + // `ColumnNotExists` regardless of stale declarations. + if !has_stable_string_form(&col_to_change.target_type) + && let Some(key) = entity_option_referencing(&self.options, change_column_name) + { + return error::InvalidAlterRequestSnafu { + table: table_name, + err: format!( + "column `{change_column_name}` is referenced by entity option \ + `{key}` and must keep a type that renders as a string, got `{}`", + col_to_change.target_type + ), + } + .fail(); + } + let column = &table_schema.column_schemas()[index]; ensure!( @@ -1408,6 +1474,14 @@ impl TableInfo { } } +fn entity_option_referencing<'a>(options: &'a TableOptions, column: &str) -> Option<&'a str> { + options.extra_options.iter().find_map(|(key, value)| { + (parse_entity_option_key(key).is_some() + && parse_entity_columns(value).iter().any(|c| c == column)) + .then_some(key.as_str()) + }) +} + /// Set column fulltext options if it passed the validation. /// /// Options allowed to modify: @@ -1738,8 +1812,12 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: " col1 ".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![( + REPARTITION_COLUMN_HINT_KEY.to_string(), + " col1 ".to_string(), + )], }; let new_meta = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1767,8 +1845,9 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: " ".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![(REPARTITION_COLUMN_HINT_KEY.to_string(), " ".to_string())], }; let err = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1791,8 +1870,12 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: "col1,col2".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![( + REPARTITION_COLUMN_HINT_KEY.to_string(), + "col1,col2".to_string(), + )], }; let err = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1815,8 +1898,12 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: "missing".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![( + REPARTITION_COLUMN_HINT_KEY.to_string(), + "missing".to_string(), + )], }; let err = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1836,8 +1923,9 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: "ts".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![(REPARTITION_COLUMN_HINT_KEY.to_string(), "ts".to_string())], }; let err = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1861,8 +1949,9 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: "col1".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![(REPARTITION_COLUMN_HINT_KEY.to_string(), "col1".to_string())], }; let err = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1891,7 +1980,13 @@ mod tests { .unwrap(); let new_meta = meta - .builder_with_alter_kind("my_table", &AlterKind::UnsetRepartitionColumnHint) + .builder_with_alter_kind( + "my_table", + &AlterKind::UnsetAnnotations { + family: AnnotationFamily::RepartitionHint, + keys: vec![REPARTITION_COLUMN_HINT_KEY.to_string()], + }, + ) .unwrap() .build() .unwrap(); @@ -2627,4 +2722,305 @@ mod tests { }; assert_eq!(actual, expected); } + + use crate::requests::{SEMANTIC_METRIC_UNIT, SEMANTIC_SIGNAL_TYPE}; + + /// `host` is a tag, `service` and `note` are string fields, `payload` is a + /// binary field. + fn semantic_test_meta() -> TableMeta { + let column_schemas = vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("payload", ConcreteDataType::binary_datatype(), true), + ColumnSchema::new("service", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("note", ConcreteDataType::string_datatype(), true), + ]; + let schema = Arc::new( + SchemaBuilder::try_from(column_schemas) + .unwrap() + .build() + .unwrap(), + ); + TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(5) + .build() + .unwrap() + } + + #[test] + fn test_set_semantic_annotations() { + let meta = semantic_test_meta(); + // `service` is a plain field: entity columns may be tags or fields. + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::Semantic, + options: vec![ + (SEMANTIC_SIGNAL_TYPE.to_string(), "trace".to_string()), + ( + "greptime.semantic.entity.service.id".to_string(), + "service".to_string(), + ), + ], + }; + let new_meta = meta + .builder_with_alter_kind("my_table", &alter_kind) + .unwrap() + .build() + .unwrap(); + assert_eq!( + new_meta.options.extra_options.get(SEMANTIC_SIGNAL_TYPE), + Some(&"trace".to_string()) + ); + assert_eq!( + new_meta + .options + .extra_options + .get("greptime.semantic.entity.service.id"), + Some(&"service".to_string()) + ); + } + + #[test] + fn test_repartition_hint_batch_must_be_single_key() { + let meta = TableMetaBuilder::empty() + .schema(Arc::new(new_test_schema())) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + + // Direct gRPC can hand the mutation layer a duplicated batch the + // converter never saw; last-write-wins must not silently apply. + let dup_set = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![ + (REPARTITION_COLUMN_HINT_KEY.to_string(), "col1".to_string()), + (REPARTITION_COLUMN_HINT_KEY.to_string(), "col2".to_string()), + ], + }; + let err = meta + .builder_with_alter_kind("my_table", &dup_set) + .err() + .unwrap(); + assert!( + err.to_string().contains("must be altered separately"), + "{err}" + ); + + let dup_unset = AlterKind::UnsetAnnotations { + family: AnnotationFamily::RepartitionHint, + keys: vec![ + REPARTITION_COLUMN_HINT_KEY.to_string(), + REPARTITION_COLUMN_HINT_KEY.to_string(), + ], + }; + let err = meta + .builder_with_alter_kind("my_table", &dup_unset) + .err() + .unwrap(); + assert!( + err.to_string().contains("must be altered separately"), + "{err}" + ); + } + + #[test] + fn test_set_semantic_annotations_rejects_invalid() { + let meta = semantic_test_meta(); + let cases = [ + ( + "greptime.semantic.unknown_key", + "x", + "unknown semantic option", + ), + (SEMANTIC_SIGNAL_TYPE, "garbage", "invalid value"), + ( + "greptime.semantic.entity.host.id", + "no_such_column", + "no_such_column", + ), + ( + "greptime.semantic.entity.host.id", + "payload", + "cannot render as a string", + ), + ]; + for (key, value, needle) in cases { + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::Semantic, + options: vec![(key.to_string(), value.to_string())], + }; + let err = meta + .builder_with_alter_kind("my_table", &alter_kind) + .err() + .unwrap(); + assert!( + err.to_string().contains(needle), + "key `{key}`: unexpected error `{err}`" + ); + } + + // ALTER keeps missing columns on the 4002 contract (pinned by sqlness); + // the shared validator must not collapse it into InvalidArguments. + let missing = meta + .builder_with_alter_kind( + "my_table", + &AlterKind::SetAnnotations { + family: AnnotationFamily::Semantic, + options: vec![( + "greptime.semantic.entity.host.id".to_string(), + "no_such_column".to_string(), + )], + }, + ) + .err() + .unwrap(); + assert_eq!( + common_error::status_code::StatusCode::TableColumnNotFound, + common_error::ext::ErrorExt::status_code(&missing) + ); + } + + #[test] + fn test_unset_semantic_annotations() { + let mut meta = semantic_test_meta(); + meta.options + .extra_options + .insert(SEMANTIC_SIGNAL_TYPE.to_string(), "trace".to_string()); + // A key this version does not recognise, still inside the namespace. + meta.options + .extra_options + .insert("greptime.semantic.future_key".to_string(), "x".to_string()); + + let alter_kind = AlterKind::UnsetAnnotations { + family: AnnotationFamily::Semantic, + keys: vec![ + SEMANTIC_SIGNAL_TYPE.to_string(), + "greptime.semantic.future_key".to_string(), + // Absent key: removal is a no-op, not an error. + SEMANTIC_METRIC_UNIT.to_string(), + ], + }; + let new_meta = meta + .builder_with_alter_kind("my_table", &alter_kind) + .unwrap() + .build() + .unwrap(); + assert!( + !new_meta + .options + .extra_options + .contains_key(SEMANTIC_SIGNAL_TYPE) + ); + assert!( + !new_meta + .options + .extra_options + .contains_key("greptime.semantic.future_key") + ); + + let outside = AlterKind::UnsetAnnotations { + family: AnnotationFamily::Semantic, + keys: vec!["ttl".to_string()], + }; + let err = meta + .builder_with_alter_kind("my_table", &outside) + .err() + .unwrap(); + assert!(err.to_string().contains("annotation namespace"), "{err}"); + } + + #[test] + fn test_modify_entity_column_type_keeps_string_form() { + let mut meta = semantic_test_meta(); + meta.options.extra_options.insert( + "greptime.semantic.entity.service.id".to_string(), + "service".to_string(), + ); + + let alter_kind = AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnTypeRequest { + column_name: "service".to_string(), + target_type: ConcreteDataType::binary_datatype(), + }], + }; + let err = meta + .builder_with_alter_kind("my_table", &alter_kind) + .err() + .unwrap(); + assert!( + err.to_string() + .contains("must keep a type that renders as a string"), + "{err}" + ); + + // An unreferenced column may still change to a non-string form. + let alter_kind = AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnTypeRequest { + column_name: "note".to_string(), + target_type: ConcreteDataType::binary_datatype(), + }], + }; + meta.builder_with_alter_kind("my_table", &alter_kind) + .unwrap(); + } + + #[test] + fn test_stale_entity_declaration_guards_readd_and_reports_missing_column() { + // A stale declaration: `gone` was dropped after being declared. + let mut meta = semantic_test_meta(); + meta.options.extra_options.insert( + "greptime.semantic.entity.service.id".to_string(), + "gone".to_string(), + ); + + // MODIFY on the missing column keeps the ColumnNotExists contract + // (4002), stale declaration or not. + let modify = AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnTypeRequest { + column_name: "gone".to_string(), + target_type: ConcreteDataType::binary_datatype(), + }], + }; + let err = meta + .builder_with_alter_kind("my_table", &modify) + .err() + .unwrap(); + assert_eq!( + common_error::status_code::StatusCode::TableColumnNotFound, + common_error::ext::ErrorExt::status_code(&err) + ); + + // Re-adding the declared column must not hand the declaration a + // non-string type... + let add = |ty: ConcreteDataType| AlterKind::AddColumns { + columns: vec![AddColumnRequest { + column_schema: ColumnSchema::new("gone", ty, true), + is_key: false, + location: None, + add_if_not_exists: false, + }], + }; + let err = meta + .builder_with_alter_kind("my_table", &add(ConcreteDataType::binary_datatype())) + .err() + .unwrap(); + assert!( + err.to_string() + .contains("must keep a type that renders as a string"), + "{err}" + ); + + // ...while a string re-add re-satisfies it. + meta.builder_with_alter_kind("my_table", &add(ConcreteDataType::string_datatype())) + .unwrap(); + } } diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index b4a7d5e052b..aea45130ed9 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -27,7 +27,7 @@ use common_time::range::TimestampRange; use datatypes::data_type::ConcreteDataType; use datatypes::prelude::VectorRef; use datatypes::schema::{ - ColumnDefaultConstraint, ColumnSchema, FulltextOptions, SkippingIndexOptions, + ColumnDefaultConstraint, ColumnSchema, FulltextOptions, Schema, SkippingIndexOptions, }; use greptime_proto::v1::region::compact_request; use once_cell::sync::Lazy; @@ -301,6 +301,211 @@ pub struct ModifyColumnTypeRequest { pub target_type: ConcreteDataType, } +/// A family of annotation table options: pure metadata markers that no region +/// consumes. Setting or unsetting them only rewrites the table's +/// `extra_options`, so the alter skips region dispatch entirely. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum AnnotationFamily { + /// `greptime.semantic.*` options (see the [`semantic`] module). + Semantic, + /// `repartition.column.hint`, consumed by the auto-repartition planner. + RepartitionHint, +} + +impl AnnotationFamily { + /// The key namespace: a prefix for [`Self::Semantic`], the exact key for + /// [`Self::RepartitionHint`]. + pub fn namespace(self) -> &'static str { + match self { + Self::Semantic => SEMANTIC_PREFIX, + Self::RepartitionHint => REPARTITION_COLUMN_HINT_KEY, + } + } + + pub fn of_key(key: &str) -> Option { + if key.starts_with(SEMANTIC_PREFIX) { + Some(Self::Semantic) + } else if key == REPARTITION_COLUMN_HINT_KEY { + Some(Self::RepartitionHint) + } else { + None + } + } + + /// Whether this family may be altered on logical metric tables. Only + /// families whose values nothing on the physical side consumes qualify; + /// the repartition hint drives physical region repartitioning. + pub fn allows_logical_tables(self) -> bool { + match self { + Self::Semantic => true, + Self::RepartitionHint => false, + } + } + + /// Whether this family's SET/UNSET batch must contain exactly one key. + /// The repartition hint is a single marker; a batch with several hint + /// entries (duplicates included) has no meaningful order. + pub fn requires_single_key(self) -> bool { + matches!(self, Self::RepartitionHint) + } + + /// The error for a SET/UNSET batch mixing this family with other options. + pub fn mixed_batch_error(self) -> String { + match self { + Self::Semantic => format!( + "`{SEMANTIC_PREFIX}*` options must be altered separately from other table options" + ), + Self::RepartitionHint => { + format!("{REPARTITION_COLUMN_HINT_KEY} must be altered separately") + } + } + } +} + +/// Table shape an annotation option is validated against. +pub struct AnnotationContext<'a> { + pub schema: &'a Schema, + pub partition_key_indices: &'a [usize], +} + +/// Why an annotation option was rejected. Typed so each DDL entry point maps +/// rules onto its existing error variants and status codes: ALTER keeps +/// missing columns as `TableColumnNotFound` (4002), CREATE keeps its +/// `InvalidArguments` family — the rules converge, the contracts do not. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AnnotationValidationError { + UnknownKey { + key: String, + }, + InvalidValue { + key: String, + value: String, + }, + ColumnNotFound { + column: String, + }, + ColumnNotStringForm { + key: String, + column: String, + ty: ConcreteDataType, + }, + NotSingleColumn, + PartitionMetadataConflict, + TimeIndexConflict, +} + +impl fmt::Display for AnnotationValidationError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::UnknownKey { key } => write!(f, "unknown semantic option `{key}`"), + Self::InvalidValue { key, value } => { + write!(f, "invalid value `{value}` for semantic option `{key}`") + } + Self::ColumnNotFound { column } => write!(f, "column `{column}` not found"), + Self::ColumnNotStringForm { key, column, ty } => write!( + f, + "entity column `{column}` (option `{key}`) has type `{ty}`, \ + which cannot render as a string" + ), + Self::NotSingleColumn => write!( + f, + "{REPARTITION_COLUMN_HINT_KEY} expects exactly one column name" + ), + Self::PartitionMetadataConflict => write!( + f, + "cannot set {REPARTITION_COLUMN_HINT_KEY} on a table with partition metadata" + ), + Self::TimeIndexConflict => write!( + f, + "cannot set {REPARTITION_COLUMN_HINT_KEY} to the time index column" + ), + } + } +} + +/// Validates one annotation option and returns the value to store — the +/// repartition hint is trimmed to the bare column name, semantic values pass +/// through unchanged. +pub(crate) fn validate_and_normalize_annotation( + family: AnnotationFamily, + cx: &AnnotationContext<'_>, + key: &str, + value: &str, +) -> std::result::Result { + match family { + AnnotationFamily::Semantic => { + if !is_semantic_option_key(key) { + return Err(AnnotationValidationError::UnknownKey { + key: key.to_string(), + }); + } + if !validate_semantic_option(key, value) { + return Err(AnnotationValidationError::InvalidValue { + key: key.to_string(), + value: value.to_string(), + }); + } + if parse_entity_option_key(key).is_some() { + for column in parse_entity_columns(value) { + let schema = cx.schema.column_schema_by_name(&column).ok_or_else(|| { + AnnotationValidationError::ColumnNotFound { + column: column.clone(), + } + })?; + if !has_stable_string_form(&schema.data_type) { + return Err(AnnotationValidationError::ColumnNotStringForm { + key: key.to_string(), + column, + ty: schema.data_type.clone(), + }); + } + } + } + Ok(value.to_string()) + } + AnnotationFamily::RepartitionHint => { + let column_name = value.trim(); + if column_name.is_empty() || column_name.contains(',') { + return Err(AnnotationValidationError::NotSingleColumn); + } + if !cx.partition_key_indices.is_empty() { + return Err(AnnotationValidationError::PartitionMetadataConflict); + } + let column_index = cx.schema.column_index_by_name(column_name).ok_or_else(|| { + AnnotationValidationError::ColumnNotFound { + column: column_name.to_string(), + } + })?; + if cx.schema.timestamp_index() == Some(column_index) { + return Err(AnnotationValidationError::TimeIndexConflict); + } + Ok(column_name.to_string()) + } + } +} + +/// CREATE-side entry: validates every annotation option present in `options` +/// and writes normalized values back in place. +pub fn validate_and_normalize_annotation_options( + options: &mut TableOptions, + cx: &AnnotationContext<'_>, +) -> std::result::Result<(), AnnotationValidationError> { + let mut normalized = Vec::new(); + for (key, value) in &options.extra_options { + let Some(family) = AnnotationFamily::of_key(key) else { + continue; + }; + let checked = validate_and_normalize_annotation(family, cx, key, value)?; + if checked != *value { + normalized.push((key.clone(), checked)); + } + } + for (key, value) in normalized { + options.extra_options.insert(key, value); + } + Ok(()) +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub enum AlterKind { AddColumns { @@ -321,10 +526,14 @@ pub enum AlterKind { UnsetTableOptions { keys: Vec, }, - SetRepartitionColumnHint { - column_name: String, + SetAnnotations { + family: AnnotationFamily, + options: Vec<(String, String)>, + }, + UnsetAnnotations { + family: AnnotationFamily, + keys: Vec, }, - UnsetRepartitionColumnHint, SetIndexes { options: Vec, }, diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index 5c0e4c3caed..2daa440c11d 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -30,6 +30,8 @@ //! [`crate::requests::validate_table_option`], so they are accepted both on the //! ingestion auto-create path and on explicit `CREATE TABLE ... WITH (...)` DDL. +use datatypes::prelude::ConcreteDataType; + /// Reserved prefix for every public semantic table-option key. pub const SEMANTIC_PREFIX: &str = "greptime.semantic."; @@ -90,8 +92,10 @@ pub const SEMANTIC_ENTITY_PREFIX: &str = "greptime.semantic.entity."; /// is the `service_name` tag column, declaring the logical `service` entity. pub const SEMANTIC_ENTITY_SERVICE_ID: &str = "greptime.semantic.entity.service.id"; -/// The role a set of columns plays for an entity: `id` (identifying attributes, -/// must be tag columns — enforced at DDL time), `descriptive`, or `scope`. +/// The role a set of columns plays for an entity: `id` (identifying +/// attributes), `descriptive`, or `scope`. Columns may be tags or fields; DDL +/// validation only requires that they exist and render as stable strings +/// ([`has_stable_string_form`]). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum EntityRole { Id, @@ -182,6 +186,24 @@ pub fn is_entity_option_key(key: &str) -> bool { parse_entity_option_key(key).is_some() } +/// Returns true if a column of `data_type` renders as a stable string — the +/// requirement for entity id/descriptive/scope columns. The read-time +/// derivation casts them to strings, so a type without a stable string form +/// would fail only when the graph is scanned; DDL validation rejects it up +/// front instead. +pub fn has_stable_string_form(data_type: &ConcreteDataType) -> bool { + !matches!( + data_type, + ConcreteDataType::Binary(_) + | ConcreteDataType::Json(_) + | ConcreteDataType::Vector(_) + | ConcreteDataType::List(_) + | ConcreteDataType::Struct(_) + | ConcreteDataType::Dictionary(_) + | ConcreteDataType::Null(_) + ) +} + /// Tokenizes an entity option's comma-separated column list (trimmed, empty /// tokens dropped). [`validate_semantic_option`] rejects empty tokens at DDL /// time, so readers only ever drop what validation already refused. diff --git a/tests/cases/standalone/common/information_schema/table_semantics.result b/tests/cases/standalone/common/information_schema/table_semantics.result index 5aa39a91fbf..b57087b3f50 100644 --- a/tests/cases/standalone/common/information_schema/table_semantics.result +++ b/tests/cases/standalone/common/information_schema/table_semantics.result @@ -86,3 +86,121 @@ DROP TABLE plain_table; Affected Rows: 0 +-- ALTER TABLE manages semantic declarations on existing tables: SET appears in +-- the view, UNSET disappears from it. +CREATE TABLE altered_semantics ( + ts TIMESTAMP TIME INDEX, + svc STRING, + payload BINARY, + val DOUBLE, +); + +Affected Rows: 0 + +ALTER TABLE altered_semantics SET 'greptime.semantic.signal_type' = 'metric', 'greptime.semantic.entity.service.id' = 'svc'; + +Affected Rows: 0 + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'altered_semantics'; + ++-------------------+-------------+-----------------------------+ +| table_name | signal_type | semantic_options | ++-------------------+-------------+-----------------------------+ +| altered_semantics | metric | {"entity.service.id":"svc"} | ++-------------------+-------------+-----------------------------+ + +-- Semantic options never share a statement with regular table options. +ALTER TABLE altered_semantics SET 'greptime.semantic.source' = 'prometheus', 'ttl' = '7d'; + +Error: 1004(InvalidArguments), Invalid table option request: `greptime.semantic.*` options must be altered separately from other table options + +-- Unknown semantic keys are rejected on SET. +ALTER TABLE altered_semantics SET 'greptime.semantic.nonsense' = 'x'; + +Error: 1004(InvalidArguments), Invalid alter table(altered_semantics) request: unknown semantic option `greptime.semantic.nonsense` + +-- Values outside the key's domain are rejected. +ALTER TABLE altered_semantics SET 'greptime.semantic.signal_type' = 'garbage'; + +Error: 1004(InvalidArguments), Invalid alter table(altered_semantics) request: invalid value `garbage` for semantic option `greptime.semantic.signal_type` + +-- Entity columns must exist. +ALTER TABLE altered_semantics SET 'greptime.semantic.entity.host.id' = 'no_such_column'; + +Error: 4002(TableColumnNotFound), Column no_such_column not exists in table altered_semantics + +-- Entity columns must render as strings. +ALTER TABLE altered_semantics SET 'greptime.semantic.entity.host.id' = 'payload'; + +Error: 1004(InvalidArguments), Invalid alter table(altered_semantics) request: entity column `payload` (option `greptime.semantic.entity.host.id`) has type `Binary`, which cannot render as a string + +ALTER TABLE altered_semantics UNSET 'greptime.semantic.entity.service.id'; + +Affected Rows: 0 + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'altered_semantics'; + ++-------------------+-------------+------------------+ +| table_name | signal_type | semantic_options | ++-------------------+-------------+------------------+ +| altered_semantics | metric | | ++-------------------+-------------+------------------+ + +DROP TABLE altered_semantics; + +Affected Rows: 0 + +-- Logical metric tables take the same metadata-only path. +CREATE TABLE phy_sem (ts TIMESTAMP TIME INDEX, val DOUBLE) engine=metric with ("physical_metric_table" = ""); + +Affected Rows: 0 + +CREATE TABLE logical_sem (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) engine=metric with ("on_physical_table" = "phy_sem"); + +Affected Rows: 0 + +ALTER TABLE logical_sem SET 'greptime.semantic.signal_type' = 'metric', 'greptime.semantic.entity.host.id' = 'host'; + +Affected Rows: 0 + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'logical_sem'; + ++-------------+-------------+---------------------------+ +| table_name | signal_type | semantic_options | ++-------------+-------------+---------------------------+ +| logical_sem | metric | {"entity.host.id":"host"} | ++-------------+-------------+---------------------------+ + +-- Regular options still cannot be altered on logical tables. +ALTER TABLE logical_sem SET 'ttl' = '7d'; + +Error: 1004(InvalidArguments), Alter logical tables invalid arguments: Only support add columns operation + +ALTER TABLE logical_sem UNSET 'greptime.semantic.entity.host.id'; + +Affected Rows: 0 + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'logical_sem'; + ++-------------+-------------+------------------+ +| table_name | signal_type | semantic_options | ++-------------+-------------+------------------+ +| logical_sem | metric | | ++-------------+-------------+------------------+ + +DROP TABLE logical_sem; + +Affected Rows: 0 + +DROP TABLE phy_sem; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/information_schema/table_semantics.sql b/tests/cases/standalone/common/information_schema/table_semantics.sql index 8fb7c78aefe..4acb15294da 100644 --- a/tests/cases/standalone/common/information_schema/table_semantics.sql +++ b/tests/cases/standalone/common/information_schema/table_semantics.sql @@ -45,3 +45,65 @@ DROP TABLE metrics_tagged; DROP TABLE traces_tagged; DROP TABLE plain_table; + +-- ALTER TABLE manages semantic declarations on existing tables: SET appears in +-- the view, UNSET disappears from it. +CREATE TABLE altered_semantics ( + ts TIMESTAMP TIME INDEX, + svc STRING, + payload BINARY, + val DOUBLE, +); + +ALTER TABLE altered_semantics SET 'greptime.semantic.signal_type' = 'metric', 'greptime.semantic.entity.service.id' = 'svc'; + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'altered_semantics'; + +-- Semantic options never share a statement with regular table options. +ALTER TABLE altered_semantics SET 'greptime.semantic.source' = 'prometheus', 'ttl' = '7d'; + +-- Unknown semantic keys are rejected on SET. +ALTER TABLE altered_semantics SET 'greptime.semantic.nonsense' = 'x'; + +-- Values outside the key's domain are rejected. +ALTER TABLE altered_semantics SET 'greptime.semantic.signal_type' = 'garbage'; + +-- Entity columns must exist. +ALTER TABLE altered_semantics SET 'greptime.semantic.entity.host.id' = 'no_such_column'; + +-- Entity columns must render as strings. +ALTER TABLE altered_semantics SET 'greptime.semantic.entity.host.id' = 'payload'; + +ALTER TABLE altered_semantics UNSET 'greptime.semantic.entity.service.id'; + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'altered_semantics'; + +DROP TABLE altered_semantics; + +-- Logical metric tables take the same metadata-only path. +CREATE TABLE phy_sem (ts TIMESTAMP TIME INDEX, val DOUBLE) engine=metric with ("physical_metric_table" = ""); + +CREATE TABLE logical_sem (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) engine=metric with ("on_physical_table" = "phy_sem"); + +ALTER TABLE logical_sem SET 'greptime.semantic.signal_type' = 'metric', 'greptime.semantic.entity.host.id' = 'host'; + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'logical_sem'; + +-- Regular options still cannot be altered on logical tables. +ALTER TABLE logical_sem SET 'ttl' = '7d'; + +ALTER TABLE logical_sem UNSET 'greptime.semantic.entity.host.id'; + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'logical_sem'; + +DROP TABLE logical_sem; + +DROP TABLE phy_sem; diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index ed6a3d3f16f..6af28cfa2b5 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -119,6 +119,61 @@ drop table graph_app_metrics; Affected Rows: 0 +-- Declarations can be added after the fact: a table created without semantic +-- options joins the graph once ALTER TABLE declares its entities, and leaves +-- it again on UNSET. +create table graph_late_metrics ( + ts timestamp time index, + svc string, + env string, + latency double +); + +Affected Rows: 0 + +insert into graph_late_metrics values (now(), 'checkout', 'eu-1', 1.5); + +Affected Rows: 1 + +select entity_type, entity_id from greptime_private.semantic_entities order by entity_type, entity_id; + +++ +++ + +alter table graph_late_metrics set 'greptime.semantic.entity.service.id' = 'svc', 'greptime.semantic.entity.service.scope' = 'env'; + +Affected Rows: 0 + +select entity_type, entity_id, scope from greptime_private.semantic_entities order by entity_type, entity_id; + ++-------------+-----------+-------+ +| entity_type | entity_id | scope | ++-------------+-----------+-------+ +| service | checkout | eu-1 | ++-------------+-----------+-------+ + +-- A column referenced by a declaration must keep a string-renderable type. +alter table graph_late_metrics modify column svc binary; + +Error: 1004(InvalidArguments), Invalid alter table(graph_late_metrics) request: column `svc` is referenced by entity option `greptime.semantic.entity.service.id` and must keep a type that renders as a string, got `Binary` + +alter table graph_late_metrics unset 'greptime.semantic.entity.service.id'; + +Affected Rows: 0 + +alter table graph_late_metrics unset 'greptime.semantic.entity.service.scope'; + +Affected Rows: 0 + +select entity_type, entity_id from greptime_private.semantic_entities order by entity_type, entity_id; + +++ +++ + +drop table graph_late_metrics; + +Affected Rows: 0 + -- Calls derivation over trace-v1 tables, all branches in one scan: the client -- span and its child server span land in different tables and still pair into -- one edge with RED metrics; the epoch-timestamped pair falls outside the diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index 181a8da66a9..bfc9ea8fd7b 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -66,6 +66,35 @@ order by rel_type, src_id; drop table graph_app_metrics; +-- Declarations can be added after the fact: a table created without semantic +-- options joins the graph once ALTER TABLE declares its entities, and leaves +-- it again on UNSET. +create table graph_late_metrics ( + ts timestamp time index, + svc string, + env string, + latency double +); + +insert into graph_late_metrics values (now(), 'checkout', 'eu-1', 1.5); + +select entity_type, entity_id from greptime_private.semantic_entities order by entity_type, entity_id; + +alter table graph_late_metrics set 'greptime.semantic.entity.service.id' = 'svc', 'greptime.semantic.entity.service.scope' = 'env'; + +select entity_type, entity_id, scope from greptime_private.semantic_entities order by entity_type, entity_id; + +-- A column referenced by a declaration must keep a string-renderable type. +alter table graph_late_metrics modify column svc binary; + +alter table graph_late_metrics unset 'greptime.semantic.entity.service.id'; + +alter table graph_late_metrics unset 'greptime.semantic.entity.service.scope'; + +select entity_type, entity_id from greptime_private.semantic_entities order by entity_type, entity_id; + +drop table graph_late_metrics; + -- Calls derivation over trace-v1 tables, all branches in one scan: the client -- span and its child server span land in different tables and still pair into -- one edge with RED metrics; the epoch-timestamped pair falls outside the From 10f587bc30071f230c979da6c0d3ebbf1d736698 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 17 Aug 2026 10:19:08 +0000 Subject: [PATCH 081/126] fix(query): preserve timestamp literal semantics in inserts (#8889) * fix(query): follow timestamp insert assignment lineage Signed-off-by: Dennis Zhuang * fix(query): fold constant insert timestamp literals at the assignment Following lineage by retyping the source column changed every output column that reads it: a string column sharing the literal was silently rewritten to a formatted timestamp, and a nanosecond column was truncated to the precision of whichever column was converted first. Resolve the constant read-only and fold it into the assignment expression instead, which leaves the source query untouched and also covers literals behind WHERE, ORDER BY and DISTINCT. VALUES rows and UNION branches carry per-row values, so they keep the in-place rewrite, now guarded against columns with more than one consumer. Signed-off-by: Dennis Zhuang * test(query): strengthen insert lineage regression Signed-off-by: Dennis Zhuang * test(query): trim redundant insert coverage Signed-off-by: Dennis Zhuang * fix(query): rebuild insert unions loosely and cover UNION distinct Per review: rebuilding a rewritten union with the strict constructor rejected legal pre-coercion plans whose untouched columns still differ across branches. Use try_new_with_loose_types, matching the SQL planner. Distinct::All joins the rewrite passthrough so UNION (distinct) literals get session-timezone parsing like UNION ALL; deduplication then keys on parsed instants instead of raw strings. The top-level rewrite path gains the same single-consumer guard as rewrite_projection for hand-built DML plans that share a source column between targets. Signed-off-by: Dennis Zhuang * style(query): tighten insert assignment comments Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/query/src/optimizer/type_conversion.rs | 160 +-------- .../type_conversion/insert_assignment.rs | 336 ++++++++++++++++++ src/query/src/planner.rs | 203 ++++++++++- .../insert/insert_default_timezone.result | 53 ++- .../common/insert/insert_default_timezone.sql | 30 +- 5 files changed, 618 insertions(+), 164 deletions(-) create mode 100644 src/query/src/optimizer/type_conversion/insert_assignment.rs diff --git a/src/query/src/optimizer/type_conversion.rs b/src/query/src/optimizer/type_conversion.rs index 366a7d8a31d..3941fedfe8a 100644 --- a/src/query/src/optimizer/type_conversion.rs +++ b/src/query/src/optimizer/type_conversion.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod insert_assignment; + use std::sync::Arc; use common_time::Timezone; @@ -19,7 +21,7 @@ use common_time::timestamp::{TimeUnit, Timestamp}; use datafusion::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; use datafusion_common::{DFSchemaRef, DataFusionError, Result, ScalarValue}; -use datafusion_expr::expr::{Cast, InList}; +use datafusion_expr::expr::InList; use datafusion_expr::{ Between, BinaryExpr, Expr, ExprSchemable, Filter, LogicalPlan, Operator, TableScan, WriteOp, }; @@ -29,6 +31,7 @@ use session::context::QueryContextRef; use crate::QueryEngineContext; use crate::optimizer::ExtensionAnalyzerRule; +use crate::optimizer::type_conversion::insert_assignment::rewrite_insert_assignments; use crate::plan::ExtractExpr; /// TypeConversionRule converts some literal values in logical plan to other types according @@ -307,120 +310,6 @@ impl TreeNodeRewriter for TypeConverter { } } -/// Rewrites timestamp assignment casts at the INSERT boundary. The source query -/// is intentionally left untouched so explicit casts keep DataFusion semantics. -fn rewrite_insert_assignments( - plan: LogicalPlan, - query_ctx: QueryContextRef, -) -> Result { - let LogicalPlan::Projection(projection) = plan else { - return Ok(plan); - }; - - // VALUES type coercion is performed inside the Values node. Only rewrite - // the direct source of the assignment projection, not nested query plans. - let input = rewrite_insert_values(projection.input.as_ref().clone(), query_ctx.clone())?; - let input_literals = projected_literals(&input); - let mut converter = InsertAssignmentConverter { - query_ctx, - input_literals, - }; - let plan = LogicalPlan::Projection(projection); - let expr = plan - .expressions_consider_join() - .into_iter() - .map(|expr| expr.rewrite(&mut converter).map(|x| x.data)) - .collect::>>()?; - - plan.with_new_exprs(expr, vec![input]) -} - -fn rewrite_insert_values(plan: LogicalPlan, query_ctx: QueryContextRef) -> Result { - let LogicalPlan::Values(mut values) = plan else { - return Ok(plan); - }; - let mut converter = InsertAssignmentConverter { - query_ctx, - input_literals: None, - }; - values.values = values - .values - .into_iter() - .map(|row| { - row.into_iter() - .map(|expr| expr.rewrite(&mut converter).map(|x| x.data)) - .collect::>>() - }) - .collect::>>()?; - - Ok(LogicalPlan::Values(values)) -} - -fn projected_literals(plan: &LogicalPlan) -> Option<(DFSchemaRef, Vec>)> { - let LogicalPlan::Projection(projection) = plan else { - return None; - }; - let literals = projection.expr.iter().map(projected_literal).collect(); - Some((projection.schema.clone(), literals)) -} - -fn projected_literal(expr: &Expr) -> Option { - match expr { - Expr::Literal(value, _) => Some(value.clone()), - Expr::Alias(alias) => projected_literal(&alias.expr), - _ => None, - } -} - -struct InsertAssignmentConverter { - query_ctx: QueryContextRef, - input_literals: Option<(DFSchemaRef, Vec>)>, -} - -impl InsertAssignmentConverter { - fn assignment_literal(&self, expr: &Expr) -> Option { - let Expr::Column(column) = expr else { - return expr.as_literal().cloned(); - }; - let (schema, literals) = self.input_literals.as_ref()?; - let index = schema.maybe_index_of_column(column)?; - literals.get(index)?.clone() - } -} - -impl TreeNodeRewriter for InsertAssignmentConverter { - type Node = Expr; - - fn f_up(&mut self, expr: Expr) -> Result> { - let Expr::Cast(Cast { - expr: inner, - data_type: target_type, - }) = expr - else { - return Ok(Transformed::no(expr)); - }; - let original = || { - Expr::Cast(Cast { - expr: inner.clone(), - data_type: target_type.clone(), - }) - }; - if !matches!(&target_type, DataType::Timestamp(_, _)) { - return Ok(Transformed::no(original())); - } - let Some(ScalarValue::Utf8(Some(value))) = self.assignment_literal(inner.as_ref()) else { - return Ok(Transformed::no(original())); - }; - - match cast_string_to_timestamp(&value, &target_type, Some(&self.query_ctx.timezone())) { - Ok(value) if !value.is_null() => Ok(Transformed::yes(Expr::Literal(value, None))), - // Preserve DataFusion's existing parser for unsupported timestamp - // syntax and out-of-range values. - Ok(_) | Err(_) => Ok(Transformed::no(original())), - } - } -} - fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr { let timestamp = match unit { TimeUnit::Second => val * 1_000, @@ -472,7 +361,7 @@ mod tests { use datafusion_common::arrow::datatypes::Field; use datafusion_common::{Column, DFSchema, NullEquality}; - use datafusion_expr::expr::Exists; + use datafusion_expr::expr::{Cast, Exists}; use datafusion_expr::{Join, JoinConstraint, JoinType, Literal, LogicalPlanBuilder, Subquery}; use datafusion_sql::TableReference; use session::context::QueryContext; @@ -594,45 +483,6 @@ mod tests { ); } - #[test] - fn test_insert_assignment_uses_query_timezone_and_target_precision() { - use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; - - let query_ctx = QueryContext::arc(); - query_ctx.set_timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()); - let mut converter = InsertAssignmentConverter { - query_ctx, - input_literals: None, - }; - - let expr = Expr::Cast(Cast { - expr: Box::new("2009-02-13 23:31:30.123456789".lit()), - data_type: DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), - }); - assert_eq!( - converter.f_up(expr).unwrap().data, - ScalarValue::TimestampNanosecond(Some(1_234_539_090_123_456_789), None).lit() - ); - } - - #[test] - fn test_insert_assignment_falls_back_for_unsupported_literal() { - use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; - - let mut converter = InsertAssignmentConverter { - query_ctx: QueryContext::arc(), - input_literals: None, - }; - for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] { - let expr = Expr::Cast(Cast { - expr: Box::new(literal.lit()), - data_type: DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), - }); - - assert_eq!(converter.f_up(expr.clone()).unwrap().data, expr); - } - } - #[test] fn test_type_converter_leaves_explicit_timestamp_cast_unchanged() { use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; diff --git a/src/query/src/optimizer/type_conversion/insert_assignment.rs b/src/query/src/optimizer/type_conversion/insert_assignment.rs new file mode 100644 index 00000000000..11ab7e7829e --- /dev/null +++ b/src/query/src/optimizer/type_conversion/insert_assignment.rs @@ -0,0 +1,336 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use datafusion_common::{Column, Result, ScalarValue}; +use datafusion_expr::expr::{Alias, Cast}; +use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Union, Values}; +use datatypes::arrow::datatypes::DataType; +use session::context::QueryContextRef; + +use crate::optimizer::type_conversion::cast_string_to_timestamp; +use crate::plan::ExtractExpr; + +/// Rewrites string literals that feed timestamp columns at an INSERT boundary. +/// Constants are folded at the assignment to avoid changing source types; +/// `VALUES` and `UNION` inputs are rewritten per row or branch. Explicit casts +/// stay on DataFusion's existing path. +pub(super) fn rewrite_insert_assignments( + plan: LogicalPlan, + query_ctx: QueryContextRef, +) -> Result { + let LogicalPlan::Projection(assignment) = plan else { + return Ok(plan); + }; + + let converter = InsertAssignmentConverter { query_ctx }; + let mut exprs = assignment.expr.clone(); + let mut input = assignment.input.as_ref().clone(); + let mut changed = false; + for (output_idx, expr) in assignment.expr.iter().enumerate() { + let target_type = assignment.schema.field(output_idx).data_type(); + if !matches!(target_type, DataType::Timestamp(_, _)) { + continue; + } + + let Some(column) = assignment_input_column(expr) else { + continue; + }; + let Some(input_idx) = input.schema().maybe_index_of_column(column) else { + continue; + }; + + let literal = lineage_literal(&input, input_idx).cloned(); + if let Some(literal) = literal + && let Some(folded) = converter.convert_literal(&literal, target_type) + { + let (qualifier, field) = assignment.schema.qualified_field(output_idx); + exprs[output_idx] = folded.alias_qualified(qualifier.cloned(), field.name()); + changed = true; + continue; + } + + // A hand-built DML plan (e.g. from flow) may share one source column + // between targets; rewriting it in place would retype every reader. + if assignment + .expr + .iter() + .enumerate() + .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) + { + continue; + } + if let Some(rewritten) = converter.rewrite_output_column(&input, input_idx, target_type)? { + input = rewritten; + changed = true; + } + } + + if !changed { + return Ok(LogicalPlan::Projection(assignment)); + } + Projection::try_new(exprs, Arc::new(input)).map(LogicalPlan::Projection) +} + +/// Resolves a literal when every row carries the same value at `output_idx`. +fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { + if output_idx >= plan.schema().fields().len() { + return None; + } + + match plan { + LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { + Expr::Literal(value, _) => Some(value), + Expr::Column(column) => { + let input_idx = projection.input.schema().maybe_index_of_column(column)?; + lineage_literal(projection.input.as_ref(), input_idx) + } + _ => None, + }, + // These nodes drop, reorder or deduplicate rows without touching the + // value a surviving row carries, and their schema stays positional. + LogicalPlan::Filter(_) + | LogicalPlan::Sort(_) + | LogicalPlan::Limit(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Distinct(Distinct::All(_)) => { + let inputs = plan.inputs(); + let [input] = inputs.as_slice() else { + return None; + }; + lineage_literal(input, output_idx) + } + _ => None, + } +} + +struct InsertAssignmentConverter { + query_ctx: QueryContextRef, +} + +impl InsertAssignmentConverter { + fn rewrite_output_column( + &self, + plan: &LogicalPlan, + output_idx: usize, + target_type: &DataType, + ) -> Result> { + if output_idx >= plan.schema().fields().len() { + return Ok(None); + } + + match plan { + // DataFusion pushes INSERT assignment casts into Values, so inspect + // them even though the Values schema already has the target type. + LogicalPlan::Values(values) => self.rewrite_values(values, output_idx, target_type), + // Once the source query has produced the target type, its casts and + // coercions belong to the source query rather than INSERT assignment. + _ if plan.schema().field(output_idx).data_type() == target_type => Ok(None), + LogicalPlan::Projection(projection) => { + self.rewrite_projection(projection, output_idx, target_type) + } + LogicalPlan::Union(union) => self.rewrite_union(union, output_idx, target_type), + // Filter and Sort are excluded here: their predicates and keys read + // the retyped column, which would change the source query. Constants + // still reach the assignment through `lineage_literal`. Distinct::All + // holds no expressions; allowing it shifts dedup keys from raw + // strings to parsed instants. + LogicalPlan::Limit(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Distinct(Distinct::All(_)) => { + self.rewrite_passthrough(plan, output_idx, target_type) + } + _ => Ok(None), + } + } + + fn rewrite_projection( + &self, + projection: &Projection, + output_idx: usize, + target_type: &DataType, + ) -> Result> { + match unalias(&projection.expr[output_idx]) { + Expr::Literal(value, _) => { + let Some(converted) = self.convert_literal(value, target_type) else { + return Ok(None); + }; + let (qualifier, field) = projection.schema.qualified_field(output_idx); + let mut exprs = projection.expr.clone(); + exprs[output_idx] = converted.alias_qualified(qualifier.cloned(), field.name()); + + Projection::try_new(exprs, projection.input.clone()) + .map(LogicalPlan::Projection) + .map(Some) + } + Expr::Column(column) => { + // Retyping the input column affects every output column reading + // it, so only follow lineage with a single consumer. + if projection + .expr + .iter() + .enumerate() + .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) + { + return Ok(None); + } + + let input = projection.input.as_ref(); + let Some(input_idx) = input.schema().maybe_index_of_column(column) else { + return Ok(None); + }; + let Some(rewritten) = self.rewrite_output_column(input, input_idx, target_type)? + else { + return Ok(None); + }; + + Projection::try_new(projection.expr.clone(), Arc::new(rewritten)) + .map(LogicalPlan::Projection) + .map(Some) + } + _ => Ok(None), + } + } + + fn rewrite_values( + &self, + values: &Values, + output_idx: usize, + target_type: &DataType, + ) -> Result> { + let mut rewritten = values.clone(); + let mut changed = false; + for row in &mut rewritten.values { + let Some(expr) = row.get_mut(output_idx) else { + return Ok(None); + }; + + if let Expr::Cast(Cast { + expr: inner, + data_type, + }) = expr + && data_type == target_type + && let Expr::Literal(value, _) = inner.as_ref() + && let Some(value) = self.convert_literal(value, target_type) + { + *expr = value; + changed = true; + continue; + } + + if expr.get_type(values.schema.as_ref()).ok().as_ref() != Some(target_type) { + return Ok(None); + } + } + + Ok(changed.then_some(LogicalPlan::Values(rewritten))) + } + + fn rewrite_union( + &self, + union: &Union, + output_idx: usize, + target_type: &DataType, + ) -> Result> { + let mut inputs = Vec::with_capacity(union.inputs.len()); + // A union has one shared schema, so rewrite it only when the same output + // column is an INSERT-assigned literal in every branch. + for input in &union.inputs { + let Some(rewritten) = self.rewrite_output_column(input, output_idx, target_type)? + else { + return Ok(None); + }; + inputs.push(Arc::new(rewritten)); + } + + // Untouched columns may still have mismatched branch types before + // TypeCoercion runs, so rebuild loosely like the SQL planner does. + Union::try_new_with_loose_types(inputs) + .map(LogicalPlan::Union) + .map(Some) + } + + fn rewrite_passthrough( + &self, + plan: &LogicalPlan, + output_idx: usize, + target_type: &DataType, + ) -> Result> { + let inputs = plan.inputs(); + let [input] = inputs.as_slice() else { + return Ok(None); + }; + let Some(rewritten_input) = self.rewrite_output_column(input, output_idx, target_type)? + else { + return Ok(None); + }; + + plan.with_new_exprs(plan.expressions_consider_join(), vec![rewritten_input]) + .map(Some) + } + + fn convert_literal(&self, value: &ScalarValue, target_type: &DataType) -> Option { + let ScalarValue::Utf8(Some(value)) = value else { + return None; + }; + cast_string_to_timestamp(value, target_type, Some(&self.query_ctx.timezone())) + .ok() + .filter(|value| !value.is_null()) + .map(|value| Expr::Literal(value, None)) + } +} + +fn assignment_input_column(expr: &Expr) -> Option<&Column> { + let expr = match unalias(expr) { + Expr::Cast(Cast { expr, .. }) => expr.as_ref(), + expr => expr, + }; + let Expr::Column(column) = expr else { + return None; + }; + Some(column) +} + +fn unalias(expr: &Expr) -> &Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias(expr), + expr => expr, + } +} + +#[cfg(test)] +mod tests { + use datafusion_common::ScalarValue; + use datafusion_common::arrow::datatypes::TimeUnit; + use session::context::QueryContext; + + use super::*; + + #[test] + fn test_convert_literal_falls_back_for_unsupported_literal() { + let converter = InsertAssignmentConverter { + query_ctx: QueryContext::arc(), + }; + let target_type = DataType::Timestamp(TimeUnit::Nanosecond, None); + + for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] { + assert_eq!( + converter + .convert_literal(&ScalarValue::Utf8(Some(literal.to_string())), &target_type,), + None + ); + } + } +} diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index c941839c0ad..a7d4a84cb28 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -711,12 +711,18 @@ mod tests { ConcreteDataType::timestamp_millisecond_datatype(), false, ), + ColumnSchema::new("note", ConcreteDataType::string_datatype(), true), + ColumnSchema::new( + "ts_ns", + ConcreteDataType::timestamp_nanosecond_datatype(), + true, + ), ]; let schema = Arc::new(Schema::new(columns)); let table_meta = TableMetaBuilder::empty() .schema(schema) .primary_key_indices(vec![]) - .value_indices(vec![0, 1]) + .value_indices(vec![0, 1, 2, 3]) .next_column_id(1024) .build() .unwrap(); @@ -869,16 +875,69 @@ mod tests { ); let engine = create_timestamp_test_engine().await; - for (sql, expected_timestamp) in [ + for (sql, expected_timestamps) in [ ( "INSERT INTO timestamps (ts, st) \ VALUES ('2026-08-02 12:00:00.001', now())", - 1_785_643_200_001_i64, + &[1_785_643_200_001_i64][..], ), ( "INSERT INTO timestamps (ts, st) \ SELECT '2026-08-03 12:00:00.001', now()", - 1_785_729_600_001_i64, + &[1_785_729_600_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-04 12:00:00.001', now() LIMIT 1", + &[1_785_816_000_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT * FROM (\ + SELECT '2026-08-05 12:00:00.001', now()\ + ) AS source", + &[1_785_902_400_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-06 12:00:00.001', now() \ + UNION ALL \ + SELECT '2026-08-07 12:00:00.001', now()", + &[1_785_988_800_001_i64, 1_786_075_200_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-16 12:00:00.001', now() \ + UNION \ + SELECT '2026-08-17 12:00:00.001', now()", + &[1_786_852_800_001_i64, 1_786_939_200_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-18 12:00:00.001', max(st) \ + FROM timestamps GROUP BY note", + &[1_787_025_600_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT c, s FROM (\ + SELECT '2026-08-13 12:00:00.001' AS c, now() AS s\ + ) AS t WHERE c > '2026-01-01'", + &[1_786_593_600_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT c, s FROM (\ + SELECT '2026-08-14 12:00:00.001' AS c, now() AS s\ + ) AS t ORDER BY c", + &[1_786_680_000_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT DISTINCT c, s FROM (\ + SELECT '2026-08-15 12:00:00.001' AS c, now() AS s\ + ) AS t", + &[1_786_766_400_001_i64][..], ), ] { let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); @@ -890,6 +949,112 @@ mod tests { .display_indent() .to_string(); + for expected_timestamp in expected_timestamps { + assert!( + plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")), + "{plan}" + ); + } + } + } + + #[tokio::test] + async fn test_insert_explicit_timestamp_cast_keeps_datafusion_semantics() { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) + .build(), + ); + let engine = create_timestamp_test_engine().await; + let sql = "INSERT INTO timestamps (ts, st) \ + VALUES (CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), now())"; + let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx) + .await + .unwrap() + .display_indent() + .to_string(); + + assert!( + plan.contains("arrow_cast(Utf8(\"2026-08-08 12:00:00.001\")"), + "{plan}" + ); + } + + #[tokio::test] + async fn test_insert_converts_source_literal_shared_by_several_columns() { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) + .build(), + ); + let engine = create_timestamp_test_engine().await; + + // Rewriting `c` in place would also retype `note` and truncate `ts_ns`. + for (sql, expected) in [ + ( + "INSERT INTO timestamps (ts, note) SELECT a, b FROM (\ + SELECT c AS a, c AS b FROM (\ + SELECT '2026-08-12 12:00:00.001' AS c\ + ) AS t1\ + ) AS t2", + &[ + "TimestampMillisecond(1786507200001, None)", + "[a:Utf8, b:Utf8]", + ][..], + ), + ( + "INSERT INTO timestamps (ts, ts_ns) SELECT a, b FROM (\ + SELECT c AS a, c AS b FROM (\ + SELECT '2026-08-12 12:00:00.123456789' AS c\ + ) AS t1\ + ) AS t2", + &[ + "TimestampMillisecond(1786507200123, None)", + "TimestampNanosecond(1786507200123456789, None)", + ][..], + ), + ] { + let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap() + .display_indent_schema() + .to_string(); + + for expected in expected { + assert!(plan.contains(expected), "{plan}"); + } + } + } + + #[tokio::test] + async fn test_insert_union_tolerates_uncoerced_untouched_column() { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) + .build(), + ); + let engine = create_timestamp_test_engine().await; + // `st` stays Timestamp vs Null across branches until TypeCoercion runs. + let sql = "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-06 12:00:00.001', now() \ + UNION ALL \ + SELECT '2026-08-07 12:00:00.001', NULL"; + let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx) + .await + .unwrap() + .display_indent() + .to_string(); + + for expected_timestamp in [1_785_988_800_001_i64, 1_786_075_200_001_i64] { assert!( plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")), "{plan}" @@ -897,6 +1062,36 @@ mod tests { } } + #[tokio::test] + async fn test_insert_mixed_union_keeps_source_coercion() { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) + .build(), + ); + let engine = create_timestamp_test_engine().await; + let sql = "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-10 12:00:00.001', now() \ + UNION ALL \ + SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now()"; + let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx) + .await + .unwrap() + .display_indent_schema() + .to_string(); + assert!( + !plan.contains("TimestampMillisecond(1786334400001, None)"), + "{plan}" + ); + assert!( + plan.contains("arrow_cast(Utf8(\"2026-08-11 12:00:00.001\")"), + "{plan}" + ); + } + #[tokio::test] async fn test_extract_placeholder_cast_types_multiple() { let plan = parse_sql_to_plan( diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.result b/tests/cases/standalone/common/insert/insert_default_timezone.result index da8b1261e7a..413eb5d46ae 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.result +++ b/tests/cases/standalone/common/insert/insert_default_timezone.result @@ -68,10 +68,48 @@ INSERT INTO test3 (ts, st) SELECT '2026-08-03 12:00:00.001', now(); Affected Rows: 1 +INSERT INTO test3 (ts, st) SELECT '2026-08-04 12:00:00.001', now() LIMIT 1; + +Affected Rows: 1 + +INSERT INTO test3 (ts, st) +SELECT '2026-08-06 12:00:00.001', now() +UNION ALL +SELECT '2026-08-07 12:00:00.001', now(); + +Affected Rows: 2 + +INSERT INTO test3 (ts, st) VALUES ( + CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), + now() +); + +Affected Rows: 1 + +INSERT INTO test3 (ts, st) +SELECT '2026-08-10 12:00:00.001', now() +UNION ALL +SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now(); + +Affected Rows: 2 + +INSERT INTO test3 (ts, st) +SELECT '2026-08-16 12:00:00.001', now() +UNION +SELECT '2026-08-17 12:00:00.001', now(); + +Affected Rows: 2 + +INSERT INTO test3 (ts, ts_ns) SELECT a, b FROM ( + SELECT c AS a, c AS b FROM (SELECT '2026-08-12 12:00:00.123456789' AS c) AS t1 +) AS t2; + +Affected Rows: 1 + INSERT INTO test3 (ts, st, ts_ns) VALUES ( - '2026-08-04 12:00:00.001', + '2026-08-09 12:00:00.001', now(), - '2026-08-04 12:00:00.123456789' + '2026-08-09 12:00:00.123456789' ); Affected Rows: 1 @@ -84,7 +122,16 @@ SELECT ts, ts_ns FROM test3 ORDER BY ts; | 2026-08-01T04:00:00.001 | | | 2026-08-02T04:00:00.001 | | | 2026-08-03T04:00:00.001 | | -| 2026-08-04T04:00:00.001 | 2026-08-04T04:00:00.123456789 | +| 2026-08-04T04:00:00.001 | | +| 2026-08-06T04:00:00.001 | | +| 2026-08-07T04:00:00.001 | | +| 2026-08-08T12:00:00.001 | | +| 2026-08-09T04:00:00.001 | 2026-08-09T04:00:00.123456789 | +| 2026-08-10T12:00:00.001 | | +| 2026-08-11T12:00:00.001 | | +| 2026-08-12T04:00:00.123 | 2026-08-12T04:00:00.123456789 | +| 2026-08-16T04:00:00.001 | | +| 2026-08-17T04:00:00.001 | | +-------------------------+-------------------------------+ SET time_zone = 'UTC'; diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.sql b/tests/cases/standalone/common/insert/insert_default_timezone.sql index fba8c134fa0..7cbef15d6d9 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.sql +++ b/tests/cases/standalone/common/insert/insert_default_timezone.sql @@ -24,10 +24,36 @@ INSERT INTO test3 (ts, st) VALUES ('2026-08-02 12:00:00.001', now()); INSERT INTO test3 (ts, st) SELECT '2026-08-03 12:00:00.001', now(); +INSERT INTO test3 (ts, st) SELECT '2026-08-04 12:00:00.001', now() LIMIT 1; + +INSERT INTO test3 (ts, st) +SELECT '2026-08-06 12:00:00.001', now() +UNION ALL +SELECT '2026-08-07 12:00:00.001', now(); + +INSERT INTO test3 (ts, st) VALUES ( + CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), + now() +); + +INSERT INTO test3 (ts, st) +SELECT '2026-08-10 12:00:00.001', now() +UNION ALL +SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now(); + +INSERT INTO test3 (ts, st) +SELECT '2026-08-16 12:00:00.001', now() +UNION +SELECT '2026-08-17 12:00:00.001', now(); + +INSERT INTO test3 (ts, ts_ns) SELECT a, b FROM ( + SELECT c AS a, c AS b FROM (SELECT '2026-08-12 12:00:00.123456789' AS c) AS t1 +) AS t2; + INSERT INTO test3 (ts, st, ts_ns) VALUES ( - '2026-08-04 12:00:00.001', + '2026-08-09 12:00:00.001', now(), - '2026-08-04 12:00:00.123456789' + '2026-08-09 12:00:00.123456789' ); SELECT ts, ts_ns FROM test3 ORDER BY ts; From d7f1233f775d54380318d6ad9ff62504a7cbcff1 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Mon, 17 Aug 2026 11:24:44 +0000 Subject: [PATCH 082/126] refactor(json2): support JSON2 storage layout settings in DDL (#8895) * refactor(json2): add JSON2 storage layout settings Signed-off-by: luofucong * resolve PR comments Signed-off-by: luofucong --------- Signed-off-by: luofucong --- src/datatypes/src/error.rs | 8 + src/datatypes/src/extension/json.rs | 1 + src/datatypes/src/json.rs | 286 +++++++++++++++--- src/sql/src/error.rs | 11 +- src/sql/src/parsers/create_parser.rs | 4 +- src/sql/src/parsers/create_parser/json.rs | 154 ++++++++-- src/sql/src/statements.rs | 18 +- src/sql/src/statements/create.rs | 144 +++++---- .../json2_storage_layout_settings/case.toml | 8 + .../json2_storage_layout_settings/setup.sql | 11 + .../verify.result | 27 ++ .../json2_storage_layout_settings/verify.sql | 4 + 12 files changed, 532 insertions(+), 144 deletions(-) create mode 100644 tests/compatibility/cases/json2_storage_layout_settings/case.toml create mode 100644 tests/compatibility/cases/json2_storage_layout_settings/setup.sql create mode 100644 tests/compatibility/cases/json2_storage_layout_settings/verify.result create mode 100644 tests/compatibility/cases/json2_storage_layout_settings/verify.sql diff --git a/src/datatypes/src/error.rs b/src/datatypes/src/error.rs index cf76812c215..b34a5ffa9bf 100644 --- a/src/datatypes/src/error.rs +++ b/src/datatypes/src/error.rs @@ -196,6 +196,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid JSON2 layout: {reason}"))] + InvalidJson2Layout { + reason: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid Vector: {}", msg))] InvalidVector { msg: String, @@ -318,6 +325,7 @@ impl ErrorExt for Error { | InvalidTimestampPrecision { .. } | InvalidPrecisionOrScale { .. } | InvalidJson { .. } + | InvalidJson2Layout { .. } | InvalidJsonb { .. } | InvalidVector { .. } | InvalidFulltextOption { .. } diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index 26d470641f5..109aa853b21 100644 --- a/src/datatypes/src/extension/json.rs +++ b/src/datatypes/src/extension/json.rs @@ -22,6 +22,7 @@ use arrow_schema::{ArrowError, DataType, Field}; use serde::{Deserialize, Serialize}; use snafu::ResultExt; +pub use crate::json::JSON2_REMAINDER_FIELD_NAME; use crate::json::JsonSettings; const LEGACY_JSON_STRUCTURE_SETTINGS_KEY: &str = "json_structure_settings"; diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index 48a9174020e..ec9ca4ac789 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -29,21 +29,47 @@ use serde_json::{Map, Value as Json}; use snafu::ResultExt; use crate::data_type::ConcreteDataType; -use crate::error::{self, Result, UnsupportedJsonTypeSnafu}; +use crate::error::{self, InvalidJson2LayoutSnafu, Result, UnsupportedJsonTypeSnafu}; use crate::json::value::{JsonValue, JsonVariant, encode_serde_json_as_jsonb}; use crate::schema::ColumnDefaultConstraint; -use crate::types::json_type::JsonNativeType; +use crate::types::json_type::{JsonNativeType, JsonObjectType}; use crate::value::{ListValue, StructValue, Value}; /// Maximum number of JSON container levels represented as nested Arrow types. pub const JSON2_MAX_STRUCTURED_DEPTH: usize = 50; +/// Reserved physical field containing unexpanded JSON2 paths. +pub const JSON2_REMAINDER_FIELD_NAME: &str = "!__remainder__!"; /// JSON2 settings stored in column schema metadata and represented through /// Arrow extension metadata. -#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)] pub struct JsonSettings { #[serde(default)] - pub type_hints: Vec, + type_hints: Vec, + /// Maximum number of unhinted JSON paths expanded into Arrow fields. + /// + /// `None` preserves the legacy unlimited layout. + #[serde(default, skip_serializing_if = "Option::is_none")] + max_auto_expanded_paths: Option, +} + +#[derive(Deserialize)] +struct JsonSettingsSerde { + #[serde(default)] + type_hints: Vec, + #[serde(default)] + max_auto_expanded_paths: Option, +} + +impl<'de> Deserialize<'de> for JsonSettings { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + let settings = JsonSettingsSerde::deserialize(deserializer)?; + Self::try_new(settings.type_hints, settings.max_auto_expanded_paths) + .map_err(serde::de::Error::custom) + } } /// Declares selected JSON2 subpaths as typed fields. @@ -79,8 +105,28 @@ pub struct JsonContext<'a> { } impl JsonSettings { - pub fn new(type_hints: Vec) -> Self { - Self { type_hints } + /// Creates and validates JSON2 settings. + pub fn try_new( + type_hints: Vec, + max_auto_expanded_paths: Option, + ) -> Result { + validate_type_hints(&type_hints)?; + Ok(Self { + type_hints, + max_auto_expanded_paths, + }) + } + + pub fn type_hints(&self) -> &[JsonTypeHint] { + &self.type_hints + } + + pub fn max_auto_expanded_paths(&self) -> Option { + self.max_auto_expanded_paths + } + + pub fn into_parts(self) -> (Vec, Option) { + (self.type_hints, self.max_auto_expanded_paths) } /// Decode an encoded StructValue back into a serde_json::Value. @@ -102,6 +148,87 @@ impl JsonSettings { } } +fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { + let mut object = JsonObjectType::new(); + for hint in type_hints { + if hint.path.len() > JSON2_MAX_STRUCTURED_DEPTH { + return InvalidJson2LayoutSnafu { + reason: format!( + "JSON2 type hint path cannot exceed {JSON2_MAX_STRUCTURED_DEPTH} segments" + ), + } + .fail(); + } + if hint + .path + .first() + .is_some_and(|x| x == JSON2_REMAINDER_FIELD_NAME) + { + return InvalidJson2LayoutSnafu { + reason: format!( + "JSON2 type hint path cannot be rooted at reserved field '{JSON2_REMAINDER_FIELD_NAME}'" + ), + } + .fail(); + } + let data_type = match &hint.data_type { + ConcreteDataType::Boolean(_) + | ConcreteDataType::UInt8(_) + | ConcreteDataType::UInt16(_) + | ConcreteDataType::UInt32(_) + | ConcreteDataType::UInt64(_) + | ConcreteDataType::Int8(_) + | ConcreteDataType::Int16(_) + | ConcreteDataType::Int32(_) + | ConcreteDataType::Int64(_) + | ConcreteDataType::Float32(_) + | ConcreteDataType::Float64(_) + | ConcreteDataType::String(_) => (&hint.data_type).into(), + data_type => { + return InvalidJson2LayoutSnafu { + reason: format!("unsupported JSON2 type hint data type: {data_type}"), + } + .fail(); + } + }; + validate_type_hint(&mut object, &hint.path, data_type)?; + } + Ok(()) +} + +fn validate_type_hint( + object: &mut JsonObjectType, + path: &[String], + data_type: JsonNativeType, +) -> Result<()> { + let Some((name, path)) = path.split_first() else { + return InvalidJson2LayoutSnafu { + reason: "JSON2 type hint path must not be empty".to_string(), + } + .fail(); + }; + if path.is_empty() { + if object.insert(name.clone(), data_type).is_some() { + return InvalidJson2LayoutSnafu { + reason: format!("duplicate JSON2 type hint path '{name}'"), + } + .fail(); + } + return Ok(()); + } + + let child = object + .entry(name.clone()) + .or_insert_with(|| JsonNativeType::Object(JsonObjectType::new())); + let JsonNativeType::Object(child) = child else { + return InvalidJson2LayoutSnafu { + reason: format!("conflicting JSON2 type hint path at '{name}'"), + } + .fail(); + }; + validate_type_hint(child, path, data_type) +} + impl<'a> JsonContext<'a> { fn type_hint(&self) -> Option<&'a JsonTypeHint> { self.settings @@ -483,7 +610,8 @@ mod tests { } #[test] - fn test_json_settings_forward_compatibility() { + fn test_json_settings_forward_compatibility() + -> std::result::Result<(), Box> { let json_str = r#"{ "type_hints": [ { @@ -512,11 +640,10 @@ mod tests { ] }"#; - let deserialized = serde_json::from_str::(json_str).unwrap(); + let deserialized = serde_json::from_str::(json_str)?; - assert_eq!( - deserialized, - JsonSettings::new(vec![ + let expected = JsonSettings::try_new( + vec![ JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), @@ -531,13 +658,16 @@ mod tests { default_constraint: None, inverted_index: false, }, - ]) - ); + ], + None, + )?; + assert_eq!(deserialized, expected); + Ok(()) } #[test] - fn test_json_settings_ser_de() { - let settings = JsonSettings::new(vec![ + fn test_json_settings_ser_de() -> std::result::Result<(), Box> { + let type_hints = vec![ JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), @@ -552,12 +682,76 @@ mod tests { default_constraint: None, inverted_index: false, }, - ]); + ]; - let serialized = serde_json::to_string(&settings).unwrap(); - let deserialized = serde_json::from_str::(&serialized).unwrap(); + for settings in [ + JsonSettings { + type_hints: type_hints.clone(), + max_auto_expanded_paths: None, + }, + JsonSettings { + type_hints, + max_auto_expanded_paths: Some(1), + }, + ] { + let serialized = serde_json::to_string(&settings)?; + let deserialized = serde_json::from_str::(&serialized)?; + assert_eq!(settings, deserialized); + } + Ok(()) + } - assert_eq!(settings, deserialized); + #[test] + fn test_json_settings_reject_invalid_type_hint_layout() { + for type_hints in [ + json!([{"path": [], "type": {"Int64": {}}, "nullable": true, "inverted_index": false}]), + json!([ + {"path": ["a"], "type": {"Int64": {}}, "nullable": true, "inverted_index": false}, + {"path": ["a", "b"], "type": {"Int64": {}}, "nullable": true, "inverted_index": false} + ]), + json!([{"path": [JSON2_REMAINDER_FIELD_NAME], "type": {"Int64": {}}, "nullable": true, "inverted_index": false}]), + ] { + let settings = json!({"type_hints": type_hints}); + assert!(serde_json::from_value::(settings).is_err()); + } + + let hint = |path, data_type| JsonTypeHint { + path, + data_type, + nullable: true, + default_constraint: None, + inverted_index: false, + }; + assert!( + JsonSettings::try_new( + vec![hint( + vec!["nested".to_string(), JSON2_REMAINDER_FIELD_NAME.to_string()], + ConcreteDataType::string_datatype(), + )], + None, + ) + .is_ok() + ); + assert!( + JsonSettings::try_new( + vec![hint( + vec!["nested".to_string(); JSON2_MAX_STRUCTURED_DEPTH + 1], + ConcreteDataType::string_datatype(), + )], + None, + ) + .is_err() + ); + assert!( + JsonSettings::try_new( + vec![hint( + vec!["date".to_string()], + ConcreteDataType::date_datatype(), + )], + None, + ) + .is_err() + ); } #[test] @@ -702,14 +896,15 @@ mod tests { } #[test] - fn test_encode_json_respects_type_hint() { - let settings = JsonSettings::new(vec![JsonTypeHint { + fn test_encode_json_respects_type_hint() -> Result<()> { + let type_hints = vec![JsonTypeHint { path: vec!["age".to_string()], data_type: ConcreteDataType::int64_datatype(), nullable: false, default_constraint: None, inverted_index: false, - }]); + }]; + let settings = JsonSettings::try_new(type_hints, None)?; let result = settings .encode(json!({ @@ -731,17 +926,19 @@ mod tests { })) .unwrap_err(); assert!(err.to_string().contains("does not match JSON2 type hint")); + Ok(()) } #[test] - fn test_encode_json_respects_unsigned_type_hint() { - let settings = JsonSettings::new(vec![JsonTypeHint { + fn test_encode_json_respects_unsigned_type_hint() -> Result<()> { + let type_hints = vec![JsonTypeHint { path: vec!["count".to_string()], data_type: ConcreteDataType::uint64_datatype(), nullable: false, default_constraint: None, inverted_index: false, - }]); + }]; + let settings = JsonSettings::try_new(type_hints, None)?; let result = settings .encode(json!({ @@ -765,17 +962,19 @@ mod tests { })) .unwrap_err(); assert!(err.to_string().contains("does not match JSON2 type hint")); + Ok(()) } #[test] - fn test_encode_json_fills_missing_type_hint_with_default() { - let settings = JsonSettings::new(vec![JsonTypeHint { + fn test_encode_json_fills_missing_type_hint_with_default() -> Result<()> { + let type_hints = vec![JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), nullable: false, default_constraint: Some(ColumnDefaultConstraint::Value(Value::Int64(7))), inverted_index: false, - }]); + }]; + let settings = JsonSettings::try_new(type_hints, None)?; let result = settings .encode(json!({})) @@ -790,17 +989,19 @@ mod tests { panic!("Expected user Struct value"); }; assert_eq!(struct_field_value(user, "age"), &Value::Int64(7)); + Ok(()) } #[test] - fn test_encode_json_fills_missing_nullable_type_hint_with_null() { - let settings = JsonSettings::new(vec![JsonTypeHint { + fn test_encode_json_fills_missing_nullable_type_hint_with_null() -> Result<()> { + let type_hints = vec![JsonTypeHint { path: vec!["user".to_string(), "name".to_string()], data_type: ConcreteDataType::string_datatype(), nullable: true, default_constraint: None, inverted_index: false, - }]); + }]; + let settings = JsonSettings::try_new(type_hints, None)?; let result = settings .encode(json!({ "user": {} })) @@ -815,28 +1016,31 @@ mod tests { panic!("Expected user Struct value"); }; assert_eq!(struct_field_value(user, "name"), &Value::Null); + Ok(()) } #[test] - fn test_encode_json_rejects_missing_non_null_type_hint() { - let settings = JsonSettings::new(vec![JsonTypeHint { + fn test_encode_json_rejects_missing_non_null_type_hint() -> Result<()> { + let type_hints = vec![JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), nullable: false, default_constraint: None, inverted_index: false, - }]); + }]; + let settings = JsonSettings::try_new(type_hints, None)?; let err = settings.encode(json!({})).unwrap_err(); assert!( err.to_string() .contains("missing non-null JSON2 type hint path user.age") ); + Ok(()) } #[test] - fn test_encode_json_merges_missing_type_hint_prefix() { - let settings = JsonSettings::new(vec![ + fn test_encode_json_merges_missing_type_hint_prefix() -> Result<()> { + let type_hints = vec![ JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), @@ -853,13 +1057,10 @@ mod tests { ))), inverted_index: false, }, - ]); + ]; + let settings = JsonSettings::try_new(type_hints, None)?; - let result = settings - .encode(json!({})) - .unwrap() - .into_json_inner() - .unwrap(); + let result = settings.encode(json!({}))?.into_json_inner().unwrap(); let Value::Struct(root) = result else { panic!("Expected Struct value"); @@ -872,6 +1073,7 @@ mod tests { struct_field_value(user, "name"), &Value::String("unknown".into()) ); + Ok(()) } #[test] diff --git a/src/sql/src/error.rs b/src/sql/src/error.rs index cc0b9f19506..dc819bda017 100644 --- a/src/sql/src/error.rs +++ b/src/sql/src/error.rs @@ -339,6 +339,13 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + + #[snafu(transparent)] + Datatypes { + #[snafu(implicit)] + location: Location, + source: datatypes::error::Error, + }, } impl ErrorExt for Error { @@ -384,7 +391,9 @@ impl ErrorExt for Error { #[cfg(feature = "enterprise")] InvalidTriggerWebhookOption { .. } => StatusCode::InvalidArguments, - SerializeColumnDefaultConstraint { source, .. } => source.status_code(), + SerializeColumnDefaultConstraint { source, .. } | Datatypes { source, .. } => { + source.status_code() + } ConvertToGrpcDataType { source, .. } => source.status_code(), SqlCommon { source, .. } => source.status_code(), diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index a257f029e35..be7ed0670dc 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -711,8 +711,8 @@ impl<'a> ParserContext<'a> { let mut extensions = ColumnExtensions::default(); let data_type = - if let Some((data_type, type_hints)) = json::parse_json2_type_and_hints(parser)? { - extensions.json_type_hints = type_hints; + if let Some((data_type, options)) = json::parse_json2_type_and_options(parser)? { + extensions.json2_options = options; data_type } else { parser.parse_data_type().context(SyntaxSnafu)? diff --git a/src/sql/src/parsers/create_parser/json.rs b/src/sql/src/parsers/create_parser/json.rs index e2f493719c1..63399f81a17 100644 --- a/src/sql/src/parsers/create_parser/json.rs +++ b/src/sql/src/parsers/create_parser/json.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use datatypes::extension::json::JSON2_REMAINDER_FIELD_NAME; use datatypes::json::JSON2_MAX_STRUCTURED_DEPTH; use snafu::{ResultExt, ensure}; use sqlparser::ast::{DataType, ExactNumberInfo, Expr, ObjectName, UnaryOperator}; @@ -22,14 +23,15 @@ use sqlparser::tokenizer::Token; use crate::ast::Ident; use crate::error::{InvalidSqlSnafu, Result, SyntaxSnafu}; use crate::parsers::create_parser::{INVERTED, SKIPPING}; -use crate::statements::create::JsonTypeHint; +use crate::statements::create::{Json2Options, JsonTypeHint}; use crate::statements::transform::type_alias::get_type_by_alias; const JSON2_TYPE_NAME: &str = "JSON2"; +const MAX_AUTO_EXPANDED_PATHS: &str = "max_auto_expanded_paths"; -pub(super) fn parse_json2_type_and_hints( +pub(super) fn parse_json2_type_and_options( parser: &mut Parser<'_>, -) -> Result)>> { +) -> Result)>> { let token = parser.peek_token(); let Token::Word(word) = &token.token else { return Ok(None); @@ -41,26 +43,62 @@ pub(super) fn parse_json2_type_and_hints( parser.next_token(); let data_type = DataType::Custom(ObjectName::from(vec![Ident::new(JSON2_TYPE_NAME)]), vec![]); - let type_hints = if parser.consume_token(&Token::LParen) { - parse_json2_type_hints(parser)? + let options = if parser.consume_token(&Token::LParen) { + parse_json2_options(parser)? } else { - vec![] + None }; - Ok(Some((data_type, type_hints))) + Ok(Some((data_type, options))) } -fn parse_json2_type_hints(parser: &mut Parser<'_>) -> Result> { - let mut hints = Vec::new(); - +fn parse_json2_options(parser: &mut Parser<'_>) -> Result> { if parser.consume_token(&Token::RParen) { - return Ok(hints); + return Ok(None); } + let mut max_auto_expanded_paths = None; + let mut type_hints = Vec::new(); loop { - let hint = parse_json2_type_hint(parser)?; - ensure_no_path_conflict(&hints, &hint.path)?; - hints.push(hint); + let token = parser.peek_token(); + let is_max_auto_expanded_paths = matches!( + &token.token, + Token::Word(word) + if word.quote_style.is_none() + && word.value.eq_ignore_ascii_case(MAX_AUTO_EXPANDED_PATHS) + ); + if is_max_auto_expanded_paths { + parser.next_token(); + ensure!( + max_auto_expanded_paths.is_none(), + InvalidSqlSnafu { + msg: format!("duplicated JSON2 option '{MAX_AUTO_EXPANDED_PATHS}'") + } + ); + parser.expect_token(&Token::Eq).context(SyntaxSnafu)?; + + let token = parser.next_token(); + let Token::Number(value, _) = token.token else { + return InvalidSqlSnafu { + msg: format!( + "JSON2 option '{MAX_AUTO_EXPANDED_PATHS}' expects a non-negative integer" + ), + } + .fail(); + }; + max_auto_expanded_paths = Some(value.parse::().map_err(|_| { + InvalidSqlSnafu { + msg: format!( + "JSON2 option '{MAX_AUTO_EXPANDED_PATHS}' expects a non-negative integer" + ), + } + .build() + })?); + } else { + let hint = parse_json2_type_hint(parser)?; + ensure_no_path_conflict(&type_hints, &hint.path)?; + type_hints.push(hint); + } if parser.consume_token(&Token::Comma) { if parser.consume_token(&Token::RParen) { @@ -72,11 +110,22 @@ fn parse_json2_type_hints(parser: &mut Parser<'_>) -> Result> } } - Ok(hints) + Ok(Some(Json2Options { + max_auto_expanded_paths, + type_hints, + })) } fn parse_json2_type_hint(parser: &mut Parser<'_>) -> Result { let path = parse_json2_path(parser)?; + ensure!( + path.first().is_none_or(|x| x != JSON2_REMAINDER_FIELD_NAME), + InvalidSqlSnafu { + msg: format!( + "JSON2 type hint path cannot be rooted at reserved field '{JSON2_REMAINDER_FIELD_NAME}'" + ) + } + ); ensure!( path.len() <= JSON2_MAX_STRUCTURED_DEPTH, InvalidSqlSnafu { @@ -309,7 +358,7 @@ CREATE TABLE traces ( column.column_def.data_type, DataType::Custom(_, _) )); - let hints = column.extensions.json_type_hints; + let hints = column.extensions.json2_options.unwrap().type_hints; assert_eq!(hints.len(), 4); assert_eq!(hints[0].path, vec!["service.name"]); @@ -339,6 +388,71 @@ CREATE TABLE traces ( assert!(hints[3].nullable); } + #[test] + fn test_parse_json2_max_auto_expanded_paths() { + let column = parse_json2_column( + r#" +CREATE TABLE traces ( + log_json_data JSON2 ( + http.method STRING, + max_auto_expanded_paths = 0 + ), + ts TIMESTAMP TIME INDEX, +)"#, + ); + + let options = column.extensions.json2_options.unwrap(); + assert_eq!(options.max_auto_expanded_paths, Some(0)); + assert_eq!(options.type_hints.len(), 1); + + let empty = parse_json2_column( + r#" +CREATE TABLE traces ( + log_json_data JSON2 (), + ts TIMESTAMP TIME INDEX, +)"#, + ); + assert!(empty.extensions.json2_options.is_none()); + + let quoted = parse_json2_column( + r#" +CREATE TABLE traces ( + log_json_data JSON2 ( + "max_auto_expanded_paths" STRING, + nested."!__remainder__!" STRING + ), + ts TIMESTAMP TIME INDEX, +)"#, + ); + let options = quoted.extensions.json2_options.unwrap(); + assert_eq!(options.max_auto_expanded_paths, None); + assert_eq!(options.type_hints.len(), 2); + } + + #[test] + fn test_parse_json2_max_auto_expanded_paths_rejects_invalid_options() { + for options in [ + "max_auto_expanded_paths = 0, max_auto_expanded_paths = 1", + "max_auto_expanded_paths = -1", + "max_auto_expanded_paths = 1.5", + "max_auto_expanded_paths = 4294967296", + r#""!__remainder__!".value STRING"#, + ] { + let sql = format!( + "CREATE TABLE traces (log_json_data JSON2 ({options}), ts TIMESTAMP TIME INDEX)" + ); + assert!( + ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default() + ) + .is_err(), + "{options}" + ); + } + } + #[test] fn test_parse_json2_type_hint_default_nullable() { let column = parse_json2_column( @@ -349,7 +463,7 @@ CREATE TABLE traces ( )"#, ); - let hints = column.extensions.json_type_hints; + let hints = column.extensions.json2_options.unwrap().type_hints; assert_eq!(hints.len(), 1); assert!(hints[0].nullable); } @@ -369,7 +483,7 @@ CREATE TABLE traces ( )"#, ); - let hints = column.extensions.json_type_hints; + let hints = column.extensions.json2_options.unwrap().type_hints; assert_eq!(hints.len(), 4); assert_eq!(hints[0].path, vec!["a", "b"]); assert_eq!(hints[1].path, vec!["x", "y"]); @@ -402,7 +516,7 @@ CREATE TABLE traces ( )"#, ); - let hints = column.extensions.json_type_hints; + let hints = column.extensions.json2_options.unwrap().type_hints; assert_eq!(hints.len(), 14); for hint in hints.iter().take(6) { assert_eq!(hint.data_type, DataType::BigInt(None)); @@ -428,7 +542,7 @@ CREATE TABLE traces ( )"#, ); - let hints = column.extensions.json_type_hints; + let hints = column.extensions.json2_options.unwrap().type_hints; assert_eq!(hints.len(), 2); assert_eq!( hints[0] diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index aafe62d2e6e..4f96b66f87f 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -718,11 +718,7 @@ mod tests { "true".to_string(), ), ])), - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], - vector_index_options: None, + ..Default::default() }, }; @@ -749,17 +745,13 @@ mod tests { options: vec![], }, extensions: ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::from([ ("metric".to_string(), "cosine".to_string()), ("connectivity".to_string(), "32".to_string()), ("expansion_add".to_string(), "200".to_string()), ("expansion_search".to_string(), "100".to_string()), ])), + ..Default::default() }, }; @@ -790,12 +782,8 @@ mod tests { options: vec![], }, extensions: ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::default()), + ..Default::default() }, }; diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index dc94c7f8889..f3fbc908362 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -136,7 +136,28 @@ pub struct ColumnExtensions { pub inverted_index_options: Option, /// Vector index options for HNSW-based vector similarity search. pub vector_index_options: Option, - pub json_type_hints: Vec, + /// JSON2-specific column options. + pub json2_options: Option, +} + +/// JSON2-specific options represented in the SQL AST. +#[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)] +pub struct Json2Options { + /// Maximum number of unhinted JSON2 paths expanded into Arrow fields. + pub(crate) max_auto_expanded_paths: Option, + /// Paths stored as explicitly typed JSON2 fields. + pub(crate) type_hints: Vec, +} + +impl Display for Json2Options { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let mut options = Vec::with_capacity(self.type_hints.len() + 1); + if let Some(max) = self.max_auto_expanded_paths { + options.push(format!("max_auto_expanded_paths = {max}")); + } + options.extend(self.type_hints.iter().map(format_json_type_hint)); + write!(f, "(\n {}\n )", options.iter().join(",\n ")) + } } #[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)] @@ -180,12 +201,8 @@ impl Display for Column { } write!(f, "{} {}", self.column_def.name, self.column_def.data_type)?; - if !self.extensions.json_type_hints.is_empty() { - write!( - f, - "{}", - format_json_type_hints(&self.extensions.json_type_hints) - )?; + if let Some(options) = &self.extensions.json2_options { + write!(f, "{options}")?; } for option in &self.column_def.options { write!(f, " {option}")?; @@ -335,29 +352,30 @@ impl ColumnExtensions { } pub fn build_json_settings(&self) -> Result> { - if self.json_type_hints.is_empty() { + let Some(options) = &self.json2_options else { return Ok(None); - } + }; - Ok(Some(JsonSettings::new( - self.json_type_hints - .iter() - .map(|hint| { - Ok(datatypes::json::JsonTypeHint { - path: hint.path.clone(), - data_type: json_type_hint_concrete_data_type(&hint.data_type)?, - nullable: hint.nullable, - default_constraint: build_json_type_hint_default_constraint(hint)?, - inverted_index: hint.inverted_index, - }) + let type_hints = options + .type_hints + .iter() + .map(|hint| { + Ok(datatypes::json::JsonTypeHint { + path: hint.path.clone(), + data_type: json_type_hint_concrete_data_type(&hint.data_type)?, + nullable: hint.nullable, + default_constraint: build_json_type_hint_default_constraint(hint)?, + inverted_index: hint.inverted_index, }) - .collect::>>()?, - ))) + }) + .collect::>>()?; + let settings = JsonSettings::try_new(type_hints, options.max_auto_expanded_paths)?; + Ok(Some(settings)) } pub fn set_json_settings(&mut self, settings: JsonSettings) -> Result<()> { - self.json_type_hints = settings - .type_hints + let (type_hints, max_auto_expanded_paths) = settings.into_parts(); + let type_hints = type_hints .into_iter() .map(|hint| { let data_type = json_type_hint_sql_data_type(&hint.data_type)?; @@ -374,6 +392,11 @@ impl ColumnExtensions { }) }) .collect::>>()?; + self.json2_options = (max_auto_expanded_paths.is_some() || !type_hints.is_empty()) + .then_some(Json2Options { + max_auto_expanded_paths, + type_hints, + }); Ok(()) } } @@ -487,13 +510,6 @@ fn format_json_type_hint(hint: &JsonTypeHint) -> String { ) } -fn format_json_type_hints(hints: &[JsonTypeHint]) -> String { - format!( - "(\n {}\n )", - hints.iter().map(format_json_type_hint).join(",\n ") - ) -} - fn format_json_path_segment(segment: &str) -> String { format!("\"{}\"", segment.replace('"', "\"\"")) } @@ -786,6 +802,7 @@ mod tests { use datatypes::schema::ColumnDefaultConstraint; use datatypes::value::Value; + use super::*; use crate::dialect::GreptimeDbDialect; use crate::error::Error; use crate::parser::{ParseOptions, ParserContext}; @@ -1066,7 +1083,7 @@ ENGINE=mito .build_json_settings() .unwrap() .unwrap(); - let hints = settings.type_hints; + let hints = settings.type_hints(); assert_eq!(hints[0].data_type, ConcreteDataType::int64_datatype()); assert_eq!( @@ -1118,10 +1135,10 @@ ENGINE=mito } #[test] - fn test_set_json_settings_normalizes_type_hint_sql_types() { + fn test_set_json_settings_normalizes_type_hint_sql_types() -> Result<()> { let mut extensions = super::ColumnExtensions::default(); - extensions - .set_json_settings(JsonSettings::new(vec![ + let settings = JsonSettings::try_new( + vec![ DatatypeJsonTypeHint { path: vec!["i".to_string()], data_type: ConcreteDataType::int32_datatype(), @@ -1157,36 +1174,51 @@ ENGINE=mito default_constraint: None, inverted_index: false, }, - ])) - .unwrap(); + ], + None, + )?; + extensions.set_json_settings(settings)?; assert_eq!( extensions - .json_type_hints + .json2_options + .unwrap() + .type_hints .iter() .map(|hint| hint.data_type.to_string()) .collect::>(), vec!["BIGINT", "DOUBLE", "BIGINT UNSIGNED", "STRING", "BOOLEAN"] ); + Ok(()) } #[test] - fn test_set_json_settings_rejects_unsupported_type_hint_type() { - let mut extensions = super::ColumnExtensions::default(); - let err = extensions - .set_json_settings(JsonSettings::new(vec![DatatypeJsonTypeHint { + fn test_set_json_settings_rejects_unsupported_type_hint_type() -> Result<()> { + let err = JsonSettings::try_new( + vec![DatatypeJsonTypeHint { path: vec!["u".to_string()], data_type: ConcreteDataType::date_datatype(), nullable: true, default_constraint: None, inverted_index: false, - }])) - .unwrap_err(); + }], + None, + ) + .unwrap_err(); assert!( err.to_string() .contains("unsupported JSON2 type hint data type") ); + Ok(()) + } + + #[test] + fn test_set_empty_json_settings_omits_json2_options() -> Result<()> { + let mut extensions = ColumnExtensions::default(); + extensions.set_json_settings(JsonSettings::default())?; + assert!(extensions.json2_options.is_none()); + Ok(()) } #[test] @@ -1364,15 +1396,11 @@ AS SELECT number FROM numbers_input where number > 10"#, // Test zero connectivity should fail let extensions = ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::from([( "connectivity".to_string(), "0".to_string(), )])), + ..Default::default() }; let result = extensions.build_vector_index_options(); assert!(result.is_err()); @@ -1385,15 +1413,11 @@ AS SELECT number FROM numbers_input where number > 10"#, // Test zero expansion_add should fail let extensions = ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::from([( "expansion_add".to_string(), "0".to_string(), )])), + ..Default::default() }; let result = extensions.build_vector_index_options(); assert!(result.is_err()); @@ -1406,15 +1430,11 @@ AS SELECT number FROM numbers_input where number > 10"#, // Test zero expansion_search should fail let extensions = ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::from([( "expansion_search".to_string(), "0".to_string(), )])), + ..Default::default() }; let result = extensions.build_vector_index_options(); assert!(result.is_err()); @@ -1427,16 +1447,12 @@ AS SELECT number FROM numbers_input where number > 10"#, // Test valid values should succeed let extensions = ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::from([ ("connectivity".to_string(), "32".to_string()), ("expansion_add".to_string(), "200".to_string()), ("expansion_search".to_string(), "100".to_string()), ])), + ..Default::default() }; let result = extensions.build_vector_index_options(); assert!(result.is_ok()); diff --git a/tests/compatibility/cases/json2_storage_layout_settings/case.toml b/tests/compatibility/cases/json2_storage_layout_settings/case.toml new file mode 100644 index 00000000000..8f1dd3fe64a --- /dev/null +++ b/tests/compatibility/cases/json2_storage_layout_settings/case.toml @@ -0,0 +1,8 @@ +name = "json2_storage_layout_settings" +reason = "Verify JSON2 storage layout settings in Arrow extension metadata survive restarts and upgrades." +introduced_by = "PR #8895" +topologies = ["distributed", "standalone"] +from_range = [">=v1.3.0"] +to_range = [">=v1.3.0"] +features = ["json2", "table", "metadata", "read"] +owner = "query" diff --git a/tests/compatibility/cases/json2_storage_layout_settings/setup.sql b/tests/compatibility/cases/json2_storage_layout_settings/setup.sql new file mode 100644 index 00000000000..2592185ef23 --- /dev/null +++ b/tests/compatibility/cases/json2_storage_layout_settings/setup.sql @@ -0,0 +1,11 @@ +CREATE TABLE t_json2_storage_layout_settings ( + ts TIMESTAMP TIME INDEX, + j JSON2(max_auto_expanded_paths = 0) +) WITH ( + 'append_mode' = 'true' +); + +INSERT INTO t_json2_storage_layout_settings VALUES + ('2026-08-17 00:00:00+0000', '{"a": 1}'); + +ADMIN FLUSH_TABLE('t_json2_storage_layout_settings'); diff --git a/tests/compatibility/cases/json2_storage_layout_settings/verify.result b/tests/compatibility/cases/json2_storage_layout_settings/verify.result new file mode 100644 index 00000000000..452ed7a5516 --- /dev/null +++ b/tests/compatibility/cases/json2_storage_layout_settings/verify.result @@ -0,0 +1,27 @@ +SELECT j.a AS a +FROM t_json2_storage_layout_settings; + ++---+ +| a | ++---+ +| 1 | ++---+ + +SHOW CREATE TABLE t_json2_storage_layout_settings; + ++---------------------------------+----------------------------------------------------------------+ +| Table | Create Table | ++---------------------------------+----------------------------------------------------------------+ +| t_json2_storage_layout_settings | CREATE TABLE IF NOT EXISTS "t_json2_storage_layout_settings" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 0 | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++---------------------------------+----------------------------------------------------------------+ diff --git a/tests/compatibility/cases/json2_storage_layout_settings/verify.sql b/tests/compatibility/cases/json2_storage_layout_settings/verify.sql new file mode 100644 index 00000000000..dae3cd8a235 --- /dev/null +++ b/tests/compatibility/cases/json2_storage_layout_settings/verify.sql @@ -0,0 +1,4 @@ +SELECT j.a AS a +FROM t_json2_storage_layout_settings; + +SHOW CREATE TABLE t_json2_storage_layout_settings; From b882e393dfadcf9406c76d1017fe64e0652f0202 Mon Sep 17 00:00:00 2001 From: sun Date: Tue, 18 Aug 2026 03:47:58 +0000 Subject: [PATCH 083/126] feat: update dashboard to v0.13.13 (#8898) --- src/servers/dashboard/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index 3d431cfd5f0..b68b26e6f78 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.13.12 +v0.13.13 From b96ea86a621e7283e77e4e37dd4e3ec8dfc44f73 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Wed, 19 Aug 2026 02:27:06 +0000 Subject: [PATCH 084/126] refactor(json2): add JSON2 v2 physical layout primitives (#8901) * refactor(json2): add JSON2 v2 physical layout primitives Signed-off-by: luofucong * resolve PR comments Signed-off-by: luofucong * fix ci Signed-off-by: luofucong --------- Signed-off-by: luofucong --- Cargo.lock | 50 +++++ Cargo.toml | 3 + src/datatypes/Cargo.toml | 3 + src/datatypes/src/extension/json.rs | 202 +++++++++++++++++- src/datatypes/src/vectors/json.rs | 1 + src/datatypes/src/vectors/json/array.rs | 248 ++++++++++++++++++++-- src/datatypes/src/vectors/json/variant.rs | 239 +++++++++++++++++++++ 7 files changed, 728 insertions(+), 18 deletions(-) create mode 100644 src/datatypes/src/vectors/json/variant.rs diff --git a/Cargo.lock b/Cargo.lock index 0f473d4e030..16b4c6b1045 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4554,6 +4554,9 @@ dependencies = [ "num", "num-traits", "ordered-float 4.6.0", + "parquet-variant", + "parquet-variant-compute", + "parquet-variant-json", "paste", "regex", "serde", @@ -10259,6 +10262,53 @@ dependencies = [ "zstd", ] +[[package]] +name = "parquet-variant" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74c8db065291f088a2aad8ab831853eae1871c0d311c8d0b83bbc3b7e735d0fc" +dependencies = [ + "arrow 58.3.0", + "arrow-schema 58.3.0", + "chrono", + "half", + "indexmap 2.13.0", + "num-traits", + "simdutf8", + "uuid", +] + +[[package]] +name = "parquet-variant-compute" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a530e8d5b5e14efcb39c9a6ec55432ad11f6afb7dc4455a79be0dc615fe3cc31" +dependencies = [ + "arrow 58.3.0", + "arrow-schema 58.3.0", + "chrono", + "half", + "indexmap 2.13.0", + "parquet-variant", + "parquet-variant-json", + "serde_json", + "uuid", +] + +[[package]] +name = "parquet-variant-json" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00ed89908289f67caa2ca078f9ff9aacd6229a313ec92b12bf4f48f613dc2b97" +dependencies = [ + "arrow-schema 58.3.0", + "base64 0.22.1", + "chrono", + "parquet-variant", + "serde_json", + "uuid", +] + [[package]] name = "parse-size" version = "1.1.0" diff --git a/Cargo.toml b/Cargo.toml index daee1bd4090..777d8b302b4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -195,6 +195,9 @@ otel-arrow-rust = { git = "https://github.com/GreptimeTeam/otel-arrow", rev = "5 ] } parking_lot = "0.12" parquet = { version = "58.3", default-features = false, features = ["arrow", "async", "object_store"] } +parquet-variant = "58.3" +parquet-variant-compute = "58.3" +parquet-variant-json = "58.3" paste = "1.0" pin-project = "1.0" pretty_assertions = "1.4.0" diff --git a/src/datatypes/Cargo.toml b/src/datatypes/Cargo.toml index 3bf2a5d8eec..124b1e0ea64 100644 --- a/src/datatypes/Cargo.toml +++ b/src/datatypes/Cargo.toml @@ -27,6 +27,9 @@ jsonb.workspace = true num = "0.4" num-traits = "0.2" ordered-float.workspace = true +parquet-variant.workspace = true +parquet-variant-compute.workspace = true +parquet-variant-json.workspace = true paste.workspace = true regex.workspace = true serde.workspace = true diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index 109aa853b21..ab4a138420e 100644 --- a/src/datatypes/src/extension/json.rs +++ b/src/datatypes/src/extension/json.rs @@ -18,31 +18,148 @@ use std::sync::Arc; use arrow_schema::extension::{ EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, }; -use arrow_schema::{ArrowError, DataType, Field}; +use arrow_schema::{ArrowError, DataType, Field, FieldRef}; +use parquet_variant_compute::VariantType; use serde::{Deserialize, Serialize}; -use snafu::ResultExt; +use snafu::{ResultExt, ensure}; +use crate::error::InvalidJson2LayoutSnafu; pub use crate::json::JSON2_REMAINDER_FIELD_NAME; use crate::json::JsonSettings; const LEGACY_JSON_STRUCTURE_SETTINGS_KEY: &str = "json_structure_settings"; +const JSON2_LAYOUT_V1: u8 = 1; +const JSON2_LAYOUT_V2: u8 = 2; -#[derive(Debug, Clone, Serialize, Deserialize, Default)] +/// Parsed physical layout of a JSON2 Arrow root field. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Json2PhysicalLayout { + version: u8, +} + +impl Json2PhysicalLayout { + /// Parses the JSON2 layout version from a root field. + /// + /// The version is read from the extension metadata; missing metadata + /// defaults to V1. + /// + /// Supported versions: + /// - V1 + /// - V2 + /// + /// Errors when the field is not a JSON2 extension or the version is not + /// one of the supported versions. + pub fn try_from_root(field: &Field) -> crate::error::Result { + ensure!( + is_json2_extension_type(field), + InvalidJson2LayoutSnafu { + reason: format!("field '{}' is not a JSON2 extension", field.name()), + } + ); + + let version = match field.extension_type_name() { + Some(Json2ExtensionType::NAME) => field + .metadata() + .get(EXTENSION_TYPE_METADATA_KEY) + .map(|x| parse_version(x)) + .transpose()? + .flatten() + .unwrap_or(JSON2_LAYOUT_V1), + _ => JSON2_LAYOUT_V1, + }; + ensure!( + matches!(version, JSON2_LAYOUT_V1 | JSON2_LAYOUT_V2), + InvalidJson2LayoutSnafu { + reason: format!("unsupported JSON2 layout version: {version}"), + } + ); + Ok(Self { version }) + } + + /// Returns whether this is the JSON2 physical layout version 2. + pub fn is_version_2(&self) -> bool { + self.version == JSON2_LAYOUT_V2 + } +} + +fn parse_version(metadata: &str) -> crate::error::Result> { + serde_json::from_str::(metadata) + .map(|x| x.layout_version) + .map_err(|e| { + InvalidJson2LayoutSnafu { + reason: format!(r#"invalid extension metadata: "{metadata}", error: {e}"#), + } + .build() + }) +} + +/// Returns the remainder field of a JSON2 v2 root. +pub(crate) fn json2_remainder_field(field: &Field) -> crate::error::Result> { + let layout = Json2PhysicalLayout::try_from_root(field)?; + if !layout.is_version_2() { + return Ok(None); + } + + let remainder = if let DataType::Struct(fields) = field.data_type() { + fields + .iter() + .find(|x| x.name() == JSON2_REMAINDER_FIELD_NAME) + } else { + None + }; + + if let Some(remainder) = remainder { + let _ = remainder.try_extension_type::().map_err(|e| { + InvalidJson2LayoutSnafu { + reason: e.to_string(), + } + .build() + })?; + } + Ok(remainder) +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct JsonMetadata { /// JSON2 settings stored in Arrow extension metadata. json_settings: JsonSettings, + /// Physical JSON2 layout used by this Arrow field. + #[serde(default, skip_serializing_if = "Option::is_none")] + layout_version: Option, } impl JsonMetadata { - /// Creates JSON2 extension metadata. + /// Creates metadata for the legacy JSON2 layout. pub fn new(json_settings: JsonSettings) -> Self { - Self { json_settings } + Self { + json_settings, + layout_version: None, + } + } + + /// Creates metadata for the new JSON2 physical v2 layout. + pub fn new_v2(json_settings: JsonSettings) -> Self { + Self { + json_settings, + layout_version: Some(JSON2_LAYOUT_V2), + } } /// Returns the JSON2 settings. pub fn json_settings(&self) -> &JsonSettings { &self.json_settings } + + /// Returns whether this metadata describes JSON2 layout version 2. + pub fn is_version_2(&self) -> bool { + self.layout_version == Some(JSON2_LAYOUT_V2) + } +} + +impl Default for JsonMetadata { + fn default() -> Self { + Self::new(JsonSettings::default()) + } } /// Arrow extension type for legacy JSONB columns. @@ -227,6 +344,7 @@ mod tests { use arrow_schema::{Field, Fields}; use super::*; + use crate::vectors::json::variant::variant_field; #[test] fn test_json2_extension_type_detection() { @@ -282,4 +400,78 @@ mod tests { assert!(JsonExtensionType::try_new(&DataType::Null, ()).is_ok()); assert!(JsonExtensionType::try_new(&DataType::Struct(Fields::empty()), ()).is_err()); } + + #[test] + fn test_json_metadata_layout_version_compatibility() -> serde_json::Result<()> { + let legacy: JsonMetadata = serde_json::from_str(r#"{"json_settings":{}}"#)?; + assert!(!legacy.is_version_2()); + + let metadata = JsonMetadata::new_v2(JsonSettings::default()); + assert!(metadata.is_version_2()); + let serialized = serde_json::to_string(&metadata)?; + let deserialized: JsonMetadata = serde_json::from_str(&serialized)?; + assert!(deserialized.is_version_2()); + assert_eq!(deserialized, metadata); + Ok(()) + } + + #[test] + fn test_parse_json2_physical_layout() -> crate::error::Result<()> { + let legacy = Field::new("data", DataType::Struct(Fields::empty()), true) + .with_extension_type(Json2ExtensionType::default()); + assert!(!Json2PhysicalLayout::try_from_root(&legacy)?.is_version_2()); + + let v2 = Field::new( + "data", + DataType::Struct( + vec![ + Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)), + Arc::new(Field::new("count", DataType::Int64, true)), + ] + .into(), + ), + true, + ) + .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2( + JsonSettings::default(), + )))); + assert!(Json2PhysicalLayout::try_from_root(&v2)?.is_version_2()); + let remainder = json2_remainder_field(&v2)?.unwrap(); + assert_eq!(JSON2_REMAINDER_FIELD_NAME, remainder.name()); + Ok(()) + } + + #[test] + fn test_reject_invalid_json2_physical_layout() { + let field = Field::new("data", DataType::Struct(Fields::empty()), true); + assert!(Json2PhysicalLayout::try_from_root(&field).is_err()); + + let metadata = + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))); + let missing = Field::new("data", DataType::Struct(Fields::empty()), true) + .with_extension_type(metadata.clone()); + assert!(json2_remainder_field(&missing).is_ok_and(|x| x.is_none())); + + let invalid = Field::new( + "data", + DataType::Struct( + vec![Arc::new(Field::new( + JSON2_REMAINDER_FIELD_NAME, + DataType::Binary, + true, + ))] + .into(), + ), + true, + ) + .with_extension_type(metadata); + assert!(json2_remainder_field(&invalid).is_err()); + + let future = Field::new("data", DataType::Struct(Fields::empty()), true) + .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata { + json_settings: JsonSettings::default(), + layout_version: Some(JSON2_LAYOUT_V2 + 1), + }))); + assert!(Json2PhysicalLayout::try_from_root(&future).is_err()); + } } diff --git a/src/datatypes/src/vectors/json.rs b/src/datatypes/src/vectors/json.rs index 3dcb8d7bc07..c68e6cf3d92 100644 --- a/src/datatypes/src/vectors/json.rs +++ b/src/datatypes/src/vectors/json.rs @@ -14,3 +14,4 @@ pub mod array; pub(crate) mod builder; +pub mod variant; diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index afe9f553056..0426334c3dd 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/src/datatypes/src/vectors/json/array.rs @@ -17,9 +17,12 @@ use std::sync::Arc; use arrow::compute::{can_cast_types, cast}; use arrow_array::cast::AsArray; -use arrow_array::types::{Float64Type, Int64Type, UInt64Type}; +use arrow_array::types::{ + Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, + UInt32Type, UInt64Type, +}; use arrow_array::{Array, ArrayRef, GenericListArray, ListArray, StructArray, new_null_array}; -use arrow_schema::{DataType, FieldRef}; +use arrow_schema::{DataType, Field, FieldRef}; use common_telemetry::trace; use serde_json::Value; use snafu::{OptionExt, ResultExt}; @@ -29,9 +32,11 @@ use crate::data_type::ConcreteDataType; use crate::error::{ AlignJsonArraySnafu, ArrowComputeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu, Result, }; +use crate::extension::json::{JSON2_REMAINDER_FIELD_NAME, json2_remainder_field}; use crate::json::value::{decode_json_variant, encode_serde_json_as_jsonb}; use crate::prelude::{DataType as _, Value as GreptimeValue}; use crate::value::{ListValue, StructValue}; +use crate::vectors::json::variant::variant_to_json_values; pub struct JsonArray<'a> { inner: &'a ArrayRef, @@ -48,8 +53,15 @@ impl JsonArray<'_> { let value = match array.data_type() { DataType::Null => Value::Null, DataType::Boolean => Value::Bool(array.as_boolean().value(i)), + DataType::Int8 => Value::from(array.as_primitive::().value(i)), + DataType::Int16 => Value::from(array.as_primitive::().value(i)), + DataType::Int32 => Value::from(array.as_primitive::().value(i)), DataType::Int64 => Value::from(array.as_primitive::().value(i)), + DataType::UInt8 => Value::from(array.as_primitive::().value(i)), + DataType::UInt16 => Value::from(array.as_primitive::().value(i)), + DataType::UInt32 => Value::from(array.as_primitive::().value(i)), DataType::UInt64 => Value::from(array.as_primitive::().value(i)), + DataType::Float32 => Value::from(array.as_primitive::().value(i)), DataType::Float64 => Value::from(array.as_primitive::().value(i)), DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { Value::String(string_array_value(array, i).to_string()) @@ -92,6 +104,60 @@ impl JsonArray<'_> { Ok(value) } + /// Projects a physical JSON2 array to a logical query type. + /// + /// TODO(LFC) Supersede `project_to_v2` to `project_to`. + pub fn project_to_v2(&self, field: &Field, target: &DataType) -> Result { + if json2_remainder_field(field)?.is_some() { + project_json_values(self.json2_values()?, target) + } else { + self.project_to(target) + } + } + + fn json2_values(&self) -> Result> { + let structs = self.inner.as_struct_opt().context(AlignJsonArraySnafu { + reason: "JSON2 layout v2 root array must be a struct", + })?; + let remainder = structs.column_by_name(JSON2_REMAINDER_FIELD_NAME); + let mut remainders = if let Some(remainder) = remainder { + variant_to_json_values(remainder)? + } else { + vec![None; structs.len()] + }; + let mut values = Vec::with_capacity(structs.len()); + let mut path = Vec::new(); + + for (i, remainder) in remainders.iter_mut().enumerate() { + if structs.is_null(i) { + values.push(Value::Null); + continue; + } + + let mut object = match remainder.take() { + None => serde_json::Map::new(), + Some(Value::Object(object)) => object, + Some(value) => { + return InvalidJsonSnafu { + value: format!("JSON2 layout v2 remainder must be an object, got {value}"), + } + .fail(); + } + }; + + for (child, column) in structs.fields().iter().zip(structs.columns()) { + if child.name() == JSON2_REMAINDER_FIELD_NAME { + continue; + } + let value = JsonArray::from(column).try_get_value(i)?; + merge_explicit_value(&mut object, child.name().clone(), value, &mut path)?; + } + values.push(Value::Object(object)); + } + + Ok(values) + } + /// Normalizes a JSON2 array to the wider `expect` data type without losing /// information. /// @@ -326,6 +392,34 @@ impl JsonArray<'_> { } } +fn merge_explicit_value( + remainder: &mut serde_json::Map, + key: String, + explicit: Value, + path: &mut Vec, +) -> Result<()> { + let Some(existing) = remainder.get_mut(&key) else { + remainder.insert(key, explicit); + return Ok(()); + }; + path.push(key); + + let (Value::Object(remainder), Value::Object(explicit)) = (existing, explicit) else { + return InvalidJsonSnafu { + value: format!( + "cannot merge '{}' in explicit fields and remainder: not both objects", + path.join("."), + ), + } + .fail(); + }; + for (key, value) in explicit { + merge_explicit_value(remainder, key, value, path)?; + } + path.pop(); + Ok(()) +} + /// Returns whether Arrow can cast between the types without JSON-aware projection. /// Binary and nested types require JSONB decoding or recursive projection. fn can_fast_cast_types(from_type: &DataType, to_type: &DataType) -> bool { @@ -453,13 +547,16 @@ impl<'a> From<&'a ArrayRef> for JsonArray<'a> { mod test { use arrow_array::types::Int64Type; use arrow_array::{ - BinaryArray, BooleanArray, Float64Array, Int32Array, Int64Array, ListArray, StringArray, - UInt64Array, + BinaryArray, BooleanArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, + Int64Array, ListArray, StringArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array, }; use arrow_schema::{Field, Fields}; use serde_json::json; use super::*; + use crate::extension::json::{Json2ExtensionType, JsonMetadata}; + use crate::json::JsonSettings; + use crate::vectors::json::variant::{json_values_to_variant, variant_field}; #[test] fn test_try_get_value() -> Result<()> { @@ -474,6 +571,20 @@ mod test { assert_eq!(JsonArray::from(&ints).try_get_value(0)?, json!(-7)); assert_eq!(JsonArray::from(&ints).try_get_value(1)?, Value::Null); + macro_rules! assert_number { + ($array:expr, $expected:expr) => {{ + let array: ArrayRef = Arc::new($array); + assert_eq!(JsonArray::from(&array).try_get_value(0)?, json!($expected)); + }}; + } + assert_number!(Int8Array::from(vec![-8]), -8); + assert_number!(Int16Array::from(vec![-16]), -16); + assert_number!(Int32Array::from(vec![-32]), -32); + assert_number!(UInt8Array::from(vec![8]), 8); + assert_number!(UInt16Array::from(vec![16]), 16); + assert_number!(UInt32Array::from(vec![32]), 32); + assert_number!(Float32Array::from(vec![1.25]), 1.25); + let floats: ArrayRef = Arc::new(Float64Array::from(vec![Some(1.5)])); assert_eq!(JsonArray::from(&floats).try_get_value(0)?, json!(1.5)); @@ -529,15 +640,6 @@ mod test { json!({"flag": null, "items": [2]}) ); - let unsupported: ArrayRef = Arc::new(Int32Array::from(vec![1])); - assert_eq!( - JsonArray::from(&unsupported) - .try_get_value(0) - .unwrap_err() - .to_string(), - "Invalid JSON: unknown JSON type Int32" - ); - Ok(()) } @@ -915,4 +1017,124 @@ mod test { Ok(()) } + + #[test] + fn test_reconstruct_json2_v2_value() -> Result<()> { + let remainders = json_values_to_variant(&[ + Some(json!({"cold": 1, "nested": {"right": true}})), + Some(json!({"!__remainder__!": "user value"})), + ])?; + let remainder = Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)); + let nested = Arc::new(Field::new_struct( + "nested", + [Arc::new(Field::new("left", DataType::Utf8, true))], + true, + )); + let nested_values: ArrayRef = Arc::new(StructArray::from(vec![( + Arc::new(Field::new("left", DataType::Utf8, true)), + Arc::new(StringArray::from(vec![Some("value"), None])) as ArrayRef, + )])); + let fields = Fields::from(vec![ + remainder, + Arc::new(Field::new("count", DataType::Int64, true)), + nested, + ]); + let array: ArrayRef = Arc::new(StructArray::new( + fields.clone(), + vec![ + remainders, + Arc::new(Int64Array::from(vec![Some(42), None])), + nested_values, + ], + None, + )); + let field = Field::new("data", DataType::Struct(fields), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))), + ); + + assert_eq!( + json!({ + "cold": 1, + "count": 42, + "nested": {"left": "value", "right": true} + }), + JsonArray::from(&array).json2_values()?[0] + ); + assert_eq!( + json!({ + "!__remainder__!": "user value", + "count": null, + "nested": {"left": null} + }), + JsonArray::from(&array).json2_values()?[1] + ); + + let target = DataType::Struct( + vec![ + Arc::new(Field::new("cold", DataType::UInt64, true)), + Arc::new(Field::new("count", DataType::Int64, true)), + ] + .into(), + ); + let projected = JsonArray::from(&array).project_to_v2(&field, &target)?; + assert_eq!( + json!({"cold": 1, "count": 42}), + JsonArray::from(&projected).try_get_value(0)? + ); + assert_eq!( + json!({"cold": null, "count": null}), + JsonArray::from(&projected).try_get_value(1)? + ); + Ok(()) + } + + #[test] + fn test_project_partial_json2_v2_without_remainder() -> Result<()> { + let fields = Fields::from(vec![Arc::new(Field::new("hot", DataType::Int64, true))]); + let array: ArrayRef = Arc::new(StructArray::new( + fields.clone(), + vec![Arc::new(Int64Array::from(vec![1, 2]))], + None, + )); + let field = Field::new("data", DataType::Struct(fields), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))), + ); + + let projected = JsonArray::from(&array).project_to_v2(&field, field.data_type())?; + assert!(Arc::ptr_eq(&array, &projected)); + Ok(()) + } + + #[test] + fn test_reject_conflict_json2_v2_path() -> Result<()> { + let remainders = json_values_to_variant(&[Some(json!({"count": 1}))])?; + let fields = Fields::from(vec![ + Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)), + Arc::new(Field::new("count", DataType::Int64, true)), + ]); + let array: ArrayRef = Arc::new(StructArray::new( + fields, + vec![remainders, Arc::new(Int64Array::from(vec![2]))], + None, + )); + let error = JsonArray::from(&array).json2_values().unwrap_err(); + assert!( + error.to_string().contains( + "cannot merge 'count' in explicit fields and remainder: not both objects" + ) + ); + + let Value::Object(mut remainder) = json!({"nested": {"count": 1}}) else { + unreachable!(); + }; + let error = merge_explicit_value( + &mut remainder, + "nested".to_string(), + json!({"count": 2}), + &mut Vec::new(), + ) + .unwrap_err(); + assert!(error.to_string().contains("cannot merge 'nested.count'")); + Ok(()) + } } diff --git a/src/datatypes/src/vectors/json/variant.rs b/src/datatypes/src/vectors/json/variant.rs new file mode 100644 index 00000000000..1a4b8a861d5 --- /dev/null +++ b/src/datatypes/src/vectors/json/variant.rs @@ -0,0 +1,239 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use arrow_array::ArrayRef; +#[cfg(test)] +use arrow_schema::ArrowError; +use arrow_schema::{DataType, Field}; +#[cfg(test)] +use parquet_variant::{ObjectFieldBuilder, Variant, VariantBuilderExt, VariantDecimal16}; +#[cfg(test)] +use parquet_variant_compute::VariantArrayBuilder; +use parquet_variant_compute::{VariantArray, VariantType}; +use parquet_variant_json::VariantToJson; +use snafu::ResultExt; + +use crate::error::{ArrowComputeSnafu, Result}; +#[cfg(test)] +use crate::json::value::{JsonNumber, JsonVariant, decode_json_variant}; + +/// Returns the canonical Arrow field for an unshredded Parquet Variant array. +pub fn variant_field(name: impl Into, nullable: bool) -> Field { + Field::new_struct( + name, + [ + Arc::new(Field::new("metadata", DataType::BinaryView, false)), + Arc::new(Field::new("value", DataType::BinaryView, false)), + ], + nullable, + ) + .with_extension_type(VariantType) +} + +#[cfg(test)] +pub(crate) fn json_values_to_variant(values: &[Option]) -> Result { + let mut builder = VariantArrayBuilder::new(values.len()); + for value in values { + match value { + Some(value) => append_json_value(&mut builder, value).context(ArrowComputeSnafu)?, + None => builder.append_null(), + } + } + Ok(ArrayRef::from(builder.build())) +} + +/// Encodes JSON variants as an unshredded Parquet Variant array. +#[cfg(test)] +fn json_variants_to_variant(values: &[Option]) -> Result { + let mut builder = VariantArrayBuilder::new(values.len()); + for value in values { + match value { + Some(value) => append_json_variant(&mut builder, value).context(ArrowComputeSnafu)?, + None => builder.append_null(), + } + } + Ok(ArrayRef::from(builder.build())) +} + +#[cfg(test)] +fn append_json_variant( + builder: &mut impl VariantBuilderExt, + value: &JsonVariant, +) -> std::result::Result<(), ArrowError> { + match value { + JsonVariant::Null => builder.append_value(Variant::Null), + JsonVariant::Bool(value) => builder.append_value(*value), + JsonVariant::Number(JsonNumber::PosInt(value)) => { + if let Ok(value) = i64::try_from(*value) { + builder.append_value(value); + } else { + append_large_u64(builder, *value)?; + } + } + JsonVariant::Number(JsonNumber::NegInt(value)) => builder.append_value(*value), + JsonVariant::Number(JsonNumber::Float(value)) => { + if value.0.is_finite() { + builder.append_value(value.0) + } else { + builder.append_value("NaN") + } + } + JsonVariant::String(value) => builder.append_value(value.as_str()), + JsonVariant::Array(values) => { + let mut list = builder.try_new_list()?; + for value in values { + append_json_variant(&mut list, value)?; + } + list.finish(); + } + JsonVariant::Object(values) => { + let mut object = builder.try_new_object()?; + for (name, value) in values { + append_json_variant(&mut ObjectFieldBuilder::new(name, &mut object), value)?; + } + object.finish(); + } + JsonVariant::Variant(value) => { + let value = decode_json_variant(value) + .map_err(|e| ArrowError::JsonError(format!("Failed to decode JSONB: {e}")))?; + append_json_value(builder, &value)?; + } + } + Ok(()) +} + +#[cfg(test)] +fn append_json_value( + builder: &mut impl VariantBuilderExt, + value: &serde_json::Value, +) -> std::result::Result<(), ArrowError> { + match value { + serde_json::Value::Null => builder.append_value(Variant::Null), + serde_json::Value::Bool(value) => builder.append_value(*value), + serde_json::Value::Number(value) => { + if let Some(value) = value.as_i64() { + builder.append_value(value); + } else if let Some(value) = value.as_u64() { + append_large_u64(builder, value)?; + } else if let Some(value) = value.as_f64() { + builder.append_value(value); + } else { + return Err(ArrowError::InvalidArgumentError(format!( + "Failed to encode JSON number as Variant: {value}" + ))); + } + } + serde_json::Value::String(value) => builder.append_value(value.as_str()), + serde_json::Value::Array(values) => { + let mut list = builder.try_new_list()?; + for value in values { + append_json_value(&mut list, value)?; + } + list.finish(); + } + serde_json::Value::Object(values) => { + let mut object = builder.try_new_object()?; + for (name, value) in values { + append_json_value(&mut ObjectFieldBuilder::new(name, &mut object), value)?; + } + object.finish(); + } + } + Ok(()) +} + +/// Parquet Variant has no unsigned integer primitive. Treat u64 as i64 first, then use Decimal16 +/// to represent large (larger than i64::MAX) u64. +#[cfg(test)] +fn append_large_u64( + builder: &mut impl VariantBuilderExt, + value: u64, +) -> std::result::Result<(), ArrowError> { + let value = VariantDecimal16::try_new(value as i128, 0).map_err(|e| { + ArrowError::InvalidArgumentError(format!( + "Failed to encode JSON large integer as Variant Decimal16: {e}" + )) + })?; + builder.append_value(value); + Ok(()) +} + +/// Decodes an unshredded Parquet Variant array into JSON values. +pub fn variant_to_json_values(array: &ArrayRef) -> Result>> { + let variants = VariantArray::try_new(array.as_ref()).context(ArrowComputeSnafu)?; + (0..variants.len()) + .map(|i| { + if variants.is_null(i) { + Ok(None) + } else { + variants + .try_value(i) + .and_then(|x| x.to_json_value()) + .context(ArrowComputeSnafu) + .map(Some) + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn test_variant_field_matches_canonical_layout() { + let expected = VariantArrayBuilder::new(0) + .build() + .field("remainder") + .with_nullable(true); + assert_eq!(expected, variant_field("remainder", true)); + } + + #[test] + fn test_variant_json_round_trip() -> Result<()> { + let values = vec![ + None, + Some(serde_json::Value::Null), + Some(json!({})), + Some(json!({"nested": {"items": [1, "two", null]}})), + Some(json!({"unicode": "\u{503c}"})), + Some(json!({"max_u64": u64::MAX})), + ]; + + let array = json_values_to_variant(&values)?; + assert_eq!(values, variant_to_json_values(&array)?); + + let variants = values + .clone() + .into_iter() + .map(|x| x.map(JsonVariant::from)) + .collect::>(); + let array = json_variants_to_variant(&variants)?; + assert_eq!(values, variant_to_json_values(&array)?); + + let variants = [Some(JsonVariant::Variant( + jsonb::parse_value(br#"{"nested": true}"#).unwrap().to_vec(), + ))]; + let array = json_variants_to_variant(&variants)?; + assert_eq!( + vec![Some(json!({"nested": true}))], + variant_to_json_values(&array)? + ); + Ok(()) + } +} From 8dbfbee611cc9d3996830ac8aeebffeb755c196f Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 19 Aug 2026 03:47:44 +0000 Subject: [PATCH 085/126] refactor: remove open metrics parser (#8905) --- Cargo.lock | 18 ------------------ src/frontend/src/instance/prom_store.rs | 10 +--------- src/servers/Cargo.toml | 1 - src/servers/src/http/prom_store.rs | 9 --------- src/servers/src/prom_store.rs | 6 ------ src/servers/src/query_handler.rs | 3 --- src/servers/tests/http/prom_store_test.rs | 6 +----- 7 files changed, 2 insertions(+), 51 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 16b4c6b1045..d7c9241afca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -977,12 +977,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "auto_ops" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7460f7dd8e100147b82a63afca1a20eb6c231ee36b90ba7272e14951cb58af59" - [[package]] name = "autocfg" version = "1.5.0" @@ -9802,17 +9796,6 @@ dependencies = [ "url", ] -[[package]] -name = "openmetrics-parser" -version = "0.4.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e40a68c62e09c5dfec2f6472af3bd5e8ddf506fcf14c78ece23794ffbb874eca" -dependencies = [ - "auto_ops", - "pest", - "pest_derive", -] - [[package]] name = "opensrv-mysql" version = "0.8.0" @@ -13461,7 +13444,6 @@ dependencies = [ "notify", "object-pool", "once_cell", - "openmetrics-parser", "opensrv-mysql", "opentelemetry-proto", "operator", diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index 428628df123..e0b17b1a89e 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -47,7 +47,7 @@ use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF, coll use servers::http::prom_store::PHYSICAL_TABLE_PARAM; use servers::interceptor::{PromStoreProtocolInterceptor, PromStoreProtocolInterceptorRef}; use servers::pending_rows_batcher::PendingRowsSchemaAlterer; -use servers::prom_store::{self, Metrics}; +use servers::prom_store; use servers::query_handler::{ PromStoreProtocolHandler, PromStoreProtocolHandlerRef, PromStoreResponse, }; @@ -648,10 +648,6 @@ impl PromStoreProtocolHandler for Instance { .fail(), } } - - async fn ingest_metrics(&self, _metrics: Metrics) -> ServerResult<()> { - todo!(); - } } impl Instance { @@ -791,10 +787,6 @@ impl PromStoreProtocolHandler for ExportMetricHandler { ) -> ServerResult { unreachable!(); } - - async fn ingest_metrics(&self, _metrics: Metrics) -> ServerResult<()> { - unreachable!(); - } } #[cfg(test)] diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index d64a1bd5a5d..59d00b90ad7 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -87,7 +87,6 @@ mysql_common = "0.38" notify.workspace = true object-pool = "0.5" once_cell.workspace = true -openmetrics-parser = "0.4" # Wait for https://github.com/databendlabs/opensrv/pull/81 opensrv-mysql = { git = "https://github.com/GreptimeTeam/opensrv", rev = "6c5a451544194b7bb60a8318d155d4f892b49f2c" } opentelemetry-proto.workspace = true diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 1f0d26789eb..cb7b6307c4b 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -770,7 +770,6 @@ mod tests { use super::*; use crate::prom_remote_write::validation::PromValidationMode; - use crate::prom_store::Metrics; use crate::query_handler::PromStoreProtocolHandler; #[test] @@ -966,10 +965,6 @@ mod tests { ) -> Result { unimplemented!() } - - async fn ingest_metrics(&self, _metrics: Metrics) -> Result<()> { - unimplemented!() - } } #[tokio::test] @@ -1067,9 +1062,5 @@ mod tests { ) -> Result { unimplemented!() } - - async fn ingest_metrics(&self, _metrics: Metrics) -> Result<()> { - unimplemented!() - } } } diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index c4f2f929cae..980504515a8 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -34,7 +34,6 @@ use datafusion::dataframe::DataFrame; use datafusion::prelude::{Expr, col, lit, regexp_match}; use datafusion_common::ScalarValue; use datafusion_expr::LogicalPlan; -use openmetrics_parser::{MetricsExposition, PrometheusType, PrometheusValue}; use snafu::{OptionExt, ResultExt, ensure}; use snap::raw::{Decoder, Encoder}; @@ -95,11 +94,6 @@ pub fn is_physical_table_selection_label(label: &str) -> bool { label == PHYSICAL_TABLE_LABEL || label == PHYSICAL_TABLE_LABEL_ALT } -/// Metrics for push gateway protocol -pub struct Metrics { - pub exposition: MetricsExposition, -} - /// Get table name from remote query pub fn table_name(q: &Query) -> Result { let mut matchers = q diff --git a/src/servers/src/query_handler.rs b/src/servers/src/query_handler.rs index a6a086e0f23..96f412ac1e1 100644 --- a/src/servers/src/query_handler.rs +++ b/src/servers/src/query_handler.rs @@ -54,7 +54,6 @@ use crate::error::Result; use crate::http::jaeger::QueryTraceParams; use crate::influxdb::InfluxdbRequest; use crate::opentsdb::codec::DataPoint; -use crate::prom_store::Metrics; pub type OpentsdbProtocolHandlerRef = Arc; pub type InfluxdbLineProtocolHandlerRef = Arc; pub type PromStoreProtocolHandlerRef = Arc; @@ -125,8 +124,6 @@ pub trait PromStoreProtocolHandler { /// Handling prometheus remote read requests async fn read(&self, request: ReadRequest, ctx: QueryContextRef) -> Result; - /// Handling push gateway requests - async fn ingest_metrics(&self, metrics: Metrics) -> Result<()>; } #[async_trait] diff --git a/src/servers/tests/http/prom_store_test.rs b/src/servers/tests/http/prom_store_test.rs index 6e8658f663f..8f3e2b0e770 100644 --- a/src/servers/tests/http/prom_store_test.rs +++ b/src/servers/tests/http/prom_store_test.rs @@ -44,7 +44,7 @@ use servers::http::{HttpOptions, HttpServerBuilder}; use servers::prom_remote_write::v2::test_util as remote_write_v2; use servers::prom_remote_write::validation::PromValidationMode; use servers::prom_store; -use servers::prom_store::{Metrics, snappy_compress}; +use servers::prom_store::snappy_compress; use servers::query_handler::sql::SqlQueryHandler; use servers::query_handler::{PromStoreProtocolHandler, PromStoreResponse}; use session::context::QueryContextRef; @@ -146,10 +146,6 @@ impl PromStoreProtocolHandler for DummyInstance { body: response.encode_to_vec(), }) } - - async fn ingest_metrics(&self, _metrics: Metrics) -> Result<()> { - unimplemented!(); - } } #[async_trait] From 3de4feddd5a5c6a282d5b2eae7541937de6bcdf8 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 19 Aug 2026 04:32:23 +0000 Subject: [PATCH 086/126] feat(otlp): report the cause of rejected trace spans (#8897) * feat(otlp): report the cause of rejected trace spans When trace-v1 ingestion cannot coerce an attribute value, it falls back to single-span writes and rejects the bad span. That behavior is correct, but the OTLP partial-success message only carried `Rejected span : (InvalidArguments)`: the column, the source value, the source type and the target type were all dropped, so locating the bad attribute required adding a detailed exporter on the collector side and replaying traffic. Two places lost the information. `prepare_trace_column_rewrites` built a message without the failing value, and the span rejection path kept only the status code from the error. Coercion errors now name the failing value, e.g. failed to coerce trace column 'span_attributes.http.response.body.size' in table 'opentelemetry_traces' from String("") to Int64 and the rejection detail carries that cause. Values are user data, so a string keeps at most 16 characters, is escaped, and binary payloads report only their length; the cause itself is bounded at 256 characters. Both truncations cut on a char boundary. Failure details now deduplicate: repeats of the same (site, cause) collapse into one entry with an occurrence count, keyed on the untruncated cause so two failures that differ past the display limit stay separate. Only four distinct entries are retained and the rest are counted, which keeps the state bounded no matter how many distinct bad values a request carries. A fully rejected request is logged at warn level and a partial success at debug level, since the latter repeats every export interval; the detail goes out as a Debug field so a newline in an attribute key cannot forge log lines. Rejection semantics are unchanged: partial success, same accepted and rejected counts, same HTTP status mapping. Signed-off-by: Dennis Zhuang * refactor(otlp): compare failure keys directly instead of hashing The dedup identity was a DefaultHasher fingerprint of `(label, key)`, which bought a fixed 8 bytes per entry at the cost of an import, four lines, and a collision argument the reader has to make. Entries are capped at four and a cause runs a couple of hundred characters, so the saving is about a kilobyte per in-flight request while the column name it avoids retaining is already held several times over by the request itself. Compare the strings instead, keeping the untruncated cause as the key so failures differing past the display limit still stay apart. Labels are metric label values and always static, so the entry borrows them. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/instance/otlp/README.md | 12 ++ .../src/instance/otlp/trace_ingest.rs | 172 ++++++++++++++++-- .../src/instance/otlp/trace_ingest/tests.rs | 151 ++++++++++++--- src/frontend/src/instance/otlp/trace_types.rs | 153 ++++++++++++++-- tests-integration/tests/http.rs | 9 + 5 files changed, 436 insertions(+), 61 deletions(-) diff --git a/src/frontend/src/instance/otlp/README.md b/src/frontend/src/instance/otlp/README.md index 8db27384183..4489ab0c3f3 100644 --- a/src/frontend/src/instance/otlp/README.md +++ b/src/frontend/src/instance/otlp/README.md @@ -158,6 +158,18 @@ Accounting follows the main-table write: accepted or rejected span counts; and - failure details are bounded before they are folded into `TraceIngestOutcome`. +A rejection detail carries the failing cause, not just a status code, so an +unusable attribute value names its column, source value, and target type. Repeats +of the same `(site, cause)` collapse into one entry with an occurrence count, and +only a fixed number of distinct entries is retained; failures past that limit are +counted but their causes are dropped, which keeps the state bounded regardless of +how many distinct bad values a request carries. + +A request that rejects every span is logged at warn level. A partial success is +logged at debug level only: it repeats every export interval while the sender +keeps emitting the value, and its detail already reaches the sender through the +OTLP partial-success message. + Finally, the HTTP handler returns: - full success when there are no rejected spans or failure details; diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs index f6ce60b2d11..671f6c6f7cf 100644 --- a/src/frontend/src/instance/otlp/trace_ingest.rs +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -25,7 +25,7 @@ use client::Output; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; use common_meta::rpc::ddl::TriggerReason; -use common_telemetry::warn; +use common_telemetry::{debug, warn}; use datatypes::prelude::ConcreteDataType; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use pipeline::{GreptimePipelineParams, PipelineWay}; @@ -50,11 +50,15 @@ use crate::instance::otlp::trace_types::{ PendingTraceColumnRewrite, PreparedTraceColumnRewrites, TraceColumnRewriteError, choose_trace_reconcile_decision, enrich_trace_reconcile_error, is_trace_reconcile_candidate_type, prepare_trace_column_rewrites, push_observed_trace_type, + truncate_for_diagnostics, }; use crate::metrics::{OTLP_TRACES_FAILURE_COUNT, OTLP_TRACES_ROWS}; const TRACE_FAILURE_MESSAGE_LIMIT: usize = 4; +/// Maximum characters of a failure cause echoed to the client and the log. +const TRACE_FAILURE_CAUSE_LIMIT: usize = 256; + /// Determines how trace ingestion responds to a failure before a write is dispatched. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ChunkFailureReaction { @@ -86,7 +90,34 @@ struct TraceChunkIngestContext<'a> { struct TraceIngestState { aux_data: TraceAuxData, outcome: TraceIngestOutcome, - failure_messages: Vec, + failure_messages: TraceFailureMessages, +} + +/// Bounded, deduplicated failure details for one trace request. +/// +/// Occurrences past [`TRACE_FAILURE_MESSAGE_LIMIT`] distinct failures are +/// counted but their keys are dropped: retaining them would let this state grow +/// with the number of distinct bad values in a request. +#[derive(Debug, Default)] +struct TraceFailureMessages { + entries: Vec, + suppressed_occurrences: usize, +} + +#[derive(Debug)] +struct TraceFailureEntry { + label: &'static str, + /// Untruncated: two causes can agree on a truncated prefix and differ + /// exactly where the actionable detail is. + key: String, + message: String, + occurrences: usize, +} + +impl TraceFailureMessages { + fn is_empty(&self) -> bool { + self.entries.is_empty() + } } /// How a v1 chunk should be reconciled when it is written. @@ -863,7 +894,7 @@ impl Instance { let mut ingest_state = TraceIngestState { aux_data: TraceAuxData::default(), outcome: TraceIngestOutcome::default(), - failure_messages: Vec::new(), + failure_messages: TraceFailureMessages::default(), }; let main_result: ServerResult<()> = async { @@ -954,6 +985,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, "aux_table_update_failed", + err.status_code().as_ref(), format!( "Auxiliary trace tables were not fully updated ({})", err.status_code().as_ref() @@ -968,6 +1000,7 @@ impl Instance { Err(err) => Self::push_trace_failure_message( &mut ingest_state.failure_messages, "aux_table_update_failed", + err.status_code().as_ref(), format!( "Auxiliary trace tables were not fully updated ({})", err.status_code().as_ref() @@ -985,6 +1018,35 @@ impl Instance { ingest_state.failure_messages, ); + if let Some(error_message) = &ingest_state.outcome.error_message { + let accepted_spans = ingest_state.outcome.accepted_spans; + let rejected_spans = ingest_state.outcome.rejected_spans; + // A partial success repeats every export interval while the sender + // keeps emitting the bad value, so only a fully rejected request is + // worth a warning. Either way the detail reaches the sender: as an + // OTLP partial success, or as the status message of a 400. + // + // The detail embeds attribute keys verbatim, so it goes out as a + // Debug field; interpolating it would let a newline in an attribute + // key forge log lines. + if accepted_spans == 0 && rejected_spans > 0 { + warn!( + table_name = ingest_ctx.table_name, + rejected_spans, + error_message = ?error_message, + "OTLP trace ingest rejected every span" + ); + } else { + debug!( + table_name = ingest_ctx.table_name, + accepted_spans, + rejected_spans, + error_message = ?error_message, + "OTLP trace ingest reported failures" + ); + } + } + Ok(ingest_state.outcome) } @@ -1019,6 +1081,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::Propagate.as_metric_label(), + err.status_code().as_ref(), format!( "Propagating chunk write failure ({})", err.status_code().as_ref() @@ -1069,13 +1132,18 @@ impl Instance { // recover valid data. let span_count = chunks.iter().map(|chunk| chunk.rows.len()).sum::(); ingest_state.outcome.rejected_spans += span_count; + let (cause, shown_cause) = Self::trace_failure_cause(&err); Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::DiscardChunk.as_metric_label(), + // Merging discards of different sizes would make the reported + // count times its occurrences wrong. + &format!("{span_count}:{cause}"), format!( - "Discarded {} spans after pre-write request failure ({})", + "Discarded {} spans after pre-write request failure ({}): {}", span_count, - err.status_code().as_ref() + err.status_code().as_ref(), + shown_cause ), ); return Ok(()); @@ -1110,6 +1178,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::RetryPerSpan.as_metric_label(), + "incompatible_binary_and_json", "Chunk fallback triggered by incompatible binary and JSON values".to_string(), ); return self @@ -1136,6 +1205,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::RetryPerSpan.as_metric_label(), + err.status_code().as_ref(), format!("Chunk fallback triggered by {}", err.status_code().as_ref()), ); return self @@ -1144,13 +1214,19 @@ impl Instance { } ChunkFailureReaction::DiscardChunk => { ingest_state.outcome.rejected_spans += span_count; + let (cause, shown_cause) = Self::trace_failure_cause(&err); Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::DiscardChunk.as_metric_label(), + // Chunks discarded for one cause differ in size; merging + // them would make the reported count times its + // occurrences wrong. + &format!("{span_count}:{cause}"), format!( - "Discarded {} spans after pre-write chunk failure ({})", + "Discarded {} spans after pre-write chunk failure ({}): {}", span_count, - err.status_code().as_ref() + err.status_code().as_ref(), + shown_cause ), ); return Ok(()); @@ -1159,6 +1235,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::Propagate.as_metric_label(), + err.status_code().as_ref(), format!( "Propagating pre-write chunk failure ({})", err.status_code().as_ref() @@ -1182,6 +1259,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::Propagate.as_metric_label(), + err.status_code().as_ref(), format!( "Propagating chunk write failure ({})", err.status_code().as_ref() @@ -1221,6 +1299,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::Propagate.as_metric_label(), + err.status_code().as_ref(), format!( "Propagating pre-write span failure for {}:{} ({})", span.trace_id, @@ -1232,14 +1311,19 @@ impl Instance { } ingest_state.outcome.rejected_spans += 1; + // Dedup on the cause, not the span id: one bad column rejects + // every span and would otherwise fill the entry list. + let (cause, shown_cause) = Self::trace_failure_cause(&err); Self::push_trace_failure_message( &mut ingest_state.failure_messages, "span_rejected", + &cause, format!( - "Rejected span {}:{} ({})", + "Rejected span {}:{} ({}): {}", span.trace_id, span.span_id, - err.status_code().as_ref() + err.status_code().as_ref(), + shown_cause ), ); continue; @@ -1256,6 +1340,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::Propagate.as_metric_label(), + err.status_code().as_ref(), format!( "Propagating span write failure for {}:{} ({})", span.trace_id, @@ -1619,24 +1704,52 @@ impl Instance { outcome.write_cost += cost; } - fn push_trace_failure_message(messages: &mut Vec, label: &str, message: String) { + /// Returns the full cause of a pre-write failure and its display form. + /// + /// `output_msg` masks internal errors and unwraps the root cause. Dedup must + /// key on the full text: two causes can agree on a truncated prefix and + /// differ exactly where the actionable detail is. + fn trace_failure_cause(err: &error::Error) -> (String, String) { + let cause = err.output_msg(); + let display = truncate_for_diagnostics(&cause, TRACE_FAILURE_CAUSE_LIMIT); + (cause, display) + } + + /// Records one failure, merging repeats of `(label, key)` into a count. + fn push_trace_failure_message( + messages: &mut TraceFailureMessages, + label: &'static str, + key: &str, + message: String, + ) { OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).inc(); - if messages.len() < TRACE_FAILURE_MESSAGE_LIMIT { - messages.push(message); - } else if messages.len() == TRACE_FAILURE_MESSAGE_LIMIT { - tracing::debug!( - label, - limit = TRACE_FAILURE_MESSAGE_LIMIT, - "Trace ingest failure message limit reached; suppressing additional failure details" - ); + if let Some(entry) = messages + .entries + .iter_mut() + .find(|entry| entry.label == label && entry.key == key) + { + entry.occurrences += 1; + return; } + + if messages.entries.len() >= TRACE_FAILURE_MESSAGE_LIMIT { + messages.suppressed_occurrences += 1; + return; + } + + messages.entries.push(TraceFailureEntry { + label, + key: key.to_string(), + message, + occurrences: 1, + }); } fn finish_trace_failure_message( accepted_spans: usize, rejected_spans: usize, - messages: Vec, + messages: TraceFailureMessages, ) -> Option { if rejected_spans == 0 && messages.is_empty() { return None; @@ -1648,8 +1761,27 @@ impl Instance { ); if !messages.is_empty() { + let details = messages + .entries + .into_iter() + .map(|entry| { + if entry.occurrences > 1 { + format!("{} (x{})", entry.message, entry.occurrences) + } else { + entry.message + } + }) + .collect::>() + .join("; "); summary.push_str(": "); - summary.push_str(&messages.join("; ")); + summary.push_str(&details); + } + + if messages.suppressed_occurrences > 0 { + summary.push_str(&format!( + "; {} additional failures suppressed", + messages.suppressed_occurrences + )); } Some(summary) diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index b3a0570db51..ec5bad89754 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -34,8 +34,8 @@ use servers::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; use servers::otlp::trace::v1::{TraceBinaryType, TraceRetryColumn}; use super::{ - ChunkFailureReaction, Instance, TraceChunkRetry, TraceChunkSchemaState, TraceRequestSchema, - TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, + ChunkFailureReaction, Instance, TraceChunkRetry, TraceChunkSchemaState, TraceFailureMessages, + TraceRequestSchema, TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, wrap_trace_alter_failure, }; use crate::metrics::OTLP_TRACES_FAILURE_COUNT; @@ -102,22 +102,28 @@ fn test_classify_trace_prewrite_failure() { #[test] fn test_finish_trace_failure_message() { - let message = Instance::finish_trace_failure_message( - 3, - 2, - vec!["Rejected span trace:span (InvalidArguments)".to_string()], - ) - .unwrap(); + let mut messages = TraceFailureMessages::default(); + Instance::push_trace_failure_message( + &mut messages, + "span_rejected", + "coercion", + "Rejected span trace:span (InvalidArguments)".to_string(), + ); + + let message = Instance::finish_trace_failure_message(3, 2, messages).unwrap(); assert!(message.contains("Accepted 3 spans, rejected 2 spans")); assert!(message.contains("Rejected span trace:span")); - assert_eq!(Instance::finish_trace_failure_message(2, 0, vec![]), None); + assert_eq!( + Instance::finish_trace_failure_message(2, 0, TraceFailureMessages::default()), + None + ); } #[test] fn test_finish_trace_failure_message_without_detail_messages() { assert_eq!( - Instance::finish_trace_failure_message(0, 2, vec![]), + Instance::finish_trace_failure_message(0, 2, TraceFailureMessages::default()), Some("Accepted 0 spans, rejected 2 spans".to_string()) ); } @@ -126,39 +132,138 @@ fn test_finish_trace_failure_message_without_detail_messages() { fn test_push_trace_failure_message_increments_labeled_counter() { let label = "retry_per_span_counter_test"; let initial = OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(); - let mut messages = Vec::new(); + let mut messages = TraceFailureMessages::default(); Instance::push_trace_failure_message( &mut messages, label, + "InvalidArguments", "Chunk fallback triggered by InvalidArguments".to_string(), ); - assert_eq!(messages.len(), 1); + assert_eq!(messages.entries.len(), 1); assert_eq!( OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(), initial + 1 ); } +/// One bad column rejects every span, so repeats must collapse instead of +/// filling the bounded entry list with the same cause. +#[test] +fn test_push_trace_failure_message_collapses_repeated_cause() { + let label = "span_rejected_dedup_test"; + let initial = OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(); + let mut messages = TraceFailureMessages::default(); + + for idx in 0..3 { + Instance::push_trace_failure_message( + &mut messages, + label, + "failed to coerce", + format!("Rejected span trace:span-{idx} (InvalidArguments): failed to coerce"), + ); + } + + assert_eq!(messages.entries.len(), 1); + assert_eq!(messages.entries[0].occurrences, 3); + assert_eq!(messages.suppressed_occurrences, 0); + // Every occurrence is still metered even though the details collapse. + assert_eq!( + OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(), + initial + 3 + ); + + let summary = Instance::finish_trace_failure_message(0, 3, messages).unwrap(); + assert!( + summary.contains("Rejected span trace:span-0 (InvalidArguments): failed to coerce (x3)"), + "unexpected summary: {summary}" + ); +} + +/// Dedup identity comes from the full cause, so two failures that only differ +/// past the display limit must stay separate instead of one silently vanishing. +#[test] +fn test_trace_failure_cause_keeps_identity_beyond_the_display_limit() { + let shared_prefix = "x".repeat(400); + let (first_cause, first_shown) = Instance::trace_failure_cause( + &servers::error::InvalidParameterSnafu { + reason: format!("{shared_prefix}-int64"), + } + .build(), + ); + let (second_cause, second_shown) = Instance::trace_failure_cause( + &servers::error::InvalidParameterSnafu { + reason: format!("{shared_prefix}-float64"), + } + .build(), + ); + + assert_eq!(first_shown, second_shown); + assert_ne!(first_cause, second_cause); + + let mut messages = TraceFailureMessages::default(); + Instance::push_trace_failure_message(&mut messages, "span_rejected", &first_cause, first_shown); + Instance::push_trace_failure_message( + &mut messages, + "span_rejected", + &second_cause, + second_shown, + ); + + assert_eq!(messages.entries.len(), 2); +} + +/// The same cause reported by two different sites must stay distinguishable. +#[test] +fn test_push_trace_failure_message_separates_labels_sharing_a_key() { + let mut messages = TraceFailureMessages::default(); + + Instance::push_trace_failure_message( + &mut messages, + "span_rejected", + "failed to coerce", + "Rejected span trace:span (InvalidArguments): failed to coerce".to_string(), + ); + Instance::push_trace_failure_message( + &mut messages, + "discard_chunk", + "failed to coerce", + "Discarded 7 spans after pre-write chunk failure (InvalidArguments): failed to coerce" + .to_string(), + ); + + assert_eq!(messages.entries.len(), 2); +} + #[test] fn test_push_trace_failure_message_caps_recorded_messages() { let label = "retry_per_span_limit_test"; - let mut messages = Vec::new(); + let mut messages = TraceFailureMessages::default(); - for idx in 0..=4 { - Instance::push_trace_failure_message(&mut messages, label, format!("failure-{idx}")); + for idx in 0..=5 { + Instance::push_trace_failure_message( + &mut messages, + label, + &format!("cause-{idx}"), + format!("failure-{idx}"), + ); } - assert_eq!(messages.len(), 4); assert_eq!( - messages, - vec![ - "failure-0".to_string(), - "failure-1".to_string(), - "failure-2".to_string(), - "failure-3".to_string() - ] + messages + .entries + .iter() + .map(|entry| entry.message.as_str()) + .collect::>(), + vec!["failure-0", "failure-1", "failure-2", "failure-3"] + ); + assert_eq!(messages.suppressed_occurrences, 2); + + let summary = Instance::finish_trace_failure_message(0, 6, messages).unwrap(); + assert!( + summary.ends_with("; 2 additional failures suppressed"), + "unexpected summary: {summary}" ); } diff --git a/src/frontend/src/instance/otlp/trace_types.rs b/src/frontend/src/instance/otlp/trace_types.rs index 0e2df25053e..603cc66c0d1 100644 --- a/src/frontend/src/instance/otlp/trace_types.rs +++ b/src/frontend/src/instance/otlp/trace_types.rs @@ -22,6 +22,39 @@ use servers::otlp::coerce::{ use crate::instance::otlp::trace_semconv::trace_semconv_fixed_type; +/// Attribute values are user data echoed back in the OTLP partial-success +/// message and the server log, so diagnostics keep at most this many characters. +const TRACE_VALUE_DIAGNOSTIC_LIMIT: usize = 16; + +/// Truncates to `limit` characters, marking the cut with `...`. +/// +/// Diagnostics carry user-controlled text; slicing by byte offset would panic in +/// the middle of a multi-byte sequence. +pub(super) fn truncate_for_diagnostics(text: &str, limit: usize) -> String { + match text.char_indices().nth(limit) { + Some((offset, _)) => format!("{}...", &text[..offset]), + None => text.to_string(), + } +} + +/// Renders a failing trace value as `Type(value)`, e.g. `String("")`. +fn describe_trace_value(value: &ValueData, request_type: ColumnDataType) -> String { + let payload = match value { + ValueData::StringValue(string_value) => format!( + "{:?}", + truncate_for_diagnostics(string_value, TRACE_VALUE_DIAGNOSTIC_LIMIT) + ), + ValueData::BoolValue(bool_value) => bool_value.to_string(), + ValueData::I64Value(int_value) => int_value.to_string(), + ValueData::F64Value(float_value) => float_value.to_string(), + ValueData::BinaryValue(bytes) => format!("{} bytes", bytes.len()), + // Other value kinds never reach trace coercion, so report the type alone. + _ => return format!("{request_type:?}"), + }; + + format!("{request_type:?}({payload})") +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) enum TraceReconcileDecision { UseExisting(ColumnDataType), @@ -191,8 +224,10 @@ pub(super) fn prepare_trace_column_rewrites( let Some(value) = row.values.get(pending_rewrite.col_idx) else { continue; }; - let Some(request_type) = value.value_data.as_ref().and_then(trace_value_datatype) - else { + let Some(request_value) = value.value_data.as_ref() else { + continue; + }; + let Some(request_type) = trace_value_datatype(request_value) else { continue; }; if request_type == pending_rewrite.target_type { @@ -201,20 +236,18 @@ pub(super) fn prepare_trace_column_rewrites( let value_data = coerce_value_data(&value.value_data, pending_rewrite.target_type, request_type) - .map_err(|_| { - TraceColumnRewriteError { - error: error::InvalidParameterSnafu { - reason: format!( - "failed to coerce trace column '{}' in table '{}' from {:?} to {:?}", - pending_rewrite.column_name, - table_name, - request_type, - pending_rewrite.target_type - ), - } - .build(), - column_name: pending_rewrite.column_name.clone(), - } + .map_err(|_| TraceColumnRewriteError { + error: error::InvalidParameterSnafu { + reason: format!( + "failed to coerce trace column '{}' in table '{}' from {} to {:?}", + pending_rewrite.column_name, + table_name, + describe_trace_value(request_value, request_type), + pending_rewrite.target_type + ), + } + .build(), + column_name: pending_rewrite.column_name.clone(), })?; values.push(PreparedTraceValueRewrite { row_idx, @@ -301,8 +334,8 @@ mod tests { use super::{ PendingTraceColumnRewrite, TraceReconcileDecision, choose_trace_reconcile_decision, - enrich_trace_reconcile_error, is_trace_reconcile_candidate_type, - prepare_trace_column_rewrites, push_observed_trace_type, + describe_trace_value, enrich_trace_reconcile_error, is_trace_reconcile_candidate_type, + prepare_trace_column_rewrites, push_observed_trace_type, truncate_for_diagnostics, }; #[test] @@ -458,6 +491,90 @@ mod tests { .unwrap_err(); assert_eq!(err.error.status_code(), StatusCode::InvalidArguments); assert_eq!(err.column_name, "span_attributes.attr_int"); + assert!( + err.error.to_string().contains( + "failed to coerce trace column 'span_attributes.attr_int' in table \ + 'trace_type_atomicity' from String(\"not_a_number\") to Int64" + ), + "unexpected error message: {}", + err.error + ); + } + + /// The PHP instrumentation case: an empty string must be distinguishable + /// from any other unparsable value in the reported diagnostics. + #[test] + fn test_prepare_trace_column_rewrites_reports_empty_string_value() { + let rows = vec![Row { + values: vec![Value { + value_data: Some(ValueData::StringValue(String::new())), + }], + }]; + let pending_rewrites = vec![PendingTraceColumnRewrite { + col_idx: 0, + target_type: ColumnDataType::Int64, + column_name: "span_attributes.http.response.body.size".to_string(), + }]; + + let err = prepare_trace_column_rewrites(&rows, pending_rewrites, "opentelemetry_traces") + .unwrap_err(); + assert!( + err.error.to_string().contains( + "'span_attributes.http.response.body.size' in table 'opentelemetry_traces' \ + from String(\"\") to Int64" + ), + "unexpected error message: {}", + err.error + ); + } + + #[test] + fn test_describe_trace_value_bounds_and_escapes_strings() { + assert_eq!( + describe_trace_value( + &ValueData::StringValue(String::new()), + ColumnDataType::String + ), + r#"String("")"# + ); + assert_eq!( + describe_trace_value( + &ValueData::StringValue("a\tb\"c".to_string()), + ColumnDataType::String + ), + r#"String("a\tb\"c")"# + ); + assert_eq!( + describe_trace_value( + &ValueData::StringValue("0123456789abcdefghij".to_string()), + ColumnDataType::String + ), + r#"String("0123456789abcdef...")"# + ); + } + + #[test] + fn test_describe_trace_value_omits_binary_content() { + assert_eq!( + describe_trace_value( + &ValueData::BinaryValue(vec![1_u8, 2, 3]), + ColumnDataType::Binary + ), + "Binary(3 bytes)" + ); + assert_eq!( + describe_trace_value(&ValueData::F64Value(1.5), ColumnDataType::Float64), + "Float64(1.5)" + ); + } + + /// Truncation runs over user-supplied text, so it must not split a + /// multi-byte character. + #[test] + fn test_truncate_for_diagnostics_cuts_on_char_boundary() { + assert_eq!(truncate_for_diagnostics("日本語テキスト", 3), "日本語..."); + assert_eq!(truncate_for_diagnostics("short", 16), "short"); + assert_eq!(truncate_for_diagnostics("exact", 5), "exact"); } #[test] diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 079439bf78c..90cbb8c93fe 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -7782,6 +7782,15 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ), "unexpected partial success body: {body:?}" ); + // The rejection must name the column, the failing value and the target type, + // otherwise locating the bad attribute needs a collector-side capture. + assert!( + partial_success.error_message.contains( + "failed to coerce trace column 'span_attributes.attr_int' in table \ + 'trace_type_abort' from String(\"not_a_number\") to Int64" + ), + "unexpected partial success body: {body:?}" + ); validate_data( "otlp_traces_v1_type_abort_rows", From 2182dccd9b76efc87e2e3ad7e7d355e62d9b1d03 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Wed, 19 Aug 2026 10:04:50 +0000 Subject: [PATCH 087/126] fix: cap default runtime sizes to a minimum of 2 threads (#8908) * fix: cap default runtime sizes to a minimum of 2 threads RuntimeOptions derived its default sizes directly from num_cpus. On single-core machines every runtime (global, compact, query, ingest) ended up with one worker thread, which can easily deadlock async code (e.g. block_on combined with spawn). Clamp all CPU-derived runtime sizes to at least 2 threads. Signed-off-by: Lei, HUANG * fix: init logging before runtimes so runtime options are logged The global runtimes were initialized before the global logging subscriber, so the "Creating runtime ..." info logs that carry the runtime sizes were silently dropped. Initialize logging first in all node start paths; common-telemetry has no dependency on common-runtime, so the reorder is safe. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- config/config.md | 10 ++--- config/datanode.example.toml | 4 +- config/frontend.example.toml | 2 +- config/metasrv.example.toml | 2 +- config/standalone.example.toml | 2 +- src/cmd/src/datanode/builder.rs | 6 +-- src/cmd/src/flownode.rs | 4 +- src/cmd/src/frontend.rs | 4 +- src/cmd/src/metasrv.rs | 4 +- src/cmd/src/standalone.rs | 4 +- src/common/runtime/src/global.rs | 67 +++++++++++++++++++++++++++----- 11 files changed, 78 insertions(+), 31 deletions(-) diff --git a/config/config.md b/config/config.md index 123eaf90a20..1b18565cc08 100644 --- a/config/config.md +++ b/config/config.md @@ -25,7 +25,7 @@ | `runtime` | -- | -- | The runtime options. | | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 1). | +| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | | `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the
`prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. | @@ -255,7 +255,7 @@ | `runtime` | -- | -- | The runtime options. | | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 1). | +| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | | `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the
`prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. | @@ -398,7 +398,7 @@ | `runtime` | -- | -- | The runtime options. | | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 1). | +| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `backend_tls` | -- | -- | TLS configuration for kv store backend (applicable for etcd, PostgreSQL, and MySQL backends)
When using etcd, PostgreSQL, or MySQL as metadata store, you can configure TLS here

Note: if TLS is configured in both this section and the `store_addrs` connection string, the
settings here will override the TLS settings in `store_addrs`. | | `backend_tls.mode` | String | `prefer` | TLS mode, refer to https://www.postgresql.org/docs/current/libpq-ssl.html
- "disable" - No TLS
- "prefer" (default) - Try TLS, fallback to plain
- "require" - Require TLS
- "verify_ca" - Require TLS and verify CA
- "verify_full" - Require TLS and verify hostname | | `backend_tls.cert_path` | String | `""` | Path to client certificate file (for client authentication)
Like "/path/to/client.crt" | @@ -509,8 +509,8 @@ | `runtime` | -- | -- | The runtime options. | | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 1). | -| `runtime.query_rt_size` | Integer | `7` | The number of threads to execute datanode query operations.
Defaults to max(num_cpus - 1, 1). | +| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.query_rt_size` | Integer | `7` | The number of threads to execute datanode query operations.
Defaults to max(num_cpus - 1, 2). | | `runtime.ingest_rt_size` | Integer | `8` | The number of threads to execute datanode ingestion operations. | | `meta_client` | -- | -- | The metasrv client options. | | `meta_client.metasrv_addrs` | Array | -- | The addresses of the metasrv. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 1f1377d91cc..5966405452a 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -84,10 +84,10 @@ watch = false ## The number of threads to execute compact operations. #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. -## Defaults to max(num_cpus / 2, 1). +## Defaults to max(num_cpus / 2, 2). #+ compact_rt_max_blocking_threads = 4 ## The number of threads to execute datanode query operations. -## Defaults to max(num_cpus - 1, 1). +## Defaults to max(num_cpus - 1, 2). #+ query_rt_size = 7 ## The number of threads to execute datanode ingestion operations. #+ ingest_rt_size = 8 diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 260984a4ab3..ec1f9c89b42 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -44,7 +44,7 @@ default_column_prefix = "greptime" ## The number of threads to execute compact operations. #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. -## Defaults to max(num_cpus / 2, 1). +## Defaults to max(num_cpus / 2, 2). #+ compact_rt_max_blocking_threads = 4 ## The HTTP server options. diff --git a/config/metasrv.example.toml b/config/metasrv.example.toml index f607211d636..6017bfcc49a 100644 --- a/config/metasrv.example.toml +++ b/config/metasrv.example.toml @@ -92,7 +92,7 @@ node_max_idle_time = "24hours" ## The number of threads to execute compact operations. #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. -## Defaults to max(num_cpus / 2, 1). +## Defaults to max(num_cpus / 2, 2). #+ compact_rt_max_blocking_threads = 4 ## TLS configuration for kv store backend (applicable for etcd, PostgreSQL, and MySQL backends) diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 068025e2a28..5ada605b92b 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -57,7 +57,7 @@ max_concurrent_queries = 0 ## The number of threads to execute compact operations. #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. -## Defaults to max(num_cpus / 2, 1). +## Defaults to max(num_cpus / 2, 2). #+ compact_rt_max_blocking_threads = 4 ## The HTTP server options. diff --git a/src/cmd/src/datanode/builder.rs b/src/cmd/src/datanode/builder.rs index afabcb85f10..3dfa03e52e5 100644 --- a/src/cmd/src/datanode/builder.rs +++ b/src/cmd/src/datanode/builder.rs @@ -61,9 +61,6 @@ impl InstanceBuilder { } async fn init(opts: &mut DatanodeOptions, plugins: &mut Plugins) -> Result> { - common_runtime::init_global_runtimes(&opts.runtime); - common_runtime::init_datanode_runtimes(&opts.runtime); - let dn_opts = &mut opts.component; let guard = common_telemetry::init_global_logging( APP_NAME, @@ -73,6 +70,9 @@ impl InstanceBuilder { None, ); + common_runtime::init_global_runtimes(&opts.runtime); + common_runtime::init_datanode_runtimes(&opts.runtime); + crate::options::flush_dropped_plugin_warnings(); log_versions(verbose_version(), short_version(), APP_NAME); maybe_activate_heap_profile(&dn_opts.memory); diff --git a/src/cmd/src/flownode.rs b/src/cmd/src/flownode.rs index ee071875c3a..eccfaa63b07 100644 --- a/src/cmd/src/flownode.rs +++ b/src/cmd/src/flownode.rs @@ -248,8 +248,6 @@ impl StartCommand { } async fn build(&self, opts: FlownodeOptions) -> Result { - common_runtime::init_global_runtimes(&opts.runtime); - let guard = common_telemetry::init_global_logging( APP_NAME, &opts.component.logging, @@ -258,6 +256,8 @@ impl StartCommand { None, ); + common_runtime::init_global_runtimes(&opts.runtime); + crate::options::flush_dropped_plugin_warnings(); log_versions(verbose_version(), short_version(), APP_NAME); maybe_activate_heap_profile(&opts.component.memory); diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index 9d5d0bab1ee..e3737088e8d 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -330,8 +330,6 @@ impl StartCommand { } async fn build(&self, opts: FrontendOptions) -> Result { - common_runtime::init_global_runtimes(&opts.runtime); - let guard = common_telemetry::init_global_logging( APP_NAME, &opts.component.logging, @@ -340,6 +338,8 @@ impl StartCommand { Some(&opts.component.slow_query), ); + common_runtime::init_global_runtimes(&opts.runtime); + crate::options::flush_dropped_plugin_warnings(); log_versions(verbose_version(), short_version(), APP_NAME); maybe_activate_heap_profile(&opts.component.memory); diff --git a/src/cmd/src/metasrv.rs b/src/cmd/src/metasrv.rs index 791efeb6e86..369f5c428d1 100644 --- a/src/cmd/src/metasrv.rs +++ b/src/cmd/src/metasrv.rs @@ -318,8 +318,6 @@ impl StartCommand { } pub async fn build(&self, opts: MetasrvOptions) -> Result { - common_runtime::init_global_runtimes(&opts.runtime); - let guard = common_telemetry::init_global_logging( APP_NAME, &opts.component.logging, @@ -328,6 +326,8 @@ impl StartCommand { None, ); + common_runtime::init_global_runtimes(&opts.runtime); + crate::options::flush_dropped_plugin_warnings(); log_versions(verbose_version(), short_version(), APP_NAME); maybe_activate_heap_profile(&opts.component.memory); diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 9f8155b9cc2..2ffcc5e729c 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -395,8 +395,6 @@ impl StartCommand { #[allow(clippy::diverging_sub_expression)] /// Build GreptimeDB instance with the loaded options. pub async fn build(&self, opts: GreptimeOptions) -> Result { - common_runtime::init_global_runtimes(&opts.runtime); - let guard = common_telemetry::init_global_logging( APP_NAME, &opts.component.logging, @@ -405,6 +403,8 @@ impl StartCommand { Some(&opts.component.slow_query), ); + common_runtime::init_global_runtimes(&opts.runtime); + crate::options::flush_dropped_plugin_warnings(); log_versions(verbose_version(), short_version(), APP_NAME); maybe_activate_heap_profile(&opts.component.memory); diff --git a/src/common/runtime/src/global.rs b/src/common/runtime/src/global.rs index e28eaf61d15..827832fca22 100644 --- a/src/common/runtime/src/global.rs +++ b/src/common/runtime/src/global.rs @@ -27,6 +27,9 @@ use crate::{Builder, JoinHandle, Runtime}; const GLOBAL_WORKERS: usize = 8; const COMPACT_WORKERS: usize = 4; const HB_WORKERS: usize = 2; +/// The minimum number of worker threads for runtimes sized by CPU count. +/// A single-threaded runtime can easily deadlock in async code. +const MIN_RUNTIME_THREADS: usize = 2; /// The options for the global runtimes. #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] @@ -44,19 +47,25 @@ pub struct RuntimeOptions { pub ingest_rt_size: usize, } -impl Default for RuntimeOptions { - fn default() -> Self { - let cpus = num_cpus::get(); +impl RuntimeOptions { + fn with_num_cpus(cpus: usize) -> Self { + let cpus = usize::max(cpus, MIN_RUNTIME_THREADS); Self { global_rt_size: cpus, - compact_rt_size: usize::max(cpus / 2, 1), - compact_rt_max_blocking_threads: usize::max(cpus / 2, 1), - query_rt_size: usize::max(cpus.saturating_sub(1), 1), + compact_rt_size: usize::max(cpus / 2, MIN_RUNTIME_THREADS), + compact_rt_max_blocking_threads: usize::max(cpus / 2, MIN_RUNTIME_THREADS), + query_rt_size: usize::max(cpus.saturating_sub(1), MIN_RUNTIME_THREADS), ingest_rt_size: cpus, } } } +impl Default for RuntimeOptions { + fn default() -> Self { + Self::with_num_cpus(num_cpus::get()) + } +} + pub fn create_runtime(runtime_name: &str, thread_name: &str, worker_threads: usize) -> Runtime { info!( "Creating runtime with runtime_name: {runtime_name}, thread_name: {thread_name}, work_threads: {worker_threads}." @@ -287,18 +296,56 @@ mod tests { #[test] fn test_datanode_runtime_options_default() { let options = RuntimeOptions::default(); - let cpus = num_cpus::get(); + let cpus = usize::max(num_cpus::get(), MIN_RUNTIME_THREADS); assert_eq!(cpus, options.global_rt_size); - assert_eq!(usize::max(cpus / 2, 1), options.compact_rt_size); assert_eq!( - usize::max(cpus / 2, 1), + usize::max(cpus / 2, MIN_RUNTIME_THREADS), + options.compact_rt_size + ); + assert_eq!( + usize::max(cpus / 2, MIN_RUNTIME_THREADS), options.compact_rt_max_blocking_threads ); - assert_eq!(usize::max(cpus.saturating_sub(1), 1), options.query_rt_size); + assert_eq!( + usize::max(cpus.saturating_sub(1), MIN_RUNTIME_THREADS), + options.query_rt_size + ); assert_eq!(cpus, options.ingest_rt_size); } + #[test] + fn test_runtime_options_min_threads() { + for cpus in [0, 1, 2] { + let options = RuntimeOptions::with_num_cpus(cpus); + assert!( + options.global_rt_size >= MIN_RUNTIME_THREADS, + "global_rt_size {} < {MIN_RUNTIME_THREADS} with {cpus} cpus", + options.global_rt_size + ); + assert!( + options.compact_rt_size >= MIN_RUNTIME_THREADS, + "compact_rt_size {} < {MIN_RUNTIME_THREADS} with {cpus} cpus", + options.compact_rt_size + ); + assert!( + options.compact_rt_max_blocking_threads >= MIN_RUNTIME_THREADS, + "compact_rt_max_blocking_threads {} < {MIN_RUNTIME_THREADS} with {cpus} cpus", + options.compact_rt_max_blocking_threads + ); + assert!( + options.query_rt_size >= MIN_RUNTIME_THREADS, + "query_rt_size {} < {MIN_RUNTIME_THREADS} with {cpus} cpus", + options.query_rt_size + ); + assert!( + options.ingest_rt_size >= MIN_RUNTIME_THREADS, + "ingest_rt_size {} < {MIN_RUNTIME_THREADS} with {cpus} cpus", + options.ingest_rt_size + ); + } + } + #[test] fn test_datanode_runtimes_fallback_to_global_runtime() { let runtimes = GlobalRuntimes::new( From 8fac712870d5f95aef6ccdec8e9ae45f872f4afb Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 19 Aug 2026 12:43:42 +0000 Subject: [PATCH 088/126] fix(meta): avoid blocking runtime on stats cache lock (#8910) Signed-off-by: WenyXu --- .../src/handler/collect_stats_handler.rs | 293 +++++++++++++++++- 1 file changed, 289 insertions(+), 4 deletions(-) diff --git a/src/meta-srv/src/handler/collect_stats_handler.rs b/src/meta-srv/src/handler/collect_stats_handler.rs index d444498798f..e3ae7ebe915 100644 --- a/src/meta-srv/src/handler/collect_stats_handler.rs +++ b/src/meta-srv/src/handler/collect_stats_handler.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::cmp::Ordering; +use std::sync::Arc; use api::v1::meta::{HeartbeatRequest, Role}; use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; @@ -24,6 +25,7 @@ use common_meta::rpc::store::PutRequest; use common_telemetry::{error, info, warn}; use dashmap::DashMap; use snafu::ResultExt; +use tokio::sync::Mutex; use crate::error::{self, Result}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; @@ -70,7 +72,7 @@ impl EpochStats { const DEFAULT_FLUSH_STATS_FACTOR: usize = 3; pub struct CollectStatsHandler { - stats_cache: DashMap, + stats_cache: DashMap>>, flush_stats_factor: usize, } @@ -106,10 +108,16 @@ impl HeartbeatHandler for CollectStatsHandler { }; let key = current_stat.stat_key(); - let mut entry = self.stats_cache.entry(key).or_default(); + let state = { + let entry = self + .stats_cache + .entry(key) + .or_insert_with(|| Arc::new(Mutex::new(EpochStats::default()))); + Arc::clone(entry.value()) + }; + let mut epoch_stats = state.lock().await; let key: Vec = key.into(); - let epoch_stats = entry.value_mut(); let refresh = if let Some(epoch) = epoch_stats.epoch() { match current_stat.node_epoch.cmp(&epoch) { @@ -170,7 +178,7 @@ impl HeartbeatHandler for CollectStatsHandler { } } -async fn rewrite_node_address(ctx: &mut Context, stat: &Stat) { +async fn rewrite_node_address(ctx: &Context, stat: &Stat) { let peer = Peer { id: stat.id, addr: stat.addr.clone(), @@ -217,10 +225,188 @@ async fn rewrite_node_address(ctx: &mut Context, stat: &Stat) { #[cfg(test)] mod tests { + use std::any::Any; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering}; + use std::sync::{Arc, Mutex as StdMutex, mpsc}; + use std::thread; + use std::time::Duration; + use common_meta::datanode::DatanodeStatKey; + use common_meta::error::{Error as MetaError, Result as MetaResult}; + use common_meta::kv_backend::{KvBackend, KvBackendRef, ResettableKvBackend, TxnService}; + use common_meta::rpc::store::{ + BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, + BatchPutRequest, BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, PutResponse, + RangeRequest, RangeResponse, + }; + use tokio::sync::Semaphore; + use tokio::time::{sleep, timeout}; use super::*; use crate::handler::test_utils::TestEnv; + use crate::service::store::cached_kv::LeaderCachedKvBackend; + + struct ControlledKvBackend { + recorded_puts: StdMutex>, + put_entered: Semaphore, + put_release: Semaphore, + block_next_put: AtomicBool, + delay_puts: AtomicBool, + active_puts: AtomicUsize, + max_active_puts: AtomicUsize, + } + + impl ControlledKvBackend { + fn new() -> Self { + Self { + recorded_puts: StdMutex::new(Vec::new()), + put_entered: Semaphore::new(0), + put_release: Semaphore::new(0), + block_next_put: AtomicBool::new(false), + delay_puts: AtomicBool::new(false), + active_puts: AtomicUsize::new(0), + max_active_puts: AtomicUsize::new(0), + } + } + + fn block_next_put(&self) { + self.block_next_put.store(true, AtomicOrdering::Relaxed); + } + + async fn wait_for_blocked_put(&self) { + self.put_entered.acquire().await.unwrap().forget(); + } + + fn release_one_put(&self) { + self.put_release.add_permits(1); + } + + fn set_delay_puts(&self, delay: bool) { + self.delay_puts.store(delay, AtomicOrdering::Relaxed); + } + + fn clear_recorded_puts(&self) { + self.recorded_puts.lock().unwrap().clear(); + } + + fn recorded_puts(&self) -> Vec { + self.recorded_puts.lock().unwrap().clone() + } + + fn max_active_puts(&self) -> usize { + self.max_active_puts.load(AtomicOrdering::Relaxed) + } + + fn start_put(&self, req: &PutRequest) -> PutGuard<'_> { + self.recorded_puts.lock().unwrap().push(req.clone()); + let active = self.active_puts.fetch_add(1, AtomicOrdering::Relaxed) + 1; + self.max_active_puts + .fetch_max(active, AtomicOrdering::Relaxed); + PutGuard { backend: self } + } + } + + struct PutGuard<'a> { + backend: &'a ControlledKvBackend, + } + + impl Drop for PutGuard<'_> { + fn drop(&mut self) { + self.backend + .active_puts + .fetch_sub(1, AtomicOrdering::Relaxed); + } + } + + #[async_trait::async_trait] + impl TxnService for ControlledKvBackend { + type Error = MetaError; + } + + #[async_trait::async_trait] + impl KvBackend for ControlledKvBackend { + fn name(&self) -> &str { + "controlled" + } + + fn as_any(&self) -> &dyn Any { + self + } + + async fn range(&self, _req: RangeRequest) -> MetaResult { + unimplemented!() + } + + async fn put(&self, req: PutRequest) -> MetaResult { + let _guard = self.start_put(&req); + if self.block_next_put.swap(false, AtomicOrdering::Relaxed) { + self.put_entered.add_permits(1); + self.put_release.acquire().await.unwrap().forget(); + } + if self.delay_puts.load(AtomicOrdering::Relaxed) { + sleep(Duration::from_millis(10)).await; + } + Ok(PutResponse::default()) + } + + async fn batch_put(&self, _req: BatchPutRequest) -> MetaResult { + unimplemented!() + } + + async fn batch_get(&self, _req: BatchGetRequest) -> MetaResult { + unimplemented!() + } + + async fn delete_range(&self, _req: DeleteRangeRequest) -> MetaResult { + unimplemented!() + } + + async fn batch_delete(&self, _req: BatchDeleteRequest) -> MetaResult { + unimplemented!() + } + } + + impl ResettableKvBackend for ControlledKvBackend { + fn reset(&self) { + self.clear_recorded_puts(); + } + + fn as_kv_backend_ref(self: Arc) -> KvBackendRef { + self + } + } + + fn stat(node_id: u64, epoch: u64, marker: u64, addr: &str) -> Stat { + Stat { + timestamp_millis: marker as i64, + id: node_id, + addr: addr.to_string(), + region_num: marker, + node_epoch: epoch, + ..Default::default() + } + } + + async fn handle_stat( + handler: Arc, + mut ctx: Context, + stat: Stat, + ) -> Result { + let mut acc = HeartbeatAccumulator { + stat: Some(stat), + ..Default::default() + }; + handler + .handle(&HeartbeatRequest::default(), &mut ctx, &mut acc) + .await + } + + fn use_controlled_address_backend(ctx: &mut Context) -> Arc { + let backend = Arc::new(ControlledKvBackend::new()); + ctx.leader_cached_kv_backend = + Arc::new(LeaderCachedKvBackend::with_always_leader(backend.clone())); + backend + } #[tokio::test] async fn test_handle_datanode_stats() { @@ -250,6 +436,105 @@ mod tests { assert_eq!(handler.flush_stats_factor, val.stats.len()); } + #[test] + fn test_same_datanode_wait_keeps_current_thread_runtime_responsive() { + let (backend_tx, backend_rx) = mpsc::sync_channel(1); + let (timer_tx, timer_rx) = mpsc::sync_channel(1); + let (done_tx, done_rx) = mpsc::sync_channel(1); + + let worker = thread::spawn(move || { + tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap() + .block_on(async move { + let env = TestEnv::new(); + let mut ctx = env.ctx(); + let address_backend = use_controlled_address_backend(&mut ctx); + address_backend.block_next_put(); + backend_tx.send(address_backend.clone()).unwrap(); + + let handler = Arc::new(CollectStatsHandler::default()); + let first = tokio::spawn(handle_stat( + handler.clone(), + ctx.clone(), + stat(101, 1, 1, "dn-101-v1"), + )); + address_backend.wait_for_blocked_put().await; + + let second = + tokio::spawn(handle_stat(handler, ctx, stat(101, 1, 2, "dn-101-v1"))); + tokio::spawn(async move { + sleep(Duration::from_millis(10)).await; + timer_tx.send(()).unwrap(); + }); + tokio::task::yield_now().await; + + first.await.unwrap().unwrap(); + second.await.unwrap().unwrap(); + done_tx.send(()).unwrap(); + }); + }); + + let address_backend = backend_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + let timer_result = timer_rx.recv_timeout(Duration::from_secs(1)); + address_backend.release_one_put(); + timer_result.expect("the current-thread runtime must remain responsive"); + done_rx + .recv_timeout(Duration::from_secs(1)) + .expect("both heartbeat handlers must complete after releasing the write"); + worker.join().unwrap(); + } + + #[tokio::test] + async fn test_concurrent_flush_persists_every_stat_once() { + let env = TestEnv::new(); + let mut ctx = env.ctx(); + let stats_backend = Arc::new(ControlledKvBackend::new()); + ctx.in_memory = stats_backend.clone(); + stats_backend.set_delay_puts(true); + + let flush_stats_factor = 3; + let handler = Arc::new(CollectStatsHandler::new(Some(flush_stats_factor))); + handle_stat(handler.clone(), ctx.clone(), stat(101, 1, 0, "dn-101")) + .await + .unwrap(); + stats_backend.clear_recorded_puts(); + + let mut tasks = Vec::with_capacity(2 * flush_stats_factor); + for marker in 1..=(2 * flush_stats_factor) { + tasks.push(tokio::spawn(handle_stat( + handler.clone(), + ctx.clone(), + stat(101, 1, marker as u64, "dn-101"), + ))); + } + for task in tasks { + timeout(Duration::from_secs(1), task) + .await + .unwrap() + .unwrap() + .unwrap(); + } + + let puts = stats_backend.recorded_puts(); + assert_eq!(2, puts.len()); + let mut markers = puts + .into_iter() + .flat_map(|put| { + let value: DatanodeStatValue = put.value.try_into().unwrap(); + value + .stats + .into_iter() + .map(|stat| stat.region_num) + .collect::>() + }) + .collect::>(); + markers.sort_unstable(); + assert_eq!((1..=6).collect::>(), markers); + assert_eq!(1, stats_backend.max_active_puts()); + } + async fn handle_request_many_times( mut ctx: Context, handler: &CollectStatsHandler, From 76f08d2b3f9f5905c9169244dbed8b8bd2178973 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Fri, 21 Aug 2026 03:28:39 +0000 Subject: [PATCH 089/126] refactor(json2): add bounded auto-expansion to the JSON2 vector builder (#8909) * refactor(json2): add bounded auto-expansion to the JSON2 vector builder Signed-off-by: luofucong * resolve PR comments Signed-off-by: luofucong * fix ci Signed-off-by: luofucong --------- Signed-off-by: luofucong --- src/datatypes/src/vectors/json/builder.rs | 841 ++++++++++++++++++++-- src/datatypes/src/vectors/json/variant.rs | 11 +- 2 files changed, 774 insertions(+), 78 deletions(-) diff --git a/src/datatypes/src/vectors/json/builder.rs b/src/datatypes/src/vectors/json/builder.rs index e6ceb6d6a5f..125dbdad58a 100644 --- a/src/datatypes/src/vectors/json/builder.rs +++ b/src/datatypes/src/vectors/json/builder.rs @@ -13,72 +13,488 @@ // limitations under the License. use std::any::Any; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; +use arrow_array::cast::AsArray; +use arrow_array::{Array, ArrayRef, StructArray}; use arrow_schema::DataType; +use parquet_variant_compute::VariantArrayBuilder; +use snafu::{ResultExt, ensure}; use crate::data_type::ConcreteDataType; -use crate::error::{Result, TryFromValueSnafu, UnexpectedSnafu, UnsupportedOperationSnafu}; +use crate::error::{ + ArrowComputeSnafu, Result, TryFromValueSnafu, UnexpectedSnafu, UnsupportedOperationSnafu, +}; +use crate::extension::json::JSON2_REMAINDER_FIELD_NAME; use crate::json::value::{JsonNumber, JsonVariant, encode_json_variant}; +use crate::json::{JSON2_MAX_STRUCTURED_DEPTH, JsonSettings}; use crate::prelude::{ValueRef, Vector, VectorRef}; use crate::types::StructType; use crate::types::json_type::{JsonNativeType, is_include}; use crate::value::{ListValue, StructValue, StructValueRef, Value}; -use crate::vectors::{MutableVector, StructVectorBuilder}; +use crate::vectors::json::variant::{append_json_variant, variant_field}; +use crate::vectors::{Helper, MutableVector, StructVectorBuilder}; +type JsonObjectValue = BTreeMap; + +/// Builds JSON2 vectors from object values. +/// +/// Legacy mode merges all observed paths into the explicit Struct schema. +/// Auto-expanding mode always materializes type-hinted paths, selects up to +/// `max_auto_expanded_paths` compatible unhinted leaf paths by frequency, and stores +/// conflicting or unselected paths in the Variant remainder field. #[derive(Clone)] pub(crate) struct JsonVectorBuilder { - merged_type: JsonNativeType, - values: Vec, + state: JsonVectorBuilderState, +} + +#[derive(Clone)] +enum JsonVectorBuilderState { + Legacy { + merged_type: JsonNativeType, + values: Vec, + }, + AutoExpanding { + /// Paths declared by type hints and always stored as dedicated Struct fields. + explicit_type: JsonNativeType, + /// Maximum number of additional paths selected from buffered values. + max_auto_expanded_paths: u32, + /// Buffered values used to infer auto-expanded paths before building the vector. + values: Vec, + }, +} + +impl JsonVectorBuilderState { + fn native_type(&self) -> JsonNativeType { + match self { + Self::Legacy { merged_type, .. } => merged_type.clone(), + Self::AutoExpanding { + explicit_type, + max_auto_expanded_paths, + values, + } => infer_expanded_type(explicit_type, *max_auto_expanded_paths, values), + } + } + + fn len(&self) -> usize { + match self { + Self::Legacy { values, .. } | Self::AutoExpanding { values, .. } => values.len(), + } + } + + fn try_build(&mut self) -> Result { + match self { + Self::Legacy { + merged_type, + values, + } => build_legacy(values, merged_type), + Self::AutoExpanding { + explicit_type, + max_auto_expanded_paths, + values, + } => { + let expanded_type = + infer_expanded_type(explicit_type, *max_auto_expanded_paths, values); + build_with_remainder(values, &expanded_type) + } + } + } + + fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> { + let ValueRef::Json(value) = value else { + return TryFromValueSnafu { + reason: format!("expected JSON value, got {value:?}"), + } + .fail(); + }; + ensure!( + value.is_object() || value.is_null(), + TryFromValueSnafu { + reason: format!("expected JSON object value, got {value:?}"), + } + ); + match self { + Self::Legacy { + merged_type, + values, + } => { + let json_type = value.json_type(); + if !is_include(merged_type, json_type.as_ref()) { + merged_type.merge(json_type.as_ref()); + } + values.push(JsonVariant::from(value.variant())); + } + Self::AutoExpanding { values, .. } => { + values.push(JsonVariant::from(value.variant())); + } + } + Ok(()) + } + + fn push_null(&mut self) { + match self { + Self::Legacy { values, .. } | Self::AutoExpanding { values, .. } => { + values.push(JsonVariant::Null) + } + } + } } impl JsonVectorBuilder { + /// Creates a builder that merges all observed paths into the explicit schema. pub(crate) fn new(initial_native_type: JsonNativeType, capacity: usize) -> Self { debug_assert!(matches!( initial_native_type, JsonNativeType::Object(_) | JsonNativeType::Null )); Self { - merged_type: initial_native_type, - values: Vec::with_capacity(capacity), + state: JsonVectorBuilderState::Legacy { + merged_type: initial_native_type, + values: Vec::with_capacity(capacity), + }, + } + } + + /// Creates a builder bounded by the JSON settings and their type hints. + #[cfg_attr(not(test), expect(dead_code))] + pub(crate) fn with_settings(settings: &JsonSettings, capacity: usize) -> Self { + let mut explicit_type = JsonNativeType::Object(Default::default()); + for hint in settings.type_hints() { + insert_dynamic_type(&mut explicit_type, &hint.path, (&hint.data_type).into()); + } + Self { + state: JsonVectorBuilderState::AutoExpanding { + explicit_type, + max_auto_expanded_paths: settings.max_auto_expanded_paths().unwrap_or(u32::MAX), + values: Vec::with_capacity(capacity), + }, } } fn try_build(&mut self) -> Result { - let DataType::Struct(fields) = self.merged_type.as_arrow_type() else { - return UnexpectedSnafu { - reason: "merged JSON2 type must map to Arrow Struct in JsonVectorBuilder", - } - .fail(); - }; - // TODO(LFC): Direct use Arrow's Struct datatype here. - let struct_type = StructType::from(&fields); - - let mut builder = - StructVectorBuilder::with_type_and_capacity(struct_type.clone(), self.values.len()); - for value in std::mem::take(&mut self.values) { - if matches!(&value, JsonVariant::Null) { - builder.push_null(); - continue; - } - let value = json_variant_into_struct_value(value, struct_type.clone())?; - builder.push_struct_value_ref(StructValueRef::Ref(&value))?; - } - Ok(builder.to_vector()) + self.state.try_build() } } -fn json_variant_into_struct_value( - value: JsonVariant, - struct_type: StructType, -) -> Result { - let JsonVariant::Object(object) = value else { - return TryFromValueSnafu { - reason: format!("expected json object value, got {value:?}"), +fn build_legacy(values: &mut Vec, merged_type: &JsonNativeType) -> Result { + build_explicit(values, merged_type, false, |value| match value { + JsonVariant::Null => Ok(None), + JsonVariant::Object(value) => Ok(Some(value)), + _ => TryFromValueSnafu { + reason: "expected json object value".to_string(), + } + .fail(), + }) +} + +fn build_with_remainder( + values: &mut Vec, + expanded_type: &JsonNativeType, +) -> Result { + let mut remainder = VariantArrayBuilder::new(values.len()); + let explicit = build_explicit(values, expanded_type, true, |value| { + if matches!(value, JsonVariant::Null) { + remainder.append_null(); + return Ok(None); + } + let (value, rest) = split_to_explicit(value, expanded_type)?; + append_json_variant(&mut remainder, &JsonVariant::Object(rest)) + .context(ArrowComputeSnafu)?; + Ok(Some(value)) + })?; + finish_vector(explicit, ArrayRef::from(remainder.build())) +} + +fn build_explicit( + values: &mut Vec, + explicit_type: &JsonNativeType, + // Temporary compatibility switch for the legacy storage layout. Once JSON2 fully switches to + // the v2 storage layout, empty objects should always be preserved instead of treated as null. + preserve_empty_structs: bool, + mut project: impl FnMut(JsonVariant) -> Result>, +) -> Result { + let DataType::Struct(fields) = explicit_type.as_arrow_type() else { + return UnexpectedSnafu { + reason: "merged JSON2 type must map to Arrow Struct in JsonVectorBuilder", } .fail(); }; + // TODO(LFC): Direct use Arrow's Struct datatype here. + let struct_type = StructType::from(&fields); + let mut builder = + StructVectorBuilder::with_type_and_capacity(struct_type.clone(), values.len()); + for value in std::mem::take(values) { + let Some(value) = project(value)? else { + builder.push_null(); + continue; + }; + let value = + json_variant_into_struct_value(value, struct_type.clone(), preserve_empty_structs)?; + builder.push_struct_value_ref(StructValueRef::Ref(&value))?; + } + Ok(builder.to_vector()) +} + +fn finish_vector(explicit: VectorRef, remainder: ArrayRef) -> Result { + let explicit = explicit.to_arrow_array(); + let explicit = explicit.as_struct(); + let mut children = explicit + .fields() + .iter() + .cloned() + .zip(explicit.columns().iter().cloned()) + .chain(std::iter::once(( + Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)), + remainder, + ))) + .collect::>(); + children.sort_unstable_by(|(x, _), (y, _)| x.name().cmp(y.name())); + let (fields, columns): (Vec<_>, Vec) = children.into_iter().unzip(); + let array: ArrayRef = Arc::new(StructArray::new( + fields.into(), + columns, + explicit.nulls().cloned(), + )); + Helper::try_into_vector(array) +} + +fn infer_expanded_type( + explicit_type: &JsonNativeType, + max_auto_expanded_paths: u32, + values: &[JsonVariant], +) -> JsonNativeType { + if max_auto_expanded_paths == 0 { + return explicit_type.clone(); + } + + let mut stats = HashMap::new(); + let mut path = Vec::new(); + init_explicit_path_stats(explicit_type, &mut path, &mut stats); + for value in values { + count_dynamic_paths(value, &mut path, &mut stats); + } + let mut candidates = stats + .iter() + // Explicit paths are already in the output schema and do not consume the dynamic + // expansion budget. + .filter(|(_, stat)| !stat.is_explicit && stat.is_leaf) + // A leaf is eligible only when both itself and every object prefix have one stable + // role and type across all observed values. + .filter(|(path, _)| { + !(1..=path.len()) + .any(|len| stats.get(&path[..len]).is_some_and(|stats| stats.conflicts)) + }) + .collect::>(); + candidates.sort_unstable_by(|(x_path, x), (y_path, y)| { + y.seen_count + .cmp(&x.seen_count) + .then_with(|| x_path.cmp(y_path)) + }); + + let mut expanded_type = explicit_type.clone(); + for (path, candidate) in candidates + .into_iter() + .take(max_auto_expanded_paths as usize) + { + insert_dynamic_type( + &mut expanded_type, + path, + candidate.expected_leaf_type.clone(), + ); + } + expanded_type +} + +/// Aggregated observations for one JSON path. +/// +/// Schema inference first seeds the map with explicit paths, then walks all input values once. +/// Objects, including empty objects, are non-leaf paths; every other non-null value is a leaf. +/// The first dynamic observation fixes the path role and exact leaf type. A later role or type +/// mismatch sets [`PathStats::conflicts`] permanently. Missing paths and null values do not affect +/// the statistics. +/// +/// After collection, dynamic leaves are ranked by [`PathStats::seen_count`]. A candidate is +/// rejected when it or any parent path conflicts, so candidate selection never rescans the input +/// values. +struct PathStats { + /// Whether the path came from a type hint and is already part of the output schema. + is_explicit: bool, + /// Whether the path is a non-object value rather than an object prefix. + is_leaf: bool, + /// Exact type required for a leaf; unused non-leaf paths keep [`JsonNativeType::Null`]. + expected_leaf_type: JsonNativeType, + /// Number of compatible observations used to rank dynamic leaves. + seen_count: usize, + /// Whether the path has ever had inconsistent roles or leaf types. + conflicts: bool, +} + +/// Seeds path statistics from the configured explicit JSON shape. +fn init_explicit_path_stats<'a>( + explicit_type: &'a JsonNativeType, + path: &mut Vec<&'a str>, + stats: &mut HashMap, PathStats>, +) { + let JsonNativeType::Object(fields) = explicit_type else { + return; + }; + for (name, data_type) in fields { + path.push(name); + let is_leaf = !matches!(data_type, JsonNativeType::Object(_)); + let expected_leaf_type = if is_leaf { + data_type.clone() + } else { + JsonNativeType::default() + }; + stats.insert( + path.clone(), + PathStats { + is_explicit: true, + is_leaf, + expected_leaf_type, + seen_count: 0, + conflicts: false, + }, + ); + init_explicit_path_stats(data_type, path, stats); + path.pop(); + } +} + +/// Collects dynamic path statistics while traversing each input value once. +fn count_dynamic_paths<'a>( + value: &'a JsonVariant, + path: &mut Vec<&'a str>, + stats: &mut HashMap, PathStats>, +) { + if matches!(value, JsonVariant::Null) || path.len() > JSON2_MAX_STRUCTURED_DEPTH { + return; + } + + if !path.is_empty() { + let is_leaf = !matches!(value, JsonVariant::Object(_)); + let conflicts = if let Some(stats) = stats.get_mut(path.as_slice()) { + if !stats.conflicts { + let role_conflict = stats.is_leaf != is_leaf; + let type_conflict = || match (&stats.expected_leaf_type, value) { + // If both objects, they are compatible. + (JsonNativeType::Null | JsonNativeType::Object(_), JsonVariant::Object(_)) => { + false + } + _ => stats.expected_leaf_type != value.native_type(), + }; + if role_conflict || type_conflict() { + stats.conflicts = true; + } else { + stats.seen_count += 1; + } + } + stats.conflicts + } else { + let expected_leaf_type = if is_leaf { + value.native_type() + } else { + JsonNativeType::default() + }; + stats.insert( + path.clone(), + PathStats { + is_explicit: false, + is_leaf, + expected_leaf_type, + seen_count: 1, + conflicts: false, + }, + ); + false + }; + if conflicts { + return; + } + } + + if let JsonVariant::Object(object) = value + && !object.is_empty() + { + for (name, value) in object { + path.push(name); + count_dynamic_paths(value, path, stats); + path.pop(); + } + } +} + +fn insert_dynamic_type>( + explicit_type: &mut JsonNativeType, + path: &[S], + data_type: JsonNativeType, +) { + let JsonNativeType::Object(fields) = explicit_type else { + return; + }; + let Some((name, path)) = path.split_first() else { + return; + }; + let name = name.as_ref().to_string(); + if path.is_empty() { + fields.insert(name, data_type); + return; + } + insert_dynamic_type( + fields + .entry(name) + .or_insert_with(|| JsonNativeType::Object(Default::default())), + path, + data_type, + ) +} + +fn split_to_explicit( + value: JsonVariant, + explicit_type: &JsonNativeType, +) -> Result<(JsonObjectValue, JsonObjectValue)> { + let JsonVariant::Object(mut remainder) = value else { + return TryFromValueSnafu { + reason: "expected json object value".to_string(), + } + .fail(); + }; + let JsonNativeType::Object(fields) = explicit_type else { + return UnexpectedSnafu { + reason: "JSON2 explicit type must be an object", + } + .fail(); + }; + let mut explicit = JsonObjectValue::new(); + + for (name, data_type) in fields { + let Some(value) = remainder.remove(name) else { + continue; + }; + if matches!(data_type, JsonNativeType::Object(_)) { + if matches!(value, JsonVariant::Null) { + explicit.insert(name.clone(), value); + continue; + } + let (child_explicit, child_remainder) = split_to_explicit(value, data_type)?; + explicit.insert(name.clone(), JsonVariant::Object(child_explicit)); + if !child_remainder.is_empty() { + remainder.insert(name.clone(), JsonVariant::Object(child_remainder)); + } + } else { + explicit.insert(name.clone(), value); + } + } + Ok((explicit, remainder)) +} + +fn json_variant_into_struct_value( + object: JsonObjectValue, + struct_type: StructType, + preserve_empty_structs: bool, +) -> Result { let mut entries = object.into_iter(); let mut entry = entries.next(); let mut values = Vec::with_capacity(struct_type.fields().len()); @@ -86,7 +502,7 @@ fn json_variant_into_struct_value( let value = match entry.take() { Some((name, value)) if name == field.name() => { entry = entries.next(); - json_variant_into_value(value, field.data_type())? + json_variant_into_value(value, field.data_type(), preserve_empty_structs)? } Some((name, _)) if name.as_str() < field.name() => { return TryFromValueSnafu { @@ -111,10 +527,19 @@ fn json_variant_into_struct_value( Ok(StructValue::new(values, struct_type)) } -fn json_variant_into_value(value: JsonVariant, expected_type: &ConcreteDataType) -> Result { +fn json_variant_into_value( + value: JsonVariant, + expected_type: &ConcreteDataType, + preserve_empty_structs: bool, +) -> Result { let value = match (value, expected_type) { (JsonVariant::Null, _) | (_, ConcreteDataType::Null(_)) => Value::Null, - (JsonVariant::Object(object), _) if object.is_empty() => Value::Null, + (JsonVariant::Object(object), _) if object.is_empty() && !preserve_empty_structs => { + Value::Null + } + (JsonVariant::Object(object), ConcreteDataType::Struct(struct_type)) => Value::Struct( + json_variant_into_struct_value(object, struct_type.clone(), preserve_empty_structs)?, + ), (JsonVariant::Bool(x), ConcreteDataType::Boolean(_)) => Value::Boolean(x), (JsonVariant::Number(x), ConcreteDataType::UInt64(_)) => { let Some(x) = x.as_u64() else { @@ -153,13 +578,10 @@ fn json_variant_into_value(value: JsonVariant, expected_type: &ConcreteDataType) let item_type = list_type.item_type().clone(); let values = array .into_iter() - .map(|v| json_variant_into_value(v, &item_type)) + .map(|v| json_variant_into_value(v, &item_type, preserve_empty_structs)) .collect::>>()?; Value::List(ListValue::new(values, Arc::new(item_type))) } - (value @ JsonVariant::Object(_), ConcreteDataType::Struct(struct_type)) => { - Value::Struct(json_variant_into_struct_value(value, struct_type.clone())?) - } (value, ConcreteDataType::Binary(_)) => Value::from(encode_json_variant(value)?), (value, expected_type) => { return TryFromValueSnafu { @@ -173,11 +595,11 @@ fn json_variant_into_value(value: JsonVariant, expected_type: &ConcreteDataType) impl MutableVector for JsonVectorBuilder { fn data_type(&self) -> ConcreteDataType { - ConcreteDataType::json2(self.merged_type.clone()) + ConcreteDataType::json2(self.state.native_type()) } fn len(&self) -> usize { - self.values.len() + self.state.len() } fn as_any(&self) -> &dyn Any { @@ -197,30 +619,11 @@ impl MutableVector for JsonVectorBuilder { } fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> { - let ValueRef::Json(value) = value else { - return TryFromValueSnafu { - reason: format!("expected json value, got {value:?}"), - } - .fail(); - }; - let json_type = value.json_type(); - let json_type = json_type.as_ref(); - if !matches!(json_type, JsonNativeType::Object(_) | JsonNativeType::Null) { - return TryFromValueSnafu { - reason: format!("expected json object value, got {value:?}"), - } - .fail(); - } - if !is_include(&self.merged_type, json_type) { - self.merged_type.merge(json_type); - } - - self.values.push(JsonVariant::from(value.variant())); - Ok(()) + self.state.try_push_value_ref(value) } fn push_null(&mut self) { - self.values.push(JsonVariant::Null) + self.state.push_null() } fn extend_slice_of(&mut self, _: &dyn Vector, _: usize, _: usize) -> Result<()> { @@ -236,13 +639,21 @@ impl MutableVector for JsonVectorBuilder { mod tests { use std::sync::Arc; + use arrow_array::cast::AsArray; + use arrow_schema::Field; use common_base::bytes::Bytes; + use serde_json::json; use super::*; use crate::data_type::ConcreteDataType; + use crate::extension::json::{Json2ExtensionType, JsonMetadata}; + use crate::json::JsonTypeHint; + use crate::json::value::decode_json_variant; use crate::types::StructField; use crate::types::json_type::JsonObjectType; use crate::value::{ListValue, StructValue, Value, ValueRef}; + use crate::vectors::json::array::JsonArray; + use crate::vectors::json::variant::variant_to_json_values; #[test] fn test_json_vector_builder() -> Result<()> { @@ -346,7 +757,7 @@ mod tests { let err = object_builder .try_push_value_ref(&boolean.as_value_ref()) .unwrap_err(); - assert!(err.to_string().contains("expected json object value")); + assert!(err.to_string().contains("expected JSON object value")); object_builder.try_push_value_ref(&object.as_value_ref())?; // Non-JSON values should be rejected at push time. @@ -355,20 +766,311 @@ mod tests { let err = invalid_builder .try_push_value_ref(&ValueRef::Boolean(true)) .unwrap_err(); - assert!(err.to_string().contains("expected json value")); + assert!(err.to_string().contains("expected JSON value")); + + Ok(()) + } + + #[test] + fn test_zero_budget_builder_uses_fixed_schema_and_remainder() -> Result<()> { + let settings = JsonSettings::try_new( + vec![ + JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }, + JsonTypeHint { + path: vec!["commit".to_string(), "operation".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }, + JsonTypeHint { + path: vec!["time_us".to_string()], + data_type: ConcreteDataType::int64_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }, + ], + Some(0), + )?; + let mut builder = JsonVectorBuilder::with_settings(&settings, 2); + let values = [ + json!({ + "kind": "record", + "commit": {"operation": "create", "collection": "post"}, + "extra": 1, + "time_us": 1 + }), + json!({"kind": "other", "dynamic": true, "time_us": 2}), + ]; + for value in values.clone() { + let value = settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + let array = builder.to_vector().to_arrow_array(); + assert_eq!(0, builder.len()); + let structs = array.as_struct(); + assert_eq!( + vec![JSON2_REMAINDER_FIELD_NAME, "commit", "kind", "time_us"], + structs + .fields() + .iter() + .map(|x| x.name().as_str()) + .collect::>() + ); + assert_eq!( + vec![ + Some(json!({"commit": {"collection": "post"}, "extra": 1})), + Some(json!({"dynamic": true})), + ], + variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())? + ); + + let field = Field::new("data", array.data_type().clone(), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(settings))), + ); + let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?; + let reconstructed = reconstructed.as_binary::(); + assert_eq!( + values[0], + decode_json_variant(reconstructed.value(0)).unwrap() + ); + assert_eq!( + json!({ + "kind": "other", + "commit": {"operation": null}, + "dynamic": true, + "time_us": 2 + }), + decode_json_variant(reconstructed.value(1)).unwrap() + ); + + let settings = JsonSettings::try_new(vec![], Some(0))?; + let mut builder = JsonVectorBuilder::with_settings(&settings, 3); + for value in [json!({}), json!({"x": 1})] { + let value = settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + builder.push_null(); + let array = builder.to_vector().to_arrow_array(); + let structs = array.as_struct(); + assert_eq!(1, structs.num_columns()); + assert_eq!( + vec![Some(json!({})), Some(json!({"x": 1})), None], + variant_to_json_values(structs.column(0))? + ); + + Ok(()) + } + + #[test] + fn test_finite_budget_selects_dynamic_paths() -> Result<()> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["hint".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }], + Some(2), + )?; + let values = [ + json!({ + "hint": "first", + "conflict": 1, + "popular": {"nested": 1}, + "tie_a": "a", + "tie_b": true, + "rare": 1 + }), + json!({ + "hint": "second", + "conflict": "string", + "popular": {"nested": 2}, + "tie_a": "b", + "tie_b": false + }), + json!({"hint": "third", "popular": "scalar"}), + ]; + let mut builder = JsonVectorBuilder::with_settings(&settings, values.len()); + for value in values.clone() { + let value = settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + + let array = builder.to_vector().to_arrow_array(); + let structs = array.as_struct(); + assert_eq!( + vec![JSON2_REMAINDER_FIELD_NAME, "hint", "tie_a", "tie_b"], + structs + .fields() + .iter() + .map(|x| x.name().as_str()) + .collect::>() + ); + assert!(structs.column_by_name("popular").is_none()); + assert_eq!( + vec![ + Some(json!({"conflict": 1, "popular": {"nested": 1}, "rare": 1})), + Some(json!({"conflict": "string", "popular": {"nested": 2}})), + Some(json!({"popular": "scalar"})), + ], + variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())? + ); + + let field = Field::new("data", array.data_type().clone(), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(settings))), + ); + let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?; + let reconstructed = reconstructed.as_binary::(); + assert_eq!( + values[0], + decode_json_variant(reconstructed.value(0)).unwrap() + ); + assert_eq!( + values[1], + decode_json_variant(reconstructed.value(1)).unwrap() + ); + assert_eq!( + json!({ + "hint": "third", + "popular": "scalar", + "tie_a": null, + "tie_b": null + }), + decode_json_variant(reconstructed.value(2)).unwrap() + ); + + Ok(()) + } + + #[test] + fn test_reconstruct_nested_remainder_only_value() -> Result<()> { + let settings = JsonSettings::try_new(vec![], Some(1))?; + let values = [ + json!({"a": {"hot": 1}}), + json!({"a": {"hot": 2}}), + json!({"a": {"cold": 3}}), + ]; + let mut builder = JsonVectorBuilder::with_settings(&settings, values.len()); + for value in values.clone() { + let value = settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + + let array = builder.to_vector().to_arrow_array(); + let field = Field::new("data", array.data_type().clone(), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(settings))), + ); + let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?; + let reconstructed = reconstructed.as_binary::(); + assert_eq!( + vec![ + values[0].clone(), + values[1].clone(), + json!({"a": {"cold": 3, "hot": null}}), + ], + (0..reconstructed.len()) + .map(|i| decode_json_variant(reconstructed.value(i)).unwrap()) + .collect::>() + ); + + Ok(()) + } + + #[test] + fn test_dynamic_paths_require_the_same_leaf_type() -> Result<()> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["nested".to_string(), "hinted".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }], + Some(8), + )?; + let values = [ + json!({ + "branch": {}, + "different": [1], + "empty": {}, + "nested": {}, + "reverse": "scalar", + "same": [1] + }), + json!({ + "branch": {"leaf": 1}, + "different": ["x"], + "empty": {}, + "nested": {"hinted": "x", "leaf": 1}, + "same": [2] + }), + json!({ + "branch": {}, + "different": [2], + "empty": {}, + "nested": {}, + "reverse": {"leaf": 1}, + "same": [3] + }), + ]; + let mut builder = JsonVectorBuilder::with_settings(&settings, values.len()); + for value in values { + let value = settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + + let JsonNativeType::Object(fields) = builder.state.native_type() else { + unreachable!(); + }; + assert!(!fields.contains_key("different")); + assert!(!fields.contains_key("empty")); + assert!(!fields.contains_key("reverse")); + assert!(matches!(fields.get("same"), Some(JsonNativeType::Array(_)))); + assert!(matches!( + fields.get("branch"), + Some(JsonNativeType::Object(fields)) if fields.contains_key("leaf") + )); + assert!(matches!( + fields.get("nested"), + Some(JsonNativeType::Object(fields)) + if fields.contains_key("hinted") && fields.contains_key("leaf") + )); Ok(()) } #[test] fn test_json_variant_into_struct_value() -> Result<()> { + let struct_type = StructType::new(Arc::new(vec![StructField::new( + "value".to_string(), + ConcreteDataType::string_datatype(), + true, + )])); assert_eq!( json_variant_into_value( JsonVariant::Object(Default::default()), - &ConcreteDataType::string_datatype(), + &ConcreteDataType::struct_datatype(struct_type.clone()), + false, )?, Value::Null ); + assert_eq!( + json_variant_into_value( + JsonVariant::Object(Default::default()), + &ConcreteDataType::struct_datatype(struct_type.clone()), + true, + )?, + Value::Struct(StructValue::new(vec![Value::Null], struct_type)) + ); let item_type = ConcreteDataType::struct_datatype(StructType::new(Arc::new(vec![StructField::new( @@ -394,22 +1096,23 @@ mod tests { true, ), ])); - let variant = JsonVariant::from([ + let variant = JsonObjectValue::from([ ( - "items", + "items".to_string(), JsonVariant::Array(vec![ JsonVariant::from([("id", JsonVariant::from(1i64))]), JsonVariant::from([("id", JsonVariant::from(2i64))]), ]), ), ( - "meta", + "meta".to_string(), JsonVariant::from([("name", JsonVariant::from("foo"))]), ), ]); let value = Value::Struct(json_variant_into_struct_value( variant, struct_type.clone(), + true, )?); assert_eq!( diff --git a/src/datatypes/src/vectors/json/variant.rs b/src/datatypes/src/vectors/json/variant.rs index 1a4b8a861d5..75d87c99861 100644 --- a/src/datatypes/src/vectors/json/variant.rs +++ b/src/datatypes/src/vectors/json/variant.rs @@ -15,10 +15,7 @@ use std::sync::Arc; use arrow_array::ArrayRef; -#[cfg(test)] -use arrow_schema::ArrowError; -use arrow_schema::{DataType, Field}; -#[cfg(test)] +use arrow_schema::{ArrowError, DataType, Field}; use parquet_variant::{ObjectFieldBuilder, Variant, VariantBuilderExt, VariantDecimal16}; #[cfg(test)] use parquet_variant_compute::VariantArrayBuilder; @@ -27,7 +24,6 @@ use parquet_variant_json::VariantToJson; use snafu::ResultExt; use crate::error::{ArrowComputeSnafu, Result}; -#[cfg(test)] use crate::json::value::{JsonNumber, JsonVariant, decode_json_variant}; /// Returns the canonical Arrow field for an unshredded Parquet Variant array. @@ -68,8 +64,7 @@ fn json_variants_to_variant(values: &[Option]) -> Result Ok(ArrayRef::from(builder.build())) } -#[cfg(test)] -fn append_json_variant( +pub(super) fn append_json_variant( builder: &mut impl VariantBuilderExt, value: &JsonVariant, ) -> std::result::Result<(), ArrowError> { @@ -115,7 +110,6 @@ fn append_json_variant( Ok(()) } -#[cfg(test)] fn append_json_value( builder: &mut impl VariantBuilderExt, value: &serde_json::Value, @@ -157,7 +151,6 @@ fn append_json_value( /// Parquet Variant has no unsigned integer primitive. Treat u64 as i64 first, then use Decimal16 /// to represent large (larger than i64::MAX) u64. -#[cfg(test)] fn append_large_u64( builder: &mut impl VariantBuilderExt, value: u64, From 8d887ddd00d67699d0fb873cda11eb50b5771555 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 21 Aug 2026 03:50:30 +0000 Subject: [PATCH 090/126] fix(query): restore columnar group-by for dictionary-encoded tags (#8902) * fix(query): restore columnar group-by for dictionary-encoded tags Bump the DataFusion fork to be93ffd85 (feat/dict-group-column-53), which backports apache/datafusion #23187: DictionaryGroupValuesColumn lets dictionary-encoded group keys use the columnar GroupValuesColumn fast path (hash distinct dictionary values once per batch, resolve rows by key index) instead of falling back to row-based GroupValuesRows. This fixes the TSBS double-groupby regression introduced by #8541 (preserve dictionary-encoded query labels): v1.2.0-beta.1 scan output changed tag columns to Dictionary(UInt32, Utf8), which DataFusion 53.1.0 did not support in GroupValuesColumn's supported_type allow-list, so GROUP BY queries silently dropped to the ~60% slower row path (time_calculating_group_ids +57%, peak_mem +50%, end-to-end +38%). Adds an end-to-end integration test (dict_groupby_sst) that flushes a flat-format SST with dictionary-encoded hostname, runs the tsbs-style double-groupby query, and asserts correct results with no CastExec inserted before the aggregate. Signed-off-by: discord9 * chore(deps): pin merged dictionary group-by support Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Co-authored-by: discord9 --- Cargo.lock | 68 +++++----- Cargo.toml | 30 ++--- tests-integration/tests/dict_groupby_sst.rs | 139 ++++++++++++++++++++ tests-integration/tests/main.rs | 1 + 4 files changed, 189 insertions(+), 49 deletions(-) create mode 100644 tests-integration/tests/dict_groupby_sst.rs diff --git a/Cargo.lock b/Cargo.lock index d7c9241afca..a389391ea90 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3738,7 +3738,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "arrow-schema 58.3.0", @@ -3792,7 +3792,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3816,7 +3816,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3838,7 +3838,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -3862,7 +3862,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "futures", "log", @@ -3872,7 +3872,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-compression", @@ -3906,7 +3906,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "arrow-ipc 58.3.0", @@ -3929,7 +3929,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3951,7 +3951,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3974,7 +3974,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4003,12 +4003,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" [[package]] name = "datafusion-execution" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4030,7 +4030,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4052,7 +4052,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4064,7 +4064,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4095,7 +4095,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4116,7 +4116,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4128,7 +4128,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "arrow-ord 58.3.0", @@ -4152,7 +4152,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4167,7 +4167,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4184,7 +4184,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4193,7 +4193,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "datafusion-doc", "quote", @@ -4203,7 +4203,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "chrono", @@ -4252,7 +4252,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4275,7 +4275,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4289,7 +4289,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4305,7 +4305,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4323,7 +4323,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4354,7 +4354,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "chrono", @@ -4381,7 +4381,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4391,7 +4391,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4407,7 +4407,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "async-trait", "datafusion-common", @@ -4420,7 +4420,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "bigdecimal 0.4.8", @@ -4438,7 +4438,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 777d8b302b4..e6b849a6cd0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -351,21 +351,21 @@ git = "https://github.com/GreptimeTeam/greptime-meter.git" rev = "5618e779cf2bb4755b499c630fba4c35e91898cb" [patch.crates-io] -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "2aefa08a8d69c96eec2d6d6703598a009bba6e4c" } # on branch v0.61.x [profile.release] diff --git a/tests-integration/tests/dict_groupby_sst.rs b/tests-integration/tests/dict_groupby_sst.rs new file mode 100644 index 00000000000..862b77ef108 --- /dev/null +++ b/tests-integration/tests/dict_groupby_sst.rs @@ -0,0 +1,139 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use common_query::Output; +use frontend::instance::Instance; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::QueryContext; +use tests_integration::standalone::GreptimeDbStandaloneBuilder; + +async fn execute_sql(instance: &Instance, sql: &str) -> Output { + SqlQueryHandler::do_query(instance, sql, QueryContext::arc()) + .await + .remove(0) + .unwrap() +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_dict_groupby_sst() { + common_telemetry::init_default_ut_logging(); + + let standalone = GreptimeDbStandaloneBuilder::new("dict_groupby_sst") + .build() + .await; + let frontend = standalone.fe_instance(); + + execute_sql( + frontend, + r#" + CREATE TABLE cpu ( + hostname STRING PRIMARY KEY, + usage_user DOUBLE, + usage_system DOUBLE, + usage_idle DOUBLE, + greptime_timestamp TIMESTAMP TIME INDEX + ) WITH (append_mode = 'true', sst_format = 'flat') + "#, + ) + .await; + + let mut rows = String::new(); + for index in 0..1200 { + if index != 0 { + rows.push(','); + } + let hostname = match index % 3 { + 0 => "host_a", + 1 => "host_b", + _ => "host_c", + }; + let usage_user = match index % 3 { + 0 => 10.0, + 1 => 20.0, + _ => 30.0, + }; + let hour = 11 + index / 600; + let minute = (index % 600) / 60; + let second = index % 60; + rows.push_str(&format!( + "('{hostname}', {usage_user}, 1.0, 0.0, '2023-06-12T{hour:02}:{minute:02}:{second:02}Z')" + )); + } + execute_sql(frontend, &format!("INSERT INTO cpu VALUES {rows}")).await; + + execute_sql(frontend, "ADMIN FLUSH_TABLE('cpu')").await; + + let count = execute_sql(frontend, "SELECT count(*) FROM cpu") + .await + .data + .pretty_print() + .await; + assert!( + count.contains("| 1200 |"), + "unexpected row count:\n{count}" + ); + + let query = r#" + SELECT hostname, avg(usage_user), date_trunc('hour', greptime_timestamp) + FROM cpu + WHERE greptime_timestamp >= '2023-06-12T00:00:00Z' + AND greptime_timestamp < '2023-06-14T00:00:00Z' + GROUP BY date_trunc('hour', greptime_timestamp), hostname + ORDER BY hostname, date_trunc('hour', greptime_timestamp) + "#; + let groups = execute_sql(frontend, query).await.data.pretty_print().await; + assert_eq!( + 6, + groups + .lines() + .filter(|line| line.starts_with("| host_")) + .count(), + "expected one group per hostname and hour:\n{groups}" + ); + for (hostname, average) in [("host_a", "10.0"), ("host_b", "20.0"), ("host_c", "30.0")] { + assert_eq!( + 2, + groups.matches(&format!("| {hostname} ")).count(), + "expected two hourly groups for {hostname}:\n{groups}" + ); + assert!( + groups.contains(average), + "expected average {average} for {hostname}:\n{groups}" + ); + } + + let explain = execute_sql( + frontend, + r#" + EXPLAIN ANALYZE VERBOSE + SELECT avg(usage_user), date_trunc('hour', greptime_timestamp) + FROM cpu + WHERE greptime_timestamp >= '2023-06-12T00:00:00Z' + AND greptime_timestamp < '2023-06-14T00:00:00Z' + GROUP BY date_trunc('hour', greptime_timestamp), hostname + "#, + ) + .await + .data + .pretty_print() + .await; + assert!( + explain.contains("AggregateExec"), + "expected aggregate in physical plan:\n{explain}" + ); + assert!( + !explain.contains("CastExec") && !explain.contains("CAST(hostname"), + "dictionary hostname must not be cast before aggregation:\n{explain}" + ); +} diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 3112cb72c32..965a3a67462 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -16,6 +16,7 @@ mod admin_function_event; mod database_ddl_event; +mod dict_groupby_sst; mod event_recorder_test_util; #[macro_use] mod grpc; From 7fd0a7bb9877b218943429296ad442efcf309f7c Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:14:47 +0000 Subject: [PATCH 091/126] refactor(json2): optimize JSON2 building without auto-expanded paths (#8928) * refactor(json2): optimize JSON2 building without auto-expanded paths Signed-off-by: luofucong * resolve PR comments Signed-off-by: luofucong * avoid panicking memtable write Signed-off-by: luofucong --------- Signed-off-by: luofucong --- src/datatypes/src/error.rs | 8 + src/datatypes/src/json.rs | 25 +++ src/datatypes/src/vectors/json/builder.rs | 254 ++++++++++++++++++++-- src/datatypes/src/vectors/json/variant.rs | 48 +++- 4 files changed, 321 insertions(+), 14 deletions(-) diff --git a/src/datatypes/src/error.rs b/src/datatypes/src/error.rs index b34a5ffa9bf..1def0d6be8d 100644 --- a/src/datatypes/src/error.rs +++ b/src/datatypes/src/error.rs @@ -77,6 +77,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Unimplemented: {feat}"))] + Unimplemented { + feat: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to parse version in schema meta, value: {}", value))] ParseSchemaVersion { value: String, @@ -312,6 +319,7 @@ impl ErrorExt for Error { use Error::*; match self { UnsupportedOperation { .. } + | Unimplemented { .. } | UnsupportedArrowType { .. } | UnsupportedJsonType { .. } | UnsupportedDefaultExpr { .. } => StatusCode::Unsupported, diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index ec9ca4ac789..79fc10b12e7 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -140,6 +140,16 @@ impl JsonSettings { /// Encode a serde_json::Value into a Value::Json using current settings. pub fn encode(&self, json: Json) -> Result { + if let Json::Object(object) = &json + && object.contains_key(JSON2_REMAINDER_FIELD_NAME) + { + return error::InvalidJsonSnafu { + value: format!( + "root object cannot contain reserved field '{JSON2_REMAINDER_FIELD_NAME}'" + ), + } + .fail(); + } let mut context = JsonContext { path: Vec::new(), settings: self, @@ -777,6 +787,21 @@ mod tests { } } + #[test] + fn test_encode_rejects_reserved_remainder_field() -> Result<()> { + let settings = JsonSettings::default(); + let err = settings + .encode(json!({"!__remainder__!": "user-value"})) + .unwrap_err(); + assert!( + err.to_string() + .contains("root object cannot contain reserved field") + ); + + settings.encode(json!({"nested": {"!__remainder__!": "user-value"}}))?; + Ok(()) + } + #[test] fn test_encode_json_object() { let json = json!({ diff --git a/src/datatypes/src/vectors/json/builder.rs b/src/datatypes/src/vectors/json/builder.rs index 125dbdad58a..ab3fd6d0436 100644 --- a/src/datatypes/src/vectors/json/builder.rs +++ b/src/datatypes/src/vectors/json/builder.rs @@ -24,17 +24,18 @@ use snafu::{ResultExt, ensure}; use crate::data_type::ConcreteDataType; use crate::error::{ - ArrowComputeSnafu, Result, TryFromValueSnafu, UnexpectedSnafu, UnsupportedOperationSnafu, + ArrowComputeSnafu, Result, TryFromValueSnafu, UnexpectedSnafu, UnimplementedSnafu, + UnsupportedOperationSnafu, }; use crate::extension::json::JSON2_REMAINDER_FIELD_NAME; -use crate::json::value::{JsonNumber, JsonVariant, encode_json_variant}; +use crate::json::value::{JsonNumber, JsonVariant, JsonVariantRef, encode_json_variant}; use crate::json::{JSON2_MAX_STRUCTURED_DEPTH, JsonSettings}; use crate::prelude::{ValueRef, Vector, VectorRef}; use crate::types::StructType; use crate::types::json_type::{JsonNativeType, is_include}; use crate::value::{ListValue, StructValue, StructValueRef, Value}; -use crate::vectors::json::variant::{append_json_variant, variant_field}; -use crate::vectors::{Helper, MutableVector, StructVectorBuilder}; +use crate::vectors::json::variant::{append_json_variant, append_json_variant_ref, variant_field}; +use crate::vectors::{Helper, MutableVector, NullVector, StructVectorBuilder}; type JsonObjectValue = BTreeMap; @@ -44,17 +45,25 @@ type JsonObjectValue = BTreeMap; /// Auto-expanding mode always materializes type-hinted paths, selects up to /// `max_auto_expanded_paths` compatible unhinted leaf paths by frequency, and stores /// conflicting or unselected paths in the Variant remainder field. -#[derive(Clone)] pub(crate) struct JsonVectorBuilder { state: JsonVectorBuilderState, } -#[derive(Clone)] enum JsonVectorBuilderState { Legacy { merged_type: JsonNativeType, values: Vec, }, + ExplicitOnly { + /// Paths declared by type hints and stored as dedicated Struct fields. + explicit_type: JsonNativeType, + /// Concrete Struct type used to append explicit values without buffering rows. + struct_type: StructType, + /// Builder for values selected by the explicit type hints. + explicit: StructVectorBuilder, + /// Builder for all values outside the explicit type hints. + remainder: VariantArrayBuilder, + }, AutoExpanding { /// Paths declared by type hints and always stored as dedicated Struct fields. explicit_type: JsonNativeType, @@ -69,6 +78,7 @@ impl JsonVectorBuilderState { fn native_type(&self) -> JsonNativeType { match self { Self::Legacy { merged_type, .. } => merged_type.clone(), + Self::ExplicitOnly { explicit_type, .. } => explicit_type.clone(), Self::AutoExpanding { explicit_type, max_auto_expanded_paths, @@ -80,6 +90,7 @@ impl JsonVectorBuilderState { fn len(&self) -> usize { match self { Self::Legacy { values, .. } | Self::AutoExpanding { values, .. } => values.len(), + Self::ExplicitOnly { explicit, .. } => explicit.len(), } } @@ -89,6 +100,14 @@ impl JsonVectorBuilderState { merged_type, values, } => build_legacy(values, merged_type), + Self::ExplicitOnly { + explicit, + remainder, + .. + } => { + let remainder = std::mem::replace(remainder, VariantArrayBuilder::new(0)).build(); + finish_vector(explicit.to_vector(), ArrayRef::from(remainder)) + } Self::AutoExpanding { explicit_type, max_auto_expanded_paths, @@ -101,6 +120,38 @@ impl JsonVectorBuilderState { } } + fn try_build_cloned(&self) -> Result { + let mut state = match self { + Self::Legacy { + merged_type, + values, + } => Self::Legacy { + merged_type: merged_type.clone(), + values: values.clone(), + }, + Self::AutoExpanding { + explicit_type, + max_auto_expanded_paths, + values, + } => Self::AutoExpanding { + explicit_type: explicit_type.clone(), + max_auto_expanded_paths: *max_auto_expanded_paths, + values: values.clone(), + }, + // Only TimeSeriesMemtable requires a non-consuming snapshot, while JSON2 targets + // BulkMemtable. We've tried our best to support it above, but if this match arm does + // not, it's OK. The only reason it doesn't is because of `VariantArrayBuilder`. We'll + // track the upstream and see. + Self::ExplicitOnly { .. } => { + return UnimplementedSnafu { + feat: "no auto expanded JSON2 array builder", + } + .fail(); + } + }; + state.try_build() + } + fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> { let ValueRef::Json(value) = value else { return TryFromValueSnafu { @@ -125,6 +176,22 @@ impl JsonVectorBuilderState { } values.push(JsonVariant::from(value.variant())); } + Self::ExplicitOnly { + explicit_type, + struct_type, + explicit, + remainder, + } => { + if value.is_null() { + explicit.push_null(); + remainder.append_null(); + } else { + let (value, rest) = + split_to_explicit_ref(value.variant(), explicit_type, struct_type)?; + explicit.push_struct_value_ref(value)?; + append_json_variant_ref(remainder, &rest).context(ArrowComputeSnafu)?; + } + } Self::AutoExpanding { values, .. } => { values.push(JsonVariant::from(value.variant())); } @@ -137,6 +204,14 @@ impl JsonVectorBuilderState { Self::Legacy { values, .. } | Self::AutoExpanding { values, .. } => { values.push(JsonVariant::Null) } + Self::ExplicitOnly { + explicit, + remainder, + .. + } => { + explicit.push_null(); + remainder.append_null(); + } } } } @@ -163,13 +238,28 @@ impl JsonVectorBuilder { for hint in settings.type_hints() { insert_dynamic_type(&mut explicit_type, &hint.path, (&hint.data_type).into()); } - Self { - state: JsonVectorBuilderState::AutoExpanding { + let state = if settings.max_auto_expanded_paths() == Some(0) { + let DataType::Struct(fields) = explicit_type.as_arrow_type() else { + unreachable!("JSON2 explicit type must map to Arrow Struct") + }; + let struct_type = StructType::from(&fields); + JsonVectorBuilderState::ExplicitOnly { + explicit_type, + explicit: StructVectorBuilder::with_type_and_capacity( + struct_type.clone(), + capacity, + ), + struct_type, + remainder: VariantArrayBuilder::new(capacity), + } + } else { + JsonVectorBuilderState::AutoExpanding { explicit_type, max_auto_expanded_paths: settings.max_auto_expanded_paths().unwrap_or(u32::MAX), values: Vec::with_capacity(capacity), - }, - } + } + }; + Self { state } } fn try_build(&mut self) -> Result { @@ -451,6 +541,132 @@ fn insert_dynamic_type>( ) } +fn split_to_explicit_ref<'a>( + value: &JsonVariantRef<'a>, + explicit_type: &JsonNativeType, + struct_type: &StructType, +) -> Result<(StructValueRef<'a>, JsonVariantRef<'a>)> { + let JsonVariantRef::Object(object) = value else { + return TryFromValueSnafu { + reason: "expected json object value".to_string(), + } + .fail(); + }; + let explicit = json_object_ref_into_struct_value_ref(object, struct_type)?; + let remainder = remainder_ref(object, explicit_type)?; + Ok((explicit, JsonVariantRef::Object(remainder))) +} + +fn json_object_ref_into_struct_value_ref<'a>( + object: &BTreeMap<&'a str, JsonVariantRef<'a>>, + struct_type: &StructType, +) -> Result> { + let mut values = Vec::with_capacity(struct_type.fields().len()); + for field in struct_type.fields().iter() { + let value = match object.get(field.name()) { + Some(value) => json_variant_ref_into_value_ref(value, field.data_type())?, + None => ValueRef::Null, + }; + values.push(value); + } + Ok(StructValueRef::RefList { + val: values, + fields: struct_type.clone(), + }) +} + +fn json_variant_ref_into_value_ref<'a>( + value: &JsonVariantRef<'a>, + expected_type: &ConcreteDataType, +) -> Result> { + let value = match (value, expected_type) { + (JsonVariantRef::Null, _) | (_, ConcreteDataType::Null(_)) => ValueRef::Null, + (JsonVariantRef::Object(object), ConcreteDataType::Struct(struct_type)) => { + ValueRef::Struct(json_object_ref_into_struct_value_ref(object, struct_type)?) + } + (JsonVariantRef::Bool(x), ConcreteDataType::Boolean(_)) => ValueRef::Boolean(*x), + (JsonVariantRef::Number(x), ConcreteDataType::UInt64(_)) => { + let Some(x) = x.as_u64() else { + return TryFromValueSnafu { + reason: format!("unable to convert {x:?} to UInt64"), + } + .fail(); + }; + ValueRef::UInt64(x) + } + (JsonVariantRef::Number(x), ConcreteDataType::Int64(_)) => { + let x = match x { + JsonNumber::PosInt(x) => i64::try_from(*x).ok(), + JsonNumber::NegInt(x) => Some(*x), + JsonNumber::Float(_) => None, + }; + let Some(x) = x else { + return TryFromValueSnafu { + reason: format!("unable to convert {x:?} to Int64"), + } + .fail(); + }; + ValueRef::Int64(x) + } + (JsonVariantRef::Number(JsonNumber::PosInt(x)), ConcreteDataType::Float64(_)) => { + ValueRef::Float64((*x as f64).into()) + } + (JsonVariantRef::Number(JsonNumber::NegInt(x)), ConcreteDataType::Float64(_)) => { + ValueRef::Float64((*x as f64).into()) + } + (JsonVariantRef::Number(JsonNumber::Float(x)), ConcreteDataType::Float64(_)) => { + ValueRef::Float64(*x) + } + (JsonVariantRef::String(x), ConcreteDataType::String(_)) => ValueRef::String(x), + (value, expected_type) => { + return TryFromValueSnafu { + reason: format!("unable to convert json value {value:?} to {expected_type}"), + } + .fail(); + } + }; + Ok(value) +} + +fn remainder_ref<'a>( + object: &BTreeMap<&'a str, JsonVariantRef<'a>>, + explicit_type: &JsonNativeType, +) -> Result>> { + let JsonNativeType::Object(fields) = explicit_type else { + return UnexpectedSnafu { + reason: "JSON2 explicit type must be an object", + } + .fail(); + }; + let mut remainder = BTreeMap::new(); + for (&name, value) in object { + match fields.get(name) { + Some(data_type @ JsonNativeType::Object(_)) => match value { + JsonVariantRef::Null => {} + JsonVariantRef::Object(object) => { + let child = remainder_ref(object, data_type)?; + if !child.is_empty() { + remainder.insert(name, JsonVariantRef::Object(child)); + } + } + _ => { + return TryFromValueSnafu { + reason: "expected json object value".to_string(), + } + .fail(); + } + }, + // A non-object entry in the explicit type tree is an explicit leaf and is + // already written to the Struct builder. So here does nothing. + Some(_) => {} + None => { + remainder.insert(name, value.clone()); + } + } + } + Ok(remainder) +} + fn split_to_explicit( value: JsonVariant, explicit_type: &JsonNativeType, @@ -611,11 +827,19 @@ impl MutableVector for JsonVectorBuilder { } fn to_vector(&mut self) -> VectorRef { - self.try_build().unwrap_or_else(|e| panic!("{:?}", e)) + self.try_build().unwrap_or_else(|e| { + // Just try to avoid panicking here. + common_telemetry::error!(e; "Unable to build JSON2 vector"); + Arc::new(NullVector::new(self.len())) + }) } fn to_vector_cloned(&self) -> VectorRef { - self.clone().to_vector() + self.state.try_build_cloned().unwrap_or_else(|e| { + // Just try to avoid panicking here. + common_telemetry::error!(e; "Unable to build JSON2 vector"); + Arc::new(NullVector::new(self.len())) + }) } fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> { @@ -772,7 +996,7 @@ mod tests { } #[test] - fn test_zero_budget_builder_uses_fixed_schema_and_remainder() -> Result<()> { + fn test_zero_budget_builder_uses_explicit_only_schema_and_remainder() -> Result<()> { let settings = JsonSettings::try_new( vec![ JsonTypeHint { @@ -800,6 +1024,10 @@ mod tests { Some(0), )?; let mut builder = JsonVectorBuilder::with_settings(&settings, 2); + assert!(matches!( + &builder.state, + JsonVectorBuilderState::ExplicitOnly { .. } + )); let values = [ json!({ "kind": "record", diff --git a/src/datatypes/src/vectors/json/variant.rs b/src/datatypes/src/vectors/json/variant.rs index 75d87c99861..9f90f0aaacc 100644 --- a/src/datatypes/src/vectors/json/variant.rs +++ b/src/datatypes/src/vectors/json/variant.rs @@ -24,7 +24,7 @@ use parquet_variant_json::VariantToJson; use snafu::ResultExt; use crate::error::{ArrowComputeSnafu, Result}; -use crate::json::value::{JsonNumber, JsonVariant, decode_json_variant}; +use crate::json::value::{JsonNumber, JsonVariant, JsonVariantRef, decode_json_variant}; /// Returns the canonical Arrow field for an unshredded Parquet Variant array. pub fn variant_field(name: impl Into, nullable: bool) -> Field { @@ -110,6 +110,52 @@ pub(super) fn append_json_variant( Ok(()) } +pub(super) fn append_json_variant_ref( + builder: &mut impl VariantBuilderExt, + value: &JsonVariantRef<'_>, +) -> std::result::Result<(), ArrowError> { + match value { + JsonVariantRef::Null => builder.append_value(Variant::Null), + JsonVariantRef::Bool(value) => builder.append_value(*value), + JsonVariantRef::Number(JsonNumber::PosInt(value)) => { + if let Ok(value) = i64::try_from(*value) { + builder.append_value(value); + } else { + append_large_u64(builder, *value)?; + } + } + JsonVariantRef::Number(JsonNumber::NegInt(value)) => builder.append_value(*value), + JsonVariantRef::Number(JsonNumber::Float(value)) => { + if value.0.is_finite() { + builder.append_value(value.0) + } else { + builder.append_value("NaN") + } + } + JsonVariantRef::String(value) => builder.append_value(*value), + JsonVariantRef::Array(values) => { + let mut list = builder.try_new_list()?; + for value in values { + append_json_variant_ref(&mut list, value)?; + } + list.finish(); + } + JsonVariantRef::Object(values) => { + let mut object = builder.try_new_object()?; + for (name, value) in values { + append_json_variant_ref(&mut ObjectFieldBuilder::new(name, &mut object), value)?; + } + object.finish(); + } + JsonVariantRef::Variant(value) => { + let value = decode_json_variant(value) + .map_err(|e| ArrowError::JsonError(format!("Failed to decode JSONB: {e}")))?; + append_json_value(builder, &value)?; + } + } + Ok(()) +} + fn append_json_value( builder: &mut impl VariantBuilderExt, value: &serde_json::Value, From a502dfdefd59cf0ff8f3b2fdf3d1a8c79236bf95 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 24 Aug 2026 06:09:44 +0000 Subject: [PATCH 092/126] fix(mito2): fence checkpoints during region transitions (#8847) * fix: fence checkpoints during region transitions Signed-off-by: WenyXu * test(datanode): fix transient downgrade setup Signed-off-by: WenyXu * test(mito2): fix checkpoint lifecycle test setup Signed-off-by: WenyXu * test(mito2): cover cancelled downgrade waiter retry Signed-off-by: WenyXu * fix(mito2): fence direct follower transitions Signed-off-by: WenyXu * test: trim checkpoint transition coverage Signed-off-by: WenyXu * refactor(mito2): clarify checkpoint task lifecycle Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .../src/heartbeat/handler/open_region.rs | 61 +++- .../src/heartbeat/handler/sync_region.rs | 50 +++ src/datanode/src/region_server.rs | 13 +- .../region_migration/open_candidate_region.rs | 8 +- .../repartition/group/sync_region.rs | 49 ++- src/meta-srv/src/procedure/test_util.rs | 19 +- src/mito2/src/engine/set_role_state_test.rs | 328 +++++++++++++++++- src/mito2/src/engine/staging_test.rs | 99 +++++- src/mito2/src/error.rs | 22 +- src/mito2/src/manifest/checkpointer.rs | 84 +++-- src/mito2/src/manifest/manager.rs | 34 +- src/mito2/src/manifest/storage/delta.rs | 18 +- src/mito2/src/manifest/tests/checkpoint.rs | 198 ++++++++++- src/mito2/src/region.rs | 25 +- src/mito2/src/test_util.rs | 133 ++++++- src/mito2/src/worker/handle_enter_staging.rs | 4 + 16 files changed, 1083 insertions(+), 62 deletions(-) diff --git a/src/datanode/src/heartbeat/handler/open_region.rs b/src/datanode/src/heartbeat/handler/open_region.rs index 89ce857ce4d..608fab0d90e 100644 --- a/src/datanode/src/heartbeat/handler/open_region.rs +++ b/src/datanode/src/heartbeat/handler/open_region.rs @@ -98,6 +98,8 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use common_error::ext::{BoxedError, RetryHint}; + use common_error::status_code::StatusCode; use common_meta::RegionIdent; use common_meta::heartbeat::handler::{HandleControl, HeartbeatResponseHandler}; use common_meta::heartbeat::mailbox::MessageMeta; @@ -105,14 +107,21 @@ mod tests { use common_meta::kv_backend::memory::MemoryKvBackend; use mito2::config::MitoConfig; use mito2::engine::MITO_ENGINE_NAME; + use mito2::error::ManifestDeltaNotFoundSnafu; use mito2::test_util::{CreateRequestBuilder, TestEnv}; + use object_store::{Error as ObjectStoreError, ErrorKind}; + use snafu::IntoError; use store_api::path_utils::table_dir; use store_api::region_request::{RegionCloseRequest, RegionRequest, RegionRequirements}; use store_api::storage::RegionId; - use crate::heartbeat::handler::RegionHeartbeatResponseHandler; + use super::OpenRegionsHandler; + use crate::error::{self, HandleRegionRequestSnafu}; use crate::heartbeat::handler::tests::HeartbeatResponseTestEnv; - use crate::tests::mock_region_server; + use crate::heartbeat::handler::{ + HandlerContext, InstructionHandler, RegionHeartbeatResponseHandler, + }; + use crate::tests::{MockRegionEngine, mock_region_server}; fn open_regions_instruction( region_ids: impl IntoIterator, @@ -198,4 +207,52 @@ mod tests { assert!(engine.is_region_exists(region_id)); assert!(engine.is_region_exists(region_id1)); } + + #[tokio::test] + async fn test_open_regions_preserves_manifest_delta_not_found_retry_hint() { + let retryable_region = RegionId::new(1024, 1); + let non_retryable_region = RegionId::new(1024, 2); + let (engine, _) = MockRegionEngine::with_mock_fn( + MITO_ENGINE_NAME, + Box::new(move |region_id, _request| { + if region_id == retryable_region { + let manifest_error = ManifestDeltaNotFoundSnafu { + version: 1_u64, + path: "manifest/00000000000000000001.json", + } + .into_error(ObjectStoreError::new( + ErrorKind::NotFound, + "mock listed manifest delta not found", + )); + return Err(HandleRegionRequestSnafu { region_id } + .into_error(BoxedError::new(manifest_error))); + } + + error::RegionNotFoundSnafu { region_id }.fail() + }), + ); + let mut region_server = mock_region_server(); + region_server.register_engine(engine); + let ctx = HandlerContext::new_for_test(region_server, Arc::new(MemoryKvBackend::new())); + let Instruction::OpenRegions(open_regions) = + open_regions_instruction([retryable_region, non_retryable_region], "test") + else { + unreachable!() + }; + + // Serial execution makes the first error selection deterministic. + let reply = OpenRegionsHandler { + open_region_parallelism: 1, + } + .handle(&ctx, open_regions) + .await + .unwrap() + .expect_open_regions_reply(); + + assert!(!reply.result); + let error = reply.error.unwrap(); + assert_eq!(StatusCode::StorageUnavailable, error.code); + assert_eq!(RetryHint::Retryable, error.retry_hint); + assert!(error.message.contains("00000000000000000001.json")); + } } diff --git a/src/datanode/src/heartbeat/handler/sync_region.rs b/src/datanode/src/heartbeat/handler/sync_region.rs index 48687541137..f8685f28470 100644 --- a/src/datanode/src/heartbeat/handler/sync_region.rs +++ b/src/datanode/src/heartbeat/handler/sync_region.rs @@ -103,11 +103,18 @@ impl SyncRegionHandler { mod tests { use std::sync::Arc; + use common_error::ext::{BoxedError, RetryHint}; + use common_error::status_code::StatusCode; use common_meta::kv_backend::memory::MemoryKvBackend; + use mito2::engine::MITO_ENGINE_NAME; + use mito2::error::ManifestDeltaNotFoundSnafu; + use object_store::{Error as ObjectStoreError, ErrorKind}; + use snafu::IntoError; use store_api::metric_engine_consts::METRIC_ENGINE_NAME; use store_api::region_engine::{RegionRole, SyncRegionFromRequest}; use store_api::storage::RegionId; + use crate::error::HandleRegionRequestSnafu; use crate::heartbeat::handler::sync_region::SyncRegionHandler; use crate::heartbeat::handler::{HandlerContext, InstructionHandler}; use crate::tests::{MockRegionEngine, mock_region_server}; @@ -201,4 +208,47 @@ mod tests { assert!(reply[0].ready); assert!(reply[0].error.is_none()); } + + #[tokio::test] + async fn test_handle_sync_region_preserves_manifest_delta_not_found_retry_hint() { + let mock_region_server = mock_region_server(); + let region_id = RegionId::new(1024, 1); + let (mock_engine, _) = MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |engine| { + engine.mock_role = Some(Some(RegionRole::Leader)); + engine.handle_sync_region_mock_fn = Some(Box::new(|region_id, _request| { + let manifest_error = ManifestDeltaNotFoundSnafu { + version: 1_u64, + path: "manifest/00000000000000000001.json", + } + .into_error(ObjectStoreError::new( + ErrorKind::NotFound, + "mock listed manifest delta not found", + )); + Err(HandleRegionRequestSnafu { region_id } + .into_error(BoxedError::new(manifest_error))) + })); + }); + mock_region_server.register_test_region(region_id, mock_engine); + + let handler_context = + HandlerContext::new_for_test(mock_region_server, Arc::new(MemoryKvBackend::new())); + let sync_region = common_meta::instruction::SyncRegion { + region_id, + request: SyncRegionFromRequest::from_manifest(Default::default()), + }; + + let reply = SyncRegionHandler + .handle(&handler_context, vec![sync_region]) + .await + .unwrap() + .expect_sync_regions_reply(); + + assert_eq!(1, reply.len()); + assert!(reply[0].exists); + assert!(!reply[0].ready); + let error = reply[0].error.as_ref().unwrap(); + assert_eq!(StatusCode::StorageUnavailable, error.code); + assert_eq!(RetryHint::Retryable, error.retry_hint); + assert!(error.message.contains("00000000000000000001.json")); + } } diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 55dc6a670a6..57facfd53ba 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -1330,11 +1330,8 @@ impl RegionServerInner { } if !errors.is_empty() { - return error::UnexpectedSnafu { - // Returns the first error. - violated: format!("Failed to open batch regions: {:?}", errors[0]), - } - .fail(); + // Preserve the first region error so callers can honor its status code and retry hint. + return Err(errors.swap_remove(0)).context(HandleBatchOpenRequestSnafu); } Ok(open_regions) @@ -1979,7 +1976,7 @@ mod tests { use std::sync::Arc; use api::v1::{Rows, SemanticType}; - use common_error::ext::ErrorExt; + use common_error::ext::{ErrorExt, RetryHint}; use common_recordbatch::RecordBatches; use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry}; use datatypes::prelude::{ConcreteDataType, VectorRef}; @@ -2514,7 +2511,9 @@ mod tests { ) .await .unwrap_err(); - assert_eq!(err.status_code(), StatusCode::Unexpected); + assert_matches!(&err, error::Error::HandleBatchOpenRequest { .. }); + assert_eq!(err.status_code(), StatusCode::RegionNotFound); + assert_eq!(err.retry_hint(), RetryHint::NonRetryable); } struct CurrentEngineTest { diff --git a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs index 0cb1131e545..bc21746a046 100644 --- a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs +++ b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs @@ -487,10 +487,14 @@ mod tests { .await; send_mock_reply(mailbox, rx, |id| { - Ok(new_open_region_reply( + Ok(new_open_region_reply_with_error( id, false, - Some("test mocked".to_string()), + Some(InstructionError { + code: StatusCode::StorageUnavailable, + message: "test mocked".to_string(), + retry_hint: RetryHint::Retryable, + }), )) }); diff --git a/src/meta-srv/src/procedure/repartition/group/sync_region.rs b/src/meta-srv/src/procedure/repartition/group/sync_region.rs index 8ce5ea3dec3..dd55db73bcc 100644 --- a/src/meta-srv/src/procedure/repartition/group/sync_region.rs +++ b/src/meta-srv/src/procedure/repartition/group/sync_region.rs @@ -350,6 +350,9 @@ impl SyncRegion { mod tests { use std::assert_matches; + use common_error::ext::RetryHint; + use common_error::status_code::StatusCode; + use common_meta::instruction::InstructionError; use common_meta::peer::Peer; use common_meta::rpc::router::{Region, RegionRoute}; use store_api::region_engine::SyncRegionFromRequest; @@ -359,7 +362,9 @@ mod tests { use crate::procedure::repartition::group::GroupPrepareResult; use crate::procedure::repartition::group::sync_region::SyncRegion; use crate::procedure::repartition::test_util::{TestingEnv, new_persistent_context}; - use crate::procedure::test_util::{new_sync_region_reply, send_mock_reply}; + use crate::procedure::test_util::{ + new_sync_region_reply, new_sync_region_reply_with_error, send_mock_reply, + }; use crate::service::mailbox::Channel; #[test] @@ -455,4 +460,46 @@ mod tests { let err = sync_region.sync_regions(&mut ctx).await.unwrap_err(); assert_matches!(err, Error::RetryLater { .. }); } + + #[tokio::test] + async fn test_sync_regions_retryable_instruction_error() { + let mut env = TestingEnv::new(); + let table_id = 1024; + let region_id = RegionId::new(table_id, 3); + let mut persistent_context = new_persistent_context(table_id, vec![], vec![]); + persistent_context.group_prepare_result = Some(test_prepare_result(table_id)); + + let (tx, rx) = tokio::sync::mpsc::channel(1); + env.mailbox_ctx + .insert_heartbeat_response_receiver(Channel::Datanode(1), tx) + .await; + send_mock_reply(env.mailbox_ctx.mailbox().clone(), rx, move |id| { + Ok(new_sync_region_reply_with_error( + id, + region_id, + false, + true, + Some(InstructionError { + code: StatusCode::StorageUnavailable, + message: "manifest delta disappeared".to_string(), + retry_hint: RetryHint::Retryable, + }), + )) + }); + + let mut ctx = env.create_context(persistent_context); + let sync_region = SyncRegion { + region_routes: vec![RegionRoute { + region: Region { + id: region_id, + ..Default::default() + }, + leader_peer: Some(Peer::empty(1)), + ..Default::default() + }], + }; + + let err = sync_region.sync_regions(&mut ctx).await.unwrap_err(); + assert_matches!(err, Error::RetryLater { .. }); + } } diff --git a/src/meta-srv/src/procedure/test_util.rs b/src/meta-srv/src/procedure/test_util.rs index ea8c3d4aec4..4208c6e52e7 100644 --- a/src/meta-srv/src/procedure/test_util.rs +++ b/src/meta-srv/src/procedure/test_util.rs @@ -288,6 +288,23 @@ pub fn new_sync_region_reply( ready: bool, exists: bool, error: Option, +) -> MailboxMessage { + new_sync_region_reply_with_error( + id, + region_id, + ready, + exists, + legacy_instruction_error(error), + ) +} + +/// Generates a [InstructionReply::SyncRegions] reply with a structured error. +pub fn new_sync_region_reply_with_error( + id: u64, + region_id: RegionId, + ready: bool, + exists: bool, + error: Option, ) -> MailboxMessage { MailboxMessage { id, @@ -301,7 +318,7 @@ pub fn new_sync_region_reply( region_id, ready, exists, - error: legacy_instruction_error(error), + error, }, ]))) .unwrap(), diff --git a/src/mito2/src/engine/set_role_state_test.rs b/src/mito2/src/engine/set_role_state_test.rs index 40e03b063a6..1d0f3c188b8 100644 --- a/src/mito2/src/engine/set_role_state_test.rs +++ b/src/mito2/src/engine/set_role_state_test.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::time::Duration; + use api::v1::Rows; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; @@ -20,12 +22,16 @@ use store_api::region_engine::{ SettableRegionRoleState, }; use store_api::region_request::{ - EnterStagingRequest, RegionPutRequest, RegionRequest, StagingPartitionDirective, + EnterStagingRequest, RegionFlushRequest, RegionPutRequest, RegionRequest, + StagingPartitionDirective, }; use store_api::storage::RegionId; use crate::config::MitoConfig; -use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema}; +use crate::region::{RegionLeaderState, RegionRoleState}; +use crate::test_util::{ + CheckpointTaskBlocker, CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema, +}; /// Helper function to assert a successful response with expected entry id fn assert_success_response(response: &SetRegionRoleStateResponse, expected_entry_id: u64) { @@ -212,6 +218,324 @@ async fn test_write_downgrading_region_with_format(flat_format: bool) { assert_eq!(err.status_code(), StatusCode::RegionNotReady) } +#[tokio::test(flavor = "multi_thread")] +async fn test_downgrading_waits_for_checkpoint_and_stops_new_checkpoints() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let mut env = TestEnv::new().await.with_mock_layer(mock_layer); + let engine = env + .create_engine(MitoConfig { + manifest_checkpoint_distance: 1, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(0, 1), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint cleanup did not start"); + + // Leave one memtable for the final flush after entering Downgrading. + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(1, 2), + }, + ) + .await; + + let region = engine.get_region(region_id).unwrap(); + let wait_started = region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .pending_checkpoint_wait_started(); + // Register before starting the transition so the notification cannot be lost. + let wait_started = wait_started.notified(); + let cloned_engine = engine.clone(); + let downgrade = tokio::spawn(async move { + cloned_engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader) + .await + }); + tokio::time::timeout(Duration::from_secs(5), wait_started) + .await + .expect("downgrade did not start waiting for the checkpoint"); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::Downgrading), + region.state() + ); + assert!( + !downgrade.is_finished(), + "downgrade returned before checkpoint cleanup finished" + ); + + blocker.release(); + assert_success_response(&downgrade.await.unwrap().unwrap(), 2); + + // Final flush publishes a normal delta, but Downgrading must not start a + // checkpoint after the barrier. + blocker.arm_next_close(); + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + assert!( + !region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .is_doing_checkpoint(), + "final flush started a checkpoint while Downgrading" + ); + assert_eq!(2, region.manifest_ctx.manifest().await.manifest_version); + assert_eq!( + 1, + region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .last_checkpoint_version() + ); + + // Leaving Downgrading removes the scheduling restriction. The next normal + // manifest update can checkpoint the accumulated deltas. + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(2, 3), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint scheduling did not resume after leaving Downgrading"); + blocker.release(); + region + .manifest_ctx + .manifest_manager + .write() + .await + .wait_for_pending_checkpoint() + .await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_direct_follower_waits_for_pending_checkpoint() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let mut env = TestEnv::new().await.with_mock_layer(mock_layer); + let engine = env + .create_engine(MitoConfig { + manifest_checkpoint_distance: 1, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 1), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint cleanup did not start"); + + let region = engine.get_region(region_id).unwrap(); + let wait_started = region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .pending_checkpoint_wait_started(); + // The no-flush downgrade path requests Follower directly. It must still + // wait for checkpoint cleanup before replying to the caller. + let wait_started = wait_started.notified(); + let cloned_engine = engine.clone(); + let set_follower = tokio::spawn(async move { + cloned_engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower) + .await + }); + tokio::time::timeout(Duration::from_secs(5), wait_started) + .await + .expect("follower transition did not start waiting for the checkpoint"); + assert_eq!(RegionRoleState::Follower, region.state()); + assert!( + !set_follower.is_finished(), + "follower transition returned before checkpoint cleanup finished" + ); + + blocker.release(); + assert_success_response(&set_follower.await.unwrap().unwrap(), 1); +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_retried_downgrade_waits_after_first_request_is_cancelled() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let mut env = TestEnv::new().await.with_mock_layer(mock_layer); + let engine = env + .create_engine(MitoConfig { + manifest_checkpoint_distance: 1, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 1), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint cleanup did not start"); + + let region = engine.get_region(region_id).unwrap(); + let first_wait_started = region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .pending_checkpoint_wait_started(); + // Register before starting the transition so the notification cannot be lost. + let first_wait_started = first_wait_started.notified(); + // Call the region transition directly so aborting this task cancels the + // actual checkpoint waiter. The engine API submits the same transition to + // a detached worker task, so aborting its caller only drops the reply receiver. + let first_region = region.clone(); + let first_downgrade = tokio::spawn(async move { + first_region + .set_role_state_gracefully(SettableRegionRoleState::DowngradingLeader) + .await + }); + tokio::time::timeout(Duration::from_secs(5), first_wait_started) + .await + .expect("first downgrade did not start waiting for the checkpoint"); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::Downgrading), + region.state() + ); + assert!( + !first_downgrade.is_finished(), + "first downgrade did not wait for checkpoint cleanup" + ); + + // Cancelling the first caller must not remove the checkpoint handle from + // the manifest manager. A retried migration request observes Downgrading + // and must wait for the same checkpoint task before it can proceed. + first_downgrade.abort(); + assert!(first_downgrade.await.unwrap_err().is_cancelled()); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::Downgrading), + region.state() + ); + + let second_wait_started = region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .pending_checkpoint_wait_started(); + // Register before spawning the retry so the notification cannot be lost. + let second_wait_started = second_wait_started.notified(); + let retry_region = region.clone(); + let retry_downgrade = tokio::spawn(async move { + retry_region + .set_role_state_gracefully(SettableRegionRoleState::DowngradingLeader) + .await + }); + tokio::time::timeout(Duration::from_secs(5), second_wait_started) + .await + .expect("retried downgrade did not start waiting for the checkpoint"); + assert!( + !retry_downgrade.is_finished(), + "retried downgrade did not wait for checkpoint cleanup" + ); + + blocker.release(); + retry_downgrade.await.unwrap().unwrap(); +} + #[tokio::test] async fn test_unified_state_transitions() { test_unified_state_transitions_with_format(false).await; diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index 71abd67eb22..12263976aa2 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -48,7 +48,9 @@ use crate::manifest::action::{ use crate::region::{RegionLeaderState, RegionRoleState, parse_partition_expr}; use crate::request::WorkerRequest; use crate::sst::FormatType; -use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema}; +use crate::test_util::{ + CheckpointTaskBlocker, CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema, +}; fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr { col(col_name) @@ -1101,6 +1103,101 @@ async fn test_write_stall_on_enter_staging_with_format(flat_format: bool) { assert_eq!(expected, batches.pretty_print().unwrap()); } +#[tokio::test(flavor = "multi_thread")] +async fn test_enter_staging_waits_for_pending_checkpoint() { + for block_last_checkpoint_write in [false, true] { + let partition_directive = + StagingPartitionDirective::UpdatePartitionExpr(default_partition_expr()); + let (blocker, mock_layer) = if block_last_checkpoint_write { + CheckpointTaskBlocker::block_last_checkpoint_write() + } else { + CheckpointTaskBlocker::block_cleanup() + }; + let mut env = TestEnv::new().await.with_mock_layer(mock_layer); + let engine = env + .create_engine(MitoConfig { + manifest_checkpoint_distance: 1, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 1), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint task did not reach the configured block point"); + + let region = engine.get_region(region_id).unwrap(); + let wait_started = region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .pending_checkpoint_wait_started(); + // Register before starting the transition so the notification cannot be lost. + let wait_started = wait_started.notified(); + let cloned_engine = engine.clone(); + let enter_staging = tokio::spawn(async move { + cloned_engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive, + }), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), wait_started) + .await + .expect("EnterStaging did not start waiting for the checkpoint"); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::EnteringStaging), + region.state() + ); + assert!( + !enter_staging.is_finished(), + "EnterStaging returned before the checkpoint task finished" + ); + + blocker.release(); + enter_staging.await.unwrap().unwrap(); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::Staging), + region.state() + ); + assert!( + !region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .is_doing_checkpoint() + ); + } +} + #[tokio::test] async fn test_enter_staging_clean_staging_manifest_error() { common_telemetry::init_default_ut_logging(); diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index a484d43a318..7708bc3a36e 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -67,6 +67,20 @@ pub enum Error { error: object_store::Error, }, + #[snafu(display( + "Manifest delta {} disappeared after it was listed, path: {}", + version, + path + ))] + ManifestDeltaNotFound { + version: ManifestVersion, + path: String, + #[snafu(source)] + error: object_store::Error, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Fail to compress object by {}, path: {}", compress_type, path))] CompressObject { compress_type: CompressionType, @@ -1359,6 +1373,7 @@ impl Error { pub(crate) fn is_object_not_found(&self) -> bool { match self { Error::OpenDal { error, .. } => error.kind() == ErrorKind::NotFound, + Error::ManifestDeltaNotFound { .. } => true, _ => false, } } @@ -1400,7 +1415,9 @@ impl ErrorExt for Error { match self { DataTypeMismatch { source, .. } => source.status_code(), - OpenDal { .. } | ReadParquet { .. } => StatusCode::StorageUnavailable, + OpenDal { .. } | ManifestDeltaNotFound { .. } | ReadParquet { .. } => { + StatusCode::StorageUnavailable + } WriteWal { source, .. } | ReadWal { source, .. } | DeleteWal { source, .. } => { source.status_code() } @@ -1601,7 +1618,8 @@ impl ErrorExt for Error { | RegionStopped { .. } | RegionBusy { .. } | ManualCompactionAlreadyRunning { .. } - | FlushableRegionState { .. } => RetryHint::Retryable, + | FlushableRegionState { .. } + | ManifestDeltaNotFound { .. } => RetryHint::Retryable, OpenDal { error, .. } | DeleteSsts { error, .. } diff --git a/src/mito2/src/manifest/checkpointer.rs b/src/mito2/src/manifest/checkpointer.rs index 9dd35e189a6..232421a2621 100644 --- a/src/mito2/src/manifest/checkpointer.rs +++ b/src/mito2/src/manifest/checkpointer.rs @@ -14,11 +14,14 @@ use std::fmt::Debug; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, Ordering}; +use common_runtime::JoinHandle; use common_telemetry::{error, info, warn}; use store_api::storage::RegionId; use store_api::{MIN_VERSION, ManifestVersion}; +#[cfg(test)] +use tokio::sync::Notify; use crate::error::Result; use crate::manifest::action::{RegionCheckpoint, RegionManifest}; @@ -31,6 +34,9 @@ use crate::metrics::MANIFEST_OP_ELAPSED; pub(crate) struct Checkpointer { manifest_options: RegionManifestOptions, inner: Arc, + checkpoint_task: Option>, + #[cfg(test)] + pending_checkpoint_wait_started: Arc, } #[derive(Debug)] @@ -38,15 +44,10 @@ struct Inner { region_id: RegionId, manifest_store: ManifestObjectStore, last_checkpoint_version: AtomicU64, - is_doing_checkpoint: AtomicBool, } impl Inner { async fn do_checkpoint(&self, checkpoint: RegionCheckpoint) { - let _guard = scopeguard::guard(&self.is_doing_checkpoint, |x| { - x.store(false, Ordering::Relaxed); - }); - let _t = MANIFEST_OP_ELAPSED .with_label_values(&["checkpoint"]) .start_timer(); @@ -90,14 +91,6 @@ impl Inner { fn region_id(&self) -> RegionId { self.region_id } - - fn is_doing_checkpoint(&self) -> bool { - self.is_doing_checkpoint.load(Ordering::Relaxed) - } - - fn set_doing_checkpoint(&self) { - self.is_doing_checkpoint.store(true, Ordering::Relaxed); - } } impl Checkpointer { @@ -113,8 +106,10 @@ impl Checkpointer { region_id, manifest_store, last_checkpoint_version: AtomicU64::new(last_checkpoint_version), - is_doing_checkpoint: AtomicBool::new(false), }), + checkpoint_task: None, + #[cfg(test)] + pending_checkpoint_wait_started: Arc::new(Notify::new()), } } @@ -140,7 +135,19 @@ impl Checkpointer { /// Check if it's needed to do checkpoint for the region by the checkpoint distance. /// If needed, and there's no currently running checkpoint task, it will start a new checkpoint /// task running in the background. - pub(crate) fn maybe_do_checkpoint(&self, manifest: &RegionManifest) { + pub(crate) async fn maybe_do_checkpoint(&mut self, manifest: &RegionManifest) { + if self + .checkpoint_task + .as_ref() + .is_some_and(|handle| !handle.is_finished()) + { + return; + } + + // Reap a completed task before checking whether to start the next one. + // This keeps the handle as the single source of truth for task state. + self.wait_for_pending_checkpoint().await; + if self.manifest_options.checkpoint_distance == 0 { return; } @@ -152,12 +159,6 @@ impl Checkpointer { return; } - // We can simply check whether there's a running checkpoint task like this, all because of - // the caller of this function is ran single threaded, inside the lock of RegionManifestManager. - if self.inner.is_doing_checkpoint() { - return; - } - let start_version = if last_checkpoint_version == 0 { // Checkpoint version can't be zero by implementation. // So last checkpoint version is zero means no last checkpoint. @@ -181,17 +182,44 @@ impl Checkpointer { self.do_checkpoint(checkpoint); } - fn do_checkpoint(&self, checkpoint: RegionCheckpoint) { - self.inner.set_doing_checkpoint(); - + fn do_checkpoint(&mut self, checkpoint: RegionCheckpoint) { let inner = self.inner.clone(); - common_runtime::spawn_global(async move { + self.checkpoint_task = Some(common_runtime::spawn_global(async move { inner.do_checkpoint(checkpoint).await; - }); + })); + } + + /// Waits for the current checkpoint task without removing its handle first. + /// + /// Keeping the handle in `self` while awaiting is important. If the caller is + /// cancelled, another lifecycle transition can still wait for the same task. + pub(crate) async fn wait_for_pending_checkpoint(&mut self) { + let Some(handle) = self.checkpoint_task.as_mut() else { + return; + }; + + #[cfg(test)] + self.pending_checkpoint_wait_started.notify_one(); + + let result = (&mut *handle).await; + // There is no cancellation point between observing completion and + // clearing the handle. + self.checkpoint_task = None; + + if let Err(e) = result { + warn!(e; "Failed to join checkpoint task for region {}", self.inner.region_id()); + } } #[cfg(test)] pub(crate) fn is_doing_checkpoint(&self) -> bool { - self.inner.is_doing_checkpoint() + self.checkpoint_task + .as_ref() + .is_some_and(|handle| !handle.is_finished()) + } + + #[cfg(test)] + pub(crate) fn pending_checkpoint_wait_started(&self) -> Arc { + self.pending_checkpoint_wait_started.clone() } } diff --git a/src/mito2/src/manifest/manager.rs b/src/mito2/src/manifest/manager.rs index fcc48395fd9..c6a0d3f3228 100644 --- a/src/mito2/src/manifest/manager.rs +++ b/src/mito2/src/manifest/manager.rs @@ -550,6 +550,28 @@ impl RegionManifestManager { &mut self, action_list: RegionMetaActionList, is_staging: bool, + ) -> Result { + self.update_inner(action_list, is_staging, !is_staging) + .await + } + + /// Updates the normal manifest without starting a checkpoint. + /// + /// This is only used while a leader is downgrading. The final flush still + /// publishes its manifest edit, but must not start cleanup after the + /// downgrade checkpoint barrier. + pub(crate) async fn update_normal_without_checkpoint( + &mut self, + action_list: RegionMetaActionList, + ) -> Result { + self.update_inner(action_list, false, false).await + } + + async fn update_inner( + &mut self, + action_list: RegionMetaActionList, + is_staging: bool, + allow_checkpoint: bool, ) -> Result { let _t = MANIFEST_OP_ELAPSED .with_label_values(&["update"]) @@ -616,13 +638,21 @@ impl RegionManifestManager { .checkpointer .update_manifest_removed_files(new_manifest)?; self.manifest = Arc::new(updated_manifest); - self.checkpointer - .maybe_do_checkpoint(self.manifest.as_ref()); + if allow_checkpoint { + self.checkpointer + .maybe_do_checkpoint(self.manifest.as_ref()) + .await; + } } Ok(version) } + /// Waits for an in-flight checkpoint to finish, including its cleanup. + pub(crate) async fn wait_for_pending_checkpoint(&mut self) { + self.checkpointer.wait_for_pending_checkpoint().await; + } + /// Clear deleted files from manifest's `removed_files` field without update version. Notice if datanode exit before checkpoint then new manifest by open region may still contain these deleted files, which is acceptable for gc process. pub fn clear_deleted_files(&mut self, deleted_files: Vec) { let mut manifest = (*self.manifest()).clone(); diff --git a/src/mito2/src/manifest/storage/delta.rs b/src/mito2/src/manifest/storage/delta.rs index 594b56ddaec..8f9962a3c69 100644 --- a/src/mito2/src/manifest/storage/delta.rs +++ b/src/mito2/src/manifest/storage/delta.rs @@ -26,7 +26,8 @@ use tokio::sync::Semaphore; use crate::cache::manifest_cache::ManifestCache; use crate::error::{ - CompressObjectSnafu, DecompressObjectSnafu, InvalidScanIndexSnafu, OpenDalSnafu, Result, + CompressObjectSnafu, DecompressObjectSnafu, InvalidScanIndexSnafu, ManifestDeltaNotFoundSnafu, + OpenDalSnafu, Result, }; use crate::manifest::storage::size_tracker::Tracker; use crate::manifest::storage::utils::{ @@ -212,11 +213,16 @@ impl DeltaStorage { // Fetch from remote object store let compress_type = file_compress_type(entry.name()); - let bytes = self - .object_store - .read(entry.path()) - .await - .context(OpenDalSnafu)?; + let bytes = match self.object_store.read(entry.path()).await { + Ok(bytes) => bytes, + Err(error) if error.kind() == ErrorKind::NotFound => { + return Err(error).context(ManifestDeltaNotFoundSnafu { + version: *v, + path: entry.path(), + }); + } + Err(error) => return Err(error).context(OpenDalSnafu), + }; let data = compress_type .decode(bytes) .await diff --git a/src/mito2/src/manifest/tests/checkpoint.rs b/src/mito2/src/manifest/tests/checkpoint.rs index 07f879701ad..86679630ac7 100644 --- a/src/mito2/src/manifest/tests/checkpoint.rs +++ b/src/mito2/src/manifest/tests/checkpoint.rs @@ -18,13 +18,16 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use common_datasource::compression::CompressionType; +use common_error::ext::{ErrorExt, RetryHint}; +use common_error::status_code::StatusCode; use object_store::layers::mock::{ - Error as MockError, ErrorKind, MockLayerBuilder, OpDelete, Result as MockResult, oio, + Buffer, Error as MockError, ErrorKind, MockLayer, MockLayerBuilder, OpDelete, + Result as MockResult, oio, }; use store_api::storage::{FileId, RegionId}; use strum::IntoEnumIterator; -use crate::error::Error::ChecksumMismatch; +use crate::error::Error::{ChecksumMismatch, ManifestDeltaNotFound}; use crate::manifest::action::{ RegionCheckpoint, RegionEdit, RegionMetaAction, RegionMetaActionList, }; @@ -33,7 +36,7 @@ use crate::manifest::storage::checkpoint::CheckpointMetadata; use crate::manifest::storage::is_delta_file; use crate::manifest::tests::utils::basic_region_metadata; use crate::sst::file::FileMeta; -use crate::test_util::TestEnv; +use crate::test_util::{CheckpointTaskBlocker, TestEnv}; async fn build_manager( checkpoint_distance: u64, @@ -85,6 +88,31 @@ fn nop_action() -> RegionMetaActionList { })]) } +struct NotFoundReader; + +impl oio::Read for NotFoundReader { + async fn read(&mut self) -> MockResult { + Err(MockError::new( + ErrorKind::NotFound, + "mock listed manifest delta not found", + )) + } +} + +fn fail_manifest_delta_reads_layer() -> MockLayer { + MockLayerBuilder::default() + .reader_factory(Arc::new(|path, _args, inner| { + let file_name = path.rsplit('/').next().unwrap_or(path); + if is_delta_file(file_name) { + Box::new(NotFoundReader) + } else { + inner + } + })) + .build() + .unwrap() +} + #[tokio::test] async fn manager_without_checkpoint() { let (_env, mut manager) = build_manager(0, CompressionType::Uncompressed).await; @@ -589,3 +617,167 @@ async fn checkpoint_advances_and_recovery_works_when_delete_fails() { .expect("manifest should be recoverable"); assert_eq!(reopened.manifest().manifest_version, 10); } + +#[tokio::test] +async fn open_preserves_listed_delta_not_found_retry_hint() { + let env = TestEnv::new() + .await + .with_mock_layer(fail_manifest_delta_reads_layer()); + let metadata = Arc::new(basic_region_metadata()); + let mut manager = env + .create_manifest_manager(CompressionType::Uncompressed, 0, Some(metadata.clone())) + .await + .unwrap() + .unwrap(); + manager.stop().await; + + let error = env + .create_manifest_manager(CompressionType::Uncompressed, 0, None) + .await + .expect_err("reopen must fail on the mocked delta read"); + assert_matches!( + &error, + ManifestDeltaNotFound { + version: 0, + path, + error, + .. + } if path.ends_with("00000000000000000000.json") + && error.kind() == object_store::ErrorKind::NotFound + ); + assert_eq!(StatusCode::StorageUnavailable, error.status_code()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); +} + +#[tokio::test] +async fn install_preserves_listed_delta_not_found_retry_hint() { + let env = TestEnv::new() + .await + .with_mock_layer(fail_manifest_delta_reads_layer()); + let metadata = Arc::new(basic_region_metadata()); + let mut manager = env + .create_manifest_manager(CompressionType::Uncompressed, 0, Some(metadata)) + .await + .unwrap() + .unwrap(); + let mut store = manager.store(); + store + .save(1, &nop_action().encode().unwrap(), false) + .await + .unwrap(); + + let error = manager.install_manifest_to(1).await.unwrap_err(); + assert_matches!( + &error, + ManifestDeltaNotFound { + version: 1, + path, + .. + } if path.ends_with("00000000000000000001.json") + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); +} + +#[tokio::test] +async fn cancelled_waiter_keeps_pending_checkpoint_handle() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let env = TestEnv::new().await.with_mock_layer(mock_layer); + let metadata = Arc::new(basic_region_metadata()); + let mut manager = env + .create_manifest_manager(CompressionType::Uncompressed, 1, Some(metadata)) + .await + .unwrap() + .unwrap(); + + manager.update(nop_action(), false).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint cleanup did not start"); + + let wait_started = manager.checkpointer().pending_checkpoint_wait_started(); + let wait_started = wait_started.notified(); + let mut first_wait = Box::pin(manager.wait_for_pending_checkpoint()); + tokio::select! { + _ = wait_started => {} + _ = &mut first_wait => panic!("checkpoint waiter returned before cleanup finished"), + } + drop(first_wait); + assert!(manager.checkpointer().is_doing_checkpoint()); + + blocker.release(); + manager.wait_for_pending_checkpoint().await; + assert!(!manager.checkpointer().is_doing_checkpoint()); + + let (version, _) = manager + .store() + .load_last_checkpoint() + .await + .unwrap() + .expect("checkpoint must be published"); + assert_eq!(1, version); +} + +#[tokio::test] +async fn running_checkpoint_prevents_scheduling_another_checkpoint() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let env = TestEnv::new().await.with_mock_layer(mock_layer); + let metadata = Arc::new(basic_region_metadata()); + let mut manager = env + .create_manifest_manager(CompressionType::Uncompressed, 1, Some(metadata)) + .await + .unwrap() + .unwrap(); + + manager.update(nop_action(), false).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint cleanup did not start"); + + // The second update is eligible for checkpointing, but the first task still + // owns the pending handle and must prevent another task from being scheduled. + manager.update(nop_action(), false).await.unwrap(); + blocker.release(); + manager.wait_for_pending_checkpoint().await; + + assert_eq!(2, manager.manifest().manifest_version); + assert_eq!(1, manager.checkpointer().last_checkpoint_version()); +} + +#[tokio::test] +async fn completed_checkpoint_is_reaped_before_scheduling_the_next_one() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let env = TestEnv::new().await.with_mock_layer(mock_layer); + let metadata = Arc::new(basic_region_metadata()); + let mut manager = env + .create_manifest_manager(CompressionType::Uncompressed, 1, Some(metadata)) + .await + .unwrap() + .unwrap(); + + manager.update(nop_action(), false).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("first checkpoint cleanup did not start"); + blocker.release(); + + tokio::time::timeout(Duration::from_secs(5), async { + while manager.checkpointer().is_doing_checkpoint() { + tokio::task::yield_now().await; + } + }) + .await + .expect("first checkpoint did not finish"); + assert_eq!(1, manager.checkpointer().last_checkpoint_version()); + + // Do not explicitly wait/reap the completed handle. The next eligible + // update must reap it before scheduling another checkpoint. + blocker.arm_next_close(); + manager.update(nop_action(), false).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("second checkpoint cleanup did not start"); + blocker.release(); + manager.wait_for_pending_checkpoint().await; + + assert_eq!(2, manager.checkpointer().last_checkpoint_version()); +} diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index c5a9021bde5..76204da9cf6 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -483,6 +483,7 @@ impl MitoRegion { let mut manager: RwLockWriteGuard<'_, RegionManifestManager> = self.manifest_ctx.manifest_manager.write().await; let current_state = self.state(); + let mut wait_for_checkpoint = false; let hook_payload: Option = match state { SettableRegionRoleState::Leader => { @@ -545,10 +546,12 @@ impl MitoRegion { ); self.exit_staging()?; self.set_role(RegionRole::Follower); + wait_for_checkpoint = true; } RegionRoleState::Leader(_) => { info!("Demoting region {} from leader to follower", self.region_id); self.set_role(RegionRole::Follower); + wait_for_checkpoint = true; } RegionRoleState::Follower => { // Already in desired state - no-op @@ -568,14 +571,17 @@ impl MitoRegion { ); self.exit_staging()?; self.set_role(RegionRole::DowngradingLeader); + wait_for_checkpoint = true; } RegionRoleState::Leader(RegionLeaderState::Writable) => { info!("Starting downgrade for region {}", self.region_id); self.set_role(RegionRole::DowngradingLeader); + wait_for_checkpoint = true; } RegionRoleState::Leader(RegionLeaderState::Downgrading) => { // Already in desired state - no-op info!("Region {} already in downgrading mode", self.region_id); + wait_for_checkpoint = true; } _ => { warn!( @@ -588,6 +594,13 @@ impl MitoRegion { } }; + // The state is changed before waiting, so no new writable-leader work + // can race with the barrier. Keep the manager lock while joining to + // serialize the barrier with checkpoint scheduling. + if wait_for_checkpoint { + manager.wait_for_pending_checkpoint().await; + } + // Hack(zhongzc): If we have just become leader (writable), persist any backfilled metadata. let mut backfill_hook_payload: Option = None; if self.state() == RegionRoleState::Leader(RegionLeaderState::Writable) { @@ -1371,10 +1384,14 @@ impl ManifestContext { // Clone before `action_list` is moved into `update` so the hook still // sees what was written. let action_list_for_hook = self.hook.as_ref().map(|_| action_list.clone()); - let version = manager - .update(action_list, is_staging) - .await - .inspect_err(|e| error!(e; "Failed to update manifest, region_id: {}", region_id))?; + let version = if !is_staging + && self.state.load() == RegionRoleState::Leader(RegionLeaderState::Downgrading) + { + manager.update_normal_without_checkpoint(action_list).await + } else { + manager.update(action_list, is_staging).await + } + .inspect_err(|e| error!(e; "Failed to update manifest, region_id: {}", region_id))?; Ok(PendingManifestHook::new( region_id, diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index 8a63d749b5b..5bf55edb07c 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -52,7 +52,9 @@ use log_store::raft_engine::log_store::RaftEngineLogStore; use log_store::test_util::log_store_util; use moka::future::CacheBuilder; use object_store::ObjectStore; -use object_store::layers::mock::MockLayer; +use object_store::layers::mock::{ + Buffer, Deleter, Metadata, MockLayer, MockLayerBuilder, OpDelete, Result as MockResult, Writer, +}; use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef}; use object_store::services::Fs; use rskafka::client::partition::{Compression, UnknownTopicHandling}; @@ -67,6 +69,7 @@ use store_api::region_request::{ RegionOpenRequest, RegionPutRequest, RegionRequest, }; use store_api::storage::{ColumnId, RegionId}; +use tokio::sync::Notify; use crate::cache::write_cache::{WriteCache, WriteCacheRef}; use crate::config::MitoConfig; @@ -75,6 +78,7 @@ use crate::engine::{MITO_ENGINE_NAME, MitoEngine}; use crate::error::Result; use crate::flush::{WriteBufferManager, WriteBufferManagerRef}; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; +use crate::manifest::storage::{is_checkpoint_file, is_delta_file}; use crate::read::{Batch, BatchBuilder, BatchReader}; use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef}; use crate::sst::FormatType; @@ -85,6 +89,133 @@ use crate::sst::index::puffin_manager::PuffinManagerFactory; use crate::time_provider::{StdTimeProvider, TimeProviderRef}; use crate::worker::WorkerGroup; +/// Controls a mock object-store layer that blocks a checkpoint task once. +#[derive(Clone)] +pub(crate) struct CheckpointTaskBlocker { + entered: Arc, + release: Arc, + armed: Arc, +} + +impl CheckpointTaskBlocker { + /// Blocks normal manifest checkpoint cleanup at batch-delete close. + pub(crate) fn block_cleanup() -> (Self, MockLayer) { + let blocker = Self { + entered: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + armed: Arc::new(AtomicBool::new(true)), + }; + let factory_blocker = blocker.clone(); + let layer = MockLayerBuilder::default() + .deleter_factory(Arc::new(move |inner| { + Box::new(BlockingCheckpointDeleter { + inner, + blocker: factory_blocker.clone(), + has_manifest_cleanup_target: false, + }) + })) + .build() + .unwrap(); + (blocker, layer) + } + + /// Blocks publication of `_last_checkpoint` at writer close. + pub(crate) fn block_last_checkpoint_write() -> (Self, MockLayer) { + let blocker = Self { + entered: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + armed: Arc::new(AtomicBool::new(true)), + }; + let factory_blocker = blocker.clone(); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _args, inner| { + Box::new(BlockingCheckpointWriter { + path: path.to_string(), + inner, + blocker: factory_blocker.clone(), + }) + })) + .build() + .unwrap(); + (blocker, layer) + } + + pub(crate) async fn wait_until_blocked(&self) { + self.entered.notified().await; + } + + pub(crate) fn release(&self) { + self.release.notify_one(); + } + + pub(crate) fn arm_next_close(&self) { + self.armed.store(true, Ordering::Release); + } + + async fn block_once(&self) { + if self.armed.swap(false, Ordering::AcqRel) { + self.entered.notify_one(); + self.release.notified().await; + } + } +} + +struct BlockingCheckpointDeleter { + inner: Deleter, + blocker: CheckpointTaskBlocker, + has_manifest_cleanup_target: bool, +} + +impl object_store::layers::mock::Delete for BlockingCheckpointDeleter { + async fn delete(&mut self, path: &str, args: OpDelete) -> MockResult<()> { + self.inner.delete(path, args).await?; + if is_manifest_checkpoint_file(path) { + self.has_manifest_cleanup_target = true; + } + Ok(()) + } + + async fn close(&mut self) -> MockResult<()> { + if self.has_manifest_cleanup_target { + self.blocker.block_once().await; + } + self.inner.close().await + } +} + +fn is_manifest_checkpoint_file(path: &str) -> bool { + // The mock deleter receives paths relative to the listed manifest + // directory, so the normal/staging directory segments are unavailable. + let path = Path::new(path); + let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else { + return false; + }; + is_delta_file(file_name) || is_checkpoint_file(file_name) +} + +struct BlockingCheckpointWriter { + path: String, + inner: Writer, + blocker: CheckpointTaskBlocker, +} + +impl object_store::layers::mock::Write for BlockingCheckpointWriter { + async fn write(&mut self, bs: Buffer) -> MockResult<()> { + self.inner.write(bs).await + } + + async fn close(&mut self) -> MockResult { + if self.path.ends_with("_last_checkpoint") { + self.blocker.block_once().await; + } + self.inner.close().await + } + + async fn abort(&mut self) -> MockResult<()> { + self.inner.abort().await + } +} + pub(crate) fn new_noop_file_purger() -> FilePurgerRef { Arc::new(NoopFilePurger) } diff --git a/src/mito2/src/worker/handle_enter_staging.rs b/src/mito2/src/worker/handle_enter_staging.rs index 000525021d3..f9e24541241 100644 --- a/src/mito2/src/worker/handle_enter_staging.rs +++ b/src/mito2/src/worker/handle_enter_staging.rs @@ -126,6 +126,10 @@ impl RegionWorkerLoop { // First step: clear all staging manifest files. { let mut manager = region.manifest_ctx.manifest_manager.write().await; + // `set_entering_staging` has already stopped new normal manifest + // publications. Wait for an older checkpoint to finish cleanup + // before exposing staging to repartition remap readers. + manager.wait_for_pending_checkpoint().await; manager .clear_staging_manifest_and_dir() .await From 174577164aaeea2c19a36d214d4eb1cbfd1e6806 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 24 Aug 2026 09:28:49 +0000 Subject: [PATCH 093/126] feat: otlp duration_nano and trace_flag signed integer coercion (#8816) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(servers): add signed→unsigned int coercion to OTLP ingest path Phase 0 of transitioning built-in data models from unsigned to signed integers (#8793): add lossless Int64→UInt64 and Int32→UInt32 coercion arms so existing UInt64/UInt32 columns (e.g. trace `duration_nano`, log `trace_flags`) keep accepting new signed ingest without an ALTER TABLE. The OTLP ingest path already reconciles every incoming column against the existing table schema and treats it as authoritative. With these arms, `choose_trace_reconcile_decision` returns `UseExisting(UInt64/Uint32)` for an existing unsigned column receiving signed data: the table keeps its type byte-for-byte and the request value is coerced. No persisted format is mutated; existing data stays readable as-is. This is the safety net that makes the actual schema flip (Phase 1) safe. Only the signed→unsigned direction is supported — the reverse would be lossy for values above the signed range and is intentionally rejected. Tests cover both new arms plus an end-to-end log test proving an existing UInt64 column coerces an incoming Int64 value while keeping its type. Signed-off-by: Ning Sun * feat: add compatibility layer for uint trace/log fields Signed-off-by: Ning Sun * fix: jaeger test Signed-off-by: Ning Sun * fix: keep trace v0 unsigned, reject negative span durations - Keep the frozen v0 data model on UInt64 duration_nano: the signed ingest compatibility layer only runs on the v1 path, so flipping v0 would break writes into every pre-existing v0 table at mito's schema check. Pin the schema with unit and integration tests. - Reject spans whose end precedes their start (or whose duration does not fit i64) on the v1 path instead of wrapping: new Int64 tables and existing UInt64 tables now fail identically, rather than storing negative durations that break the Jaeger query API. - Extract is_supported_signed_to_unsigned_coercion so the trace and log ingest paths share one supported-pair predicate and cannot drift. Signed-off-by: Ning Sun * fix: clamp negative span durations to zero, revert semantic_graph comment - Record duration 0 for spans whose end precedes their start instead of erroring: a malformed span no longer fails the request, and the value written is always a non-negative, in-range i64 so new Int64 tables and existing UInt64 tables (via the checked coercion) behave identically. Durations above i64::MAX saturate rather than wrap. - Revert the doc-comment tweak on the semantic_graph test fixture; the file is untouched by this PR again. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- .../src/instance/otlp/trace_ingest.rs | 23 +- .../src/instance/otlp/trace_ingest/tests.rs | 28 +++ src/frontend/src/instance/otlp/trace_types.rs | 54 ++++ src/servers/src/otlp/coerce.rs | 108 +++++++- src/servers/src/otlp/logs.rs | 231 +++++++++++++++++- src/servers/src/otlp/trace/v0.rs | 27 +- src/servers/src/otlp/trace/v1.rs | 67 ++++- tests-integration/tests/http.rs | 23 +- 8 files changed, 532 insertions(+), 29 deletions(-) diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs index 671f6c6f7cf..5951728e6f4 100644 --- a/src/frontend/src/instance/otlp/trace_ingest.rs +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -31,7 +31,7 @@ use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use pipeline::{GreptimePipelineParams, PipelineWay}; use servers::error::{self, Result as ServerResult}; use servers::otlp; -use servers::otlp::coerce::{coerce_value_data, trace_value_datatype}; +use servers::otlp::coerce::{coerce_value_data, is_supported_trace_coercion, trace_value_datatype}; use servers::otlp::trace::span::{TraceSpan, TraceSpanGroup}; use servers::otlp::trace::v1::{TraceBatchSchema, TraceBinaryType, TraceRetryColumns}; use servers::otlp::trace::{SERVICE_NAME_COLUMN, TraceAuxData}; @@ -2025,10 +2025,23 @@ fn trace_logical_types_incompatible( right_datatype: ColumnDataType, right_concrete_type: &ConcreteDataType, ) -> bool { - left_concrete_type != right_concrete_type - && (left_datatype == right_datatype - || !is_trace_reconcile_candidate_type(left_datatype) - || !is_trace_reconcile_candidate_type(right_datatype)) + if left_concrete_type == right_concrete_type { + return false; + } + // A supported coercion in either direction means the two types can be + // reconciled, so they are not logically incompatible. This lets a signed + // request reach the coercion path instead of being excluded up front when + // the existing column is unsigned (e.g. trace `duration_nano` written as + // Int64 into an existing UInt64 column during the unsigned -> signed + // transition). + if is_supported_trace_coercion(left_datatype, right_datatype) + || is_supported_trace_coercion(right_datatype, left_datatype) + { + return false; + } + left_datatype == right_datatype + || !is_trace_reconcile_candidate_type(left_datatype) + || !is_trace_reconcile_candidate_type(right_datatype) } fn chunk_owned(items: Vec, chunk_size: usize) -> Vec> { diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index ec5bad89754..2445da95bb8 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -659,6 +659,34 @@ fn test_trace_request_schema_isolates_non_candidate_batch() { ); } +#[test] +fn test_trace_request_schema_keeps_signed_request_into_existing_unsigned_column() { + // Regression for the unsigned -> signed transition: an existing + // `duration_nano` UInt64 column must NOT exclude a new signed (Int64) + // batch as "logically incompatible". The batch must reach the + // reconciliation path, which coerces Int64 -> UInt64 in place (no ALTER). + // Excluding it up front made ingestion error after the flip. + let column_name = "duration_nano"; + let int64_schema = field_schema(column_name, ColumnDataType::Int64); + let existing_uint64_schema = + DatatypesSchemaBuilder::try_from_columns(vec![DatatypesColumnSchema::new( + column_name, + ConcreteDataType::uint64_datatype(), + true, + )]) + .unwrap() + .build() + .unwrap(); + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &int64_schema, Some(ColumnDataType::Int64)); + + assert_eq!( + request_schema.incompatible_schema_observations(Some(&existing_uint64_schema)), + HashMap::new() + ); +} + #[test] fn test_trace_request_schema_isolates_non_candidate_new_column_batch() { let column_name = "span_attributes.payload"; diff --git a/src/frontend/src/instance/otlp/trace_types.rs b/src/frontend/src/instance/otlp/trace_types.rs index 603cc66c0d1..0295d12e685 100644 --- a/src/frontend/src/instance/otlp/trace_types.rs +++ b/src/frontend/src/instance/otlp/trace_types.rs @@ -351,6 +351,24 @@ mod tests { ); } + #[test] + fn test_choose_trace_reconcile_decision_existing_uint64_keeps_uint64() { + // Backward-compat for the unsigned -> signed transition: an existing + // table whose `duration_nano` is still UInt64 must keep that type (no + // ALTER) when new signed (Int64) ingest arrives, coercing the value in + // place. This is the no-ALTER guarantee for the trace path; it relies on + // the Int64 -> UInt64 coercion arm added in Phase 0. + assert_eq!( + choose_trace_reconcile_decision( + "duration_nano", + &[ColumnDataType::Int64], + Some(ColumnDataType::Uint64) + ) + .unwrap(), + Some(TraceReconcileDecision::UseExisting(ColumnDataType::Uint64)) + ); + } + #[test] fn test_choose_trace_reconcile_decision_existing_int64_widens_to_float64() { assert_eq!( @@ -612,6 +630,42 @@ mod tests { ); } + #[test] + fn test_prepare_trace_column_rewrites_coerces_int64_into_existing_uint64() { + // Existing-table backward-compat for the trace path: new signed ingest + // arrives as Int64, but the existing `duration_nano` column is UInt64, so + // the rewrite coerces the value into the existing type in place (no + // ALTER). Mirrors what happens for a table created before the + // unsigned -> signed flip. + let mut rows = Rows { + schema: vec![ColumnSchema { + datatype: ColumnDataType::Int64 as i32, + ..Default::default() + }], + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::I64Value(42)), + }], + }], + }; + let pending_rewrites = vec![PendingTraceColumnRewrite { + col_idx: 0, + target_type: ColumnDataType::Uint64, + column_name: "duration_nano".to_string(), + }]; + + let prepared = + prepare_trace_column_rewrites(&rows.rows, pending_rewrites, "trace_type_atomicity") + .unwrap(); + + prepared.apply(&mut rows); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Uint64 as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::U64Value(42)) + ); + } + #[test] fn test_prepare_trace_column_rewrites_boolean_rejects_invalid_string_parse() { let rows = vec![Row { diff --git a/src/servers/src/otlp/coerce.rs b/src/servers/src/otlp/coerce.rs index 2d87744c5c2..2e7dd36ed7a 100644 --- a/src/servers/src/otlp/coerce.rs +++ b/src/servers/src/otlp/coerce.rs @@ -30,6 +30,37 @@ pub enum TraceCoerceError { // - String to Int64 // - String to Float64 // - String to Boolean +// +// Lossless signed-to-unsigned integer casts. These let the built-in data +// models move from unsigned to signed integers while existing unsigned tables +// keep accepting new signed ingest without an `ALTER TABLE`: an existing +// UInt64/UInt32 column coerces an incoming Int64/Int32 request into the +// existing type: +// - Int64 to UInt64 (e.g. trace `duration_nano` on the v1 path): checked, +// so negative values are rejected rather than silently wrapping. Counts and +// durations are non-negative by construction; a negative request is a +// malformed value (e.g. a span whose end precedes its start). +// - Int32 to UInt32 (e.g. log `trace_flags`): bit-preserving, because +// `trace_flags` is a bit field whose high bits may legitimately be set +// (e.g. the W3C sampled flag); bit patterns must round-trip exactly. + +/// The signed→unsigned integer coercions that let the built-in data models +/// move from unsigned to signed integers while existing unsigned tables keep +/// accepting new signed ingest without an `ALTER TABLE` (see the pair +/// descriptions in the module-level comment above). Kept as one predicate so +/// the trace and log ingest paths share the same supported pair set and +/// cannot drift. +pub fn is_supported_signed_to_unsigned_coercion( + request_type: ColumnDataType, + target_type: ColumnDataType, +) -> bool { + matches!( + (request_type, target_type), + (ColumnDataType::Int64, ColumnDataType::Uint64) + | (ColumnDataType::Int32, ColumnDataType::Uint32) + ) +} + pub fn is_supported_trace_coercion( request_type: ColumnDataType, target_type: ColumnDataType, @@ -43,7 +74,7 @@ pub fn is_supported_trace_coercion( | (ColumnDataType::String, ColumnDataType::Int64) | (ColumnDataType::String, ColumnDataType::Float64) | (ColumnDataType::String, ColumnDataType::Boolean) - ) + ) || is_supported_signed_to_unsigned_coercion(request_type, target_type) } pub fn coerce_value_data( @@ -88,6 +119,18 @@ pub fn coerce_non_null_value( (ColumnDataType::String, ColumnDataType::Boolean, ValueData::StringValue(s)) => { s.parse::().ok().map(ValueData::BoolValue) } + // Checked signed -> unsigned cast for built-in fields moving to signed + // types (durations, counts are always non-negative). Negative values + // are rejected instead of wrapping so the coercion stays lossless. + (ColumnDataType::Int64, ColumnDataType::Uint64, ValueData::I64Value(n)) => { + u64::try_from(*n).ok().map(ValueData::U64Value) + } + // Bit-preserving cast for the `trace_flags` bit field: high bits may + // legitimately be set, and the bit pattern must survive the round-trip + // into an existing UInt32 column unchanged. + (ColumnDataType::Int32, ColumnDataType::Uint32, ValueData::I32Value(n)) => { + Some(ValueData::U32Value(*n as u32)) + } _ => None, } } @@ -240,6 +283,55 @@ mod tests { assert_eq!(result, Err(TraceCoerceError::Unsupported)); } + #[test] + fn test_coerce_int64_to_uint64() { + // Non-negative durations coerce losslessly into an existing UInt64 + // column, so built-in fields moving to signed keep accepting writes. + let result = coerce_value_data( + &Some(ValueData::I64Value(123)), + ColumnDataType::Uint64, + ColumnDataType::Int64, + ); + assert_eq!(result, Ok(Some(ValueData::U64Value(123)))); + } + + #[test] + fn test_coerce_negative_int64_to_uint64_rejected() { + // Negative values must not wrap into the unsigned column: the + // signed -> unsigned coercion is only lossless for non-negative + // inputs, so anything else is rejected instead of silently wrapping + // to a huge number. + let result = coerce_value_data( + &Some(ValueData::I64Value(-1)), + ColumnDataType::Uint64, + ColumnDataType::Int64, + ); + assert_eq!(result, Err(TraceCoerceError::Unsupported)); + } + + #[test] + fn test_coerce_int32_to_uint32() { + let result = coerce_value_data( + &Some(ValueData::I32Value(7)), + ColumnDataType::Uint32, + ColumnDataType::Int32, + ); + assert_eq!(result, Ok(Some(ValueData::U32Value(7)))); + } + + #[test] + fn test_coerce_uint_to_int_not_supported() { + // Only the signed -> unsigned direction is supported (the direction + // the no-ALTER transition needs); the reverse would be lossy for + // values above the signed range and is intentionally rejected. + let result = coerce_value_data( + &Some(ValueData::U64Value(9)), + ColumnDataType::Int64, + ColumnDataType::Uint64, + ); + assert_eq!(result, Err(TraceCoerceError::Unsupported)); + } + #[test] fn test_coerce_none_value() { let result = coerce_value_data(&None, ColumnDataType::Float64, ColumnDataType::Int64); @@ -280,6 +372,20 @@ mod tests { ColumnDataType::Binary, ColumnDataType::Json )); + // Signed -> unsigned casts are supported (built-in no-ALTER transition). + assert!(is_supported_trace_coercion( + ColumnDataType::Int64, + ColumnDataType::Uint64 + )); + assert!(is_supported_trace_coercion( + ColumnDataType::Int32, + ColumnDataType::Uint32 + )); + // The reverse direction is intentionally not supported (lossy). + assert!(!is_supported_trace_coercion( + ColumnDataType::Uint64, + ColumnDataType::Int64 + )); } #[test] diff --git a/src/servers/src/otlp/logs.rs b/src/servers/src/otlp/logs.rs index 33d4f4eb5e5..74f1a3d03a6 100644 --- a/src/servers/src/otlp/logs.rs +++ b/src/servers/src/otlp/logs.rs @@ -44,7 +44,7 @@ use crate::error::{ UnsupportedJsonDataTypeForTagSnafu, }; use crate::http::event::PipelineIngestRequest; -use crate::otlp::coerce::coerce_value_data; +use crate::otlp::coerce::{coerce_value_data, is_supported_signed_to_unsigned_coercion}; use crate::otlp::trace::attributes::OtlpAnyValue; use crate::otlp::utils::{bytes_to_hex_string, key_value_to_jsonb}; use crate::pipeline::run_pipeline; @@ -257,7 +257,7 @@ fn build_otlp_logs_identity_schema() -> Vec { ), ( "trace_flags", - ColumnDataType::Uint32, + ColumnDataType::Int32, SemanticType::Field, None, None, @@ -440,7 +440,7 @@ fn build_otlp_build_in_row( value_data: Some(ValueData::BinaryValue(log_attr.to_vec())), }, GreptimeValue { - value_data: Some(ValueData::U32Value(log.flags)), + value_data: Some(ValueData::I32Value(log.flags as i32)), }, GreptimeValue { value_data: parse_ctx.scope_name.clone().map(ValueData::StringValue), @@ -549,14 +549,17 @@ fn decide_column_schema_and_convert_value( key: column_name, } .fail(), - JsonbNumber::UInt64(u) => Ok(Some(( - GreptimeValue { - value_data: Some(ValueData::U64Value(u)), - }, - ColumnDataType::Uint64, - SemanticType::Tag, - None, - ))), + JsonbNumber::UInt64(u) => { + let value = jsonb_uint64_to_log_value(u, column_name)?; + Ok(Some(( + GreptimeValue { + value_data: Some(ValueData::I64Value(value)), + }, + ColumnDataType::Int64, + SemanticType::Tag, + None, + ))) + } }, JsonbValue::Bool(b) => Ok(Some(( GreptimeValue { @@ -615,6 +618,21 @@ fn decide_existing_column_schema_and_convert_value( ))) } +/// Converts a JSON `UInt64` number into the signed value used by built-in +/// log columns. `JsonbNumber::UInt64` only arises for values that do not fit +/// in `i64` (see `decide_column_schema_and_convert_value`), so the conversion +/// is fallible: wrapping to a negative number would silently corrupt the +/// value (a counter reading back as negative), so out-of-range values are +/// rejected instead. +fn jsonb_uint64_to_log_value(u: u64, column_name: &str) -> Result { + i64::try_from(u).map_err(|_| InvalidParameterSnafu { + reason: format!( + "uint64 value {u} in column '{column_name}' exceeds the i64 range supported by built-in log columns" + ), + } + .build()) +} + fn jsonb_value_to_log_value_data( column_name: &str, value: JsonbValue, @@ -635,7 +653,10 @@ fn jsonb_value_to_log_value_data( key: column_name, } .fail(), - JsonbNumber::UInt64(u) => Ok(Some((ValueData::U64Value(u), ColumnDataType::Uint64))), + JsonbNumber::UInt64(u) => Ok(Some(( + ValueData::I64Value(jsonb_uint64_to_log_value(u, column_name)?), + ColumnDataType::Int64, + ))), }, JsonbValue::Bool(b) => Ok(Some((ValueData::BoolValue(b), ColumnDataType::Boolean))), JsonbValue::Array(_) | JsonbValue::Object(_) => UnsupportedJsonDataTypeForTagSnafu { @@ -718,6 +739,14 @@ fn coerce_log_value_data( return align_timestamp_value(value_data, target_unit, column_name, table_name).map(Some); } + // Lossless signed -> unsigned integer cast for built-in fields moving from + // unsigned to signed types (e.g. an existing UInt64/UInt32 log column + // receiving new Int64/Int32 ingest). Lets existing tables keep their + // unsigned columns without an `ALTER`. In the mutually-exclusive `else` of + // the String branch so the move stays local to non-String targets; the + // shared predicate keeps the supported pairs identical to the trace path + // (the generic String -> numeric coercions keep their pre-existing + // rejection on this path). if target_type == ColumnDataType::String { if let Ok(value_data) = coerce_value_data(&Some(value_data.clone()), target_type, request_type) @@ -727,6 +756,11 @@ fn coerce_log_value_data( if let Some(value_data) = stringify_scalar_value(value_data) { return Ok(Some(value_data)); } + } else if is_supported_signed_to_unsigned_coercion(request_type, target_type) + && let Ok(Some(value_data)) = + coerce_value_data(&Some(value_data), target_type, request_type) + { + return Ok(Some(value_data)); } InvalidParameterSnafu { @@ -1089,6 +1123,16 @@ mod tests { ExistingLogSchema::try_from_schema_parts(&columns, primary_key_indices).unwrap() } + fn existing_uint32_trace_flags_schema() -> ExistingLogSchema { + existing_schema( + vec![ + time_column(ConcreteDataType::timestamp_nanosecond_datatype()), + column("trace_flags", ConcreteDataType::uint32_datatype()), + ], + &[], + ) + } + fn kv(key: &str, value: OtlpValue) -> KeyValue { KeyValue { key: key.to_string(), @@ -1098,12 +1142,20 @@ mod tests { } fn request_with_log_attrs(attrs: Vec) -> ExportLogsServiceRequest { + request_with_log_attrs_and_flags(attrs, 0) + } + + fn request_with_log_attrs_and_flags( + attrs: Vec, + flags: u32, + ) -> ExportLogsServiceRequest { ExportLogsServiceRequest { resource_logs: vec![ResourceLogs { scope_logs: vec![ScopeLogs { log_records: vec![LogRecord { time_unix_nano: 1_234_000_000, trace_id: vec![1; 16], + flags, attributes: attrs, ..Default::default() }], @@ -1151,6 +1203,83 @@ mod tests { ); } + #[test] + fn test_fresh_table_trace_flags_is_int32() { + // Phase 1 of the unsigned -> signed transition: a fresh log table now + // creates `trace_flags` as Int32. Existing tables keep UInt32 (covered + // by the coercion test above); new tables start signed. + let rows = parse_with_select(request_with_log_attrs(vec![]), "", None).unwrap(); + let idx = column_index(&rows, "trace_flags"); + + assert_eq!(rows.schema[idx].datatype, ColumnDataType::Int32 as i32); + assert_eq!( + rows.rows[0].values[idx].value_data, + Some(ValueData::I32Value(0)) + ); + } + + #[test] + fn test_existing_uint32_trace_flags_keeps_type_and_coerces_int32_request() { + // An existing UInt32 `trace_flags` column keeps its type when new + // signed (Int32) ingest arrives, without an `ALTER`. + let existing = existing_uint32_trace_flags_schema(); + + let rows = parse_with_select(request_with_log_attrs(vec![]), "", Some(&existing)).unwrap(); + let idx = column_index(&rows, "trace_flags"); + + assert_eq!(rows.schema[idx].datatype, ColumnDataType::Uint32 as i32); + assert_eq!( + rows.rows[0].values[idx].value_data, + Some(ValueData::U32Value(0)) + ); + } + + #[test] + fn test_existing_uint32_trace_flags_preserves_nonzero_flags_value() { + // A non-zero trace_flags bit pattern (e.g. the W3C sampled flag as sent + // by common encoders) round-trips exactly through the Int32 -> UInt32 + // coercion into an existing unsigned column. Guards that the cast is + // bit-preserving for realistic values, not just the default 0. + let existing = existing_uint32_trace_flags_schema(); + + let rows = parse_with_select( + request_with_log_attrs_and_flags(vec![], 256), + "", + Some(&existing), + ) + .unwrap(); + let idx = column_index(&rows, "trace_flags"); + + assert_eq!(rows.schema[idx].datatype, ColumnDataType::Uint32 as i32); + assert_eq!( + rows.rows[0].values[idx].value_data, + Some(ValueData::U32Value(256)) + ); + } + + #[test] + fn test_existing_uint32_trace_flags_preserves_high_bit_flags_value() { + // The full u32 range round-trips: a value with the sign bit set proves + // the Int32 -> UInt32 coercion is bit-exact (`as i32` then `as u32` + // preserve the pattern), so no flags value is corrupted on the + // unsigned -> signed transition. + let existing = existing_uint32_trace_flags_schema(); + + let rows = parse_with_select( + request_with_log_attrs_and_flags(vec![], 0x8000_0000), + "", + Some(&existing), + ) + .unwrap(); + let idx = column_index(&rows, "trace_flags"); + + assert_eq!(rows.schema[idx].datatype, ColumnDataType::Uint32 as i32); + assert_eq!( + rows.rows[0].values[idx].value_data, + Some(ValueData::U32Value(0x8000_0000)) + ); + } + #[test] fn test_existing_primary_key_updates_builtin_column_semantic_type() { let existing = existing_schema( @@ -1326,4 +1455,82 @@ mod tests { SemanticType::Field as i32 ); } + + #[test] + fn test_existing_uint64_column_keeps_type_and_coerces_int64_request() { + // An existing UInt64 column keeps its type when new signed (Int64) + // ingest arrives, without an `ALTER`. The built-in models move to + // signed while existing unsigned tables stay byte-for-byte unchanged. + let existing = existing_schema( + vec![ + time_column(ConcreteDataType::timestamp_nanosecond_datatype()), + column("counter", ConcreteDataType::uint64_datatype()), + ], + &[], + ); + + let rows = parse_with_select( + request_with_log_attrs(vec![kv("counter", OtlpValue::IntValue(42))]), + "counter", + Some(&existing), + ) + .unwrap(); + let idx = column_index(&rows, "counter"); + + assert_eq!(rows.schema[idx].datatype, ColumnDataType::Uint64 as i32); + assert_eq!( + rows.rows[0].values[idx].value_data, + Some(ValueData::U64Value(42)) + ); + } + + #[test] + fn test_jsonb_uint64_exceeding_i64_range_rejected() { + // A JSON number that only fits in u64 must be rejected, not wrapped: + // storing it as a negative i64 would silently corrupt the value. The + // unsigned -> signed transition is only lossless for values that fit + // the signed range. + let err = jsonb_value_to_log_value_data( + "counter", + JsonbValue::Number(JsonbNumber::UInt64(u64::MAX)), + false, + ) + .unwrap_err(); + + assert!( + err.to_string() + .contains("exceeds the i64 range supported by built-in log columns") + ); + } + + #[test] + fn test_existing_int64_column_rejects_numeric_string_value() { + // Regression guard for the signed -> unsigned transition scope: even a + // parseable numeric string must keep its pre-existing rejection when + // targeting a non-String column. Only the signed -> unsigned integer + // pairs gained coercion on this path, not the generic String -> + // numeric ones. + let existing = existing_schema( + vec![ + time_column(ConcreteDataType::timestamp_nanosecond_datatype()), + column("counter", ConcreteDataType::int64_datatype()), + ], + &[], + ); + + let err = parse_with_select( + request_with_log_attrs(vec![kv( + "counter", + OtlpValue::StringValue("42".to_string()), + )]), + "counter", + Some(&existing), + ) + .unwrap_err(); + + assert!( + err.to_string() + .contains("failed to align log column 'counter'") + ); + } } diff --git a/src/servers/src/otlp/trace/v0.rs b/src/servers/src/otlp/trace/v0.rs index fa10dcc00f8..1bc6752e900 100644 --- a/src/servers/src/otlp/trace/v0.rs +++ b/src/servers/src/otlp/trace/v0.rs @@ -108,6 +108,11 @@ pub fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> span.end_in_nanosecond as i64, )), ), + // The v0 data model is frozen: `duration_nano` stays UInt64. The + // signed→unsigned ingest compatibility layer only runs on the v1 path + // (see `trace_ingest.rs`), so flipping v0 here would break writes into + // every pre-existing v0 table at mito's schema check. New v1 tables + // are signed; v0 tables keep the legacy unsigned column indefinitely. make_column_data( DURATION_NANO_COLUMN, ColumnDataType::Uint64, @@ -231,10 +236,13 @@ fn write_trace_operations_to_row( #[cfg(test)] mod tests { + use api::v1::ColumnDataType; + use api::v1::value::ValueData; + use super::{build_aux_table_requests, build_trace_table_data}; - use crate::otlp::trace::TraceAuxData; use crate::otlp::trace::attributes::Attributes; use crate::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; + use crate::otlp::trace::{DURATION_NANO_COLUMN, TraceAuxData}; fn make_span(service_name: &str, trace_id: &str, span_id: &str) -> TraceSpan { TraceSpan { @@ -271,6 +279,23 @@ mod tests { assert_eq!(rows.len(), 1); } + #[test] + fn test_v0_duration_nano_stays_uint64() { + // The v0 data model is frozen on the unsigned `duration_nano`. This + // pins the schema so the unsigned→signed transition of the built-in + // models (which flips only v1) cannot silently leak into v0 and break + // writes into pre-existing v0 tables. + let writer = build_trace_table_data(&[make_span("svc-a", "trace-a", "span-a")]).unwrap(); + let (schema, rows) = writer.into_schema_and_rows(); + + let idx = schema + .iter() + .position(|c| c.column_name == DURATION_NANO_COLUMN) + .unwrap(); + assert_eq!(schema[idx].datatype, ColumnDataType::Uint64 as i32); + assert_eq!(rows[0].values[idx].value_data, Some(ValueData::U64Value(1))); + } + #[test] fn test_build_aux_table_requests_deduplicates_services_and_operations() { let spans = vec![ diff --git a/src/servers/src/otlp/trace/v1.rs b/src/servers/src/otlp/trace/v1.rs index 8ddeac24f33..f55ab48251e 100644 --- a/src/servers/src/otlp/trace/v1.rs +++ b/src/servers/src/otlp/trace/v1.rs @@ -79,7 +79,7 @@ impl FixedTraceColumnIndexes { Ok(Self { timestamp, timestamp_end: field("timestamp_end", ColumnDataType::TimestampNanosecond)?, - duration_nano: field(DURATION_NANO_COLUMN, ColumnDataType::Uint64)?, + duration_nano: field(DURATION_NANO_COLUMN, ColumnDataType::Int64)?, parent_span_id: field(PARENT_SPAN_ID_COLUMN, ColumnDataType::String)?, trace_id: field(TRACE_ID_COLUMN, ColumnDataType::String)?, span_id: field(SPAN_ID_COLUMN, ColumnDataType::String)?, @@ -342,6 +342,22 @@ pub fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> write_span_to_row_inner(writer, span, row_index, &fixed_columns, None) } +/// Computes the span duration as the signed `duration_nano` value written by +/// the v1 data model. +/// +/// A span whose end precedes its start carries no meaningful duration; it +/// clamps to 0 instead of wrapping — a negative `duration_nano` would fail +/// the checked Int64→UInt64 coercion into pre-existing unsigned tables and +/// would break unsigned readers such as the Jaeger query API. A duration +/// that does not fit `i64` saturates at `i64::MAX`. Clamping at the source +/// keeps new Int64 tables and existing UInt64 tables behaving identically: +/// the written value is always a non-negative, in-range `i64`. +fn span_duration_nano(span: &TraceSpan) -> i64 { + span.end_in_nanosecond + .saturating_sub(span.start_in_nanosecond) + .min(i64::MAX as u64) as i64 +} + /// Writes one span and optionally records its dynamic columns for reconciliation. fn write_span_to_row_inner( writer: &mut TableData, @@ -367,9 +383,7 @@ fn write_span_to_row_inner( ), ( fixed_columns.duration_nano, - Some(ValueData::U64Value( - span.end_in_nanosecond - span.start_in_nanosecond, - )), + Some(ValueData::I64Value(span_duration_nano(&span))), ), ( fixed_columns.parent_span_id, @@ -689,6 +703,49 @@ mod tests { } } + #[test] + fn test_span_end_before_start_records_zero_duration() { + // A span whose end precedes its start carries no meaningful duration; + // it clamps to 0 instead of wrapping. A negative value would fail the + // checked Int64→UInt64 coercion into pre-existing unsigned tables and + // break unsigned readers such as the Jaeger query API. + let mut span = make_span("svc", "trace", "span"); + span.start_in_nanosecond = 200; + span.end_in_nanosecond = 100; + + let (schema, rows) = build_trace_table_data(&[span]) + .unwrap() + .into_schema_and_rows(); + + let idx = schema + .iter() + .position(|c| c.column_name == DURATION_NANO_COLUMN) + .unwrap(); + assert_eq!(rows[0].values[idx].value_data, Some(ValueData::I64Value(0))); + } + + #[test] + fn test_span_duration_above_i64_max_saturates() { + // The companion clamp: a duration that fits u64 but not i64 saturates + // at i64::MAX rather than wrapping to a negative number. + let mut span = make_span("svc", "trace", "span"); + span.start_in_nanosecond = 0; + span.end_in_nanosecond = u64::MAX; + + let (schema, rows) = build_trace_table_data(&[span]) + .unwrap() + .into_schema_and_rows(); + + let idx = schema + .iter() + .position(|c| c.column_name == DURATION_NANO_COLUMN) + .unwrap(); + assert_eq!( + rows[0].values[idx].value_data, + Some(ValueData::I64Value(i64::MAX)) + ); + } + #[test] fn test_fixed_trace_columns_keep_schema_and_values_aligned() { let table_data = build_trace_table_data(&[make_span("svc", "trace", "span")]).unwrap(); @@ -702,7 +759,7 @@ mod tests { "timestamp_end", Some(ValueData::TimestampNanosecondValue(2)), ), - (DURATION_NANO_COLUMN, Some(ValueData::U64Value(1))), + (DURATION_NANO_COLUMN, Some(ValueData::I64Value(1))), (PARENT_SPAN_ID_COLUMN, None), ( TRACE_ID_COLUMN, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 90cbb8c93fe..f89d75e9ef0 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -6860,6 +6860,19 @@ pub async fn test_otlp_traces_v0(store_type: StorageType) { ) .await; + // The v0 data model is frozen on the unsigned `duration_nano`: pin the + // schema so the unsigned→signed flip of the built-in models (which flips + // only v1) cannot leak into v0 and break writes into pre-existing v0 + // tables. + validate_data( + "otlp_traces_v0_schema", + &client, + "select column_name, lower(data_type) from information_schema.columns \ + where table_name = 'opentelemetry_traces' and column_name = 'duration_nano';", + r#"[["duration_nano","bigint unsigned"]]"#, + ) + .await; + // drop table let res = client .get("/v1/sql?sql=drop table opentelemetry_traces;") @@ -7075,7 +7088,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; - let expected_ddl = r#"[["mytable","CREATE TABLE IF NOT EXISTS \"mytable\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '1',\n trace_id >= '1' AND trace_id < '2',\n trace_id >= '2' AND trace_id < '3',\n trace_id >= '3' AND trace_id < '4',\n trace_id >= '4' AND trace_id < '5',\n trace_id >= '5' AND trace_id < '6',\n trace_id >= '6' AND trace_id < '7',\n trace_id >= '7' AND trace_id < '8',\n trace_id >= '8' AND trace_id < '9',\n trace_id >= '9' AND trace_id < 'a',\n trace_id >= 'a' AND trace_id < 'b',\n trace_id >= 'b' AND trace_id < 'c',\n trace_id >= 'c' AND trace_id < 'd',\n trace_id >= 'd' AND trace_id < 'e',\n trace_id >= 'e' AND trace_id < 'f',\n trace_id >= 'f'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["mytable","CREATE TABLE IF NOT EXISTS \"mytable\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '1',\n trace_id >= '1' AND trace_id < '2',\n trace_id >= '2' AND trace_id < '3',\n trace_id >= '3' AND trace_id < '4',\n trace_id >= '4' AND trace_id < '5',\n trace_id >= '5' AND trace_id < '6',\n trace_id >= '6' AND trace_id < '7',\n trace_id >= '7' AND trace_id < '8',\n trace_id >= '8' AND trace_id < '9',\n trace_id >= '9' AND trace_id < 'a',\n trace_id >= 'a' AND trace_id < 'b',\n trace_id >= 'b' AND trace_id < 'c',\n trace_id >= 'c' AND trace_id < 'd',\n trace_id >= 'd' AND trace_id < 'e',\n trace_id >= 'e' AND trace_id < 'f',\n trace_id >= 'f'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -7128,7 +7141,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part1","CREATE TABLE IF NOT EXISTS \"trace_table_part1\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\n\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part1","CREATE TABLE IF NOT EXISTS \"trace_table_part1\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\n\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -7165,7 +7178,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part4","CREATE TABLE IF NOT EXISTS \"trace_table_part4\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '4',\n trace_id >= '4' AND trace_id < '8',\n trace_id >= '8' AND trace_id < 'c',\n trace_id >= 'c'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part4","CREATE TABLE IF NOT EXISTS \"trace_table_part4\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '4',\n trace_id >= '4' AND trace_id < '8',\n trace_id >= '8' AND trace_id < 'c',\n trace_id >= 'c'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -7202,7 +7215,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part32","CREATE TABLE IF NOT EXISTS \"trace_table_part32\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '08',\n trace_id >= '08' AND trace_id < '10',\n trace_id >= '10' AND trace_id < '18',\n trace_id >= '18' AND trace_id < '20',\n trace_id >= '20' AND trace_id < '28',\n trace_id >= '28' AND trace_id < '30',\n trace_id >= '30' AND trace_id < '38',\n trace_id >= '38' AND trace_id < '40',\n trace_id >= '40' AND trace_id < '48',\n trace_id >= '48' AND trace_id < '50',\n trace_id >= '50' AND trace_id < '58',\n trace_id >= '58' AND trace_id < '60',\n trace_id >= '60' AND trace_id < '68',\n trace_id >= '68' AND trace_id < '70',\n trace_id >= '70' AND trace_id < '78',\n trace_id >= '78' AND trace_id < '80',\n trace_id >= '80' AND trace_id < '88',\n trace_id >= '88' AND trace_id < '90',\n trace_id >= '90' AND trace_id < '98',\n trace_id >= '98' AND trace_id < 'a0',\n trace_id >= 'a0' AND trace_id < 'a8',\n trace_id >= 'a8' AND trace_id < 'b0',\n trace_id >= 'b0' AND trace_id < 'b8',\n trace_id >= 'b8' AND trace_id < 'c0',\n trace_id >= 'c0' AND trace_id < 'c8',\n trace_id >= 'c8' AND trace_id < 'd0',\n trace_id >= 'd0' AND trace_id < 'd8',\n trace_id >= 'd8' AND trace_id < 'e0',\n trace_id >= 'e0' AND trace_id < 'e8',\n trace_id >= 'e8' AND trace_id < 'f0',\n trace_id >= 'f0' AND trace_id < 'f8',\n trace_id >= 'f8'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part32","CREATE TABLE IF NOT EXISTS \"trace_table_part32\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '08',\n trace_id >= '08' AND trace_id < '10',\n trace_id >= '10' AND trace_id < '18',\n trace_id >= '18' AND trace_id < '20',\n trace_id >= '20' AND trace_id < '28',\n trace_id >= '28' AND trace_id < '30',\n trace_id >= '30' AND trace_id < '38',\n trace_id >= '38' AND trace_id < '40',\n trace_id >= '40' AND trace_id < '48',\n trace_id >= '48' AND trace_id < '50',\n trace_id >= '50' AND trace_id < '58',\n trace_id >= '58' AND trace_id < '60',\n trace_id >= '60' AND trace_id < '68',\n trace_id >= '68' AND trace_id < '70',\n trace_id >= '70' AND trace_id < '78',\n trace_id >= '78' AND trace_id < '80',\n trace_id >= '80' AND trace_id < '88',\n trace_id >= '88' AND trace_id < '90',\n trace_id >= '90' AND trace_id < '98',\n trace_id >= '98' AND trace_id < 'a0',\n trace_id >= 'a0' AND trace_id < 'a8',\n trace_id >= 'a8' AND trace_id < 'b0',\n trace_id >= 'b0' AND trace_id < 'b8',\n trace_id >= 'b8' AND trace_id < 'c0',\n trace_id >= 'c0' AND trace_id < 'c8',\n trace_id >= 'c8' AND trace_id < 'd0',\n trace_id >= 'd0' AND trace_id < 'd8',\n trace_id >= 'd8' AND trace_id < 'e0',\n trace_id >= 'e0' AND trace_id < 'e8',\n trace_id >= 'e8' AND trace_id < 'f0',\n trace_id >= 'f0' AND trace_id < 'f8',\n trace_id >= 'f8'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -9601,7 +9614,7 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { .await; assert_eq!(StatusCode::OK, res.status()); - let trace_table_sql = "[[\"mytable\",\"CREATE TABLE IF NOT EXISTS \\\"mytable\\\" (\\n \\\"timestamp\\\" TIMESTAMP(9) NOT NULL,\\n \\\"timestamp_end\\\" TIMESTAMP(9) NULL,\\n \\\"duration_nano\\\" BIGINT UNSIGNED NULL,\\n \\\"parent_span_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"trace_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_id\\\" STRING NULL,\\n \\\"span_kind\\\" STRING NULL,\\n \\\"span_name\\\" STRING NULL,\\n \\\"span_status_code\\\" STRING NULL,\\n \\\"span_status_message\\\" STRING NULL,\\n \\\"trace_state\\\" STRING NULL,\\n \\\"scope_name\\\" STRING NULL,\\n \\\"scope_version\\\" STRING NULL,\\n \\\"service_name\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_attributes.operation.type\\\" STRING NULL,\\n \\\"span_attributes.net.peer.ip\\\" STRING NULL,\\n \\\"span_attributes.peer.service\\\" STRING NULL,\\n \\\"span_events\\\" JSON NULL,\\n \\\"span_links\\\" JSON NULL,\\n TIME INDEX (\\\"timestamp\\\"),\\n PRIMARY KEY (\\\"service_name\\\")\\n)\\nPARTITION ON COLUMNS (\\\"trace_id\\\") (\\n trace_id < '1',\\n trace_id >= '1' AND trace_id < '2',\\n trace_id >= '2' AND trace_id < '3',\\n trace_id >= '3' AND trace_id < '4',\\n trace_id >= '4' AND trace_id < '5',\\n trace_id >= '5' AND trace_id < '6',\\n trace_id >= '6' AND trace_id < '7',\\n trace_id >= '7' AND trace_id < '8',\\n trace_id >= '8' AND trace_id < '9',\\n trace_id >= '9' AND trace_id < 'a',\\n trace_id >= 'a' AND trace_id < 'b',\\n trace_id >= 'b' AND trace_id < 'c',\\n trace_id >= 'c' AND trace_id < 'd',\\n trace_id >= 'd' AND trace_id < 'e',\\n trace_id >= 'e' AND trace_id < 'f',\\n trace_id >= 'f'\\n)\\nENGINE=mito\\nWITH(\\n 'comment' = 'Created on insertion',\\n append_mode = 'true',\\n 'greptime.semantic.entity.service.id' = 'service_name',\\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\\n 'greptime.semantic.signal_type' = 'trace',\\n 'greptime.semantic.source' = 'opentelemetry',\\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\\n table_data_model = 'greptime_trace_v1',\\n ttl = '7days'\\n)\"]]"; + let trace_table_sql = "[[\"mytable\",\"CREATE TABLE IF NOT EXISTS \\\"mytable\\\" (\\n \\\"timestamp\\\" TIMESTAMP(9) NOT NULL,\\n \\\"timestamp_end\\\" TIMESTAMP(9) NULL,\\n \\\"duration_nano\\\" BIGINT NULL,\\n \\\"parent_span_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"trace_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_id\\\" STRING NULL,\\n \\\"span_kind\\\" STRING NULL,\\n \\\"span_name\\\" STRING NULL,\\n \\\"span_status_code\\\" STRING NULL,\\n \\\"span_status_message\\\" STRING NULL,\\n \\\"trace_state\\\" STRING NULL,\\n \\\"scope_name\\\" STRING NULL,\\n \\\"scope_version\\\" STRING NULL,\\n \\\"service_name\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_attributes.operation.type\\\" STRING NULL,\\n \\\"span_attributes.net.peer.ip\\\" STRING NULL,\\n \\\"span_attributes.peer.service\\\" STRING NULL,\\n \\\"span_events\\\" JSON NULL,\\n \\\"span_links\\\" JSON NULL,\\n TIME INDEX (\\\"timestamp\\\"),\\n PRIMARY KEY (\\\"service_name\\\")\\n)\\nPARTITION ON COLUMNS (\\\"trace_id\\\") (\\n trace_id < '1',\\n trace_id >= '1' AND trace_id < '2',\\n trace_id >= '2' AND trace_id < '3',\\n trace_id >= '3' AND trace_id < '4',\\n trace_id >= '4' AND trace_id < '5',\\n trace_id >= '5' AND trace_id < '6',\\n trace_id >= '6' AND trace_id < '7',\\n trace_id >= '7' AND trace_id < '8',\\n trace_id >= '8' AND trace_id < '9',\\n trace_id >= '9' AND trace_id < 'a',\\n trace_id >= 'a' AND trace_id < 'b',\\n trace_id >= 'b' AND trace_id < 'c',\\n trace_id >= 'c' AND trace_id < 'd',\\n trace_id >= 'd' AND trace_id < 'e',\\n trace_id >= 'e' AND trace_id < 'f',\\n trace_id >= 'f'\\n)\\nENGINE=mito\\nWITH(\\n 'comment' = 'Created on insertion',\\n append_mode = 'true',\\n 'greptime.semantic.entity.service.id' = 'service_name',\\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\\n 'greptime.semantic.signal_type' = 'trace',\\n 'greptime.semantic.source' = 'opentelemetry',\\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\\n table_data_model = 'greptime_trace_v1',\\n ttl = '7days'\\n)\"]]"; validate_data( "trace_v1_create_table", &client, From 4ac3423261488fe9d84ed1e2068c79180821072e Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Mon, 24 Aug 2026 10:02:32 +0000 Subject: [PATCH 094/126] fix(mito2): split SSTs at primary key series boundaries (#8888) * fix(mito2): split SSTs at series boundaries Signed-off-by: Lei, HUANG * test(mito2): cover SST splitting without primary key Also document the sortedness precondition and the series boundary split semantics on write_all_flat/write_all_flat_as_primary_key and the new split helpers. Signed-off-by: Lei, HUANG * test(mito2): avoid per-row Vec allocation for empty primary key Signed-off-by: Lei, HUANG * Update src/mito2/src/sst/parquet/writer.rs Co-authored-by: Yingwen --------- Signed-off-by: Lei, HUANG Co-authored-by: Yingwen --- src/mito2/src/sst/parquet.rs | 223 ++++++++++++++++++++++++++- src/mito2/src/sst/parquet/writer.rs | 229 +++++++++++++++++++++------- 2 files changed, 395 insertions(+), 57 deletions(-) diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 5c40bc36bed..e67dcf2bee5 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -152,6 +152,7 @@ mod tests { use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder; use crate::sst::index::{IndexBuildType, Indexer, IndexerBuilder, IndexerBuilderImpl}; use crate::sst::parquet::flat_format::FlatWriteFormat; + use crate::sst::parquet::metadata::extract_primary_key_range; use crate::sst::parquet::reader::{ParquetReader, ParquetReaderBuilder, ReaderMetrics}; use crate::sst::parquet::row_selection::RowGroupSelection; use crate::sst::parquet::writer::ParquetWriter; @@ -687,12 +688,13 @@ mod tests { let object_store = env.init_object_store_manager(); let metadata = Arc::new(sst_region_metadata()); let batches = vec![ - new_record_batch_by_range(&["a", "d"], 0, 1000), - new_record_batch_by_range(&["b", "f"], 0, 1000), - new_record_batch_by_range(&["c", "g"], 0, 1000), - new_record_batch_by_range(&["b", "h"], 100, 200), - new_record_batch_by_range(&["b", "h"], 200, 300), - new_record_batch_by_range(&["b", "h"], 300, 1000), + new_record_batch_by_range(&["a", "a"], 0, 1000), + new_record_batch_by_range(&["b", "b"], 0, 1000), + new_record_batch_by_range(&["c", "c"], 0, 1000), + new_record_batch_by_range(&["d", "d"], 100, 200), + new_record_batch_by_range(&["d", "d"], 200, 300), + new_record_batch_by_range(&["d", "d"], 300, 1000), + new_record_batch_by_range(&["e", "e"], 0, 100), ]; let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); @@ -745,6 +747,159 @@ mod tests { assert_eq!(total_rows, rows_read); } + #[tokio::test] + async fn test_split_file_at_series_boundary_inside_batch() { + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let metadata = Arc::new(sst_region_metadata()); + let first_batch_rows = (0..1000).map(|ts| ("a", "a", ts)).collect::>(); + let second_batch_rows = (0..1000).map(|ts| ("b", "b", ts)).collect::>(); + let mut third_batch_rows = (1000..2000).map(|ts| ("b", "b", ts)).collect::>(); + third_batch_rows.extend((0..1000).map(|ts| ("c", "c", ts))); + let batches = vec![ + new_record_batch_from_rows(&first_batch_rows), + new_record_batch_from_rows(&second_batch_rows), + new_record_batch_from_rows(&third_batch_rows), + ]; + let total_rows = batches.iter().map(RecordBatch::num_rows).sum::(); + let source = new_flat_source_from_record_batches(batches); + let write_opts = WriteOptions { + row_group_size: 50, + max_file_size: Some(1), + ..Default::default() + }; + let path_provider = RegionFilePathFactory { + table_dir: "test_series_boundary".to_string(), + path_type: PathType::Bare, + }; + let mut metrics = Metrics::new(WriteType::Compaction); + let mut writer = ParquetWriter::new_with_object_store( + object_store, + metadata.clone(), + IndexConfig::default(), + NoopIndexBuilder, + path_provider, + &mut metrics, + ) + .await; + + let files = writer + .write_all_flat(source, None, &write_opts) + .await + .unwrap(); + + assert!(files.len() > 1); + assert_eq!( + total_rows, + files.iter().map(|file| file.num_rows).sum::() + ); + let primary_key_ranges = files + .iter() + .map(|file| { + extract_primary_key_range(file.file_metadata.as_ref().unwrap(), metadata.as_ref()) + .unwrap() + }) + .collect::>(); + assert!( + primary_key_ranges + .windows(2) + .all(|ranges| ranges[0].1 < ranges[1].0) + ); + } + + #[tokio::test] + async fn test_oversized_single_series_stays_in_one_file() { + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let metadata = Arc::new(sst_region_metadata()); + let first_batch_rows = (0..1000).map(|ts| ("a", "a", ts)).collect::>(); + let second_batch_rows = (1000..2000).map(|ts| ("a", "a", ts)).collect::>(); + let third_batch_rows = (2000..3000).map(|ts| ("a", "a", ts)).collect::>(); + let batches = vec![ + new_record_batch_from_rows(&first_batch_rows), + new_record_batch_from_rows(&second_batch_rows), + new_record_batch_from_rows(&third_batch_rows), + ]; + let total_rows = batches.iter().map(RecordBatch::num_rows).sum::(); + let source = new_flat_source_from_record_batches(batches); + let write_opts = WriteOptions { + row_group_size: 50, + max_file_size: Some(1), + ..Default::default() + }; + let path_provider = RegionFilePathFactory { + table_dir: "test_oversized_series".to_string(), + path_type: PathType::Bare, + }; + let mut metrics = Metrics::new(WriteType::Compaction); + let mut writer = ParquetWriter::new_with_object_store( + object_store, + metadata, + IndexConfig::default(), + NoopIndexBuilder, + path_provider, + &mut metrics, + ) + .await; + + let files = writer + .write_all_flat_as_primary_key(source, None, &write_opts) + .await + .unwrap(); + + assert_eq!(1, files.len()); + assert_eq!(total_rows, files[0].num_rows); + assert!(files[0].file_size > write_opts.max_file_size.unwrap() as u64); + } + + #[tokio::test] + async fn test_write_multiple_files_without_primary_key() { + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let metadata = Arc::new(sst_region_metadata_without_primary_key()); + let batch_rows = 1000; + let batches = vec![ + new_record_batch_without_primary_key(0, batch_rows), + new_record_batch_without_primary_key(batch_rows, 2 * batch_rows), + new_record_batch_without_primary_key(2 * batch_rows, 3 * batch_rows), + ]; + let total_rows = batches.iter().map(RecordBatch::num_rows).sum::(); + let source = new_flat_source_from_record_batches(batches); + let write_opts = WriteOptions { + row_group_size: 50, + max_file_size: Some(1), + ..Default::default() + }; + let path_provider = RegionFilePathFactory { + table_dir: "test_no_primary_key".to_string(), + path_type: PathType::Bare, + }; + let mut metrics = Metrics::new(WriteType::Compaction); + let mut writer = ParquetWriter::new_with_object_store( + object_store, + metadata, + IndexConfig::default(), + NoopIndexBuilder, + path_provider, + &mut metrics, + ) + .await; + + let files = writer + .write_all_flat(source, None, &write_opts) + .await + .unwrap(); + + // Regions without a primary key keep splitting at batch boundaries: + // the limit produces multiple files and no batch is sliced. + assert!(files.len() > 1); + assert!(files.iter().all(|file| file.num_rows % batch_rows == 0)); + assert_eq!( + total_rows, + files.iter().map(|file| file.num_rows).sum::() + ); + } + #[tokio::test] async fn test_write_read_with_index() { let mut env = TestEnv::new().await; @@ -1145,6 +1300,62 @@ mod tests { assert!(reader.next_record_batch().await.unwrap().is_none()); } + /// Creates a new region metadata without primary key for testing SSTs. + /// + /// Schema: field_0, ts + fn sst_region_metadata_without_primary_key() -> RegionMetadata { + let mut builder = RegionMetadataBuilder::new(REGION_ID); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "field_0".to_string(), + ConcreteDataType::uint64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .primary_key(vec![]); + builder.build().unwrap() + } + + /// Creates a flat format RecordBatch for regions without a primary key. + fn new_record_batch_without_primary_key(start: usize, end: usize) -> RecordBatch { + assert!(end >= start); + let metadata = Arc::new(sst_region_metadata_without_primary_key()); + let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()); + + let num_rows = end - start; + let mut pk_builder = BinaryDictionaryBuilder::::new(); + // Regions without a primary key encode it as empty bytes. + for _ in 0..num_rows { + pk_builder.append([]).unwrap(); + } + + RecordBatch::try_new( + flat_schema, + vec![ + Arc::new(UInt64Array::from_iter_values(start as u64..end as u64)) as ArrayRef, + Arc::new(TimestampMillisecondArray::from_iter_values( + start as i64..end as i64, + )) as ArrayRef, + Arc::new(pk_builder.finish()) as ArrayRef, + Arc::new(UInt64Array::from_value(1000, num_rows)) as ArrayRef, + Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)) as ArrayRef, + ], + ) + .unwrap() + } + fn new_record_batch_from_rows(rows: &[(&str, &str, i64)]) -> RecordBatch { let metadata = Arc::new(sst_region_metadata()); let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()); diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index a6a4cd5cd06..b9d5016173e 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -23,11 +23,12 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::task::{Context, Poll}; use std::time::Instant; +use bytes::Bytes; use common_telemetry::debug; use common_time::Timestamp; use datatypes::arrow::array::{ - ArrayRef, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, - TimestampSecondArray, + ArrayRef, BinaryArray, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, UInt32Array, }; use datatypes::arrow::compute::{max, min}; use datatypes::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; @@ -40,7 +41,7 @@ use parquet::file::metadata::KeyValue; use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder}; use parquet::schema::types::ColumnPath; use smallvec::smallvec; -use snafu::ResultExt; +use snafu::{OptionExt, ResultExt}; use store_api::metadata::RegionMetadataRef; use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME}; use store_api::storage::{FileId, SequenceNumber}; @@ -50,13 +51,16 @@ use tokio_util::compat::{Compat, FuturesAsyncWriteCompatExt}; use crate::access_layer::{FilePathProvider, Metrics, SstInfoArray, TempFileCleaner}; use crate::config::{IndexBuildMode, IndexConfig}; use crate::error::{ - InvalidMetadataSnafu, OpenDalSnafu, Result, UnexpectedSnafu, WriteParquetSnafu, + InvalidMetadataSnafu, InvalidRecordBatchSnafu, OpenDalSnafu, Result, UnexpectedSnafu, + WriteParquetSnafu, }; use crate::read::FlatSource; use crate::sst::file::RegionFileId; use crate::sst::index::{IndexOutput, Indexer, IndexerBuilder}; -use crate::sst::parquet::flat_format::{FlatWriteFormat, time_index_column_index}; -use crate::sst::parquet::format::PrimaryKeyWriteFormat; +use crate::sst::parquet::flat_format::{ + FlatWriteFormat, primary_key_column_index, time_index_column_index, +}; +use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, WriteOptions}; use crate::sst::{ DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, @@ -85,6 +89,71 @@ impl FlatBatchConverter { } } +/// Result of splitting a batch at the next series boundary. +enum SeriesBoundarySplit { + /// The whole batch belongs to the current series and stays in the current file. + Continue(RecordBatch), + /// The batch crosses a series boundary. + Split { + /// Remaining rows of the current series, appended to the current file. + current_file_tail: Option, + /// Rows of the following series that start the next file. + next_file_head: RecordBatch, + }, +} + +/// Returns the dictionary keys and values of the encoded `__primary_key` column. +fn encoded_primary_keys(batch: &RecordBatch) -> Result<(&UInt32Array, &BinaryArray)> { + let column = batch.column(primary_key_column_index(batch.num_columns())); + let primary_keys = column + .as_any() + .downcast_ref::() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "expected dictionary primary key column, got {:?}", + column.data_type() + ), + })?; + let values = primary_keys + .values() + .as_any() + .downcast_ref::() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "expected binary primary key values, got {:?}", + primary_keys.values().data_type() + ), + })?; + Ok((primary_keys.keys(), values)) +} + +/// Splits `batch` at the first row whose encoded primary key differs from +/// `current_primary_key`, the last primary key written to the current file. +/// +/// The writer finishes an oversized file at such a series boundary so that the +/// primary key ranges of output files never overlap, which allows pickers like +/// TWCS to detect overlapping files by their time and primary key ranges. A +/// series is never split in the middle: if no row starts a new series, the whole +/// batch stays in the current file even if it already exceeds the size limit. +fn split_at_next_series( + batch: RecordBatch, + current_primary_key: &[u8], +) -> Result { + let (keys, values) = encoded_primary_keys(&batch)?; + let Some(offset) = (0..batch.num_rows()) + .find(|&row| values.value(keys.value(row) as usize) != current_primary_key) + else { + return Ok(SeriesBoundarySplit::Continue(batch)); + }; + + let next_file_head = batch.slice(offset, batch.num_rows() - offset); + let current_file_tail = (offset > 0).then(|| batch.slice(0, offset)); + Ok(SeriesBoundarySplit::Split { + current_file_tail, + next_file_head, + }) +} + /// Parquet SST writer. pub struct ParquetWriter<'a, F: WriterFactory, I: IndexerBuilder, P: FilePathProvider> { /// Path provider that creates SST and index file paths according to file id. @@ -269,6 +338,13 @@ where /// Iterates FlatSource and writes all RecordBatch in flat format to Parquet file. /// + /// The source must yield batches globally sorted by the encoded primary key. + /// `opts.max_file_size` is a soft limit for regions with a primary key: the + /// current file is finished at the next series boundary after exceeding the + /// limit (see [split_at_next_series]), so a series larger than the limit + /// stays in one file. Regions without a primary key are split at batch + /// boundaries. + /// /// Returns the [SstInfo] if the SST is written. pub async fn write_all_flat( &mut self, @@ -304,6 +380,9 @@ where /// Iterates FlatSource and writes all RecordBatch in primary-key format to Parquet file. /// + /// The source must yield batches globally sorted by the encoded primary key. + /// See [Self::write_all_flat] for the `opts.max_file_size` semantics. + /// /// Returns the [SstInfo] if the SST is written. pub async fn write_all_flat_as_primary_key( &mut self, @@ -336,36 +415,55 @@ where let mut results = smallvec![]; let mut stats = SourceStats::default(); - while let Some(record_batch) = self - .write_next_flat_batch(&mut source, converter, opts) - .await - .transpose() - { - match record_batch { - Ok(batch) => { - stats.update_flat(&batch)?; - if matches!(self.index_config.build_mode, IndexBuildMode::Sync) { - let start = Instant::now(); - // safety: self.current_indexer must be set when first batch has been written. - self.current_indexer - .as_mut() - .unwrap() - .update_flat(&batch) - .await; - self.metrics.update_index += start.elapsed(); - } - if let Some(max_file_size) = opts.max_file_size - && self.bytes_written.load(Ordering::Relaxed) > max_file_size - { - self.finish_current_file(&mut results, &mut stats).await?; - } - } + loop { + let start = Instant::now(); + let batch = match source.next_batch().await { + Ok(Some(batch)) => batch, + Ok(None) => break, Err(e) => { - if let Some(indexer) = &mut self.current_indexer { - indexer.abort().await; - } + self.abort_current_indexer().await; return Err(e); } + }; + self.metrics.iter_source += start.elapsed(); + + if self.metadata.primary_key.is_empty() { + self.append_flat_batch(&batch, converter, opts, &mut stats) + .await?; + if self.exceeds_max_file_size(opts) { + self.finish_current_file(&mut results, &mut stats).await?; + } + } else if self.exceeds_max_file_size(opts) + && let Some(current_primary_key) = stats.last_primary_key.as_deref() + { + let series_split = match split_at_next_series(batch, current_primary_key) { + Ok(series_split) => series_split, + Err(e) => { + self.abort_current_indexer().await; + return Err(e); + } + }; + match series_split { + SeriesBoundarySplit::Continue(batch) => { + self.append_flat_batch(&batch, converter, opts, &mut stats) + .await?; + } + SeriesBoundarySplit::Split { + current_file_tail, + next_file_head, + } => { + if let Some(tail) = current_file_tail { + self.append_flat_batch(&tail, converter, opts, &mut stats) + .await?; + } + self.finish_current_file(&mut results, &mut stats).await?; + self.append_flat_batch(&next_file_head, converter, opts, &mut stats) + .await?; + } + } + } else { + self.append_flat_batch(&batch, converter, opts, &mut stats) + .await?; } } @@ -398,29 +496,53 @@ where .set_column_compression(op_type_col, Compression::UNCOMPRESSED) } - async fn write_next_flat_batch( + async fn append_flat_batch( &mut self, - source: &mut FlatSource, + batch: &RecordBatch, converter: &FlatBatchConverter, opts: &WriteOptions, - ) -> Result> { - let start = Instant::now(); - let Some(record_batch) = source.next_batch().await? else { - return Ok(None); - }; - self.metrics.iter_source += start.elapsed(); + stats: &mut SourceStats, + ) -> Result<()> { + let result = async { + let arrow_batch = converter.convert_batch(batch)?; + let start = Instant::now(); + self.maybe_init_writer(arrow_batch.schema_ref(), opts) + .await? + .write(&arrow_batch) + .await + .context(WriteParquetSnafu)?; + self.metrics.write_batch += start.elapsed(); - let arrow_batch = converter.convert_batch(&record_batch)?; + stats.update_flat(batch)?; + if matches!(self.index_config.build_mode, IndexBuildMode::Sync) { + let start = Instant::now(); + // safety: self.current_indexer must be set when first batch has been written. + self.current_indexer + .as_mut() + .unwrap() + .update_flat(batch) + .await; + self.metrics.update_index += start.elapsed(); + } + Ok(()) + } + .await; - let start = Instant::now(); - self.maybe_init_writer(arrow_batch.schema_ref(), opts) - .await? - .write(&arrow_batch) - .await - .context(WriteParquetSnafu)?; - self.metrics.write_batch += start.elapsed(); - // Return original flat batch for stats/indexer which use flat layout. - Ok(Some(record_batch)) + if result.is_err() { + self.abort_current_indexer().await; + } + result + } + + fn exceeds_max_file_size(&self, opts: &WriteOptions) -> bool { + opts.max_file_size + .is_some_and(|max_size| self.bytes_written.load(Ordering::Relaxed) >= max_size) + } + + async fn abort_current_indexer(&mut self) { + if let Some(indexer) = &mut self.current_indexer { + indexer.abort().await; + } } async fn maybe_init_writer( @@ -481,6 +603,8 @@ struct SourceStats { num_rows: usize, /// Time range of fetched batches. time_range: Option<(Timestamp, Timestamp)>, + /// Last primary key written to the current file. + last_primary_key: Option, /// Series estimator for computing num_series. series_estimator: SeriesEstimator, } @@ -493,6 +617,9 @@ impl SourceStats { self.num_rows += record_batch.num_rows(); self.series_estimator.update_flat(record_batch); + let (keys, values) = encoded_primary_keys(record_batch)?; + let key = keys.value(record_batch.num_rows() - 1); + self.last_primary_key = Some(Bytes::copy_from_slice(values.value(key as usize))); // Get the timestamp column by index let time_index_col_idx = time_index_column_index(record_batch.num_columns()); From 82444635f501113cccb56a2bf49372073120e456 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:45:35 +0000 Subject: [PATCH 095/126] feat: support quantile and fraction queries on mixed histograms (#8874) * fix: test Signed-off-by: shuiyisong * fix: remove try_build_float_literal Signed-off-by: shuiyisong * fix: remove NonCommutative Signed-off-by: shuiyisong * fix: NaN Signed-off-by: shuiyisong * chore: add comments Signed-off-by: shuiyisong * test(query): cover frontend-only histogram fold planning Signed-off-by: shuiyisong * fix(promql): ignore unparseable histogram bucket labels Signed-off-by: shuiyisong * fix: typo Signed-off-by: shuiyisong * fix: sqlness Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/common/query/src/native_histogram.rs | 123 ++- .../substrait/src/extension_serializer.rs | 2 +- src/promql/src/extension_plan.rs | 4 +- .../src/extension_plan/histogram_fold.rs | 769 ++++++++++++++---- src/promql/src/functions/native_histogram.rs | 102 ++- src/query/src/dist_plan/analyzer/test.rs | 71 +- src/query/src/dist_plan/commutativity.rs | 26 +- src/query/src/promql/planner.rs | 709 +++++++++++++--- .../distributed/flow-tql/tsid_on_phy.result | 2 +- .../promql/histogram_multi_partition.result | 2 +- .../tsid_histogram_quantile_regression.result | 2 +- 11 files changed, 1532 insertions(+), 280 deletions(-) diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index b090ac33fcd..1c7bb09868e 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -215,6 +215,15 @@ struct Bucket { boundary_rule: BoundaryRule, } +/// Additional information produced while estimating a native histogram quantile. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum NativeHistogramQuantileInfo { + /// NaN observations made the estimated quantile skew higher. + NaNSkew, + /// The requested quantile fell beyond all populated buckets because of NaN observations. + NaNResult, +} + /// Query-time representation of a Prometheus native histogram. /// /// Bucket counts are absolute `f64` values even when the persisted payload used @@ -635,19 +644,27 @@ impl NativeHistogram { /// Returns negative or positive infinity outside `[0, 1]`, and `NaN` when /// the quantile cannot be estimated. pub fn quantile(&self, q: f64) -> f64 { + self.quantile_with_info(q).0 + } + + /// Estimates the value at quantile `q` and reports the effect of NaN observations. + pub fn quantile_with_info(&self, q: f64) -> (f64, Option) { if q < 0.0 { - return f64::NEG_INFINITY; + return (f64::NEG_INFINITY, None); } if q > 1.0 { - return f64::INFINITY; + return (f64::INFINITY, None); } if self.count == 0.0 || q.is_nan() { - return f64::NAN; + return (f64::NAN, None); } let Some(mut buckets) = self.all_buckets() else { - return f64::NAN; + return (f64::NAN, None); }; + let bucket_total = buckets.iter().map(|bucket| bucket.count).sum::(); + let has_nan_observations = self.sum.is_nan() && bucket_total < self.count; + let info = has_nan_observations.then_some(NativeHistogramQuantileInfo::NaNSkew); let rank = q * self.count; let mut count = 0.0; for bucket in &mut buckets { @@ -668,42 +685,65 @@ impl NativeHistogram { } else if self.uses_custom_buckets() { if bucket.lower == f64::NEG_INFINITY { if bucket.upper <= 0.0 { - return bucket.upper; + return (bucket.upper, info); } bucket.lower = 0.0; } else if bucket.upper == f64::INFINITY { - return bucket.lower; + return (bucket.lower, info); } } let rank_in_bucket = rank - (count - bucket.count); let fraction = rank_in_bucket / bucket.count; if self.uses_custom_buckets() || (bucket.lower <= 0.0 && bucket.upper >= 0.0) { - return bucket.lower + (bucket.upper - bucket.lower) * fraction; + return ( + bucket.lower + (bucket.upper - bucket.lower) * fraction, + info, + ); } let log_lower = bucket.lower.abs().log2(); let log_upper = bucket.upper.abs().log2(); if bucket.lower > 0.0 { - return 2.0_f64.powf(log_lower + (log_upper - log_lower) * fraction); + return ( + 2.0_f64.powf(log_lower + (log_upper - log_lower) * fraction), + info, + ); } - return -2.0_f64.powf(log_upper + (log_lower - log_upper) * (1.0 - fraction)); + return ( + -2.0_f64.powf(log_upper + (log_lower - log_upper) * (1.0 - fraction)), + info, + ); } - f64::NAN + let info = if self.sum.is_nan() && count < self.count { + Some(if count < rank { + NativeHistogramQuantileInfo::NaNResult + } else { + NativeHistogramQuantileInfo::NaNSkew + }) + } else { + None + }; + (f64::NAN, info) } /// Estimates the fraction of observations between `lower` and `upper`. pub fn fraction(&self, lower: f64, upper: f64) -> f64 { + self.fraction_with_info(lower, upper).0 + } + + /// Estimates a fraction and reports whether NaN observations were excluded. + pub fn fraction_with_info(&self, lower: f64, upper: f64) -> (f64, bool) { if self.count == 0.0 || lower.is_nan() || upper.is_nan() { - return f64::NAN; + return (f64::NAN, false); } if lower >= upper { - return 0.0; + return (0.0, false); } let Some(mut buckets) = self.all_buckets() else { - return f64::NAN; + return (f64::NAN, false); }; let count = if self.sum.is_nan() { buckets.iter().map(|bucket| bucket.count).sum() @@ -759,7 +799,10 @@ impl NativeHistogram { upper_rank = count; } - (upper_rank - lower_rank) / self.count + ( + (upper_rank - lower_rank) / self.count, + self.sum.is_nan() && count < self.count, + ) } /// Converts populated buckets to `(boundary rule, lower, upper, count)` strings. @@ -2141,7 +2184,57 @@ mod tests { histogram.count = 10.0; histogram.sum = f64::NAN; - assert_eq!(histogram.fraction(f64::NEG_INFINITY, f64::INFINITY), 0.8); + assert_eq!( + histogram.fraction_with_info(f64::NEG_INFINITY, f64::INFINITY), + (0.8, true) + ); + } + + #[test] + fn quantile_reports_nan_observation_effect() { + let mut histogram = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![8.0], + ); + histogram.count = 10.0; + histogram.sum = f64::NAN; + + let (skewed, info) = histogram.quantile_with_info(0.5); + assert!(skewed.is_finite()); + assert_eq!(info, Some(NativeHistogramQuantileInfo::NaNSkew)); + + let (nan, info) = histogram.quantile_with_info(0.9); + assert!(nan.is_nan()); + assert_eq!(info, Some(NativeHistogramQuantileInfo::NaNResult)); + + histogram.count = 8.0; + assert_eq!(histogram.quantile_with_info(0.5).1, None); + + histogram.count = 3.0; + histogram.positive_spans.clear(); + histogram.positive_buckets.clear(); + let (nan, info) = histogram.quantile_with_info(0.0); + assert!(nan.is_nan()); + assert_eq!(info, Some(NativeHistogramQuantileInfo::NaNSkew)); + + histogram.schema = CUSTOM_BUCKETS_SCHEMA; + histogram.count = 10.0; + histogram.positive_spans = vec![Span { + offset: 0, + length: 1, + }]; + histogram.positive_buckets = vec![8.0]; + for (bound, offset, expected) in [(-1.0, 0, -1.0), (1.0, 1, 1.0)] { + histogram.custom_values = vec![bound]; + histogram.positive_spans[0].offset = offset; + assert_eq!( + histogram.quantile_with_info(0.5), + (expected, Some(NativeHistogramQuantileInfo::NaNSkew)) + ); + } } #[test] diff --git a/src/common/substrait/src/extension_serializer.rs b/src/common/substrait/src/extension_serializer.rs index e2e7da524c1..5e94a1a57ae 100644 --- a/src/common/substrait/src/extension_serializer.rs +++ b/src/common/substrait/src/extension_serializer.rs @@ -78,7 +78,7 @@ impl SerializerRegistry for ExtensionSerializer { .as_any() .downcast_ref::() .expect("Failed to downcast to HistogramFold"); - Ok(histogram_fold.serialize()) + histogram_fold.serialize() } name if name == UnionDistinctOn::name() => { let union_distinct_on = node diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index 973758032c2..29b31b7ca03 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -34,7 +34,9 @@ use datafusion::common::DFSchemaRef; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datatypes::data_type::DataType as _; pub use empty_metric::{EmptyMetric, EmptyMetricExec, EmptyMetricStream, build_special_time_expr}; -pub use histogram_fold::{HistogramFold, HistogramFoldExec, HistogramFoldStream}; +pub use histogram_fold::{ + HistogramFold, HistogramFoldExec, HistogramFoldOperation, HistogramFoldStream, +}; pub use instant_manipulate::{InstantManipulate, InstantManipulateExec, InstantManipulateStream}; pub use normalize::{SeriesNormalize, SeriesNormalizeExec, SeriesNormalizeStream}; pub use planner::PromExtensionPlanner; diff --git a/src/promql/src/extension_plan/histogram_fold.rs b/src/promql/src/extension_plan/histogram_fold.rs index 462f577fc67..5a5f2f8d466 100644 --- a/src/promql/src/extension_plan/histogram_fold.rs +++ b/src/promql/src/extension_plan/histogram_fold.rs @@ -33,7 +33,7 @@ use datafusion::physical_expr::{ EquivalenceProperties, LexRequirement, OrderingRequirements, PhysicalSortRequirement, }; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion::physical_plan::expressions::{CastExpr as PhyCast, Column as PhyColumn}; +use datafusion::physical_plan::expressions::{Column as PhyColumn, TryCastExpr as PhyTryCast}; use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, @@ -78,11 +78,42 @@ pub struct HistogramFold { ts_column: String, input: LogicalPlan, field_column: String, - quantile: OrderedF64, + /// Native histogram companion column for a mixed classic/native input. + histogram_column: Option, + operation: HistogramFoldOperation, output_schema: DFSchemaRef, unfix: Option, } +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd)] +pub enum HistogramFoldOperation { + Quantile(OrderedF64), + Fraction { + lower: OrderedF64, + upper: OrderedF64, + }, +} + +impl HistogramFoldOperation { + pub const fn function_name(self) -> &'static str { + match self { + Self::Quantile(_) => "histogram_quantile", + Self::Fraction { .. } => "histogram_fraction", + } + } +} + +impl std::fmt::Display for HistogramFoldOperation { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Quantile(quantile) => write!(f, "quantile={quantile}"), + Self::Fraction { lower, upper } => { + write!(f, "fraction=[{lower}, {upper}]") + } + } + } +} + #[derive(Debug, PartialEq, Eq, Hash, PartialOrd)] struct UnfixIndices { pub le_column_idx: u64, @@ -137,6 +168,21 @@ impl UserDefinedLogicalNodeCore for HistogramFold { return Some(vec![indices]); } + if let Some(histogram_column) = &self.histogram_column { + let mut necessary_indices = output_columns.to_vec(); + for column in [ + &self.le_column, + &self.ts_column, + &self.field_column, + histogram_column, + ] { + necessary_indices.push(input_schema.index_of_column_by_name(None, column)?); + } + necessary_indices.sort_unstable(); + necessary_indices.dedup(); + return Some(vec![necessary_indices]); + } + let mut necessary_indices = output_columns .iter() .map(|&output_column| { @@ -156,9 +202,13 @@ impl UserDefinedLogicalNodeCore for HistogramFold { fn fmt_for_explain(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, - "HistogramFold: le={}, field={}, quantile={}", - self.le_column, self.field_column, self.quantile - ) + "HistogramFold: le={}, field={}", + self.le_column, self.field_column + )?; + if let Some(histogram) = &self.histogram_column { + write!(f, ", histogram={histogram}")?; + } + write!(f, ", {}", self.operation) } fn with_exprs_and_inputs( @@ -194,7 +244,8 @@ impl UserDefinedLogicalNodeCore for HistogramFold { ts_column, input, field_column, - quantile: self.quantile, + histogram_column: None, + operation: self.operation, output_schema, unfix: None, }) @@ -204,7 +255,8 @@ impl UserDefinedLogicalNodeCore for HistogramFold { ts_column: self.ts_column.clone(), input, field_column: self.field_column.clone(), - quantile: self.quantile, + histogram_column: self.histogram_column.clone(), + operation: self.operation, output_schema: self.output_schema.clone(), unfix: None, }) @@ -219,16 +271,42 @@ impl HistogramFold { ts_column: String, quantile: f64, input: LogicalPlan, + ) -> DataFusionResult { + Self::new_with_operation( + le_column, + field_column, + ts_column, + HistogramFoldOperation::Quantile(quantile.into()), + None, + input, + ) + } + + pub fn new_with_operation( + le_column: String, + field_column: String, + ts_column: String, + operation: HistogramFoldOperation, + histogram_column: Option, + input: LogicalPlan, ) -> DataFusionResult { let input_schema = input.schema(); Self::check_schema(input_schema, &le_column, &field_column, &ts_column)?; - let output_schema = Self::convert_schema(input_schema, &le_column)?; + if let Some(histogram_column) = &histogram_column { + Self::check_column(input_schema, histogram_column)?; + } + let output_schema = if histogram_column.is_some() { + input_schema.clone() + } else { + Self::convert_schema(input_schema, &le_column)? + }; Ok(Self { le_column, ts_column, input, field_column, - quantile: quantile.into(), + histogram_column, + operation, output_schema, unfix: None, }) @@ -244,23 +322,22 @@ impl HistogramFold { field_column: &str, ts_column: &str, ) -> DataFusionResult<()> { - let check_column = |col| { - if !input_schema.has_column_with_unqualified_name(col) { - Err(DataFusionError::SchemaError( - Box::new(datafusion::common::SchemaError::FieldNotFound { - field: Box::new(Column::new(None::, col)), - valid_fields: input_schema.columns(), - }), - Box::new(None), - )) - } else { - Ok(()) - } - }; + Self::check_column(input_schema, le_column)?; + Self::check_column(input_schema, ts_column)?; + Self::check_column(input_schema, field_column) + } - check_column(le_column)?; - check_column(ts_column)?; - check_column(field_column) + fn check_column(input_schema: &DFSchemaRef, column: &str) -> DataFusionResult<()> { + if !input_schema.has_column_with_unqualified_name(column) { + return Err(DataFusionError::SchemaError( + Box::new(datafusion::common::SchemaError::FieldNotFound { + field: Box::new(Column::new(None::, column)), + valid_fields: input_schema.columns(), + }), + Box::new(None), + )); + } + Ok(()) } pub fn to_execution_plan(&self, exec_input: Arc) -> Arc { @@ -272,6 +349,10 @@ impl HistogramFold { let field_column_index = input_schema .index_of_column_by_name(None, &self.field_column) .unwrap(); + let histogram_column_index = self + .histogram_column + .as_ref() + .map(|column| input_schema.index_of_column_by_name(None, column).unwrap()); let ts_column_index = input_schema .index_of_column_by_name(None, &self.ts_column) .unwrap(); @@ -282,7 +363,11 @@ impl HistogramFold { .iter() .enumerate() .filter_map(|(idx, field)| { - if idx == le_column_index || idx == field_column_index || idx == ts_column_index { + if idx == le_column_index + || idx == field_column_index + || Some(idx) == histogram_column_index + || idx == ts_column_index + { None } else { Some(Arc::new(PhyColumn::new(field.name(), idx)) as _) @@ -309,11 +394,12 @@ impl HistogramFold { Arc::new(HistogramFoldExec { le_column_index, field_column_index, + histogram_column_index, ts_column_index, input: exec_input, tag_columns, partition_exprs, - quantile: self.quantile.into(), + operation: self.operation, output_schema, metric: ExecutionPlanMetricsSet::new(), properties, @@ -343,18 +429,28 @@ impl HistogramFold { )?)) } - pub fn serialize(&self) -> Vec { + pub fn serialize(&self) -> DataFusionResult> { + if self.histogram_column.is_some() { + return Err(DataFusionError::NotImplemented( + "mixed HistogramFold is frontend-only".to_string(), + )); + } + let HistogramFoldOperation::Quantile(quantile) = self.operation else { + return Err(DataFusionError::NotImplemented( + "HistogramFold fraction is frontend-only".to_string(), + )); + }; let le_column_idx = serialize_column_index(self.input.schema(), &self.le_column); let ts_column_idx = serialize_column_index(self.input.schema(), &self.ts_column); let field_column_idx = serialize_column_index(self.input.schema(), &self.field_column); - pb::HistogramFold { + Ok(pb::HistogramFold { le_column_idx, ts_column_idx, field_column_idx, - quantile: self.quantile.into(), + quantile: quantile.into(), } - .encode_to_vec() + .encode_to_vec()) } pub fn deserialize(bytes: &[u8]) -> Result { @@ -375,7 +471,8 @@ impl HistogramFold { ts_column: String::new(), input: placeholder_plan, field_column: String::new(), - quantile: pb_histogram_fold.quantile.into(), + histogram_column: None, + operation: HistogramFoldOperation::Quantile(pb_histogram_fold.quantile.into()), output_schema: Arc::new(DFSchema::empty()), unfix: Some(unfix), }) @@ -401,7 +498,11 @@ impl PartialOrd for HistogramFold { Some(core::cmp::Ordering::Equal) => {} ord => return ord, } - self.quantile.partial_cmp(&other.quantile) + match self.histogram_column.partial_cmp(&other.histogram_column) { + Some(core::cmp::Ordering::Equal) => {} + ord => return ord, + } + self.operation.partial_cmp(&other.operation) } } @@ -413,11 +514,13 @@ pub struct HistogramFoldExec { output_schema: SchemaRef, /// Index for field column in the schema of input. field_column_index: usize, + /// Index for the native histogram companion column on mixed inputs. + histogram_column_index: Option, ts_column_index: usize, /// Tag columns are all columns except `le`, `field` and `ts` columns. tag_columns: Vec>, partition_exprs: Vec>, - quantile: f64, + operation: HistogramFoldOperation, metric: ExecutionPlanMetricsSet, properties: Arc, } @@ -448,21 +551,22 @@ impl ExecutionPlan for HistogramFoldExec { )), options: None, }); - // add le ASC - cols.push(PhysicalSortRequirement { - expr: Arc::new(PhyCast::new( - Arc::new(PhyColumn::new( - self.input.schema().field(self.le_column_index).name(), - self.le_column_index, + if self.histogram_column_index.is_none() { + // add le ASC + cols.push(PhysicalSortRequirement { + expr: Arc::new(PhyTryCast::new( + Arc::new(PhyColumn::new( + self.input.schema().field(self.le_column_index).name(), + self.le_column_index, + )), + DataType::Float64, )), - DataType::Float64, - None, - )), - options: Some(SortOptions { - descending: false, // +INF in the last - nulls_first: false, // not nullable - }), - }); + options: Some(SortOptions { + descending: false, // numeric bounds ascending + nulls_first: false, // unparsable bounds last + }), + }); + } // Safety: `cols` is not empty let requirement = LexRequirement::new(cols).unwrap(); @@ -475,7 +579,7 @@ impl ExecutionPlan for HistogramFoldExec { } fn maintains_input_order(&self) -> Vec { - vec![true; self.children().len()] + vec![self.histogram_column_index.is_none(); self.children().len()] } fn children(&self) -> Vec<&Arc> { @@ -505,9 +609,10 @@ impl ExecutionPlan for HistogramFoldExec { ts_column_index: self.ts_column_index, tag_columns: self.tag_columns.clone(), partition_exprs: self.partition_exprs.clone(), - quantile: self.quantile, + operation: self.operation, output_schema: self.output_schema.clone(), field_column_index: self.field_column_index, + histogram_column_index: self.histogram_column_index, properties, })) } @@ -526,25 +631,31 @@ impl ExecutionPlan for HistogramFoldExec { let mut normal_indices = (0..input.schema().fields().len()).collect::>(); normal_indices.remove(&self.field_column_index); normal_indices.remove(&self.le_column_index); + if let Some(histogram_column_index) = self.histogram_column_index { + normal_indices.remove(&histogram_column_index); + } + let mode = if self.histogram_column_index.is_some() { + FoldMode::Safe + } else { + FoldMode::Optimistic + }; Ok(Box::pin(HistogramFoldStream { le_column_index: self.le_column_index, field_column_index: self.field_column_index, - quantile: self.quantile, + histogram_column_index: self.histogram_column_index, + operation: self.operation, normal_indices: normal_indices.into_iter().collect(), bucket_size: None, input_buffer: vec![], input, output_schema, input_schema: self.input.schema(), - mode: FoldMode::Optimistic, + mode, safe_group: None, metric: baseline_metric, batch_size, input_buffered_rows: 0, - output_buffer: HistogramFoldStream::empty_output_buffer( - &self.output_schema, - self.le_column_index, - )?, + output_buffer: HistogramFoldStream::empty_output_buffer(&self.input.schema())?, output_buffered_rows: 0, })) } @@ -574,9 +685,13 @@ impl DisplayAs for HistogramFoldExec { | DisplayFormatType::TreeRender => { write!( f, - "HistogramFoldExec: le=@{}, field=@{}, quantile={}", - self.le_column_index, self.field_column_index, self.quantile - ) + "HistogramFoldExec: le=@{}, field=@{}", + self.le_column_index, self.field_column_index + )?; + if let Some(histogram) = self.histogram_column_index { + write!(f, ", histogram=@{histogram}")?; + } + write!(f, ", {}", self.operation) } } } @@ -592,7 +707,8 @@ pub struct HistogramFoldStream { // internal states le_column_index: usize, field_column_index: usize, - quantile: f64, + histogram_column_index: Option, + operation: HistogramFoldOperation, /// Columns need not folding. This indices is based on input schema normal_indices: Vec, bucket_size: Option, @@ -619,6 +735,7 @@ struct SafeGroup { tag_values: Vec, buckets: Vec, counters: Vec, + native_samples: Vec<(Value, Value)>, } impl RecordBatchStream for HistogramFoldStream { @@ -686,25 +803,16 @@ impl HistogramFoldStream { self.maybe_take_output() } - /// Generate a group of empty [MutableVector]s from the output schema. - /// - /// For simplicity, this method will insert a placeholder for `le`. So that - /// the output buffers has the same schema with input. This placeholder needs - /// to be removed before returning the output batch. + /// Generate input-aligned output builders. Classic-only output drops `le` later. pub fn empty_output_buffer( schema: &SchemaRef, - le_column_index: usize, ) -> DataFusionResult>> { - let mut builders = Vec::with_capacity(schema.fields().len() + 1); + let mut builders = Vec::with_capacity(schema.fields().len()); for field in schema.fields() { let concrete_datatype = ConcreteDataType::try_from(field.data_type()).unwrap(); let mutable_vector = concrete_datatype.create_mutable_vector(0); builders.push(mutable_vector); } - builders.insert( - le_column_index, - ConcreteDataType::float64_datatype().create_mutable_vector(0), - ); Ok(builders) } @@ -803,7 +911,7 @@ impl HistogramFoldStream { counters.push(counter); } // ignore invalid data - let result = Self::evaluate_row(self.quantile, &bucket, &counters).unwrap_or(f64::NAN); + let result = Self::evaluate_row(self.operation, &bucket, &counters).unwrap_or(f64::NAN); self.output_buffer[self.field_column_index].push_value_ref(&ValueRef::from(result)); cursor += bucket_num; remaining_rows -= bucket_num; @@ -869,6 +977,13 @@ impl HistogramFoldStream { if !Self::is_positive_infinity(le_array, inf_index) { return false; } + if (cursor..=inf_index).any(|row| { + string_array_value_at_index(le_array, row) + .and_then(|value| value.parse::().ok()) + .is_none() + }) { + return false; + } for offset in 1..bucket_num { let row = cursor + offset; @@ -903,27 +1018,88 @@ impl HistogramFoldStream { self.output_buffered_rows += 1; } - fn finalize_safe_group(&mut self) -> DataFusionResult<()> { - if let Some(group) = self.safe_group.take() { - if group.tag_values.is_empty() { - return Ok(()); - } - - let has_inf = group - .buckets - .last() - .map(|v| v.is_infinite() && v.is_sign_positive()) - .unwrap_or(false); - let result = if group.buckets.len() < 2 || !has_inf { - f64::NAN - } else { - Self::evaluate_row(self.quantile, &group.buckets, &group.counters) - .unwrap_or(f64::NAN) - }; - let mut tag_value_refs = Vec::with_capacity(group.tag_values.len()); - tag_value_refs.extend(group.tag_values.iter().map(|v| v.as_value_ref())); - self.push_output_row(&tag_value_refs, result); + fn push_mixed_output_row( + &mut self, + tag_values: &[Value], + le: &Value, + result: Option, + histogram: &Value, + ) { + let histogram_column_index = self.histogram_column_index.unwrap(); + for (idx, value) in self.normal_indices.iter().zip(tag_values) { + self.output_buffer[*idx].push_value_ref(&value.as_value_ref()); } + self.output_buffer[self.le_column_index].push_value_ref(&le.as_value_ref()); + self.output_buffer[self.field_column_index] + .push_value_ref(&result.map_or(ValueRef::Null, ValueRef::from)); + self.output_buffer[histogram_column_index].push_value_ref(&histogram.as_value_ref()); + self.output_buffered_rows += 1; + } + + fn finalize_safe_group(&mut self) -> DataFusionResult<()> { + let Some(group) = self.safe_group.take() else { + return Ok(()); + }; + if group.tag_values.is_empty() { + return Ok(()); + } + + if self.histogram_column_index.is_some() { + let classic_result = if group.buckets.is_empty() { + None + } else { + let mut buckets = group + .buckets + .into_iter() + .zip(group.counters) + .collect::>(); + buckets.sort_by(|lhs, rhs| lhs.0.total_cmp(&rhs.0)); + let (bounds, counters): (Vec<_>, Vec<_>) = buckets.into_iter().unzip(); + let has_inf = bounds + .last() + .is_some_and(|value| value.is_infinite() && value.is_sign_positive()); + Some(if has_inf { + Self::evaluate_row(self.operation, &bounds, &counters).unwrap_or(f64::NAN) + } else { + f64::NAN + }) + }; + let has_null_native = group.native_samples.iter().any(|(le, _)| le.is_null()); + for (le, histogram) in group.native_samples.iter().filter(|(le, _)| !le.is_null()) { + self.push_mixed_output_row(&group.tag_values, le, None, histogram); + } + for (le, histogram) in group.native_samples.iter().filter(|(le, _)| le.is_null()) { + self.push_mixed_output_row(&group.tag_values, le, classic_result, histogram); + } + if !has_null_native && let Some(result) = classic_result { + self.push_mixed_output_row( + &group.tag_values, + &Value::Null, + Some(result), + &Value::Null, + ); + } + return Ok(()); + } + if group.buckets.is_empty() { + return Ok(()); + } + + let has_inf = group + .buckets + .last() + .is_some_and(|value| value.is_infinite() && value.is_sign_positive()); + let result = if has_inf { + Self::evaluate_row(self.operation, &group.buckets, &group.counters).unwrap_or(f64::NAN) + } else { + f64::NAN + }; + let tag_value_refs = group + .tag_values + .iter() + .map(Value::as_value_ref) + .collect::>(); + self.push_output_row(&tag_value_refs, result); Ok(()) } @@ -955,6 +1131,7 @@ impl HistogramFoldStream { tag_values: tag_values_buf.iter().cloned().map(Value::from).collect(), buckets: Vec::new(), counters: Vec::new(), + native_samples: Vec::new(), }); } @@ -962,17 +1139,28 @@ impl HistogramFoldStream { continue; }; + let mixed = self.histogram_column_index.is_some(); let bucket = string_array_value_at_index(le_array, row) - .and_then(|value| value.parse::().ok()) - .unwrap_or(f64::NAN); - let counter = if field_array.is_valid(row) { - field_array.value(row) - } else { - f64::NAN - }; - - group.buckets.push(bucket); - group.counters.push(counter); + .and_then(|value| value.parse::().ok()); + if let Some(bucket) = bucket + && (!mixed || field_array.is_valid(row)) + { + let counter = if field_array.is_valid(row) { + field_array.value(row) + } else { + f64::NAN + }; + group.buckets.push(bucket); + group.counters.push(counter); + } + if let Some(histogram_column_index) = self.histogram_column_index { + let histogram = vectors[histogram_column_index].get(row); + if !histogram.is_null() { + group + .native_samples + .push((vectors[self.le_column_index].get(row), histogram)); + } + } } Ok(()) @@ -998,14 +1186,15 @@ impl HistogramFoldStream { return Ok(None); } - let mut output_buf = Self::empty_output_buffer(&self.output_schema, self.le_column_index)?; + let mut output_buf = Self::empty_output_buffer(&self.input_schema)?; std::mem::swap(&mut self.output_buffer, &mut output_buf); let mut columns = Vec::with_capacity(output_buf.len()); for builder in output_buf.iter_mut() { columns.push(builder.to_vector().to_arrow_array()); } - // remove the placeholder column for `le` - columns.remove(self.le_column_index); + if self.histogram_column_index.is_none() { + columns.remove(self.le_column_index); + } self.output_buffered_rows = 0; RecordBatch::try_new(self.output_schema.clone(), columns) @@ -1040,12 +1229,18 @@ impl HistogramFoldStream { } /// Evaluate the field column and return the result - fn evaluate_row(quantile: f64, bucket: &[f64], counter: &[f64]) -> DataFusionResult { + fn evaluate_row( + operation: HistogramFoldOperation, + bucket: &[f64], + counter: &[f64], + ) -> DataFusionResult { // check bucket - if bucket.len() <= 1 { + if bucket.is_empty() + || matches!(operation, HistogramFoldOperation::Quantile(_)) && bucket.len() == 1 + { return Ok(f64::NAN); } - if bucket.last().unwrap().is_finite() { + if bucket.last() != Some(&f64::INFINITY) { return Err(DataFusionError::Execution( "last bucket should be +Inf".to_string(), )); @@ -1055,39 +1250,55 @@ impl HistogramFoldStream { "bucket and counter should have the same length".to_string(), )); } - // check quantile - if quantile < 0.0 { - return Ok(f64::NEG_INFINITY); - } else if quantile > 1.0 { - return Ok(f64::INFINITY); - } else if quantile.is_nan() { - return Ok(f64::NAN); + if let HistogramFoldOperation::Quantile(quantile) = operation { + let quantile = f64::from(quantile); + if quantile < 0.0 { + return Ok(f64::NEG_INFINITY); + } else if quantile > 1.0 { + return Ok(f64::INFINITY); + } else if quantile.is_nan() { + return Ok(f64::NAN); + } } // check input value if !bucket.windows(2).all(|w| w[0] <= w[1]) { return Ok(f64::NAN); } - let counter = { - let needs_fix = - counter.iter().any(|v| !v.is_finite()) || !counter.windows(2).all(|w| w[0] <= w[1]); - if !needs_fix { - Cow::Borrowed(counter) - } else { - let mut fixed = Vec::with_capacity(counter.len()); - let mut prev = 0.0; - for (idx, &v) in counter.iter().enumerate() { - let mut val = if v.is_finite() { v } else { prev }; - if idx > 0 && val < prev { - val = prev; + let counter = match operation { + HistogramFoldOperation::Quantile(_) => { + let needs_fix = counter.iter().any(|v| !v.is_finite()) + || !counter.windows(2).all(|w| w[0] <= w[1]); + if !needs_fix { + Cow::Borrowed(counter) + } else { + let mut fixed = Vec::with_capacity(counter.len()); + let mut prev = 0.0; + for (idx, &v) in counter.iter().enumerate() { + let mut val = if v.is_finite() { v } else { prev }; + if idx > 0 && val < prev { + val = prev; + } + fixed.push(val); + prev = val; } - fixed.push(val); - prev = val; + Cow::Owned(fixed) } - Cow::Owned(fixed) } + HistogramFoldOperation::Fraction { .. } => Cow::Borrowed(counter), }; + Ok(match operation { + HistogramFoldOperation::Quantile(quantile) => { + Self::evaluate_quantile(quantile.into(), bucket, &counter) + } + HistogramFoldOperation::Fraction { lower, upper } => { + Self::evaluate_fraction(lower.into(), upper.into(), bucket, &counter) + } + }) + } + + fn evaluate_quantile(quantile: f64, bucket: &[f64], counter: &[f64]) -> f64 { let total = *counter.last().unwrap(); let expected_pos = total * quantile; let mut fit_bucket_pos = 0; @@ -1095,7 +1306,7 @@ impl HistogramFoldStream { fit_bucket_pos += 1; } if fit_bucket_pos >= bucket.len() - 1 { - Ok(bucket[bucket.len() - 2]) + bucket[bucket.len() - 2] } else { let upper_bound = bucket[fit_bucket_pos]; let upper_count = counter[fit_bucket_pos]; @@ -1106,13 +1317,100 @@ impl HistogramFoldStream { lower_count = counter[fit_bucket_pos - 1]; } if (upper_count - lower_count).abs() < 1e-10 { - return Ok(f64::NAN); + return f64::NAN; } - Ok(lower_bound + lower_bound + (upper_bound - lower_bound) / (upper_count - lower_count) - * (expected_pos - lower_count)) + * (expected_pos - lower_count) } } + + fn evaluate_fraction(lower: f64, upper: f64, bucket: &[f64], counter: &[f64]) -> f64 { + let coalesced = bucket + .windows(2) + .any(|bounds| bounds[0] == bounds[1]) + .then(|| { + let mut bounds = Vec::with_capacity(bucket.len()); + let mut counts = Vec::with_capacity(counter.len()); + for (&bound, &count) in bucket.iter().zip(counter) { + if bounds.last() == Some(&bound) { + *counts.last_mut().unwrap() += count; + } else { + bounds.push(bound); + counts.push(count); + } + } + (bounds, counts) + }); + let (bucket, counter) = match &coalesced { + Some((bounds, counts)) => (bounds.as_slice(), counts.as_slice()), + None => (bucket, counter), + }; + let total = *counter.last().unwrap(); + if total == 0.0 || lower.is_nan() || upper.is_nan() { + return f64::NAN; + } + if lower >= upper { + return 0.0; + } + + let mut rank = 0.0; + let mut lower_rank = 0.0; + let mut upper_rank = 0.0; + let mut lower_set = false; + let mut upper_set = false; + let mut lower_bound = if bucket[0] > 0.0 { + 0.0 + } else { + f64::NEG_INFINITY + }; + + for (idx, (&upper_bound, &upper_count)) in bucket.iter().zip(counter).enumerate() { + if idx > 0 { + lower_bound = bucket[idx - 1]; + } + let interpolate = |value: f64| { + if lower_bound == f64::NEG_INFINITY { + upper_count + } else { + rank + (upper_count - rank) * (value - lower_bound) + / (upper_bound - lower_bound) + } + }; + + if !lower_set && lower_bound >= lower { + lower_rank = rank; + lower_set = true; + } + if !upper_set && lower_bound >= upper { + upper_rank = rank; + upper_set = true; + } + if lower_set && upper_set { + break; + } + if !lower_set && lower_bound < lower && upper_bound > lower { + lower_rank = interpolate(lower); + lower_set = true; + } + if !upper_set && lower_bound < upper && upper_bound > upper { + upper_rank = interpolate(upper); + upper_set = true; + } + if lower_set && upper_set { + break; + } + rank = upper_count; + } + + if !lower_set || lower_rank > total { + lower_rank = total; + } + if !upper_set || upper_rank > total { + upper_rank = total; + } + (upper_rank - lower_rank) / total + } } #[cfg(test)] @@ -1209,6 +1507,20 @@ mod test { schema: SchemaRef, quantile: f64, ts_column_index: usize, + ) -> Arc { + build_fold_exec_from_batches_with_operation( + batches, + schema, + HistogramFoldOperation::Quantile(quantile.into()), + ts_column_index, + ) + } + + fn build_fold_exec_from_batches_with_operation( + batches: Vec, + schema: SchemaRef, + operation: HistogramFoldOperation, + ts_column_index: usize, ) -> Arc { let input: Arc = Arc::new(DataSourceExec::new(Arc::new( MemorySourceConfig::try_new(&[batches], schema.clone(), None).unwrap(), @@ -1226,7 +1538,8 @@ mod test { Arc::new(HistogramFoldExec { le_column_index: 1, field_column_index: 2, - quantile, + histogram_column_index: None, + operation, ts_column_index, input, output_schema, @@ -1301,7 +1614,8 @@ mod test { let fold_exec = Arc::new(HistogramFoldExec { le_column_index: 1, field_column_index: 2, - quantile: 0.4, + histogram_column_index: None, + operation: HistogramFoldOperation::Quantile(0.4.into()), ts_column_index: 0, input: memory_exec, output_schema, @@ -1563,6 +1877,46 @@ mod test { assert_eq!(result_literal, expected); } + #[tokio::test] + async fn ignore_unparsable_bucket_bounds() { + let schema = Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, true), + Field::new("le", DataType::Utf8, true), + Field::new("val", DataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["a", "a", "a", "b"])), + Arc::new(StringArray::from(vec![ + Some("bad"), + Some("1"), + Some("+Inf"), + None, + ])), + Arc::new(Float64Array::from(vec![1.0, 2.0, 4.0, 1.0])), + ], + ) + .unwrap(); + let fold_exec = build_fold_exec_from_batches(vec![batch], schema, 0.5, 0); + + let result = + datafusion::physical_plan::collect(fold_exec, SessionContext::default().task_ctx()) + .await + .unwrap(); + + assert_eq!(result.len(), 1); + assert_eq!(result[0].num_rows(), 1); + assert_eq!( + string_array_value_at_index(result[0].column(0), 0), + Some("a") + ); + assert_eq!( + result[0].column(1).as_primitive::().value(0), + 1.0 + ); + } + #[tokio::test] async fn safe_mode_handles_misaligned_groups() { let schema = Arc::new(Schema::new(vec![ @@ -1738,8 +2092,12 @@ mod test { ]; for case in cases { - let actual = - HistogramFoldStream::evaluate_row(case.quantile, &bucket, &case.counters).unwrap(); + let actual = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(case.quantile.into()), + &bucket, + &case.counters, + ) + .unwrap(); assert_eq!( format!("{actual}"), format!("{}", case.expected), @@ -1753,7 +2111,12 @@ mod test { fn evaluate_out_of_order_input() { let bucket = [0.0, 1.0, 2.0, 3.0, 4.0, f64::INFINITY]; let counters = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]; - let result = HistogramFoldStream::evaluate_row(0.5, &bucket, &counters).unwrap(); + let result = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(0.5.into()), + &bucket, + &counters, + ) + .unwrap(); assert_eq!(0.0, result); } @@ -1761,7 +2124,11 @@ mod test { fn evaluate_wrong_bucket() { let bucket = [0.0, 1.0, 2.0, 3.0, 4.0, f64::INFINITY, 5.0]; let counters = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]; - let result = HistogramFoldStream::evaluate_row(0.5, &bucket, &counters); + let result = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(0.5.into()), + &bucket, + &counters, + ); assert!(result.is_err()); } @@ -1769,7 +2136,12 @@ mod test { fn evaluate_small_fraction() { let bucket = [0.0, 2.0, 4.0, 6.0, f64::INFINITY]; let counters = [0.0, 1.0 / 300.0, 2.0 / 300.0, 0.01, 0.01]; - let result = HistogramFoldStream::evaluate_row(0.5, &bucket, &counters).unwrap(); + let result = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(0.5.into()), + &bucket, + &counters, + ) + .unwrap(); assert_eq!(3.0, result); } @@ -1777,7 +2149,12 @@ mod test { fn evaluate_non_monotonic_counter() { let bucket = [0.0, 1.0, 2.0, 3.0, f64::INFINITY]; let counters = [0.1, 0.2, 0.4, 0.17, 0.5]; - let result = HistogramFoldStream::evaluate_row(0.5, &bucket, &counters).unwrap(); + let result = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(0.5.into()), + &bucket, + &counters, + ) + .unwrap(); assert!((result - 1.25).abs() < 1e-10, "{result}"); } @@ -1785,10 +2162,96 @@ mod test { fn evaluate_nan_counter() { let bucket = [0.0, 1.0, 2.0, 3.0, f64::INFINITY]; let counters = [f64::NAN, 1.0, 2.0, 3.0, 3.0]; - let result = HistogramFoldStream::evaluate_row(0.5, &bucket, &counters).unwrap(); + let result = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(0.5.into()), + &bucket, + &counters, + ) + .unwrap(); assert!((result - 1.5).abs() < 1e-10, "{result}"); } + #[test] + fn evaluate_classic_histogram_fraction() { + let buckets = [1.0, 2.0, f64::INFINITY]; + let counters = [2.0, 4.0, 4.0]; + let fraction = |lower, upper| { + HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Fraction { + lower: OrderedF64::from(lower), + upper: OrderedF64::from(upper), + }, + &buckets, + &counters, + ) + .unwrap() + }; + + assert_eq!(fraction(0.0, 1.0), 0.5); + assert_eq!(fraction(f64::NEG_INFINITY, f64::INFINITY), 1.0); + assert_eq!(fraction(2.0, 1.0), 0.0); + + assert_eq!( + HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Fraction { + lower: 0.0.into(), + upper: 1.0.into(), + }, + &[1.0, 1.0, f64::INFINITY], + &[1.0, 2.0, 4.0], + ) + .unwrap(), + 0.75 + ); + + assert_eq!( + HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Fraction { + lower: f64::NEG_INFINITY.into(), + upper: f64::INFINITY.into(), + }, + &[f64::INFINITY], + &[4.0], + ) + .unwrap(), + 1.0 + ); + } + + #[tokio::test] + async fn fraction_handles_single_inf_bucket_after_safe_fallback() { + let schema = Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, false), + Field::new("le", DataType::Utf8, false), + Field::new("val", DataType::Float64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["a", "a", "b"])), + Arc::new(StringArray::from(vec!["1", "+Inf", "+Inf"])), + Arc::new(Float64Array::from(vec![2.0, 4.0, 4.0])), + ], + ) + .unwrap(); + let fold = build_fold_exec_from_batches_with_operation( + vec![batch], + schema, + HistogramFoldOperation::Fraction { + lower: f64::NEG_INFINITY.into(), + upper: f64::INFINITY.into(), + }, + 0, + ); + + let batches = + datafusion::physical_plan::collect(fold, SessionContext::default().task_ctx()) + .await + .unwrap(); + let values = batches[0].column(1).as_primitive::(); + assert_eq!(values.values(), &[1.0, 1.0]); + } + fn build_empty_relation(schema: &Arc) -> LogicalPlan { LogicalPlan::EmptyRelation(EmptyRelation { produce_one_row: false, @@ -1812,8 +2275,21 @@ mod test { input_plan.clone(), ) .unwrap(); + let fraction_node = HistogramFold::new_with_operation( + "le".to_string(), + "val".to_string(), + "ts".to_string(), + HistogramFoldOperation::Fraction { + lower: 0.0.into(), + upper: 1.0.into(), + }, + None, + input_plan.clone(), + ) + .unwrap(); + assert!(fraction_node.serialize().is_err()); - let bytes = plan_node.serialize(); + let bytes = plan_node.serialize().unwrap(); let histogram_fold = HistogramFold::deserialize(&bytes).unwrap(); // need fix @@ -1824,7 +2300,10 @@ mod test { assert_eq!(histogram_fold.le_column, "le"); assert_eq!(histogram_fold.ts_column, "ts"); assert_eq!(histogram_fold.field_column, "val"); - assert_eq!(histogram_fold.quantile, OrderedF64::from(0.8)); + assert_eq!( + histogram_fold.operation, + HistogramFoldOperation::Quantile(OrderedF64::from(0.8)) + ); assert_eq!(histogram_fold.output_schema.fields().len(), 2); assert_eq!(histogram_fold.output_schema.field(0).name(), "ts"); assert_eq!(histogram_fold.output_schema.field(1).name(), "val"); diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index 18363c8f771..ffdd0062665 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -326,11 +326,17 @@ fn record_counter_reset_contradiction_warning( ); } -fn scalar_histogram_udf( +fn scalar_histogram_udf( name: &'static str, extra_input_types: Vec, - calc: fn(&NativeHistogram, &[ColumnarValue], usize, usize, &'static str) -> DfResult, -) -> ScalarUDF { + calc: F, +) -> ScalarUDF +where + F: Fn(&NativeHistogram, &[ColumnarValue], usize, usize, &'static str) -> DfResult + + Send + + Sync + + 'static, +{ let mut input_types = vec![native_histogram_arrow_type()]; input_types.extend(extra_input_types); create_udf( @@ -795,11 +801,28 @@ impl NativeHistogramQuantile { } pub fn scalar_udf() -> ScalarUDF { + Self::scalar_udf_with_collector(None) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { scalar_histogram_udf( Self::name(), vec![DataType::Float64], - |histogram, input, row, len, name| { - Ok(histogram.quantile(read_scalar_f64_arg(&input[1], row, len, name)?)) + move |histogram, input, row, len, name| { + let q = read_scalar_f64_arg(&input[1], row, len, name)?; + let (value, info) = histogram.quantile_with_info(q); + if let Some(info) = info { + let message = match info { + NativeHistogramQuantileInfo::NaNSkew => { + "input to histogram_quantile has NaN observations, result is skewed higher" + } + NativeHistogramQuantileInfo::NaNResult => { + "input to histogram_quantile has NaN observations, result is NaN" + } + }; + record_info(&collector, message); + } + Ok(value) }, ) } @@ -813,13 +836,24 @@ impl NativeHistogramFraction { } pub fn scalar_udf() -> ScalarUDF { + Self::scalar_udf_with_collector(None) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { scalar_histogram_udf( Self::name(), vec![DataType::Float64, DataType::Float64], - |histogram, input, row, len, name| { + move |histogram, input, row, len, name| { let lower = read_scalar_f64_arg(&input[1], row, len, name)?; let upper = read_scalar_f64_arg(&input[2], row, len, name)?; - Ok(histogram.fraction(lower, upper)) + let (value, excluded_nans) = histogram.fraction_with_info(lower, upper); + if excluded_nans { + record_info( + &collector, + "input to histogram_fraction has NaN observations, which are excluded from all fractions", + ); + } + Ok(value) }, ) } @@ -3076,6 +3110,60 @@ mod tests { infos } + #[test] + fn quantile_and_fraction_report_nan_observations() { + let histogram = sample_histogram(10.0, f64::NAN, vec![8.0]); + let histogram_arg = + || ColumnarValue::Array(build_histogram_array(&[Some(histogram.clone())])); + + let quantile_collector = PromqlAnnotationCollector::default(); + let skewed = run_scalar_udf( + NativeHistogramQuantile::scalar_udf_with_collector(Some(quantile_collector.clone())), + vec![ + histogram_arg(), + ColumnarValue::Scalar(ScalarValue::Float64(Some(0.5))), + ], + ); + assert!(skewed.is_finite()); + let nan = run_scalar_udf( + NativeHistogramQuantile::scalar_udf_with_collector(Some(quantile_collector.clone())), + vec![ + histogram_arg(), + ColumnarValue::Scalar(ScalarValue::Float64(Some(0.9))), + ], + ); + assert!(nan.is_nan()); + let infos = collected_infos(&quantile_collector); + assert!( + infos + .iter() + .any(|info| info.ends_with("result is skewed higher")) + ); + assert!(infos.iter().any(|info| info.ends_with("result is NaN"))); + + let fraction_collector = PromqlAnnotationCollector::default(); + assert_eq!( + run_scalar_udf( + NativeHistogramFraction::scalar_udf_with_collector(Some( + fraction_collector.clone(), + )), + vec![ + histogram_arg(), + ColumnarValue::Scalar(ScalarValue::Float64(Some(f64::NEG_INFINITY))), + ColumnarValue::Scalar(ScalarValue::Float64(Some(f64::INFINITY))), + ], + ), + 0.8 + ); + assert_eq!( + collected_infos(&fraction_collector), + vec![ + "input to histogram_fraction has NaN observations, which are excluded from all fractions" + .to_string() + ] + ); + } + #[test] fn mixed_ranges_follow_prometheus_sample_type_semantics() { let first = sample_histogram(1.0, 1.0, vec![1.0]); diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 83ccc273791..c42eb0e285a 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -33,8 +33,8 @@ use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ExprSchema, JoinType, ScalarValue}; use datafusion_expr::expr::{Exists, ScalarFunction}; use datafusion_expr::{ - AggregateUDF, Expr, ExprSchemable as _, LogicalPlanBuilder, Operator, Subquery, binary_expr, - col, lit, + AggregateUDF, Expr, ExprSchemable as _, Extension, LogicalPlanBuilder, Operator, Subquery, + binary_expr, col, lit, }; use datafusion_functions::datetime::date_bin; use datafusion_functions::datetime::expr_fn::now; @@ -44,6 +44,7 @@ use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef}; use futures::Stream; use futures::task::{Context, Poll}; use pretty_assertions::assert_eq; +use promql::extension_plan::{HistogramFold, HistogramFoldOperation}; use regex::Regex; use store_api::data_source::DataSource; use store_api::storage::ScanRequest; @@ -300,6 +301,72 @@ fn find_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { plan.inputs().into_iter().find_map(find_merge_scan) } +#[test] +fn frontend_only_histogram_folds_stay_above_merge_scan() { + let table = TestTable::table_with_name(0, "t".to_string()); + let table_source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(table), + ))); + let input = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) + .unwrap() + .build() + .unwrap(); + + for (name, operation, histogram_column) in [ + ( + "mixed", + HistogramFoldOperation::Quantile(0.5.into()), + Some("pk2".to_string()), + ), + ( + "fraction", + HistogramFoldOperation::Fraction { + lower: 0.0.into(), + upper: 1.0.into(), + }, + None, + ), + ] { + let fold = HistogramFold::new_with_operation( + "pk1".to_string(), + "number".to_string(), + "ts".to_string(), + operation, + histogram_column, + input.clone(), + ) + .unwrap(); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(fold), + }); + assert!( + DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .is_err(), + "{name}" + ); + + let result = DistPlannerAnalyzer {} + .analyze(plan, &ConfigOptions::default()) + .unwrap(); + let result_text = result.to_string(); + assert!( + result_text.contains("HistogramFold:"), + "{name}\n{result_text}" + ); + + let remote_input = find_merge_scan(&result).unwrap().input(); + let remote_text = remote_input.to_string(); + assert!( + !remote_text.contains("HistogramFold:"), + "{name}\n{remote_text}" + ); + DFLogicalSubstraitConvertor + .encode(remote_input, DefaultSerializer) + .unwrap(); + } +} + #[test] fn dictionary_literals_are_unwrapped_before_pushdown() { let table = Arc::new(Table::new( diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index b380fbba05e..1a70fd8e1d9 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -315,7 +315,10 @@ impl Categorizer { | Expr::Exists(_) | Expr::InList(_) | Expr::Case(_) => Commutativity::Commutative, - Expr::ScalarFunction(_udf) => Commutativity::Commutative, + // Annotation collection must not affect distribution. Keep scalar UDFs pushdownable; + // TODO: Preserve datanode PromQL annotations through plan decoding, return them in + // Flight stream metadata, and merge them on the frontend. + Expr::ScalarFunction(_) => Commutativity::Commutative, Expr::AggregateFunction(_udaf) => Commutativity::Commutative, Expr::Like(_) @@ -376,6 +379,11 @@ mod tests { use datafusion_common::Column; use datafusion_expr::LogicalPlanBuilder; + use datafusion_expr::expr::ScalarFunction; + use datafusion_functions::core::coalesce; + use promql::functions::{ + NativeHistogramDrop, NativeHistogramFraction, NativeHistogramQuantile, + }; use super::*; @@ -396,6 +404,22 @@ mod tests { let commutativity = Categorizer::check_extension_plan(&series_divide, &partition_cols); assert!(matches!(commutativity, Commutativity::Commutative)); } + + #[test] + fn annotated_histogram_helpers_do_not_block_pushdown() { + for udf in [ + NativeHistogramQuantile::scalar_udf(), + NativeHistogramFraction::scalar_udf(), + NativeHistogramDrop::warning_bool_false_udf(String::new(), None), + ] { + let helper = Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(udf), vec![])); + let expr = Expr::ScalarFunction(ScalarFunction::new_udf(coalesce(), vec![helper])); + assert!(matches!( + Categorizer::check_expr(&expr), + Commutativity::Commutative + )); + } + } } pub type Transformer = Arc Option>; diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 5dfb1a537a8..de0181af7f6 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -41,6 +41,7 @@ use datafusion::logical_expr::{ BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator, ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition, }; +use datafusion::optimizer::simplify_expressions::ExprSimplifier; use datafusion::prelude as df_prelude; use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; use datafusion::scalar::ScalarValue; @@ -49,6 +50,7 @@ use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; use datafusion_common::{DFSchema, NullEquality}; use datafusion_expr::expr::WindowFunctionParams; use datafusion_expr::expr_fn::when; +use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{conjunction, disjunction}; use datafusion_expr::{ ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit, @@ -59,8 +61,9 @@ use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType}; use itertools::Itertools; use once_cell::sync::Lazy; use promql::extension_plan::{ - Absent, EmptyMetric, HistogramFold, InstantManipulate, Millisecond, RangeManipulate, - ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, build_special_time_expr, + Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, InstantManipulate, Millisecond, + RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, + build_special_time_expr, }; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, @@ -118,6 +121,8 @@ const SCALAR_FUNCTION: &str = "scalar"; const SPECIAL_ABSENT_FUNCTION: &str = "absent"; /// `histogram_quantile` function in PromQL const SPECIAL_HISTOGRAM_QUANTILE: &str = "histogram_quantile"; +/// `histogram_fraction` function in PromQL +const SPECIAL_HISTOGRAM_FRACTION: &str = "histogram_fraction"; /// `vector` function in PromQL const SPECIAL_VECTOR_FUNCTION: &str = "vector"; /// `le` column for conventional histogram. @@ -856,7 +861,11 @@ impl PromPlanner { .context(DataFusionPlanningSnafu); } - let val = Self::get_param_as_literal_expr(param, Some(*op), Some(ArrowDataType::Float64))?; + let val = Self::get_param_as_literal_expr( + param.as_deref(), + Some(*op), + Some(ArrowDataType::Float64), + )?; // convert op and value columns to window exprs. let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?; @@ -2074,8 +2083,10 @@ impl PromPlanner { let Call { func, args } = call_expr; // some special functions that are not expression but a plan match func.name { - SPECIAL_HISTOGRAM_QUANTILE => { - return self.create_histogram_plan(args, query_engine_state).await; + SPECIAL_HISTOGRAM_QUANTILE | SPECIAL_HISTOGRAM_FRACTION => { + return self + .create_histogram_plan(func.name, args, query_engine_state) + .await; } SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await, SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await, @@ -3221,7 +3232,7 @@ impl PromPlanner { } _ => { - let expr = Self::get_param_as_literal_expr(&Some(arg.clone()), None, None)?; + let expr = Self::get_param_as_literal_expr(Some(arg.as_ref()), None, None)?; result.literals.push(expr); } } @@ -3563,9 +3574,6 @@ impl PromPlanner { "histogram_stdvar" => { ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf())) } - "histogram_fraction" => { - ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramFraction::scalar_udf())) - } "time" => { exprs.push(build_special_time_expr( self.ctx.time_index_column.as_ref().unwrap(), @@ -4304,8 +4312,11 @@ impl PromPlanner { let expr = match op.id() { token::T_SUM => sum_udaf().call(vec![input]), token::T_QUANTILE => { - let q = - Self::get_param_as_literal_expr(param, Some(op), Some(ArrowDataType::Float64))?; + let q = Self::get_param_as_literal_expr( + param.as_deref(), + Some(op), + Some(ArrowDataType::Float64), + )?; quantile_udaf().call(vec![q, input]) } token::T_AVG => avg_udaf().call(vec![input]), @@ -4580,11 +4591,11 @@ impl PromPlanner { } fn get_param_as_literal_expr( - param: &Option>, + param: Option<&PromExpr>, op: Option, expected_type: Option, ) -> Result { - let prom_param = param.as_deref().with_context(|| { + let prom_param = param.with_context(|| { if let Some(op) = op { FunctionInvalidArgumentSnafu { fn_name: op.to_string(), @@ -4682,63 +4693,80 @@ impl PromPlanner { Ok(normalized_exprs) } - /// Try to build a [f64] from [PromExpr]. - #[deprecated( - note = "use `Self::get_param_as_literal_expr` instead. This is only for `create_histogram_plan`" - )] - fn try_build_float_literal(expr: &PromExpr) -> Option { - match expr { - PromExpr::NumberLiteral(NumberLiteral { val }) => Some(*val), - PromExpr::Paren(ParenExpr { expr }) => Self::try_build_float_literal(expr), - PromExpr::Unary(UnaryExpr { expr, .. }) => { - Self::try_build_float_literal(expr).map(|f| -f) - } - PromExpr::StringLiteral(_) - | PromExpr::Binary(_) - | PromExpr::VectorSelector(_) - | PromExpr::MatrixSelector(_) - | PromExpr::Call(_) - | PromExpr::Extension(_) - | PromExpr::Aggregate(_) - | PromExpr::Subquery(_) => None, - } - } - - /// Create a [SPECIAL_HISTOGRAM_QUANTILE] plan. + /// Create a classic, native, or mixed histogram helper plan. async fn create_histogram_plan( &mut self, + function_name: &str, args: &PromFunctionArgs, query_engine_state: &QueryEngineState, ) -> Result { - if args.args.len() != 2 { - return FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(), + let float_literal = |param: &PromExpr| -> Result { + let value = (|| { + let expr = Self::get_param_as_literal_expr( + Some(param), + None, + Some(ArrowDataType::Float64), + ) + .ok()?; + let simplifier = ExprSimplifier::new(SimplifyContext::default()); + let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?; + let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else { + return None; + }; + let ScalarValue::Float64(Some(value)) = + value.cast_to(&ArrowDataType::Float64).ok()? + else { + return None; + }; + Some(value) + })() + .with_context(|| FunctionInvalidArgumentSnafu { + fn_name: function_name.to_string(), + })?; + Ok(value) + }; + let (function, input) = match (function_name, args.args.as_slice()) { + (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => ( + HistogramFoldOperation::Quantile(float_literal(quantile)?.into()), + input.as_ref().clone(), + ), + (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => ( + HistogramFoldOperation::Fraction { + lower: float_literal(lower)?.into(), + upper: float_literal(upper)?.into(), + }, + input.as_ref().clone(), + ), + _ => { + return FunctionInvalidArgumentSnafu { + fn_name: function_name.to_string(), + } + .fail(); } - .fail(); - } - #[allow(deprecated)] - let phi = Self::try_build_float_literal(&args.args[0]).with_context(|| { - FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(), - } - })?; + }; - let input = args.args[1].as_ref().clone(); let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?; - // `histogram_quantile` folds buckets across `le`, so `__tsid` (which includes `le`) is not - // a stable series identifier anymore. Also, HistogramFold infers label columns from the - // input schema and must not treat `__tsid` as a label column. + // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a + // stable series identifier anymore. HistogramFold must not treat it as a label column. let input_plan = self.strip_tsid_column(input_plan)?; self.ctx.use_tsid = false; - if let Some(histogram_field) = + if let Some((float_field, histogram_field)) = Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns) - .map(|(_, histogram)| histogram.to_string()) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())) { + if self.ctx.has_le_tag() { + return self.create_mixed_histogram_plan( + function, + input_plan, + float_field, + histogram_field, + ); + } self.ctx.field_columns = vec![histogram_field]; } if self.all_field_columns_are_native_histograms(input_plan.schema()) { - return self.create_native_histogram_quantile_plan(phi, input_plan); + return self.create_native_histogram_plan(function, input_plan); } if !self.ctx.has_le_tag() { @@ -4764,51 +4792,68 @@ impl PromPlanner { .field_columns .first() .with_context(|| FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(), + fn_name: function.function_name().to_string(), })? .clone(); // remove le column from tag columns self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME); + let fold = HistogramFold::new_with_operation( + LE_COLUMN_NAME.to_string(), + field_column, + time_index_column, + function, + None, + input_plan, + ) + .context(DataFusionPlanningSnafu)?; Ok(LogicalPlan::Extension(Extension { - node: Arc::new( - HistogramFold::new( - LE_COLUMN_NAME.to_string(), - field_column, - time_index_column, - phi, - input_plan, - ) - .context(DataFusionPlanningSnafu)?, - ), + node: Arc::new(fold), })) } - fn create_native_histogram_quantile_plan( + fn create_native_histogram_expr( + &self, + function: HistogramFoldOperation, + field_column: &str, + ) -> DfExpr { + let field = DfExpr::Column(Column::from_name(field_column)); + let (func, args) = match function { + HistogramFoldOperation::Quantile(quantile) => ( + Arc::new(NativeHistogramQuantile::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + vec![field, lit(f64::from(quantile))], + ), + HistogramFoldOperation::Fraction { lower, upper } => ( + Arc::new(NativeHistogramFraction::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + vec![field, lit(f64::from(lower)), lit(f64::from(upper))], + ), + }; + DfExpr::ScalarFunction(ScalarFunction { func, args }) + } + + fn create_native_histogram_plan( &mut self, - phi: f64, + function: HistogramFoldOperation, input_plan: LogicalPlan, ) -> Result { ensure!( self.ctx.field_columns.len() == 1, MultiFieldsNotSupportedSnafu { - operator: SPECIAL_HISTOGRAM_QUANTILE + operator: function.function_name() }, ); let field_column = self.ctx.field_columns[0].clone(); - let quantile_expr = DfExpr::ScalarFunction(ScalarFunction { - func: Arc::new(NativeHistogramQuantile::scalar_udf()), - args: vec![ - DfExpr::Column(Column::from_name(field_column)), - DfExpr::Literal(ScalarValue::Float64(Some(phi)), None), - ], - }); - let display_name = quantile_expr.schema_name().to_string(); + let function_expr = self.create_native_histogram_expr(function, &field_column); + let display_name = function_expr.schema_name().to_string(); self.ctx.field_columns = vec![display_name.clone()]; let project_exprs = std::iter::once(self.create_time_index_column_expr()?) - .chain(std::iter::once(quantile_expr.alias(display_name))) + .chain(std::iter::once(function_expr.alias(display_name))) .chain(self.create_tag_column_exprs()?) .collect::>(); @@ -4821,6 +4866,68 @@ impl PromPlanner { .context(DataFusionPlanningSnafu) } + fn create_mixed_histogram_plan( + &mut self, + function: HistogramFoldOperation, + input_plan: LogicalPlan, + float_field: String, + histogram_field: String, + ) -> Result { + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let tag_columns = self.ctx.tag_columns.clone(); + let folded = HistogramFold::new_with_operation( + LE_COLUMN_NAME.to_string(), + float_field.clone(), + time_index_column.clone(), + function, + Some(histogram_field.clone()), + input_plan, + ) + .context(DataFusionPlanningSnafu)?; + let record_collision = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::warning_bool_false_udf( + "vector contains a mix of classic and native histograms".to_string(), + self.promql_annotations.clone(), + )), + args: vec![col(&float_field), col(&histogram_field)], + }); + let keep = when( + col(&float_field) + .is_not_null() + .and(col(&histogram_field).is_not_null()), + record_collision, + ) + .otherwise(lit(true)) + .context(DataFusionPlanningSnafu)?; + + let native_expr = self.create_native_histogram_expr(function, &histogram_field); + let output_field = native_expr.schema_name().to_string(); + let value = DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: vec![col(&float_field), native_expr], + }); + self.ctx.field_columns = vec![output_field.clone()]; + LogicalPlanBuilder::from(LogicalPlan::Extension(Extension { + node: Arc::new(folded), + })) + .filter(keep) + .context(DataFusionPlanningSnafu)? + .project( + std::iter::once(col(&time_index_column)) + .chain(std::iter::once(value.alias(output_field))) + .chain(tag_columns.iter().map(col)), + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + /// Create a [SPECIAL_VECTOR_FUNCTION] plan async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result { if args.args.len() != 1 { @@ -4829,7 +4936,7 @@ impl PromPlanner { } .fail(); } - let lit = Self::get_param_as_literal_expr(&Some(args.args[0].clone()), None, None)?; + let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?; // reuse `SPECIAL_TIME_FUNCTION` as name of time index column self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); @@ -5019,8 +5126,9 @@ impl PromPlanner { } } PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr), - // TODO(ruihang): support Unary operator - PromExpr::Unary(UnaryExpr { expr, .. }) => Self::try_build_literal_expr(expr), + PromExpr::Unary(UnaryExpr { expr, .. }) => Some(DfExpr::Negative(Box::new( + Self::try_build_literal_expr(expr)?, + ))), PromExpr::Binary(PromBinaryExpr { lhs, rhs, @@ -6904,6 +7012,7 @@ mod test { name: &str, table_id: u32, tag: &str, + le: Option<&str>, value: DirectOrValue, ) -> table::TableRef { let value_type = match &value { @@ -6912,12 +7021,20 @@ mod test { DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(), DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(), }; - let schema = Arc::new(Schema::new(vec![ - ColumnSchema::new( - "tag".to_string(), + let tag_count = 1 + usize::from(le.is_some()); + let mut columns = vec![ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + )]; + if le.is_some() { + columns.push(ColumnSchema::new( + LE_COLUMN_NAME.to_string(), ConcreteDataType::string_datatype(), false, - ), + )); + } + columns.extend([ ColumnSchema::new( "ts".to_string(), ConcreteDataType::timestamp_millisecond_datatype(), @@ -6925,16 +7042,17 @@ mod test { ) .with_time_index(true), ColumnSchema::new("v".to_string(), value_type, true), - ])); - let batch = RecordBatch::try_new( - schema.arrow_schema().clone(), - vec![ - Arc::new(StringArray::from(vec![tag])) as Arc, - Arc::new(TimestampMillisecondArray::from(vec![1_000])), - value.array(), - ], - ) - .unwrap(); + ]); + let schema = Arc::new(Schema::new(columns)); + let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc]; + if let Some(le) = le { + arrays.push(Arc::new(StringArray::from(vec![le]))); + } + arrays.extend([ + Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc, + value.array(), + ]); + let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); let backing = GreptimeMemTable::new_with_catalog( name, GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), @@ -6942,11 +7060,12 @@ mod test { DEFAULT_CATALOG_NAME.to_string(), DEFAULT_SCHEMA_NAME.to_string(), ); + let value_index = tag_count + 1; let meta = TableMetaBuilder::empty() .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) - .next_column_id(3) + .primary_key_indices((0..tag_count).collect()) + .value_indices(vec![value_index]) + .next_column_id((value_index + 1) as u32) .build() .unwrap(); let info = Arc::new( @@ -6967,21 +7086,37 @@ mod test { fn operator_table_provider() -> DfTableSourceProvider { let catalog = MemoryCatalogManager::with_default_setup(); let tables = [ - operator_metric_table("lf", 2_001, "a", DirectOrValue::Float64(2.0)), + operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)), operator_metric_table( "lh", 2_002, "b", + None, DirectOrValue::NativeHistogram(direct_or_histogram()), ), - operator_metric_table("rf", 2_003, "b", DirectOrValue::Float64(3.0)), + operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)), operator_metric_table( "rh", 2_004, "a", + None, + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)), + operator_metric_table( + "bad_classic", + 2_006, + "d", + Some("broken"), + DirectOrValue::Float64(1.0), + ), + operator_metric_table( + "bad_native", + 2_007, + "d", + None, DirectOrValue::NativeHistogram(direct_or_histogram()), ), - operator_metric_table("fallback", 2_005, "c", DirectOrValue::Float64(7.0)), ]; for table in tables { let info = table.table_info(); @@ -7621,6 +7756,11 @@ mod test { ConcreteDataType::string_datatype(), false, ), + ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + true, + ), ColumnSchema::new( "timestamp".to_string(), ConcreteDataType::timestamp_millisecond_datatype(), @@ -7636,8 +7776,8 @@ mod test { let schema = Arc::new(Schema::new(columns)); let table_meta = TableMetaBuilder::empty() .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) + .primary_key_indices(vec![0, 1]) + .value_indices(vec![3]) .next_column_id(1024) .build() .unwrap(); @@ -7816,6 +7956,122 @@ mod test { ) } + fn classic_and_native_histogram_table_provider( + native_tag: &str, + native_le: Option<&str>, + native_histogram: NativeHistogram, + ) -> DfTableSourceProvider { + let table_name = "mixed_histogram"; + let catalog = MemoryCatalogManager::with_default_setup(); + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + true, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ])); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![0, 1]) + .value_indices(vec![3, 4]) + .next_column_id(5) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from(vec![ + "classic", "classic", native_tag, "classic", "classic", native_tag, + ])), + Arc::new(StringArray::from(vec![ + Some("1"), + Some("+Inf"), + native_le, + Some("1"), + Some("+Inf"), + native_le, + ])), + Arc::new(TimestampMillisecondArray::from(vec![ + 1_000, 1_000, 1_000, 2_000, 2_000, 2_000, + ])), + build_histogram_array(&[ + None, + None, + Some(native_histogram.clone()), + None, + None, + Some(native_histogram), + ]), + Arc::new(Float64Array::from(vec![ + Some(2.0), + Some(4.0), + None, + Some(2.0), + Some(4.0), + None, + ])), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + table_name, + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 2_200, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 2_200, + table, + }) + .unwrap(); + + DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) + } + async fn build_test_table_provider_with_tsid( table_name_tuples: &[(String, String)], num_tag: usize, @@ -10091,7 +10347,7 @@ mod test { let plan = native_histogram_plan("histogram_count(some_metric)").await; assert!(plan.contains("prom_native_histogram_count"), "{plan}"); - assert!(!plan.contains("PromHistogramFold"), "{plan}"); + assert!(!plan.contains("HistogramFold:"), "{plan}"); } #[tokio::test] @@ -10102,6 +10358,7 @@ mod test { "stale_histogram", 2_100, "a", + None, DirectOrValue::NativeHistogram(stale), ); let catalog = MemoryCatalogManager::with_default_setup(); @@ -10211,7 +10468,8 @@ mod test { let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await; assert!(plan.contains("prom_native_histogram_quantile"), "{plan}"); - assert!(!plan.contains("PromHistogramFold"), "{plan}"); + assert!(!plan.contains("HistogramFold:"), "{plan}"); + assert!(plan.contains("some_metric.le"), "{plan}"); // The phi literal is threaded into the native quantile UDF as its second argument. assert!(plan.contains("Float64(0.9)"), "{plan}"); // The empty-values filter drops NULL quantile results so the output is empty @@ -10239,6 +10497,213 @@ mod test { assert!(!plan.contains("EmptyRelation"), "{plan}"); } + #[tokio::test] + async fn mixed_histogram_helpers_execute_classic_and_native_samples() { + let state = build_query_engine_state(); + for (query, expected) in [ + ( + "histogram_quantile(0.5, mixed_histogram)", + vec![("classic", 1.0), ("native", 0.0)], + ), + ( + "histogram_fraction(-Inf, +Inf, mixed_histogram)", + vec![("classic", 1.0), ("native", 1.0)], + ), + ] { + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider("native", None, direct_or_histogram()), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!(plan_text.contains("HistogramFold:"), "{plan_text}"); + assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}"); + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + let mut actual = batches + .iter() + .flat_map(|batch| { + let tags = batch + .column_by_name("tag") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name(&value_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row))) + }) + .collect::>(); + actual.sort_by_key(|(tag, _)| *tag); + assert_eq!(actual, expected, "{query}"); + } + } + + #[tokio::test] + async fn mixed_histogram_helpers_report_annotations() { + let state = build_query_engine_state(); + let mut native_histogram = direct_or_histogram(); + native_histogram.count = 2.0; + native_histogram.sum = f64::NAN; + for (native_tag, expected_rows, expected_warnings, expected_infos) in [ + ( + "classic", + 0, + vec!["vector contains a mix of classic and native histograms"], + vec![], + ), + ( + "native", + 2, + vec![], + vec!["input to histogram_quantile has NaN observations, result is skewed higher"], + ), + ] { + let collector = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + classic_and_native_histogram_table_provider( + native_tag, + None, + native_histogram.clone(), + ), + &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected_rows + ); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert_eq!(warnings, expected_warnings); + assert_eq!(infos, expected_infos); + } + } + + #[tokio::test] + async fn mixed_histogram_helper_preserves_native_le_and_scans_once() { + let state = build_query_engine_state(); + let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"); + stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap(); + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider( + "classic", + Some("native"), + direct_or_histogram(), + ), + &stmt, + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert_eq!( + plan_text.matches("TableScan: mixed_histogram").count(), + 1, + "{plan_text}" + ); + + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let (_, batches) = execute(plan, &state).await; + let mut actual = batches + .iter() + .flat_map(|batch| { + let le = batch + .column_by_name(LE_COLUMN_NAME) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let timestamps = batch + .column_by_name("timestamp") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name(&value_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(|row| { + ( + timestamps.value(row), + (!le.is_null(row)).then(|| le.value(row).to_string()), + values.value(row), + ) + }) + }) + .collect::>(); + actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1))); + assert_eq!( + actual, + vec![ + (1_000, None, 1.0), + (1_000, Some("native".to_string()), 0.0), + (2_000, None, 1.0), + (2_000, Some("native".to_string()), 0.0), + ] + ); + } + + #[tokio::test] + async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() { + let state = build_query_engine_state(); + for native_le in [None, Some("native")] { + for query in [ + "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))", + "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))", + ] { + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider( + "native", + native_le, + direct_or_histogram(), + ), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "native_le={native_le:?}, query={query}" + ); + } + } + } + #[tokio::test] async fn native_histogram_quantile_rejects_multi_field_input() { let table_provider = build_test_multi_histogram_table_provider("some_metric").await; @@ -10461,7 +10926,7 @@ mod test { "prom_native_histogram_stdvar", ), ( - "histogram_fraction(0, 1, some_metric)", + "histogram_fraction(-2 + 1, 2 / 2, some_metric)", "prom_native_histogram_fraction", ), ]; @@ -10469,6 +10934,9 @@ mod test { for (query, expected_udf) in cases { let plan = native_histogram_plan(query).await; assert!(plan.contains(expected_udf), "{query}\n{plan}"); + if query.starts_with("histogram_fraction") { + assert!(plan.contains("Float64(-1)"), "{query}\n{plan}"); + } } } @@ -12319,6 +12787,37 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom } } + #[tokio::test] + async fn malformed_classic_bucket_does_not_drop_native_histogram() { + let state = build_query_engine_state(); + let collector = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(values(&batches, &value_field), vec![0.0]); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert!(infos.is_empty()); + } + #[tokio::test] async fn test_mixed_binary_operator_aligns_both_alternative_inputs() { let state = build_query_engine_state(); diff --git a/tests/cases/distributed/flow-tql/tsid_on_phy.result b/tests/cases/distributed/flow-tql/tsid_on_phy.result index 68dd01374c1..0a1755b5c2b 100644 --- a/tests/cases/distributed/flow-tql/tsid_on_phy.result +++ b/tests/cases/distributed/flow-tql/tsid_on_phy.result @@ -120,7 +120,7 @@ TQL EXPLAIN ( | | ]] | | physical_plan | HistogramFoldExec: le=@0, field=@4, quantile=0.5 | | | RepartitionExec: REDACTED -| | SortExec: expr=[tag4@1 ASC NULLS LAST, tag5@2 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] | +| | SortExec: expr=[tag4@1 ASC NULLS LAST, tag5@2 ASC NULLS LAST, ts@3 ASC NULLS LAST, TRY_CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] | | | MergeScanExec: REDACTED | | | +---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ diff --git a/tests/cases/standalone/common/promql/histogram_multi_partition.result b/tests/cases/standalone/common/promql/histogram_multi_partition.result index 93168ce449f..3fba2ba8465 100644 --- a/tests/cases/standalone/common/promql/histogram_multi_partition.result +++ b/tests/cases/standalone/common/promql/histogram_multi_partition.result @@ -42,7 +42,7 @@ tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bu | stage | node | plan_| +-+-+-+ | 0_| 0_|_HistogramFoldExec: le=@0, field=@2, quantile=0.5 REDACTED -|_|_|_SortExec: expr=[ts@1 ASC NULLS LAST, CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[true] REDACTED +|_|_|_SortExec: expr=[ts@1 ASC NULLS LAST, TRY_CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[true] REDACTED |_|_|_RepartitionExec: partitioning=Hash([ts@1],REDACTED |_|_|_AggregateExec: mode=FinalPartitioned, gby=[le@0 as le, ts@1 as ts], aggr=[sum(histogram_gap_bucket.val)] REDACTED |_|_|_RepartitionExec: partitioning=Hash([le@0, ts@1],REDACTED diff --git a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result index af5f3a9a4f8..e50fa43f10e 100644 --- a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result +++ b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result @@ -110,7 +110,7 @@ TQL ANALYZE (0, 10, '5s') histogram_quantile(0.5, tsid_no_aggr_histogram_bucket) +-+-+-+ | 0_| 0_|_HistogramFoldExec: le=@2, field=@0, quantile=0.5 REDACTED |_|_|_RepartitionExec: partitioning=REDACTED -|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@2 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] REDACTED +|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@3 ASC NULLS LAST, TRY_CAST(le@2 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[val@0 as val, job@1 as job, le@2 as le, ts@4 as ts] REDACTED From 1c5eabcbbfc11f8eb75ccad6411f2a0ee055412c Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:59:44 +0000 Subject: [PATCH 096/126] feat(otlp): support cumulative exponential histograms (#8900) * feat: implement exponential histogram Signed-off-by: shuiyisong * chore: remove duplicate tests Signed-off-by: shuiyisong * fix(otlp): enforce exponential histogram ingestion safety Signed-off-by: shuiyisong * chore: update rfc Signed-off-by: shuiyisong * fix: test Signed-off-by: shuiyisong * fix(otlp): remove protocol-coupled histogram checks Signed-off-by: shuiyisong * perf(otlp): reuse native histogram schema across data points Signed-off-by: shuiyisong * fix: merge repeated OTLP histogram fragments Signed-off-by: shuiyisong * fix(otlp): build rejection messages lazily Signed-off-by: shuiyisong * fix: add doc Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- config/config.md | 2 + config/frontend.example.toml | 2 + config/standalone.example.toml | 2 + docs/rfcs/2026-08-04-native-histograms.md | 109 ++- src/cmd/src/standalone.rs | 4 +- src/cmd/tests/load_config_test.rs | 49 + src/frontend/src/frontend.rs | 4 +- src/frontend/src/instance/otlp.rs | 31 +- src/frontend/src/server.rs | 7 +- src/frontend/src/service_config/otlp.rs | 13 + src/operator/src/insert.rs | 15 +- src/servers/src/http.rs | 9 +- src/servers/src/http/otlp.rs | 75 +- src/servers/src/lib.rs | 1 + src/servers/src/native_histogram.rs | 524 ++++++++++ src/servers/src/otel_arrow.rs | 108 ++- src/servers/src/otlp/metrics.rs | 1051 ++++++++++++++++++++- src/servers/src/prom_remote_write/v2.rs | 562 +---------- src/servers/src/query_handler.rs | 11 +- src/session/src/protocol_ctx.rs | 2 + src/standalone/src/options.rs | 5 +- tests-integration/src/test_util.rs | 22 +- tests-integration/tests/grpc.rs | 199 +++- tests-integration/tests/http.rs | 145 ++- 24 files changed, 2322 insertions(+), 630 deletions(-) create mode 100644 src/servers/src/native_histogram.rs diff --git a/config/config.md b/config/config.md index 1b18565cc08..91cb4fe7669 100644 --- a/config/config.md +++ b/config/config.md @@ -74,6 +74,7 @@ | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | +| `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | @@ -316,6 +317,7 @@ | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | +| `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index ec1f9c89b42..3635a63f063 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -237,6 +237,8 @@ enable = true [otlp] ## Whether to enable OpenTelemetry protocol in HTTP API. enable = true +## Experimental: enable cumulative OTLP exponential histogram ingestion. +experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 5ada605b92b..72ba992cce0 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -204,6 +204,8 @@ enable = true [otlp] ## Whether to enable OpenTelemetry protocol in HTTP API. enable = true +## Experimental: enable cumulative OTLP exponential histogram ingestion. +experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 diff --git a/docs/rfcs/2026-08-04-native-histograms.md b/docs/rfcs/2026-08-04-native-histograms.md index 1e7692bf516..c714b35351b 100644 --- a/docs/rfcs/2026-08-04-native-histograms.md +++ b/docs/rfcs/2026-08-04-native-histograms.md @@ -7,22 +7,23 @@ Author: codex # Summary -GreptimeDB stores a Prometheus native histogram in one Struct-valued field and -evaluates it as a first-class PromQL sample. This document records the supported -protocols, the storage invariant, and the compatibility decisions needed to make -that path predictable. +GreptimeDB stores native histograms in one Prometheus-compatible Struct-valued +field and evaluates them as first-class PromQL samples. Prometheus Remote Write +2.0 supplies native histograms directly. OTLP `ExponentialHistogram` is an +ingestion transport: accepted cumulative points are normalized into the same +Struct before persistence and are queried only as native histograms. -Native histograms are experimental. Prometheus Remote Write 2.0 is the only -supported native-histogram ingestion protocol. Remote Write 1.0 histogram -payloads are rejected instead of being acknowledged and dropped. +Native histograms are experimental. Prometheus Remote Write 2.0 and cumulative +OTLP exponential histograms are supported behind separate configuration gates. +Remote Write 1.0 histogram payloads are rejected instead of being acknowledged and dropped. Native-histogram Remote Read is deferred; the existing Remote Read path continues to return scalar samples only. # Goals 1. Prevent silent native-histogram data loss at protocol boundaries. -2. Keep one stable persisted representation for integer, float, exponential, - and custom-bucket native histograms. +2. Keep one stable Prometheus-compatible persisted representation across + Prometheus and OTLP ingestion. 3. Match Prometheus query behavior where it is observable and practical. 4. State intentional limitations explicitly so incomplete behavior is not mistaken for support. @@ -32,20 +33,37 @@ continues to return scalar samples only. - Supporting native histograms in Remote Write 1.0. - Returning native histograms through Prometheus Remote Read. - Persisting Remote Write metric metadata. -- Ingesting OTLP exponential histograms. - Persisting exemplars. +- Providing an OTLP-specific histogram query surface or reconstructing and + re-exporting the original OTLP point after persistence. - Removing mixed float/histogram handling from PromQL expressions. - Propagating PromQL annotations produced on datanodes back to the frontend. # Data Model and Invariants -Each native histogram is stored in the configured native-histogram field -(`greptime_native_histogram` by default), a Struct whose children preserve the -wire-level schema, zero threshold, sum, reset hint, start timestamp, custom -bounds, spans, and either the integer or float count family. -Integer bucket deltas are converted to absolute integer counts for storage; -float bucket counts are already absolute. No separate sample-kind discriminator -is stored because the populated count family identifies it. +Each accepted native histogram, whether received directly through Remote Write +2.0 or normalized from OTLP `ExponentialHistogram`, is stored in the configured +native-histogram field (`greptime_native_histogram` by default) as the same +canonical Struct. The persistence boundary admits only values accepted by the +shared native-histogram validation. + +## Prometheus Remote Write 2.0 + +Remote Write 2.0 native histograms enter the persistence path in Prometheus +format. The Struct preserves the validated schema, zero threshold, sum, reset +hint, start timestamp, custom bounds, spans, and either the integer or float +count family. Integer bucket deltas are converted to absolute integer counts +for storage; float bucket counts are already absolute. No separate sample-kind +discriminator is stored because the populated count family identifies it. + +## OTLP ExponentialHistogram + +OTLP `ExponentialHistogram` does not introduce a second persisted format. +Accepted cumulative points are normalized into a valid Prometheus native +histogram and then pass through the same validator and Struct encoder. The +Struct does not retain the raw OTLP point or a source-protocol discriminator, +and queries use only the native-histogram PromQL behavior. Detailed conversion +and rejection rules are listed under [OTLP](#otlp). Within one resolved catalog, schema, and physical-table routing context, a metric name has exactly one persisted sample kind: @@ -78,7 +96,7 @@ that send native histograms must use Remote Write 2.0. ## Remote Write 2.0 Remote Write 2.0 accepts integer and float native histograms while -`http.experimental_enable_prometheus_native_histogram` is enabled. Supported +`prom_store.experimental_enable_prometheus_native_histogram` is enabled. Supported exponential schemas are `-4` through `8`; schema `-53` represents native histograms with custom buckets. @@ -103,11 +121,46 @@ explicit if sampled responses are implemented. ## OTLP -OTLP exponential histograms are not converted in this version. The ingestion -branch intentionally remains deferred until it can map temporality, scale, -reset behavior, attributes, and rejected-point reporting into the canonical -native-histogram path. The code carries an explicit TODO rather than a partial -encoder. +OTLP exponential histograms are accepted when +`otlp.experimental_enable_exponential_histogram` is enabled. The option defaults +to false and applies to OTLP/HTTP. Disabled points are rejected rather than +silently acknowledged. OTel Arrow exponential histograms are rejected because +the current Arrow wire format omits `zero_threshold`; accepting them would +silently change the distribution. Cumulative temporality is required; delta and +unspecified exponential histograms are rejected before their points are +converted. Explicit OTLP histograms keep their existing `_bucket`, `_sum`, and +`_count` representation, including their existing delta behavior. + +OTLP scales `-4` through `8` map directly to Prometheus schemas. Higher scales +are downscaled to schema `8`: dense counts that collide are merged before the +OTLP lower-bound index is shifted by one to the Prometheus upper-bound index. +This preserves count mass but irreversibly loses distinctions between source +buckets that merge. Lower scales are rejected. OTLP counts always populate the +integer histogram family. The transmitted non-negative finite zero threshold +and zero count are preserved; the Struct start timestamp is stored in +milliseconds and the reset hint is unknown. Point timestamps and attributes +retain their existing mode-specific conversion rules. Legacy mode retains its +normalized name, attribute rules, and nanosecond row timestamp. Non-legacy mode +retains Prometheus-compatible translation and its millisecond row timestamp. +Both modes write the same canonical Struct. + +An absent sum is stored as an ordinary quiet NaN, so the sample remains +selectable while `histogram_sum` is unknown. Any NaN sum on a point without +`NoRecordedValue` is normalized to the same value, even if its payload has the +Prometheus stale-marker bits. Only an exponential-histogram point carrying OTLP +`NoRecordedValue` becomes an empty schema-0 integer histogram with the canonical +Prometheus stale-NaN sum; its attributes and timestamps remain. This +interpretation intentionally does not change gauges, sums, or explicit histograms. + +Invalid points are skipped while unrelated valid points continue. Mixed +accepted/rejected OTLP/HTTP requests return partial success; a request with only +rejected points returns `InvalidArgument`. OTel Arrow uses an `OK` batch status +for mixed batches and `INVALID_ARGUMENT` when all points are rejected. Rejection +details are bounded, and metric metadata is emitted only for a metric that +produced an accepted row. + +Minimum, maximum, and exemplars are not persisted. Delta accumulation, +zero-run span compaction, and a dedicated rejection metric remain deferred. # PromQL Compatibility Decisions @@ -208,20 +261,22 @@ boundary already recorded in # Testing and Compatibility -Compatibility coverage includes protocol rejection, kind exclusivity, +Compatibility coverage includes protocol rejection, OTLP scale conversion and +partial-success behavior, kind exclusivity, stale-marker selector semantics, synthetic-zero rates, incompatible empty layouts, overflow-safe averages, layout-sensitive equality, infinite custom midpoints, and exponential overflow indices for schemas `-4`, `0`, and `8`. Behavioral coverage also verifies frontend PromQL warning and info responses. -This work does not change GreptimeDB's persisted Struct, protobuf dependencies, -or public configuration. Existing native-histogram data remains readable. +This work does not change GreptimeDB's persisted Struct or protobuf +dependencies. It adds the disabled-by-default public OTLP gate; existing +native-histogram data remains readable. # Future Work - Native-histogram Remote Read, including exact integer round-trips and streamed chunks with start timestamps. -- OTLP exponential-histogram conversion with partial-success reporting. +- OTLP exponential-histogram delta accumulation. - Persistent Remote Write metadata and accurate help/unit updates. - Native-histogram exemplars and exemplar query APIs. - Start-timestamp overlap annotations. diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 2ffcc5e729c..2b5bd5e8cec 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -1129,7 +1129,9 @@ mod tests { fn test_toml() { let opts = StandaloneOptions::default(); let toml_string = toml::to_string(&opts).unwrap(); - let _parsed: StandaloneOptions = toml::from_str(&toml_string).unwrap(); + assert!(toml_string.contains("experimental_enable_exponential_histogram = false")); + let parsed: StandaloneOptions = toml::from_str(&toml_string).unwrap(); + assert_eq!(parsed.otlp, opts.otlp); } #[test] diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index fd32680f4c6..70594d0653f 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -153,6 +153,12 @@ fn test_load_frontend_example_config() { let options = GreptimeOptions::::load_layered_options(example_config.to_str(), "") .unwrap(); + assert!( + !options + .component + .otlp + .experimental_enable_exponential_histogram + ); let expected = GreptimeOptions:: { component: FrontendOptions { default_timezone: Some("UTC".to_string()), @@ -329,6 +335,12 @@ fn test_load_standalone_example_config() { let options = GreptimeOptions::::load_layered_options(example_config.to_str(), "") .unwrap(); + assert!( + !options + .component + .otlp + .experimental_enable_exponential_histogram + ); let expected = GreptimeOptions:: { component: StandaloneOptions { default_timezone: Some("UTC".to_string()), @@ -380,6 +392,43 @@ fn test_load_standalone_example_config() { similar_asserts::assert_eq!(options, expected); } +#[test] +fn test_load_otlp_exponential_histogram_option() { + let config = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + config.path(), + "[otlp]\nexperimental_enable_exponential_histogram = true\n", + ) + .unwrap(); + + let frontend = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + assert!( + frontend + .component + .otlp + .experimental_enable_exponential_histogram + ); + + let standalone = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + assert!( + standalone + .component + .otlp + .experimental_enable_exponential_histogram + ); + assert!( + standalone + .component + .frontend_options() + .otlp + .experimental_enable_exponential_histogram + ); +} + #[test] fn test_load_standalone_user_provider_from_config() { let config = tempfile::NamedTempFile::new().unwrap(); diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index b8fea7b6043..18cedc0545c 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -218,7 +218,9 @@ mod tests { fn test_toml() { let opts = FrontendOptions::default(); let toml_string = toml::to_string(&opts).unwrap(); - let _parsed: FrontendOptions = toml::from_str(&toml_string).unwrap(); + assert!(toml_string.contains("experimental_enable_exponential_histogram = false")); + let parsed: FrontendOptions = toml::from_str(&toml_string).unwrap(); + assert_eq!(parsed.otlp, opts.otlp); } #[test] diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 459a05321cb..08086198536 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -27,7 +27,7 @@ use client::Output; use common_catalog::consts::{trace_operations_table_name, trace_services_table_name}; use common_error::ext::BoxedError; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; -use common_telemetry::tracing; +use common_telemetry::{tracing, warn}; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; @@ -38,7 +38,7 @@ use servers::interceptor::{OpenTelemetryProtocolInterceptor, OpenTelemetryProtoc use servers::otlp; use servers::otlp::trace::span::TraceSpanGroup; use servers::query_handler::{ - OpenTelemetryProtocolHandler, PipelineHandlerRef, TraceIngestOutcome, + MetricsIngestOutcome, OpenTelemetryProtocolHandler, PipelineHandlerRef, TraceIngestOutcome, }; use session::context::QueryContextRef; use snafu::ResultExt; @@ -90,7 +90,7 @@ impl OpenTelemetryProtocolHandler for Instance { &self, request: ExportMetricsServiceRequest, ctx: QueryContextRef, - ) -> ServerResult { + ) -> ServerResult { self.plugins .get::() .as_ref() @@ -120,8 +120,19 @@ impl OpenTelemetryProtocolHandler for Instance { .unwrap_or_default(); metric_ctx.is_legacy = is_legacy; - let (requests, rows, semantic_index) = + let (requests, rows, semantic_index, mut outcome) = otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; + if outcome.rejected_data_points > 0 { + warn!( + "Rejected {} OTLP exponential histogram data points: {}", + outcome.rejected_data_points, + outcome.error_message.as_deref().unwrap_or_default() + ); + } + if outcome.accepted_data_points == 0 { + return Ok(outcome); + } + self.check_row_insert_permission(&requests, &ctx, PermissionReq::Action(OTLP_WRITE)) .context(AuthSnafu)?; self.cache_otlp_legacy(&input_names, &ctx, is_legacy)?; @@ -143,9 +154,9 @@ impl OpenTelemetryProtocolHandler for Instance { Arc::new(c) }; - // If the user uses the legacy path, it is by default without metric engine. - if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { - self.handle_row_inserts(requests, ctx, false, false) + // OTLP tables have one sample field in both the legacy and physical paths. + let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { + self.handle_row_inserts(requests, ctx, false, true) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) @@ -154,11 +165,13 @@ impl OpenTelemetryProtocolHandler for Instance { .extension(PHYSICAL_TABLE_PARAM) .unwrap_or(GREPTIME_PHYSICAL_TABLE) .to_string(); - self.handle_metric_row_inserts(requests, ctx, physical_table.clone()) + self.handle_metric_row_inserts(requests, ctx, physical_table) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) - } + }?; + outcome.write_cost = output.meta.cost; + Ok(outcome) } #[tracing::instrument(skip_all)] diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 8fdebd8683c..f0dc226ba85 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -158,8 +158,11 @@ where } if opts.otlp.enable { - builder = builder - .with_otlp_handler(self.instance.clone(), opts.prom_store.with_metric_engine); + builder = builder.with_otlp_handler( + self.instance.clone(), + opts.prom_store.with_metric_engine, + opts.otlp.experimental_enable_exponential_histogram, + ); } if opts.jaeger.enable { diff --git a/src/frontend/src/service_config/otlp.rs b/src/frontend/src/service_config/otlp.rs index 3b01f422894..5dcab2c8ce2 100644 --- a/src/frontend/src/service_config/otlp.rs +++ b/src/frontend/src/service_config/otlp.rs @@ -20,6 +20,7 @@ const DEFAULT_TRACE_INGEST_CHUNK_SIZE: usize = 512; #[serde(default)] pub struct OtlpOptions { pub enable: bool, + pub experimental_enable_exponential_histogram: bool, /// Maximum spans per trace ingest chunk. Set to 0 to disable splitting. pub trace_ingest_chunk_size: usize, } @@ -28,6 +29,7 @@ impl Default for OtlpOptions { fn default() -> Self { Self { enable: true, + experimental_enable_exponential_histogram: false, trace_ingest_chunk_size: DEFAULT_TRACE_INGEST_CHUNK_SIZE, } } @@ -41,10 +43,12 @@ mod tests { fn test_otlp_options() { let default = OtlpOptions::default(); assert!(default.enable); + assert!(!default.experimental_enable_exponential_histogram); assert_eq!(default.trace_ingest_chunk_size, 512); let options: OtlpOptions = toml::from_str("enable = false").unwrap(); assert!(!options.enable); + assert!(!options.experimental_enable_exponential_histogram); assert_eq!( options.trace_ingest_chunk_size, DEFAULT_TRACE_INGEST_CHUNK_SIZE @@ -52,6 +56,15 @@ mod tests { let options: OtlpOptions = toml::from_str("trace_ingest_chunk_size = 0").unwrap(); assert!(options.enable); + assert!(!options.experimental_enable_exponential_histogram); assert_eq!(options.trace_ingest_chunk_size, 0); + + let options: OtlpOptions = + toml::from_str("experimental_enable_exponential_histogram = true").unwrap(); + assert!(options.experimental_enable_exponential_histogram); + + let serialized = toml::to_string(&options).unwrap(); + assert!(serialized.contains("experimental_enable_exponential_histogram = true")); + assert_eq!(toml::from_str::(&serialized).unwrap(), options); } } diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index fabc1b3778b..ead517ab80f 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -982,8 +982,8 @@ impl Inserter { /// When `accommodate_existing_schema` is true, it may modify the input `req` to /// accommodate it with existing schema. See [`create_or_alter_tables_on_demand`](Self::create_or_alter_tables_on_demand) /// for more details. - /// When `accommodate_existing_schema` is true and `is_single_value` is true, it also consider fields when modifying the - /// input `req`. + /// When `is_single_value` is true, it also rejects native-histogram/float kind changes. + /// When both options are true, it considers fields when modifying the input `req`. fn get_alter_table_expr_on_demand( &self, req: &mut RowInsertRequest, @@ -1018,8 +1018,7 @@ impl Inserter { return Ok(None); }; - // If accommodate_existing_schema is true, update request schema for Timestamp/Field columns - if accommodate_existing_schema { + if is_single_value { let request_is_native_histogram = request_is_native_histogram(request_schema); let table_is_native_histogram = table_is_native_histogram(table); ensure!( @@ -1030,6 +1029,10 @@ impl Inserter { ), } ); + } + + // If accommodate_existing_schema is true, update request schema for Timestamp/Field columns + if accommodate_existing_schema { let table_schema = table.schema(); // Find timestamp column name let ts_col_name = table_schema.timestamp_column().map(|c| c.name.clone()); @@ -1638,7 +1641,7 @@ mod tests { }), }; let error = inserter - .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, true, true, true) + .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, false, true, true) .unwrap_err(); assert!( error @@ -1666,7 +1669,7 @@ mod tests { &mut sample_req, &histogram_table, &ctx, - true, + false, true, true, ) diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index bf7478eacfd..c0f6af81843 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -724,11 +724,16 @@ impl HttpServerBuilder { self, handler: OpenTelemetryProtocolHandlerRef, with_metric_engine: bool, + experimental_enable_exponential_histogram: bool, ) -> Self { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/otlp"), - HttpServer::route_otlp(handler, with_metric_engine), + HttpServer::route_otlp( + handler, + with_metric_engine, + experimental_enable_exponential_histogram, + ), ), ..self } @@ -1386,6 +1391,7 @@ impl HttpServer { fn route_otlp( otlp_handler: OpenTelemetryProtocolHandlerRef, with_metric_engine: bool, + experimental_enable_exponential_histogram: bool, ) -> Router { Router::new() .route("/v1/metrics", routing::post(otlp::metrics)) @@ -1397,6 +1403,7 @@ impl HttpServer { ) .with_state(OtlpState { with_metric_engine, + experimental_enable_exponential_histogram, handler: otlp_handler, }) } diff --git a/src/servers/src/http/otlp.rs b/src/servers/src/http/otlp.rs index 6540f6d4c48..d04022cdf88 100644 --- a/src/servers/src/http/otlp.rs +++ b/src/servers/src/http/otlp.rs @@ -27,7 +27,9 @@ use mime_guess::mime; use opentelemetry_proto::tonic::collector::logs::v1::{ ExportLogsServiceRequest, ExportLogsServiceResponse, }; -use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceResponse; +use opentelemetry_proto::tonic::collector::metrics::v1::{ + ExportMetricsPartialSuccess, ExportMetricsServiceResponse, +}; use opentelemetry_proto::tonic::collector::trace::v1::{ ExportTracePartialSuccess, ExportTraceServiceRequest, ExportTraceServiceResponse, }; @@ -44,7 +46,9 @@ use crate::http::extractor::{ }; use crate::http::header::{CONTENT_TYPE_PROTOBUF, write_cost_header_map}; use crate::metrics::METRIC_HTTP_OPENTELEMETRY_LOGS_ELAPSED; -use crate::query_handler::{OpenTelemetryProtocolHandlerRef, PipelineHandler, TraceIngestOutcome}; +use crate::query_handler::{ + MetricsIngestOutcome, OpenTelemetryProtocolHandlerRef, PipelineHandler, TraceIngestOutcome, +}; #[derive(Clone, prost::Message)] pub struct GoogleRpcStatus { @@ -73,6 +77,7 @@ fn content_type_to_string(content_type: Option<&TypedHeader>) -> St #[derive(Clone)] pub struct OtlpState { pub with_metric_engine: bool, + pub experimental_enable_exponential_histogram: bool, pub handler: OpenTelemetryProtocolHandlerRef, } @@ -84,7 +89,7 @@ pub async fn metrics( http_opts: OtlpMetricOptions, content_type: Option>, bytes: Bytes, -) -> Result> { +) -> Result { if is_json_content_type(content_type.as_ref().map(|h| &h.0)) { return error::UnsupportedJsonContentTypeSnafu {}.fail(); } @@ -103,6 +108,7 @@ pub async fn metrics( let OtlpState { with_metric_engine, + experimental_enable_exponential_histogram, handler, } = state; @@ -111,6 +117,7 @@ pub async fn metrics( resource_attrs: http_opts.resource_attrs, promote_scope_attrs: http_opts.promote_scope_attrs, with_metric_engine, + experimental_enable_exponential_histogram, // set is_legacy later is_legacy: false, metric_type: MetricType::Init, @@ -118,15 +125,15 @@ pub async fn metrics( })); let query_ctx = Arc::new(query_ctx); - handler - .metrics(request, query_ctx) - .await - .map(|o| OtlpResponse { - resp_body: ExportMetricsServiceResponse { - partial_success: None, - }, - write_cost: o.meta.cost, - }) + handler.metrics(request, query_ctx).await.map(|outcome| { + if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 { + OtlpMetricsResponse::Failure(outcome) + } else if outcome.rejected_data_points > 0 { + OtlpMetricsResponse::PartialSuccess(outcome) + } else { + OtlpMetricsResponse::FullSuccess(outcome) + } + }) } #[axum_macros::debug_handler] @@ -259,6 +266,50 @@ pub struct OtlpResponse { write_cost: usize, } +pub enum OtlpMetricsResponse { + FullSuccess(MetricsIngestOutcome), + PartialSuccess(MetricsIngestOutcome), + Failure(MetricsIngestOutcome), +} + +impl IntoResponse for OtlpMetricsResponse { + fn into_response(self) -> axum::response::Response { + match self { + OtlpMetricsResponse::FullSuccess(outcome) => { + let mut header_map = write_cost_header_map(outcome.write_cost); + header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone()); + let body = ExportMetricsServiceResponse { + partial_success: None, + }; + (header_map, body.encode_to_vec()).into_response() + } + OtlpMetricsResponse::PartialSuccess(outcome) => { + let mut header_map = write_cost_header_map(outcome.write_cost); + header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone()); + let body = ExportMetricsServiceResponse { + partial_success: Some(ExportMetricsPartialSuccess { + rejected_data_points: outcome.rejected_data_points, + error_message: outcome.error_message.unwrap_or_default(), + }), + }; + (header_map, body.encode_to_vec()).into_response() + } + OtlpMetricsResponse::Failure(outcome) => { + let status = GoogleRpcStatus { + code: tonic::Code::InvalidArgument as i32, + message: outcome.error_message.unwrap_or_default(), + }; + ( + StatusCode::BAD_REQUEST, + [(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.as_ref())], + status.encode_to_vec(), + ) + .into_response() + } + } + } +} + impl IntoResponse for OtlpResponse { fn into_response(self) -> axum::response::Response { let mut header_map = write_cost_header_map(self.write_cost); diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index 2a3796aab21..79fd392a1a9 100644 --- a/src/servers/src/lib.rs +++ b/src/servers/src/lib.rs @@ -34,6 +34,7 @@ pub mod interceptor; pub mod metrics; pub mod metrics_handler; pub mod mysql; +pub(crate) mod native_histogram; pub mod opentsdb; pub mod otel_arrow; pub mod otlp; diff --git a/src/servers/src/native_histogram.rs b/src/servers/src/native_histogram.rs new file mode 100644 index 00000000000..f1dae526993 --- /dev/null +++ b/src/servers/src/native_histogram.rs @@ -0,0 +1,524 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; +use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; +use api::helper::ColumnDataTypeWrapper; +use api::v1::value::ValueData; +use api::v1::{ColumnSchema, ListValue, SemanticType, Value}; +use common_query::native_histogram::*; +use common_query::prelude::greptime_native_histogram; +use snafu::{Snafu, ensure}; + +const MAX_NATIVE_HISTOGRAM_SCHEMA: i32 = 8; +const MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA: i32 = 52; + +#[derive(Debug, Snafu)] +#[snafu(display("{message}"))] +pub(crate) struct NativeHistogramError { + message: String, +} + +type Result = std::result::Result; + +/// Returns the canonical column schema for a native histogram value. +pub(crate) fn native_histogram_column_schema() -> Result { + let (datatype, datatype_extension) = + ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone()) + .map_err(|error| NativeHistogramError { + message: format!("native histogram type cannot be encoded: {error}"), + })? + .into_parts(); + + Ok(ColumnSchema { + column_name: greptime_native_histogram().to_string(), + datatype: datatype as i32, + semantic_type: SemanticType::Field as i32, + datatype_extension, + options: None, + }) +} + +/// Validates and encodes a Prometheus histogram into the canonical Struct value. +pub(crate) fn encode_native_histogram(histogram: &Histogram) -> Result { + let uses_float_counts = native_histogram_uses_float_counts(histogram)?; + validate_native_histogram(histogram, uses_float_counts)?; + + let mut items = Vec::with_capacity(NATIVE_HISTOGRAM_FIELD_NAMES.len()); + let positive_span_lengths = i32_span_lengths("positive", &histogram.positive_spans)?; + let negative_span_lengths = i32_span_lengths("negative", &histogram.negative_spans)?; + items.extend([ + pb_value(ValueData::I32Value(histogram.schema)), + pb_value(ValueData::F64Value(histogram.zero_threshold)), + pb_value(ValueData::F64Value(histogram.sum)), + pb_value(ValueData::I32Value(histogram.reset_hint)), + optional_pb_value((histogram.start_timestamp != 0).then_some( + ValueData::TimestampMillisecondValue(histogram.start_timestamp), + )), + f64_list_value(histogram.custom_values.iter().copied()), + i32_list_value(histogram.positive_spans.iter().map(|span| span.offset)), + i32_list_value(positive_span_lengths), + i32_list_value(histogram.negative_spans.iter().map(|span| span.offset)), + i32_list_value(negative_span_lengths), + ]); + + if uses_float_counts { + validate_float_native_histogram_counts(histogram)?; + let count = match histogram.count.as_ref() { + Some(Count::CountFloat(count)) => *count, + _ => 0.0, + }; + let zero_count = match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count, + _ => 0.0, + }; + items.extend([ + null_pb_value(), + null_pb_value(), + i64_list_value(std::iter::empty()), + i64_list_value(std::iter::empty()), + pb_value(ValueData::F64Value(count)), + pb_value(ValueData::F64Value(zero_count)), + f64_list_value(histogram.positive_counts.iter().copied()), + f64_list_value(histogram.negative_counts.iter().copied()), + ]); + } else { + let count = match histogram.count.as_ref() { + Some(Count::CountInt(count)) => *count, + _ => 0, + }; + let zero_count = match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count, + _ => 0, + }; + let positive_buckets = bucket_counts_from_deltas(&histogram.positive_deltas)?; + let negative_buckets = bucket_counts_from_deltas(&histogram.negative_deltas)?; + validate_integer_native_histogram_counts(histogram, &positive_buckets, &negative_buckets)?; + let count = i64::try_from(count).map_err(|_| NativeHistogramError { + message: format!("native histogram integer count {count} overflows i64"), + })?; + let zero_count = i64::try_from(zero_count).map_err(|_| NativeHistogramError { + message: format!("native histogram integer zero_count {zero_count} overflows i64"), + })?; + items.extend([ + pb_value(ValueData::I64Value(count)), + pb_value(ValueData::I64Value(zero_count)), + i64_list_value(positive_buckets), + i64_list_value(negative_buckets), + null_pb_value(), + null_pb_value(), + f64_list_value(std::iter::empty()), + f64_list_value(std::iter::empty()), + ]); + } + + Ok(ValueData::StructValue(api::v1::StructValue { items })) +} + +fn validate_native_histogram(histogram: &Histogram, uses_float_counts: bool) -> Result<()> { + let exponential_overflow_index = validate_native_histogram_schema(histogram.schema)?; + validate_native_histogram_custom_values(histogram)?; + + if histogram.schema == CUSTOM_BUCKETS_SCHEMA { + ensure!( + histogram.zero_threshold == 0.0 && native_histogram_zero_count_is_zero(histogram), + NativeHistogramSnafu { + message: "custom native histogram must not use a zero bucket" + } + ); + ensure!( + histogram.negative_spans.is_empty() + && histogram.negative_deltas.is_empty() + && histogram.negative_counts.is_empty(), + NativeHistogramSnafu { + message: "custom native histogram must not use negative buckets" + } + ); + } + + let (positive_buckets, negative_buckets) = if uses_float_counts { + ( + histogram.positive_counts.len(), + histogram.negative_counts.len(), + ) + } else { + ( + histogram.positive_deltas.len(), + histogram.negative_deltas.len(), + ) + }; + let bucket_index_range = if let Some(overflow_index) = exponential_overflow_index { + (i32::MIN, overflow_index) + } else { + ( + 0, + i32::try_from(histogram.custom_values.len()).map_err(|_| NativeHistogramError { + message: "custom native histogram has too many custom_values".to_string(), + })?, + ) + }; + validate_native_histogram_spans( + "positive", + &histogram.positive_spans, + positive_buckets, + bucket_index_range, + )?; + validate_native_histogram_spans( + "negative", + &histogram.negative_spans, + negative_buckets, + bucket_index_range, + )?; + + Ok(()) +} + +fn validate_native_histogram_schema(schema: i32) -> Result> { + if schema == CUSTOM_BUCKETS_SCHEMA { + return Ok(None); + } + + if let Some(overflow_index) = exponential_overflow_bucket_index(schema) { + return Ok(Some(overflow_index)); + } + + if (MAX_NATIVE_HISTOGRAM_SCHEMA + 1..=MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA).contains(&schema) { + Err(NativeHistogramError { + message: format!("native histogram schema {schema} must be reduced before ingestion"), + }) + } else { + Err(NativeHistogramError { + message: format!("native histogram schema {schema} is unsupported"), + }) + } +} + +fn validate_native_histogram_custom_values(histogram: &Histogram) -> Result<()> { + if histogram.schema != CUSTOM_BUCKETS_SCHEMA { + ensure!( + histogram.custom_values.is_empty(), + NativeHistogramSnafu { + message: "standard native histogram must not use custom_values" + } + ); + return Ok(()); + } + + for value in &histogram.custom_values { + ensure!( + !value.is_nan() && *value != f64::INFINITY, + NativeHistogramSnafu { + message: "custom native histogram custom_values must not contain +Inf or NaN" + } + ); + } + for values in histogram.custom_values.windows(2) { + ensure!( + values[0] < values[1], + NativeHistogramSnafu { + message: "custom native histogram custom_values must be sorted" + } + ); + } + + Ok(()) +} + +fn validate_native_histogram_spans( + name: &str, + spans: &[BucketSpan], + bucket_count: usize, + bucket_index_range: (i32, i32), +) -> Result<()> { + let span_len = spans.iter().try_fold(0usize, |sum, span| { + let length = usize::try_from(span.length).map_err(|_| NativeHistogramError { + message: format!("native histogram {name} span length exceeds usize"), + })?; + sum.checked_add(length).ok_or_else(|| NativeHistogramError { + message: format!("native histogram {name} spans overflow"), + }) + })?; + ensure!( + span_len == bucket_count, + NativeHistogramSnafu { + message: format!( + "native histogram {name} spans describe {span_len} buckets, found {bucket_count}" + ) + } + ); + + let mut current_index = 0i32; + for (span_index, span) in spans.iter().enumerate() { + ensure!( + span.offset >= 0 || (span_index == 0 && bucket_index_range.0 == i32::MIN), + NativeHistogramSnafu { + message: format!( + "native histogram {name} span {} has negative offset {}", + span_index + 1, + span.offset + ) + } + ); + current_index = if span_index == 0 { + span.offset + } else { + current_index + .checked_add(span.offset) + .ok_or_else(|| NativeHistogramError { + message: format!("native histogram {name} span index overflows i32"), + })? + }; + + for _ in 0..span.length { + ensure!( + (bucket_index_range.0..=bucket_index_range.1).contains(¤t_index), + NativeHistogramSnafu { + message: format!( + "native histogram {name} bucket index {current_index} is out of range" + ) + } + ); + current_index = current_index + .checked_add(1) + .ok_or_else(|| NativeHistogramError { + message: format!("native histogram {name} span index overflows i32"), + })?; + } + } + + Ok(()) +} + +fn validate_float_native_histogram_counts(histogram: &Histogram) -> Result<()> { + let count = match histogram.count.as_ref() { + Some(Count::CountFloat(count)) => *count, + _ => 0.0, + }; + ensure!( + count >= 0.0 || count.is_nan(), + NativeHistogramSnafu { + message: "native histogram float count must not be negative" + } + ); + + let zero_count = match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count, + _ => 0.0, + }; + ensure!( + zero_count >= 0.0 || zero_count.is_nan(), + NativeHistogramSnafu { + message: "native histogram float zero_count must not be negative" + } + ); + + for (name, counts) in [ + ("positive", &histogram.positive_counts), + ("negative", &histogram.negative_counts), + ] { + for (index, count) in counts.iter().enumerate() { + ensure!( + *count >= 0.0 || count.is_nan(), + NativeHistogramSnafu { + message: format!( + "native histogram {name} bucket {} count must not be negative", + index + 1 + ) + } + ); + } + } + + Ok(()) +} + +fn validate_integer_native_histogram_counts( + histogram: &Histogram, + positive_buckets: &[i64], + negative_buckets: &[i64], +) -> Result<()> { + let count = match histogram.count.as_ref() { + Some(Count::CountInt(count)) => *count, + _ => 0, + }; + let zero_count = match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count, + _ => 0, + }; + let bucket_count = + positive_buckets + .iter() + .chain(negative_buckets) + .try_fold(zero_count, |total, bucket| { + let bucket = u64::try_from(*bucket).map_err(|_| NativeHistogramError { + message: "native histogram bucket count is negative".to_string(), + })?; + total + .checked_add(bucket) + .ok_or_else(|| NativeHistogramError { + message: "native histogram bucket total overflows u64".to_string(), + }) + })?; + ensure!( + if histogram.sum.is_nan() { + bucket_count <= count + } else { + bucket_count == count + }, + NativeHistogramSnafu { + message: format!( + "native histogram has {bucket_count} observations in buckets, count is {count}" + ) + } + ); + + Ok(()) +} + +fn native_histogram_zero_count_is_zero(histogram: &Histogram) -> bool { + match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count == 0, + Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count == 0.0, + None => true, + } +} + +fn native_histogram_uses_float_counts(histogram: &Histogram) -> Result { + let uses_float_count = matches!(histogram.count, Some(Count::CountFloat(_))) + || matches!(histogram.zero_count, Some(ZeroCount::ZeroCountFloat(_))); + let uses_int_count = matches!(histogram.count, Some(Count::CountInt(_))) + || matches!(histogram.zero_count, Some(ZeroCount::ZeroCountInt(_))); + let uses_float_buckets = + !histogram.positive_counts.is_empty() || !histogram.negative_counts.is_empty(); + let uses_int_buckets = + !histogram.positive_deltas.is_empty() || !histogram.negative_deltas.is_empty(); + + ensure!( + !matches!( + (&histogram.count, &histogram.zero_count), + (Some(Count::CountInt(_)), Some(ZeroCount::ZeroCountFloat(_))) + | (Some(Count::CountFloat(_)), Some(ZeroCount::ZeroCountInt(_))) + ), + NativeHistogramSnafu { + message: "native histogram count and zero_count must use the same integer or float family" + } + ); + ensure!( + !(uses_float_buckets && uses_int_buckets), + NativeHistogramSnafu { + message: "native histogram bucket counts must use either integer deltas or float counts" + } + ); + ensure!( + !(uses_float_count && uses_int_buckets), + NativeHistogramSnafu { + message: "float native histogram must not use integer bucket deltas" + } + ); + ensure!( + !(uses_int_count && uses_float_buckets), + NativeHistogramSnafu { + message: "integer native histogram must not use float bucket counts" + } + ); + + Ok(uses_float_count || uses_float_buckets) +} + +fn pb_value(value_data: ValueData) -> Value { + optional_pb_value(Some(value_data)) +} + +fn null_pb_value() -> Value { + optional_pb_value(None) +} + +fn optional_pb_value(value_data: Option) -> Value { + Value { value_data } +} + +fn list_value(values: impl IntoIterator) -> Value { + pb_value(ValueData::ListValue(ListValue { + items: values.into_iter().map(pb_value).collect(), + })) +} + +fn i32_list_value(values: impl IntoIterator) -> Value { + list_value(values.into_iter().map(ValueData::I32Value)) +} + +fn i32_span_lengths(name: &str, spans: &[BucketSpan]) -> Result> { + spans + .iter() + .map(|span| { + i32::try_from(span.length).map_err(|_| NativeHistogramError { + message: format!( + "native histogram {name} span length {} overflows i32", + span.length + ), + }) + }) + .collect() +} + +fn i64_list_value(values: impl IntoIterator) -> Value { + list_value(values.into_iter().map(ValueData::I64Value)) +} + +fn f64_list_value(values: impl IntoIterator) -> Value { + list_value(values.into_iter().map(ValueData::F64Value)) +} + +fn bucket_counts_from_deltas(deltas: &[i64]) -> Result> { + let mut count = 0_i64; + let mut buckets = Vec::with_capacity(deltas.len()); + + for delta in deltas { + count = count + .checked_add(*delta) + .ok_or_else(|| NativeHistogramError { + message: "native histogram bucket count overflows i64".to_string(), + })?; + ensure!( + count >= 0, + NativeHistogramSnafu { + message: "native histogram bucket count is negative" + } + ); + buckets.push(count); + } + + Ok(buckets) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn shared_errors_are_protocol_neutral() { + let error = encode_native_histogram(&Histogram { + schema: 9, + ..Default::default() + }) + .unwrap_err(); + + assert_eq!( + error.to_string(), + "native histogram schema 9 must be reduced before ingestion" + ); + assert!(!error.to_string().contains("remote write")); + assert!(!error.to_string().contains("OTLP")); + } +} diff --git a/src/servers/src/otel_arrow.rs b/src/servers/src/otel_arrow.rs index f905b39f5bc..320ef87dafe 100644 --- a/src/servers/src/otel_arrow.rs +++ b/src/servers/src/otel_arrow.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use auth::UserProviderRef; use common_error::ext::ErrorExt; use common_error::status_code::status_to_tonic_code; @@ -19,14 +21,20 @@ use common_telemetry::error; use futures::SinkExt; use otel_arrow_rust::Consumer; use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_server::ArrowMetricsService; -use otel_arrow_rust::proto::opentelemetry::arrow::v1::{BatchArrowRecords, BatchStatus}; +use otel_arrow_rust::proto::opentelemetry::arrow::v1::{ + BatchArrowRecords, BatchStatus, StatusCode as ArrowStatusCode, +}; +use otel_arrow_rust::proto::opentelemetry::metrics::v1::metric; +use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; use tonic::metadata::{Entry, MetadataValue}; use tonic::service::Interceptor; use tonic::{Request, Response, Status, Streaming}; use crate::error; use crate::grpc::context_auth; -use crate::query_handler::OpenTelemetryProtocolHandlerRef; +use crate::query_handler::{MetricsIngestOutcome, OpenTelemetryProtocolHandlerRef}; + +const EXPONENTIAL_HISTOGRAM_UNSUPPORTED: &str = "OTel Arrow exponential histograms are unsupported because the Arrow wire format omits zero_threshold"; pub struct OtelArrowServiceHandler { handler: T, @@ -42,6 +50,31 @@ impl OtelArrowServiceHandler { } } +fn batch_status( + batch_id: i64, + outcome: MetricsIngestOutcome, + has_exponential_histogram_data_points: bool, +) -> BatchStatus { + let status_code = if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 { + ArrowStatusCode::InvalidArgument + } else { + ArrowStatusCode::Ok + }; + let status_message = match outcome.error_message { + // Arrow keeps the feature gate off, so these fail before per-point validation. + Some(_) if has_exponential_histogram_data_points => { + EXPONENTIAL_HISTOGRAM_UNSUPPORTED.to_string() + } + Some(message) => message, + None => String::new(), + }; + BatchStatus { + batch_id, + status_code: status_code as i32, + status_message, + } +} + #[async_trait::async_trait] impl ArrowMetricsService for OtelArrowServiceHandler { type ArrowMetricsStream = futures::channel::mpsc::Receiver>; @@ -55,6 +88,11 @@ impl ArrowMetricsService for OtelArrowServiceHandler outcome, + Err(e) => { + let _ = sender + .send(Err(Status::new( + status_to_tonic_code(e.status_code()), + e.to_string(), + ))) + .await; + error!(e; "Failed to ingest metrics from otel-arrow"); + return; + } + }; + let batch_status = + batch_status(batch_id, outcome, has_exponential_histogram_data_points); let _ = sender.send(Ok(batch_status)).await; } }); @@ -131,3 +181,25 @@ impl Interceptor for HeaderInterceptor { Ok(request) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn batch_status_explains_arrow_exponential_histogram_limit() { + let status = batch_status( + 7, + MetricsIngestOutcome { + rejected_data_points: 1, + error_message: Some("internal OTLP rejection detail".to_string()), + ..Default::default() + }, + true, + ); + + assert_eq!(7, status.batch_id); + assert_eq!(ArrowStatusCode::InvalidArgument as i32, status.status_code); + assert_eq!(EXPONENTIAL_HISTOGRAM_UNSUPPORTED, status.status_message); + } +} diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 6d94b3a139f..332fa46f1d1 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -13,9 +13,16 @@ // limitations under the License. use ahash::HashSet; -use api::v1::{RowInsertRequests, Value}; +use api::greptime_proto::io::prometheus::write::v2::histogram::{ + Count as PromCount, ResetHint, ZeroCount as PromZeroCount, +}; +use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram as PromHistogram}; +use api::v1::value::ValueData; +use api::v1::{RowInsertRequests, SemanticType, Value}; use common_grpc::precision::Precision; +use common_query::native_histogram::native_histogram_value_type; use common_query::prelude::{GREPTIME_COUNT, greptime_timestamp, greptime_value}; +use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use lazy_static::lazy_static; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; @@ -26,8 +33,10 @@ use table::requests::{ SEMANTIC_METRIC_TEMPORALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, }; -use crate::error::Result; +use crate::error::{self, Result}; +use crate::native_histogram::{encode_native_histogram, native_histogram_column_schema}; use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME}; +use crate::query_handler::MetricsIngestOutcome; use crate::row_writer::{self, MultiTableData, TableData}; pub use crate::semantic::SemanticIndex; use crate::semantic::{ @@ -82,6 +91,9 @@ lazy_static! { const OTEL_SCOPE_NAME: &str = "name"; const OTEL_SCOPE_VERSION: &str = "version"; const OTEL_SCOPE_SCHEMA_URL: &str = "schema_url"; +const MIN_EXPONENTIAL_HISTOGRAM_SCALE: i32 = -4; +const MAX_EXPONENTIAL_HISTOGRAM_SCALE: i32 = 8; +const MAX_REJECTION_MESSAGE_BYTES: usize = 512; /// Convert OpenTelemetry metrics to GreptimeDB insert requests /// @@ -94,9 +106,15 @@ const OTEL_SCOPE_SCHEMA_URL: &str = "schema_url"; pub fn to_grpc_insert_requests( request: ExportMetricsServiceRequest, metric_ctx: &mut OtlpMetricCtx, -) -> Result<(RowInsertRequests, usize, SemanticIndex)> { +) -> Result<( + RowInsertRequests, + usize, + SemanticIndex, + MetricsIngestOutcome, +)> { let mut table_writer = MultiTableData::default(); let mut semantic_index = SemanticIndex::default(); + let mut outcome = MetricsIngestOutcome::default(); for resource in &request.resource_metrics { let resource_attrs = resource.resource.as_ref().map(|r| { @@ -123,13 +141,46 @@ pub fn to_grpc_insert_requests( scope_attrs.as_ref(), metric_ctx, &mut semantic_index, + &mut outcome, )?; } } } let (requests, rows) = table_writer.into_row_insert_requests(); - Ok((requests, rows, semantic_index)) + validate_sample_kinds(&requests)?; + Ok((requests, rows, semantic_index, outcome)) +} + +fn validate_sample_kinds(requests: &RowInsertRequests) -> Result<()> { + for request in &requests.inserts { + let Some(rows) = &request.rows else { + continue; + }; + let field_count = rows + .schema + .iter() + .filter(|column| column.semantic_type == SemanticType::Field as i32) + .count(); + let has_native_histogram = rows.schema.iter().any(|column| { + column.semantic_type == SemanticType::Field as i32 + && api::helper::is_column_type_value_eq( + column.datatype, + column.datatype_extension.clone(), + native_histogram_value_type(), + ) + }); + if has_native_histogram && field_count != 1 { + return Err(error::InvalidParameterSnafu { + reason: format!( + "OTLP metric `{}` cannot mix native histogram and float sample fields", + request.table_name + ), + } + .build()); + } + } + Ok(()) } /// The tables a metric emits and their per-table `metric.type`. Histogram fans @@ -159,8 +210,10 @@ fn emitted_semantic_tables( (format!("{base}{COUNT_TABLE_SUFFIX}"), METRIC_TYPE_COUNTER), (format!("{base}{SUM_TABLE_SUFFIX}"), METRIC_TYPE_COUNTER), ], - // ExponentialHistogram is a no-op today; Init never reaches encoding. - MetricType::ExponentialHistogram | MetricType::Init => vec![], + MetricType::ExponentialHistogram => { + vec![(base.to_string(), METRIC_TYPE_HISTOGRAM)] + } + MetricType::Init => vec![], } } @@ -170,6 +223,7 @@ fn temporality_value(data: &metric::Data) -> Option<&'static str> { let raw = match data { metric::Data::Sum(sum) => sum.aggregation_temporality, metric::Data::Histogram(hist) => hist.aggregation_temporality, + metric::Data::ExponentialHistogram(hist) => hist.aggregation_temporality, _ => return None, }; match AggregationTemporality::try_from(raw) { @@ -311,6 +365,7 @@ fn encode_metrics( scope_attrs: Option<&Vec>, metric_ctx: &OtlpMetricCtx, semantic_index: &mut SemanticIndex, + outcome: &mut MetricsIngestOutcome, ) -> Result<()> { let name = if metric_ctx.is_legacy { legacy_normalize_otlp_name(&metric.name) @@ -322,12 +377,7 @@ fn encode_metrics( ) }; - // Stamp semantic metadata against the same table name(s) the data is written - // to below. `unit` is captured here (it is otherwise discarded by the row - // encoders) along with the declared type/temporality. - record_metric_semantics(semantic_index, metric, &name, metric_ctx); - - if let Some(data) = &metric.data { + let emitted = if let Some(data) = &metric.data { match data { metric::Data::Gauge(gauge) => { encode_gauge( @@ -338,6 +388,8 @@ fn encode_metrics( scope_attrs, metric_ctx, )?; + add_accepted_data_points(outcome, gauge.data_points.len())?; + !gauge.data_points.is_empty() } metric::Data::Sum(sum) => { encode_sum( @@ -348,6 +400,8 @@ fn encode_metrics( scope_attrs, metric_ctx, )?; + add_accepted_data_points(outcome, sum.data_points.len())?; + !sum.data_points.is_empty() } metric::Data::Summary(summary) => { encode_summary( @@ -358,6 +412,8 @@ fn encode_metrics( scope_attrs, metric_ctx, )?; + add_accepted_data_points(outcome, summary.data_points.len())?; + !summary.data_points.is_empty() } metric::Data::Histogram(hist) => { encode_histogram( @@ -368,16 +424,384 @@ fn encode_metrics( scope_attrs, metric_ctx, )?; + add_accepted_data_points(outcome, hist.data_points.len())?; + !hist.data_points.is_empty() } - // TODO: Convert OTLP exponential histograms into the canonical native-histogram - // Struct. See docs/rfcs/2026-08-04-native-histograms.md. - metric::Data::ExponentialHistogram(_hist) => {} + metric::Data::ExponentialHistogram(hist) => encode_exponential_histogram( + table_writer, + &name, + hist, + resource_attrs, + scope_attrs, + metric_ctx, + outcome, + )?, } + } else { + false + }; + + if emitted { + // Stamp semantic metadata only after at least one row was accepted. + record_metric_semantics(semantic_index, metric, &name, metric_ctx); } Ok(()) } +fn add_accepted_data_points(outcome: &mut MetricsIngestOutcome, count: usize) -> Result<()> { + let count = i64::try_from(count).map_err(|_| { + error::InvalidParameterSnafu { + reason: "OTLP metrics data-point count exceeds i64", + } + .build() + })?; + outcome.accepted_data_points = + outcome + .accepted_data_points + .checked_add(count) + .ok_or_else(|| { + error::InvalidParameterSnafu { + reason: "OTLP accepted data-point count overflows i64", + } + .build() + })?; + Ok(()) +} + +fn reject_data_points( + outcome: &mut MetricsIngestOutcome, + count: usize, + reason: impl FnOnce() -> String, +) -> Result<()> { + if count == 0 { + return Ok(()); + } + let count = i64::try_from(count).map_err(|_| { + error::InvalidParameterSnafu { + reason: "OTLP rejected data-point count exceeds i64", + } + .build() + })?; + outcome.rejected_data_points = + outcome + .rejected_data_points + .checked_add(count) + .ok_or_else(|| { + error::InvalidParameterSnafu { + reason: "OTLP rejected data-point count overflows i64", + } + .build() + })?; + append_rejection_message(&mut outcome.error_message, reason); + Ok(()) +} + +fn append_rejection_message(message: &mut Option, reason: impl FnOnce() -> String) { + let message = message.get_or_insert_with(String::new); + let separator = if message.is_empty() { "" } else { "; " }; + let Some(available) = MAX_REJECTION_MESSAGE_BYTES.checked_sub(message.len()) else { + return; + }; + if available <= separator.len() { + return; + } + let reason = reason(); + message.push_str(separator); + + let available = MAX_REJECTION_MESSAGE_BYTES - message.len(); + if reason.len() <= available { + message.push_str(&reason); + return; + } + + const ELLIPSIS: &str = "..."; + let mut end = available.saturating_sub(ELLIPSIS.len()); + while !reason.is_char_boundary(end) { + end -= 1; + } + message.push_str(&reason[..end]); + if available >= ELLIPSIS.len() { + message.push_str(ELLIPSIS); + } +} + +fn encode_exponential_histogram( + table_writer: &mut MultiTableData, + name: &str, + histogram: &ExponentialHistogram, + resource_attrs: Option<&Vec>, + scope_attrs: Option<&Vec>, + metric_ctx: &OtlpMetricCtx, + outcome: &mut MetricsIngestOutcome, +) -> Result { + if !metric_ctx.experimental_enable_exponential_histogram { + reject_data_points(outcome, histogram.data_points.len(), || { + format!( + "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" + ) + })?; + return Ok(false); + } + + match AggregationTemporality::try_from(histogram.aggregation_temporality) { + Ok(AggregationTemporality::Cumulative) => {} + Ok(AggregationTemporality::Delta) => { + reject_data_points(outcome, histogram.data_points.len(), || { + format!( + "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" + ) + })?; + return Ok(false); + } + _ => { + reject_data_points(outcome, histogram.data_points.len(), || { + format!( + "metric `{name}` has unspecified OTLP exponential histogram temporality; cumulative temporality is required" + ) + })?; + return Ok(false); + } + } + + let column_schema = native_histogram_column_schema().map_err(|error| { + error::InternalSnafu { + err_msg: error.to_string(), + } + .build() + })?; + let mut emitted = false; + for (index, data_point) in histogram.data_points.iter().enumerate() { + let (value, timestamp_nanos) = match exponential_histogram_value(data_point) { + Ok(value) => value, + Err(reason) => { + reject_data_points(outcome, 1, || { + format!("metric `{name}` data point {index}: {reason}") + })?; + continue; + } + }; + + let table = table_writer.get_or_default_table_data( + name, + APPROXIMATE_COLUMN_COUNT, + histogram.data_points.len(), + ); + let mut row = table.alloc_one_row(); + write_tags_and_timestamp( + table, + &mut row, + resource_attrs, + scope_attrs, + Some(data_point.attributes.as_ref()), + timestamp_nanos, + metric_ctx, + )?; + row_writer::write_by_schema( + table, + std::iter::once((column_schema.clone(), Some(value))), + &mut row, + )?; + table.add_row(row); + add_accepted_data_points(outcome, 1)?; + emitted = true; + } + + Ok(emitted) +} + +fn exponential_histogram_value( + data_point: &ExponentialHistogramDataPoint, +) -> std::result::Result<(ValueData, i64), String> { + if data_point.start_time_unix_nano > data_point.time_unix_nano { + return Err(format!( + "start_time_unix_nano {} exceeds time_unix_nano {}", + data_point.start_time_unix_nano, data_point.time_unix_nano + )); + } + + let timestamp_nanos = i64::try_from(data_point.time_unix_nano) + .map_err(|_| format!("time_unix_nano {} overflows i64", data_point.time_unix_nano))?; + let timestamp = timestamp_millis(data_point.time_unix_nano, "time_unix_nano")?; + let start_timestamp = + timestamp_millis(data_point.start_time_unix_nano, "start_time_unix_nano")?; + + let no_recorded_value = data_point.flags & DataPointFlags::NoRecordedValueMask as u32 != 0; + let histogram = if no_recorded_value { + PromHistogram { + sum: f64::from_bits(PROMETHEUS_STALE_NAN_BITS), + schema: 0, + zero_threshold: 0.0, + reset_hint: ResetHint::Unspecified as i32, + timestamp, + start_timestamp, + count: Some(PromCount::CountInt(0)), + zero_count: Some(PromZeroCount::ZeroCountInt(0)), + ..Default::default() + } + } else { + if data_point.scale < MIN_EXPONENTIAL_HISTOGRAM_SCALE { + return Err(format!( + "scale {} is unsupported; minimum supported scale is {}", + data_point.scale, MIN_EXPONENTIAL_HISTOGRAM_SCALE + )); + } + if !data_point.zero_threshold.is_finite() || data_point.zero_threshold < 0.0 { + return Err(format!( + "zero_threshold {} must be finite and non-negative", + data_point.zero_threshold + )); + } + + let downscale_shift = if data_point.scale > MAX_EXPONENTIAL_HISTOGRAM_SCALE { + let shift = data_point + .scale + .checked_sub(MAX_EXPONENTIAL_HISTOGRAM_SCALE) + .ok_or_else(|| "downscale shift overflows i32".to_string())?; + u32::try_from(shift).map_err(|_| "downscale shift exceeds u32".to_string())? + } else { + 0 + }; + let (positive_spans, positive_deltas, positive_count) = + convert_bucket_range("positive", data_point.positive.as_ref(), downscale_shift)?; + let (negative_spans, negative_deltas, negative_count) = + convert_bucket_range("negative", data_point.negative.as_ref(), downscale_shift)?; + let bucket_count = data_point + .zero_count + .checked_add(positive_count) + .and_then(|count| count.checked_add(negative_count)) + .ok_or_else(|| "bucket observation total overflows u64".to_string())?; + if bucket_count != data_point.count { + return Err(format!( + "buckets contain {bucket_count} observations, declared count is {}", + data_point.count + )); + } + + i64::try_from(data_point.count) + .map_err(|_| format!("count {} overflows i64", data_point.count))?; + i64::try_from(data_point.zero_count) + .map_err(|_| format!("zero_count {} overflows i64", data_point.zero_count))?; + + let sum = match data_point.sum { + Some(sum) if sum.is_nan() => f64::NAN, + Some(sum) => sum, + None => f64::NAN, + }; + PromHistogram { + sum, + schema: data_point.scale.min(MAX_EXPONENTIAL_HISTOGRAM_SCALE), + zero_threshold: data_point.zero_threshold, + negative_spans, + negative_deltas, + positive_spans, + positive_deltas, + reset_hint: ResetHint::Unspecified as i32, + timestamp, + start_timestamp, + count: Some(PromCount::CountInt(data_point.count)), + zero_count: Some(PromZeroCount::ZeroCountInt(data_point.zero_count)), + ..Default::default() + } + }; + + encode_native_histogram(&histogram) + .map(|value| (value, timestamp_nanos)) + .map_err(|error| format!("OTLP exponential histogram cannot be encoded: {error}")) +} + +fn timestamp_millis(timestamp_nanos: u64, name: &str) -> std::result::Result { + i64::try_from(timestamp_nanos / 1_000_000) + .map_err(|_| format!("{name} {timestamp_nanos} milliseconds overflow i64")) +} + +fn convert_bucket_range( + name: &str, + buckets: Option<&exponential_histogram_data_point::Buckets>, + downscale_shift: u32, +) -> std::result::Result<(Vec, Vec, u64), String> { + let Some(buckets) = buckets else { + return Ok((Vec::new(), Vec::new(), 0)); + }; + if buckets.bucket_counts.is_empty() { + return Ok((Vec::new(), Vec::new(), 0)); + } + + let mut merged = Vec::<(i32, u64)>::with_capacity(buckets.bucket_counts.len()); + let mut total = 0u64; + for (position, count) in buckets.bucket_counts.iter().copied().enumerate() { + let position = + i32::try_from(position).map_err(|_| format!("{name} bucket length exceeds i32"))?; + let source_index = buckets.offset.checked_add(position).ok_or_else(|| { + format!( + "{name} bucket index overflows i32 at offset {} and position {position}", + buckets.offset + ) + })?; + let target_index = downscale_bucket_index(source_index, downscale_shift)? + .checked_add(1) + .ok_or_else(|| format!("{name} shifted bucket index overflows i32"))?; + if let Some((last_index, last_count)) = merged.last_mut() { + if *last_index == target_index { + *last_count = last_count + .checked_add(count) + .ok_or_else(|| format!("{name} merged bucket count overflows u64"))?; + total = total + .checked_add(count) + .ok_or_else(|| format!("{name} bucket count total overflows u64"))?; + continue; + } + let next_index = last_index + .checked_add(1) + .ok_or_else(|| format!("{name} target bucket index overflows i32"))?; + if target_index != next_index { + return Err(format!( + "{name} bucket indexes are not contiguous after downscaling" + )); + } + } + total = total + .checked_add(count) + .ok_or_else(|| format!("{name} bucket count total overflows u64"))?; + merged.push((target_index, count)); + } + + let length = u32::try_from(merged.len()) + .map_err(|_| format!("{name} bucket span length exceeds u32"))?; + let span = BucketSpan { + offset: merged[0].0, + length, + }; + let mut deltas = Vec::with_capacity(merged.len()); + let mut previous = 0i64; + for (_, count) in merged { + let count = i64::try_from(count) + .map_err(|_| format!("{name} bucket count {count} overflows i64"))?; + let delta = count + .checked_sub(previous) + .ok_or_else(|| format!("{name} bucket delta overflows i64"))?; + deltas.push(delta); + previous = count; + } + + Ok((vec![span], deltas, total)) +} + +fn downscale_bucket_index(index: i32, downscale_shift: u32) -> std::result::Result { + if downscale_shift == 0 { + return Ok(index); + } + if downscale_shift >= i32::BITS { + return Ok(if index < 0 { -1 } else { 0 }); + } + + let divisor = 1i64 + .checked_shl(downscale_shift) + .ok_or_else(|| format!("downscale shift {downscale_shift} is invalid"))?; + i32::try_from(i64::from(index).div_euclid(divisor)) + .map_err(|_| format!("downscaled bucket index {index} overflows i32")) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum AttributeType { Resource, @@ -620,17 +1044,17 @@ fn encode_histogram( let count_table_name = format!("{}{}", normalized_name, COUNT_TABLE_SUFFIX); let data_points_len = hist.data_points.len(); - // Note that the row and columns number here is approximate - let mut bucket_table = TableData::new(APPROXIMATE_COLUMN_COUNT, data_points_len * 3); - let mut sum_table = TableData::new(APPROXIMATE_COLUMN_COUNT, data_points_len); - let mut count_table = TableData::new(APPROXIMATE_COLUMN_COUNT, data_points_len); - for data_point in &hist.data_points { + let bucket_table = table_writer.get_or_default_table_data( + &bucket_table_name, + APPROXIMATE_COLUMN_COUNT, + data_points_len * 3, + ); let mut accumulated_count = 0; for (idx, count) in data_point.bucket_counts.iter().enumerate() { let mut bucket_row = bucket_table.alloc_one_row(); write_tags_and_timestamp( - &mut bucket_table, + bucket_table, &mut bucket_row, resource_attrs, scope_attrs, @@ -641,7 +1065,7 @@ fn encode_histogram( if let Some(upper_bounds) = data_point.explicit_bounds.get(idx) { row_writer::write_tag( - &mut bucket_table, + bucket_table, HISTOGRAM_LE_COLUMN, upper_bounds, &mut bucket_row, @@ -649,7 +1073,7 @@ fn encode_histogram( } else if idx == data_point.explicit_bounds.len() { // The last bucket row_writer::write_tag( - &mut bucket_table, + bucket_table, HISTOGRAM_LE_COLUMN, f64::INFINITY, &mut bucket_row, @@ -658,7 +1082,7 @@ fn encode_histogram( accumulated_count += count; row_writer::write_f64( - &mut bucket_table, + bucket_table, greptime_value(), accumulated_count as f64, &mut bucket_row, @@ -668,9 +1092,14 @@ fn encode_histogram( } if let Some(sum) = data_point.sum { + let sum_table = table_writer.get_or_default_table_data( + &sum_table_name, + APPROXIMATE_COLUMN_COUNT, + data_points_len, + ); let mut sum_row = sum_table.alloc_one_row(); write_tags_and_timestamp( - &mut sum_table, + sum_table, &mut sum_row, resource_attrs, scope_attrs, @@ -679,13 +1108,18 @@ fn encode_histogram( metric_ctx, )?; - row_writer::write_f64(&mut sum_table, greptime_value(), sum, &mut sum_row)?; + row_writer::write_f64(sum_table, greptime_value(), sum, &mut sum_row)?; sum_table.add_row(sum_row); } + let count_table = table_writer.get_or_default_table_data( + &count_table_name, + APPROXIMATE_COLUMN_COUNT, + data_points_len, + ); let mut count_row = count_table.alloc_one_row(); write_tags_and_timestamp( - &mut count_table, + count_table, &mut count_row, resource_attrs, scope_attrs, @@ -695,7 +1129,7 @@ fn encode_histogram( )?; row_writer::write_f64( - &mut count_table, + count_table, greptime_value(), data_point.count as f64, &mut count_row, @@ -703,10 +1137,6 @@ fn encode_histogram( count_table.add_row(count_row); } - table_writer.add_table_data(bucket_table_name, bucket_table); - table_writer.add_table_data(sum_table_name, sum_table); - table_writer.add_table_data(count_table_name, count_table); - Ok(()) } @@ -851,6 +1281,7 @@ mod tests { use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ AggregationTemporality, HistogramDataPoint, NumberDataPoint, SummaryDataPoint, }; + use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; use super::*; @@ -1361,4 +1792,560 @@ mod tests { ); } } + + fn exponential_buckets( + offset: i32, + bucket_counts: Vec, + ) -> exponential_histogram_data_point::Buckets { + exponential_histogram_data_point::Buckets { + offset, + bucket_counts, + } + } + + fn exponential_point() -> ExponentialHistogramDataPoint { + ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + count: 28, + sum: None, + scale: 9, + zero_count: 7, + positive: Some(exponential_buckets(-3, vec![1, 2, 3, 4, 5, 6])), + zero_threshold: 0.0, + ..Default::default() + } + } + + fn native_field(value: &ValueData, name: &str) -> Option { + let ValueData::StructValue(value) = value else { + panic!("expected native histogram Struct value"); + }; + let index = common_query::native_histogram::NATIVE_HISTOGRAM_FIELD_NAMES + .iter() + .position(|field| *field == name) + .unwrap(); + value.items[index].value_data.clone() + } + + fn i32_list(value: Option) -> Vec { + let Some(ValueData::ListValue(value)) = value else { + panic!("expected i32 list"); + }; + value + .items + .into_iter() + .map(|item| match item.value_data { + Some(ValueData::I32Value(value)) => value, + _ => panic!("expected i32 value"), + }) + .collect() + } + + fn i64_list(value: Option) -> Vec { + let Some(ValueData::ListValue(value)) = value else { + panic!("expected i64 list"); + }; + value + .items + .into_iter() + .map(|item| match item.value_data { + Some(ValueData::I64Value(value)) => value, + _ => panic!("expected i64 value"), + }) + .collect() + } + + #[test] + fn test_downscale_bucket_index_uses_signed_floor_division() { + assert_eq!(downscale_bucket_index(-3, 1).unwrap(), -2); + assert_eq!(downscale_bucket_index(-2, 1).unwrap(), -1); + assert_eq!(downscale_bucket_index(-1, 1).unwrap(), -1); + assert_eq!(downscale_bucket_index(0, 1).unwrap(), 0); + assert_eq!(downscale_bucket_index(1, 1).unwrap(), 0); + assert_eq!(downscale_bucket_index(2, 1).unwrap(), 1); + assert_eq!(downscale_bucket_index(i32::MIN, 32).unwrap(), -1); + assert_eq!(downscale_bucket_index(i32::MAX, 32).unwrap(), 0); + } + + #[test] + fn test_convert_bucket_range_downscales_before_prometheus_shift() { + let buckets = exponential_buckets(-3, vec![1, 2, 3, 4, 5, 6]); + let (spans, deltas, total) = convert_bucket_range("positive", Some(&buckets), 1).unwrap(); + + assert_eq!( + spans, + vec![BucketSpan { + offset: -1, + length: 4 + }] + ); + assert_eq!(deltas, vec![1, 4, 4, -3]); + assert_eq!(total, 21); + } + + #[test] + fn test_exponential_histogram_value_uses_integer_family() { + use common_query::native_histogram::{ + COUNT_F64_FIELD, COUNT_I64_FIELD, POSITIVE_BUCKETS_I64_FIELD, + POSITIVE_SPAN_LENGTHS_FIELD, POSITIVE_SPAN_OFFSETS_FIELD, SCHEMA_FIELD, SUM_FIELD, + ZERO_COUNT_I64_FIELD, + }; + + let (value, timestamp_nanos) = exponential_histogram_value(&exponential_point()).unwrap(); + + assert_eq!(timestamp_nanos, 2_000_000); + assert_eq!( + native_field(&value, SCHEMA_FIELD), + Some(ValueData::I32Value(8)) + ); + assert_eq!( + native_field(&value, COUNT_I64_FIELD), + Some(ValueData::I64Value(28)) + ); + assert_eq!( + native_field(&value, ZERO_COUNT_I64_FIELD), + Some(ValueData::I64Value(7)) + ); + assert_eq!(native_field(&value, COUNT_F64_FIELD), None); + assert_eq!( + i32_list(native_field(&value, POSITIVE_SPAN_OFFSETS_FIELD)), + vec![-1] + ); + assert_eq!( + i32_list(native_field(&value, POSITIVE_SPAN_LENGTHS_FIELD)), + vec![4] + ); + assert_eq!( + i64_list(native_field(&value, POSITIVE_BUCKETS_I64_FIELD)), + vec![1, 5, 9, 6] + ); + let Some(ValueData::F64Value(sum)) = native_field(&value, SUM_FIELD) else { + panic!("expected histogram sum"); + }; + assert_eq!(sum.to_bits(), f64::NAN.to_bits()); + assert_ne!(sum.to_bits(), PROMETHEUS_STALE_NAN_BITS); + } + + #[test] + fn test_no_recorded_value_ignores_other_value_fields() { + use common_query::native_histogram::{ + COUNT_I64_FIELD, POSITIVE_BUCKETS_I64_FIELD, SCHEMA_FIELD, SUM_FIELD, + ZERO_COUNT_I64_FIELD, ZERO_THRESHOLD_FIELD, + }; + + let point = ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + count: u64::MAX, + sum: Some(1.0), + scale: i32::MIN, + zero_count: u64::MAX, + positive: Some(exponential_buckets(i32::MAX, vec![u64::MAX])), + flags: DataPointFlags::NoRecordedValueMask as u32, + zero_threshold: f64::NAN, + ..Default::default() + }; + let (value, _) = exponential_histogram_value(&point).unwrap(); + + assert_eq!( + native_field(&value, SCHEMA_FIELD), + Some(ValueData::I32Value(0)) + ); + assert_eq!( + native_field(&value, ZERO_THRESHOLD_FIELD), + Some(ValueData::F64Value(0.0)) + ); + assert_eq!( + native_field(&value, COUNT_I64_FIELD), + Some(ValueData::I64Value(0)) + ); + assert_eq!( + native_field(&value, ZERO_COUNT_I64_FIELD), + Some(ValueData::I64Value(0)) + ); + assert!(i64_list(native_field(&value, POSITIVE_BUCKETS_I64_FIELD)).is_empty()); + let Some(ValueData::F64Value(sum)) = native_field(&value, SUM_FIELD) else { + panic!("expected histogram sum"); + }; + assert_eq!(sum.to_bits(), PROMETHEUS_STALE_NAN_BITS); + } + + #[test] + fn test_non_flag_nan_sum_is_normalized() { + use common_query::native_histogram::SUM_FIELD; + + let point = ExponentialHistogramDataPoint { + sum: Some(f64::from_bits(PROMETHEUS_STALE_NAN_BITS)), + ..Default::default() + }; + let (value, _) = exponential_histogram_value(&point).unwrap(); + let Some(ValueData::F64Value(sum)) = native_field(&value, SUM_FIELD) else { + panic!("expected histogram sum"); + }; + assert_eq!(sum.to_bits(), f64::NAN.to_bits()); + assert_ne!(sum.to_bits(), PROMETHEUS_STALE_NAN_BITS); + } + + #[test] + fn test_exponential_histogram_rejects_invalid_values() { + let mut cases = Vec::new(); + + let mut point = exponential_point(); + point.scale = -5; + cases.push((point, "scale -5 is unsupported")); + + let mut point = exponential_point(); + point.zero_threshold = f64::INFINITY; + cases.push((point, "must be finite and non-negative")); + + let mut point = exponential_point(); + point.start_time_unix_nano = point.time_unix_nano + 1; + cases.push((point, "start_time_unix_nano")); + + let mut point = exponential_point(); + point.count = 27; + cases.push((point, "declared count is 27")); + + let point = ExponentialHistogramDataPoint { + count: u64::MAX, + zero_count: u64::MAX, + ..Default::default() + }; + cases.push((point, "count 18446744073709551615 overflows i64")); + + let point = ExponentialHistogramDataPoint { + count: 1, + scale: 8, + positive: Some(exponential_buckets(i32::MAX, vec![1])), + ..Default::default() + }; + cases.push((point, "shifted bucket index overflows i32")); + + for (point, expected) in cases { + let error = exponential_histogram_value(&point).unwrap_err(); + assert!( + error.contains(expected), + "expected {expected:?}, got {error}" + ); + } + + let buckets = exponential_buckets(0, vec![u64::MAX, 1]); + let error = convert_bucket_range("positive", Some(&buckets), 1).unwrap_err(); + assert!( + error.contains("merged bucket count overflows u64"), + "{error}" + ); + + let buckets = exponential_buckets(i32::MAX, vec![1, 1]); + let error = convert_bucket_range("positive", Some(&buckets), 1).unwrap_err(); + assert!(error.contains("bucket index overflows i32"), "{error}"); + } + + fn metrics_request(metrics: Vec) -> ExportMetricsServiceRequest { + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics, + ..Default::default() + }], + ..Default::default() + }], + } + } + + fn exponential_metric( + name: impl Into, + data_points: Vec, + temporality: AggregationTemporality, + ) -> Metric { + Metric { + name: name.into(), + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points, + aggregation_temporality: temporality as i32, + })), + ..Default::default() + } + } + + fn histogram_metric(name: impl Into) -> Metric { + Metric { + name: name.into(), + data: Some(metric::Data::Histogram(Histogram { + data_points: vec![HistogramDataPoint { + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + count: 1, + sum: Some(1.0), + bucket_counts: vec![1], + ..Default::default() + }], + aggregation_temporality: AggregationTemporality::Cumulative as i32, + })), + ..Default::default() + } + } + + #[test] + fn test_exponential_histogram_gate_and_partial_outcome() { + let request = metrics_request(vec![ + Metric { + name: "temperature".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint::default()], + })), + ..Default::default() + }, + exponential_metric( + "latency", + vec![exponential_point()], + AggregationTemporality::Cumulative, + ), + ]); + let (requests, _, semantic_index, outcome) = + to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + + assert_eq!(outcome.accepted_data_points, 1); + assert_eq!(outcome.rejected_data_points, 1); + assert!( + outcome + .error_message + .as_deref() + .unwrap() + .contains("otlp.experimental_enable_exponential_histogram") + ); + assert_eq!(requests.inserts.len(), 1); + assert_eq!(requests.inserts[0].table_name, "temperature"); + let semantics = decode(&semantic_index); + assert!(semantics.contains_key("temperature")); + assert!(!semantics.contains_key("latency")); + + let empty = metrics_request(vec![exponential_metric( + "empty", + vec![], + AggregationTemporality::Cumulative, + )]); + let (_, _, _, outcome) = + to_grpc_insert_requests(empty, &mut OtlpMetricCtx::default()).unwrap(); + assert_eq!(outcome.rejected_data_points, 0); + assert_eq!(outcome.error_message, None); + } + + #[test] + fn test_exponential_histogram_cannot_share_table_with_scalar_metric() { + let request = metrics_request(vec![ + Metric { + name: "latency".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + value: Some(number_data_point::Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }, + exponential_metric( + "latency", + vec![exponential_point()], + AggregationTemporality::Cumulative, + ), + ]); + let mut ctx = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + + let error = to_grpc_insert_requests(request, &mut ctx).unwrap_err(); + assert!( + error + .to_string() + .contains("cannot mix native histogram and float sample fields") + ); + } + + #[test] + fn test_histogram_cannot_replace_exponential_histogram_table() { + let request = metrics_request(vec![ + exponential_metric( + "latency_bucket", + vec![exponential_point()], + AggregationTemporality::Cumulative, + ), + histogram_metric("latency"), + ]); + let mut ctx = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + + let error = to_grpc_insert_requests(request, &mut ctx).unwrap_err(); + assert!( + error + .to_string() + .contains("cannot mix native histogram and float sample fields") + ); + } + + #[test] + fn test_histograms_with_same_name_across_resources_are_merged() { + let metric = histogram_metric("latency"); + let request = ExportMetricsServiceRequest { + resource_metrics: ["service-a", "service-b"] + .into_iter() + .map(|service| ResourceMetrics { + resource: Some(Resource { + attributes: vec![keyvalue("service.name", service)], + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![metric.clone()], + ..Default::default() + }], + ..Default::default() + }) + .collect(), + }; + + let (requests, rows, _, outcome) = + to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + + assert_eq!(outcome.accepted_data_points, 2); + assert_eq!(rows, 6); + assert_eq!(requests.inserts.len(), 3); + for request in requests.inserts { + assert_eq!( + request.rows.unwrap().rows.len(), + 2, + "{}", + request.table_name + ); + } + } + + #[test] + fn test_exponential_histogram_rejects_temporality_before_stale_point() { + let stale = ExponentialHistogramDataPoint { + flags: DataPointFlags::NoRecordedValueMask as u32, + ..Default::default() + }; + for temporality in [ + AggregationTemporality::Delta, + AggregationTemporality::Unspecified, + ] { + let request = metrics_request(vec![exponential_metric( + "latency", + vec![stale.clone()], + temporality, + )]); + let mut ctx = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + let (requests, rows, semantic_index, outcome) = + to_grpc_insert_requests(request, &mut ctx).unwrap(); + + assert_eq!(outcome.accepted_data_points, 0); + assert_eq!(outcome.rejected_data_points, 1); + assert_eq!(rows, 0); + assert!(requests.inserts.is_empty()); + assert!(semantic_index.is_empty()); + } + } + + #[test] + fn test_exponential_histogram_legacy_and_new_modes_share_struct() { + use common_query::prelude::greptime_native_histogram; + + let mut point = exponential_point(); + point.sum = Some(42.0); + let request = metrics_request(vec![exponential_metric( + "request.duration", + vec![point], + AggregationTemporality::Cumulative, + )]); + let mut new_ctx = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + let (new_requests, _, _, _) = + to_grpc_insert_requests(request.clone(), &mut new_ctx).unwrap(); + let mut legacy_ctx = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + is_legacy: true, + ..Default::default() + }; + let (legacy_requests, _, _, _) = to_grpc_insert_requests(request, &mut legacy_ctx).unwrap(); + + let new_insert = &new_requests.inserts[0]; + let legacy_insert = &legacy_requests.inserts[0]; + assert_eq!(new_insert.table_name, "request_duration"); + assert_eq!(legacy_insert.table_name, "request_duration"); + let new_rows = new_insert.rows.as_ref().unwrap(); + let legacy_rows = legacy_insert.rows.as_ref().unwrap(); + let field = greptime_native_histogram(); + let new_histogram = new_rows.rows[0].values[new_rows + .schema + .iter() + .position(|column| column.column_name == field) + .unwrap()] + .clone(); + let legacy_histogram = legacy_rows.rows[0].values[legacy_rows + .schema + .iter() + .position(|column| column.column_name == field) + .unwrap()] + .clone(); + assert_eq!(new_histogram, legacy_histogram); + assert!(matches!( + new_rows.rows[0].values[new_rows + .schema + .iter() + .position(|column| column.column_name == greptime_timestamp()) + .unwrap()] + .value_data, + Some(ValueData::TimestampMillisecondValue(2)) + )); + assert!(matches!( + legacy_rows.rows[0].values[legacy_rows + .schema + .iter() + .position(|column| column.column_name == greptime_timestamp()) + .unwrap()] + .value_data, + Some(ValueData::TimestampNanosecondValue(2_000_000)) + )); + } + + #[test] + fn test_rejection_message_is_bounded() { + let request = metrics_request(vec![exponential_metric( + "x".repeat(1_000), + vec![ExponentialHistogramDataPoint::default()], + AggregationTemporality::Cumulative, + )]); + let (_, _, _, outcome) = + to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + + assert_eq!(outcome.rejected_data_points, 1); + assert!(outcome.error_message.unwrap().len() <= MAX_REJECTION_MESSAGE_BYTES); + } + + #[test] + fn test_rejection_message_reason_is_lazy_after_cap() { + let mut outcome = MetricsIngestOutcome { + error_message: Some("x".repeat(MAX_REJECTION_MESSAGE_BYTES)), + ..Default::default() + }; + let mut reason_built = false; + + reject_data_points(&mut outcome, 1, || { + reason_built = true; + "unused".to_string() + }) + .unwrap(); + + assert_eq!(outcome.rejected_data_points, 1); + assert!(!reason_built); + } } diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 1bcec98d5f5..846ef3ee6be 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -15,19 +15,23 @@ use std::collections::hash_map::Entry; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; +#[cfg(test)] +use api::greptime_proto::io::prometheus::write::v2::BucketSpan; +#[cfg(test)] use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; use api::greptime_proto::io::prometheus::write::v2::{ - BucketSpan, Exemplar, Histogram, Metadata, Sample, metadata, + Exemplar, Histogram, Metadata, Sample, metadata, }; #[cfg(test)] use api::greptime_proto::io::prometheus::write::v2::{Request, TimeSeries}; -use api::helper::ColumnDataTypeWrapper; +#[cfg(test)] +use api::v1::ColumnSchema; use api::v1::value::ValueData; -use api::v1::{ - ColumnDataType, ColumnSchema, ListValue, RowInsertRequest, Rows, SemanticType, Value, -}; +use api::v1::{ColumnDataType, RowInsertRequest, Rows, SemanticType, Value}; use bytes::{Buf, Bytes}; use common_grpc::precision::Precision; +use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; +#[cfg(test)] use common_query::native_histogram::*; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use pipeline::{ContextOpt, ContextReq}; @@ -42,6 +46,7 @@ use table::requests::{ }; use crate::error::{self, Result}; +use crate::native_histogram::{encode_native_histogram, native_histogram_column_schema}; use crate::prom_remote_write::row_builder::PromCtx; use crate::prom_remote_write::validation::validate_label_name; use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, try_decompress}; @@ -59,8 +64,6 @@ use crate::semantic::{ type PromTags<'a> = Vec<(&'a str, String)>; type ResolvedSeriesLabels<'a> = (PromCtx, String, PromTags<'a>); -const MAX_REMOTE_WRITE_V2_SCHEMA: i32 = 8; -const MAX_REDUCIBLE_REMOTE_WRITE_V2_SCHEMA: i32 = 52; const TIME_SERIES_LABELS_REFS_TAG: u32 = 1; const TIME_SERIES_SAMPLES_TAG: u32 = 2; const TIME_SERIES_HISTOGRAMS_TAG: u32 = 3; @@ -597,6 +600,19 @@ fn write_native_histogram<'a>( histogram: &Histogram, tags: impl Iterator, ) -> Result<()> { + let value = encode_native_histogram(histogram).map_err(|error| { + error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {error}"), + } + .build() + })?; + let column_schema = native_histogram_column_schema().map_err(|error| { + error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {error}"), + } + .build() + })?; + // Persist both int and float families into the logical table schema. Only one // family is populated per row; the other is written as NULL so PromQL can // infer the original histogram flavor without a separate type column. @@ -608,8 +624,11 @@ fn write_native_histogram<'a>( Precision::Millisecond, &mut row, )?; - - write_native_histogram_value(table_data, histogram, &mut row)?; + row_writer::write_by_schema( + table_data, + std::iter::once((column_schema, Some(value))), + &mut row, + )?; row_writer::write_tags(table_data, tags, &mut row)?; table_data.add_row(row); @@ -617,512 +636,6 @@ fn write_native_histogram<'a>( Ok(()) } -fn write_native_histogram_value( - table_data: &mut TableData, - histogram: &Histogram, - row: &mut Vec, -) -> Result<()> { - let column_schema = native_histogram_column_schema(); - let value = native_histogram_struct_value(histogram)?; - - row_writer::write_by_schema( - table_data, - std::iter::once((column_schema, Some(value))), - row, - ) -} - -fn native_histogram_column_schema() -> ColumnSchema { - let (datatype, datatype_extension) = - ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone()) - .expect("native histogram type is convertible to protobuf") - .into_parts(); - - ColumnSchema { - column_name: greptime_native_histogram().to_string(), - datatype: datatype as i32, - semantic_type: SemanticType::Field as i32, - datatype_extension, - options: None, - } -} - -fn native_histogram_struct_value(histogram: &Histogram) -> Result { - let uses_float_counts = native_histogram_uses_float_counts(histogram)?; - validate_native_histogram(histogram, uses_float_counts)?; - - let mut items = Vec::with_capacity(NATIVE_HISTOGRAM_FIELD_NAMES.len()); - let positive_span_lengths = i32_span_lengths("positive", &histogram.positive_spans)?; - let negative_span_lengths = i32_span_lengths("negative", &histogram.negative_spans)?; - items.extend([ - pb_value(ValueData::I32Value(histogram.schema)), - pb_value(ValueData::F64Value(histogram.zero_threshold)), - pb_value(ValueData::F64Value(histogram.sum)), - pb_value(ValueData::I32Value(histogram.reset_hint)), - optional_pb_value((histogram.start_timestamp != 0).then_some( - ValueData::TimestampMillisecondValue(histogram.start_timestamp), - )), - f64_list_value(histogram.custom_values.iter().copied()), - i32_list_value(histogram.positive_spans.iter().map(|span| span.offset)), - i32_list_value(positive_span_lengths.iter().copied()), - i32_list_value(histogram.negative_spans.iter().map(|span| span.offset)), - i32_list_value(negative_span_lengths.iter().copied()), - ]); - - if uses_float_counts { - validate_float_native_histogram_counts(histogram)?; - let count = match histogram.count.as_ref() { - Some(Count::CountFloat(count)) => *count, - _ => 0.0, - }; - let zero_count = match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count, - _ => 0.0, - }; - items.extend([ - null_pb_value(), - null_pb_value(), - i64_list_value(std::iter::empty()), - i64_list_value(std::iter::empty()), - pb_value(ValueData::F64Value(count)), - pb_value(ValueData::F64Value(zero_count)), - f64_list_value(histogram.positive_counts.iter().copied()), - f64_list_value(histogram.negative_counts.iter().copied()), - ]); - } else { - let count = match histogram.count.as_ref() { - Some(Count::CountInt(count)) => *count, - _ => 0, - }; - let zero_count = match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count, - _ => 0, - }; - let positive_buckets = bucket_counts_from_deltas(&histogram.positive_deltas)?; - let negative_buckets = bucket_counts_from_deltas(&histogram.negative_deltas)?; - validate_integer_native_histogram_counts(histogram, &positive_buckets, &negative_buckets)?; - let count = i64::try_from(count) - .ok() - .context(error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram integer count {count} overflows i64" - ), - })?; - let zero_count = i64::try_from(zero_count).ok().context( - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram integer zero_count {zero_count} overflows i64" - ), - }, - )?; - items.extend([ - pb_value(ValueData::I64Value(count)), - pb_value(ValueData::I64Value(zero_count)), - i64_list_value(positive_buckets.iter().copied()), - i64_list_value(negative_buckets.iter().copied()), - null_pb_value(), - null_pb_value(), - f64_list_value(std::iter::empty()), - f64_list_value(std::iter::empty()), - ]); - } - - Ok(ValueData::StructValue(api::v1::StructValue { items })) -} - -fn validate_native_histogram(histogram: &Histogram, uses_float_counts: bool) -> Result<()> { - let exponential_overflow_index = validate_native_histogram_schema(histogram.schema)?; - validate_native_histogram_custom_values(histogram)?; - - if histogram.schema == CUSTOM_BUCKETS_SCHEMA { - ensure!( - histogram.zero_threshold == 0.0 && native_histogram_zero_count_is_zero(histogram), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 custom native histogram must not use a zero bucket" - .to_string(), - } - ); - ensure!( - histogram.negative_spans.is_empty() - && histogram.negative_deltas.is_empty() - && histogram.negative_counts.is_empty(), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 custom native histogram must not use negative buckets" - .to_string(), - } - ); - } - - let (positive_buckets, negative_buckets) = if uses_float_counts { - ( - histogram.positive_counts.len(), - histogram.negative_counts.len(), - ) - } else { - ( - histogram.positive_deltas.len(), - histogram.negative_deltas.len(), - ) - }; - let bucket_index_range = if let Some(overflow_index) = exponential_overflow_index { - (i32::MIN, overflow_index) - } else { - ( - 0, - i32::try_from(histogram.custom_values.len()).ok().context( - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 custom native histogram has too many custom_values" - .to_string(), - }, - )?, - ) - }; - validate_native_histogram_spans( - "positive", - &histogram.positive_spans, - positive_buckets, - bucket_index_range, - )?; - validate_native_histogram_spans( - "negative", - &histogram.negative_spans, - negative_buckets, - bucket_index_range, - )?; - - Ok(()) -} - -fn validate_native_histogram_schema(schema: i32) -> Result> { - if schema == CUSTOM_BUCKETS_SCHEMA { - return Ok(None); - } - - if let Some(overflow_index) = exponential_overflow_bucket_index(schema) { - return Ok(Some(overflow_index)); - } - - if (MAX_REMOTE_WRITE_V2_SCHEMA + 1..=MAX_REDUCIBLE_REMOTE_WRITE_V2_SCHEMA).contains(&schema) { - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram schema {schema} must be reduced before ingestion" - ), - } - .fail() - } else { - error::InvalidPromRemoteRequestSnafu { - msg: format!("remote write v2 native histogram schema {schema} is unsupported"), - } - .fail() - } -} - -fn validate_native_histogram_custom_values(histogram: &Histogram) -> Result<()> { - if histogram.schema != CUSTOM_BUCKETS_SCHEMA { - ensure!( - histogram.custom_values.is_empty(), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 standard native histogram must not use custom_values" - .to_string(), - } - ); - return Ok(()); - } - - for value in &histogram.custom_values { - ensure!( - !value.is_nan() && *value != f64::INFINITY, - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 custom native histogram custom_values must not contain +Inf or NaN" - .to_string(), - } - ); - } - for values in histogram.custom_values.windows(2) { - ensure!( - values[0] < values[1], - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 custom native histogram custom_values must be sorted" - .to_string(), - } - ); - } - - Ok(()) -} - -fn validate_native_histogram_spans( - name: &str, - spans: &[BucketSpan], - bucket_count: usize, - bucket_index_range: (i32, i32), -) -> Result<()> { - let span_len = spans - .iter() - .try_fold(0usize, |sum, span| sum.checked_add(span.length as usize)) - .with_context(|| error::InvalidPromRemoteRequestSnafu { - msg: format!("remote write v2 native histogram {name} spans overflow"), - })?; - ensure!( - span_len == bucket_count, - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} spans describe {span_len} buckets, found {bucket_count}" - ), - } - ); - - let mut current_index = 0i32; - for (span_index, span) in spans.iter().enumerate() { - ensure!( - span.offset >= 0 || (span_index == 0 && bucket_index_range.0 == i32::MIN), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} span {} has negative offset {}", - span_index + 1, - span.offset - ), - } - ); - current_index = if span_index == 0 { - span.offset - } else { - current_index.checked_add(span.offset).with_context(|| { - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} span index overflows i32" - ), - } - })? - }; - - for _ in 0..span.length { - ensure!( - (bucket_index_range.0..=bucket_index_range.1).contains(¤t_index), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} bucket index {current_index} is out of range" - ), - } - ); - current_index = - current_index - .checked_add(1) - .context(error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} span index overflows i32" - ), - })?; - } - } - - Ok(()) -} - -fn validate_float_native_histogram_counts(histogram: &Histogram) -> Result<()> { - let count = match histogram.count.as_ref() { - Some(Count::CountFloat(count)) => *count, - _ => 0.0, - }; - ensure!( - count >= 0.0 || count.is_nan(), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram float count must not be negative".to_string(), - } - ); - - let zero_count = match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count, - _ => 0.0, - }; - ensure!( - zero_count >= 0.0 || zero_count.is_nan(), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram float zero_count must not be negative" - .to_string(), - } - ); - - for (name, counts) in [ - ("positive", &histogram.positive_counts), - ("negative", &histogram.negative_counts), - ] { - for (index, count) in counts.iter().enumerate() { - ensure!( - *count >= 0.0 || count.is_nan(), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} bucket {} count must not be negative", - index + 1 - ), - } - ); - } - } - - Ok(()) -} - -fn validate_integer_native_histogram_counts( - histogram: &Histogram, - positive_buckets: &[i64], - negative_buckets: &[i64], -) -> Result<()> { - let count = match histogram.count.as_ref() { - Some(Count::CountInt(count)) => *count, - _ => 0, - }; - let zero_count = match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count, - _ => 0, - }; - let bucket_count = positive_buckets - .iter() - .chain(negative_buckets) - .try_fold(zero_count, |total, bucket| { - total.checked_add(*bucket as u64) - }) - .with_context(|| error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram bucket total overflows u64".to_string(), - })?; - ensure!( - if histogram.sum.is_nan() { - bucket_count <= count - } else { - bucket_count == count - }, - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram has {bucket_count} observations in buckets, count is {count}" - ), - } - ); - - Ok(()) -} - -fn native_histogram_zero_count_is_zero(histogram: &Histogram) -> bool { - match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count == 0, - Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count == 0.0, - None => true, - } -} - -fn native_histogram_uses_float_counts(histogram: &Histogram) -> Result { - let uses_float_count = matches!(histogram.count, Some(Count::CountFloat(_))) - || matches!(histogram.zero_count, Some(ZeroCount::ZeroCountFloat(_))); - let uses_int_count = matches!(histogram.count, Some(Count::CountInt(_))) - || matches!(histogram.zero_count, Some(ZeroCount::ZeroCountInt(_))); - let uses_float_buckets = - !histogram.positive_counts.is_empty() || !histogram.negative_counts.is_empty(); - let uses_int_buckets = - !histogram.positive_deltas.is_empty() || !histogram.negative_deltas.is_empty(); - - if matches!( - (&histogram.count, &histogram.zero_count), - (Some(Count::CountInt(_)), Some(ZeroCount::ZeroCountFloat(_))) - | (Some(Count::CountFloat(_)), Some(ZeroCount::ZeroCountInt(_))) - ) { - return error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram count and zero_count must use the same integer or float family".to_string(), - } - .fail(); - } - - ensure!( - !(uses_float_buckets && uses_int_buckets), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram bucket counts must use either integer deltas or float counts".to_string(), - } - ); - ensure!( - !(uses_float_count && uses_int_buckets), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 float native histogram must not use integer bucket deltas" - .to_string(), - } - ); - ensure!( - !(uses_int_count && uses_float_buckets), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 integer native histogram must not use float bucket counts" - .to_string(), - } - ); - - Ok(uses_float_count || uses_float_buckets) -} - -fn pb_value(value_data: ValueData) -> Value { - optional_pb_value(Some(value_data)) -} - -fn null_pb_value() -> Value { - optional_pb_value(None) -} - -fn optional_pb_value(value_data: Option) -> Value { - Value { value_data } -} - -fn list_value(values: impl IntoIterator) -> Value { - pb_value(ValueData::ListValue(ListValue { - items: values.into_iter().map(pb_value).collect(), - })) -} - -fn i32_list_value(values: impl IntoIterator) -> Value { - list_value(values.into_iter().map(ValueData::I32Value)) -} - -fn i32_span_lengths(name: &str, spans: &[BucketSpan]) -> Result> { - spans - .iter() - .map(|span| { - i32::try_from(span.length) - .ok() - .context(error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} span length {} overflows i32", - span.length - ), - }) - }) - .collect() -} - -fn i64_list_value(values: impl IntoIterator) -> Value { - list_value(values.into_iter().map(ValueData::I64Value)) -} - -fn f64_list_value(values: impl IntoIterator) -> Value { - list_value(values.into_iter().map(ValueData::F64Value)) -} - -fn bucket_counts_from_deltas(deltas: &[i64]) -> Result> { - let mut count = 0_i64; - let mut buckets = Vec::with_capacity(deltas.len()); - - for delta in deltas { - count = - count - .checked_add(*delta) - .with_context(|| error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram bucket count overflows i64".to_string(), - })?; - ensure!( - count >= 0, - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram bucket count is negative".to_string(), - } - ); - buckets.push(count); - } - - Ok(buckets) -} - fn ensure_no_internal_histogram_labels(tags: &PromTags<'_>) -> Result<()> { // The histogram field column is generated from the protobuf payload. for (name, _) in tags { @@ -2371,6 +1884,25 @@ mod tests { ); } + #[test] + fn test_into_context_req_preserves_exponential_zero_threshold() { + for zero_threshold in [-1.0, f64::NAN] { + let ctx_req = decode_test_request(request_with_histogram(Histogram { + zero_threshold, + ..Default::default() + })) + .unwrap(); + let rows = ctx_req.histograms.all_req().next().unwrap().rows.unwrap(); + let Some(ValueData::F64Value(actual)) = + histogram_field_value(&rows, 0, ZERO_THRESHOLD_FIELD) + else { + panic!("expected zero threshold"); + }; + + assert_eq!(zero_threshold.to_bits(), actual.to_bits()); + } + } + #[test] fn test_into_context_req_rejects_internal_histogram_labels() { let mut request = test_util::request_with_labels_and_samples( diff --git a/src/servers/src/query_handler.rs b/src/servers/src/query_handler.rs index 96f412ac1e1..35ec1cc2ff5 100644 --- a/src/servers/src/query_handler.rs +++ b/src/servers/src/query_handler.rs @@ -70,6 +70,15 @@ pub struct TraceIngestOutcome { pub error_message: Option, } +/// Result of ingesting one OTLP metrics request or Arrow batch. +#[derive(Debug, Default, Clone)] +pub struct MetricsIngestOutcome { + pub write_cost: usize, + pub accepted_data_points: i64, + pub rejected_data_points: i64, + pub error_message: Option, +} + #[async_trait] pub trait InfluxdbLineProtocolHandler { /// A successful request will not return a response. @@ -133,7 +142,7 @@ pub trait OpenTelemetryProtocolHandler: PipelineHandler { &self, request: ExportMetricsServiceRequest, ctx: QueryContextRef, - ) -> Result; + ) -> Result; /// Handling opentelemetry traces request async fn traces( diff --git a/src/session/src/protocol_ctx.rs b/src/session/src/protocol_ctx.rs index 6e84761c2d9..537dec5c414 100644 --- a/src/session/src/protocol_ctx.rs +++ b/src/session/src/protocol_ctx.rs @@ -42,6 +42,7 @@ impl ProtocolCtx { /// If true, all scope attributes will be promoted to the final table schema. /// Along with the scope name, scope version and scope schema URL. /// - `with_metric_engine` +/// - `experimental_enable_exponential_histogram` /// - `is_legacy` /// If the user uses OTLP metrics ingestion before v0.16, it uses the old path. /// So we call this path 'legacy'. @@ -53,6 +54,7 @@ pub struct OtlpMetricCtx { pub resource_attrs: HashSet, pub promote_scope_attrs: bool, pub with_metric_engine: bool, + pub experimental_enable_exponential_histogram: bool, pub is_legacy: bool, pub metric_type: MetricType, pub metric_translation_strategy: OtlpMetricTranslationStrategy, diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index bedf9eebc70..4f31ac2ebdf 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -24,7 +24,7 @@ use file_engine::config::EngineConfig as FileEngineConfig; use flow::FlowConfig; use frontend::frontend::FrontendOptions; use frontend::service_config::{ - InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, PostgresOptions, + InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions, PostgresOptions, PromStoreOptions, }; use mito2::config::MitoConfig; @@ -56,6 +56,7 @@ pub struct StandaloneOptions { pub opentsdb: OpentsdbOptions, pub influxdb: InfluxdbOptions, pub jaeger: JaegerOptions, + pub otlp: OtlpOptions, pub prom_store: PromStoreOptions, pub wal: DatanodeWalConfig, pub storage: StorageConfig, @@ -94,6 +95,7 @@ impl Default for StandaloneOptions { opentsdb: OpentsdbOptions::default(), influxdb: InfluxdbOptions::default(), jaeger: JaegerOptions::default(), + otlp: OtlpOptions::default(), prom_store: PromStoreOptions::default(), wal: DatanodeWalConfig::default(), storage: StorageConfig::default(), @@ -153,6 +155,7 @@ impl StandaloneOptions { opentsdb: cloned_opts.opentsdb, influxdb: cloned_opts.influxdb, jaeger: cloned_opts.jaeger, + otlp: cloned_opts.otlp, prom_store: cloned_opts.prom_store, meta_client: None, logging: cloned_opts.logging, diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 30b9a07348d..cb958a84e5c 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -490,7 +490,7 @@ pub async fn setup_test_http_app_with_frontend_and_slow_query_threshold( .with_log_ingest_handler(instance.fe_instance().clone(), None, None) .with_logs_handler(instance.fe_instance().clone()) .with_influxdb_handler(instance.fe_instance().clone()) - .with_otlp_handler(instance.fe_instance().clone(), true) + .with_otlp_handler(instance.fe_instance().clone(), true, false) .with_jaeger_handler(instance.fe_instance().clone()) .with_greptime_config_options(instance.opts.to_toml().unwrap()) .build(); @@ -510,6 +510,18 @@ pub async fn setup_test_http_app_with_frontend_and_user_provider( user_provider, None, None, + false, + ) + .await +} + +pub async fn setup_test_http_app_with_otlp_exponential_histogram( + store_type: StorageType, + name: &str, + enabled: bool, +) -> (Router, TestGuard) { + setup_test_http_app_with_frontend_and_custom_options( + store_type, name, None, None, None, enabled, ) .await } @@ -520,6 +532,7 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( user_provider: Option, http_opts: Option, memory_limiter: Option, + experimental_enable_exponential_histogram: bool, ) -> (Router, TestGuard) { let plugins = Plugins::new(); if let Some(user_provider) = user_provider.clone() { @@ -541,7 +554,12 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( .with_log_ingest_handler(instance.fe_instance().clone(), None, None) .with_logs_handler(instance.fe_instance().clone()) .with_influxdb_handler(instance.fe_instance().clone()) - .with_otlp_handler(instance.fe_instance().clone(), true) + .with_otlp_handler( + instance.fe_instance().clone(), + true, + experimental_enable_exponential_histogram, + ) + .with_prometheus_handler(instance.fe_instance().clone()) .with_jaeger_handler(instance.fe_instance().clone()) .with_dashboard_handler(instance.fe_instance().clone()) .with_greptime_config_options(instance.opts.to_toml().unwrap()); diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 23862870dba..9d0e13dddc9 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use api::v1::alter_table_expr::Kind; use api::v1::promql_request::Promql; use api::v1::value::ValueData; @@ -32,8 +34,20 @@ use common_recordbatch::RecordBatches; use common_runtime::Runtime; use common_runtime::runtime::{BuilderBuild, RuntimeTrait}; use common_test_util::find_workspace_path; -use otel_arrow_rust::proto::opentelemetry::arrow::v1::BatchArrowRecords; +use datatypes::arrow::array::{ + Array, ArrayRef, Float64Array, Int32Array, ListBuilder, StringArray, StructArray, + TimestampNanosecondArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array, UInt64Builder, +}; +use datatypes::arrow::datatypes::{DataType, Field}; +use datatypes::arrow::ipc::writer::StreamWriter; +use datatypes::arrow::record_batch::RecordBatch as ArrowRecordBatch; +use otel_arrow_rust::otlp::metrics::MetricType as ArrowMetricType; use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_client::ArrowMetricsServiceClient; +use otel_arrow_rust::proto::opentelemetry::arrow::v1::{ + ArrowPayload, ArrowPayloadType, BatchArrowRecords, StatusCode as ArrowStatusCode, +}; +use otel_arrow_rust::proto::opentelemetry::metrics::v1::AggregationTemporality; +use otel_arrow_rust::schema::consts as arrow_consts; use servers::grpc::GrpcServerConfig; use servers::grpc::builder::GrpcServerBuilder; use servers::http::prometheus::{ @@ -85,6 +99,7 @@ macro_rules! grpc_tests { test_auto_create_table_with_hints, test_auto_create_table_disabled_by_config, test_otel_arrow_auth, + test_otel_arrow_exponential_histogram, test_insert_and_select, test_dbname, test_grpc_message_size_ok, @@ -378,6 +393,188 @@ pub async fn test_otel_arrow_auth(store_type: StorageType) { let _ = fe_grpc_server.shutdown().await; } +// The pinned otel-arrow Producer cannot hash List-typed bucket schemas yet, so +// serialize these test-only record batches directly into the same Arrow stream format. +fn serialize_arrow_record_batch(record_batch: &ArrowRecordBatch) -> Vec { + let mut bytes = Vec::new(); + let mut writer = StreamWriter::try_new(&mut bytes, record_batch.schema_ref()).unwrap(); + writer.write(record_batch).unwrap(); + writer.finish().unwrap(); + drop(writer); + bytes +} + +fn exponential_histogram_arrow_batch(batch_id: i64, scales: &[i32]) -> BatchArrowRecords { + let resource = StructArray::from(vec![( + Arc::new(Field::new(arrow_consts::ID, DataType::UInt16, true)), + Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef, + )]); + let scope = StructArray::from(vec![( + Arc::new(Field::new(arrow_consts::ID, DataType::UInt16, true)), + Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef, + )]); + let metrics = ArrowRecordBatch::try_from_iter(vec![ + ( + arrow_consts::ID, + Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef, + ), + (arrow_consts::RESOURCE, Arc::new(resource) as ArrayRef), + (arrow_consts::SCOPE, Arc::new(scope) as ArrayRef), + ( + arrow_consts::METRIC_TYPE, + Arc::new(UInt8Array::from(vec![ + ArrowMetricType::ExponentialHistogram as u8, + ])) as ArrayRef, + ), + ( + arrow_consts::NAME, + Arc::new(StringArray::from(vec!["otel.arrow.exponential.latency"])) as ArrayRef, + ), + ( + arrow_consts::AGGREGATION_TEMPORALITY, + Arc::new(Int32Array::from(vec![ + AggregationTemporality::Cumulative as i32, + ])) as ArrayRef, + ), + ]) + .unwrap(); + + let point_count = scales.len(); + let mut positive_counts = ListBuilder::new(UInt64Builder::new()); + let mut negative_counts = ListBuilder::new(UInt64Builder::new()); + for _ in scales { + positive_counts.values().append_slice(&[1, 2]); + positive_counts.append(true); + negative_counts.append(true); + } + let positive_counts = positive_counts.finish(); + let negative_counts = negative_counts.finish(); + let positive = StructArray::from(vec![ + ( + Arc::new(Field::new( + arrow_consts::EXP_HISTOGRAM_OFFSET, + DataType::Int32, + true, + )), + Arc::new(Int32Array::from(vec![-1; point_count])) as ArrayRef, + ), + ( + Arc::new(Field::new( + arrow_consts::EXP_HISTOGRAM_BUCKET_COUNTS, + positive_counts.data_type().clone(), + true, + )), + Arc::new(positive_counts) as ArrayRef, + ), + ]); + let negative = StructArray::from(vec![ + ( + Arc::new(Field::new( + arrow_consts::EXP_HISTOGRAM_OFFSET, + DataType::Int32, + true, + )), + Arc::new(Int32Array::from(vec![0; point_count])) as ArrayRef, + ), + ( + Arc::new(Field::new( + arrow_consts::EXP_HISTOGRAM_BUCKET_COUNTS, + negative_counts.data_type().clone(), + true, + )), + Arc::new(negative_counts) as ArrayRef, + ), + ]); + let data_points = ArrowRecordBatch::try_from_iter(vec![ + ( + arrow_consts::ID, + Arc::new(UInt32Array::from_iter_values( + (0..point_count).map(|id| u32::try_from(id).unwrap()), + )) as ArrayRef, + ), + ( + arrow_consts::PARENT_ID, + Arc::new(UInt16Array::from(vec![0_u16; point_count])) as ArrayRef, + ), + ( + arrow_consts::START_TIME_UNIX_NANO, + Arc::new(TimestampNanosecondArray::from(vec![ + 1_000_000_000; + point_count + ])) as ArrayRef, + ), + ( + arrow_consts::TIME_UNIX_NANO, + Arc::new(TimestampNanosecondArray::from(vec![ + 3_000_000_000; + point_count + ])) as ArrayRef, + ), + ( + arrow_consts::HISTOGRAM_COUNT, + Arc::new(UInt64Array::from(vec![4_u64; point_count])) as ArrayRef, + ), + ( + arrow_consts::HISTOGRAM_SUM, + Arc::new(Float64Array::from(vec![8.0; point_count])) as ArrayRef, + ), + ( + arrow_consts::EXP_HISTOGRAM_SCALE, + Arc::new(Int32Array::from(scales.to_vec())) as ArrayRef, + ), + ( + arrow_consts::EXP_HISTOGRAM_ZERO_COUNT, + Arc::new(UInt64Array::from(vec![1_u64; point_count])) as ArrayRef, + ), + ( + arrow_consts::EXP_HISTOGRAM_POSITIVE, + Arc::new(positive) as ArrayRef, + ), + ( + arrow_consts::EXP_HISTOGRAM_NEGATIVE, + Arc::new(negative) as ArrayRef, + ), + ( + arrow_consts::FLAGS, + Arc::new(UInt32Array::from(vec![0_u32; point_count])) as ArrayRef, + ), + ]) + .unwrap(); + BatchArrowRecords { + batch_id, + arrow_payloads: vec![ + ArrowPayload { + schema_id: format!("metrics-{batch_id}"), + r#type: ArrowPayloadType::UnivariateMetrics as i32, + record: serialize_arrow_record_batch(&metrics), + }, + ArrowPayload { + schema_id: format!("exp-histogram-{batch_id}"), + r#type: ArrowPayloadType::ExpHistogramDataPoints as i32, + record: serialize_arrow_record_batch(&data_points), + }, + ], + headers: vec![], + } +} + +pub async fn test_otel_arrow_exponential_histogram(store_type: StorageType) { + let (_instance, server) = + setup_grpc_server(store_type, "test_otel_arrow_exponential_histogram").await; + let addr = server.bind_addr().unwrap().to_string(); + let mut client = ArrowMetricsServiceClient::connect(format!("http://{addr}")) + .await + .unwrap(); + let batch = exponential_histogram_arrow_batch(0, &[0]); + let request = Request::new(futures::stream::once(async { batch })); + let mut response = client.arrow_metrics(request).await.unwrap().into_inner(); + let status = response.message().await.unwrap().unwrap(); + assert_eq!(0, status.batch_id); + assert_eq!(ArrowStatusCode::InvalidArgument as i32, status.status_code); + assert!(status.status_message.contains("omits zero_threshold")); + let _ = server.shutdown().await; +} + fn basic_auth(username: &str, password: &str) -> String { format!("Basic {}", basic_auth_credentials(username, password)) } diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index f89d75e9ef0..2a66c4353e7 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -50,7 +50,9 @@ use log_query::{AggFunc, Context, Limit, LogExpr, LogQuery, TimeFilter}; use loki_proto::logproto::{EntryAdapter, LabelPairAdapter, PushRequest, StreamAdapter}; use loki_proto::prost_types::Timestamp; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; -use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest; +use opentelemetry_proto::tonic::collector::metrics::v1::{ + ExportMetricsServiceRequest, ExportMetricsServiceResponse, +}; use opentelemetry_proto::tonic::collector::trace::v1::{ ExportTraceServiceRequest, ExportTraceServiceResponse, }; @@ -160,6 +162,7 @@ macro_rules! http_tests { test_pipeline_index_options, test_otlp_metrics_new, + test_otlp_exponential_histogram, test_otlp_metric_translation_strategies, test_otlp_traces_v0, test_otlp_traces_v1, @@ -2201,6 +2204,11 @@ default_merge_mode = "last_non_null" [jaeger] enable = true +[otlp] +enable = true +experimental_enable_exponential_histogram = false +trace_ingest_chunk_size = 512 + [prom_store] enable = true with_metric_engine = true @@ -6709,6 +6717,140 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_otlp_exponential_histogram(store_type: StorageType) { + use opentelemetry_proto::tonic::metrics::v1::{ + AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Metric, + ResourceMetrics, ScopeMetrics, exponential_histogram_data_point, metric, + }; + use tests_integration::test_util::setup_test_http_app_with_otlp_exponential_histogram; + + common_telemetry::init_default_ut_logging(); + let req = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: "otlp.exponential.latency".to_string(), + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points: vec![ + ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000_000, + time_unix_nano: 3_000_000_000, + count: 4, + sum: Some(8.0), + scale: 0, + zero_count: 1, + positive: Some(exponential_histogram_data_point::Buckets { + offset: -1, + bucket_counts: vec![1, 2], + }), + ..Default::default() + }, + ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000_000, + time_unix_nano: 4_000_000_000, + scale: -5, + ..Default::default() + }, + ], + aggregation_temporality: AggregationTemporality::Cumulative as i32, + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + let body = req.encode_to_vec(); + let headers = || { + vec![( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + )] + }; + + let (app, mut guard) = setup_test_http_app_with_otlp_exponential_histogram( + store_type, + "test_otlp_exponential_histogram_disabled", + false, + ) + .await; + let client = TestClient::new(app).await; + let res = send_req( + &client, + headers(), + "/v1/otlp/v1/metrics", + body.clone(), + false, + ) + .await; + assert_eq!(StatusCode::BAD_REQUEST, res.status()); + let status = GoogleRpcStatus::decode(res.bytes().await.as_ref()).unwrap(); + assert_eq!(3, status.code); + assert!( + status + .message + .contains("otlp.experimental_enable_exponential_histogram") + ); + validate_data( + "otlp_exponential_histogram_disabled_no_table", + &client, + "select count(*) from information_schema.tables where table_name = 'otlp_exponential_latency';", + "[[0]]", + ) + .await; + guard.remove_all().await; + + let (app, mut guard) = setup_test_http_app_with_otlp_exponential_histogram( + store_type, + "test_otlp_exponential_histogram_enabled", + true, + ) + .await; + let client = TestClient::new(app).await; + let res = send_req(&client, headers(), "/v1/otlp/v1/metrics", body, false).await; + assert_eq!(StatusCode::OK, res.status()); + let response = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + let partial_success = response.partial_success.unwrap(); + assert_eq!(1, partial_success.rejected_data_points); + assert!(partial_success.error_message.contains("scale -5")); + + validate_data( + "otlp_exponential_histogram_row", + &client, + "select greptime_timestamp, greptime_native_histogram from otlp_exponential_latency;", + "[[3000,{\"count_f64\":null,\"count_i64\":4,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,2],\"positive_span_lengths\":[2],\"positive_span_offsets\":[0],\"reset_hint\":0,\"schema\":0,\"start_timestamp\":1000,\"sum\":8.0,\"zero_count_f64\":null,\"zero_count_i64\":1,\"zero_threshold\":0.0}]]", + ) + .await; + validate_data( + "otlp_exponential_histogram_metadata", + &client, + "select count(*) from information_schema.tables where table_name = 'otlp_exponential_latency' and create_options like '%greptime.semantic.metric.type=histogram%' and create_options like '%greptime.semantic.metric.temporality=cumulative%';", + "[[1]]", + ) + .await; + + for query in [ + "histogram_count(otlp_exponential_latency)", + "histogram_sum(otlp_exponential_latency)", + "histogram_quantile(0.5,otlp_exponential_latency)", + "histogram_count(rate(otlp_exponential_latency[5s]))", + ] { + let res = client + .get(&format!( + "/v1/prometheus/api/v1/query?query={}&time=3", + encode(query) + )) + .send() + .await; + assert_eq!(StatusCode::OK, res.status(), "query: {query}"); + let body = res.text().await; + assert!(body.contains("otlp_exponential_latency"), "{query}: {body}"); + } + + guard.remove_all().await; +} + pub async fn test_otlp_metric_translation_strategies(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = @@ -10646,6 +10788,7 @@ pub async fn test_http_memory_limit(store_type: StorageType) { None, Some(http_opts), Some(memory_limiter), + false, ) .await; From 0cc83c45708e442d929c22084fc0e6dda9c76f4a Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 24 Aug 2026 13:34:30 +0000 Subject: [PATCH 097/126] feat(cmd): add parquet development tools (#8939) * feat(cmd): add parquet metadata development tool Signed-off-by: evenyag * feat(cmd): add parquet rewrite development tool Signed-off-by: evenyag * feat(cmd): add SST replacement development tool Signed-off-by: evenyag * feat(cmd): support local parquet files in parquetbench Signed-off-by: evenyag * fix(cmd): clean up parquet development tools Signed-off-by: evenyag * refactor(cmd): share datanode tool utilities Signed-off-by: evenyag * fix(cmd): satisfy parquet tool lints Signed-off-by: evenyag * docs: document parquet development tools Signed-off-by: evenyag * refactor(cmd): rename SST replacement module Signed-off-by: evenyag * fix(cmd): validate parquet rewrite options Signed-off-by: evenyag --------- Signed-off-by: evenyag --- docs/parquet-tools.md | 253 +++++++++++ src/cmd/Cargo.toml | 2 +- src/cmd/src/bin/greptime.rs | 6 + src/cmd/src/datanode.rs | 46 ++ src/cmd/src/datanode/objbench.rs | 107 +---- src/cmd/src/datanode/parquet_meta.rs | 340 ++++++++++++++ src/cmd/src/datanode/parquet_rewrite.rs | 579 ++++++++++++++++++++++++ src/cmd/src/datanode/parquetbench.rs | 399 +++++++++++----- src/cmd/src/datanode/scanbench.rs | 60 +-- src/cmd/src/datanode/sst_replace.rs | 361 +++++++++++++++ src/cmd/src/datanode/tool_util.rs | 236 ++++++++++ 11 files changed, 2126 insertions(+), 263 deletions(-) create mode 100644 docs/parquet-tools.md create mode 100644 src/cmd/src/datanode/parquet_meta.rs create mode 100644 src/cmd/src/datanode/parquet_rewrite.rs create mode 100644 src/cmd/src/datanode/sst_replace.rs create mode 100644 src/cmd/src/datanode/tool_util.rs diff --git a/docs/parquet-tools.md b/docs/parquet-tools.md new file mode 100644 index 00000000000..8528f969431 --- /dev/null +++ b/docs/parquet-tools.md @@ -0,0 +1,253 @@ +# Parquet Development Tools + +GreptimeDB provides the following datanode CLI commands for inspecting, +rewriting, benchmarking, and replacing Parquet SST files: + +- `parquet-meta`: inspect a local Parquet file. +- `parquet-rewrite`: rewrite a local Parquet file with different writer + properties. +- `parquetbench`: benchmark reads from a local GreptimeDB SST or an SST in a + configured object store. +- `sst-replace`: replace an existing Mito SST object and update its manifest + metadata. + +These are development and recovery tools, not stable user-facing interfaces. +Always test a workflow on disposable data before using it on an important +region. + +## Build + +`parquet-meta`, `parquet-rewrite`, and `sst-replace` require the `dev-tools` +feature. The same build also includes `parquetbench`: + +```bash +cargo build -p cmd --bin greptime --features dev-tools +``` + +The examples below use: + +```bash +GREPTIME=./target/debug/greptime +``` + +Run `$GREPTIME datanode --help` for the authoritative argument list. + +## Inspect a Parquet file + +`parquet-meta` reads a local Parquet footer and optional page indexes. It reports +file, row-group, and column metadata, including compression, encodings, sizes, +page offsets, statistics presence, and bloom-filter/index offsets. + +Text output: + +```bash +$GREPTIME datanode parquet-meta \ + --input /tmp/source.parquet +``` + +Machine-readable JSON output: + +```bash +$GREPTIME datanode parquet-meta \ + --input /tmp/source.parquet \ + --format json > /tmp/source-meta.json +``` + +`--format` accepts `text` (the default) or `json`. This command can inspect a +general Parquet file; it does not require GreptimeDB SST key-value metadata. + +## Rewrite a Parquet file + +`parquet-rewrite` has two modes. First, dump a TOML properties file inferred +from an existing file: + +```bash +$GREPTIME datanode parquet-rewrite \ + --input /tmp/source.parquet \ + --dump-properties /tmp/writer-properties.toml +``` + +Review and edit the generated file. Its shape is: + +```toml +[writer] +dictionary_enabled = true +compression = "zstd" +max_row_group_row_count = 8192 + +[[columns]] +path = ["host"] +dictionary_enabled = true +compression = "zstd" +encoding = "plain" +``` + +Supported compression names are `uncompressed`, `snappy`, `gzip`, `lzo`, +`brotli`, `lz4`, `zstd`, and `lz4-raw`. Supported encodings are `plain`, +`delta-binary-packed`, `delta-length-byte-array`, `delta-byte-array`, and +`byte-stream-split`. + +The `[writer]` table also accepts `compression_level`, +`data_page_size_limit`, `data_page_row_count_limit`, and +`dictionary_page_size_limit`. A `[[columns]]` entry overrides dictionary, +compression, compression level, or encoding for its column path. Unknown TOML +fields are rejected. Compression levels are supported for `gzip`, `brotli`, and +`zstd`. A column-level `compression_level` without `compression` inherits the +writer compression codec; it is rejected when there is no writer codec to +inherit. + +Rewrite the data using the edited properties: + +```bash +$GREPTIME datanode parquet-rewrite \ + --input /tmp/source.parquet \ + --properties /tmp/writer-properties.toml \ + --output /tmp/rewritten.parquet +``` + +Use `--batch-size ` to control reader batch size. Output and dumped +properties files are not replaced unless `--overwrite` is passed. The output or +dump path must not be exactly the same path as the input, even with +`--overwrite`. This check does not resolve symbolic links, hard links, or other +spellings of the same path. + +The rewrite decodes and writes the Arrow record batches and copies the Parquet +key-value metadata, but it creates a new physical Parquet layout. Inspect and +validate the result before using it as an SST: + +```bash +$GREPTIME datanode parquet-meta \ + --input /tmp/rewritten.parquet +``` + +The dumped properties are inferred primarily from the first row group. Review +them when the source uses different properties across row groups. + +## Benchmark an SST + +`parquetbench` expects GreptimeDB region metadata embedded in the SST. It can +read a local file with the direct reader: + +```bash +$GREPTIME datanode parquetbench \ + --file-path /tmp/rewritten.parquet \ + --reader direct \ + --iterations 5 \ + --batch-size 8192 +``` + +Local-file mode cannot be combined with `--config`, `--region-id`, +`--table-dir`, or `--file-id`, and it does not support the `flat-prune` reader. + +To benchmark an SST in the object store configured for a datanode or standalone +deployment: + +```bash +$GREPTIME datanode parquetbench \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --path-type bare \ + --reader flat-prune \ + --iterations 5 +``` + +Region mode requires all four of `--config`, `--region-id`, `--table-dir`, and +`--file-id`. `--region-id` accepts either the packed unsigned integer or +`:`. `--path-type` accepts `bare`, `data`, or +`metadata` and defaults to `bare`. + +An optional scan configuration selects columns and row groups: + +```json +{ + "projection_names": ["host", "value", "ts"], + "row_groups": [0, 2] +} +``` + +```bash +$GREPTIME datanode parquetbench \ + --file-path /tmp/rewritten.parquet \ + --scan-config /tmp/parquet-scan.json \ + --iterations 5 +``` + +Use `--pk-as-binary` to expose `__primary_key` as binary with the direct reader. +On Unix, `--pprof-file ` writes a flamegraph. Add +`--pprof-after-warmup` and use at least two iterations to exclude the first +iteration from profiling. + +## Replace an existing region SST + +> **Warning:** `sst-replace` mutates both an SST object and its region manifest. +> Stop the datanode that owns the region and back up the target SST and manifest +> before using `--confirm`. The SST write happens before the manifest update, so +> an interrupted or failed operation may require restoring the backup. + +`sst-replace` replaces the contents of an existing SST file ID. It does not add +a new file ID to a region. The command requires the replacement to have the same +row count and row-group count recorded in the manifest when those manifest +values are nonzero. It does not validate schema or row contents. The replacement +is loaded into memory in full, so ensure the machine has enough memory for the +SST. + +Start with the default dry run using a local replacement file: + +```bash +$GREPTIME datanode sst-replace \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --replacement-file /tmp/rewritten.parquet +``` + +The dry run locates the manifest and target SST, reads and validates the +replacement footer, and prints the old and new sizes without writing anything. +`--path-type` defaults to `auto`, which probes `bare`, `data`, and `metadata`. +Specify the path type if the file ID is visible in more than one manifest. + +The replacement can instead be read from the configured object store: + +```bash +$GREPTIME datanode sst-replace \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --replacement-object staging/rewritten.parquet \ + --path-type bare +``` + +After reviewing the dry-run output and confirming that the datanode is stopped, +repeat the exact command with `--confirm`: + +```bash +$GREPTIME datanode sst-replace \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --replacement-file /tmp/rewritten.parquet \ + --path-type bare \ + --confirm +``` + +The confirmed operation overwrites the existing SST object, recalculates file +size and row-group statistics, and appends a manifest edit for the existing file +ID. Restart the datanode and validate queries against the region before removing +the backup. + +## Recommended rewrite and replacement workflow + +1. Back up the source SST and its region manifest. +2. Inspect the source with `parquet-meta`. +3. Dump and edit writer properties with `parquet-rewrite`. +4. Rewrite to a new local file; never rewrite directly over the source SST. +5. Inspect the output and benchmark it with `parquetbench`. +6. Stop the owning datanode. +7. Run `sst-replace` without `--confirm` and review its resolved target and + statistics. +8. Repeat with `--confirm`, restart the datanode, and validate the region. diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 10865b48a10..579e5b1454c 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -27,7 +27,7 @@ default = [ "meta-srv/mysql_kvbackend", ] enterprise = ["common-meta/enterprise", "frontend/enterprise", "meta-srv/enterprise"] -# Developer-only helper binary gate for `query_perf_fixture`. +# Developer-only helper binaries and diagnostic datanode commands. # Kept out of `default` so normal/release builds don't compile them. dev-tools = [] mysql-object-store = ["object-store/mysql-object-store"] diff --git a/src/cmd/src/bin/greptime.rs b/src/cmd/src/bin/greptime.rs index 957e227355a..f5f4a590487 100644 --- a/src/cmd/src/bin/greptime.rs +++ b/src/cmd/src/bin/greptime.rs @@ -118,6 +118,12 @@ async fn start(cli: Command) -> Result<()> { datanode::SubCommand::Objbench(ref bench) => bench.run().await, datanode::SubCommand::Scanbench(ref bench) => bench.run().await, datanode::SubCommand::Parquetbench(ref bench) => bench.run().await, + #[cfg(feature = "dev-tools")] + datanode::SubCommand::ParquetMeta(ref cmd) => cmd.run().await, + #[cfg(feature = "dev-tools")] + datanode::SubCommand::ParquetRewrite(ref cmd) => cmd.run().await, + #[cfg(feature = "dev-tools")] + datanode::SubCommand::SstReplace(ref cmd) => cmd.run().await, }, SubCommand::Flownode(cmd) => { cmd.build(cmd.load_options(&cli.global_options)?) diff --git a/src/cmd/src/datanode.rs b/src/cmd/src/datanode.rs index c33be6c6f33..9cb8a8c66c1 100644 --- a/src/cmd/src/datanode.rs +++ b/src/cmd/src/datanode.rs @@ -15,10 +15,20 @@ pub mod builder; #[allow(clippy::print_stdout)] pub(crate) mod objbench; +#[cfg(feature = "dev-tools")] +#[allow(clippy::print_stdout)] +mod parquet_meta; +#[cfg(feature = "dev-tools")] +#[allow(clippy::print_stdout)] +mod parquet_rewrite; #[allow(clippy::print_stdout)] pub mod parquetbench; #[allow(clippy::print_stdout)] pub mod scanbench; +#[cfg(feature = "dev-tools")] +#[allow(clippy::print_stdout)] +mod sst_replace; +mod tool_util; use std::path::Path; use std::time::Duration; @@ -39,8 +49,14 @@ use tracing_appender::non_blocking::WorkerGuard; use crate::App; use crate::datanode::builder::InstanceBuilder; use crate::datanode::objbench::ObjbenchCommand; +#[cfg(feature = "dev-tools")] +use crate::datanode::parquet_meta::ParquetMetaCommand; +#[cfg(feature = "dev-tools")] +use crate::datanode::parquet_rewrite::ParquetRewriteCommand; use crate::datanode::parquetbench::ParquetbenchCommand; use crate::datanode::scanbench::ScanbenchCommand; +#[cfg(feature = "dev-tools")] +use crate::datanode::sst_replace::SstReplaceCommand; use crate::error::{ LoadLayeredConfigSnafu, MissingConfigSnafu, Result, ShutdownDatanodeSnafu, StartDatanodeSnafu, }; @@ -114,6 +130,12 @@ impl Command { // Bench commands are standalone utilities and don't need to load DatanodeOptions. SubCommand::Objbench(_) | SubCommand::Scanbench(_) => Self::default_bench_options(), SubCommand::Parquetbench(_) => Self::default_bench_options(), + #[cfg(feature = "dev-tools")] + SubCommand::ParquetMeta(_) => Self::default_bench_options(), + #[cfg(feature = "dev-tools")] + SubCommand::ParquetRewrite(_) => Self::default_bench_options(), + #[cfg(feature = "dev-tools")] + SubCommand::SstReplace(_) => Self::default_bench_options(), } } @@ -139,6 +161,15 @@ pub enum SubCommand { Scanbench(ScanbenchCommand), /// Benchmark scanning a single parquet SST. Parquetbench(ParquetbenchCommand), + /// Display metadata of a parquet file. + #[cfg(feature = "dev-tools")] + ParquetMeta(ParquetMetaCommand), + /// Rewrite a parquet file with different writer properties. + #[cfg(feature = "dev-tools")] + ParquetRewrite(ParquetRewriteCommand), + /// Replace a Mito region SST and update its manifest metadata. + #[cfg(feature = "dev-tools")] + SstReplace(SstReplaceCommand), } impl SubCommand { @@ -160,6 +191,21 @@ impl SubCommand { cmd.run().await?; std::process::exit(0); } + #[cfg(feature = "dev-tools")] + SubCommand::ParquetMeta(cmd) => { + cmd.run().await?; + std::process::exit(0); + } + #[cfg(feature = "dev-tools")] + SubCommand::ParquetRewrite(cmd) => { + cmd.run().await?; + std::process::exit(0); + } + #[cfg(feature = "dev-tools")] + SubCommand::SstReplace(cmd) => { + cmd.run().await?; + std::process::exit(0); + } } } } diff --git a/src/cmd/src/datanode/objbench.rs b/src/cmd/src/datanode/objbench.rs index f33734687dc..a01cad8468a 100644 --- a/src/cmd/src/datanode/objbench.rs +++ b/src/cmd/src/datanode/objbench.rs @@ -18,8 +18,6 @@ use std::time::Instant; use clap::Parser; use colored::Colorize; -use datanode::config::RegionEngineConfig; -use datanode::store; use futures::stream; use mito2::access_layer::{ AccessLayer, AccessLayerRef, Metrics, OperationType, SstWriteRequest, WriteType, @@ -32,19 +30,20 @@ use mito2::sst::file::{FileHandle, FileMeta}; use mito2::sst::file_purger::{FilePurger, FilePurgerRef}; use mito2::sst::index::intermediate::IntermediateManager; use mito2::sst::index::puffin_manager::PuffinManagerFactory; +use mito2::sst::parquet::WriteOptions; use mito2::sst::parquet::reader::ParquetReaderBuilder; -use mito2::sst::parquet::{PARQUET_METADATA_KEY, WriteOptions}; use mito2::worker::write_cache_from_config; use object_store::ObjectStore; -use parquet::file::metadata::{FooterTail, KeyValue}; +use parquet::file::metadata::FooterTail; use regex::Regex; use snafu::OptionExt; -use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::path_utils::region_name; use store_api::region_request::PathType; use store_api::storage::FileId; -use crate::datanode::{StorageConfig, StorageConfigWrapper}; +use crate::datanode::tool_util::{ + build_object_store, extract_region_metadata, max_row_group_uncompressed_size, parse_config, +}; use crate::error; /// Object storage benchmark command @@ -67,46 +66,6 @@ pub struct ObjbenchCommand { pub pprof_file: Option, } -pub(crate) fn parse_config( - config_path: &PathBuf, -) -> error::Result<( - StorageConfig, - MitoConfig, - common_wal::config::DatanodeWalConfig, -)> { - let cfg_str = std::fs::read_to_string(config_path).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("failed to read config {}: {e}", config_path.display()), - } - .build() - })?; - - let store_cfg: StorageConfigWrapper = toml::from_str(&cfg_str).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("failed to parse config {}: {e}", config_path.display()), - } - .build() - })?; - - let wal_config = store_cfg.wal; - let storage_config = store_cfg.storage; - let mito_engine_config = store_cfg - .region_engine - .into_iter() - .filter_map(|c| { - if let RegionEngineConfig::Mito(mito) = c { - Some(mito) - } else { - None - } - }) - .next() - .with_context(|| error::IllegalConfigSnafu { - msg: format!("Engine config not found in {:?}", config_path), - })?; - Ok((storage_config, mito_engine_config, wal_config)) -} - impl ObjbenchCommand { pub async fn run(&self) -> error::Result<()> { if self.verbose { @@ -154,17 +113,7 @@ impl ObjbenchCommand { let region_meta = extract_region_metadata(&self.source, &parquet_meta)?; let num_rows = parquet_meta.file_metadata().num_rows() as u64; let num_row_groups = parquet_meta.num_row_groups() as u64; - let max_row_group_uncompressed_size: u64 = parquet_meta - .row_groups() - .iter() - .map(|rg| { - rg.columns() - .iter() - .map(|c| c.uncompressed_size() as u64) - .sum::() - }) - .max() - .unwrap_or(0); + let max_row_group_uncompressed_size = max_row_group_uncompressed_size(&parquet_meta); println!( "{} Metadata loaded - rows: {}, size: {} bytes", @@ -485,47 +434,6 @@ fn parse_file_dir_components(path: &str) -> error::Result { }) } -pub(crate) fn extract_region_metadata( - file_path: &str, - meta: &parquet::file::metadata::ParquetMetaData, -) -> error::Result { - let kvs: Option<&Vec> = meta.file_metadata().key_value_metadata(); - let Some(kvs) = kvs else { - return Err(error::IllegalConfigSnafu { - msg: format!("{file_path}: missing parquet key_value metadata"), - } - .build()); - }; - let json = kvs - .iter() - .find(|kv| kv.key == PARQUET_METADATA_KEY) - .and_then(|kv| kv.value.as_ref()) - .ok_or_else(|| { - error::IllegalConfigSnafu { - msg: format!("{file_path}: key {PARQUET_METADATA_KEY} not found or empty"), - } - .build() - })?; - let region: RegionMetadata = RegionMetadata::from_json(json).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region metadata json: {e}"), - } - .build() - })?; - Ok(Arc::new(region)) -} - -pub(crate) async fn build_object_store(sc: &StorageConfig) -> error::Result { - store::new_object_store(sc.store.clone(), &sc.data_home) - .await - .map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("Failed to build object store: {e:?}"), - } - .build() - }) -} - async fn build_access_layer_simple( components: &FileDirComponents, object_store: ObjectStore, @@ -666,7 +574,8 @@ mod tests { use common_base::readable_size::ReadableSize; use store_api::region_request::PathType; - use crate::datanode::objbench::{parse_config, parse_file_dir_components}; + use crate::datanode::objbench::parse_file_dir_components; + use crate::datanode::tool_util::parse_config; #[test] fn test_parse_dir() { diff --git a/src/cmd/src/datanode/parquet_meta.rs b/src/cmd/src/datanode/parquet_meta.rs new file mode 100644 index 00000000000..425da45e3c0 --- /dev/null +++ b/src/cmd/src/datanode/parquet_meta.rs @@ -0,0 +1,340 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; + +use clap::{Parser, ValueEnum}; +use parquet::file::metadata::ParquetMetaData; +use serde::Serialize; +use snafu::ResultExt; + +use crate::datanode::tool_util::{compression_name, load_local_parquet_metadata}; +use crate::error; + +/// Display parquet file metadata. +#[derive(Debug, Parser)] +pub struct ParquetMetaCommand { + /// Path to input parquet file. + #[clap(long, value_name = "FILE")] + input: PathBuf, + + /// Output format. + #[clap(long, value_enum, default_value = "text")] + format: MetaOutputFormat, +} + +#[derive(Debug, Clone, Copy, ValueEnum)] +enum MetaOutputFormat { + Text, + Json, +} + +#[derive(Debug, Serialize)] +struct FileMetaView { + input: String, + num_rows: i64, + num_row_groups: usize, + num_columns: usize, + key_value_metadata: BTreeMap>, + row_groups: Vec, +} + +#[derive(Debug, Serialize)] +struct RowGroupMetaView { + index: usize, + num_rows: i64, + uncompressed_size: i64, + compressed_size: i64, + compression_ratio: Option, + data_pages: Option, + dictionary_page_bytes: i64, + columns: Vec, +} + +#[derive(Debug, Serialize)] +struct ColumnChunkMetaView { + index: usize, + path: String, + physical_type: String, + encodings: Vec, + compression: String, + num_values: i64, + uncompressed_size: i64, + compressed_size: i64, + compression_ratio: Option, + data_page_offset: i64, + dictionary_page_offset: Option, + dictionary_page_bytes: Option, + data_pages: Option, + has_statistics: bool, + column_index_offset: Option, + column_index_length: Option, + offset_index_offset: Option, + offset_index_length: Option, + bloom_filter_offset: Option, + bloom_filter_length: Option, +} + +impl ParquetMetaCommand { + pub async fn run(&self) -> error::Result<()> { + let metadata = load_local_parquet_metadata(&self.input)?; + let view = build_file_meta_view(&self.input, &metadata); + + match self.format { + MetaOutputFormat::Text => print_meta_text(&view), + MetaOutputFormat::Json => { + let json = serde_json::to_string_pretty(&view).context(error::SerdeJsonSnafu)?; + println!("{json}"); + } + } + + Ok(()) + } +} + +fn build_file_meta_view(path: &Path, metadata: &ParquetMetaData) -> FileMetaView { + let key_value_metadata = metadata + .file_metadata() + .key_value_metadata() + .into_iter() + .flatten() + .map(|kv| (kv.key.clone(), kv.value.clone())) + .collect(); + let offset_index = metadata.offset_index(); + let row_groups = metadata + .row_groups() + .iter() + .enumerate() + .map(|(row_group_idx, row_group)| { + let columns: Vec<_> = row_group + .columns() + .iter() + .enumerate() + .map(|(column_idx, column)| { + let data_pages = offset_index + .and_then(|index| index.get(row_group_idx)) + .and_then(|columns| columns.get(column_idx)) + .map(|index| index.page_locations().len()); + let dictionary_page_bytes = dictionary_page_bytes( + column.dictionary_page_offset(), + column.data_page_offset(), + ); + + ColumnChunkMetaView { + index: column_idx, + path: column.column_path().string(), + physical_type: format!("{:?}", column.column_type()), + encodings: column.encodings().map(|enc| format!("{enc:?}")).collect(), + compression: compression_name(column.compression()).to_string(), + num_values: column.num_values(), + uncompressed_size: column.uncompressed_size(), + compressed_size: column.compressed_size(), + compression_ratio: compression_ratio( + column.uncompressed_size(), + column.compressed_size(), + ), + data_page_offset: column.data_page_offset(), + dictionary_page_offset: column.dictionary_page_offset(), + dictionary_page_bytes, + data_pages, + has_statistics: column.statistics().is_some(), + column_index_offset: column.column_index_offset(), + column_index_length: column.column_index_length(), + offset_index_offset: column.offset_index_offset(), + offset_index_length: column.offset_index_length(), + bloom_filter_offset: column.bloom_filter_offset(), + bloom_filter_length: column.bloom_filter_length(), + } + }) + .collect(); + let data_pages = if columns.iter().all(|column| column.data_pages.is_some()) { + Some( + columns + .iter() + .map(|column| column.data_pages.unwrap_or_default()) + .sum(), + ) + } else { + None + }; + let dictionary_page_bytes = columns + .iter() + .map(|column| column.dictionary_page_bytes.unwrap_or_default()) + .sum(); + + RowGroupMetaView { + index: row_group_idx, + num_rows: row_group.num_rows(), + uncompressed_size: row_group.total_byte_size(), + compressed_size: row_group.compressed_size(), + compression_ratio: compression_ratio( + row_group.total_byte_size(), + row_group.compressed_size(), + ), + data_pages, + dictionary_page_bytes, + columns, + } + }) + .collect(); + + FileMetaView { + input: path.display().to_string(), + num_rows: metadata.file_metadata().num_rows(), + num_row_groups: metadata.num_row_groups(), + num_columns: metadata.file_metadata().schema_descr().num_columns(), + key_value_metadata, + row_groups, + } +} + +fn print_meta_text(view: &FileMetaView) { + println!("file: {}", view.input); + println!("rows: {}", view.num_rows); + println!("row_groups: {}", view.num_row_groups); + println!("columns: {}", view.num_columns); + println!("key_value_metadata: {}", view.key_value_metadata.len()); + for row_group in &view.row_groups { + println!( + "row_group[{}]: rows={}, uncompressed_size={}, compressed_size={}, compression_ratio={}, data_pages={}, dictionary_page_bytes={}", + row_group.index, + row_group.num_rows, + row_group.uncompressed_size, + row_group.compressed_size, + format_ratio(row_group.compression_ratio), + format_optional_usize(row_group.data_pages), + row_group.dictionary_page_bytes, + ); + for column in &row_group.columns { + println!( + " column[{}] path={}: type={}, compression={}, encodings=[{}], values={}, uncompressed_size={}, compressed_size={}, compression_ratio={}, data_pages={}, dictionary_page_offset={}, dictionary_page_bytes={}, data_page_offset={}, statistics={}, column_index={}/{}, offset_index={}/{}, bloom_filter={}/{}", + column.index, + column.path, + column.physical_type, + column.compression, + column.encodings.join(","), + column.num_values, + column.uncompressed_size, + column.compressed_size, + format_ratio(column.compression_ratio), + format_optional_usize(column.data_pages), + format_optional_i64(column.dictionary_page_offset), + column + .dictionary_page_bytes + .map(|value| value.to_string()) + .unwrap_or_else(|| "none".to_string()), + column.data_page_offset, + column.has_statistics, + format_optional_i64(column.column_index_offset), + format_optional_i32(column.column_index_length), + format_optional_i64(column.offset_index_offset), + format_optional_i32(column.offset_index_length), + format_optional_i64(column.bloom_filter_offset), + format_optional_i32(column.bloom_filter_length), + ); + } + } +} + +fn dictionary_page_bytes( + dictionary_page_offset: Option, + data_page_offset: i64, +) -> Option { + dictionary_page_offset.map(|offset| data_page_offset.saturating_sub(offset)) +} + +fn compression_ratio(uncompressed: i64, compressed: i64) -> Option { + (uncompressed > 0).then(|| compressed as f64 / uncompressed as f64) +} + +fn format_ratio(ratio: Option) -> String { + ratio + .map(|value| format!("{value:.4}")) + .unwrap_or_else(|| "unknown".to_string()) +} + +fn format_optional_usize(value: Option) -> String { + value + .map(|value| value.to_string()) + .unwrap_or_else(|| "unknown".to_string()) +} + +fn format_optional_i64(value: Option) -> String { + value + .map(|value| value.to_string()) + .unwrap_or_else(|| "none".to_string()) +} + +fn format_optional_i32(value: Option) -> String { + value + .map(|value| value.to_string()) + .unwrap_or_else(|| "none".to_string()) +} + +#[cfg(test)] +mod tests { + use std::fs::File; + use std::sync::Arc; + + use datatypes::arrow::array::{Int32Array, RecordBatch, StringArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use parquet::arrow::ArrowWriter; + use parquet::file::metadata::KeyValue; + use parquet::file::properties::WriterProperties; + use tempfile::tempdir; + + use super::*; + + #[test] + fn test_meta_view_unknown_page_count() { + let dir = tempdir().unwrap(); + let path = dir.path().join("input.parquet"); + write_test_parquet(&path); + + let metadata = load_local_parquet_metadata(&path).unwrap(); + let view = build_file_meta_view(&path, &metadata); + + assert_eq!(view.num_row_groups, 1); + assert!(view.row_groups[0].uncompressed_size > 0); + assert!(view.row_groups[0].compressed_size > 0); + assert!(view.row_groups[0].compression_ratio.is_some()); + assert_eq!(view.row_groups[0].columns.len(), 2); + } + + fn write_test_parquet(path: &Path) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("host", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(StringArray::from(vec!["a", "a", "b", "b"])), + ], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_key_value_metadata(Some(vec![KeyValue::new( + "greptime:test".to_string(), + "value".to_string(), + )])) + .build(); + let mut writer = + ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } +} diff --git a/src/cmd/src/datanode/parquet_rewrite.rs b/src/cmd/src/datanode/parquet_rewrite.rs new file mode 100644 index 00000000000..47346aac889 --- /dev/null +++ b/src/cmd/src/datanode/parquet_rewrite.rs @@ -0,0 +1,579 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fs::File; +use std::path::{Path, PathBuf}; + +use clap::Parser; +use datatypes::arrow::record_batch::RecordBatchReader; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::basic::{BrotliLevel, Compression, Encoding, GzipLevel, ZstdLevel}; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::properties::WriterProperties; +use parquet::file::reader::FileReader; +use parquet::file::serialized_reader::SerializedFileReader; +use parquet::schema::types::ColumnPath; +use serde::{Deserialize, Serialize}; +use snafu::ResultExt; + +use crate::datanode::tool_util::load_local_parquet_metadata; +use crate::error; + +/// Read and rewrite a parquet file with different writer properties. +#[derive(Debug, Parser)] +pub struct ParquetRewriteCommand { + /// Path to input parquet file. + #[clap(long, value_name = "FILE")] + input: PathBuf, + + /// Path to output parquet file in rewrite mode. + #[clap(long, value_name = "FILE")] + output: Option, + + /// Path to writer properties TOML in rewrite mode. + #[clap(long, value_name = "FILE")] + properties: Option, + + /// Dump writer properties TOML inferred from the input parquet file. + #[clap(long, value_name = "FILE")] + dump_properties: Option, + + /// Number of rows per record batch. + #[clap(long)] + batch_size: Option, + + /// Overwrite output files. + #[clap(long, default_value_t = false)] + overwrite: bool, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default, deny_unknown_fields)] +struct RewriteProperties { + writer: WriterConfig, + columns: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] +#[serde(default, deny_unknown_fields)] +struct WriterConfig { + dictionary_enabled: Option, + compression: Option, + compression_level: Option, + encoding: Option, + max_row_group_row_count: Option, + data_page_size_limit: Option, + data_page_row_count_limit: Option, + dictionary_page_size_limit: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default, deny_unknown_fields)] +struct ColumnConfig { + path: Vec, + dictionary_enabled: Option, + compression: Option, + compression_level: Option, + encoding: Option, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "kebab-case")] +enum CompressionConfig { + Uncompressed, + Snappy, + Gzip, + Lzo, + Brotli, + Lz4, + Zstd, + Lz4Raw, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "kebab-case")] +enum EncodingConfig { + Plain, + DeltaBinaryPacked, + DeltaLengthByteArray, + DeltaByteArray, + ByteStreamSplit, +} + +impl ParquetRewriteCommand { + pub async fn run(&self) -> error::Result<()> { + match (&self.dump_properties, &self.output, &self.properties) { + (Some(path), None, None) => self.dump_properties(path), + (None, Some(output), Some(properties)) => self.rewrite(output, properties), + (Some(_), Some(_), _) | (Some(_), _, Some(_)) => illegal_config( + "use either --dump-properties or rewrite mode, not both".to_string(), + ), + (None, _, _) => illegal_config( + "rewrite mode requires --output and --properties; config dump mode requires --dump-properties".to_string(), + ), + } + } + + fn dump_properties(&self, path: &Path) -> error::Result<()> { + ensure_distinct_paths(&self.input, path)?; + ensure_can_write(path, self.overwrite)?; + let metadata = load_local_parquet_metadata(&self.input)?; + let properties = infer_rewrite_properties(&metadata); + let content = toml::to_string_pretty(&properties).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to serialize writer properties: {e}"), + } + .build() + })?; + std::fs::write(path, content).context(error::FileIoSnafu)?; + println!("Wrote writer properties to {}", path.display()); + Ok(()) + } + + fn rewrite(&self, output: &Path, properties: &Path) -> error::Result<()> { + ensure_distinct_paths(&self.input, output)?; + ensure_can_write(output, self.overwrite)?; + + let props = load_rewrite_properties(properties)?; + let input_reader = SerializedFileReader::new(open_file(&self.input)?) + .map_err(|e| parquet_error("read source parquet metadata", &self.input, e))?; + let key_value_metadata = input_reader + .metadata() + .file_metadata() + .key_value_metadata() + .cloned(); + + let mut reader_builder = ParquetRecordBatchReaderBuilder::try_new(open_file(&self.input)?) + .map_err(|e| parquet_error("open source parquet", &self.input, e))?; + if let Some(batch_size) = self.batch_size { + if batch_size == 0 { + return illegal_config("--batch-size must be greater than 0".to_string()); + } + reader_builder = reader_builder.with_batch_size(batch_size); + } + let reader = reader_builder + .build() + .map_err(|e| parquet_error("build parquet reader", &self.input, e))?; + let schema = reader.schema(); + + let writer_props = build_writer_properties(props, key_value_metadata)?; + let mut writer = ArrowWriter::try_new(create_file(output)?, schema, Some(writer_props)) + .map_err(|e| parquet_error("create parquet writer", output, e))?; + + for batch in reader { + let batch = + batch.map_err(|e| parquet_error("read parquet batch", &self.input, e.into()))?; + writer + .write(&batch) + .map_err(|e| parquet_error("write parquet batch", output, e))?; + } + writer + .close() + .map_err(|e| parquet_error("close parquet writer", output, e))?; + + println!("Wrote parquet file to {}", output.display()); + Ok(()) + } +} + +fn infer_rewrite_properties(metadata: &ParquetMetaData) -> RewriteProperties { + let first_column = metadata + .row_groups() + .first() + .and_then(|row_group| row_group.columns().first()); + let compression = first_column.map(|column| compression_to_config(column.compression())); + let max_row_group_row_count = metadata + .row_groups() + .first() + .and_then(|row_group| usize::try_from(row_group.num_rows()).ok()); + let dictionary_enabled = first_column + .map(|column| column.dictionary_page_offset().is_some()) + .or(Some(true)); + + let columns = metadata + .file_metadata() + .schema_descr() + .columns() + .iter() + .enumerate() + .map(|(idx, column)| { + let first_chunk = metadata + .row_groups() + .first() + .and_then(|row_group| row_group.columns().get(idx)); + ColumnConfig { + path: column.path().parts().to_vec(), + dictionary_enabled: first_chunk + .map(|chunk| chunk.dictionary_page_offset().is_some()), + compression: first_chunk.map(|chunk| compression_to_config(chunk.compression())), + compression_level: None, + encoding: first_chunk.and_then(infer_data_encoding), + } + }) + .collect(); + + RewriteProperties { + writer: WriterConfig { + dictionary_enabled, + compression, + compression_level: None, + encoding: first_column.and_then(infer_data_encoding), + max_row_group_row_count, + data_page_size_limit: None, + data_page_row_count_limit: None, + dictionary_page_size_limit: None, + }, + columns, + } +} + +fn load_rewrite_properties(path: &Path) -> error::Result { + let content = std::fs::read_to_string(path).context(error::FileIoSnafu)?; + toml::from_str(&content).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to parse writer properties {}: {e}", path.display()), + } + .build() + }) +} + +fn build_writer_properties( + config: RewriteProperties, + key_value_metadata: Option>, +) -> error::Result { + let mut builder = WriterProperties::builder().set_key_value_metadata(key_value_metadata); + let writer_compression = config.writer.compression; + if let Some(dictionary_enabled) = config.writer.dictionary_enabled { + builder = builder.set_dictionary_enabled(dictionary_enabled); + } + match (writer_compression, config.writer.compression_level) { + (Some(compression), level) => { + builder = builder.set_compression(to_parquet_compression(compression, level)?); + } + (None, Some(_)) => { + return illegal_config( + "writer compression_level requires writer compression".to_string(), + ); + } + (None, None) => {} + } + if let Some(encoding) = config.writer.encoding { + builder = builder.set_encoding(to_parquet_encoding(encoding)); + } + if let Some(max_row_group_row_count) = config.writer.max_row_group_row_count { + if max_row_group_row_count == 0 { + return illegal_config("max_row_group_row_count must be greater than 0".to_string()); + } + builder = builder.set_max_row_group_row_count(Some(max_row_group_row_count)); + } + if let Some(data_page_size_limit) = config.writer.data_page_size_limit { + builder = builder.set_data_page_size_limit(data_page_size_limit); + } + if let Some(data_page_row_count_limit) = config.writer.data_page_row_count_limit { + builder = builder.set_data_page_row_count_limit(data_page_row_count_limit); + } + if let Some(dictionary_page_size_limit) = config.writer.dictionary_page_size_limit { + builder = builder.set_dictionary_page_size_limit(dictionary_page_size_limit); + } + + for column in config.columns { + if column.path.is_empty() { + return illegal_config("column path must not be empty".to_string()); + } + let path = ColumnPath::new(column.path); + if let Some(dictionary_enabled) = column.dictionary_enabled { + builder = builder.set_column_dictionary_enabled(path.clone(), dictionary_enabled); + } + if column.compression.is_some() || column.compression_level.is_some() { + let Some(compression) = column.compression.or(writer_compression) else { + return illegal_config(format!( + "compression_level for column {} requires column or writer compression", + path.string() + )); + }; + builder = builder.set_column_compression( + path.clone(), + to_parquet_compression(compression, column.compression_level)?, + ); + } + if let Some(encoding) = column.encoding { + builder = builder.set_column_encoding(path, to_parquet_encoding(encoding)); + } + } + + Ok(builder.build()) +} + +fn ensure_distinct_paths(input: &Path, output: &Path) -> error::Result<()> { + if input == output { + return illegal_config(format!( + "input and output paths must be different: {}", + input.display() + )); + } + Ok(()) +} + +fn ensure_can_write(path: &Path, overwrite: bool) -> error::Result<()> { + if !overwrite && path.exists() { + return illegal_config(format!( + "{} already exists; pass --overwrite to replace it", + path.display() + )); + } + Ok(()) +} + +fn open_file(path: &Path) -> error::Result { + File::open(path).context(error::FileIoSnafu) +} + +fn create_file(path: &Path) -> error::Result { + File::create(path).context(error::FileIoSnafu) +} + +fn illegal_config(msg: String) -> error::Result { + error::IllegalConfigSnafu { msg }.fail() +} + +fn parquet_error( + action: &'static str, + path: &Path, + error: parquet::errors::ParquetError, +) -> error::Error { + error::IllegalConfigSnafu { + msg: format!("{action} failed for {}: {error}", path.display()), + } + .build() +} + +fn to_parquet_compression( + compression: CompressionConfig, + level: Option, +) -> error::Result { + if level.is_some() + && !matches!( + compression, + CompressionConfig::Gzip | CompressionConfig::Brotli | CompressionConfig::Zstd + ) + { + return illegal_config(format!( + "compression level is not supported for {}", + compression.name() + )); + } + + Ok(match compression { + CompressionConfig::Uncompressed => Compression::UNCOMPRESSED, + CompressionConfig::Snappy => Compression::SNAPPY, + CompressionConfig::Gzip => Compression::GZIP(match level { + Some(level) => GzipLevel::try_new(level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid gzip compression level {level}: {e}"), + } + .build() + })?, + None => GzipLevel::default(), + }), + CompressionConfig::Lzo => Compression::LZO, + CompressionConfig::Brotli => Compression::BROTLI(match level { + Some(level) => BrotliLevel::try_new(level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid brotli compression level {level}: {e}"), + } + .build() + })?, + None => BrotliLevel::default(), + }), + CompressionConfig::Lz4 => Compression::LZ4, + CompressionConfig::Zstd => Compression::ZSTD(match level { + Some(level) => { + let converted_level = i32::try_from(level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid zstd compression level {level}: {e}"), + } + .build() + })?; + ZstdLevel::try_new(converted_level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid zstd compression level {level}: {e}"), + } + .build() + })? + } + None => ZstdLevel::default(), + }), + CompressionConfig::Lz4Raw => Compression::LZ4_RAW, + }) +} + +impl CompressionConfig { + fn name(self) -> &'static str { + match self { + CompressionConfig::Uncompressed => "uncompressed", + CompressionConfig::Snappy => "snappy", + CompressionConfig::Gzip => "gzip", + CompressionConfig::Lzo => "lzo", + CompressionConfig::Brotli => "brotli", + CompressionConfig::Lz4 => "lz4", + CompressionConfig::Zstd => "zstd", + CompressionConfig::Lz4Raw => "lz4-raw", + } + } +} + +fn to_parquet_encoding(encoding: EncodingConfig) -> Encoding { + match encoding { + EncodingConfig::Plain => Encoding::PLAIN, + EncodingConfig::DeltaBinaryPacked => Encoding::DELTA_BINARY_PACKED, + EncodingConfig::DeltaLengthByteArray => Encoding::DELTA_LENGTH_BYTE_ARRAY, + EncodingConfig::DeltaByteArray => Encoding::DELTA_BYTE_ARRAY, + EncodingConfig::ByteStreamSplit => Encoding::BYTE_STREAM_SPLIT, + } +} + +fn infer_data_encoding( + column: &parquet::file::metadata::ColumnChunkMetaData, +) -> Option { + let encodings: Vec<_> = column.encodings().collect(); + if encodings.contains(&Encoding::DELTA_BINARY_PACKED) { + Some(EncodingConfig::DeltaBinaryPacked) + } else if encodings.contains(&Encoding::DELTA_LENGTH_BYTE_ARRAY) { + Some(EncodingConfig::DeltaLengthByteArray) + } else if encodings.contains(&Encoding::DELTA_BYTE_ARRAY) { + Some(EncodingConfig::DeltaByteArray) + } else if encodings.contains(&Encoding::BYTE_STREAM_SPLIT) { + Some(EncodingConfig::ByteStreamSplit) + } else if encodings.contains(&Encoding::PLAIN) { + Some(EncodingConfig::Plain) + } else { + None + } +} + +fn compression_to_config(compression: Compression) -> CompressionConfig { + match compression { + Compression::UNCOMPRESSED => CompressionConfig::Uncompressed, + Compression::SNAPPY => CompressionConfig::Snappy, + Compression::GZIP(_) => CompressionConfig::Gzip, + Compression::LZO => CompressionConfig::Lzo, + Compression::BROTLI(_) => CompressionConfig::Brotli, + Compression::LZ4 => CompressionConfig::Lz4, + Compression::ZSTD(_) => CompressionConfig::Zstd, + Compression::LZ4_RAW => CompressionConfig::Lz4Raw, + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datatypes::arrow::array::{Int32Array, RecordBatch, StringArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use parquet::file::metadata::KeyValue; + use tempfile::tempdir; + + use super::*; + + #[test] + fn test_dump_and_rewrite_preserves_key_value_metadata_and_disables_dictionary() { + let dir = tempdir().unwrap(); + let input = dir.path().join("input.parquet"); + let output = dir.path().join("output.parquet"); + write_test_parquet(&input, true); + + let metadata = load_local_parquet_metadata(&input).unwrap(); + let mut properties = infer_rewrite_properties(&metadata); + properties.columns[1].dictionary_enabled = Some(false); + let props_path = dir.path().join("props.toml"); + std::fs::write(&props_path, toml::to_string(&properties).unwrap()).unwrap(); + + let command = ParquetRewriteCommand { + input: input.clone(), + output: Some(output.clone()), + properties: Some(props_path), + dump_properties: None, + batch_size: Some(2), + overwrite: false, + }; + command + .rewrite(&output, command.properties.as_ref().unwrap()) + .unwrap(); + + let rewritten = load_local_parquet_metadata(&output).unwrap(); + assert_eq!(rewritten.file_metadata().num_rows(), 4); + let key_values = rewritten + .file_metadata() + .key_value_metadata() + .cloned() + .unwrap_or_default(); + assert!(key_values.iter().any(|kv| kv.key == "greptime:test")); + assert!( + rewritten + .row_groups() + .iter() + .all(|row_group| row_group.column(1).dictionary_page_offset().is_none()) + ); + } + + #[test] + fn test_column_compression_level_inherits_writer_compression() { + let path = ColumnPath::new(vec!["host".to_string()]); + let config = RewriteProperties { + writer: WriterConfig { + compression: Some(CompressionConfig::Zstd), + compression_level: Some(1), + ..Default::default() + }, + columns: vec![ColumnConfig { + path: vec!["host".to_string()], + compression_level: Some(3), + ..Default::default() + }], + }; + + let properties = build_writer_properties(config, None).unwrap(); + assert_eq!( + properties.compression(&path), + Compression::ZSTD(ZstdLevel::try_new(3).unwrap()) + ); + } + + fn write_test_parquet(path: &Path, dictionary_enabled: bool) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("host", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(StringArray::from(vec!["a", "a", "b", "b"])), + ], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(dictionary_enabled) + .set_key_value_metadata(Some(vec![KeyValue::new( + "greptime:test".to_string(), + "value".to_string(), + )])) + .build(); + let mut writer = + ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } +} diff --git a/src/cmd/src/datanode/parquetbench.rs b/src/cmd/src/datanode/parquetbench.rs index 2b46a775f76..09ec50f48ef 100644 --- a/src/cmd/src/datanode/parquetbench.rs +++ b/src/cmd/src/datanode/parquetbench.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -42,7 +42,10 @@ use store_api::region_request::PathType; use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, is_internal_column}; use store_api::storage::{ColumnId, FileId, RegionId}; -use crate::datanode::objbench::{build_object_store, extract_region_metadata, parse_config}; +use crate::datanode::tool_util::{ + build_object_store, extract_region_metadata, format_bytes, parse_config, parse_file_id, + parse_path_type, parse_region_id, +}; use crate::error; const DEFAULT_READ_BATCH_SIZE: usize = 8 * 1024; @@ -52,19 +55,23 @@ const DEFAULT_READ_BATCH_SIZE: usize = 8 * 1024; pub struct ParquetbenchCommand { /// Path to config TOML file (same format as standalone/datanode config) #[clap(long, value_name = "FILE")] - config: PathBuf, + config: Option, /// Region ID: either numeric u64 (e.g. "4398046511104") or "table_id:region_num" (e.g. "1024:0") #[clap(long)] - region_id: String, + region_id: Option, /// Table directory relative to data home (e.g. "data/greptime/public/1024/") #[clap(long)] - table_dir: String, + table_dir: Option, /// SST file id to benchmark. #[clap(long)] - file_id: String, + file_id: Option, + + /// Local parquet SST file to benchmark with the direct reader. + #[clap(long, value_name = "FILE")] + file_path: Option, /// Path to scan request JSON config file (supports projection_names only) #[clap(long, value_name = "FILE")] @@ -129,6 +136,32 @@ struct IterationStats { elapsed: Duration, } +#[derive(Debug)] +enum ParquetbenchInput { + LocalFile { + file_path: PathBuf, + }, + Region { + config: PathBuf, + region_id: String, + table_dir: String, + file_id: String, + path_type: PathType, + }, +} + +struct ParquetbenchSource { + object_store: object_store::ObjectStore, + file_path: String, + display_path: String, + region_id_label: String, + region_id: RegionId, + file_id: FileId, + region_file_id: RegionFileId, + path_type: Option, + table_dir: Option, +} + impl ParquetbenchCommand { pub async fn run(&self) -> error::Result<()> { if self.verbose { @@ -145,41 +178,40 @@ impl ParquetbenchCommand { .fail(); } - let region_id = parse_region_id(&self.region_id)?; - let path_type = parse_path_type(&self.path_type)?; - let file_id = FileId::parse_str(&self.file_id).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid file_id '{}': {}", self.file_id, e), - } - .build() - })?; - let region_file_id = RegionFileId::new(region_id, file_id); - let file_path = sst_file_path(&self.table_dir, region_file_id, path_type); + let input = self.resolve_input()?; + let mut source = build_source(input).await?; - let (store_cfg, _mito_config, _wal_config) = parse_config(&self.config)?; - let object_store = build_object_store(&store_cfg).await?; - - let file_size = object_store - .stat(&file_path) + let file_size = source + .object_store + .stat(&source.file_path) .await .map_err(|e| { error::IllegalConfigSnafu { - msg: format!("stat failed for {}: {}", file_path, e), + msg: format!("stat failed for {}: {}", source.display_path, e), } .build() })? .content_length(); let mut metadata_metrics = MetadataCacheMetrics::default(); - let parquet_meta = MetadataLoader::new(object_store.clone(), &file_path, file_size) - .load(&mut metadata_metrics) - .await - .map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("read parquet metadata failed for {}: {:?}", file_path, e), - } - .build() - })?; - let region_meta = extract_region_metadata(&file_path, &parquet_meta)?; + let parquet_meta = + MetadataLoader::new(source.object_store.clone(), &source.file_path, file_size) + .load(&mut metadata_metrics) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!( + "read parquet metadata failed for {}: {:?}", + source.display_path, e + ), + } + .build() + })?; + let region_meta = extract_region_metadata(&source.display_path, &parquet_meta)?; + if source.table_dir.is_none() { + source.region_id = region_meta.region_id; + source.region_id_label = source.region_id.as_u64().to_string(); + source.region_file_id = RegionFileId::new(source.region_id, source.file_id); + } let scan_config = self.load_scan_config().await?; let projection = if self.reader == ReaderMode::Direct { resolve_projection_names(&scan_config, ®ion_meta)? @@ -216,10 +248,10 @@ impl ParquetbenchCommand { println!( "{} Region ID: {} (u64: {})", "✓".green(), - self.region_id, - region_id.as_u64() + source.region_id_label, + source.region_id.as_u64() ); - println!("{} File path: {}", "✓".green(), file_path.cyan()); + println!("{} File path: {}", "✓".green(), source.display_path.cyan()); println!( "{} Columns: {}", "✓".green(), @@ -307,8 +339,8 @@ impl ParquetbenchCommand { let mut schema_printed = false; let file_handle = FileHandle::new( FileMeta { - region_id, - file_id, + region_id: source.region_id, + file_id: source.file_id, time_range: Default::default(), level: 0, file_size, @@ -332,9 +364,9 @@ impl ParquetbenchCommand { let stats = match self.reader { ReaderMode::Direct => { run_direct_iteration( - object_store.clone(), - file_path.clone(), - region_file_id, + source.object_store.clone(), + source.file_path.clone(), + source.region_file_id, parquet_meta.clone(), projection.clone(), row_groups.clone(), @@ -345,9 +377,19 @@ impl ParquetbenchCommand { } ReaderMode::FlatPrune => { run_flat_prune_iteration( - object_store.clone(), - self.table_dir.clone(), - path_type, + source.object_store.clone(), + source.table_dir.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "flat-prune reader requires --table-dir".to_string(), + } + .build() + })?, + source.path_type.ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "flat-prune reader requires --path-type".to_string(), + } + .build() + })?, file_handle.clone(), region_meta.clone(), projection_column_ids.clone(), @@ -450,6 +492,65 @@ impl ParquetbenchCommand { Ok(()) } + fn resolve_input(&self) -> error::Result { + let has_region_args = self.config.is_some() + || self.region_id.is_some() + || self.table_dir.is_some() + || self.file_id.is_some(); + + if let Some(file_path) = &self.file_path { + if self.reader == ReaderMode::FlatPrune { + return Err(error::IllegalConfigSnafu { + msg: "--file-path currently supports only --reader direct".to_string(), + } + .build()); + } + if has_region_args { + return Err(error::IllegalConfigSnafu { + msg: "--file-path cannot be used with --config, --region-id, --table-dir, or --file-id".to_string(), + } + .build()); + } + return Ok(ParquetbenchInput::LocalFile { + file_path: file_path.clone(), + }); + } + + let config = self.config.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --config unless --file-path is specified".to_string(), + } + .build() + })?; + let region_id = self.region_id.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --region-id unless --file-path is specified".to_string(), + } + .build() + })?; + let table_dir = self.table_dir.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --table-dir unless --file-path is specified".to_string(), + } + .build() + })?; + let file_id = self.file_id.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --file-id unless --file-path is specified".to_string(), + } + .build() + })?; + let path_type = parse_path_type(&self.path_type)?; + + Ok(ParquetbenchInput::Region { + config, + region_id, + table_dir, + file_id, + path_type, + }) + } + async fn load_scan_config(&self) -> error::Result { if let Some(path) = &self.scan_config { let content = tokio::fs::read_to_string(path) @@ -462,6 +563,117 @@ impl ParquetbenchCommand { } } +async fn build_source(input: ParquetbenchInput) -> error::Result { + match input { + ParquetbenchInput::LocalFile { file_path } => build_local_file_source(&file_path), + ParquetbenchInput::Region { + config, + region_id, + table_dir, + file_id, + path_type, + } => build_region_source(&config, ®ion_id, table_dir, &file_id, path_type).await, + } +} + +fn build_local_file_source(file_path: &Path) -> error::Result { + let file_path = std::fs::canonicalize(file_path).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid --file-path {}: {e}", file_path.display()), + } + .build() + })?; + if !file_path.is_file() { + return Err(error::IllegalConfigSnafu { + msg: format!("--file-path {} is not a file", file_path.display()), + } + .build()); + } + let parent = file_path.parent().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!( + "--file-path {} has no parent directory", + file_path.display() + ), + } + .build() + })?; + let file_name = file_path + .file_name() + .and_then(|name| name.to_str()) + .ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!("invalid UTF-8 file name in {}", file_path.display()), + } + .build() + })? + .to_string(); + let object_store = object_store::ObjectStore::new(object_store::services::Fs::default().root( + parent.to_str().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!("invalid UTF-8 parent directory in {}", file_path.display()), + } + .build() + })?, + )) + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to build local file object store: {e:?}"), + } + .build() + })? + .finish(); + + let file_id = file_path + .file_stem() + .and_then(|stem| stem.to_str()) + .and_then(|stem| FileId::parse_str(stem).ok()) + .unwrap_or_else(FileId::random); + let region_id = RegionId::new(0, 0); + let region_file_id = RegionFileId::new(region_id, file_id); + let display_path = file_path.display().to_string(); + + Ok(ParquetbenchSource { + object_store, + file_path: file_name, + display_path, + region_id_label: region_id.as_u64().to_string(), + region_id, + file_id, + region_file_id, + path_type: None, + table_dir: None, + }) +} + +async fn build_region_source( + config: &Path, + region_id: &str, + table_dir: String, + file_id: &str, + path_type: PathType, +) -> error::Result { + let region = parse_region_id(region_id)?; + let file_id = parse_file_id(file_id)?; + let region_file_id = RegionFileId::new(region, file_id); + let file_path = sst_file_path(&table_dir, region_file_id, path_type); + + let (store_cfg, _mito_config, _wal_config) = parse_config(config)?; + let object_store = build_object_store(&store_cfg).await?; + + Ok(ParquetbenchSource { + object_store, + display_path: file_path.clone(), + file_path, + region_id_label: region_id.to_string(), + region_id: region, + file_id, + region_file_id, + path_type: Some(path_type), + table_dir: Some(table_dir), + }) +} + #[allow(clippy::too_many_arguments)] async fn run_direct_iteration( object_store: object_store::ObjectStore, @@ -773,60 +985,6 @@ fn parse_batch_size(s: &str) -> Result { Ok(batch_size) } -fn parse_region_id(s: &str) -> error::Result { - if s.contains(':') { - let parts: Vec<&str> = s.splitn(2, ':').collect(); - let table_id: u32 = parts[0].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid table_id in region_id '{}': {}", s, e), - } - .build() - })?; - let region_num: u32 = parts[1].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_num in region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::new(table_id, region_num)) - } else { - let id: u64 = s.parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::from_u64(id)) - } -} - -fn parse_path_type(s: &str) -> error::Result { - match s.to_lowercase().as_str() { - "bare" => Ok(PathType::Bare), - "data" => Ok(PathType::Data), - "metadata" => Ok(PathType::Metadata), - _ => Err(error::IllegalConfigSnafu { - msg: format!("invalid path_type '{}', expected: bare, data, metadata", s), - } - .build()), - } -} - -fn format_bytes(bytes: u64) -> String { - const KIB: u64 = 1024; - const MIB: u64 = 1024 * KIB; - const GIB: u64 = 1024 * MIB; - if bytes >= GIB { - format!("{:.2} GiB", bytes as f64 / GIB as f64) - } else if bytes >= MIB { - format!("{:.2} MiB", bytes as f64 / MIB as f64) - } else if bytes >= KIB { - format!("{:.2} KiB", bytes as f64 / KIB as f64) - } else { - format!("{} B", bytes) - } -} - fn format_rate(rate: f64) -> String { if !rate.is_finite() { return "inf rows".to_string(); @@ -851,6 +1009,25 @@ mod tests { use super::*; + fn test_command() -> ParquetbenchCommand { + ParquetbenchCommand { + config: None, + region_id: None, + table_dir: None, + file_id: None, + file_path: None, + scan_config: None, + iterations: 1, + batch_size: DEFAULT_READ_BATCH_SIZE, + path_type: "bare".to_string(), + verbose: false, + pprof_file: None, + pprof_after_warmup: false, + pk_as_binary: false, + reader: ReaderMode::Direct, + } + } + fn new_test_metadata() -> RegionMetadata { let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 0)); builder @@ -882,19 +1059,27 @@ mod tests { } #[test] - fn test_parse_region_id() { - assert_eq!(parse_region_id("1024:7").unwrap(), RegionId::new(1024, 7)); - assert_eq!( - parse_region_id(&RegionId::new(1, 2).as_u64().to_string()).unwrap(), - RegionId::new(1, 2) - ); + fn test_resolve_input_accepts_direct_file_for_direct_reader() { + let mut command = test_command(); + command.file_path = Some(PathBuf::from("/tmp/source.parquet")); + + let input = command.resolve_input().unwrap(); + match input { + ParquetbenchInput::LocalFile { file_path } => { + assert_eq!(file_path, PathBuf::from("/tmp/source.parquet")); + } + ParquetbenchInput::Region { .. } => panic!("expected local file input"), + } } #[test] - fn test_parse_path_type() { - assert_eq!(parse_path_type("bare").unwrap(), PathType::Bare); - assert_eq!(parse_path_type("data").unwrap(), PathType::Data); - assert_eq!(parse_path_type("metadata").unwrap(), PathType::Metadata); + fn test_resolve_input_rejects_mixed_input_modes() { + let mut command = test_command(); + command.file_path = Some(PathBuf::from("/tmp/source.parquet")); + command.config = Some(PathBuf::from("config.toml")); + + let err = command.resolve_input().unwrap_err(); + assert!(err.to_string().contains("--file-path cannot be used with")); } #[test] diff --git a/src/cmd/src/datanode/scanbench.rs b/src/cmd/src/datanode/scanbench.rs index 4909467f481..1dbee025a7f 100644 --- a/src/cmd/src/datanode/scanbench.rs +++ b/src/cmd/src/datanode/scanbench.rs @@ -52,11 +52,13 @@ use sqlparser::parser::Parser as SqlParser; use store_api::metadata::RegionMetadata; use store_api::path_utils::WAL_DIR; use store_api::region_engine::{PrepareRequest, QueryScanContext, RegionEngine}; -use store_api::region_request::{PathType, RegionOpenRequest, RegionRequest}; +use store_api::region_request::{RegionOpenRequest, RegionRequest}; use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector}; use tokio::fs; -use crate::datanode::objbench::{build_object_store, parse_config}; +use crate::datanode::tool_util::{ + build_object_store, format_bytes, parse_config, parse_path_type, parse_region_id, +}; use crate::error; /// Scan benchmark command - benchmarks scanning a region directly from storage. @@ -164,60 +166,6 @@ fn resolve_projection( Ok(None) } -fn format_bytes(bytes: u64) -> String { - const KIB: u64 = 1024; - const MIB: u64 = 1024 * KIB; - const GIB: u64 = 1024 * MIB; - if bytes >= GIB { - format!("{:.2} GiB", bytes as f64 / GIB as f64) - } else if bytes >= MIB { - format!("{:.2} MiB", bytes as f64 / MIB as f64) - } else if bytes >= KIB { - format!("{:.2} KiB", bytes as f64 / KIB as f64) - } else { - format!("{} B", bytes) - } -} - -fn parse_region_id(s: &str) -> error::Result { - if s.contains(':') { - let parts: Vec<&str> = s.splitn(2, ':').collect(); - let table_id: u32 = parts[0].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid table_id in region_id '{}': {}", s, e), - } - .build() - })?; - let region_num: u32 = parts[1].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_num in region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::new(table_id, region_num)) - } else { - let id: u64 = s.parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::from_u64(id)) - } -} - -fn parse_path_type(s: &str) -> error::Result { - match s.to_lowercase().as_str() { - "bare" => Ok(PathType::Bare), - "data" => Ok(PathType::Data), - "metadata" => Ok(PathType::Metadata), - _ => Err(error::IllegalConfigSnafu { - msg: format!("invalid path_type '{}', expected: bare, data, metadata", s), - } - .build()), - } -} - /// Rewrites literal values in comparison expressions to match the column's arrow type. struct LiteralTypeCaster { schema: DFSchemaRef, diff --git a/src/cmd/src/datanode/sst_replace.rs b/src/cmd/src/datanode/sst_replace.rs new file mode 100644 index 00000000000..df13fcd6f25 --- /dev/null +++ b/src/cmd/src/datanode/sst_replace.rs @@ -0,0 +1,361 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::path::PathBuf; + +use clap::Parser; +use colored::Colorize; +use mito2::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; +use mito2::manifest::manager::{RegionManifestManager, RegionManifestOptions}; +use mito2::region::ManifestStats; +use mito2::sst::file::RegionFileId; +use mito2::sst::location::{region_dir_from_table_dir, sst_file_path}; +use object_store::ObjectStore; +use parquet::file::FOOTER_SIZE; +use parquet::file::metadata::{FooterTail, ParquetMetaDataReader}; +use store_api::region_request::PathType; + +use crate::datanode::tool_util::{ + build_object_store, max_row_group_uncompressed_size, parse_config, parse_file_id, + parse_path_type, parse_region_id, +}; +use crate::error; + +/// Replace a mito region SST and update the corresponding manifest metadata. +#[derive(Debug, Parser)] +pub struct SstReplaceCommand { + /// Path to config TOML file (same format as standalone/datanode config). + #[clap(long, value_name = "FILE")] + config: PathBuf, + + /// Region ID: either numeric u64 (e.g. "4398046511104") or "table_id:region_num" (e.g. "1024:0"). + #[clap(long)] + region_id: String, + + /// Table directory relative to data home (e.g. "data/greptime/public/1024/"). + #[clap(long)] + table_dir: String, + + /// SST file id to replace. + #[clap(long)] + file_id: String, + + /// Local parquet file used as the replacement. + #[clap(long, value_name = "FILE", conflicts_with = "replacement_object")] + replacement_file: Option, + + /// Object-store parquet path used as the replacement. + #[clap(long, value_name = "PATH", conflicts_with = "replacement_file")] + replacement_object: Option, + + /// Path type for the region: auto, bare, data, metadata. + #[clap(long, default_value = "auto")] + path_type: String, + + /// Actually overwrite the SST object and append a manifest delta. + #[clap(long, default_value_t = false)] + confirm: bool, + + /// Verbose output. + #[clap(short, long, default_value_t = false)] + verbose: bool, +} + +impl SstReplaceCommand { + pub async fn run(&self) -> error::Result<()> { + if self.verbose { + common_telemetry::init_default_ut_logging(); + } + + println!("{}", "Starting sst-replace...".cyan().bold()); + + let region_id = parse_region_id(&self.region_id)?; + let file_id = parse_file_id(&self.file_id)?; + let path_type = parse_optional_path_type(&self.path_type)?; + let replacement = self.replacement()?; + + let (store_cfg, mito_config, _wal_config) = parse_config(&self.config)?; + let object_store = build_object_store(&store_cfg).await?; + println!("{} Object store initialized", "[ok]".green()); + + let candidates = path_type + .map(|path_type| vec![path_type]) + .unwrap_or_else(|| vec![PathType::Bare, PathType::Data, PathType::Metadata]); + let mut found = None; + let mut existing_manifests = Vec::new(); + for candidate in candidates { + let stats = ManifestStats::default(); + let region_dir = region_dir_from_table_dir(&self.table_dir, region_id, candidate); + let manifest_opts = + RegionManifestOptions::new(&mito_config, ®ion_dir, &object_store); + let Some(manifest_manager) = RegionManifestManager::open(manifest_opts, &stats) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("open {} manifest failed: {e:?}", path_type_name(candidate)), + } + .build() + })? + else { + continue; + }; + existing_manifests.push(path_type_name(candidate)); + let manifest = manifest_manager.manifest(); + if let Some(file) = manifest.files.get(&file_id) { + if found.is_some() { + return error::IllegalConfigSnafu { + msg: format!( + "file {} exists in multiple manifests under region {}; specify --path-type", + file_id, region_id + ), + } + .fail(); + } + found = Some((candidate, manifest_manager, file.clone())); + } + } + let Some((path_type, mut manifest_manager, old_file)) = found else { + let suffix = if existing_manifests.is_empty() { + "no manifest found".to_string() + } else { + format!( + "manifests found for [{}], but none contains file {}", + existing_manifests.join(", "), + file_id + ) + }; + return error::IllegalConfigSnafu { + msg: format!("region manifest not found for {}: {}", region_id, suffix), + } + .fail(); + }; + + let region_file_id = RegionFileId::new(old_file.region_id, old_file.file_id); + let target_path = sst_file_path(&self.table_dir, region_file_id, path_type); + let target_stat = object_store.stat(&target_path).await.map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("stat target SST {} failed: {e}", target_path), + } + .build() + })?; + let replacement_bytes = load_replacement_bytes(&object_store, replacement).await?; + let new_file_size = replacement_bytes.len() as u64; + let parquet_meta = decode_parquet_metadata(&replacement_bytes).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read replacement parquet metadata failed: {e}"), + } + .build() + })?; + let new_num_rows = parquet_meta.file_metadata().num_rows() as u64; + let new_num_row_groups = parquet_meta.num_row_groups() as u64; + validate_replacement(&old_file, new_num_rows, new_num_row_groups)?; + + let mut new_file = old_file.clone(); + new_file.file_size = new_file_size; + new_file.num_rows = new_num_rows; + new_file.num_row_groups = new_num_row_groups; + new_file.max_row_group_uncompressed_size = max_row_group_uncompressed_size(&parquet_meta); + + println!("{} Region: {}", "[ok]".green(), region_id); + println!("{} Target SST: {}", "[ok]".green(), target_path.cyan()); + println!( + "{} Size: {} -> {} bytes", + "[ok]".green(), + target_stat.content_length(), + new_file_size + ); + println!( + "{} Rows: {}, row groups: {}", + "[ok]".green(), + new_num_rows, + new_num_row_groups + ); + + if !self.confirm { + println!( + "{} Dry run only. Re-run with --confirm to overwrite the SST and update the manifest.", + "[dry-run]".yellow() + ); + return Ok(()); + } + + object_store + .write(&target_path, replacement_bytes) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("write replacement SST {} failed: {e}", target_path), + } + .build() + })?; + + let edit = RegionEdit { + files_to_add: vec![new_file], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }; + let version = manifest_manager + .update( + RegionMetaActionList::with_action(RegionMetaAction::Edit(edit)), + false, + ) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("update manifest failed: {e:?}"), + } + .build() + })?; + + println!( + "{} Replacement committed, manifest version {}", + "[ok]".green(), + version + ); + Ok(()) + } + + fn replacement(&self) -> error::Result> { + match (&self.replacement_file, &self.replacement_object) { + (Some(path), None) => Ok(Replacement::Local(path)), + (None, Some(path)) => Ok(Replacement::Object(path)), + _ => error::IllegalConfigSnafu { + msg: "specify exactly one of --replacement-file or --replacement-object" + .to_string(), + } + .fail(), + } + } +} + +enum Replacement<'a> { + Local(&'a PathBuf), + Object(&'a str), +} + +async fn load_replacement_bytes( + object_store: &ObjectStore, + replacement: Replacement<'_>, +) -> error::Result> { + match replacement { + Replacement::Local(path) => tokio::fs::read(path).await.map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read replacement file {} failed: {e}", path.display()), + } + .build() + }), + Replacement::Object(path) => { + object_store + .read(path) + .await + .map(|b| b.to_vec()) + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read replacement object {} failed: {e}", path), + } + .build() + }) + } + } +} + +fn validate_replacement( + old_file: &mito2::sst::file::FileMeta, + new_num_rows: u64, + new_num_row_groups: u64, +) -> error::Result<()> { + if old_file.num_rows != 0 && old_file.num_rows != new_num_rows { + return error::IllegalConfigSnafu { + msg: format!( + "replacement row count mismatch: manifest has {}, replacement has {}", + old_file.num_rows, new_num_rows + ), + } + .fail(); + } + if old_file.num_row_groups != 0 && old_file.num_row_groups != new_num_row_groups { + return error::IllegalConfigSnafu { + msg: format!( + "replacement row group count mismatch: manifest has {}, replacement has {}", + old_file.num_row_groups, new_num_row_groups + ), + } + .fail(); + } + Ok(()) +} + +fn decode_parquet_metadata( + data: &[u8], +) -> Result> { + if data.len() < FOOTER_SIZE { + return Err("file too small".into()); + } + let footer_start = data.len() - FOOTER_SIZE; + let mut footer = [0; FOOTER_SIZE]; + footer.copy_from_slice(&data[footer_start..]); + let footer = FooterTail::try_new(&footer)?; + let metadata_len = footer.metadata_length(); + if footer_start < metadata_len { + return Err("invalid footer/metadata length".into()); + } + let metadata_start = footer_start - metadata_len; + Ok(ParquetMetaDataReader::decode_metadata( + &data[metadata_start..footer_start], + )?) +} + +fn parse_optional_path_type(value: &str) -> error::Result> { + match value.to_lowercase().as_str() { + "auto" => Ok(None), + _ => parse_path_type(value).map(Some).map_err(|_| { + error::IllegalConfigSnafu { + msg: format!("invalid path_type '{value}', expected: auto, bare, data, metadata"), + } + .build() + }), + } +} + +fn path_type_name(path_type: PathType) -> &'static str { + match path_type { + PathType::Bare => "bare", + PathType::Data => "data", + PathType::Metadata => "metadata", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_validate_replacement_counts() { + let old_file = mito2::sst::file::FileMeta { + num_rows: 4, + num_row_groups: 1, + ..Default::default() + }; + assert!(validate_replacement(&old_file, 4, 1).is_ok()); + assert!(validate_replacement(&old_file, 3, 1).is_err()); + assert!(validate_replacement(&old_file, 4, 2).is_err()); + + let legacy_file = mito2::sst::file::FileMeta::default(); + assert!(validate_replacement(&legacy_file, 4, 1).is_ok()); + } +} diff --git a/src/cmd/src/datanode/tool_util.rs b/src/cmd/src/datanode/tool_util.rs new file mode 100644 index 00000000000..3b16415e89d --- /dev/null +++ b/src/cmd/src/datanode/tool_util.rs @@ -0,0 +1,236 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#[cfg(feature = "dev-tools")] +use std::fs::File; +use std::path::Path; +use std::sync::Arc; + +use common_wal::config::DatanodeWalConfig; +use datanode::config::RegionEngineConfig; +use datanode::store; +use mito2::config::MitoConfig; +use object_store::ObjectStore; +#[cfg(feature = "dev-tools")] +use parquet::basic::Compression; +use parquet::file::metadata::{KeyValue, ParquetMetaData}; +#[cfg(feature = "dev-tools")] +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; +use snafu::OptionExt; +#[cfg(feature = "dev-tools")] +use snafu::ResultExt; +use store_api::metadata::{RegionMetadata, RegionMetadataRef}; +use store_api::region_request::PathType; +use store_api::storage::{FileId, RegionId}; + +use crate::datanode::{StorageConfig, StorageConfigWrapper}; +use crate::error; + +pub(crate) fn parse_config( + config_path: &Path, +) -> error::Result<(StorageConfig, MitoConfig, DatanodeWalConfig)> { + let cfg_str = std::fs::read_to_string(config_path).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to read config {}: {e}", config_path.display()), + } + .build() + })?; + + let store_cfg: StorageConfigWrapper = toml::from_str(&cfg_str).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to parse config {}: {e}", config_path.display()), + } + .build() + })?; + + let wal_config = store_cfg.wal; + let storage_config = store_cfg.storage; + let mito_engine_config = store_cfg + .region_engine + .into_iter() + .find_map(|config| match config { + RegionEngineConfig::Mito(mito) => Some(mito), + _ => None, + }) + .with_context(|| error::IllegalConfigSnafu { + msg: format!("Engine config not found in {:?}", config_path), + })?; + + Ok((storage_config, mito_engine_config, wal_config)) +} + +pub(crate) async fn build_object_store(config: &StorageConfig) -> error::Result { + store::new_object_store(config.store.clone(), &config.data_home) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("Failed to build object store: {e:?}"), + } + .build() + }) +} + +pub(crate) fn extract_region_metadata( + file_path: &str, + metadata: &ParquetMetaData, +) -> error::Result { + let key_values: Option<&Vec> = metadata.file_metadata().key_value_metadata(); + let Some(key_values) = key_values else { + return Err(error::IllegalConfigSnafu { + msg: format!("{file_path}: missing parquet key_value metadata"), + } + .build()); + }; + let json = key_values + .iter() + .find(|key_value| key_value.key == mito2::sst::parquet::PARQUET_METADATA_KEY) + .and_then(|key_value| key_value.value.as_ref()) + .ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!( + "{file_path}: key {} not found or empty", + mito2::sst::parquet::PARQUET_METADATA_KEY + ), + } + .build() + })?; + let region = RegionMetadata::from_json(json).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid region metadata json: {e}"), + } + .build() + })?; + Ok(Arc::new(region)) +} + +pub(crate) fn parse_region_id(value: &str) -> error::Result { + if let Some((table_id, region_number)) = value.split_once(':') { + let table_id = table_id.parse().map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid table_id in region_id '{value}': {e}"), + } + .build() + })?; + let region_number = region_number.parse().map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid region_num in region_id '{value}': {e}"), + } + .build() + })?; + Ok(RegionId::new(table_id, region_number)) + } else { + value.parse().map(RegionId::from_u64).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid region_id '{value}': {e}"), + } + .build() + }) + } +} + +pub(crate) fn parse_file_id(value: &str) -> error::Result { + FileId::parse_str(value).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid file_id '{value}': {e}"), + } + .build() + }) +} + +pub(crate) fn parse_path_type(value: &str) -> error::Result { + match value.to_lowercase().as_str() { + "bare" => Ok(PathType::Bare), + "data" => Ok(PathType::Data), + "metadata" => Ok(PathType::Metadata), + _ => Err(error::IllegalConfigSnafu { + msg: format!("invalid path_type '{value}', expected: bare, data, metadata"), + } + .build()), + } +} + +pub(crate) fn format_bytes(bytes: u64) -> String { + const KIB: u64 = 1024; + const MIB: u64 = 1024 * KIB; + const GIB: u64 = 1024 * MIB; + if bytes >= GIB { + format!("{:.2} GiB", bytes as f64 / GIB as f64) + } else if bytes >= MIB { + format!("{:.2} MiB", bytes as f64 / MIB as f64) + } else if bytes >= KIB { + format!("{:.2} KiB", bytes as f64 / KIB as f64) + } else { + format!("{bytes} B") + } +} + +pub(crate) fn max_row_group_uncompressed_size(metadata: &ParquetMetaData) -> u64 { + metadata + .row_groups() + .iter() + .map(|row_group| { + row_group + .columns() + .iter() + .map(|column| column.uncompressed_size() as u64) + .sum::() + }) + .max() + .unwrap_or(0) +} + +#[cfg(feature = "dev-tools")] +pub(crate) fn load_local_parquet_metadata(path: &Path) -> error::Result { + let file = File::open(path).context(error::FileIoSnafu)?; + ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .parse_and_finish(&file) + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read parquet metadata failed for {}: {e}", path.display()), + } + .build() + }) +} + +#[cfg(feature = "dev-tools")] +pub(crate) fn compression_name(compression: Compression) -> &'static str { + match compression { + Compression::UNCOMPRESSED => "uncompressed", + Compression::SNAPPY => "snappy", + Compression::GZIP(_) => "gzip", + Compression::LZO => "lzo", + Compression::BROTLI(_) => "brotli", + Compression::LZ4 => "lz4", + Compression::ZSTD(_) => "zstd", + Compression::LZ4_RAW => "lz4-raw", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_region_and_path_type() { + assert_eq!(parse_region_id("1024:7").unwrap(), RegionId::new(1024, 7)); + assert_eq!( + parse_region_id(&RegionId::new(1, 2).as_u64().to_string()).unwrap(), + RegionId::new(1, 2) + ); + assert_eq!(parse_path_type("bare").unwrap(), PathType::Bare); + assert_eq!(parse_path_type("data").unwrap(), PathType::Data); + assert_eq!(parse_path_type("metadata").unwrap(), PathType::Metadata); + } +} From 3493d2d0fbfd7fe05bf9e3276b98bdf727ee9aaf Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 25 Aug 2026 05:16:50 +0000 Subject: [PATCH 098/126] perf(servers)!: speed up Prometheus JSON response building with ryu and per-series entry reuse (#8815) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(perf): align direct-SST CREATE TABLE with baked index metadata The offline fixture generator (query_perf_fixture::direct_sst:: build_region_metadata) bakes greptime:inverted_index / greptime:skipping_index field metadata into the region manifest for tag/field columns, but create_table_sql emitted a bare CREATE TABLE without those declarations. MergeScan's remote-schema validation then failed on any tag/field projection (HTTP 500 'advertised remote stream schema field mismatch'), breaking direct_readable_sst perf cases. CREATE TABLE now declares the matching SKIPPING INDEX WITH (granularity='1') / INVERTED INDEX column options. A round-trip test proves the emitted SQL is parser-valid and yields the exact catalog metadata. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(servers): speed up Prometheus JSON response building with ryu and per-series entry reuse PrometheusJsonResponse::record_batches_to_data spends ~47% of its CPU in f64::to_string() per sample and ~32% in IndexMap::entry() per row (60s profile of concurrent query_range workloads, ~800k series). - Replace f64::to_string() with ryu::Buffer::format_finite for finite values (shortest round-trip, 3-5x faster); NaN/+Inf/-Inf keep the previous std formatting so wire output is unchanged. - Remember the previous row's label vector and entry index; query output is clustered by series, so consecutive rows reuse the same IndexMap entry via get_index_mut instead of rebuilding and hashing the label vector (worst case adds one Vec comparison per series transition). Also adds a query-regression case (prom_json_response) that measures the real Prometheus HTTP range API path (/v1/prometheus/api/v1/query_range), which is the only frontend path that builds the Prometheus JSON response (TQL ANALYZE formats the SQL JSON shape instead), plus a prom_http query kind in the regression runner. Perf (aligned base d90cca4b75, 256 series x 481 points): - prom_range_2h (JSON response path): 29.31ms -> 21.46ms (-26.8%) - tql_range_2h_control (non-JSON path): +1.87% (noise) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(servers): keep Prometheus wire format for integral floats with ryu ryu::Buffer::format_finite prints integral values as "1.0", but the Prometheus JSON wire format (matching std f64::to_string) expects "1". Strip the trailing ".0" that ryu only emits for integral values; extreme values keep ryu scientific notation, and NaN/Inf keep std output. Adds wire-format tests covering 1.0, 0.0, -0.0, 1.5, 0.1, 1e21, 1e30, 1e-7, f64::MAX, NaN, ±Inf. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(servers): address Prometheus response review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(servers)!: use ryu for Prometheus sample values Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(cmd): skip Prometheus execution time extraction Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 + Cargo.toml | 1 + .../src/bin/query_regression_runner/direct.rs | 106 ++++++- .../bin/query_regression_runner/measure.rs | 64 +++- .../src/bin/query_regression_runner/model.rs | 13 + .../src/bin/query_regression_runner/sql.rs | 72 ++++- src/servers/Cargo.toml | 1 + .../src/http/result/prometheus_resp.rs | 295 +++++++++++++++++- tests-integration/tests/grpc.rs | 8 +- tests-integration/tests/http.rs | 6 +- tests/perf/fixture-format.md | 14 + .../query_cases/prom_json_response/case.toml | 89 ++++++ 12 files changed, 644 insertions(+), 26 deletions(-) create mode 100644 tests/perf/query_cases/prom_json_response/case.toml diff --git a/Cargo.lock b/Cargo.lock index a389391ea90..534f484d4fc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -13470,6 +13470,7 @@ dependencies = [ "rustls", "rustls-pemfile", "rustls-pki-types", + "ryu", "serde", "serde_json", "servers", diff --git a/Cargo.toml b/Cargo.toml index e6b849a6cd0..7c3fccfbbf8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -231,6 +231,7 @@ rstest_reuse = "0.7" rust_decimal = "1.33" rustc-hash = "2.0" rustls = { version = "0.23.25", default-features = false } +ryu = "1.0" sea-query = "0.32" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } diff --git a/src/cmd/src/bin/query_regression_runner/direct.rs b/src/cmd/src/bin/query_regression_runner/direct.rs index 04c191c8c55..b003b6ee5d3 100644 --- a/src/cmd/src/bin/query_regression_runner/direct.rs +++ b/src/cmd/src/bin/query_regression_runner/direct.rs @@ -166,7 +166,28 @@ fn create_table_sql(table: &Table) -> Result { let columns = table .columns .iter() - .map(|column| format!("{} {}", sql_ident(&column.name), column.data_type)) + .map(|column| { + // The direct-SST fixture generator (`query_perf_fixture::direct_sst:: + // build_region_metadata`) bakes index metadata into the region + // manifest's column schemas: tag columns get an inverted index and + // a skipping index with granularity 1, field columns get a skipping + // index with granularity 1. The CREATE TABLE must declare the same + // column options, otherwise the table schema (frontend) and the + // region schema (datanode, from the manifest) disagree and the + // MergeScan remote schema validation fails with "advertised remote + // stream schema field mismatch". + let mut sql = format!("{} {}", sql_ident(&column.name), column.data_type); + match column.semantic.as_deref() { + Some("tag") => { + sql.push_str(" SKIPPING INDEX WITH (granularity='1') INVERTED INDEX"); + } + Some("field") => { + sql.push_str(" SKIPPING INDEX WITH (granularity='1')"); + } + _ => {} + } + sql + }) .collect::>() .join(",\n "); let primary_key = table @@ -455,10 +476,12 @@ mod tests { Column { name: "host".to_string(), data_type: "STRING".to_string(), + semantic: Some("tag".to_string()), }, Column { name: "ts".to_string(), data_type: "TIMESTAMP(9)".to_string(), + semantic: Some("timestamp".to_string()), }, ], primary_key: vec!["host".to_string()], @@ -469,7 +492,86 @@ mod tests { }; assert_eq!( create_table_sql(&table).unwrap(), - "CREATE TABLE \"metric\"\"name\" (\n \"host\" STRING,\n \"ts\" TIMESTAMP(9),\n TIME INDEX (\"ts\"),\n PRIMARY KEY (\"host\")\n) ENGINE=mito\nWITH ('append_mode'='true', 'sst_format'='flat');" + "CREATE TABLE \"metric\"\"name\" (\n \"host\" STRING SKIPPING INDEX WITH (granularity='1') INVERTED INDEX,\n \"ts\" TIMESTAMP(9),\n TIME INDEX (\"ts\"),\n PRIMARY KEY (\"host\")\n) ENGINE=mito\nWITH ('append_mode'='true', 'sst_format'='flat');" + ); + } + + #[test] + fn create_table_sql_declares_indexes_matching_fixture_metadata() { + // Mirrors the index metadata baked by + // query_perf_fixture::direct_sst::build_region_metadata: tag columns + // get SKIPPING INDEX WITH (granularity='1') + INVERTED INDEX, field + // columns get SKIPPING INDEX WITH (granularity='1'), timestamp columns + // get no column options. Without these options the frontend table + // schema and the datanode region schema (from the fixture manifest) + // disagree and the MergeScan remote schema validation fails. + let table = Table { + database: "public".to_string(), + name: "metric".to_string(), + engine: "mito".to_string(), + columns: vec![ + Column { + name: "host".to_string(), + data_type: "STRING".to_string(), + semantic: Some("tag".to_string()), + }, + Column { + name: "instance".to_string(), + data_type: "STRING".to_string(), + semantic: Some("tag".to_string()), + }, + Column { + name: "value".to_string(), + data_type: "DOUBLE".to_string(), + semantic: Some("field".to_string()), + }, + Column { + name: "ts".to_string(), + data_type: "TIMESTAMP(9)".to_string(), + semantic: Some("timestamp".to_string()), + }, + ], + primary_key: vec!["host".to_string(), "instance".to_string()], + time_index: Some("ts".to_string()), + append_mode: Some(true), + sst_format: Some("flat".to_string()), + validate_show_create_engine: true, + }; + assert_eq!( + create_table_sql(&table).unwrap(), + "CREATE TABLE \"metric\" (\n \"host\" STRING SKIPPING INDEX WITH (granularity='1') INVERTED INDEX,\n \"instance\" STRING SKIPPING INDEX WITH (granularity='1') INVERTED INDEX,\n \"value\" DOUBLE SKIPPING INDEX WITH (granularity='1'),\n \"ts\" TIMESTAMP(9),\n TIME INDEX (\"ts\"),\n PRIMARY KEY (\"host\", \"instance\")\n) ENGINE=mito\nWITH ('append_mode'='true', 'sst_format'='flat');" + ); + } + + #[test] + fn create_table_sql_without_semantic_keeps_bare_columns() { + // Columns without a semantic type must keep the historical bare form so + // tables that predate the index-metadata fix are unaffected. + let table = Table { + database: "public".to_string(), + name: "metric".to_string(), + engine: "mito".to_string(), + columns: vec![ + Column { + name: "host".to_string(), + data_type: "STRING".to_string(), + semantic: None, + }, + Column { + name: "ts".to_string(), + data_type: "TIMESTAMP(9)".to_string(), + semantic: None, + }, + ], + primary_key: vec!["host".to_string()], + time_index: Some("ts".to_string()), + append_mode: None, + sst_format: None, + validate_show_create_engine: true, + }; + assert_eq!( + create_table_sql(&table).unwrap(), + "CREATE TABLE \"metric\" (\n \"host\" STRING,\n \"ts\" TIMESTAMP(9),\n TIME INDEX (\"ts\"),\n PRIMARY KEY (\"host\")\n) ENGINE=mito;" ); } diff --git a/src/cmd/src/bin/query_regression_runner/measure.rs b/src/cmd/src/bin/query_regression_runner/measure.rs index 2b389f7d782..fd80813f40e 100644 --- a/src/cmd/src/bin/query_regression_runner/measure.rs +++ b/src/cmd/src/bin/query_regression_runner/measure.rs @@ -21,7 +21,7 @@ use serde_json::{Map, Value, json}; use crate::query_regression_runner::model::{Measurement, Query, QueryResult, Scenario, Table}; use crate::query_regression_runner::plan::{load_plan, normalize_scenario}; -use crate::query_regression_runner::sql::{http_post_sql, sql_ident}; +use crate::query_regression_runner::sql::{http_post_prom_range_query, http_post_sql, sql_ident}; use crate::query_regression_runner::{MeasureArgs, Result}; pub(super) async fn run_measure(args: MeasureArgs) -> Result<()> { @@ -106,6 +106,26 @@ fn target_report(name: &str, http_port: u16, result: QueryResult) -> Value { }) } +/// Routes a configured query to the right endpoint: `prom_http` queries hit +/// the Prometheus HTTP range API (which exercises the Prometheus JSON response +/// builder), everything else goes through `/v1/sql` (including `TQL ANALYZE`). +async fn post_query(client: &Client, port: u16, query: &Query, db: &str) -> Value { + if query.kind.as_deref() == Some("prom_http") { + http_post_prom_range_query( + client, + port, + &query.query, + query.start.as_deref(), + query.end.as_deref(), + query.step.as_deref(), + db, + ) + .await + } else { + http_post_sql(client, port, &query.query, db).await + } +} + async fn run_target( port: u16, tables: &[Table], @@ -118,6 +138,9 @@ async fn run_target( name: Some("count_all".to_string()), kind: Some("sql".to_string()), query: format!("SELECT count(*) FROM {}", sql_ident(&tables[0].name)), + start: None, + end: None, + step: None, warmup: 0, iterations: 1, thresholds: Map::new(), @@ -147,7 +170,7 @@ async fn run_target( } validation.push(sample); } - let first = http_post_sql(client, port, &queries[0].query, db).await; + let first = post_query(client, port, &queries[0], db).await; if !first["ok"].as_bool().unwrap_or(false) { validation_errors.push(json!({ "sql": queries[0].query, @@ -160,7 +183,7 @@ async fn run_target( let mut measurements = Vec::with_capacity(queries.len()); for query in &queries { for _ in 0..query.warmup { - let warmup = http_post_sql(client, port, &query.query, db).await; + let warmup = post_query(client, port, query, db).await; if !warmup["ok"].as_bool().unwrap_or(false) { validation_errors.push(json!({ "sql": query.query, @@ -173,10 +196,12 @@ async fn run_target( let mut samples = Vec::with_capacity(query.iterations); let mut good_latencies = Vec::with_capacity(query.iterations); for _ in 0..query.iterations { - let mut sample = http_post_sql(client, port, &query.query, db).await; + let mut sample = post_query(client, port, query, db).await; let execution_time = sample .get("response") - .and_then(extract_execution_time) + .and_then(|response| { + extract_execution_time_for_kind(query.kind.as_deref(), response) + }) .cloned() .unwrap_or(Value::Null); sample @@ -242,6 +267,14 @@ fn response_text(body: &Value) -> String { .unwrap_or_else(|| serde_json::to_string(body).unwrap_or_default()) } +fn extract_execution_time_for_kind<'a>(kind: Option<&str>, body: &'a Value) -> Option<&'a Value> { + if kind == Some("prom_http") { + None + } else { + extract_execution_time(body) + } +} + fn extract_execution_time(body: &Value) -> Option<&Value> { match body { Value::Object(map) => { @@ -379,12 +412,33 @@ mod tests { ); } + #[test] + fn prom_http_does_not_extract_execution_time_from_response() { + let response = json!({ + "data": { + "result": [{"metric": {"job": "api", "elapsed": "label"}}], + "elapsed": 42 + } + }); + assert_eq!( + extract_execution_time_for_kind(Some("prom_http"), &response), + None + ); + assert_eq!( + extract_execution_time_for_kind(Some("sql"), &response), + Some(&json!(42)) + ); + } + #[test] fn threshold_rejects_unknown_keys_and_zero_base() { let query = Query { name: Some("q".to_string()), kind: None, query: "SELECT 1".to_string(), + start: None, + end: None, + step: None, warmup: 0, iterations: 1, thresholds: Map::from_iter([ diff --git a/src/cmd/src/bin/query_regression_runner/model.rs b/src/cmd/src/bin/query_regression_runner/model.rs index 50d41b3f85f..6831a5256a4 100644 --- a/src/cmd/src/bin/query_regression_runner/model.rs +++ b/src/cmd/src/bin/query_regression_runner/model.rs @@ -90,6 +90,12 @@ pub(super) struct Column { pub(super) name: String, #[serde(rename = "type")] pub(super) data_type: String, + /// Semantic type ("tag", "field", "timestamp") from the case TOML. The + /// direct-SST fixture generator bakes index metadata into the region + /// manifest based on this semantic, so the CREATE TABLE must declare the + /// matching column options (see `create_table_sql` in `direct.rs`). + #[serde(default)] + pub(super) semantic: Option, } const fn default_show_create_engine() -> bool { @@ -203,6 +209,13 @@ pub(super) struct Query { #[serde(default)] pub(super) kind: Option, pub(super) query: String, + /// Prometheus HTTP range query parameters (`kind = "prom_http"` only). + #[serde(default)] + pub(super) start: Option, + #[serde(default)] + pub(super) end: Option, + #[serde(default)] + pub(super) step: Option, #[serde(default)] pub(super) warmup: usize, #[serde(default = "one")] diff --git a/src/cmd/src/bin/query_regression_runner/sql.rs b/src/cmd/src/bin/query_regression_runner/sql.rs index 4d8baab2c8c..b33c0a1c22b 100644 --- a/src/cmd/src/bin/query_regression_runner/sql.rs +++ b/src/cmd/src/bin/query_regression_runner/sql.rs @@ -117,10 +117,62 @@ pub(super) fn value_f64(value: Option<&Value>) -> Option { } pub(super) async fn http_post_sql(client: &Client, port: u16, sql: &str, db: &str) -> Value { + let mut sample = post_form( + client, + format!("http://127.0.0.1:{port}/v1/sql"), + &[("sql", sql), ("db", db), ("format", "json")], + ) + .await; + sample + .as_object_mut() + .expect("HTTP samples are objects") + .insert("sql".to_string(), Value::String(sql.to_string())); + sample +} + +/// Posts a Prometheus HTTP API range query (`/v1/prometheus/api/v1/query_range`) +/// and measures the full request-to-body latency. This exercises the Prometheus +/// JSON response building path (`PrometheusJsonResponse::record_batches_to_data`) +/// that `/v1/sql` (including `TQL ANALYZE`) does not go through. +pub(super) async fn http_post_prom_range_query( + client: &Client, + port: u16, + query: &str, + start: Option<&str>, + end: Option<&str>, + step: Option<&str>, + db: &str, +) -> Value { + let mut sample = post_form( + client, + prom_range_query_url(port, db), + &[ + ("query", query), + ("start", start.unwrap_or_default()), + ("end", end.unwrap_or_default()), + ("step", step.unwrap_or_default()), + ], + ) + .await; + sample + .as_object_mut() + .expect("HTTP samples are objects") + .insert("query".to_string(), Value::String(query.to_string())); + sample +} + +fn prom_range_query_url(port: u16, db: &str) -> String { + let mut url = reqwest::Url::parse(&format!( + "http://127.0.0.1:{port}/v1/prometheus/api/v1/query_range" + )) + .expect("fixed Prometheus range query URL must be valid"); + url.query_pairs_mut().append_pair("db", db); + url.into() +} + +async fn post_form(client: &Client, url: String, form: &[(&str, &str)]) -> Value { let started = Instant::now(); - let request = client - .post(format!("http://127.0.0.1:{port}/v1/sql")) - .form(&[("sql", sql), ("db", db), ("format", "json")]); + let request = client.post(url).form(form); match request.send().await { Ok(response) => { let status = response.status().as_u16(); @@ -133,7 +185,6 @@ pub(super) async fn http_post_sql(client: &Client, port: u16, sql: &str, db: &st "status": status, "latency_ms": started.elapsed().as_secs_f64() * 1000.0, "response": body, - "sql": sql, }); if status >= 400 { sample @@ -148,7 +199,6 @@ pub(super) async fn http_post_sql(client: &Client, port: u16, sql: &str, db: &st "status": status, "latency_ms": started.elapsed().as_secs_f64() * 1000.0, "error": error.to_string(), - "sql": sql, }), } } @@ -157,7 +207,6 @@ pub(super) async fn http_post_sql(client: &Client, port: u16, sql: &str, db: &st "status": Value::Null, "latency_ms": started.elapsed().as_secs_f64() * 1000.0, "error": error.to_string(), - "sql": sql, }), } } @@ -208,6 +257,17 @@ pub(super) fn sql_ident(name: &str) -> String { mod tests { use super::*; + #[test] + fn prom_range_query_url_places_database_in_query_parameters() { + let url = reqwest::Url::parse(&prom_range_query_url(4000, "catalog-schema name")) + .expect("generated URL must be valid"); + assert_eq!(url.path(), "/v1/prometheus/api/v1/query_range"); + assert_eq!( + url.query_pairs().collect::>(), + vec![("db".into(), "catalog-schema name".into())] + ); + } + #[test] fn top_level_errors_do_not_inspect_rows() { assert!(response_has_error(&json!({"error_code": 7}))); diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 59d00b90ad7..7936c96fbda 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -115,6 +115,7 @@ rust_decimal = { workspace = true, features = ["db-postgres"] } rustls = { workspace = true, default-features = false, features = ["aws_lc_rs", "logging", "std", "tls12"] } rustls-pemfile = "2.0" rustls-pki-types = "1.0" +ryu.workspace = true serde.workspace = true serde_json.workspace = true session.workspace = true diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index e79628fac0d..bc1d0849037 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -38,6 +38,7 @@ use datatypes::prelude::ConcreteDataType; use indexmap::IndexMap; use promql_parser::label::METRIC_NAME; use promql_parser::parser::value::ValueType; +use ryu::Buffer; use serde::{Deserialize, Serialize}; use serde_json::Value; use snafu::{OptionExt, ResultExt}; @@ -70,6 +71,19 @@ fn prometheus_native_histogram(histogram: &NativeHistogram) -> Result String { + if value.is_finite() { + Buffer::new().format_finite(value).to_string() + } else { + value.to_string() + } +} + #[derive(Debug, Default, Serialize, Deserialize, PartialEq)] pub struct PrometheusJsonResponse { pub status: String, @@ -307,6 +321,15 @@ impl PrometheusJsonResponse { // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. let mut buffer = IndexMap::, PromSeriesSamples>::new(); + // Query output is clustered by series (the range plan sorts by series + // key + timestamp), so consecutive rows usually belong to the same + // series. Remember the index of the previous row's entry in `buffer`, + // and reuse it directly when its tags are unchanged. This avoids + // building and hashing the label vector on every row; the worst case + // adds one `Vec` comparison per series transition before falling back + // to the map lookup. + let mut last_entry_index = None; + let schema = batches.schema(); for batch in batches.iter() { // prepare things... @@ -380,15 +403,32 @@ impl PrometheusJsonResponse { } } - let entry = buffer.entry(tags).or_default(); + let reuse = last_entry_index.filter(|index| { + buffer + .get_index(*index) + .is_some_and(|(key, _)| key == &tags) + }); + let samples = if let Some(index) = reuse { + buffer + .get_index_mut(index) + .map(|(_, samples)| samples) + .with_context(|| UnexpectedResultSnafu { + reason: "reused series entry must exist", + })? + } else { + let entry = buffer.entry(tags); + last_entry_index = Some(entry.index()); + entry.or_default() + }; if let Some((timestamp_millis, histogram)) = histogram { - entry + samples .histograms .push((timestamp_millis as f64 / 1000.0, histogram)); } else if let Some((timestamp_millis, value)) = value { - entry - .values - .push((timestamp_millis as f64 / 1000.0, value.to_string())); + samples.values.push(( + timestamp_millis as f64 / 1000.0, + format_prometheus_sample_value(value), + )); } } } @@ -564,6 +604,67 @@ mod tests { Arc::new(StructVector::try_new(histogram_type, histogram_array).unwrap()) } + #[test] + fn format_prometheus_sample_value_uses_ryu_for_finite_values() { + let values = [ + 1.5, + 0.1, + 1.0, + 0.0, + -0.0, + 100.0, + 1e-6, + 1e-7, + 1e21, + 1e30, + f64::MAX, + f64::MIN_POSITIVE, + ]; + + for value in values { + let output = format_prometheus_sample_value(value); + assert_eq!(output.parse::().unwrap().to_bits(), value.to_bits()); + } + + // Representative integral values use ryu's explicit .0 form. + assert_eq!(format_prometheus_sample_value(1.0), "1.0"); + assert_eq!(format_prometheus_sample_value(-0.0), "-0.0"); + assert_eq!(format_prometheus_sample_value(100.0), "100.0"); + assert_eq!(format_prometheus_sample_value(1e-6), "1e-6"); + assert_eq!(format_prometheus_sample_value(1e-7), "1e-7"); + assert_eq!(format_prometheus_sample_value(1e21), "1e21"); + + // These known shortest-roundtrip tie cases have different text but + // remain numerically equivalent to Rust's representation. + for value in [ + f64::from_bits(0x42374876e8000400), + f64::from_bits(0x3ff0000800000000), + f64::from_bits(0x430a8e5672bc7312), + ] { + let ryu_output = format_prometheus_sample_value(value); + let std_output = value.to_string(); + assert_ne!(ryu_output, std_output); + assert_eq!(ryu_output.parse::().unwrap(), value); + assert_eq!(std_output.parse::().unwrap(), value); + } + } + + #[test] + fn format_prometheus_sample_value_preserves_nonfinite_values() { + assert_eq!(format_prometheus_sample_value(f64::NAN), "NaN"); + assert_eq!(format_prometheus_sample_value(f64::INFINITY), "inf"); + assert_eq!(format_prometheus_sample_value(f64::NEG_INFINITY), "-inf"); + + // Parsing a NaN does not preserve its payload bits, so NaN is checked + // by its required semantic spelling rather than by to_bits(). + assert!( + format_prometheus_sample_value(f64::NAN) + .parse::() + .unwrap() + .is_nan() + ); + } + #[test] fn matrix_response_preserves_ordinary_nan_and_filters_stale_markers() { let schema = Arc::new(Schema::new(vec![ @@ -604,10 +705,190 @@ mod tests { assert_eq!(series.len(), 1); assert_eq!( series[0].values, - vec![(1.0, "1".to_string()), (2.0, "NaN".to_string())] + vec![(1.0, "1.0".to_string()), (2.0, "NaN".to_string())] ); } + #[test] + fn record_batches_to_data_formats_values_with_ryu() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ])); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![ + 1_000, 2_000, 3_000, 4_000, 5_000, 6_000, + ])) as _, + Arc::new(Float64Vector::from(vec![ + Some(0.0), + Some(-0.0), + Some(1.25), + Some(1e30), + Some(1e-7), + Some(f64::MAX), + ])) as _, + ], + ) + .unwrap(); + let batches = RecordBatches::try_new(schema, vec![batch]).unwrap(); + + let response = + PrometheusJsonResponse::record_batches_to_data(batches, None, ValueType::Matrix) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = response + else { + panic!("expected matrix response"); + }; + + assert_eq!(series.len(), 1); + let input_values = [0.0, -0.0, 1.25, 1e30, 1e-7, f64::MAX]; + // Keep this expected-value generation independent from the production + // formatter while still asserting the Arrow/batch-to-Prometheus path. + let expected = input_values + .into_iter() + .enumerate() + .map(|(index, value)| { + let expected_value = if value.is_finite() { + let mut buffer = Buffer::new(); + buffer.format_finite(value).to_string() + } else { + value.to_string() + }; + ((index + 1) as f64, expected_value) + }) + .collect::>(); + assert_eq!(series[0].values, expected); + } + + #[test] + fn record_batches_to_data_preserves_infinity_output() { + // NaN and infinities use Rust's `f64::to_string()` output. + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ])); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![ + 1_000, 2_000, 3_000, + ])) as _, + Arc::new(Float64Vector::from(vec![ + Some(f64::INFINITY), + Some(f64::NEG_INFINITY), + Some(f64::NAN), + ])) as _, + ], + ) + .unwrap(); + let batches = RecordBatches::try_new(schema, vec![batch]).unwrap(); + + let response = + PrometheusJsonResponse::record_batches_to_data(batches, None, ValueType::Matrix) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = response + else { + panic!("expected matrix response"); + }; + + assert_eq!(series.len(), 1); + assert_eq!( + series[0].values, + vec![ + (1.0, "inf".to_string()), + (2.0, "-inf".to_string()), + (3.0, "NaN".to_string()), + ] + ); + } + + #[test] + fn record_batches_to_data_reuses_entries_for_clustered_series() { + // Rows are clustered by series (a, b, a, a, b, c): consecutive rows of + // the same series exercise the entry-reuse fast path, while series + // transitions fall back to the map lookup. The result must keep the + // first-occurrence order and accumulate values per series as before. + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ])); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![ + 1_000, 2_000, 3_000, 4_000, 5_000, 6_000, + ])) as _, + Arc::new(StringVector::from(vec![ + Some("a"), + Some("b"), + Some("a"), + Some("a"), + Some("b"), + Some("c"), + ])) as _, + Arc::new(Float64Vector::from(vec![ + Some(1.0), + Some(2.0), + Some(3.0), + Some(4.0), + Some(5.0), + Some(6.0), + ])) as _, + ], + ) + .unwrap(); + let batches = RecordBatches::try_new(schema, vec![batch]).unwrap(); + + let response = PrometheusJsonResponse::record_batches_to_data( + batches, + Some("metric".to_string()), + ValueType::Vector, + ) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Vector(series), + .. + }) = response + else { + panic!("expected vector response"); + }; + + assert_eq!(series.len(), 3); + // Output order is first-occurrence order: a, b, c. + assert_eq!( + series + .iter() + .map(|series| series.metric["host"].as_str()) + .collect::>(), + vec!["a", "b", "c"] + ); + // Vector results keep the last sample of each series. + assert_eq!(series[0].value, Some((4.0, "4.0".to_string()))); + assert_eq!(series[1].value, Some((5.0, "5.0".to_string()))); + assert_eq!(series[2].value, Some((6.0, "6.0".to_string()))); + } + #[test] fn record_batches_to_data_preserves_mixed_float_and_histogram_rows() { let schema = Arc::new(Schema::new(vec![ @@ -713,7 +994,7 @@ mod tests { assert_eq!(series[0].metric["__name__"], "label_replace_repro"); assert_eq!(series[0].metric["host"], "server-01"); assert_eq!(series[0].metric["host_copy"], "server-01"); - assert_eq!(series[0].value, Some((1.0, "1".to_string()))); + assert_eq!(series[0].value, Some((1.0, "1.0".to_string()))); } #[test] diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 9d0e13dddc9..737b4fb6c1d 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -1071,7 +1071,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { ] .into_iter() .collect(), - value: Some((5.0, "1".to_string())), + value: Some((5.0, "1.0".to_string())), ..Default::default() }, PromSeriesVector { @@ -1081,7 +1081,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { ] .into_iter() .collect(), - value: Some((5.0, "2".to_string())), + value: Some((5.0, "2.0".to_string())), ..Default::default() }, ] @@ -1133,7 +1133,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { ] .into_iter() .collect(), - values: vec![(5.0, "1".to_string()), (10.0, "1".to_string())], + values: vec![(5.0, "1.0".to_string()), (10.0, "1.0".to_string())], ..Default::default() }, PromSeriesMatrix { @@ -1143,7 +1143,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { ] .into_iter() .collect(), - values: vec![(5.0, "2".to_string()), (10.0, "2".to_string())], + values: vec![(5.0, "2.0".to_string()), (10.0, "2.0".to_string())], ..Default::default() }, ] diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 2a66c4353e7..7a787833376 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -935,7 +935,7 @@ pub async fn test_prom_http_api(store_type: StorageType) { assert_eq!( body.data, serde_json::from_value::( - json!({"resultType":"scalar","result":[1.0,"2"]}) + json!({"resultType":"scalar","result":[1.0,"2.0"]}) ) .unwrap() ); @@ -2967,6 +2967,8 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag .await; assert_eq!(res.status(), StatusCode::OK); let body = res.json::().await; + // Finite vector sample values use the HTTP response's ryu contract, so + // integral f64 values retain their explicit `.0`; timestamps remain JSON numbers. assert_eq!( serde_json::to_string_pretty(&body).unwrap(), r#"{ @@ -2982,7 +2984,7 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag }, "value": [ 4.0, - "6" + "6.0" ] } ] diff --git a/tests/perf/fixture-format.md b/tests/perf/fixture-format.md index ca60ae54d72..361f0ec79fd 100644 --- a/tests/perf/fixture-format.md +++ b/tests/perf/fixture-format.md @@ -58,6 +58,20 @@ warmup = 0 iterations = 1 ``` +### Query kinds + +`kind = "prom_http"` runs a Prometheus range query by POSTing form fields to +`/v1/prometheus/api/v1/query_range`. `query`, `start`, `end`, and `step` are +sent as form fields; `start` and `end` define the range, and `step` defines its +evaluation interval. `start`, `end`, and `step` are optional in the case model +and default to empty form values when omitted. The table database is sent as the +`db` URL query parameter, rather than as a form field. + +The measurement starts before the HTTP request is sent and ends after the full +response body is read and parsed as JSON. It therefore includes request send, +server-side Prometheus JSON response building, and response-body read and JSON +parsing; it is not server-side-only latency. + `series_layout = "round_robin"` advances the timestamp once per generated row and cycles series labels across rows. `series_layout = "timestamp_major"` writes all series for one timestamp before advancing to the next timestamp; use it for diff --git a/tests/perf/query_cases/prom_json_response/case.toml b/tests/perf/query_cases/prom_json_response/case.toml new file mode 100644 index 00000000000..17472cd5f32 --- /dev/null +++ b/tests/perf/query_cases/prom_json_response/case.toml @@ -0,0 +1,89 @@ +# Prometheus JSON response building benchmark (issue #8805). +# +# Exercises the serial per-request +# `PrometheusJsonResponse::record_batches_to_data` +# (src/servers/src/http/result/prometheus_resp.rs) path, ryu float formatting, +# and its per-series IndexMap entry reuse. +# +# 256 series x 481 evaluation timestamps (2h at 15s) = ~123k output samples +# per request. The primary `prom_http` query builds the Prometheus JSON +# response. The TQL ANALYZE control runs the same workload without building the +# Prometheus JSON response, to measure query-engine/scan variance. + +[case] +name = "prom_json_response" +description = "Prometheus HTTP API JSON response building for large range queries" + +[scenario] +kind = "direct_readable_sst" +seed = 8805 + +[[scenario.tables]] +database = "public" +name = "prom_json_response" +engine = "mito" +append_mode = true +sst_format = "flat" +primary_key = ["host", "instance"] +time_index = "ts" + +[[scenario.tables.columns]] +name = "host" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 16, prefix = "host" } + +[[scenario.tables.columns]] +name = "instance" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 256, prefix = "instance" } + +[[scenario.tables.columns]] +name = "value" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 1000.0 } + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(9)" +semantic = "timestamp" + +[scenario.layout] +regions = 1 +sst_count = 16 +rows_per_sst = 7680 +row_group_size = 1920 +series_count = 256 +start_unix_nanos = 1704067200000000000 +step_nanos = 15000000000 +time_range_layout = "non_overlapping_per_sst" +series_layout = "timestamp_major" + +# Primary case: Prometheus HTTP range query over the full 2h window. +# 256 series x 481 points -> ~123k samples built and serialized per request. +[[scenario.queries]] +name = "prom_range_2h" +kind = "prom_http" +query = "prom_json_response{host=~\"host.*\"}" +start = "1704067200" +end = "1704074400" +step = "15s" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +# Control: same workload through the SQL/TQL frontend path, which does not +# build the Prometheus JSON response. Isolates query-engine/scan variance. +[[scenario.queries]] +name = "tql_range_2h_control" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704074400, '15s') prom_json_response{host=~'host.*'}" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 From 04614175feff271fdf3608462c84e4ecef36b601 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Tue, 25 Aug 2026 06:26:09 +0000 Subject: [PATCH 099/126] refactor: centralize native histogram encoding in common-query (#8945) * refactor: centralize native histogram encoding in common-query Signed-off-by: shuiyisong * fix: fmt Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/common/query/src/native_histogram.rs | 3 +++ .../query/src/native_histogram/encoding.rs} | 18 +++++++++++------- src/servers/src/lib.rs | 1 - src/servers/src/otlp/metrics.rs | 5 +++-- src/servers/src/prom_remote_write/v2.rs | 5 +++-- 5 files changed, 20 insertions(+), 12 deletions(-) rename src/{servers/src/native_histogram.rs => common/query/src/native_histogram/encoding.rs} (96%) diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index 1c7bb09868e..d8a559a2dc5 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -19,6 +19,8 @@ //! [`NativeHistogram`] is the query-time representation and therefore normalizes //! integer and floating-point payloads to absolute `f64` counts. +mod encoding; + use std::collections::BTreeMap; use std::sync::Arc; @@ -34,6 +36,7 @@ use datafusion::arrow::datatypes::{ use datafusion_common::{DataFusionError, Result as DfResult}; use datatypes::data_type::{ConcreteDataType, DataType}; use datatypes::types::{StructField, StructType}; +pub use encoding::{NativeHistogramError, encode_native_histogram, native_histogram_column_schema}; use once_cell::sync::Lazy; use crate::prelude::greptime_native_histogram; diff --git a/src/servers/src/native_histogram.rs b/src/common/query/src/native_histogram/encoding.rs similarity index 96% rename from src/servers/src/native_histogram.rs rename to src/common/query/src/native_histogram/encoding.rs index f1dae526993..be7b35ec454 100644 --- a/src/servers/src/native_histogram.rs +++ b/src/common/query/src/native_histogram/encoding.rs @@ -17,23 +17,27 @@ use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; use api::helper::ColumnDataTypeWrapper; use api::v1::value::ValueData; use api::v1::{ColumnSchema, ListValue, SemanticType, Value}; -use common_query::native_histogram::*; -use common_query::prelude::greptime_native_histogram; use snafu::{Snafu, ensure}; -const MAX_NATIVE_HISTOGRAM_SCHEMA: i32 = 8; +use crate::native_histogram::{ + CUSTOM_BUCKETS_SCHEMA, MAX_EXPONENTIAL_SCHEMA, NATIVE_HISTOGRAM_FIELD_NAMES, + exponential_overflow_bucket_index, native_histogram_value_type, +}; +use crate::prelude::greptime_native_histogram; + const MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA: i32 = 52; +/// Error returned while validating or encoding a native histogram. #[derive(Debug, Snafu)] #[snafu(display("{message}"))] -pub(crate) struct NativeHistogramError { +pub struct NativeHistogramError { message: String, } type Result = std::result::Result; /// Returns the canonical column schema for a native histogram value. -pub(crate) fn native_histogram_column_schema() -> Result { +pub fn native_histogram_column_schema() -> Result { let (datatype, datatype_extension) = ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone()) .map_err(|error| NativeHistogramError { @@ -51,7 +55,7 @@ pub(crate) fn native_histogram_column_schema() -> Result { } /// Validates and encodes a Prometheus histogram into the canonical Struct value. -pub(crate) fn encode_native_histogram(histogram: &Histogram) -> Result { +pub fn encode_native_histogram(histogram: &Histogram) -> Result { let uses_float_counts = native_histogram_uses_float_counts(histogram)?; validate_native_histogram(histogram, uses_float_counts)?; @@ -193,7 +197,7 @@ fn validate_native_histogram_schema(schema: i32) -> Result> { return Ok(Some(overflow_index)); } - if (MAX_NATIVE_HISTOGRAM_SCHEMA + 1..=MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA).contains(&schema) { + if (MAX_EXPONENTIAL_SCHEMA + 1..=MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA).contains(&schema) { Err(NativeHistogramError { message: format!("native histogram schema {schema} must be reduced before ingestion"), }) diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index 79fd392a1a9..2a3796aab21 100644 --- a/src/servers/src/lib.rs +++ b/src/servers/src/lib.rs @@ -34,7 +34,6 @@ pub mod interceptor; pub mod metrics; pub mod metrics_handler; pub mod mysql; -pub(crate) mod native_histogram; pub mod opentsdb; pub mod otel_arrow; pub mod otlp; diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 332fa46f1d1..bd0bfb2ce77 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -20,7 +20,9 @@ use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram as Pr use api::v1::value::ValueData; use api::v1::{RowInsertRequests, SemanticType, Value}; use common_grpc::precision::Precision; -use common_query::native_histogram::native_histogram_value_type; +use common_query::native_histogram::{ + encode_native_histogram, native_histogram_column_schema, native_histogram_value_type, +}; use common_query::prelude::{GREPTIME_COUNT, greptime_timestamp, greptime_value}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use lazy_static::lazy_static; @@ -34,7 +36,6 @@ use table::requests::{ }; use crate::error::{self, Result}; -use crate::native_histogram::{encode_native_histogram, native_histogram_column_schema}; use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME}; use crate::query_handler::MetricsIngestOutcome; use crate::row_writer::{self, MultiTableData, TableData}; diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 846ef3ee6be..b388071106f 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -30,9 +30,11 @@ use api::v1::value::ValueData; use api::v1::{ColumnDataType, RowInsertRequest, Rows, SemanticType, Value}; use bytes::{Buf, Bytes}; use common_grpc::precision::Precision; -use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; #[cfg(test)] use common_query::native_histogram::*; +use common_query::native_histogram::{ + NATIVE_HISTOGRAM_FIELD, encode_native_histogram, native_histogram_column_schema, +}; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use pipeline::{ContextOpt, ContextReq}; use prost::encoding::{ @@ -46,7 +48,6 @@ use table::requests::{ }; use crate::error::{self, Result}; -use crate::native_histogram::{encode_native_histogram, native_histogram_column_schema}; use crate::prom_remote_write::row_builder::PromCtx; use crate::prom_remote_write::validation::validate_label_name; use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, try_decompress}; From a3a0db63b80a5ef00f52610c070f4b8d9707cd23 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Tue, 25 Aug 2026 07:01:10 +0000 Subject: [PATCH 100/126] feat(mito2): add series index searcher (#8926) * feat(mito2): add series index searcher Signed-off-by: evenyag * refactor(mito2): use parquet push decoder for series index Signed-off-by: evenyag * fix(mito2): handle evolved series index schemas Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/AGENTS.md | 2 +- src/mito2/src/read/series_candidate.rs | 19 +- src/mito2/src/read/series_reader.rs | 3 +- src/mito2/src/series_index.rs | 26 +- src/mito2/src/series_index/searcher.rs | 734 +++++++++++++++++++++++++ src/mito2/src/series_index/writer.rs | 8 +- src/mito2/src/sst/parquet/format.rs | 50 +- src/mito2/src/sst/parquet/prefilter.rs | 38 +- 8 files changed, 832 insertions(+), 48 deletions(-) create mode 100644 src/mito2/src/series_index/searcher.rs diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index e68ca586096..20ca9b53ca0 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -26,7 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | -| `series_index` | `src/mito2/src/series_index/` | Incremental writer for aggregate series index files | +| `series_index` | `src/mito2/src/series_index/` | Incremental writer and predicate searcher for aggregate series index files | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | | `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits | | `cache` | `src/mito2/src/cache.rs` | Write/file/page caches | diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index 2933327a38e..73591d4ffa3 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -30,7 +30,6 @@ use datatypes::arrow::compute::SortOptions; use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; -use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; use mito_codec::row_converter::{PrimaryKeyFilter, SparsePrimaryKeyCodec}; use snafu::{OptionExt, ResultExt, ensure}; @@ -51,6 +50,7 @@ use crate::read::range_cache::{ }; use crate::read::scan_region::StreamContext; use crate::read::scan_util::{PartitionMetrics, new_filter_metrics, scan_flat_mem_ranges}; +use crate::series_index::{METRIC_SERIES_ID_BATCH_SIZE, MetricSeriesId, MetricSeriesIdStream}; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::format::PrimaryKeyArray; use crate::sst::parquet::prefilter::{ @@ -59,17 +59,6 @@ use crate::sst::parquet::prefilter::{ use crate::sst::parquet::reader::ReaderMetrics; use crate::sst::parquet::row_group::ParquetFetchMetrics; -const CANDIDATE_SERIES_BATCH_SIZE: usize = 500; - -/// Identifies one series in a physical metric region. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub(crate) struct MetricSeriesId { - pub(crate) table_id: u32, - pub(crate) tsid: u64, -} - -pub(crate) type MetricSeriesIdStream = BoxStream<'static, Result>>; - /// Builds candidate metric series from the ranges assigned to a [`SeriesScan`](super::series_scan::SeriesScan). pub(crate) struct SeriesCandidateScanner { stream_ctx: Arc, @@ -537,7 +526,7 @@ fn decode_metric_series( let codec = SparsePrimaryKeyCodec::new(&metadata); Ok(Box::pin(try_stream! { let mut last_series = None; - let mut output = Vec::with_capacity(CANDIDATE_SERIES_BATCH_SIZE); + let mut output = Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE); while let Some(batch) = input.try_next().await? { let array = batch .column(0) @@ -556,10 +545,10 @@ fn decode_metric_series( } last_series = Some(series); output.push(series); - if output.len() == CANDIDATE_SERIES_BATCH_SIZE { + if output.len() == METRIC_SERIES_ID_BATCH_SIZE { yield std::mem::replace( &mut output, - Vec::with_capacity(CANDIDATE_SERIES_BATCH_SIZE), + Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE), ); } } diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index 0a3e008f27a..15e47924323 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -40,7 +40,8 @@ use crate::read::scan_util::{ should_split_flat_batches_for_merge, }; use crate::read::seq_scan::SeqScan; -use crate::read::series_candidate::{MetricSeriesId, validate_metric_metadata}; +use crate::read::series_candidate::validate_metric_metadata; +use crate::series_index::MetricSeriesId; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::prefilter::prefilter_flat_batch_by_primary_key; diff --git a/src/mito2/src/series_index.rs b/src/mito2/src/series_index.rs index 21ced00cf64..e35769bec4c 100644 --- a/src/mito2/src/series_index.rs +++ b/src/mito2/src/series_index.rs @@ -12,10 +12,34 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Series index writer. +//! Series index writer and searcher. +mod searcher; mod writer; +use futures::stream::BoxStream; +pub use searcher::SeriesIndexSearcher; pub use writer::{ SeriesIndexWriter, SeriesIndexWriterMetrics, SeriesIndexWriterOptions, series_index_schema, }; + +use crate::error::Result; + +pub(crate) const MIN_TS_COLUMN: &str = "__series_min_ts"; +pub(crate) const MAX_TS_COLUMN: &str = "__series_max_ts"; +pub(crate) const ROW_COUNT_COLUMN: &str = "__series_row_count"; +pub(crate) const TABLE_ID_COLUMN: &str = "__table_id"; +pub(crate) const TSID_COLUMN: &str = "__tsid"; +pub(crate) const METRIC_SERIES_ID_BATCH_SIZE: usize = 500; + +/// Identifies one series in a physical metric region. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct MetricSeriesId { + /// Logical table ID inside the physical metric region. + pub table_id: u32, + /// Time-series ID inside the logical table. + pub tsid: u64, +} + +/// Stream of bounded batches of matching metric-series IDs. +pub type MetricSeriesIdStream = BoxStream<'static, Result>>; diff --git a/src/mito2/src/series_index/searcher.rs b/src/mito2/src/series_index/searcher.rs new file mode 100644 index 00000000000..a303bcaebda --- /dev/null +++ b/src/mito2/src/series_index/searcher.rs @@ -0,0 +1,734 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashSet; +use std::ops::Range; +use std::sync::Arc; + +use api::v1::SemanticType; +use async_stream::try_stream; +use bytes::Bytes; +use common_recordbatch::filter::SimpleFilterEvaluator; +use common_time::range::TimestampRange; +use datafusion_common::pruning::PruningStatistics; +use datafusion_common::{Column, ScalarValue}; +use datafusion_expr::{Expr, col, lit}; +use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt32Array, UInt64Array}; +use datatypes::arrow::buffer::BooleanBuffer; +use datatypes::arrow::datatypes::{DataType, SchemaRef}; +use object_store::ObjectStore; +use parquet::DecodeResult; +use parquet::arrow::ProjectionMask; +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; +use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metadata::RegionMetadataRef; +use table::predicate::Predicate; + +use crate::error::{ + InvalidMetaSnafu, InvalidRecordBatchSnafu, OpenDalSnafu, ReadParquetSnafu, RecordBatchSnafu, + Result, UnexpectedSnafu, +}; +use crate::series_index::{ + MAX_TS_COLUMN, METRIC_SERIES_ID_BATCH_SIZE, MIN_TS_COLUMN, MetricSeriesId, + MetricSeriesIdStream, ROW_COUNT_COLUMN, TABLE_ID_COLUMN, TSID_COLUMN, series_index_schema, +}; +use crate::sst::parquet::format::{column_null_counts, column_values_by_type}; +use crate::sst::parquet::helper::fetch_byte_ranges; +use crate::sst::parquet::metadata::MetadataLoader; +use crate::sst::parquet::prefilter::simple_tag_filters; +use crate::sst::parquet::reader::MetadataCacheMetrics; + +#[derive(Clone)] +struct SeriesIndexRangeFetcher { + object_store: ObjectStore, +} + +impl SeriesIndexRangeFetcher { + async fn fetch(&self, path: &str, ranges: &[Range]) -> Result> { + fetch_byte_ranges(path, self.object_store.clone(), ranges) + .await + .context(OpenDalSnafu) + } +} + +#[derive(Clone)] +struct SeriesIndexMetadataProvider { + object_store: ObjectStore, +} + +impl SeriesIndexMetadataProvider { + async fn load(&self, path: &str) -> Result> { + let mut metrics = MetadataCacheMetrics::default(); + MetadataLoader::new(self.object_store.clone(), path, 0) + .load(&mut metrics) + .await + .map(Arc::new) + } +} + +/// Searches a series-index file for metric series matching query predicates. +pub struct SeriesIndexSearcher { + range_fetcher: SeriesIndexRangeFetcher, + metadata_provider: SeriesIndexMetadataProvider, + filters: Vec<(Expr, SimpleFilterEvaluator)>, + empty_time_range: bool, +} + +impl SeriesIndexSearcher { + /// Creates a searcher reusable across series-index files of `metadata`. + pub fn try_new( + metadata: RegionMetadataRef, + object_store: ObjectStore, + predicate: Option<&Predicate>, + time_range: Option, + ) -> Result { + // Keep search-time metadata validation identical to the writer. + series_index_schema(&metadata)?; + + let mut filters = simple_tag_filters(&metadata, None, predicate); + let (empty_time_range, time_exprs) = time_range_filters(&metadata, time_range)?; + for expr in time_exprs { + let filter = SimpleFilterEvaluator::try_new(&expr).context(UnexpectedSnafu { + reason: "failed to build an internal series-index time filter", + })?; + filters.push((expr, filter)); + } + + Ok(Self { + range_fetcher: SeriesIndexRangeFetcher { + object_store: object_store.clone(), + }, + metadata_provider: SeriesIndexMetadataProvider { object_store }, + filters, + empty_time_range, + }) + } + + /// Searches `path` and returns sorted batches of matching metric-series IDs. + pub async fn search(&self, path: &str) -> Result { + if self.empty_time_range { + return Ok(Box::pin(futures::stream::empty())); + } + + let parquet_metadata = self.metadata_provider.load(path).await?; + let arrow_metadata = + ArrowReaderMetadata::try_new(parquet_metadata, ArrowReaderOptions::new()) + .with_context(|_| ReadParquetSnafu { + path: path.to_string(), + })?; + validate_index_schema(arrow_metadata.schema())?; + + // An older index file may not contain tags added by schema evolution. + // Ignore filters on those tags to preserve a conservative candidate set. + let (pruning_predicate, filters) = self.filters_for_schema(arrow_metadata.schema()); + let row_groups = row_groups_to_read( + arrow_metadata.metadata().row_groups(), + arrow_metadata.schema().clone(), + &pruning_predicate, + ); + let projection = projection_mask( + arrow_metadata.parquet_schema(), + arrow_metadata.schema(), + &filters, + )?; + let mut decoder = ParquetPushDecoderBuilder::new_with_metadata(arrow_metadata) + .with_row_groups(row_groups) + .with_projection(projection) + .build() + .with_context(|_| ReadParquetSnafu { + path: path.to_string(), + })?; + let path = path.to_string(); + let range_fetcher = self.range_fetcher.clone(); + + Ok(Box::pin(try_stream! { + let mut last_series = None; + let mut output = Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE); + loop { + let batch = match decoder + .try_decode() + .with_context(|_| ReadParquetSnafu { path: path.clone() })? + { + DecodeResult::NeedsData(ranges) => { + let data = range_fetcher.fetch(&path, &ranges).await?; + decoder + .push_ranges(ranges, data) + .with_context(|_| ReadParquetSnafu { path: path.clone() })?; + continue; + } + DecodeResult::Data(batch) => batch, + DecodeResult::Finished => break, + }; + + let mut mask = BooleanBuffer::new_set(batch.num_rows()); + for filter in &filters { + let column = column(&batch, filter.column_name())?; + let evaluated = filter.evaluate_array(column).context(RecordBatchSnafu)?; + mask = &mask & &evaluated; + } + + let table_ids = column(&batch, TABLE_ID_COLUMN)? + .as_any() + .downcast_ref::() + .context(InvalidRecordBatchSnafu { + reason: "series index __table_id is not UInt32", + })?; + let tsids = column(&batch, TSID_COLUMN)? + .as_any() + .downcast_ref::() + .context(InvalidRecordBatchSnafu { + reason: "series index __tsid is not UInt64", + })?; + + for (row, matched) in mask.iter().enumerate() { + if !matched { + continue; + } + let series = MetricSeriesId { + table_id: table_ids.value(row), + tsid: tsids.value(row), + }; + if last_series == Some(series) { + continue; + } + last_series = Some(series); + output.push(series); + if output.len() == METRIC_SERIES_ID_BATCH_SIZE { + yield std::mem::replace( + &mut output, + Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE), + ); + } + } + } + if !output.is_empty() { + yield output; + } + })) + } + + fn filters_for_schema(&self, schema: &SchemaRef) -> (Predicate, Vec) { + let (exprs, filters): (Vec<_>, Vec<_>) = self + .filters + .iter() + .filter(|(_, filter)| schema.field_with_name(filter.column_name()).is_ok()) + .cloned() + .unzip(); + (Predicate::new(exprs), filters) + } +} + +fn time_range_filters( + metadata: &RegionMetadataRef, + time_range: Option, +) -> Result<(bool, Vec)> { + let Some(time_range) = time_range else { + return Ok((false, Vec::new())); + }; + if time_range.is_empty() { + return Ok((true, Vec::new())); + } + + let time_index = metadata.time_index_column(); + ensure!( + time_index.semantic_type == SemanticType::Timestamp, + InvalidMetaSnafu { + reason: "series index metadata has no timestamp time index", + } + ); + let timestamp_type = + time_index + .column_schema + .data_type + .as_timestamp() + .context(InvalidMetaSnafu { + reason: "series index time index is not a timestamp", + })?; + let unit = timestamp_type.unit(); + let mut exprs = Vec::with_capacity(2); + // A series overlaps [start, end) only if its maximum is at least start. + // Round start up so a series ending before an unaligned start is pruned. + if let Some(start) = time_range + .start() + .and_then(|start| start.convert_to_ceil(unit)) + { + exprs.push(col(MAX_TS_COLUMN).gt_eq(lit(start.value()))); + } + // A series overlaps [start, end) only if its minimum is less than end. + // Round the exclusive end up to avoid pruning the containing unit interval. + if let Some(end) = time_range.end().and_then(|end| end.convert_to_ceil(unit)) { + exprs.push(col(MIN_TS_COLUMN).lt(lit(end.value()))); + } + Ok((false, exprs)) +} + +fn validate_index_schema(schema: &SchemaRef) -> Result<()> { + for (name, data_type) in [ + (MIN_TS_COLUMN, DataType::Int64), + (MAX_TS_COLUMN, DataType::Int64), + (ROW_COUNT_COLUMN, DataType::UInt64), + (TABLE_ID_COLUMN, DataType::UInt32), + (TSID_COLUMN, DataType::UInt64), + ] { + let field = schema + .field_with_name(name) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("series index is missing internal column {name}"), + })?; + ensure!( + field.data_type() == &data_type && !field.is_nullable(), + InvalidRecordBatchSnafu { + reason: format!( + "series index internal column {name} must be non-nullable {data_type:?}, got {:?}", + field.data_type() + ), + } + ); + } + Ok(()) +} + +fn projection_mask( + parquet_schema: &parquet::schema::types::SchemaDescriptor, + arrow_schema: &SchemaRef, + filters: &[SimpleFilterEvaluator], +) -> Result { + let mut indices = HashSet::new(); + for name in [TABLE_ID_COLUMN, TSID_COLUMN] { + let index = arrow_schema + .index_of(name) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("series index is missing internal column {name}"), + })?; + indices.insert(index); + } + for filter in filters { + let index = arrow_schema + .index_of(filter.column_name()) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "series index is missing predicate column {}", + filter.column_name() + ), + })?; + indices.insert(index); + } + Ok(ProjectionMask::roots(parquet_schema, indices)) +} + +fn column<'a>( + batch: &'a datatypes::arrow::record_batch::RecordBatch, + name: &str, +) -> Result<&'a ArrayRef> { + let index = batch + .schema() + .index_of(name) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("series index batch is missing column {name}"), + })?; + Ok(batch.column(index)) +} + +struct SeriesIndexPruningStats<'a> { + row_groups: &'a [RowGroupMetaData], + schema: SchemaRef, +} + +impl PruningStatistics for SeriesIndexPruningStats<'_> { + fn min_values(&self, column: &Column) -> Option { + self.column_values(column, true) + } + + fn max_values(&self, column: &Column) -> Option { + self.column_values(column, false) + } + + fn num_containers(&self) -> usize { + self.row_groups.len() + } + + fn null_counts(&self, column: &Column) -> Option { + let column_index = self.schema.index_of(&column.name).ok()?; + column_null_counts(self.row_groups, column_index) + } + + fn row_counts(&self, _column: &Column) -> Option { + None + } + + fn contained(&self, _column: &Column, _values: &HashSet) -> Option { + None + } +} + +impl SeriesIndexPruningStats<'_> { + fn column_values(&self, column: &Column, is_min: bool) -> Option { + let column_index = self.schema.index_of(&column.name).ok()?; + let data_type = self.schema.field(column_index).data_type(); + column_values_by_type(self.row_groups, data_type, column_index, is_min) + } +} + +fn row_groups_to_read( + row_groups: &[RowGroupMetaData], + schema: SchemaRef, + predicate: &Predicate, +) -> Vec { + let stats = SeriesIndexPruningStats { row_groups, schema }; + predicate + .prune_with_stats(&stats, &stats.schema) + .into_iter() + .enumerate() + .filter_map(|(row_group, keep)| keep.then_some(row_group)) + .collect() +} + +#[cfg(test)] +mod tests { + use datafusion_expr::{col, lit}; + use datatypes::arrow::array::{BinaryArray, TimestampMillisecondArray, UInt8Array}; + use datatypes::arrow::datatypes::{Field, Schema}; + use datatypes::arrow::record_batch::RecordBatch; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::ColumnSchema; + use futures::TryStreamExt; + use object_store::services::Memory; + use store_api::codec::PrimaryKeyEncoding; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + + use super::*; + use crate::series_index::{SeriesIndexWriter, SeriesIndexWriterOptions}; + use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; + + fn object_store() -> ObjectStore { + ObjectStore::new(Memory::default()).unwrap().finish() + } + + fn flat_batch(primary_keys: &[Vec], timestamps: &[i64]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("__primary_key", DataType::Binary, false), + Field::new("__sequence", DataType::UInt64, false), + Field::new("__op_type", DataType::UInt8, false), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())), + Arc::new(BinaryArray::from_iter_values( + primary_keys.iter().map(Vec::as_slice), + )), + Arc::new(UInt64Array::from(vec![1; timestamps.len()])), + Arc::new(UInt8Array::from(vec![0; timestamps.len()])), + ], + ) + .unwrap() + } + + async fn write_index( + metadata: RegionMetadataRef, + object_store: ObjectStore, + path: &str, + rows: &[(u32, u64, &str, &str, i64)], + row_group_size: usize, + ) { + let primary_keys = rows + .iter() + .map(|(table_id, tsid, tag_0, tag_1, _)| { + new_sparse_primary_key(&[*tag_0, *tag_1], &metadata, *table_id, *tsid) + }) + .collect::>(); + let timestamps = rows.iter().map(|row| row.4).collect::>(); + let mut writer = SeriesIndexWriter::try_new( + metadata, + object_store, + path, + SeriesIndexWriterOptions { row_group_size }, + ) + .await + .unwrap(); + writer + .write(&flat_batch(&primary_keys, ×tamps)) + .await + .unwrap(); + writer.finish().await.unwrap(); + } + + async fn collect_ids(stream: MetricSeriesIdStream) -> Vec { + stream + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .collect() + } + + #[tokio::test] + async fn search_applies_candidate_tag_filters_and_time_overlap() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let path = "search.parquet"; + write_index( + metadata.clone(), + object_store.clone(), + path, + &[ + (1, 10, "a", "x", 10), + (1, 20, "b", "x", 20), + (1, 30, "a", "y", 30), + ], + 2, + ) + .await; + + // The field filter is not available in the series index and is ignored, + // matching candidate-primary-key filter behavior. + let predicate = Predicate::new(vec![ + col("tag_0").eq(lit("a")), + col("field_0").gt(lit(0_u64)), + ]); + let time_range = TimestampRange::new( + common_time::Timestamp::new_millisecond(20), + common_time::Timestamp::new_millisecond(31), + ) + .unwrap(); + let searcher = SeriesIndexSearcher::try_new( + metadata.clone(), + object_store.clone(), + Some(&predicate), + Some(time_range), + ) + .unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 30 + }] + ); + + // Both bounds fall between millisecond ticks. Rounding the inclusive + // start and exclusive end upward leaves only the 30 ms series. + let time_range = TimestampRange::new( + common_time::Timestamp::new_microsecond(20_001), + common_time::Timestamp::new_microsecond(30_001), + ) + .unwrap(); + let searcher = SeriesIndexSearcher::try_new( + metadata.clone(), + object_store.clone(), + None, + Some(time_range), + ) + .unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 30 + }] + ); + + // The stored maximum equal to the query start intersects, while the + // stored minimum equal to the exclusive query end does not. + let time_range = TimestampRange::new( + common_time::Timestamp::new_millisecond(20), + common_time::Timestamp::new_millisecond(30), + ) + .unwrap(); + let searcher = + SeriesIndexSearcher::try_new(metadata, object_store, None, Some(time_range)).unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 20 + }] + ); + } + + #[tokio::test] + async fn search_skips_filters_for_columns_missing_from_older_index() { + let old_metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let path = "schema-evolution.parquet"; + write_index( + old_metadata.clone(), + object_store.clone(), + path, + &[ + (1, 10, "a", "x", 10), + (1, 20, "b", "x", 20), + (1, 30, "a", "y", 30), + ], + 2, + ) + .await; + + let mut builder = RegionMetadataBuilder::from_existing(old_metadata.as_ref().clone()); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("tag_2", ConcreteDataType::string_datatype(), true), + semantic_type: SemanticType::Tag, + column_id: 4, + }); + let mut primary_key = old_metadata.primary_key.clone(); + primary_key.push(4); + builder.primary_key(primary_key); + let current_metadata = Arc::new(builder.build().unwrap()); + + let predicate = + Predicate::new(vec![col("tag_0").eq(lit("a")), col("tag_2").eq(lit("new"))]); + let searcher = + SeriesIndexSearcher::try_new(current_metadata, object_store, Some(&predicate), None) + .unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![ + MetricSeriesId { + table_id: 1, + tsid: 10, + }, + MetricSeriesId { + table_id: 1, + tsid: 30, + }, + ] + ); + } + + #[tokio::test] + async fn search_streams_fixed_size_batches() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let rows = (0..501_u64) + .map(|tsid| (1, tsid, "a", "x", tsid as i64)) + .collect::>(); + write_index( + metadata.clone(), + object_store.clone(), + "batching.parquet", + &rows, + 100, + ) + .await; + + let searcher = SeriesIndexSearcher::try_new(metadata, object_store, None, None).unwrap(); + let batches = searcher + .search("batching.parquet") + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(batches.iter().map(Vec::len).collect::>(), [500, 1]); + assert_eq!( + batches[0][0], + MetricSeriesId { + table_id: 1, + tsid: 0 + } + ); + assert_eq!( + batches[1][0], + MetricSeriesId { + table_id: 1, + tsid: 500 + } + ); + } + + #[tokio::test] + async fn search_prunes_row_groups_and_empty_ranges() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let path = "pruning.parquet"; + write_index( + metadata.clone(), + object_store.clone(), + path, + &[ + (1, 0, "a", "x", 0), + (1, 1, "b", "x", 1), + (1, 2, "m", "x", 2), + (1, 3, "m", "x", 3), + (1, 4, "y", "x", 4), + (1, 5, "z", "x", 5), + ], + 2, + ) + .await; + + let predicate = Predicate::new(vec![col("tag_0").eq(lit("m"))]); + let searcher = SeriesIndexSearcher::try_new( + metadata.clone(), + object_store.clone(), + Some(&predicate), + None, + ) + .unwrap(); + let parquet_metadata = searcher.metadata_provider.load(path).await.unwrap(); + let arrow_metadata = + ArrowReaderMetadata::try_new(parquet_metadata, ArrowReaderOptions::new()).unwrap(); + let (pruning_predicate, _) = searcher.filters_for_schema(arrow_metadata.schema()); + assert_eq!( + row_groups_to_read( + arrow_metadata.metadata().row_groups(), + arrow_metadata.schema().clone(), + &pruning_predicate, + ), + vec![1] + ); + + let empty = SeriesIndexSearcher::try_new( + metadata, + object_store, + None, + Some(TimestampRange::empty()), + ) + .unwrap(); + assert!( + empty + .search("does-not-need-to-exist.parquet") + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + .is_empty() + ); + } +} diff --git a/src/mito2/src/series_index/writer.rs b/src/mito2/src/series_index/writer.rs index ce5b7405f21..9ebba289fb4 100644 --- a/src/mito2/src/series_index/writer.rs +++ b/src/mito2/src/series_index/writer.rs @@ -45,15 +45,13 @@ use crate::error::{ DecodeSnafu, InvalidMetaSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, OpenDalSnafu, Result, UnexpectedSnafu, WriteParquetSnafu, }; +use crate::series_index::{ + MAX_TS_COLUMN, MIN_TS_COLUMN, ROW_COUNT_COLUMN, TABLE_ID_COLUMN, TSID_COLUMN, +}; use crate::sst::parquet::DEFAULT_ROW_GROUP_SIZE; use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index}; use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; -const MIN_TS_COLUMN: &str = "__series_min_ts"; -const MAX_TS_COLUMN: &str = "__series_max_ts"; -const ROW_COUNT_COLUMN: &str = "__series_row_count"; -const TABLE_ID_COLUMN: &str = "__table_id"; -const TSID_COLUMN: &str = "__tsid"; const WRITE_BATCH_SIZE: usize = 1024; type ParquetWriter = AsyncArrowWriter; diff --git a/src/mito2/src/sst/parquet/format.rs b/src/mito2/src/sst/parquet/format.rs index f3f0a34a682..4ccdcb41089 100644 --- a/src/mito2/src/sst/parquet/format.rs +++ b/src/mito2/src/sst/parquet/format.rs @@ -36,7 +36,7 @@ use datafusion_common::ScalarValue; use datatypes::arrow::array::{ ArrayRef, BinaryArray, BinaryDictionaryBuilder, DictionaryArray, UInt64Array, }; -use datatypes::arrow::datatypes::{SchemaRef, UInt32Type}; +use datatypes::arrow::datatypes::{DataType as ArrowDataType, SchemaRef, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::DataType; use datatypes::types::json_type::JsonNativeType; @@ -136,12 +136,22 @@ pub(crate) fn column_values( column_index: usize, is_min: bool, ) -> Option { - let null_scalar: ScalarValue = column - .column_schema - .data_type - .as_arrow_type() - .try_into() - .ok()?; + column_values_by_type( + row_groups, + &column.column_schema.data_type.as_arrow_type(), + column_index, + is_min, + ) +} + +/// Returns min/max values of a parquet column with the given Arrow data type. +pub(crate) fn column_values_by_type( + row_groups: &[impl Borrow], + data_type: &ArrowDataType, + column_index: usize, + is_min: bool, +) -> Option { + let null_scalar: ScalarValue = data_type.try_into().ok()?; let scalar_values = row_groups .iter() .map(|meta| { @@ -152,16 +162,22 @@ pub(crate) fn column_values( } else { *s.max_opt()? }))), - Statistics::Int32(s) => Some(ScalarValue::Int32(Some(if is_min { - *s.min_opt()? - } else { - *s.max_opt()? - }))), - Statistics::Int64(s) => Some(ScalarValue::Int64(Some(if is_min { - *s.min_opt()? - } else { - *s.max_opt()? - }))), + Statistics::Int32(s) => { + let value = if is_min { *s.min_opt()? } else { *s.max_opt()? }; + if data_type == &ArrowDataType::UInt32 { + Some(ScalarValue::UInt32(Some(value as u32))) + } else { + Some(ScalarValue::Int32(Some(value))) + } + } + Statistics::Int64(s) => { + let value = if is_min { *s.min_opt()? } else { *s.max_opt()? }; + if data_type == &ArrowDataType::UInt64 { + Some(ScalarValue::UInt64(Some(value as u64))) + } else { + Some(ScalarValue::Int64(Some(value))) + } + } Statistics::Int96(_) => None, Statistics::Float(s) => Some(ScalarValue::Float32(Some(if is_min { *s.min_opt()? diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index fcaaae4a54d..3af3b20d9df 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -24,6 +24,7 @@ use std::sync::Arc; use api::v1::SemanticType; use common_recordbatch::filter::SimpleFilterEvaluator; +use datafusion_expr::Expr; use datatypes::arrow::array::{Array, BinaryArray, BooleanArray, BooleanBufferBuilder}; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::datatypes::SchemaRef; @@ -272,15 +273,9 @@ pub(crate) fn build_primary_key_filter( expected_metadata: Option<&RegionMetadata>, predicate: Option<&Predicate>, ) -> Option { - let filters = predicate + let filters = simple_tag_filters(sst_metadata, expected_metadata, predicate) .into_iter() - .flat_map(|predicate| predicate.exprs()) - .filter_map(|expr| SimpleFilterContext::new_opt(sst_metadata, expected_metadata, expr)) - .filter_map(|filter_ctx| { - (filter_ctx.semantic_type() == SemanticType::Tag) - .then(|| filter_ctx.filter().as_filter().cloned()) - .flatten() - }) + .map(|(_, filter)| filter) .collect::>(); if filters.is_empty() { return None; @@ -291,6 +286,33 @@ pub(crate) fn build_primary_key_filter( Some(CachedPrimaryKeyFilter::new(filter)) } +/// Extracts simple tag filters that can be applied to encoded primary keys or series indexes. +pub(crate) fn simple_tag_filters( + sst_metadata: &RegionMetadataRef, + expected_metadata: Option<&RegionMetadata>, + predicate: Option<&Predicate>, +) -> Vec<(Expr, SimpleFilterEvaluator)> { + predicate + .into_iter() + .flat_map(|predicate| predicate.exprs()) + .filter_map(|expr| { + SimpleFilterContext::new_opt(sst_metadata, expected_metadata, expr) + .map(|filter_ctx| (expr, filter_ctx)) + }) + .filter_map(|(expr, filter_ctx)| { + (filter_ctx.semantic_type() == SemanticType::Tag) + .then(|| { + filter_ctx + .filter() + .as_filter() + .cloned() + .map(|filter| (expr.clone(), filter)) + }) + .flatten() + }) + .collect() +} + /// How the parquet reader should apply each predicate. /// /// The reader runs in two phases. Predicates routed into `prefilter_builder` From 9bfabca72b56b15cdb9872a0fe9defe31c896262 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 25 Aug 2026 09:02:12 +0000 Subject: [PATCH 101/126] chore(deps): bump postgres-protocol from 0.6.8 to 0.6.12 (#8948) Bumps [postgres-protocol](https://github.com/rust-postgres/rust-postgres) from 0.6.8 to 0.6.12. - [Release notes](https://github.com/rust-postgres/rust-postgres/releases) - [Commits](https://github.com/rust-postgres/rust-postgres/compare/postgres-protocol-v0.6.8...postgres-protocol-v0.6.12) --- updated-dependencies: - dependency-name: postgres-protocol dependency-version: 0.6.12 dependency-type: indirect ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- Cargo.lock | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 534f484d4fc..ec477e5c9a3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10906,19 +10906,19 @@ dependencies = [ [[package]] name = "postgres-protocol" -version = "0.6.8" +version = "0.6.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76ff0abab4a9b844b93ef7b81f1efc0a366062aaef2cd702c76256b5dc075c54" +checksum = "08808e3c483c46e999108051c78334f473d5adb59d78bb80a1268c7e6aa6c514" dependencies = [ "base64 0.22.1", "byteorder", "bytes", "fallible-iterator", - "hmac 0.12.1", - "md-5 0.10.6", + "hmac 0.13.0", + "md-5 0.11.0", "memchr", - "rand 0.9.4", - "sha2 0.10.9", + "rand 0.10.1", + "sha2 0.11.0", "stringprep", ] @@ -11299,7 +11299,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.5.0", + "heck 0.4.1", "itertools 0.12.1", "log", "multimap", @@ -11319,7 +11319,7 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "itertools 0.14.0", "log", "multimap", @@ -13785,7 +13785,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.117", From 046862f067847c2f37d40f4c3601354c79b7f339 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Tue, 25 Aug 2026 09:22:53 +0000 Subject: [PATCH 102/126] feat(cmd): improve parquet rewrite fidelity and scanbench output (#8947) * feat(cmd): preserve parquet rewrite metadata settings Signed-off-by: evenyag * feat(cmd): print scanner metrics in verbose scanbench Signed-off-by: evenyag * fix(cmd): preserve SST truncation settings in property dumps Signed-off-by: evenyag --------- Signed-off-by: evenyag --- docs/parquet-tools.md | 4 +- src/cmd/src/datanode/parquet_meta.rs | 72 +++++++++++++++++++++---- src/cmd/src/datanode/parquet_rewrite.rs | 65 +++++++++++++++++++--- src/cmd/src/datanode/scanbench.rs | 19 +++++++ 4 files changed, 144 insertions(+), 16 deletions(-) diff --git a/docs/parquet-tools.md b/docs/parquet-tools.md index 8528f969431..9dfe404626e 100644 --- a/docs/parquet-tools.md +++ b/docs/parquet-tools.md @@ -89,7 +89,9 @@ Supported compression names are `uncompressed`, `snappy`, `gzip`, `lzo`, The `[writer]` table also accepts `compression_level`, `data_page_size_limit`, `data_page_row_count_limit`, and -`dictionary_page_size_limit`. A `[[columns]]` entry overrides dictionary, +`dictionary_page_size_limit`. It also accepts +`column_index_truncate_length` and `statistics_truncate_length`; set either to +zero to disable that truncation. A `[[columns]]` entry overrides dictionary, compression, compression level, or encoding for its column path. Unknown TOML fields are rejected. Compression levels are supported for `gzip`, `brotli`, and `zstd`. A column-level `compression_level` without `compression` inherits the diff --git a/src/cmd/src/datanode/parquet_meta.rs b/src/cmd/src/datanode/parquet_meta.rs index 425da45e3c0..bbd89a5bb20 100644 --- a/src/cmd/src/datanode/parquet_meta.rs +++ b/src/cmd/src/datanode/parquet_meta.rs @@ -17,6 +17,7 @@ use std::path::{Path, PathBuf}; use clap::{Parser, ValueEnum}; use parquet::file::metadata::ParquetMetaData; +use parquet::file::page_index::offset_index::PageLocation; use serde::Serialize; use snafu::ResultExt; @@ -78,6 +79,7 @@ struct ColumnChunkMetaView { dictionary_page_offset: Option, dictionary_page_bytes: Option, data_pages: Option, + max_data_page_row_count: Option, has_statistics: bool, column_index_offset: Option, column_index_length: Option, @@ -123,10 +125,14 @@ fn build_file_meta_view(path: &Path, metadata: &ParquetMetaData) -> FileMetaView .iter() .enumerate() .map(|(column_idx, column)| { - let data_pages = offset_index + let page_locations = offset_index .and_then(|index| index.get(row_group_idx)) .and_then(|columns| columns.get(column_idx)) - .map(|index| index.page_locations().len()); + .map(|index| index.page_locations()); + let data_pages = page_locations.map(|locations| locations.len()); + let max_data_page_row_count = page_locations.and_then(|locations| { + max_data_page_row_count(locations, row_group.num_rows()) + }); let dictionary_page_bytes = dictionary_page_bytes( column.dictionary_page_offset(), column.data_page_offset(), @@ -149,6 +155,7 @@ fn build_file_meta_view(path: &Path, metadata: &ParquetMetaData) -> FileMetaView dictionary_page_offset: column.dictionary_page_offset(), dictionary_page_bytes, data_pages, + max_data_page_row_count, has_statistics: column.statistics().is_some(), column_index_offset: column.column_index_offset(), column_index_length: column.column_index_length(), @@ -219,7 +226,7 @@ fn print_meta_text(view: &FileMetaView) { ); for column in &row_group.columns { println!( - " column[{}] path={}: type={}, compression={}, encodings=[{}], values={}, uncompressed_size={}, compressed_size={}, compression_ratio={}, data_pages={}, dictionary_page_offset={}, dictionary_page_bytes={}, data_page_offset={}, statistics={}, column_index={}/{}, offset_index={}/{}, bloom_filter={}/{}", + " column[{}] path={}: type={}, compression={}, encodings=[{}], values={}, uncompressed_size={}, compressed_size={}, compression_ratio={}, data_pages={}, max_data_page_row_count={}, dictionary_page_offset={}, dictionary_page_bytes={}, data_page_offset={}, statistics={}, column_index={}/{}, offset_index={}/{}, bloom_filter={}/{}", column.index, column.path, column.physical_type, @@ -230,6 +237,7 @@ fn print_meta_text(view: &FileMetaView) { column.compressed_size, format_ratio(column.compression_ratio), format_optional_usize(column.data_pages), + format_optional_i64(column.max_data_page_row_count), format_optional_i64(column.dictionary_page_offset), column .dictionary_page_bytes @@ -248,6 +256,33 @@ fn print_meta_text(view: &FileMetaView) { } } +fn max_data_page_row_count( + page_locations: &[PageLocation], + row_group_row_count: i64, +) -> Option { + if page_locations.is_empty() { + return Some(0); + } + if page_locations.first()?.first_row_index != 0 { + return None; + } + + let mut max_row_count = 0; + for locations in page_locations.windows(2) { + let row_count = locations[1].first_row_index - locations[0].first_row_index; + if row_count < 0 { + return None; + } + max_row_count = max_row_count.max(row_count); + } + + let last_page_row_count = row_group_row_count - page_locations.last()?.first_row_index; + if last_page_row_count < 0 { + return None; + } + Some(max_row_count.max(last_page_row_count)) +} + fn dictionary_page_bytes( dictionary_page_offset: Option, data_page_offset: i64, @@ -301,7 +336,7 @@ mod tests { fn test_meta_view_unknown_page_count() { let dir = tempdir().unwrap(); let path = dir.path().join("input.parquet"); - write_test_parquet(&path); + write_test_parquet(&path, None); let metadata = load_local_parquet_metadata(&path).unwrap(); let view = build_file_meta_view(&path, &metadata); @@ -313,7 +348,21 @@ mod tests { assert_eq!(view.row_groups[0].columns.len(), 2); } - fn write_test_parquet(path: &Path) { + #[test] + fn test_meta_view_reports_max_data_page_row_count() { + let dir = tempdir().unwrap(); + let path = dir.path().join("input.parquet"); + write_test_parquet(&path, Some(2)); + + let metadata = load_local_parquet_metadata(&path).unwrap(); + let view = build_file_meta_view(&path, &metadata); + + for column in &view.row_groups[0].columns { + assert_eq!(Some(2), column.max_data_page_row_count); + } + } + + fn write_test_parquet(path: &Path, data_page_row_count_limit: Option) { let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), Field::new("host", DataType::Utf8, false), @@ -326,12 +375,17 @@ mod tests { ], ) .unwrap(); - let props = WriterProperties::builder() - .set_key_value_metadata(Some(vec![KeyValue::new( + let mut props = + WriterProperties::builder().set_key_value_metadata(Some(vec![KeyValue::new( "greptime:test".to_string(), "value".to_string(), - )])) - .build(); + )])); + if let Some(limit) = data_page_row_count_limit { + props = props + .set_write_batch_size(limit) + .set_data_page_row_count_limit(limit); + } + let props = props.build(); let mut writer = ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap(); writer.write(&batch).unwrap(); diff --git a/src/cmd/src/datanode/parquet_rewrite.rs b/src/cmd/src/datanode/parquet_rewrite.rs index 47346aac889..b04db93edcb 100644 --- a/src/cmd/src/datanode/parquet_rewrite.rs +++ b/src/cmd/src/datanode/parquet_rewrite.rs @@ -77,6 +77,8 @@ struct WriterConfig { data_page_size_limit: Option, data_page_row_count_limit: Option, dictionary_page_size_limit: Option, + column_index_truncate_length: Option, + statistics_truncate_length: Option, } #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] @@ -189,6 +191,15 @@ impl ParquetRewriteCommand { } fn infer_rewrite_properties(metadata: &ParquetMetaData) -> RewriteProperties { + let is_greptime_sst = metadata + .file_metadata() + .key_value_metadata() + .is_some_and(|key_values| { + key_values + .iter() + .any(|key_value| key_value.key == mito2::sst::parquet::PARQUET_METADATA_KEY) + }); + let truncate_length = is_greptime_sst.then_some(0); let first_column = metadata .row_groups() .first() @@ -234,6 +245,8 @@ fn infer_rewrite_properties(metadata: &ParquetMetaData) -> RewriteProperties { data_page_size_limit: None, data_page_row_count_limit: None, dictionary_page_size_limit: None, + column_index_truncate_length: truncate_length, + statistics_truncate_length: truncate_length, }, columns, } @@ -287,6 +300,16 @@ fn build_writer_properties( if let Some(dictionary_page_size_limit) = config.writer.dictionary_page_size_limit { builder = builder.set_dictionary_page_size_limit(dictionary_page_size_limit); } + if let Some(column_index_truncate_length) = config.writer.column_index_truncate_length { + builder = builder.set_column_index_truncate_length( + (column_index_truncate_length > 0).then_some(column_index_truncate_length), + ); + } + if let Some(statistics_truncate_length) = config.writer.statistics_truncate_length { + builder = builder.set_statistics_truncate_length( + (statistics_truncate_length > 0).then_some(statistics_truncate_length), + ); + } for column in config.columns { if column.path.is_empty() { @@ -492,7 +515,7 @@ mod tests { let dir = tempdir().unwrap(); let input = dir.path().join("input.parquet"); let output = dir.path().join("output.parquet"); - write_test_parquet(&input, true); + write_test_parquet(&input, true, Some("greptime:test")); let metadata = load_local_parquet_metadata(&input).unwrap(); let mut properties = infer_rewrite_properties(&metadata); @@ -528,6 +551,37 @@ mod tests { ); } + #[test] + fn test_infer_truncation_settings_for_greptime_sst() { + let dir = tempdir().unwrap(); + let greptime_sst = dir.path().join("greptime.parquet"); + let generic_parquet = dir.path().join("generic.parquet"); + write_test_parquet( + &greptime_sst, + true, + Some(mito2::sst::parquet::PARQUET_METADATA_KEY), + ); + write_test_parquet(&generic_parquet, true, None); + + let metadata = load_local_parquet_metadata(&greptime_sst).unwrap(); + let config = infer_rewrite_properties(&metadata); + assert_eq!(Some(0), config.writer.column_index_truncate_length); + assert_eq!(Some(0), config.writer.statistics_truncate_length); + + let serialized = toml::to_string(&config).unwrap(); + assert!(serialized.contains("column_index_truncate_length = 0")); + assert!(serialized.contains("statistics_truncate_length = 0")); + + let properties = build_writer_properties(config, None).unwrap(); + assert_eq!(None, properties.column_index_truncate_length()); + assert_eq!(None, properties.statistics_truncate_length()); + + let metadata = load_local_parquet_metadata(&generic_parquet).unwrap(); + let config = infer_rewrite_properties(&metadata); + assert_eq!(None, config.writer.column_index_truncate_length); + assert_eq!(None, config.writer.statistics_truncate_length); + } + #[test] fn test_column_compression_level_inherits_writer_compression() { let path = ColumnPath::new(vec!["host".to_string()]); @@ -551,7 +605,7 @@ mod tests { ); } - fn write_test_parquet(path: &Path, dictionary_enabled: bool) { + fn write_test_parquet(path: &Path, dictionary_enabled: bool, metadata_key: Option<&str>) { let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), Field::new("host", DataType::Utf8, false), @@ -564,12 +618,11 @@ mod tests { ], ) .unwrap(); + let key_value_metadata = + metadata_key.map(|key| vec![KeyValue::new(key.to_string(), "value".to_string())]); let props = WriterProperties::builder() .set_dictionary_enabled(dictionary_enabled) - .set_key_value_metadata(Some(vec![KeyValue::new( - "greptime:test".to_string(), - "value".to_string(), - )])) + .set_key_value_metadata(key_value_metadata) .build(); let mut writer = ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap(); diff --git a/src/cmd/src/datanode/scanbench.rs b/src/cmd/src/datanode/scanbench.rs index 1dbee025a7f..fd6c08fe6db 100644 --- a/src/cmd/src/datanode/scanbench.rs +++ b/src/cmd/src/datanode/scanbench.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashMap; +use std::fmt; use std::path::PathBuf; use std::sync::Arc; use std::time::Instant; @@ -31,6 +32,7 @@ use datafusion::logical_expr::{BinaryExpr, Expr as DfExpr, ExprSchemable, Operat use datafusion_common::tree_node::{Transformed, TreeNodeRewriter}; use datafusion_common::{DFSchemaRef, ScalarValue, ToDFSchema}; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_physical_plan::{DisplayAs, DisplayFormatType}; use datatypes::arrow::compute; use futures::StreamExt; use futures::stream::FuturesUnordered; @@ -61,6 +63,15 @@ use crate::datanode::tool_util::{ }; use crate::error; +/// Displays a scanner using DataFusion's verbose explain format. +struct VerboseScannerDisplay<'a, T: ?Sized>(&'a T); + +impl fmt::Display for VerboseScannerDisplay<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt_as(DisplayFormatType::Verbose, f) + } +} + /// Scan benchmark command - benchmarks scanning a region directly from storage. #[derive(Debug, Parser)] pub struct ScanbenchCommand { @@ -688,6 +699,14 @@ impl ScanbenchCommand { format_bytes(total_estimated_size), ); + if self.verbose { + println!( + " {} Scanner explain: {}", + "ℹ".blue(), + VerboseScannerDisplay(scanner.as_ref()) + ); + } + // Start profiling after the first iteration (warmup) if pprof_after_warmup is set #[cfg(unix)] if iteration == 0 From 8a473c5bf0d3cb88d0bd328c14f33d5f2fea2224 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 25 Aug 2026 10:30:59 +0000 Subject: [PATCH 103/126] fix(meta): allow manual migration from offline datanodes (#8934) * fix(meta): allow migration from offline datanodes Signed-off-by: WenyXu * test: fix offline migration event actor Signed-off-by: WenyXu * test: read migration routes from metadata Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/meta-srv/src/service/procedure.rs | 3 +- tests-integration/tests/region_migration.rs | 135 +++++++++++++++++++- 2 files changed, 134 insertions(+), 4 deletions(-) diff --git a/src/meta-srv/src/service/procedure.rs b/src/meta-srv/src/service/procedure.rs index df27e2cb32b..be4eedc2ab7 100644 --- a/src/meta-srv/src/service/procedure.rs +++ b/src/meta-srv/src/service/procedure.rs @@ -26,6 +26,7 @@ use api::v1::meta::{ use common_event_recorder::{PersistentEventContext, ProcedureEventInput}; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_name::TableNameKey; +use common_meta::peer::Peer; use common_meta::procedure_executor::ExecutorContext; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, @@ -181,7 +182,7 @@ impl procedure_service_server::ProcedureService for Metasrv { let from_peer = self .lookup_datanode_peer(from_peer) .await? - .context(error::PeerUnavailableSnafu { peer_id: from_peer })?; + .unwrap_or_else(|| Peer::empty(from_peer)); let to_peer = self .lookup_datanode_peer(to_peer) .await? diff --git a/tests-integration/tests/region_migration.rs b/tests-integration/tests/region_migration.rs index 17a58de2eed..0df42ee1e88 100644 --- a/tests-integration/tests/region_migration.rs +++ b/tests-integration/tests/region_migration.rs @@ -25,8 +25,10 @@ use common_event_recorder::{ DEFAULT_EVENTS_TABLE_NAME, DEFAULT_FLUSH_INTERVAL_SECONDS, EVENTS_TABLE_TIMESTAMP_COLUMN_NAME, EVENTS_TABLE_TYPE_COLUMN_NAME, PersistentEventContext, TriggerReason, }; +use common_meta::distributed_time_constants::default_distributed_time_constants; use common_meta::key::{RegionDistribution, RegionRoleSet, TableMetadataManagerRef}; use common_meta::peer::Peer; +use common_meta::rpc::store::BatchDeleteRequest; use common_procedure::ProcedureContext; use common_procedure::event::{ EVENTS_TABLE_PROCEDURE_ID_COLUMN_NAME, EVENTS_TABLE_PROCEDURE_STATE_COLUMN_NAME, @@ -46,6 +48,7 @@ use futures::future::BoxFuture; use meta_srv::error; use meta_srv::error::Result as MetaResult; use meta_srv::event::region_migration::REGION_MIGRATION_EVENT_TYPE; +use meta_srv::key::DatanodeLeaseKey; use meta_srv::metasrv::SelectorContext; use meta_srv::procedure::region_migration::{ RegionMigrationProcedureTask, RegionMigrationTriggerReason, @@ -102,6 +105,7 @@ macro_rules! region_migration_tests { test_region_migration, test_region_migration_by_sql, + test_region_migration_with_offline_source_by_sql, test_region_migration_multiple_regions, test_region_migration_all_regions, test_region_migration_incorrect_from_peer, @@ -419,8 +423,24 @@ pub async fn test_metric_table_region_migration_by_sql( .await; } -/// A naive region migration test by SQL function +/// A naive region migration test by SQL function. pub async fn test_region_migration_by_sql(store_type: StorageType, endpoints: Vec) { + test_region_migration_by_sql_inner(store_type, endpoints, false).await; +} + +/// A region migration test by SQL function with an offline source datanode. +pub async fn test_region_migration_with_offline_source_by_sql( + store_type: StorageType, + endpoints: Vec, +) { + test_region_migration_by_sql_inner(store_type, endpoints, true).await; +} + +async fn test_region_migration_by_sql_inner( + store_type: StorageType, + endpoints: Vec, + simulate_offline_source: bool, +) { let cluster_name = "test_region_migration"; let peer_factory = |id| Peer { id, @@ -437,7 +457,7 @@ pub async fn test_region_migration_by_sql(store_type: StorageType, endpoints: Ve peer_factory(2), peer_factory(3), ])); - let cluster = builder + let mut cluster = builder .with_datanodes(datanodes as u32) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Kafka(DatanodeKafkaConfig { @@ -463,6 +483,7 @@ pub async fn test_region_migration_by_sql(store_type: StorageType, endpoints: Ve .with_meta_selector(const_selector.clone()) .build(true) .await; + let table_metadata_manager = cluster.metasrv.table_metadata_manager().clone(); let (actor_db, _actor_grpc_server) = setup_authenticated_grpc_database( cluster.fe_instance().clone(), PROCEDURE_ACTOR, @@ -487,7 +508,7 @@ pub async fn test_region_migration_by_sql(store_type: StorageType, endpoints: Ve let old_distribution = distribution.clone(); // Selecting target of region migration. - let region_migration_manager = cluster.metasrv.region_migration_manager(); + let region_migration_manager = cluster.metasrv.region_migration_manager().clone(); let (from_peer_id, from_regions) = distribution.pop_first().unwrap(); info!( "Selecting from peer: {from_peer_id}, and regions: {:?}", @@ -545,6 +566,114 @@ pub async fn test_region_migration_by_sql(store_type: StorageType, endpoints: Ve // Asserts the writes. assert_values(cluster.fe_instance()).await; + if simulate_offline_source { + let mut expected_distribution = + find_region_distribution(&table_metadata_manager, table_id).await; + let (offline_from_peer_id, offline_region_number) = expected_distribution + .iter() + .find(|(peer_id, regions)| { + **peer_id != from_peer_id + && **peer_id != to_peer_id + && !regions.leader_regions.is_empty() + }) + .map(|(peer_id, regions)| (*peer_id, regions.leader_regions[0])) + .unwrap(); + let offline_region_id = RegionId::new(table_id, offline_region_number); + + // Simulates scale-in: the source datanode stops and its lease disappears. + cluster + .datanode_instances + .get_mut(&offline_from_peer_id) + .unwrap() + .shutdown() + .await + .unwrap(); + let source_lease_key: Vec = DatanodeLeaseKey { + node_id: offline_from_peer_id, + } + .try_into() + .unwrap(); + cluster + .metasrv + .in_memory() + .batch_delete(BatchDeleteRequest { + keys: vec![source_lease_key], + prev_kv: false, + }) + .await + .unwrap(); + + let offline_procedure_id = trigger_migration_by_sql( + &cluster, + offline_region_id.as_u64(), + offline_from_peer_id, + from_peer_id, + ) + .await; + let frontend = cluster.fe_instance().clone(); + let procedure_id_for_closure = offline_procedure_id.clone(); + wait_condition( + default_distributed_time_constants().region_lease + Duration::from_secs(10), + Box::pin(async move { + loop { + let state = query_procedure_by_sql(&frontend, &procedure_id_for_closure).await; + if state == "{\"status\":\"Done\"}" { + info!("Offline-source migration done: {state}"); + break; + } + info!("Offline-source migration not finished: {state}"); + tokio::time::sleep(Duration::from_millis(200)).await; + } + }), + ) + .await; + + check_region_migration_events_system_table( + cluster.fe_instance(), + &offline_procedure_id, + offline_region_id.as_u64(), + offline_from_peer_id, + from_peer_id, + Some("greptime"), + ) + .await; + + let remove_offline_source = { + let source_regions = expected_distribution + .get_mut(&offline_from_peer_id) + .unwrap(); + source_regions + .leader_regions + .retain(|region_number| *region_number != offline_region_number); + source_regions.leader_regions.is_empty() && source_regions.follower_regions.is_empty() + }; + if remove_offline_source { + expected_distribution.remove(&offline_from_peer_id); + } + let target_regions = expected_distribution.entry(from_peer_id).or_default(); + target_regions.add_leader_region(offline_region_number); + target_regions.sort(); + + let table_metadata_manager = table_metadata_manager.clone(); + wait_condition( + Duration::from_secs(10), + Box::pin(async move { + loop { + let distribution = + find_region_distribution(&table_metadata_manager, table_id).await; + if distribution == expected_distribution { + break; + } + info!("Offline-source migration has unexpected distribution: {distribution:?}"); + tokio::time::sleep(Duration::from_millis(200)).await; + } + }), + ) + .await; + + assert_values(cluster.fe_instance()).await; + } + // Triggers again. let err = region_migration_manager .submit_procedure( From 28398138ec03d154524a8217672d463f28e7f466 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 25 Aug 2026 10:31:06 +0000 Subject: [PATCH 104/126] fix(flight): bound DoGet response wait (#8943) * fix(flight): defer datanode query initialization Signed-off-by: WenyXu * fix(client): retain Flight stream peer context Signed-off-by: WenyXu * fix(client): improve Flight stream diagnostics Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/client/src/region.rs | 116 +++++++++++++++++++++----- src/datanode/src/region_server.rs | 20 +++-- src/servers/src/grpc/flight.rs | 6 +- src/servers/src/grpc/flight/stream.rs | 104 +++++++++++++++++++---- tests-integration/src/grpc/flight.rs | 103 ++++++++++++++++++++++- 5 files changed, 299 insertions(+), 50 deletions(-) diff --git a/src/client/src/region.rs b/src/client/src/region.rs index f1730a93c8e..2bb5f7dd0a6 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::sync::Arc; +use std::time::Duration; use api::region::RegionResponse; use api::v1::ResponseHeader; @@ -24,7 +25,7 @@ use arc_swap::ArcSwapOption; use arrow_flight::Ticket; use async_stream::stream; use async_trait::async_trait; -use common_error::ext::BoxedError; +use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; use common_grpc::flight::{FlightDecoder, FlightMessage}; use common_meta::error::{self as meta_error, Result as MetaResult}; @@ -49,6 +50,8 @@ use crate::error::{ use crate::flight::decode_flight_data; use crate::{Client, metrics}; +const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); + #[derive(Debug)] pub struct RegionRequester { client: Client, @@ -109,20 +112,24 @@ impl RegionRequester { let mut flight_client = self .client .make_flight_client(self.send_compression, self.accept_compression)?; + // Limit Flight DoGet response time without limiting query stream execution. + let addr = flight_client.addr().to_string(); + let mut request = tonic::Request::new(ticket); + request.set_timeout(FLIGHT_DO_GET_TIMEOUT); let response = flight_client .mut_inner() - .do_get(ticket) + .do_get(request) .await .or_else(|e| { let tonic_code = e.code(); let e: error::Error = e.into(); error!( e; "Failed to do Flight get, addr: {}, code: {}", - flight_client.addr(), + addr, tonic_code ); Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu { - addr: flight_client.addr().to_string(), + addr: addr.clone(), tonic_code, }) })?; @@ -133,7 +140,7 @@ impl RegionRequester { let flight_message_stream = flight_data_stream .filter_map(move |flight_data| decode_flight_data(&mut decoder, flight_data)); - recordbatches_from_flight_message_stream(flight_message_stream).await + recordbatches_from_flight_message_stream(addr, flight_message_stream).await } async fn handle_inner(&self, request: RegionRequest) -> Result { @@ -195,6 +202,7 @@ impl RegionRequester { } async fn recordbatches_from_flight_message_stream( + addr: String, mut flight_message_stream: S, ) -> Result where @@ -204,13 +212,17 @@ where return IllegalFlightMessagesSnafu { reason: "Expect the response not to be empty", } - .fail(); + .fail() + .map_err(|error| flight_stream_error(&addr, error)); }; - let FlightMessage::Schema(schema) = first_flight_message? else { + let FlightMessage::Schema(schema) = + first_flight_message.map_err(|e| flight_stream_error(&addr, e))? + else { return IllegalFlightMessagesSnafu { reason: "Expect schema to be the first flight message", } - .fail(); + .fail() + .map_err(|error| flight_stream_error(&addr, error)); }; let metrics = Arc::new(ArcSwapOption::from(None)); @@ -221,6 +233,7 @@ where let schema = Arc::new(datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?); let schema_cloned = schema.clone(); + let stream_addr = addr.clone(); let stream = Box::pin(stream!({ let _span = tracing_context.attach(common_telemetry::tracing::info_span!( "poll_flight_data_stream" @@ -240,7 +253,8 @@ where let flight_message = match flight_message_item { Some(Ok(message)) => message, Some(Err(e)) => { - yield Err(BoxedError::new(e)).context(ExternalSnafu); + yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + .context(ExternalSnafu); break; } None => break, @@ -273,7 +287,8 @@ where break; } Err(e) => { - yield Err(BoxedError::new(e)).context(ExternalSnafu); + yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + .context(ExternalSnafu); break; } } @@ -312,6 +327,23 @@ where Ok(Box::pin(record_batch_stream)) } +fn flight_stream_error(addr: &str, error: error::Error) -> error::Error { + let tonic_code = error.tonic_code().unwrap_or(tonic::Code::Unknown); + if error.status_code().should_log_error() { + error!( + error; "Failed to receive Flight data, addr: {}, code: {}", + addr, + tonic_code + ); + } + + error::Error::FlightGet { + addr: addr.to_string(), + tonic_code, + source: BoxedError::new(error), + } +} + pub fn build_remote_dyn_filter_update_request( query_id: impl Into, update: RemoteDynFilterUpdate, @@ -389,6 +421,44 @@ mod test { use super::*; use crate::Error::{self, IllegalDatabaseResponse, Server}; + #[test] + fn test_flight_stream_error_preserves_peer_address() { + let error = flight_stream_error( + "127.0.0.1:4001", + tonic::Status::unavailable("datanode unavailable").into(), + ); + + assert!(matches!( + error, + error::Error::FlightGet { + addr, + tonic_code: tonic::Code::Unavailable, + .. + } if addr == "127.0.0.1:4001" + )); + } + + #[tokio::test] + async fn test_empty_flight_stream_preserves_peer_address() { + let Err(error) = recordbatches_from_flight_message_stream( + "127.0.0.1:4001".to_string(), + stream::empty::>(), + ) + .await + else { + panic!("expected empty Flight stream to fail"); + }; + + assert!(matches!( + error, + error::Error::FlightGet { + addr, + tonic_code: tonic::Code::Unknown, + .. + } if addr == "127.0.0.1:4001" + )); + } + fn test_schema() -> Arc { Arc::new(Schema::new(vec![ColumnSchema::new( "v", @@ -509,11 +579,14 @@ mod test { ) .unwrap(); - let mut recordbatches = recordbatches_from_flight_message_stream(stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::Metrics(test_metrics_json())), - Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), - ])) + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::Metrics(test_metrics_json())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + ]), + ) .await .unwrap(); @@ -528,11 +601,14 @@ mod test { #[tokio::test] async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() { let schema = test_schema(); - let mut recordbatches = recordbatches_from_flight_message_stream(stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::Metrics(test_metrics_json())), - Err(Error::from(Status::internal("boom after metrics"))), - ])) + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::Metrics(test_metrics_json())), + Err(Error::from(Status::internal("boom after metrics"))), + ]), + ) .await .unwrap(); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 57facfd53ba..0d9398b0889 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -66,7 +66,9 @@ use servers::error::{ self as servers_error, ExecuteGrpcRequestSnafu, Result as ServerResult, SuspendedSnafu, }; use servers::grpc::FlightCompression; -use servers::grpc::flight::{FlightCraft, FlightRecordBatchStream, TonicStream}; +use servers::grpc::flight::{ + FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, TonicStream, +}; use servers::grpc::region_server::RegionServerHandler; use session::context::{ FLIGHT_METRICS_HEARTBEAT_INTERVAL, QueryContext, QueryContextBuilder, QueryContextRef, @@ -976,13 +978,19 @@ impl FlightCraft for RegionServer { .map(|h| Arc::new(QueryContext::from(h))) .unwrap_or(QueryContext::arc()); - let result = self - .handle_remote_read(request, query_ctx.clone()) - .trace(tracing_context.attach(info_span!("RegionServer::handle_read"))) - .await?; + let region_server = self.clone(); + let initializer_query_ctx = query_ctx.clone(); + let initializer_tracing_context = tracing_context.clone(); + let initializer = async move { + region_server + .handle_remote_read(request, initializer_query_ctx) + .trace(initializer_tracing_context.attach(info_span!("RegionServer::handle_read"))) + .await + .map_err(Into::into) + }; let stream = Box::pin(FlightRecordBatchStream::new( - result, + FlightRecordBatchSource::initializer(initializer), tracing_context, self.flight_compression, query_ctx, diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index 97a579d952c..e6f35e17b29 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -53,7 +53,7 @@ use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status, Streaming}; use crate::error::{InvalidParameterSnafu, Result, ToJsonSnafu}; -pub use crate::grpc::flight::stream::FlightRecordBatchStream; +pub use crate::grpc::flight::stream::{FlightRecordBatchSource, FlightRecordBatchStream}; use crate::grpc::greptime_handler::{ GreptimeRequestHandler, create_query_context, get_request_type, }; @@ -583,7 +583,7 @@ fn to_flight_data_stream( match output.data { OutputData::Stream(stream) => { let stream = FlightRecordBatchStream::new( - stream, + FlightRecordBatchSource::RecordBatches(stream), tracing_context, flight_compression, query_ctx, @@ -592,7 +592,7 @@ fn to_flight_data_stream( } OutputData::RecordBatches(x) => { let stream = FlightRecordBatchStream::new( - x.as_stream(), + FlightRecordBatchSource::RecordBatches(x.as_stream()), tracing_context, flight_compression, query_ctx, diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index c222c72e1cd..319a797d9ca 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/src/servers/src/grpc/flight/stream.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::VecDeque; +use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; @@ -40,6 +41,30 @@ use crate::error; use crate::grpc::FlightCompression; use crate::grpc::flight::TonicResult; +pub type FlightRecordBatchStreamInitializer = Pin< + Box< + dyn Future> + + Send + + 'static, + >, +>; + +pub enum FlightRecordBatchSource { + RecordBatches(SendableRecordBatchStream), + Initializer(FlightRecordBatchStreamInitializer), +} + +impl FlightRecordBatchSource { + pub fn initializer(initializer: F) -> Self + where + F: Future> + + Send + + 'static, + { + Self::Initializer(Box::pin(initializer)) + } +} + /// Metrics collector for Flight stream with RAII logging pattern struct StreamMetrics { send_schema_duration: Duration, @@ -136,7 +161,7 @@ impl FlightRecordBatchStream { } pub fn new( - recordbatches: SendableRecordBatchStream, + source: FlightRecordBatchSource, tracing_context: TracingContext, compression: FlightCompression, query_ctx: QueryContextRef, @@ -148,17 +173,43 @@ impl FlightRecordBatchStream { .remote_query_id() .zip(query_ctx.extension(SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY)) .is_some_and(|(remote_query_id, capability)| capability == remote_query_id); - let (tx, rx) = mpsc::channel::>(1); - let join_handle = common_runtime::spawn_global(async move { - Self::flight_data_stream( - recordbatches, - tx, - should_send_partial_metrics, - can_send_metrics_before_batch, - ) - .trace(tracing_context.attach(info_span!("flight_data_stream"))) - .await - }); + let (mut tx, rx) = mpsc::channel::>(1); + let source_type = match &source { + FlightRecordBatchSource::RecordBatches(_) => "record_batches", + FlightRecordBatchSource::Initializer(_) => "initializer", + }; + let initializer_tracing_context = tracing_context.clone(); + let join_handle = common_runtime::spawn_global( + async move { + let recordbatches = async move { + match source { + FlightRecordBatchSource::RecordBatches(recordbatches) => Ok(recordbatches), + FlightRecordBatchSource::Initializer(initializer) => initializer.await, + } + } + .trace( + initializer_tracing_context + .attach(info_span!("flight_data_stream_init", source_type)), + ) + .await; + + match recordbatches { + Ok(recordbatches) => { + Self::flight_data_stream( + recordbatches, + tx, + should_send_partial_metrics, + can_send_metrics_before_batch, + ) + .await; + } + Err(status) => { + let _ = tx.send(Err(status)).await; + } + } + } + .trace(tracing_context.attach(info_span!("flight_data_stream"))), + ); let encoder = if compression.arrow_compression() { FlightEncoder::default() } else { @@ -434,7 +485,7 @@ mod test { .unwrap() .as_stream(); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), QueryContext::arc(), @@ -468,6 +519,23 @@ mod test { } } + #[tokio::test] + async fn test_flight_record_batch_stream_forwards_initializer_error() { + let mut stream = FlightRecordBatchStream::new( + FlightRecordBatchSource::initializer(async { + Err(tonic::Status::unavailable( + "remote read initialization failed", + )) + }), + TracingContext::default(), + FlightCompression::default(), + QueryContext::arc(), + ); + + let error = stream.next().await.unwrap().unwrap_err(); + assert_eq!(tonic::Code::Unavailable, error.code()); + assert!(stream.next().await.is_none()); + } #[tokio::test] async fn test_flight_record_batch_stream_emits_metrics_while_pending() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( @@ -486,7 +554,7 @@ mod test { let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -541,7 +609,7 @@ mod test { let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -595,7 +663,7 @@ mod test { let query_ctx = query_context_with_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -638,7 +706,7 @@ mod test { let query_ctx = Arc::new(query_ctx); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -677,7 +745,7 @@ mod test { }); let query_ctx = query_context_with_matching_capability(); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index 520635518a4..cf75ae7fff1 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -17,35 +17,132 @@ mod test { use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Arc; - use std::time::Duration; + use std::time::{Duration, Instant}; use api::v1::auth_header::AuthScheme; use api::v1::query_request::Query; use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, QueryRequest, SemanticType}; - use arrow_flight::FlightDescriptor; + use arrow_flight::flight_service_server::FlightServiceServer; + use arrow_flight::{FlightData, FlightDescriptor, Ticket}; use auth::user_provider_from_option; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_grpc::flight::do_put::DoPutMetadata; use common_grpc::flight::{FlightEncoder, FlightMessage}; use common_query::OutputData; - use common_recordbatch::RecordBatch; use common_recordbatch::adapter::RegionWatermarkEntry; + use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; + use common_telemetry::tracing_context::TracingContext; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector}; use futures_util::StreamExt; + use hyper_util::rt::TokioIo; use itertools::Itertools; use servers::grpc::builder::GrpcServerBuilder; + use servers::grpc::flight::{ + FlightCraft, FlightCraftWrapper, FlightRecordBatchSource, FlightRecordBatchStream, + TonicStream, + }; use servers::grpc::greptime_handler::GreptimeRequestHandler; use servers::grpc::{FlightCompression, GrpcServerConfig}; use servers::server::Server; + use tonic::Response; + use tonic::transport::Server as TonicServer; + use tower::service_fn; use crate::cluster::GreptimeDbClusterBuilder; use crate::grpc::query_and_expect; use crate::test_util::{StorageType, setup_grpc_server}; use crate::tests::test_util::MockInstance; + struct SlowFlightCraft; + + fn slow_recordbatch_stream() -> SendableRecordBatchStream { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "value", + ConcreteDataType::int32_datatype(), + false, + )])); + let recordbatch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_vec(vec![1])) as VectorRef], + ) + .unwrap(); + + RecordBatches::try_new(schema, vec![recordbatch]) + .unwrap() + .as_stream() + } + + #[async_trait::async_trait] + impl FlightCraft for SlowFlightCraft { + async fn do_get( + &self, + _: tonic::Request, + ) -> std::result::Result>, tonic::Status> { + let stream = FlightRecordBatchStream::new( + FlightRecordBatchSource::initializer(async { + tokio::time::sleep(Duration::from_secs(2)).await; + Ok(slow_recordbatch_stream()) + }), + TracingContext::default(), + FlightCompression::default(), + session::context::QueryContext::arc(), + ); + + Ok(Response::new(Box::pin(stream))) + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { + let (client_io, server_io) = tokio::io::duplex(1024); + tokio::spawn(async move { + TonicServer::builder() + .add_service(FlightServiceServer::new(FlightCraftWrapper( + SlowFlightCraft, + ))) + .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>( + server_io, + )])) + .await + .unwrap(); + }); + + let channel_manager = ChannelManager::with_config(ChannelConfig::new().timeout(None), None); + let mut client_io = Some(client_io); + channel_manager + .reset_with_connector( + "slow-flight", + service_fn(move |_| { + let client_io = client_io.take(); + + async move { + client_io + .map(TokioIo::new) + .ok_or_else(|| std::io::Error::other("Client already taken")) + } + }), + ) + .unwrap(); + let client = Client::with_manager_and_urls(channel_manager, ["slow-flight"]); + let mut flight_client = client.make_flight_client(false, false).unwrap(); + + let start = Instant::now(); + let mut request = tonic::Request::new(Ticket::default()); + request.set_timeout(Duration::from_secs(1)); + let response = flight_client.mut_inner().do_get(request).await.unwrap(); + assert!(start.elapsed() < Duration::from_secs(1)); + let mut stream = response.into_inner(); + + let start = Instant::now(); + assert!(stream.message().await.unwrap().is_some()); + assert!(start.elapsed() >= Duration::from_secs(1)); + + assert!(stream.message().await.unwrap().is_some()); + } #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put() { common_telemetry::init_default_ut_logging(); From 932f87f7a82d6d7af270b93e3bc2420f0f5cc7ab Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:11:06 +0000 Subject: [PATCH 105/126] refactor(json2): support querying v2 storage layout (#8940) * feat(json2): support querying v2 storage layout - route missing JSON2 paths to the v2 remainder field - reconstruct complete values from explicit fields and remainder data - preserve root JSON2 columns across projections and filters - support nested JSON values in json_get string results - add and reorganize JSON2 sqlness coverage Signed-off-by: luofucong * resolve PR comments Signed-off-by: luofucong --------- Signed-off-by: luofucong --- .../function/src/scalars/json/json_get.rs | 70 ++++++-- src/mito2/src/read/flat_projection.rs | 4 +- src/mito2/src/sst/parquet/read_columns.rs | 167 +++++++++++++++++- .../src/optimizer/json_type_concretize.rs | 59 +++++-- .../standalone/common/types/json/json2.result | 25 +++ .../standalone/common/types/json/json2.sql | 4 + .../common/types/json/json2_join.result | 68 +++++++ .../common/types/json/json2_join.sql | 37 ++++ .../common/types/json/json2_limit.result | 77 -------- .../common/types/json/json2_limit.sql | 43 ----- 10 files changed, 392 insertions(+), 162 deletions(-) create mode 100644 tests/cases/standalone/common/types/json/json2_join.result create mode 100644 tests/cases/standalone/common/types/json/json2_join.sql diff --git a/src/common/function/src/scalars/json/json_get.rs b/src/common/function/src/scalars/json/json_get.rs index ce49f9d59fc..2357bc88f91 100644 --- a/src/common/function/src/scalars/json/json_get.rs +++ b/src/common/function/src/scalars/json/json_get.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::borrow::Cow; use std::sync::Arc; use arrow::array::{ArrayRef, BinaryViewArray, new_null_array}; @@ -54,12 +55,17 @@ trait JsonGetResultBuilder { fn build(&mut self) -> ArrayRef; } -fn result_builder(len: usize, with_type: &DataType) -> Result> { +fn result_builder( + len: usize, + with_type: &DataType, + is_json2: bool, +) -> Result> { let builder = match with_type { - DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { - Box::new(StringResultBuilder(StringViewBuilder::with_capacity(len))) - as Box - } + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Box::new(StringResultBuilder { + inner: StringViewBuilder::with_capacity(len), + is_json2, + }) + as Box, DataType::Int64 => Box::new(IntResultBuilder(Int64Builder::with_capacity(len))), DataType::Float64 => Box::new(FloatResultBuilder(Float64Builder::with_capacity(len))), DataType::Boolean => Box::new(BoolResultBuilder(BooleanBuilder::with_capacity(len))), @@ -71,20 +77,30 @@ fn result_builder(len: usize, with_type: &DataType) -> Result Result<()> { - self.0.append_option(jsonb::to_str(value).ok()); + // Scalar casts stay unquoted and map JSON null to SQL NULL; only containers + // use `to_string` to preserve their JSON representation. + let value = if self.is_json2 && (jsonb::is_array(value) || jsonb::is_object(value)) { + Some(jsonb::to_string(value)) + } else { + jsonb::to_str(value).ok() + }; + self.inner.append_option(value); Ok(()) } fn append_null(&mut self) { - self.0.append_null(); + self.inner.append_null(); } fn build(&mut self) -> ArrayRef { - Arc::new(self.0.finish()) + Arc::new(self.inner.finish()) } } @@ -404,18 +420,21 @@ impl Function for JsonGetWithType { let result = match arg0.data_type() { DataType::Binary | DataType::LargeBinary | DataType::BinaryView => { let arg0 = compute::cast(&arg0, &DataType::BinaryView)?; + let is_json2 = args.arg_fields.first().is_some_and(is_json2_extension_type); - if args.arg_fields.first().is_some_and(is_json2_extension_type) { - // Query concretization projects nested JSON2 paths as Struct arrays. A binary - // JSON2 argument is therefore an already-selected scalar or root value that - // only needs conversion from its JSONB representation to the requested type. + if is_json2 && path.trim_start_matches('$').split('.').all(str::is_empty) { JsonArray::from(&arg0) .project_to(&with_type) .map_err(|e| exec_datafusion_err!("{e:?}"))? } else { let jsons = arg0.as_binary_view(); - let mut builder = result_builder(len, &with_type)?; - jsonb_get(jsons, path, builder.as_mut())?; + let path = if is_json2 && !path.starts_with('$') { + Cow::Owned(format!("$.{path}")) + } else { + Cow::Borrowed(path) + }; + let mut builder = result_builder(len, &with_type, is_json2)?; + jsonb_get(jsons, &path, builder.as_mut())?; builder.build() } } @@ -511,6 +530,7 @@ mod tests { use datafusion_common::ScalarValue; use datafusion_common::arrow::array::{BinaryArray, BinaryViewArray, StringArray}; use datafusion_common::arrow::datatypes::{Float64Type, Int64Type}; + use datatypes::extension::json::Json2ExtensionType; use datatypes::types::parse_string_to_jsonb; use serde_json::json; @@ -566,6 +586,15 @@ mod tests { )) } + fn test_json_field(json: &ArrayRef, is_json2: bool) -> Arc { + let field = Field::new("json", json.data_type().clone(), true); + Arc::new(if is_json2 { + field.with_extension_type(Json2ExtensionType::default()) + } else { + field + }) + } + #[test] fn test_json_get_int() { let json_get_int = JsonGetInt::default(); @@ -850,7 +879,10 @@ mod tests { ColumnarValue::Array(json.clone()), ColumnarValue::Scalar(path.into()), ], - arg_fields: vec![], + arg_fields: vec![ + test_json_field(json, i >= json_strings.len()), + Arc::new(Field::new("path", DataType::Utf8, false)), + ], number_rows: 1, return_field: Arc::new(Field::new("x", DataType::Utf8View, false)), config_options: Arc::new(Default::default()), @@ -984,7 +1016,11 @@ mod tests { ColumnarValue::Scalar(path.into()), ColumnarValue::Scalar(ScalarValue::Utf8View(None)), ], - arg_fields: vec![], + arg_fields: vec![ + test_json_field(json, i >= json_strings.len()), + Arc::new(Field::new("path", DataType::Utf8, false)), + Arc::new(Field::new("with_type", DataType::Utf8View, true)), + ], number_rows: 1, return_field: Arc::new(Field::new("x", DataType::Utf8View, false)), config_options: Arc::new(Default::default()), diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index e631cf45db4..5178778ba9c 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -348,9 +348,9 @@ impl FlatProjectionMapper { } let field = &self.output_schema.arrow_schema().fields()[output_idx]; - if is_json2_extension_type(field) { + if is_json2_extension_type(field) && array.data_type() != field.data_type() { array = JsonArray::from(&array) - .project_to(field.data_type()) + .project_to_v2(batch.schema_ref().field(*index), field.data_type()) .context(DataTypesSnafu)?; } diff --git a/src/mito2/src/sst/parquet/read_columns.rs b/src/mito2/src/sst/parquet/read_columns.rs index 6f632aee300..68d7cae619d 100644 --- a/src/mito2/src/sst/parquet/read_columns.rs +++ b/src/mito2/src/sst/parquet/read_columns.rs @@ -14,6 +14,7 @@ use std::collections::{HashMap, HashSet}; +use datatypes::extension::json::JSON2_REMAINDER_FIELD_NAME; use parquet::arrow::ProjectionMask; use parquet::basic::{ConvertedType, Type as PhysicalType}; use parquet::schema::types::{ColumnDescriptor, SchemaDescriptor}; @@ -179,7 +180,7 @@ pub struct ProjectionMaskPlan { /// returned plan keeps `k` in the projection mask and marks `j` as /// not present in the output, so it can be synthesized during /// post-processing. -pub fn build_projection_plan( +pub(crate) fn build_projection_plan( parquet_read_cols: &ParquetReadColumns, parquet_schema_desc: &SchemaDescriptor, ) -> ProjectionMaskPlan { @@ -261,13 +262,30 @@ fn build_parquet_leaves_indices( } } - // Then fallback prefix misses to their nearest variant parent. + // Then include v2 remainder leaves or fallback prefix misses to their nearest variant parent. // TODO(fys): Gate fallback planning on the root being JSON2. A raw Binary // leaf is a JSONB variant only under a JSON2 root; plain struct Binary // children should not enter this fallback path. for col in &projection.cols { - for (path_idx, nested_path) in col.nested_paths.iter().enumerate() { - if prefix_matched[&col.root_index][path_idx] { + let path_matches = &prefix_matched[&col.root_index]; + let needs_remainder = col + .nested_paths + .iter() + .zip(path_matches) + .any(|(path, matched)| { + !*matched || path_points_to_struct(parquet_schema_desc, col.root_index, path) + }); + if needs_remainder { + let remainder_leaves = find_remainder_leaves(parquet_schema_desc, col.root_index); + if !remainder_leaves.is_empty() { + matched_leaves.extend(remainder_leaves); + matched_roots.insert(col.root_index); + continue; + } + } + + for (matched, nested_path) in path_matches.iter().zip(&col.nested_paths) { + if *matched { continue; } @@ -287,6 +305,49 @@ fn build_parquet_leaves_indices( (matched_leaves, matched_roots) } +/// Returns whether a nested path points to an explicitly materialized object. +/// +/// JSON2 v2 can split an object's children between its Struct field and the remainder, +/// so reading the Struct leaves alone may produce an incomplete object. +fn path_points_to_struct( + parquet_schema_desc: &SchemaDescriptor, + root_idx: usize, + path: &[String], +) -> bool { + let Some(mut field) = parquet_schema_desc.root_schema().get_fields().get(root_idx) else { + return false; + }; + for name in path.iter().skip(1) { + if !field.is_group() { + return false; + } + let Some(child) = field.get_fields().iter().find(|field| field.name() == name) else { + return false; + }; + field = child; + } + field.is_group() +} + +/// Finds the Parquet leaves backing a JSON2 v2 remainder field. +/// +/// The remainder is a sibling of explicitly materialized fields, so prefix matching a +/// requested path cannot find it. These leaves are needed when an explicit path is absent +/// or an explicitly materialized object may have additional children in the remainder. +fn find_remainder_leaves(parquet_schema_desc: &SchemaDescriptor, root_idx: usize) -> Vec { + parquet_schema_desc + .columns() + .iter() + .enumerate() + .filter_map(|(i, column)| { + let path = column.path().parts(); + (parquet_schema_desc.get_column_root_idx(i) == root_idx + && path.get(1).is_some_and(|x| x == JSON2_REMAINDER_FIELD_NAME)) + .then_some(i) + }) + .collect::>() +} + fn find_nearest_variant_parent( parquet_schema_desc: &SchemaDescriptor, root_idx: usize, @@ -329,6 +390,7 @@ mod tests { use std::sync::Arc; use parquet::basic::{ConvertedType, LogicalType, Repetition}; + use parquet::errors::ParquetError; use parquet::schema::types::Type; use super::*; @@ -441,6 +503,54 @@ mod tests { ); } + #[test] + fn test_v2_routes_missing_path_to_remainder() -> Result<(), ParquetError> { + let parquet = build_test_v2_schema()?; + let projection = + ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths( + vec![ + vec!["j".to_string(), "cold".to_string()], + vec!["j".to_string(), "another".to_string()], + ], + )]); + + let plan = build_projection_plan(&projection, &parquet); + + assert_eq!(vec![true], plan.projected_root_presence); + assert_eq!(ProjectionMask::leaves(&parquet, [0, 1]), plan.mask); + Ok(()) + } + + #[test] + fn test_v2_explicit_path_does_not_read_remainder() -> Result<(), ParquetError> { + let parquet = build_test_v2_schema()?; + let projection = ParquetReadColumns::from_deduped(vec![ + ParquetReadColumn::new(0) + .with_nested_paths(vec![vec!["j".to_string(), "hot".to_string()]]), + ]); + + let plan = build_projection_plan(&projection, &parquet); + + assert_eq!(vec![true], plan.projected_root_presence); + assert_eq!(ProjectionMask::leaves(&parquet, [3]), plan.mask); + Ok(()) + } + + #[test] + fn test_v2_container_path_reads_remainder() -> Result<(), ParquetError> { + let parquet = build_test_v2_schema()?; + let projection = ParquetReadColumns::from_deduped(vec![ + ParquetReadColumn::new(0) + .with_nested_paths(vec![vec!["j".to_string(), "commit".to_string()]]), + ]); + + let plan = build_projection_plan(&projection, &parquet); + + assert_eq!(vec![true], plan.projected_root_presence); + assert_eq!(ProjectionMask::leaves(&parquet, [0, 1, 2]), plan.mask); + Ok(()) + } + #[test] fn test_merges_mixed_paths() { let parquet_schema_desc = build_test_nested_parquet_schema(); @@ -673,6 +783,55 @@ mod tests { SchemaDescriptor::new(schema) } + fn build_test_v2_schema() -> Result { + let metadata = Arc::new( + Type::primitive_type_builder("metadata", parquet::basic::Type::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .build()?, + ); + let value = Arc::new( + Type::primitive_type_builder("value", parquet::basic::Type::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .build()?, + ); + let remainder = Arc::new( + Type::group_type_builder(JSON2_REMAINDER_FIELD_NAME) + .with_repetition(Repetition::OPTIONAL) + .with_logical_type(Some(LogicalType::Variant { + specification_version: None, + })) + .with_fields(vec![metadata, value]) + .build()?, + ); + let operation = Arc::new( + Type::primitive_type_builder("operation", parquet::basic::Type::INT64) + .with_repetition(Repetition::OPTIONAL) + .build()?, + ); + let commit = Arc::new( + Type::group_type_builder("commit") + .with_repetition(Repetition::OPTIONAL) + .with_fields(vec![operation]) + .build()?, + ); + let hot = Arc::new( + Type::primitive_type_builder("hot", parquet::basic::Type::INT64) + .with_repetition(Repetition::OPTIONAL) + .build()?, + ); + let root = Arc::new( + Type::group_type_builder("j") + .with_repetition(Repetition::OPTIONAL) + .with_fields(vec![remainder, commit, hot]) + .build()?, + ); + Ok(SchemaDescriptor::new(Arc::new( + Type::group_type_builder("schema") + .with_fields(vec![root]) + .build()?, + ))) + } + // Test schema: // schema // `- j diff --git a/src/query/src/optimizer/json_type_concretize.rs b/src/query/src/optimizer/json_type_concretize.rs index 4625b619784..c9964e119d0 100644 --- a/src/query/src/optimizer/json_type_concretize.rs +++ b/src/query/src/optimizer/json_type_concretize.rs @@ -111,13 +111,26 @@ fn apply_json_type_hint( fn deduce_json_types(plan: &LogicalPlan) -> Result> { let mut json_types = HashMap::::new(); + // JSON2 columns in the final output must retain their complete values even when + // predicates or other expressions access only specific paths. + // For example, `SELECT j FROM t WHERE json_get(j, 'a') = 1`. + plan.schema() + .fields() + .iter() + .filter(|field| is_json2_extension_type(field)) + .for_each(|field| { + json_types.insert(field.name().clone(), JsonNativeType::Variant); + }); + plan.apply(|plan| { for expr in plan.expressions() { expr.apply(|expr| { if let Some((column, json_type)) = deduce_json_type(expr)? { json_types.entry(column).or_default().merge(&json_type); + Ok(TreeNodeRecursion::Jump) + } else { + Ok(TreeNodeRecursion::Continue) } - Ok(TreeNodeRecursion::Continue) })?; } Ok(TreeNodeRecursion::Continue) @@ -128,6 +141,7 @@ fn deduce_json_types(plan: &LogicalPlan) -> Result Result> { let f = match expr { Expr::ScalarFunction(f) if f.name().eq_ignore_ascii_case(JsonGetWithType::NAME) => f, + Expr::Column(c) => return Ok(Some((c.name.clone(), JsonNativeType::Variant))), _ => return Ok(None), }; @@ -360,25 +374,13 @@ mod tests { .rewrite(plan, &OptimizerContext::default())? .transformed ); - assert!(provider.scan_request().json_type_hint.contains_key("j")); - Ok(()) - } - - #[test] - fn test_allow_json2_passthrough_for_later_projection() -> Result<()> { - let json_get = json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?; - let (provider, plan) = build_json2_scan()?; - let plan = plan - .project(vec![json_get.alias("__common_expr"), col("j")])? - .aggregate(Vec::::new(), vec![count(lit(1))])? - .build()?; - - assert!( - JsonTypeConcretizeRule - .rewrite(plan, &OptimizerContext::default())? - .transformed + assert_eq!( + Some(&JsonNativeType::Object(JsonObjectType::from([( + "a".to_string(), + JsonNativeType::i64(), + )]))), + provider.scan_request().json_type_hint.get("j") ); - assert!(provider.scan_request().json_type_hint.contains_key("j")); Ok(()) } @@ -402,6 +404,25 @@ mod tests { Ok(()) } + #[test] + fn test_allow_json2_filter_with_root_projection() -> Result<()> { + let predicate = + json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?.eq(lit(1_i64)); + let (provider, plan) = build_json2_scan()?; + let plan = plan.filter(predicate)?.build()?; + + assert!( + JsonTypeConcretizeRule + .rewrite(plan, &OptimizerContext::default())? + .transformed + ); + assert_eq!( + Some(&JsonNativeType::Variant), + provider.scan_request().json_type_hint.get("j") + ); + Ok(()) + } + #[test] fn test_deduce_json_type_with_non_column_base() -> Result<()> { let expr = json_get_expr( diff --git a/tests/cases/standalone/common/types/json/json2.result b/tests/cases/standalone/common/types/json/json2.result index 2e9c9c8ee83..6afa0a759ca 100644 --- a/tests/cases/standalone/common/types/json/json2.result +++ b/tests/cases/standalone/common/types/json/json2.result @@ -134,6 +134,31 @@ select j.a, j.a.x from json2_table order by ts; | {"b":10,"x":null} | | +-----------------------------------+-------------------------------------+ +select j, j.a from json2_table order by ts; + ++--------------------------------------------------------------------+-----------------------------------+ +| j | json_get(json2_table.j,Utf8("a")) | ++--------------------------------------------------------------------+-----------------------------------+ +| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | {"b":1} | +| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2,"g":null}}]} | {"b":-2} | +| {"a":{"b":3},"c":"s3","d":null} | {"b":3} | +| {"a":{"b":-4},"c":null,"d":[{"e":{"f":null,"g":-0.4}}]} | {"b":-4} | +| {"a":null,"c":"s5","d":null} | | +| {"a":null,"c":"s6","d":null} | | +| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | {"b":"s7"} | +| {"a":{"b":8},"c":"s8","d":null} | {"b":8} | +| {"a":{"b":null,"x":true},"c":"s9","d":[{"e":{"g":-0.9}}],"y":null} | {"b":null,"x":true} | +| {"a":{"b":10,"x":null},"c":null,"d":null,"y":false} | {"b":10,"x":null} | ++--------------------------------------------------------------------+-----------------------------------+ + +select j from json2_table where j.a.b = 1; + ++-------------------------------------------------------+ +| j | ++-------------------------------------------------------+ +| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | ++-------------------------------------------------------+ + select j.c, j.y from json2_table order by ts; +-----------------------------------+-----------------------------------+ diff --git a/tests/cases/standalone/common/types/json/json2.sql b/tests/cases/standalone/common/types/json/json2.sql index f0bb5286b4c..406cc4af54d 100644 --- a/tests/cases/standalone/common/types/json/json2.sql +++ b/tests/cases/standalone/common/types/json/json2.sql @@ -44,6 +44,10 @@ select j.a.b from json2_table order by ts; select j.a, j.a.x from json2_table order by ts; +select j, j.a from json2_table order by ts; + +select j from json2_table where j.a.b = 1; + select j.c, j.y from json2_table order by ts; select j from json2_table order by ts; diff --git a/tests/cases/standalone/common/types/json/json2_join.result b/tests/cases/standalone/common/types/json/json2_join.result new file mode 100644 index 00000000000..cd5b215a980 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_join.result @@ -0,0 +1,68 @@ +create table json2_join_same_name_left ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +create table json2_join_same_name_right ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +insert into json2_join_same_name_left values + (1, 'a', '{"a": 1, "left_only": "kept"}'); + +Affected Rows: 1 + +insert into json2_join_same_name_right values + (1, 'a', '{"a": "right", "right_only": "should be kept"}'); + +Affected Rows: 1 + +admin flush_table('json2_join_same_name_left'); + ++------------------------------------------------+ +| ADMIN flush_table('json2_join_same_name_left') | ++------------------------------------------------+ +| 0 | ++------------------------------------------------+ + +admin flush_table('json2_join_same_name_right'); + ++-------------------------------------------------+ +| ADMIN flush_table('json2_join_same_name_right') | ++-------------------------------------------------+ +| 0 | ++-------------------------------------------------+ + +-- Conflicting hints for same-named JSON2 columns preserve both values as Variant. +select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 +from json2_join_same_name_left l +join json2_join_same_name_right r +on l.k = r.k; + ++-------------------------+---------------------------------------------------+ +| json_get(r.j,Utf8("a")) | arrow_cast(json_get(l.j,Utf8("a")),Utf8("Int64")) | ++-------------------------+---------------------------------------------------+ +| right | 1 | ++-------------------------+---------------------------------------------------+ + +drop table json2_join_same_name_left; + +Affected Rows: 0 + +drop table json2_join_same_name_right; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_join.sql b/tests/cases/standalone/common/types/json/json2_join.sql new file mode 100644 index 00000000000..3d49d740f62 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_join.sql @@ -0,0 +1,37 @@ +create table json2_join_same_name_left ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +create table json2_join_same_name_right ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +insert into json2_join_same_name_left values + (1, 'a', '{"a": 1, "left_only": "kept"}'); + +insert into json2_join_same_name_right values + (1, 'a', '{"a": "right", "right_only": "should be kept"}'); + +admin flush_table('json2_join_same_name_left'); + +admin flush_table('json2_join_same_name_right'); + +-- Conflicting hints for same-named JSON2 columns preserve both values as Variant. +select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 +from json2_join_same_name_left l +join json2_join_same_name_right r +on l.k = r.k; + +drop table json2_join_same_name_left; + +drop table json2_join_same_name_right; diff --git a/tests/cases/standalone/common/types/json/json2_limit.result b/tests/cases/standalone/common/types/json/json2_limit.result index 035c4ed07bf..0394bfca66e 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.result +++ b/tests/cases/standalone/common/types/json/json2_limit.result @@ -65,87 +65,10 @@ order by json_get(j, 'a.b'); | 2 | 1 | +---------------------------------------------------+----------+ -select j, j.a from json2_whole_and_path_read; - -Error: 3001(EngineExecuteQuery), Invalid argument error: column types must match schema types, expected Binary but found Struct("a": Utf8View) at column index 0 - -select j from json2_whole_and_path_read where j.a.b = 1; - -Error: 3001(EngineExecuteQuery), Invalid argument error: column types must match schema types, expected Binary but found Struct("a": Struct("b": Int64)) at column index 0 - drop table json2_whole_and_path_read; Affected Rows: 0 -create table json2_join_same_name_left ( - ts timestamp time index, - k string, - j json2 -) -with ( - 'append_mode' = 'true' -); - -Affected Rows: 0 - -create table json2_join_same_name_right ( - ts timestamp time index, - k string, - j json2 -) -with ( - 'append_mode' = 'true' -); - -Affected Rows: 0 - -insert into json2_join_same_name_left values - (1, 'a', '{"a": 1, "left_only": "kept"}'); - -Affected Rows: 1 - -insert into json2_join_same_name_right values - (1, 'a', '{"a": "right", "right_only": "should be kept"}'); - -Affected Rows: 1 - -admin flush_table('json2_join_same_name_left'); - -+------------------------------------------------+ -| ADMIN flush_table('json2_join_same_name_left') | -+------------------------------------------------+ -| 0 | -+------------------------------------------------+ - -admin flush_table('json2_join_same_name_right'); - -+-------------------------------------------------+ -| ADMIN flush_table('json2_join_same_name_right') | -+-------------------------------------------------+ -| 0 | -+-------------------------------------------------+ - --- FIXME: This should return `right` and `1`. The current NULL values are caused --- by JSON type hints losing the table qualifier in joins. -select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 -from json2_join_same_name_left l -join json2_join_same_name_right r -on l.k = r.k; - -+---------------------------------------------+---------------------------------------------------+ -| json_get(r.j,Utf8("a")) | arrow_cast(json_get(l.j,Utf8("a")),Utf8("Int64")) | -+---------------------------------------------+---------------------------------------------------+ -| {"a":"right","right_only":"should be kept"} | | -+---------------------------------------------+---------------------------------------------------+ - -drop table json2_join_same_name_left; - -Affected Rows: 0 - -drop table json2_join_same_name_right; - -Affected Rows: 0 - create table json2_without_append_mode ( ts timestamp time index, j json2 diff --git a/tests/cases/standalone/common/types/json/json2_limit.sql b/tests/cases/standalone/common/types/json/json2_limit.sql index 986ed58eb2a..fc80e97cd03 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.sql +++ b/tests/cases/standalone/common/types/json/json2_limit.sql @@ -38,51 +38,8 @@ from json2_whole_and_path_read group by json_get(j, 'a.b') order by json_get(j, 'a.b'); -select j, j.a from json2_whole_and_path_read; - -select j from json2_whole_and_path_read where j.a.b = 1; - drop table json2_whole_and_path_read; -create table json2_join_same_name_left ( - ts timestamp time index, - k string, - j json2 -) -with ( - 'append_mode' = 'true' -); - -create table json2_join_same_name_right ( - ts timestamp time index, - k string, - j json2 -) -with ( - 'append_mode' = 'true' -); - -insert into json2_join_same_name_left values - (1, 'a', '{"a": 1, "left_only": "kept"}'); - -insert into json2_join_same_name_right values - (1, 'a', '{"a": "right", "right_only": "should be kept"}'); - -admin flush_table('json2_join_same_name_left'); - -admin flush_table('json2_join_same_name_right'); - --- FIXME: This should return `right` and `1`. The current NULL values are caused --- by JSON type hints losing the table qualifier in joins. -select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 -from json2_join_same_name_left l -join json2_join_same_name_right r -on l.k = r.k; - -drop table json2_join_same_name_left; - -drop table json2_join_same_name_right; - create table json2_without_append_mode ( ts timestamp time index, j json2 From 73c4140938f4035c0782e18f6e766558041b9c63 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:31:47 +0000 Subject: [PATCH 106/126] feat(function): expose uddsketch rank (#8929) * feat(function): expose uddsketch rank Signed-off-by: Lei, HUANG * test(function): cover uddsketch rank in sqlness Signed-off-by: Lei, HUANG * fix(function): support legacy uddsketch rank Signed-off-by: Lei, HUANG * docs(function): document uddsketch rank behavior Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/common/function/src/function_registry.rs | 7 + src/common/function/src/scalars.rs | 1 + .../function/src/scalars/uddsketch_rank.rs | 224 ++++++++++++ src/common/function/src/uddsketch_compat.rs | 338 ++++++++++++++++-- .../common/aggregate/uddsketch.result | 36 ++ .../standalone/common/aggregate/uddsketch.sql | 18 + .../cases/uddsketch_legacy_state/case.toml | 2 +- .../uddsketch_legacy_state/verify.result | 22 +- .../cases/uddsketch_legacy_state/verify.sql | 14 +- 9 files changed, 633 insertions(+), 29 deletions(-) create mode 100644 src/common/function/src/scalars/uddsketch_rank.rs diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs index ef2eb3658d4..15721ccabc6 100644 --- a/src/common/function/src/function_registry.rs +++ b/src/common/function/src/function_registry.rs @@ -41,6 +41,7 @@ use crate::scalars::primary_key::DecodePrimaryKeyFunction; use crate::scalars::string::register_string_functions; use crate::scalars::timestamp::TimestampFunction; use crate::scalars::uddsketch_calc::UddSketchCalcFunction; +use crate::scalars::uddsketch_rank::UddSketchRankFunction; use crate::scalars::vector::VectorFunction as VectorScalarFunction; use crate::system::SystemFunction; @@ -209,6 +210,7 @@ pub static FUNCTION_REGISTRY: LazyLock> = LazyLock::new(|| DateFunction::register(&function_registry); ExpressionFunction::register(&function_registry); UddSketchCalcFunction::register(&function_registry); + UddSketchRankFunction::register(&function_registry); HllCalcFunction::register(&function_registry); DecodePrimaryKeyFunction::register(&function_registry); @@ -363,6 +365,11 @@ mod tests { assert_eq!(1, registry.scalar_functions().len()); } + #[test] + fn test_uddsketch_rank_registered() { + assert!(FUNCTION_REGISTRY.get_function("uddsketch_rank").is_some()); + } + #[test] fn test_register_if_absent_registers_new_function() { let registry = FunctionRegistry::default(); diff --git a/src/common/function/src/scalars.rs b/src/common/function/src/scalars.rs index 0c089dedbf1..75dec84ea62 100644 --- a/src/common/function/src/scalars.rs +++ b/src/common/function/src/scalars.rs @@ -31,4 +31,5 @@ pub mod ip; pub(crate) mod test; pub(crate) mod timestamp; pub(crate) mod uddsketch_calc; +pub(crate) mod uddsketch_rank; pub mod udf; diff --git a/src/common/function/src/scalars/uddsketch_rank.rs b/src/common/function/src/scalars/uddsketch_rank.rs new file mode 100644 index 00000000000..d746b212ed7 --- /dev/null +++ b/src/common/function/src/scalars/uddsketch_rank.rs @@ -0,0 +1,224 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Implementation of the scalar function `uddsketch_rank`. + +use std::fmt; +use std::fmt::Display; +use std::sync::Arc; + +use datafusion_common::DataFusionError; +use datafusion_common::arrow::array::{Array, AsArray, Float64Builder}; +use datafusion_common::arrow::datatypes::{DataType, Float64Type}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility}; + +use crate::function::{Function, extract_args}; +use crate::function_registry::FunctionRegistry; + +const NAME: &str = "uddsketch_rank"; + +/// Implements the scalar function `uddsketch_rank`. +/// +/// It accepts a value and a serialized UDDSketch state, then estimates the +/// value's quantile rank by counting half of the matching bucket. Both current +/// and legacy state encodings are supported. Null arguments, empty or invalid +/// states, and invalid values produce null results. +#[derive(Debug)] +pub(crate) struct UddSketchRankFunction { + signature: Signature, +} + +impl UddSketchRankFunction { + pub fn register(registry: &FunctionRegistry) { + registry.register_scalar(UddSketchRankFunction::default()); + } +} + +impl Default for UddSketchRankFunction { + fn default() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Float64, DataType::Binary], + Volatility::Immutable, + ), + } + } +} + +impl Display for UddSketchRankFunction { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{}", NAME.to_ascii_uppercase()) + } +} + +impl Function for UddSketchRankFunction { + fn name(&self) -> &str { + NAME + } + + fn return_type(&self, _: &[DataType]) -> datafusion_common::Result { + Ok(DataType::Float64) + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + let [arg0, arg1] = extract_args(self.name(), &args)?; + + let Some(values) = arg0.as_primitive_opt::() else { + return Err(DataFusionError::Execution(format!( + "'{}' expects 1st argument to be Float64 datatype, got {}", + self.name(), + arg0.data_type() + ))); + }; + let Some(sketches) = arg1.as_binary_opt::() else { + return Err(DataFusionError::Execution(format!( + "'{}' expects 2nd argument to be Binary datatype, got {}", + self.name(), + arg1.data_type() + ))); + }; + let mut builder = Float64Builder::with_capacity(sketches.len()); + + for i in 0..sketches.len() { + if values.is_null(i) || sketches.is_null(i) { + builder.append_null(); + continue; + } + + match crate::uddsketch_compat::rank(sketches.value(i), values.value(i)) { + Ok(Some(rank)) => builder.append_value(rank), + Ok(None) => builder.append_null(), + Err(error) => { + common_telemetry::trace!("Failed to calculate UDDSketch rank: {}", error); + builder.append_null(); + } + } + } + + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_schema::Field; + use datafusion_common::arrow::array::{Array, AsArray, BinaryArray, Float64Array}; + use datafusion_common::arrow::datatypes::{DataType, Float64Type}; + use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; + use uddsketch::UddSketch; + + use super::UddSketchRankFunction; + use crate::function::Function; + use crate::uddsketch_compat; + + fn invoke( + function: &UddSketchRankFunction, + values: Float64Array, + states: BinaryArray, + ) -> Float64Array { + let number_rows = values.len(); + let result = function + .invoke_with_args(ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(values)), + ColumnarValue::Array(Arc::new(states)), + ], + arg_fields: vec![], + number_rows, + return_field: Arc::new(Field::new("x", DataType::Float64, true)), + config_options: Arc::new(Default::default()), + }) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + result.as_primitive::().clone() + } + + #[test] + fn test_uddsketch_rank_function_name_and_return_type() { + let function = UddSketchRankFunction::default(); + + assert_eq!("uddsketch_rank", function.name()); + assert_eq!( + DataType::Float64, + function + .return_type(&[DataType::Float64, DataType::Binary]) + .unwrap() + ); + } + + #[test] + fn test_uddsketch_rank_function_canonical_state() { + let function = UddSketchRankFunction::default(); + let mut sketch = UddSketch::new(128, 0.01).unwrap(); + sketch + .add_batch(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0]) + .unwrap(); + let values = [5.0, 10.0, 55.0, 100.0, 110.0]; + let expected = values.map(|value| sketch.rank(value).unwrap().unwrap()); + let encoded = sketch.encode().unwrap(); + let states = BinaryArray::from_iter_values((0..values.len()).map(|_| encoded.as_slice())); + + let result = invoke(&function, Float64Array::from(values.to_vec()), states); + + assert_eq!(result.values(), &expected); + } + + #[test] + fn test_uddsketch_rank_function_reads_legacy_state() { + let function = UddSketchRankFunction::default(); + let values = [f64::NEG_INFINITY, 0.0, 0.99, f64::INFINITY]; + let states = BinaryArray::from_iter_values( + (0..values.len()).map(|_| uddsketch_compat::LEGACY_STATE), + ); + + let result = invoke(&function, Float64Array::from(values.to_vec()), states); + + assert_eq!(result.null_count(), 0); + assert_eq!(result.values(), &[0.0, 0.375, 0.625, 1.0]); + } + + #[test] + fn test_uddsketch_rank_function_returns_null_for_invalid_rows() { + let function = UddSketchRankFunction::default(); + let empty = UddSketch::new(128, 0.01).unwrap().encode().unwrap(); + let mut populated = UddSketch::new(128, 0.01).unwrap(); + populated.add(1.0).unwrap(); + let populated = populated.encode().unwrap(); + let malformed = [1, 2, 3]; + let states = BinaryArray::from_iter([ + Some(populated.as_slice()), + None, + Some(empty.as_slice()), + Some(malformed.as_slice()), + Some(populated.as_slice()), + ]); + let values = + Float64Array::from(vec![None, Some(1.0), Some(1.0), Some(1.0), Some(f64::NAN)]); + + let result = invoke(&function, values, states); + + assert_eq!(result.null_count(), 5); + } +} diff --git a/src/common/function/src/uddsketch_compat.rs b/src/common/function/src/uddsketch_compat.rs index 733b4ab043a..3b14f3b91d5 100644 --- a/src/common/function/src/uddsketch_compat.rs +++ b/src/common/function/src/uddsketch_compat.rs @@ -61,6 +61,13 @@ struct LegacyUddSketchState { initial_error: f64, } +struct ValidatedLegacyUddSketch { + buckets: Vec<(LegacyBucketKey, u64)>, + alpha: f64, + gamma: f64, + count: u64, +} + pub(crate) fn decode(raw: &[u8]) -> Result { match UddSketch::decode(raw) { Ok(sketch) => Ok(sketch), @@ -85,6 +92,19 @@ pub(crate) fn quantile(raw: &[u8], quantile: f64) -> Result, String> } } +pub(crate) fn rank(raw: &[u8], value: f64) -> Result, String> { + match UddSketchRef::parse(raw) { + Ok(sketch) => sketch.rank(value).map_err(|error| error.to_string()), + Err(current_error) => decode_legacy_sketch(raw) + .and_then(|sketch| sketch.rank(value)) + .map_err(|legacy_error| { + format!( + "canonical decode failed: {current_error}; legacy decode failed: {legacy_error}" + ) + }), + } +} + fn validate_legacy_input(raw: &[u8]) -> Result<(), String> { if raw.len() > MAX_BYTES { return Err("input exceeds the legacy decode byte limit".to_string()); @@ -185,6 +205,69 @@ impl LegacyUddSketch { if !quantile.is_finite() || !(0.0..=1.0).contains(&quantile) { return Err("invalid quantile".to_string()); } + let validated = self.validate()?; + let count = validated.count; + if count == 0 { + return Ok(None); + } + + let target = if quantile == 1.0 { + count + } else { + ((count as f64 * quantile) as u64) + .saturating_add(1) + .min(count) + }; + let mut seen = 0_u64; + for (key, bucket_count) in validated.buckets { + seen += bucket_count; + if seen >= target { + return Ok(Some(legacy_bucket_value( + validated.alpha, + validated.gamma, + key, + )?)); + } + } + Err("legacy bucket counts do not cover the quantile rank".to_string()) + } + + fn rank(self, value: f64) -> Result, String> { + let validated = self.validate()?; + let count = validated.count; + if count == 0 { + return Ok(None); + } + if value.is_nan() { + return Err("invalid rank value".to_string()); + } + + let index = if value.is_infinite() { + i64::MAX + } else { + value.abs().log(validated.gamma).ceil() as i64 + }; + let target = if value == 0.0 { + LegacyBucketKey::Zero + } else if value.is_sign_negative() { + LegacyBucketKey::Negative(index) + } else { + LegacyBucketKey::Positive(index) + }; + let mut below = 0.0; + for (key, bucket_count) in validated.buckets { + if key == target { + return Ok(Some((below + bucket_count as f64 / 2.0) / count as f64)); + } + if key_lt(target, key) { + return Ok(Some(below / count as f64)); + } + below += bucket_count as f64; + } + Ok(Some(1.0)) + } + + fn validate(self) -> Result { if self.compactions >= 64 { return Err("legacy compaction count must be below 64".to_string()); } @@ -193,9 +276,6 @@ impl LegacyUddSketch { return Err("legacy maximum bucket count is outside supported limits".to_string()); } - let count = self.count; - let alpha = self.alpha; - let gamma = self.gamma; let buckets = self.buckets.ordered()?; if buckets.len() > self.max_buckets as usize { return Err("legacy populated bucket count exceeds its configured limit".to_string()); @@ -210,28 +290,16 @@ impl LegacyUddSketch { .checked_add(*count) .ok_or_else(|| "legacy bucket count sum overflows u64".to_string()) })?; - if decoded_count != count { + if decoded_count != self.count { return Err("legacy bucket counts do not match the value count".to_string()); } - if count == 0 { - return Ok(None); - } - let target = if quantile == 1.0 { - count - } else { - ((count as f64 * quantile) as u64) - .saturating_add(1) - .min(count) - }; - let mut seen = 0_u64; - for (key, bucket_count) in buckets { - seen += bucket_count; - if seen >= target { - return Ok(Some(legacy_bucket_value(alpha, gamma, key)?)); - } - } - Err("legacy bucket counts do not cover the quantile rank".to_string()) + Ok(ValidatedLegacyUddSketch { + buckets, + alpha: self.alpha, + gamma: self.gamma, + count: self.count, + }) } } @@ -471,6 +539,232 @@ pub(crate) const COMPACTED_LEGACY_SKETCH: &[u8] = &[ mod tests { use super::*; + fn legacy_rank_sketch() -> LegacyUddSketch { + LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([ + ( + LegacyBucketKey::Negative(2), + LegacyBucket { + count: 2, + next: LegacyBucketKey::Zero, + }, + ), + ( + LegacyBucketKey::Zero, + LegacyBucket { + count: 1, + next: LegacyBucketKey::Positive(2), + }, + ), + ( + LegacyBucketKey::Positive(2), + LegacyBucket { + count: 1, + next: LegacyBucketKey::Invalid, + }, + ), + ]), + head: LegacyBucketKey::Negative(2), + }, + alpha: 0.01, + gamma: (1.0 + 0.01) / (1.0 - 0.01), + compactions: 0, + max_buckets: 128, + count: 4, + sum: 0.0, + } + } + + #[test] + fn legacy_rank_matches_canonical_rank() { + let mut sketch = UddSketch::new(128, 0.01).unwrap(); + sketch + .add_batch(&[-10.0, -1.0, 0.0, 1.0, 10.0, f64::INFINITY]) + .unwrap(); + let encoded = sketch.encode().unwrap(); + + for value in [ + f64::NEG_INFINITY, + -10.0, + -0.0, + 0.0, + 1.0, + 10.0, + f64::INFINITY, + ] { + assert_eq!(rank(&encoded, value).unwrap(), sketch.rank(value).unwrap()); + } + } + + #[test] + fn legacy_rank_reads_wrapped_and_bare_legacy_states() { + let bare = &LEGACY_STATE[..LEGACY_STATE.len() - std::mem::size_of::()]; + + for value in [f64::NEG_INFINITY, 0.0, 0.99, f64::INFINITY] { + assert_eq!( + rank(LEGACY_STATE, value).unwrap(), + rank(bare, value).unwrap() + ); + } + assert_eq!(rank(LEGACY_STATE, f64::NEG_INFINITY).unwrap(), Some(0.0)); + assert_eq!(rank(LEGACY_STATE, f64::INFINITY).unwrap(), Some(1.0)); + } + + #[test] + fn legacy_rank_reads_compacted_bare_sketch() { + assert_eq!( + rank(COMPACTED_LEGACY_SKETCH, f64::NEG_INFINITY).unwrap(), + Some(0.0) + ); + assert!(rank(COMPACTED_LEGACY_SKETCH, 1.0).unwrap().is_some()); + assert_eq!( + rank(COMPACTED_LEGACY_SKETCH, f64::INFINITY).unwrap(), + Some(1.0) + ); + } + + #[test] + fn legacy_rank_handles_bucket_boundaries_and_empty_sketch() { + let gamma = legacy_rank_sketch().gamma; + for (value, expected) in [ + (f64::NEG_INFINITY, 0.0), + (-1.0, 0.5), + (-0.0, 0.625), + (0.0, 0.625), + (1.0, 0.75), + (gamma.powi(2), 0.875), + (f64::INFINITY, 1.0), + ] { + assert_eq!(legacy_rank_sketch().rank(value).unwrap(), Some(expected)); + } + + let mut empty = legacy_rank_sketch(); + empty.buckets.map.clear(); + empty.buckets.head = LegacyBucketKey::Invalid; + empty.count = 0; + assert_eq!(empty.rank(0.0).unwrap(), None); + } + + #[test] + fn legacy_rank_empty_sketch_returns_none_for_nan() { + let mut empty = legacy_rank_sketch(); + empty.buckets.map.clear(); + empty.buckets.head = LegacyBucketKey::Invalid; + empty.count = 0; + + assert_eq!(empty.rank(f64::NAN).unwrap(), None); + } + + #[test] + fn legacy_rank_accumulates_large_counts_with_canonical_rounding() { + let large_count = 1_u64 << 53; + let total_count = large_count + 3; + let sketch = LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([ + ( + LegacyBucketKey::Negative(3), + LegacyBucket { + count: large_count, + next: LegacyBucketKey::Negative(2), + }, + ), + ( + LegacyBucketKey::Negative(2), + LegacyBucket { + count: 1, + next: LegacyBucketKey::Negative(1), + }, + ), + ( + LegacyBucketKey::Negative(1), + LegacyBucket { + count: 1, + next: LegacyBucketKey::Zero, + }, + ), + ( + LegacyBucketKey::Zero, + LegacyBucket { + count: 1, + next: LegacyBucketKey::Invalid, + }, + ), + ]), + head: LegacyBucketKey::Negative(3), + }, + alpha: 0.01, + gamma: (1.0 + 0.01) / (1.0 - 0.01), + compactions: 0, + max_buckets: 128, + count: total_count, + sum: 0.0, + }; + let mut below = 0.0; + for count in [large_count, 1, 1] { + below += count as f64; + } + let expected = (below + 0.5) / total_count as f64; + + assert_eq!(sketch.rank(0.0).unwrap(), Some(expected)); + } + + #[test] + fn legacy_rank_rejects_nan_and_malformed_data() { + let bare_len = LEGACY_STATE.len() - std::mem::size_of::(); + let bare = &LEGACY_STATE[..bare_len]; + assert!(rank(bare, f64::NAN).is_err()); + + let mut invalid_mapping = bare.to_vec(); + let alpha_offset = bare_len - 44; + invalid_mapping[alpha_offset..alpha_offset + 8].copy_from_slice(&f64::NAN.to_le_bytes()); + let error = rank(&invalid_mapping, 0.0).unwrap_err(); + assert!(error.contains("canonical decode failed")); + assert!(error.contains("legacy decode failed")); + + let mut invalid_count = legacy_rank_sketch(); + invalid_count.count += 1; + assert!(invalid_count.rank(0.0).is_err()); + + let mut invalid_bucket = legacy_rank_sketch(); + let bucket = invalid_bucket + .buckets + .map + .remove(&LegacyBucketKey::Negative(2)) + .unwrap(); + invalid_bucket + .buckets + .map + .insert(LegacyBucketKey::Negative(i64::MIN), bucket); + invalid_bucket.buckets.head = LegacyBucketKey::Negative(i64::MIN); + assert!(invalid_bucket.rank(0.0).is_err()); + } + + #[test] + fn legacy_rank_accepts_saturated_mapping() { + let sketch = LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([( + LegacyBucketKey::Positive(0), + LegacyBucket { + count: 1, + next: LegacyBucketKey::Invalid, + }, + )]), + head: LegacyBucketKey::Positive(0), + }, + alpha: 1.0, + gamma: f64::INFINITY, + compactions: 63, + max_buckets: 7, + count: 1, + sum: f64::INFINITY, + }; + + assert_eq!(sketch.rank(1.0).unwrap(), Some(0.5)); + } + #[test] fn current_format_is_decoded_without_legacy_fallback() { let mut sketch = UddSketch::new(128, 0.01).unwrap(); diff --git a/tests/cases/standalone/common/aggregate/uddsketch.result b/tests/cases/standalone/common/aggregate/uddsketch.result index 6304ea93625..41f459a36af 100644 --- a/tests/cases/standalone/common/aggregate/uddsketch.result +++ b/tests/cases/standalone/common/aggregate/uddsketch.result @@ -52,6 +52,23 @@ select uddsketch_calc(0.95, uddsketch_state(128, 0.01, `value`)) from test_uddsk | 100.49456770856492 | +----------------------------------------------------------------------------------------------+ +SELECT + ROUND(uddsketch_rank(5, `state`), 4) AS below_min_rank, + ROUND(uddsketch_rank(10, `state`), 4) AS equal_min_rank, + ROUND(uddsketch_rank(55, `state`), 4) AS middle_rank, + ROUND(uddsketch_rank(100, `state`), 4) AS equal_max_rank, + ROUND(uddsketch_rank(110, `state`), 4) AS above_max_rank +FROM ( + SELECT uddsketch_state(128, 0.01, `value`) AS `state` + FROM test_uddsketch +); + ++----------------+----------------+-------------+----------------+----------------+ +| below_min_rank | equal_min_rank | middle_rank | equal_max_rank | above_max_rank | ++----------------+----------------+-------------+----------------+----------------+ +| 0.0 | 0.05 | 0.5 | 0.95 | 1.0 | ++----------------+----------------+-------------+----------------+----------------+ + CREATE TABLE grouped_uddsketch ( `state` BINARY, id_group INT PRIMARY KEY, @@ -72,6 +89,15 @@ SELECT uddsketch_calc(0.1, uddsketch_merge(128, 0.01, `state`)) FROM grouped_udd | 19.886670240866184 | +------------------------------------------------------------------------------------------------+ +SELECT ROUND(uddsketch_rank(55, uddsketch_merge(128, 0.01, `state`)), 4) AS merged_rank +FROM grouped_uddsketch; + ++-------------+ +| merged_rank | ++-------------+ +| 0.5 | ++-------------+ + -- should fail SELECT uddsketch_calc(0.1, uddsketch_merge(128, 0.1, `state`)) FROM grouped_uddsketch; @@ -82,6 +108,16 @@ SELECT uddsketch_calc(0.1, uddsketch_merge(64, 0.01, `state`)) FROM grouped_udds Error: 3001(EngineExecuteQuery), Error during planning: Merging UDDSketch with different parameters: arguments=(64, 0.01) vs actual input=(128, 0.01) +SELECT uddsketch_rank(55, uddsketch_state(128, 0.01, `value`)) IS NULL AS empty_rank_is_null +FROM test_uddsketch +WHERE `id` < 0; + ++--------------------+ +| empty_rank_is_null | ++--------------------+ +| true | ++--------------------+ + drop table test_uddsketch; Affected Rows: 0 diff --git a/tests/cases/standalone/common/aggregate/uddsketch.sql b/tests/cases/standalone/common/aggregate/uddsketch.sql index c959f33ee99..13deb88ddb1 100644 --- a/tests/cases/standalone/common/aggregate/uddsketch.sql +++ b/tests/cases/standalone/common/aggregate/uddsketch.sql @@ -24,6 +24,17 @@ select uddsketch_calc(0.75, uddsketch_state(128, 0.01, `value`)) from test_uddsk select uddsketch_calc(0.95, uddsketch_state(128, 0.01, `value`)) from test_uddsketch; +SELECT + ROUND(uddsketch_rank(5, `state`), 4) AS below_min_rank, + ROUND(uddsketch_rank(10, `state`), 4) AS equal_min_rank, + ROUND(uddsketch_rank(55, `state`), 4) AS middle_rank, + ROUND(uddsketch_rank(100, `state`), 4) AS equal_max_rank, + ROUND(uddsketch_rank(110, `state`), 4) AS above_max_rank +FROM ( + SELECT uddsketch_state(128, 0.01, `value`) AS `state` + FROM test_uddsketch +); + CREATE TABLE grouped_uddsketch ( `state` BINARY, id_group INT PRIMARY KEY, @@ -34,12 +45,19 @@ INSERT INTO grouped_uddsketch (`state`, id_group) SELECT uddsketch_state(128, 0. SELECT uddsketch_calc(0.1, uddsketch_merge(128, 0.01, `state`)) FROM grouped_uddsketch; +SELECT ROUND(uddsketch_rank(55, uddsketch_merge(128, 0.01, `state`)), 4) AS merged_rank +FROM grouped_uddsketch; + -- should fail SELECT uddsketch_calc(0.1, uddsketch_merge(128, 0.1, `state`)) FROM grouped_uddsketch; -- should fail SELECT uddsketch_calc(0.1, uddsketch_merge(64, 0.01, `state`)) FROM grouped_uddsketch; +SELECT uddsketch_rank(55, uddsketch_state(128, 0.01, `value`)) IS NULL AS empty_rank_is_null +FROM test_uddsketch +WHERE `id` < 0; + drop table test_uddsketch; drop table grouped_uddsketch; diff --git a/tests/compatibility/cases/uddsketch_legacy_state/case.toml b/tests/compatibility/cases/uddsketch_legacy_state/case.toml index b22670037d1..e68fdf3804b 100644 --- a/tests/compatibility/cases/uddsketch_legacy_state/case.toml +++ b/tests/compatibility/cases/uddsketch_legacy_state/case.toml @@ -1,5 +1,5 @@ name = "uddsketch_legacy_state" -reason = "Verify UDDSketch state bytes persisted by old binaries (wrapped legacy bincode UddSketchState) are read directly by uddsketch_calc and uddsketch_merge on the new binary, merged into canonical v1, and merged together with newly generated canonical states." +reason = "Verify UDDSketch state bytes persisted by old binaries (wrapped legacy bincode UddSketchState) are read directly by uddsketch_calc, uddsketch_rank, and uddsketch_merge on the new binary, merged into canonical v1, and merged together with newly generated canonical states." introduced_by = "PR #8867" topologies = ["distributed", "standalone"] from_range = [">=v1.0.0"] diff --git a/tests/compatibility/cases/uddsketch_legacy_state/verify.result b/tests/compatibility/cases/uddsketch_legacy_state/verify.result index b4c8a3d1566..d1a40d5a15b 100644 --- a/tests/compatibility/cases/uddsketch_legacy_state/verify.result +++ b/tests/compatibility/cases/uddsketch_legacy_state/verify.result @@ -11,7 +11,23 @@ ORDER BY grp; | 10 | 100.49456770856492 | +-----+--------------------+ --- (2) uddsketch_merge over the persisted legacy rows, then calc. +-- (2) direct uddsketch_rank of each persisted legacy row +SELECT grp, + ROUND(uddsketch_rank(30.0, state), 4) AS rank_at_30, + ROUND(uddsketch_rank(70.0, state), 4) AS rank_at_70, + ROUND(uddsketch_rank(100.0, state), 4) AS rank_at_100 +FROM uddsketch_states +ORDER BY grp; + ++-----+------------+------------+-------------+ +| grp | rank_at_30 | rank_at_70 | rank_at_100 | ++-----+------------+------------+-------------+ +| 0 | 0.625 | 1.0 | 1.0 | +| 5 | 0.0 | 0.5 | 1.0 | +| 10 | 0.0 | 0.0 | 0.5 | ++-----+------------+------------+-------------+ + +-- (3) uddsketch_merge over the persisted legacy rows, then calc. -- Succeeds only if every legacy row decodes to the exact declared -- parameters (128, 0.01); the merge parameter check is the guard. SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 @@ -23,7 +39,7 @@ FROM uddsketch_states; | 59.745049810145126 | +--------------------+ --- (3) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows +-- (4) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows CREATE TABLE uddsketch_new_states ( state BINARY, grp INT PRIMARY KEY, @@ -52,7 +68,7 @@ FROM ( | 59.745049810145126 | +--------------------+ --- (4) canonical-only control +-- (5) canonical-only control SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 FROM uddsketch_new_states; diff --git a/tests/compatibility/cases/uddsketch_legacy_state/verify.sql b/tests/compatibility/cases/uddsketch_legacy_state/verify.sql index 8cf84cfbbc5..dd3188f6712 100644 --- a/tests/compatibility/cases/uddsketch_legacy_state/verify.sql +++ b/tests/compatibility/cases/uddsketch_legacy_state/verify.sql @@ -3,13 +3,21 @@ SELECT grp, uddsketch_calc(0.5, state) AS p50 FROM uddsketch_states ORDER BY grp; --- (2) uddsketch_merge over the persisted legacy rows, then calc. +-- (2) direct uddsketch_rank of each persisted legacy row +SELECT grp, + ROUND(uddsketch_rank(30.0, state), 4) AS rank_at_30, + ROUND(uddsketch_rank(70.0, state), 4) AS rank_at_70, + ROUND(uddsketch_rank(100.0, state), 4) AS rank_at_100 +FROM uddsketch_states +ORDER BY grp; + +-- (3) uddsketch_merge over the persisted legacy rows, then calc. -- Succeeds only if every legacy row decodes to the exact declared -- parameters (128, 0.01); the merge parameter check is the guard. SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 FROM uddsketch_states; --- (3) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows +-- (4) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows CREATE TABLE uddsketch_new_states ( state BINARY, grp INT PRIMARY KEY, @@ -28,6 +36,6 @@ FROM ( SELECT state FROM uddsketch_new_states ) AS all_states; --- (4) canonical-only control +-- (5) canonical-only control SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 FROM uddsketch_new_states; From 1851f6bf4de7b6ae94fc9b0cb3048aa24877aea9 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 25 Aug 2026 12:34:53 +0000 Subject: [PATCH 107/126] fix(query): keep INSERT timestamp conversion out of the source query (#8911) * fix(query): keep INSERT timestamp conversion out of the source query Interpreting an INSERT's string timestamps used to work by pushing the conversion down into the source query, which changed what that query means. Two consequences: - Pushing through a UNION's DISTINCT moved the dedup key from the raw strings to parsed instants, so rows spelling the same instant differently collapsed into one. On an append-only table that is a silently dropped row. - A UNION branch that needed no conversion (a NULL, or an explicit cast) made the whole column give up, leaving sibling branches on UTC while the rest of the row used the session timezone. Convert at the assignment instead, by routing its cast through a timezone-carrying timestamp type and back. Arrow applies the timezone when a cast target carries one, and stripping it afterwards preserves the value. The source query is no longer touched, so both cases go away and the tree-walking rewrite (roughly 160 lines) is deleted. The rewrite reads source types, so it now runs TypeCoercion first: a UNION still carries its loose per-branch schema before coercion, and retargeting a cast whose input later becomes a timestamp would shift the value rather than reinterpret it. Signed-off-by: Dennis Zhuang * fix(query): address review on INSERT assignment rewrite - Clone the input `Arc` instead of the whole subtree, and only rebuild it when a `Values` row actually changes. - Defer cloning the cast source until the literal-folding path has been ruled out. - Move the UTC check onto `Timezone::is_utc`, replacing a bare string compare. - Cover a prepared `INSERT ... VALUES (?)`: an untyped placeholder types as `Null`, so the assignment cast is left for parameter substitution. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/time/src/timezone.rs | 9 + src/query/src/optimizer.rs | 1 + src/query/src/optimizer/insert_assignment.rs | 307 ++++++++++++++++ src/query/src/optimizer/type_conversion.rs | 11 +- .../type_conversion/insert_assignment.rs | 336 ------------------ src/query/src/planner.rs | 126 +++---- .../insert/insert_default_timezone.result | 58 +++ .../common/insert/insert_default_timezone.sql | 28 ++ 8 files changed, 472 insertions(+), 404 deletions(-) create mode 100644 src/query/src/optimizer/insert_assignment.rs delete mode 100644 src/query/src/optimizer/type_conversion/insert_assignment.rs diff --git a/src/common/time/src/timezone.rs b/src/common/time/src/timezone.rs index 41cc1f78428..dbd8dffcd49 100644 --- a/src/common/time/src/timezone.rs +++ b/src/common/time/src/timezone.rs @@ -108,6 +108,15 @@ impl Timezone { } } + /// A named zone merely sitting at +00:00 today is not UTC: it may have been + /// elsewhere at the timestamp being converted. + pub fn is_utc(&self) -> bool { + match self { + Self::Offset(offset) => offset.local_minus_utc() == 0, + Self::Named(tz) => matches!(tz, Tz::UTC), + } + } + /// Returns the number of seconds to add to convert from UTC to the local time. pub fn local_minus_utc(&self) -> i64 { match self { diff --git a/src/query/src/optimizer.rs b/src/query/src/optimizer.rs index 8827d48ed63..a6b36e6fdae 100644 --- a/src/query/src/optimizer.rs +++ b/src/query/src/optimizer.rs @@ -17,6 +17,7 @@ pub mod constant_term; pub mod count_nest_aggr; pub mod count_wildcard; pub mod global_limit; +pub(crate) mod insert_assignment; pub(crate) mod json_type_concretize; pub mod parallelize_scan; pub mod pass_distribution; diff --git a/src/query/src/optimizer/insert_assignment.rs b/src/query/src/optimizer/insert_assignment.rs new file mode 100644 index 00000000000..ad77736ce0a --- /dev/null +++ b/src/query/src/optimizer/insert_assignment.rs @@ -0,0 +1,307 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use common_time::Timezone; +use datafusion::config::ConfigOptions; +use datafusion_common::{DFSchemaRef, Result, ScalarValue}; +use datafusion_expr::expr::{Alias, Cast}; +use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Values}; +use datafusion_optimizer::analyzer::AnalyzerRule; +use datafusion_optimizer::analyzer::type_coercion::TypeCoercion; +use datatypes::arrow::datatypes::{DataType, TimeUnit}; +use session::context::QueryContextRef; + +use crate::optimizer::type_conversion::cast_string_to_timestamp; + +/// Interprets strings assigned to timestamp columns at an `INSERT` boundary +/// using the session timezone. `plan` is the assignment projection under a +/// `WriteOp::Insert`. +/// +/// DataFusion plans `INSERT` as a projection casting each source column to its +/// target column type, and that cast reads a naive string as UTC. Arrow does +/// apply a timezone when the cast target carries one, so the assignment is +/// routed through `Timestamp(unit, Some(tz))` and back. Stripping the timezone +/// afterwards is value-preserving — arrow only shifts values in the opposite +/// direction. +/// +/// The source query is left untouched. Reinterpreting a value where it is +/// *produced* would change what the source query means: pushing the conversion +/// below a `UNION`'s `DISTINCT`, for instance, moves the dedup key from the raw +/// strings to parsed instants and silently drops rows. +/// +/// # Why `TypeCoercion` runs here +/// +/// The rewrite reads source types, and those are only settled once a `UNION`'s +/// branch types have been reconciled: before coercion a union carries its loose +/// schema (the first branch's types), so a mixed +/// `SELECT 'string' UNION ALL SELECT CAST(.. AS TIMESTAMP)` still looks like a +/// string. Retargeting that cast would leave `Timestamp(None) -> +/// Timestamp(Some(tz))` behind once coercion retypes the union — the one +/// direction in which arrow shifts the value instead of relabelling it. +/// +/// Coercing here rather than deferring to the analyzer is forced by where an +/// INSERT is still identifiable: `exec_dml_statement` strips the `Dml` node and +/// executes its input, so by the time the analyzer runs, an assignment +/// projection is indistinguishable from any other projection. +/// +/// Explicit casts stay out of this: the SQL layer turns a user's +/// `CAST(x AS TIMESTAMP)` into an `arrow_cast` call, which only becomes an +/// `Expr::Cast` in the optimizer's `SimplifyExpressions`. Assignment casts are +/// therefore the only `Expr::Cast` reaching a timestamp column here. +/// +/// # Reach +/// +/// The emitted cast only carries its timezone where the expression is evaluated +/// on this node. Substrait drops the timezone name when a plan is pushed down — +/// it encodes any zoned timestamp as `PrecisionTimestampTz` and decodes it back +/// as UTC — so a source reading from a table falls back to UTC, the behaviour it +/// had before this rule existed. Sources that never leave this node (literals, +/// `VALUES`, and `UNION`s of them) keep the session timezone, and those are what +/// an INSERT's timestamp assignment is in practice. +pub(crate) fn rewrite_insert_assignments( + plan: LogicalPlan, + query_ctx: &QueryContextRef, + config: &ConfigOptions, +) -> Result { + let Some(timezone) = session_timezone(query_ctx) else { + return Ok(plan); + }; + + let plan = TypeCoercion::new().analyze(plan, config)?; + rewrite_assignment(plan, &timezone) +} + +/// Session timezone, in both forms the rewrite needs. +struct SessionTimezone { + /// Parses literals, matching the plain `INSERT ... VALUES` path. + parsed: Timezone, + /// Names the intermediate arrow cast target. + name: Arc, +} + +fn session_timezone(query_ctx: &QueryContextRef) -> Option { + let parsed = query_ctx.timezone(); + + // A UTC session already gets UTC semantics from the plain assignment cast. + if parsed.is_utc() { + return None; + } + + Some(SessionTimezone { + name: Arc::from(parsed.to_string()), + parsed, + }) +} + +fn rewrite_assignment(plan: LogicalPlan, timezone: &SessionTimezone) -> Result { + let LogicalPlan::Projection(assignment) = plan else { + return Ok(plan); + }; + + let mut exprs = assignment.expr.clone(); + let mut changed = false; + for expr in &mut exprs { + changed |= retarget_assignment_cast( + expr, + assignment.input.schema(), + Some(assignment.input.as_ref()), + timezone, + )?; + } + + // The planner types `VALUES` against the target table, so the assignment + // cast lands inside the `Values` rows instead of on the projection above. + let mut input = assignment.input.clone(); + if let LogicalPlan::Values(values) = assignment.input.as_ref() + && let Some(rewritten) = rewrite_values(values, timezone)? + { + input = Arc::new(LogicalPlan::Values(rewritten)); + changed = true; + } + + if !changed { + return Ok(LogicalPlan::Projection(assignment)); + } + Projection::try_new(exprs, input).map(LogicalPlan::Projection) +} + +fn rewrite_values(values: &Values, timezone: &SessionTimezone) -> Result> { + let mut rewritten = values.clone(); + let mut changed = false; + for row in &mut rewritten.values { + for expr in row.iter_mut() { + changed |= retarget_assignment_cast(expr, &values.schema, None, timezone)?; + } + } + + Ok(changed.then_some(rewritten)) +} + +/// Reinterprets one assignment cast, returning whether it was rewritten. +/// +/// `source_plan` is the projection's input, used to resolve a literal behind a +/// column reference; `Values` rows carry their expression inline and pass `None`. +fn retarget_assignment_cast( + expr: &mut Expr, + schema: &DFSchemaRef, + source_plan: Option<&LogicalPlan>, + timezone: &SessionTimezone, +) -> Result { + let expr = unalias_mut(expr); + let Expr::Cast(Cast { + expr: source, + data_type: DataType::Timestamp(unit, None), + }) = expr + else { + return Ok(false); + }; + let unit = *unit; + + if !matches!( + source.get_type(schema)?, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) { + return Ok(false); + } + + // Fold literals with the same parser the plain `INSERT ... VALUES` path + // uses, so a given string means the same thing however it reaches a column. + // The parsers disagree on ambiguous local times: this one resolves them, + // arrow rejects them. + let folded = source_literal(source.as_ref(), source_plan) + .and_then(|literal| convert_literal(&literal, unit, &timezone.parsed)); + if let Some(folded) = folded { + *expr = folded; + return Ok(true); + } + + let source = source.as_ref().clone(); + *expr = Expr::Cast(Cast::new( + Box::new(Expr::Cast(Cast::new( + Box::new(source), + DataType::Timestamp(unit, Some(timezone.name.clone())), + ))), + DataType::Timestamp(unit, None), + )); + Ok(true) +} + +fn source_literal(source: &Expr, source_plan: Option<&LogicalPlan>) -> Option { + match source { + Expr::Literal(value, _) => Some(value.clone()), + Expr::Column(column) => { + let plan = source_plan?; + let index = plan.schema().maybe_index_of_column(column)?; + lineage_literal(plan, index).cloned() + } + _ => None, + } +} + +/// Resolves a literal when every row carries the same value at `output_idx`. +/// +/// Read-only: the literal is folded into the assignment above, so nodes that +/// drop, reorder or deduplicate rows can be traversed — none of them changes +/// the value a surviving row carries, and folding above them leaves their keys +/// on the original strings. +fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { + if output_idx >= plan.schema().fields().len() { + return None; + } + + match plan { + LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { + Expr::Literal(value, _) => Some(value), + Expr::Column(column) => { + let input_idx = projection.input.schema().maybe_index_of_column(column)?; + lineage_literal(projection.input.as_ref(), input_idx) + } + _ => None, + }, + LogicalPlan::Filter(_) + | LogicalPlan::Sort(_) + | LogicalPlan::Limit(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Distinct(Distinct::All(_)) => { + let inputs = plan.inputs(); + let [input] = inputs.as_slice() else { + return None; + }; + lineage_literal(input, output_idx) + } + _ => None, + } +} + +fn convert_literal(value: &ScalarValue, unit: TimeUnit, timezone: &Timezone) -> Option { + let ScalarValue::Utf8(Some(value)) = value else { + return None; + }; + cast_string_to_timestamp(value, &DataType::Timestamp(unit, None), Some(timezone)) + .ok() + .filter(|value| !value.is_null()) + .map(|value| Expr::Literal(value, None)) +} + +fn unalias(expr: &Expr) -> &Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias(expr), + expr => expr, + } +} + +fn unalias_mut(expr: &mut Expr) -> &mut Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias_mut(expr), + expr => expr, + } +} + +#[cfg(test)] +mod tests { + use datafusion_common::DFSchema; + use datafusion_expr::expr::Placeholder; + + use super::*; + + fn shanghai() -> SessionTimezone { + let parsed = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + SessionTimezone { + name: Arc::from(parsed.to_string()), + parsed, + } + } + + /// A prepared `INSERT ... VALUES (?)` arrives here as a cast over an untyped + /// placeholder, which must survive for parameter substitution. + #[test] + fn test_untyped_placeholder_assignment_is_left_alone() { + let schema = Arc::new(DFSchema::empty()); + let mut expr = Expr::Cast(Cast::new( + Box::new(Expr::Placeholder(Placeholder::new_with_field( + "$1".to_string(), + None, + ))), + DataType::Timestamp(TimeUnit::Millisecond, None), + )); + let original = expr.clone(); + + let changed = retarget_assignment_cast(&mut expr, &schema, None, &shanghai()).unwrap(); + + assert!(!changed); + assert_eq!(expr, original); + } +} diff --git a/src/query/src/optimizer/type_conversion.rs b/src/query/src/optimizer/type_conversion.rs index 3941fedfe8a..40bb43e791e 100644 --- a/src/query/src/optimizer/type_conversion.rs +++ b/src/query/src/optimizer/type_conversion.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -mod insert_assignment; - use std::sync::Arc; use common_time::Timezone; @@ -31,7 +29,7 @@ use session::context::QueryContextRef; use crate::QueryEngineContext; use crate::optimizer::ExtensionAnalyzerRule; -use crate::optimizer::type_conversion::insert_assignment::rewrite_insert_assignments; +use crate::optimizer::insert_assignment::rewrite_insert_assignments; use crate::plan::ExtractExpr; /// TypeConversionRule converts some literal values in logical plan to other types according @@ -46,7 +44,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule { &self, plan: LogicalPlan, ctx: &QueryEngineContext, - _config: &ConfigOptions, + config: &ConfigOptions, ) -> Result { plan.transform_up_with_subqueries(|plan| match plan { LogicalPlan::Filter(filter) => { @@ -129,7 +127,8 @@ impl ExtensionAnalyzerRule for TypeConversionRule { LogicalPlan::Dml(mut dml) if matches!(dml.op, WriteOp::Insert(_)) => { dml.input = Arc::new(rewrite_insert_assignments( dml.input.as_ref().clone(), - ctx.query_ctx(), + &ctx.query_ctx(), + config, )?); Ok(Transformed::yes(LogicalPlan::Dml(dml))) } @@ -324,7 +323,7 @@ fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr { ) } -fn cast_string_to_timestamp( +pub(crate) fn cast_string_to_timestamp( string: &str, target_type: &DataType, timezone: Option<&Timezone>, diff --git a/src/query/src/optimizer/type_conversion/insert_assignment.rs b/src/query/src/optimizer/type_conversion/insert_assignment.rs deleted file mode 100644 index 11ab7e7829e..00000000000 --- a/src/query/src/optimizer/type_conversion/insert_assignment.rs +++ /dev/null @@ -1,336 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::sync::Arc; - -use datafusion_common::{Column, Result, ScalarValue}; -use datafusion_expr::expr::{Alias, Cast}; -use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Union, Values}; -use datatypes::arrow::datatypes::DataType; -use session::context::QueryContextRef; - -use crate::optimizer::type_conversion::cast_string_to_timestamp; -use crate::plan::ExtractExpr; - -/// Rewrites string literals that feed timestamp columns at an INSERT boundary. -/// Constants are folded at the assignment to avoid changing source types; -/// `VALUES` and `UNION` inputs are rewritten per row or branch. Explicit casts -/// stay on DataFusion's existing path. -pub(super) fn rewrite_insert_assignments( - plan: LogicalPlan, - query_ctx: QueryContextRef, -) -> Result { - let LogicalPlan::Projection(assignment) = plan else { - return Ok(plan); - }; - - let converter = InsertAssignmentConverter { query_ctx }; - let mut exprs = assignment.expr.clone(); - let mut input = assignment.input.as_ref().clone(); - let mut changed = false; - for (output_idx, expr) in assignment.expr.iter().enumerate() { - let target_type = assignment.schema.field(output_idx).data_type(); - if !matches!(target_type, DataType::Timestamp(_, _)) { - continue; - } - - let Some(column) = assignment_input_column(expr) else { - continue; - }; - let Some(input_idx) = input.schema().maybe_index_of_column(column) else { - continue; - }; - - let literal = lineage_literal(&input, input_idx).cloned(); - if let Some(literal) = literal - && let Some(folded) = converter.convert_literal(&literal, target_type) - { - let (qualifier, field) = assignment.schema.qualified_field(output_idx); - exprs[output_idx] = folded.alias_qualified(qualifier.cloned(), field.name()); - changed = true; - continue; - } - - // A hand-built DML plan (e.g. from flow) may share one source column - // between targets; rewriting it in place would retype every reader. - if assignment - .expr - .iter() - .enumerate() - .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) - { - continue; - } - if let Some(rewritten) = converter.rewrite_output_column(&input, input_idx, target_type)? { - input = rewritten; - changed = true; - } - } - - if !changed { - return Ok(LogicalPlan::Projection(assignment)); - } - Projection::try_new(exprs, Arc::new(input)).map(LogicalPlan::Projection) -} - -/// Resolves a literal when every row carries the same value at `output_idx`. -fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { - if output_idx >= plan.schema().fields().len() { - return None; - } - - match plan { - LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { - Expr::Literal(value, _) => Some(value), - Expr::Column(column) => { - let input_idx = projection.input.schema().maybe_index_of_column(column)?; - lineage_literal(projection.input.as_ref(), input_idx) - } - _ => None, - }, - // These nodes drop, reorder or deduplicate rows without touching the - // value a surviving row carries, and their schema stays positional. - LogicalPlan::Filter(_) - | LogicalPlan::Sort(_) - | LogicalPlan::Limit(_) - | LogicalPlan::SubqueryAlias(_) - | LogicalPlan::Distinct(Distinct::All(_)) => { - let inputs = plan.inputs(); - let [input] = inputs.as_slice() else { - return None; - }; - lineage_literal(input, output_idx) - } - _ => None, - } -} - -struct InsertAssignmentConverter { - query_ctx: QueryContextRef, -} - -impl InsertAssignmentConverter { - fn rewrite_output_column( - &self, - plan: &LogicalPlan, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - if output_idx >= plan.schema().fields().len() { - return Ok(None); - } - - match plan { - // DataFusion pushes INSERT assignment casts into Values, so inspect - // them even though the Values schema already has the target type. - LogicalPlan::Values(values) => self.rewrite_values(values, output_idx, target_type), - // Once the source query has produced the target type, its casts and - // coercions belong to the source query rather than INSERT assignment. - _ if plan.schema().field(output_idx).data_type() == target_type => Ok(None), - LogicalPlan::Projection(projection) => { - self.rewrite_projection(projection, output_idx, target_type) - } - LogicalPlan::Union(union) => self.rewrite_union(union, output_idx, target_type), - // Filter and Sort are excluded here: their predicates and keys read - // the retyped column, which would change the source query. Constants - // still reach the assignment through `lineage_literal`. Distinct::All - // holds no expressions; allowing it shifts dedup keys from raw - // strings to parsed instants. - LogicalPlan::Limit(_) - | LogicalPlan::SubqueryAlias(_) - | LogicalPlan::Distinct(Distinct::All(_)) => { - self.rewrite_passthrough(plan, output_idx, target_type) - } - _ => Ok(None), - } - } - - fn rewrite_projection( - &self, - projection: &Projection, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - match unalias(&projection.expr[output_idx]) { - Expr::Literal(value, _) => { - let Some(converted) = self.convert_literal(value, target_type) else { - return Ok(None); - }; - let (qualifier, field) = projection.schema.qualified_field(output_idx); - let mut exprs = projection.expr.clone(); - exprs[output_idx] = converted.alias_qualified(qualifier.cloned(), field.name()); - - Projection::try_new(exprs, projection.input.clone()) - .map(LogicalPlan::Projection) - .map(Some) - } - Expr::Column(column) => { - // Retyping the input column affects every output column reading - // it, so only follow lineage with a single consumer. - if projection - .expr - .iter() - .enumerate() - .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) - { - return Ok(None); - } - - let input = projection.input.as_ref(); - let Some(input_idx) = input.schema().maybe_index_of_column(column) else { - return Ok(None); - }; - let Some(rewritten) = self.rewrite_output_column(input, input_idx, target_type)? - else { - return Ok(None); - }; - - Projection::try_new(projection.expr.clone(), Arc::new(rewritten)) - .map(LogicalPlan::Projection) - .map(Some) - } - _ => Ok(None), - } - } - - fn rewrite_values( - &self, - values: &Values, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let mut rewritten = values.clone(); - let mut changed = false; - for row in &mut rewritten.values { - let Some(expr) = row.get_mut(output_idx) else { - return Ok(None); - }; - - if let Expr::Cast(Cast { - expr: inner, - data_type, - }) = expr - && data_type == target_type - && let Expr::Literal(value, _) = inner.as_ref() - && let Some(value) = self.convert_literal(value, target_type) - { - *expr = value; - changed = true; - continue; - } - - if expr.get_type(values.schema.as_ref()).ok().as_ref() != Some(target_type) { - return Ok(None); - } - } - - Ok(changed.then_some(LogicalPlan::Values(rewritten))) - } - - fn rewrite_union( - &self, - union: &Union, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let mut inputs = Vec::with_capacity(union.inputs.len()); - // A union has one shared schema, so rewrite it only when the same output - // column is an INSERT-assigned literal in every branch. - for input in &union.inputs { - let Some(rewritten) = self.rewrite_output_column(input, output_idx, target_type)? - else { - return Ok(None); - }; - inputs.push(Arc::new(rewritten)); - } - - // Untouched columns may still have mismatched branch types before - // TypeCoercion runs, so rebuild loosely like the SQL planner does. - Union::try_new_with_loose_types(inputs) - .map(LogicalPlan::Union) - .map(Some) - } - - fn rewrite_passthrough( - &self, - plan: &LogicalPlan, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let inputs = plan.inputs(); - let [input] = inputs.as_slice() else { - return Ok(None); - }; - let Some(rewritten_input) = self.rewrite_output_column(input, output_idx, target_type)? - else { - return Ok(None); - }; - - plan.with_new_exprs(plan.expressions_consider_join(), vec![rewritten_input]) - .map(Some) - } - - fn convert_literal(&self, value: &ScalarValue, target_type: &DataType) -> Option { - let ScalarValue::Utf8(Some(value)) = value else { - return None; - }; - cast_string_to_timestamp(value, target_type, Some(&self.query_ctx.timezone())) - .ok() - .filter(|value| !value.is_null()) - .map(|value| Expr::Literal(value, None)) - } -} - -fn assignment_input_column(expr: &Expr) -> Option<&Column> { - let expr = match unalias(expr) { - Expr::Cast(Cast { expr, .. }) => expr.as_ref(), - expr => expr, - }; - let Expr::Column(column) = expr else { - return None; - }; - Some(column) -} - -fn unalias(expr: &Expr) -> &Expr { - match expr { - Expr::Alias(Alias { expr, .. }) => unalias(expr), - expr => expr, - } -} - -#[cfg(test)] -mod tests { - use datafusion_common::ScalarValue; - use datafusion_common::arrow::datatypes::TimeUnit; - use session::context::QueryContext; - - use super::*; - - #[test] - fn test_convert_literal_falls_back_for_unsupported_literal() { - let converter = InsertAssignmentConverter { - query_ctx: QueryContext::arc(), - }; - let target_type = DataType::Timestamp(TimeUnit::Nanosecond, None); - - for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] { - assert_eq!( - converter - .convert_literal(&ScalarValue::Utf8(Some(literal.to_string())), &target_type,), - None - ); - } - } -} diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index a7d4a84cb28..c3871a96afa 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -866,6 +866,30 @@ mod tests { engine.planner().plan(&stmt, query_ctx).await.unwrap() } + /// Plans `sql` and runs the DataFusion analyzer, which is where + /// `InsertAssignmentRule` sits. Planning alone stops short of it, so these + /// assertions would not see the assignment rewrite at all. + async fn analyze_insert( + engine: &QueryEngineRef, + sql: &str, + query_ctx: &QueryContextRef, + ) -> String { + let stmt = QueryLanguageParser::parse_sql(sql, query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap(); + let context = engine.engine_context(query_ctx.clone()); + let state = context.state(); + state + .analyzer() + .execute_and_check(plan, state.config_options(), |_, _| {}) + .unwrap() + .display_indent_schema() + .to_string() + } + #[tokio::test] async fn test_insert_timestamp_literals_use_query_timezone() { let query_ctx = Arc::new( @@ -898,20 +922,6 @@ mod tests { ) AS source", &[1_785_902_400_001_i64][..], ), - ( - "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-06 12:00:00.001', now() \ - UNION ALL \ - SELECT '2026-08-07 12:00:00.001', now()", - &[1_785_988_800_001_i64, 1_786_075_200_001_i64][..], - ), - ( - "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-16 12:00:00.001', now() \ - UNION \ - SELECT '2026-08-17 12:00:00.001', now()", - &[1_786_852_800_001_i64, 1_786_939_200_001_i64][..], - ), ( "INSERT INTO timestamps (ts, st) \ SELECT '2026-08-18 12:00:00.001', max(st) \ @@ -940,14 +950,7 @@ mod tests { &[1_786_766_400_001_i64][..], ), ] { - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx.clone()) - .await - .unwrap() - .display_indent() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; for expected_timestamp in expected_timestamps { assert!( @@ -968,19 +971,19 @@ mod tests { let engine = create_timestamp_test_engine().await; let sql = "INSERT INTO timestamps (ts, st) \ VALUES (CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), now())"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; + // An explicit cast reaches the analyzer as an `arrow_cast` call rather + // than an `Expr::Cast`, which is how it stays out of the rewrite. assert!( plan.contains("arrow_cast(Utf8(\"2026-08-08 12:00:00.001\")"), "{plan}" ); + // 12:00:00.001 read as Shanghai local time; the source query keeps UTC. + assert!( + !plan.contains("TimestampMillisecond(1786104000001, None)"), + "{plan}" + ); } #[tokio::test] @@ -1017,14 +1020,7 @@ mod tests { ][..], ), ] { - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx.clone()) - .await - .unwrap() - .display_indent_schema() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; for expected in expected { assert!(plan.contains(expected), "{plan}"); @@ -1033,30 +1029,37 @@ mod tests { } #[tokio::test] - async fn test_insert_union_tolerates_uncoerced_untouched_column() { + async fn test_insert_union_converts_via_assignment_cast() { let query_ctx = Arc::new( QueryContextBuilder::default() .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) .build(), ); let engine = create_timestamp_test_engine().await; - // `st` stays Timestamp vs Null across branches until TypeCoercion runs. - let sql = "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-06 12:00:00.001', now() \ - UNION ALL \ - SELECT '2026-08-07 12:00:00.001', NULL"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent() - .to_string(); + // Branches disagree, so the conversion stays as a cast on the + // assignment instead of folding. One cast covers every branch, which is + // why a NULL branch no longer cancels the conversion for the column and + // why UNION's dedup keys stay on the original strings. + for sql in [ + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-06 12:00:00.001', now() \ + UNION ALL \ + SELECT '2026-08-07 12:00:00.001', NULL", + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-16 12:00:00.001', now() \ + UNION \ + SELECT '2026-08-17 12:00:00.001', now()", + ] { + let plan = analyze_insert(&engine, sql, &query_ctx).await; - for expected_timestamp in [1_785_988_800_001_i64, 1_786_075_200_001_i64] { assert!( - plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")), + plan.contains("AS Timestamp(ms, \"Asia/Shanghai\")"), + "{plan}" + ); + // The branches themselves are untouched. + assert!( + plan.contains("Utf8(\"2026-08-07 12:00:00.001\")") + || plan.contains("Utf8(\"2026-08-17 12:00:00.001\")"), "{plan}" ); } @@ -1074,14 +1077,8 @@ mod tests { SELECT '2026-08-10 12:00:00.001', now() \ UNION ALL \ SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now()"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent_schema() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; + assert!( !plan.contains("TimestampMillisecond(1786334400001, None)"), "{plan}" @@ -1090,6 +1087,11 @@ mod tests { plan.contains("arrow_cast(Utf8(\"2026-08-11 12:00:00.001\")"), "{plan}" ); + // TypeCoercion has already settled this union to timestamp, so the + // assignment has nothing left to reinterpret. Retargeting the cast here + // would leave a Timestamp(None) -> Timestamp(Some(tz)) step behind, + // which shifts the value instead of relabelling it. + assert!(!plan.contains("Asia/Shanghai"), "{plan}"); } #[tokio::test] diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.result b/tests/cases/standalone/common/insert/insert_default_timezone.result index 413eb5d46ae..41d8ca6cba8 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.result +++ b/tests/cases/standalone/common/insert/insert_default_timezone.result @@ -114,6 +114,27 @@ INSERT INTO test3 (ts, st, ts_ns) VALUES ( Affected Rows: 1 +-- a NULL branch must not cancel the conversion for the whole column +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-13 12:00:00.001' AS a, '2026-08-13 12:00:00.123456789' AS b +UNION ALL +SELECT '2026-08-14 12:00:00.001', NULL; + +Affected Rows: 2 + +-- NULL in the first branch: the union's schema starts out as Null +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-15 12:00:00.001' AS a, NULL AS b +UNION ALL +SELECT '2026-08-19 12:00:00.001', '2026-08-19 12:00:00.123456789'; + +Affected Rows: 2 + +-- the assignment cast also lands on non-literal VALUES expressions +INSERT INTO test3 (ts, st) VALUES (concat('2026-08-20 ', '12:00:00.001'), now()); + +Affected Rows: 1 + SELECT ts, ts_ns FROM test3 ORDER BY ts; +-------------------------+-------------------------------+ @@ -130,10 +151,43 @@ SELECT ts, ts_ns FROM test3 ORDER BY ts; | 2026-08-10T12:00:00.001 | | | 2026-08-11T12:00:00.001 | | | 2026-08-12T04:00:00.123 | 2026-08-12T04:00:00.123456789 | +| 2026-08-13T04:00:00.001 | 2026-08-13T04:00:00.123456789 | +| 2026-08-14T04:00:00.001 | | +| 2026-08-15T04:00:00.001 | | | 2026-08-16T04:00:00.001 | | | 2026-08-17T04:00:00.001 | | +| 2026-08-19T04:00:00.001 | 2026-08-19T04:00:00.123456789 | +| 2026-08-20T04:00:00.001 | | +-------------------------+-------------------------------+ +-- UNION dedup keys must stay on the source strings: these two spell the same +-- instant differently, so the source query yields two rows and both are kept. +CREATE TABLE test4 (ts TIMESTAMP TIME INDEX) WITH ('append_mode'='true'); + +Affected Rows: 0 + +INSERT INTO test4 (ts) +SELECT '2026-08-06 04:00:00' UNION SELECT '2026-08-06 04:00:00.000'; + +Affected Rows: 2 + +SELECT count(*) FROM test4; + ++----------+ +| count(*) | ++----------+ +| 2 | ++----------+ + +SELECT ts FROM test4 ORDER BY ts; + ++---------------------+ +| ts | ++---------------------+ +| 2026-08-05T20:00:00 | +| 2026-08-05T20:00:00 | ++---------------------+ + SET time_zone = 'UTC'; Affected Rows: 0 @@ -150,3 +204,7 @@ DROP TABLE test3; Affected Rows: 0 +DROP TABLE test4; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.sql b/tests/cases/standalone/common/insert/insert_default_timezone.sql index 7cbef15d6d9..f7c196699b7 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.sql +++ b/tests/cases/standalone/common/insert/insert_default_timezone.sql @@ -56,8 +56,34 @@ INSERT INTO test3 (ts, st, ts_ns) VALUES ( '2026-08-09 12:00:00.123456789' ); +-- a NULL branch must not cancel the conversion for the whole column +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-13 12:00:00.001' AS a, '2026-08-13 12:00:00.123456789' AS b +UNION ALL +SELECT '2026-08-14 12:00:00.001', NULL; + +-- NULL in the first branch: the union's schema starts out as Null +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-15 12:00:00.001' AS a, NULL AS b +UNION ALL +SELECT '2026-08-19 12:00:00.001', '2026-08-19 12:00:00.123456789'; + +-- the assignment cast also lands on non-literal VALUES expressions +INSERT INTO test3 (ts, st) VALUES (concat('2026-08-20 ', '12:00:00.001'), now()); + SELECT ts, ts_ns FROM test3 ORDER BY ts; +-- UNION dedup keys must stay on the source strings: these two spell the same +-- instant differently, so the source query yields two rows and both are kept. +CREATE TABLE test4 (ts TIMESTAMP TIME INDEX) WITH ('append_mode'='true'); + +INSERT INTO test4 (ts) +SELECT '2026-08-06 04:00:00' UNION SELECT '2026-08-06 04:00:00.000'; + +SELECT count(*) FROM test4; + +SELECT ts FROM test4 ORDER BY ts; + SET time_zone = 'UTC'; DROP TABLE test1; @@ -65,3 +91,5 @@ DROP TABLE test1; DROP TABLE test2; DROP TABLE test3; + +DROP TABLE test4; From 6d86e6ff068fe8471c3d7a28e70f9144e7684e15 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 25 Aug 2026 13:08:44 +0000 Subject: [PATCH 108/126] feat: synthesize OTLP resource descriptor for the semantic entity graph (#8904) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(servers): compose OTLP metrics job from service.namespace/service.name The OTel Prometheus compatibility spec defines job as "/" when the namespace is present. The OTLP metrics path only used the bare service.name, so the job tag diverged from target_info produced by Prometheus-side exporters for the same resource. Compose the namespace form, and keep not fabricating a job when service.name is absent. Behavior change: resources carrying service.namespace now get "namespace/name" as their job tag value. Signed-off-by: Dennis Zhuang * feat(otlp): synthesize otel_resource_info at OTLP metrics ingestion Ordinary OTLP metrics scatter filtered resource attributes as tags over every logical metric table, so metrics-only services contribute nothing to the semantic entity graph. Each request now also projects its distinct resources into one info-metric-shaped mito table, otel_resource_info: a fixed allowlist of identity-relevant attributes under their raw OTel keys (independent of the label translation strategy and the promote/ignore headers) plus derived job/instance compatibility columns, value 1.0, and the newest data-point timestamp. The descriptor is written after the main insert is committed; a failure there (conflicting pre-existing table, auto-create disabled) degrades to an OTLP partial_success warning with rejected_data_points = 0 instead of failing the request and triggering client retries of already-accepted data. A request writing a metric named otel_resource_info suppresses synthesis. Legacy mode is unchanged. Signed-off-by: Dennis Zhuang * feat(operator): otel info-metric conventions with host/container entities Whitelist the ingestion-synthesized otel_resource_info descriptor via a new otel_info_metrics conventions map, gated on source=opentelemetry (the existing gate hardcoded source=prometheus). Its declarations use explicit descriptive lists instead of descriptive_rest so identifying attributes of other entities do not leak into service.instance. Conventions tightened per the Astronomy Shop findings: host identity is host.id with host.name descriptive only (host.name is not stable across SDKs and resource detectors), a generic container entity (new entity type) is declared only when container.id is present, and trace-v1 tables now synthesize host/container from their flattened resource attributes too. New co-declared edges: service.instance runs_on container, container runs_on host. Signed-off-by: Dennis Zhuang * test(otlp): cover the resource descriptor in integration tests Covers the descriptor's raw-key columns and info-metric options through the HTTP path, the namespace/name job composition end-to-end, column names being independent of the translation strategy, the allowlist excluding unlisted resource attributes, auto-create after a drop, the metric-name collision suppressing synthesis, and the partial-success warning (rejected_data_points = 0) when a pre-existing incompatible table fails the descriptor write while metric data is accepted. Signed-off-by: Dennis Zhuang * chore: cargo fmt Signed-off-by: Dennis Zhuang * fix(frontend): degrade descriptor permission denial to a warning A table-level permission policy denying otel_resource_info would have failed the whole OTLP metrics request because the descriptor's permission check ran before the main insert. The descriptor is derived enrichment: check its permission in the degrade path so a denial skips the write and surfaces as the partial-success warning, like any other descriptor write failure. Signed-off-by: Dennis Zhuang * fix(otlp): guard descriptor writes with semantic ownership markers A pre-existing schema-compatible table named otel_resource_info would silently receive descriptor rows while its missing semantic stamps kept it out of the entity graph. The descriptor write now requires the auto-created table's ownership markers (mito engine + signal_type + source + metric.type=info + metadata_quality=declared) and otherwise degrades to the partial-success warning; the entity-graph gate for the otel whitelist likewise requires metric.type=info, so a user table stamped with only signal/source no longer picks up implicit declarations. Also fold the descriptor write cost into the response and surface the degrade warning through the otel-arrow BatchStatus status_message. Integration tests pin the full marker set on auto-create and that an existing owned descriptor keeps accepting writes without degrading — a missing marker would otherwise silently stop every descriptor write after the first request. Signed-off-by: Dennis Zhuang * perf(otlp): build descriptor rows without the per-resource BTreeMap Projecting a resource allocated a BTreeMap and then collected it into the row key, and every attribute was matched against the allowlist by linear scan. Collect the tags into a Vec and sort once, and match the allowlist instead of scanning it. Measured on the conversion path: descriptor work drops 16-18%, from 10.6% to 8.9% of conversion CPU on the worst shape (1000 resources with 4 data points each), where the cost tracks resource count rather than data-point count. Also trims the comments and tests added with the descriptor to what carries information. Signed-off-by: Dennis Zhuang * test(otlp): pin the descriptor permission-denial degrade path A policy denying the descriptor table must not fail the metrics request, which the fix in 2401b3dd9c does but nothing covered. Verified as a regression guard by mutation: moving the permission check back before the main write makes this test fail. Signed-off-by: Dennis Zhuang * test(otlp): keep legacy mode covered after trimming the unit tests Trimming the descriptor tests dropped the only assertion that legacy mode skips the job/instance remap and the promote filter. Both alter the columns of tables already in use, so fold the check into the legacy conversion test rather than leaving it uncovered. Signed-off-by: Dennis Zhuang * fix(semantic-graph): stop encoding column names into composite entity ids A composite entity id rendered the identifying columns as sorted `col=value` pairs, so the same identity split into one entity per signal: a trace table names its columns service_name and resource_attributes.service.instance.id where a metric table names them job and instance. One service instance became two nodes with two parallel edge sets, breaking the walk from a trace to that instance's metrics. Render an id as its values in declared order instead, escaping the separator so components stay distinguishable, which is what single-column ids already did by keeping only the value. entity_id_attrs still carries the structured form. Values alone are not enough for a namespaced service: the metric side folds service.namespace into job while traces keep the bare name. Add qualified_by to the conventions so the trace declarations compose the namespace the same way, per the OTel rule that job is / or the bare name when the namespace is empty. A table without the namespace column keeps the unqualified identity rather than losing the declaration. Conventions validation now rejects one entity type declared with a different number of id columns by two sources, which would silently produce ids that can never match. Signed-off-by: Dennis Zhuang * style(otlp): import the parent module by crate path check-super-imports.py, part of the CI format gate, rejects a file-level `use super::`. Signed-off-by: Dennis Zhuang * feat(otlp): gate the resource descriptor, and fix what review found Synthesizing greptime_otel_resource_info creates and writes a table the user never sent, so it is now off unless otlp.experimental_enable_resource_info says otherwise. With it off the request costs exactly what it did before the descriptor existed: nothing is projected, no table is created, no write and no permission check happen. Tests run with it on. StandaloneOptions carried no otlp field, so the whole [otlp] section was silently dropped in standalone mode; map it through, or the new option (and trace_ingest_chunk_size before it) would do nothing there. Renamed from otel_resource_info: the greptime_ prefix marks the table as engine-managed and makes a collision with a user metric unlikely, which is what the pre-existing-table ownership check and its per-request catalog lookup were defending against. Both are gone. A request may carry data for several graph windows, but the descriptor folded every data-point time into one row at the newest of them, leaving the earlier windows with metric rows and no entities. Key the rows by window as well, and take the times from the data points the encoder actually writes: it drops exponential histograms, and a resource carrying nothing else was being described as an entity with no measurements. Projecting a resource cloned its attributes once per data point. Nest the windows under the attributes instead, so they are moved once per resource, and walk the data-point times through a visitor rather than collecting a Vec per metric. Also documents what the two maps key and hold, and lifts the projected attribute names to constants beside KEY_SERVICE_NAME. Signed-off-by: Dennis Zhuang * fix(otlp): skip the descriptor's work entirely when it is disabled The collision scan over the request's output tables ran even with the feature off. Short-circuit on the option instead, and update the config snapshot the new [otlp] section changed. Also drops the doc comment orphaned by the deleted ownership check: it had attached itself to the trait impl and described a check that no longer exists. Signed-off-by: Dennis Zhuang * refactor(semantic-graph): drop the expect and name the service identity The CASE is built through Case directly rather than the fallible when().otherwise() builder, so the non-test path no longer carries an expect (architecture-invariants $4). service_identity returned two same-typed Options that both call sites destructured positionally; a named struct makes a swap fail to compile. Also records that id-column order is part of the identity, where the option docs and the conventions authors will read it. Signed-off-by: Dennis Zhuang * chore(semantic-graph): drop comments that narrate the code Signed-off-by: Dennis Zhuang * fix(semantic-graph): cast duration_nano before the trace-table union Trace tables written before the signed-integer ingest change hold duration_nano as UInt64 and later ones as Int64. The calls derivation unions the per-table selects, and the two have no common integer type, so a deployment holding both shapes could not build the plan. The cross-table test now spans both. Signed-off-by: Dennis Zhuang * refactor(semantic-graph): drop the redundant duration_nano casts Signed-off-by: Dennis Zhuang * fix(otlp): decide exponential histogram acceptance in one place The resource descriptor mirrored only the experimental gate, so with both experimental flags on a resource whose only metric is a delta exponential histogram was described as an entity with no measurements. The encoder's whole-metric rules move into exponential_histogram_gate, which both call, and the descriptor takes its timestamps through exponential_histogram_value so per-point rejections drop out too. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- config/config.md | 2 + config/frontend.example.toml | 3 + config/standalone.example.toml | 3 + ...25-entity-relationships-and-graph-query.md | 8 +- src/common/catalog/src/consts.rs | 5 + src/frontend/src/instance.rs | 1 + src/frontend/src/instance/builder.rs | 1 + src/frontend/src/instance/entity_graph.rs | 230 ++++++++- src/frontend/src/instance/otlp.rs | 47 +- src/frontend/src/metrics.rs | 8 + src/frontend/src/service_config/otlp.rs | 7 + src/operator/src/statement/semantic_graph.rs | 202 ++++++-- .../statement/semantic_graph/conventions.rs | 74 ++- .../statement/semantic_graph/conventions.yaml | 50 +- .../statement/semantic_graph/relationships.rs | 79 ++- src/servers/src/http/otlp.rs | 5 +- src/servers/src/otlp/metrics.rs | 467 ++++++++++++++---- src/servers/src/otlp/metrics/resource_info.rs | 393 +++++++++++++++ src/servers/src/otlp/trace.rs | 8 + src/session/src/protocol_ctx.rs | 3 + src/table/src/requests/semantic.rs | 4 + tests-integration/src/instance.rs | 115 +++++ tests-integration/src/standalone.rs | 5 + tests-integration/tests/http.rs | 220 ++++++++- .../common/system/semantic_graph.result | 261 ++++++++-- .../common/system/semantic_graph.sql | 119 +++++ tests/conf/frontend-test.toml.template | 3 + tests/conf/standalone-test.toml.template | 3 + 28 files changed, 2093 insertions(+), 233 deletions(-) create mode 100644 src/servers/src/otlp/metrics/resource_info.rs diff --git a/config/config.md b/config/config.md index 91cb4fe7669..992296d4827 100644 --- a/config/config.md +++ b/config/config.md @@ -76,6 +76,7 @@ | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | @@ -319,6 +320,7 @@ | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 3635a63f063..80534855721 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -241,6 +241,9 @@ enable = true experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric +## resource attributes, so metrics-only services reach the semantic graph. +experimental_enable_resource_info = false ## Prometheus remote storage options [prom_store] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 72ba992cce0..70f74c7f946 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -208,6 +208,9 @@ enable = true experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric +## resource attributes, so metrics-only services reach the semantic graph. +experimental_enable_resource_info = false ## Prometheus remote storage options [prom_store] diff --git a/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md b/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md index 995cf6cab7e..1c78e91f252 100644 --- a/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md +++ b/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md @@ -149,7 +149,7 @@ The decisions behind the shape: - **`provenance` is part of edge identity**, so a declared edge and a derived edge between the same pair coexist, and a user-declared edge survives when derivation is off. - **`confidence` expresses derivation certainty, not statistical completeness**: `1.0` for a successfully paired or declared edge, `< 1.0` for virtual-node or agent-inferred edges. It does not correct for trace sampling. - **RED metrics describe the observed span-pair population** — the `service_graph` connector semantics. Under sampling, counts understate true traffic, and ratios (error rate, mean duration) are representative only when sampling is unbiased with respect to status and latency — tail sampling that keeps errors and slow traces skews both. -- **Endpoint encoding (v1).** A single-attribute identity is the value verbatim; a composite identity is the attributes sorted by key, rendered `k1=v1,k2=v2` — a human-readable traversal key, not a collision-free canonical encoding. Nodes carry the structured `entity_id_attrs`; edges carry only the string key. Known v1 limitation: composite components containing unescaped `,` or `=` are not guaranteed to be distinguished. +- **Endpoint encoding.** An identity renders as its values in declared order (broad to narrow), joined by `,` with `,` and `\` escaped in each value. The identifying attribute *keys* are deliberately absent: the same identity reaches the graph under different column names per signal — a trace table's `service_name` against a metric table's `job` — and encoding the keys would split one entity into one entity per signal. Where a source carries a qualifier separately (`service.namespace` beside `service.name`) the convention folds it in as `/`, reproducing the identity a source that pre-composed it (Prometheus `job`) already emits. Nodes carry the structured `entity_id_attrs`; edges carry only the string key. - **Hand-declared edges** (`provenance = 'declared'`) live in a physical table of the same shape (plus a business validity window) in `greptime_private` and are unioned into the computed `semantic_relationships`. The relationship vocabulary is small, typed, and inverse-paired; stored direction is `src -> dst`, the inverse is a query concern: @@ -197,7 +197,7 @@ GROUP BY 1, src_id, dst_id; Edge endpoints are built from each trace table's `service` entity declaration (the same id expression the registry uses), so edges land on exactly the entity ids the registry emits, self-calls are judged on the full endpoint id, and a table whose declaration is unusable contributes no edges rather than silently switching identity. The join is bounded by a time-proximity window and the results of **all** trace tables combine into one edge per window with merged RED metrics. One assumption to note: a deployment can route traces to multiple tables (`x-greptime-trace-table-name`), so one distributed trace is not guaranteed to live in a single table — pairing must account for that, and how is an implementation concern. Uninstrumented peers become **virtual nodes**: a client span with no matching server span yields an edge to a synthetic node named from `peer.service`/`db.name`/`server.address`. -**The entity registry.** Per declaring table: filter to the window, project (time bin, casts, JSON assembly), `DISTINCT`. The per-table results are unioned without any cross-source merge: no descriptive-conflict or JSON-merging policy lives in the derivation. Rows whose identity columns are NULL identify nothing and are filtered out; single-column ids are used verbatim, composite ids render the sorted `k=v` form plus the `entity_id_attrs` JSON. Consumers deduplicate with `SELECT DISTINCT entity_type, entity_id`; a unique node set is the snapshot relation's job (see *Querying the graph*). +**The entity registry.** Per declaring table: filter to the window, project (time bin, casts, JSON assembly), `DISTINCT`. The per-table results are unioned without any cross-source merge: no descriptive-conflict or JSON-merging policy lives in the derivation. Rows whose identity columns are NULL identify nothing and are filtered out; ids render as their escaped values joined in declared order, with the `entity_id_attrs` JSON alongside once an id is assembled from more than one column. Consumers deduplicate with `SELECT DISTINCT entity_type, entity_id`; a unique node set is the snapshot relation's job (see *Querying the graph*). **Attribute edges.** Shared attributes provide join keys; co-declaration or a relationship template provides the relationship semantics. A shared value alone determines neither direction nor edge type, so the derivation is rule-based, not join-everything. A table declaring both a `service.instance` and a `host` identity on the same rows derives `runs_on` between them: the row itself witnesses the relationship, and the built-in vocabulary fixes the direction. Cross-table edges require a declared relationship template (endpoint mappings, `rel_type`, direction): a derivation *rule* declared on schema, not per-instance edge data. @@ -284,7 +284,7 @@ Drawbacks: - Read-time derivation costs a self-join over trace tables on every scan of the computed tables. The default one-hour window bounds it, but very large deployments will want materialisation (Future Work). - `semantic_entities` exposes per-observation rows; deduplication falls on the consumer (`DISTINCT`) until the snapshot relation lands in M2. -- The v1 `k=v` id encoding can collide on hostile values (unescaped `,`/`=`); hardening is an open question. +- The id encoding is readable rather than canonical: escaping keeps components distinguishable, but a value containing the qualifier separator (`/`) is still ambiguous against a qualified name. - The graph is only as connected as the identity columns tables actually share: the `job`/`service_name` misalignment above is the canonical example, and enrichment discipline gates connectivity. # Implementation Plan @@ -306,7 +306,7 @@ Drawbacks: # Open Questions -1. **Endpoint encoding hardening.** The v1 `k=v` rendering (with `entity_id_attrs` JSON as the source of truth) is readable but collides when values contain `,`/`=`. Open: a fixed-width `entity_id_hash` join key (New-Relic-style); typed value encoding; the criteria for introducing the `semantic_key=column` mapping form; `schema_url`-gated entity merge. Related and unresolved: whether the identifying attribute *key* should participate in entity equality — a single-attribute id keeps only the value, so `job=api` and `service_name=api` currently name the same entity even when the keys differ semantically. +1. **Endpoint encoding hardening.** The rendering (with `entity_id_attrs` JSON as the source of truth) is readable rather than canonical. Open: a fixed-width `entity_id_hash` join key (New-Relic-style); typed value encoding; the criteria for introducing the `semantic_key=column` mapping form; `schema_url`-gated entity merge. Settled: identifying attribute *keys* do not participate in entity equality, so `job=api` and `service_name=api` name the same entity — the property that lets one service reach the graph from several signals. 2. **Identity reconciliation across signals.** Trace edges key on the trace table's service declaration; host edges key on `host.id`/`k8s.pod.uid`. Linking them relies on tables carrying both columns; how aggressively to enrich (a `k8sattributes`-equivalent) gates how connected the graph is. 3. **`GRAPH_TABLE` surface scope.** How much of SQL/PGQ to implement vs. a minimal single-hop subset. Start minimal and grow with demand. 4. **Visibility beyond table grants.** The derivation contract filters by table-read permission; whether topology needs finer policies (row-level, tenant-scoped visibility of cross-schema edges) is open. diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index b976cbc256e..63aa103b83d 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -220,6 +220,11 @@ pub const SEMANTIC_RELATIONSHIPS_TABLE_NAME: &str = "semantic_relationships"; /// `semantic_relationships`; see [`is_ddl_reserved_table`] for its lifecycle. pub const SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME: &str = "semantic_relationships_declared"; +/// Width of the window the graph derivation bins observations into. Sources +/// synthesizing observations must land a row in every window they describe, +/// so this is shared rather than restated per crate. +pub const SEMANTIC_GRAPH_WINDOW_NANOS: i64 = 60 * 1_000_000_000; + // Column names of the graph tables: `catalog` exposes these schemas and the // read-time plans in `operator` must project exactly them. pub const OBSERVED_AT_COLUMN: &str = "observed_at"; diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 78b701a3442..412ff495f65 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -134,6 +134,7 @@ pub struct Instance { slow_query_options: SlowQueryOptions, influxdb_default_merge_mode: InfluxdbMergeMode, trace_ingest_chunk_size: usize, + otlp_resource_info: bool, suspend: Arc, // cache for otlp metrics diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index c49a8394beb..0f87dc424f0 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -367,6 +367,7 @@ impl FrontendBuilder { slow_query_options: self.options.slow_query.clone(), influxdb_default_merge_mode: self.options.influxdb.default_merge_mode, trace_ingest_chunk_size: self.options.otlp.trace_ingest_chunk_size, + otlp_resource_info: self.options.otlp.experimental_enable_resource_info, suspend: Arc::new(AtomicBool::new(false)), }) } diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs index 14f18ed1b79..a0817cf2806 100644 --- a/src/frontend/src/instance/entity_graph.rs +++ b/src/frontend/src/instance/entity_graph.rs @@ -23,7 +23,7 @@ //! the query engine. Injected into the catalog manager after the engine is built, //! breaking the `catalog -> query` cycle. -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Weak; use async_trait::async_trait; @@ -60,8 +60,9 @@ use table::TableRef; use table::metadata::TableInfo; use table::predicate::{TimeRangeExtraction, extract_time_range_strict}; use table::requests::{ - EntityRole, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, SOURCE_PROMETHEUS, - is_trace_v1_table, parse_entity_columns, parse_entity_option_key, + EntityRole, SEMANTIC_METRIC_TYPE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, + SOURCE_OPENTELEMETRY, SOURCE_PROMETHEUS, is_trace_v1_table, parse_entity_columns, + parse_entity_option_key, }; use crate::error; @@ -177,6 +178,7 @@ impl EntityGraphProviderImpl { time_index: time_index.clone(), entity_type, id_columns, + id_qualifier: None, descriptive_columns, scope_columns, }) @@ -188,10 +190,11 @@ impl EntityGraphProviderImpl { /// All entity declarations of one table: the explicit options plus the /// zero-configuration conventions (`otlp_trace_entities` for trace-v1 /// tables — including the `service` identity of tables created before the - /// ingest-side auto-stamp — and the Prometheus descriptor whitelist). An - /// explicit declaration of a type always suppresses the implicit one, even - /// when the explicit declaration is invalid and skipped: silently falling - /// back would change entity identity behind the user's back. + /// ingest-side auto-stamp — and the Prometheus/OTel descriptor + /// whitelists). An explicit declaration of a type always suppresses the + /// implicit one, even when the explicit declaration is invalid and + /// skipped: silently falling back would change entity identity behind the + /// user's back. fn declarations_for( table_info: &TableInfo, conventions: &Conventions, @@ -204,7 +207,20 @@ impl EntityGraphProviderImpl { &mut declarations, ); } - Self::extend_with_prometheus_conventions(table_info, conventions, &mut declarations); + Self::extend_with_info_metric_conventions( + table_info, + &conventions.prometheus_info_metrics, + SOURCE_PROMETHEUS, + None, + &mut declarations, + ); + Self::extend_with_info_metric_conventions( + table_info, + &conventions.otel_info_metrics, + SOURCE_OPENTELEMETRY, + Some(servers::semantic::METRIC_TYPE_INFO), + &mut declarations, + ); declarations } @@ -216,29 +232,34 @@ impl EntityGraphProviderImpl { }) } - /// Implicit declarations of the well-known Prometheus entity-descriptor - /// metrics (the `prometheus_info_metrics` whitelist of `conventions.yaml`), - /// gated on the ingest-stamped `signal_type=metric` + `source=prometheus` - /// options and keyed by table name. The metric engine's physical table + /// Implicit declarations of the well-known entity-descriptor metrics + /// (the `prometheus_info_metrics` / `otel_info_metrics` whitelists of + /// `conventions.yaml`), gated on the ingest-stamped `signal_type=metric` + /// option plus the whitelist's expected `source`; OTel descriptors also + /// require `metric.type=info`. The metric engine's physical table /// aggregates every logical table's columns and must not contribute a /// duplicate source. - fn extend_with_prometheus_conventions( + fn extend_with_info_metric_conventions( table_info: &TableInfo, - conventions: &Conventions, + whitelist: &BTreeMap>, + expected_source: &str, + expected_metric_type: Option<&str>, declarations: &mut Vec, ) { - let Some(implicit_entities) = conventions.prometheus_info_metrics.get(&table_info.name) - else { + let Some(implicit_entities) = whitelist.get(&table_info.name) else { return; }; let options = &table_info.meta.options.extra_options; if options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str) != Some(SIGNAL_TYPE_METRIC) - || options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(SOURCE_PROMETHEUS) + || options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(expected_source) + || expected_metric_type.is_some_and(|expected| { + options.get(SEMANTIC_METRIC_TYPE).map(String::as_str) != Some(expected) + }) || table_info.is_physical_table() { debug!( - "Table `{}` matches the info-metric whitelist but is not a prometheus logical \ - metric table; skipping its implicit declarations", + "Table `{}` matches the info-metric whitelist but is not an eligible \ + `{expected_source}` info-metric source; skipping its implicit declarations", table_info.name ); return; @@ -296,12 +317,19 @@ impl EntityGraphProviderImpl { .cloned() .collect() }; + // A table predating the qualifier column keeps the unqualified + // identity rather than losing the declaration. + let id_qualifier = implicit + .qualified_by + .clone() + .filter(|c| schema.column_schema_by_name(c).is_some()); declarations.push(EntityDeclaration { schema: table_info.schema_name.clone(), table: table_info.name.clone(), time_index: time_index.clone(), entity_type: implicit.entity.clone(), id_columns: implicit.id.clone(), + id_qualifier, descriptive_columns, scope_columns: vec![], }); @@ -832,6 +860,24 @@ mod tests { declarations[2].id_columns, vec!["service_name", "resource_attributes.service.instance.id"] ); + assert_eq!(declarations[1].id_qualifier, None); + + let namespaced = table_info( + &[ + "service_name", + "resource_attributes.service.namespace", + "resource_attributes.service.instance.id", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + for declaration in sorted_declarations(&namespaced) { + assert_eq!( + declaration.id_qualifier.as_deref(), + Some("resource_attributes.service.namespace"), + "{} must qualify its identity like the metric side's job", + declaration.entity_type + ); + } // Missing uid column: no pod entity synthesized, no name-based guess. let no_uid = table_info( @@ -856,6 +902,47 @@ mod tests { assert!(sorted_declarations(&invalid_explicit).is_empty()); } + #[test] + fn trace_table_host_and_container_require_stable_ids() { + let with_ids = table_info( + &[ + "service_name", + "resource_attributes.host.id", + "resource_attributes.host.name", + "resource_attributes.container.id", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + let declarations = sorted_declarations(&with_ids); + let types: Vec<&str> = declarations + .iter() + .map(|d| d.entity_type.as_str()) + .collect(); + assert_eq!(types, vec!["container", "host", "service"]); + assert_eq!( + declarations[1].id_columns, + vec!["resource_attributes.host.id"] + ); + assert_eq!( + declarations[1].descriptive_columns, + vec!["resource_attributes.host.name"] + ); + + let names_only = table_info( + &[ + "service_name", + "resource_attributes.host.name", + "resource_attributes.container.name", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + let types: Vec = sorted_declarations(&names_only) + .into_iter() + .map(|d| d.entity_type) + .collect(); + assert_eq!(types, vec!["service"]); + } + #[test] fn prometheus_implicit_declarations_are_gated() { let labels: &[&str] = &["namespace", "pod", "node"]; @@ -900,6 +987,111 @@ mod tests { assert_eq!(declarations[1].id_columns, vec!["pod"]); } + const OTEL_STAMPS: &[(&str, &str)] = &[ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY), + (SEMANTIC_METRIC_TYPE, servers::semantic::METRIC_TYPE_INFO), + ]; + + #[test] + fn greptime_otel_resource_info_gets_implicit_declarations() { + let info = prom_table_info( + "greptime_otel_resource_info", + &[ + "job", + "instance", + "service.name", + "service.namespace", + "host.id", + "host.name", + "container.id", + "container.name", + "k8s.pod.uid", + "k8s.pod.name", + "k8s.namespace.name", + ], + OTEL_STAMPS, + ); + let declarations = sorted_declarations(&info); + let types: Vec<&str> = declarations + .iter() + .map(|d| d.entity_type.as_str()) + .collect(); + assert_eq!( + types, + vec![ + "container", + "host", + "k8s.pod", + "service", + "service.instance" + ] + ); + assert_eq!(declarations[1].id_columns, vec!["host.id"]); + assert_eq!(declarations[1].descriptive_columns, vec!["host.name"]); + assert_eq!(declarations[3].id_columns, vec!["job"]); + assert_eq!( + declarations[3].descriptive_columns, + vec!["service.name", "service.namespace"] + ); + assert_eq!(declarations[4].id_columns, vec!["job", "instance"]); + assert!(declarations[4].descriptive_columns.is_empty()); + + let partial = prom_table_info( + "greptime_otel_resource_info", + &["job", "service.name", "host.id"], + OTEL_STAMPS, + ); + let types: Vec = sorted_declarations(&partial) + .into_iter() + .map(|d| d.entity_type) + .collect(); + assert_eq!(types, vec!["host", "service"]); + } + + #[test] + fn otel_implicit_declarations_are_gated() { + let labels: &[&str] = &["job", "instance", "host.id"]; + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + PROM_STAMPS + )) + .is_empty() + ); + let mut stamps = OTEL_STAMPS.to_vec(); + stamps.push((PHYSICAL_TABLE_METADATA_KEY, "true")); + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + &stamps + )) + .is_empty() + ); + assert!( + conventions() + .unwrap() + .otel_info_metrics + .contains_key(servers::otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME) + ); + + let wrong_type = [ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY), + (SEMANTIC_METRIC_TYPE, servers::semantic::METRIC_TYPE_GAUGE), + ]; + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + &wrong_type + )) + .is_empty() + ); + } + #[test] fn target_info_descriptive_rest_covers_remaining_tags() { let info = prom_table_info( diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 08086198536..ccc09c9978c 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -50,7 +50,7 @@ use table::requests::{ use self::trace_ingest::trace_conventions; use crate::instance::Instance; -use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS}; +use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS}; fn trace_permission_targets( table_name: &str, @@ -119,9 +119,15 @@ impl OpenTelemetryProtocolHandler for Instance { .cloned() .unwrap_or_default(); metric_ctx.is_legacy = is_legacy; + metric_ctx.resource_info = self.otlp_resource_info; - let (requests, rows, semantic_index, mut outcome) = - otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; + let otlp::metrics::MetricsConversion { + requests, + rows, + semantic_index, + resource_info, + mut outcome, + } = otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; if outcome.rejected_data_points > 0 { warn!( "Rejected {} OTLP exponential histogram data points: {}", @@ -156,7 +162,7 @@ impl OpenTelemetryProtocolHandler for Instance { // OTLP tables have one sample field in both the legacy and physical paths. let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { - self.handle_row_inserts(requests, ctx, false, true) + self.handle_row_inserts(requests, ctx.clone(), false, true) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) @@ -165,12 +171,43 @@ impl OpenTelemetryProtocolHandler for Instance { .extension(PHYSICAL_TABLE_PARAM) .unwrap_or(GREPTIME_PHYSICAL_TABLE) .to_string(); - self.handle_metric_row_inserts(requests, ctx, physical_table) + self.handle_metric_row_inserts(requests, ctx.clone(), physical_table) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) }?; outcome.write_cost = output.meta.cost; + + // Derived enrichment, written after the metric data is committed: + // failing here would make the client retry data the server already + // accepted, so every failure degrades to a warning instead. + if let Some(resource_info) = resource_info { + let written = match self.check_row_insert_permission( + &resource_info, + &ctx, + PermissionReq::Action(OTLP_WRITE), + ) { + Ok(_) => self + .handle_row_inserts(resource_info, ctx, false, false) + .await + .map_err(BoxedError::new) + .map_err(|e| e.to_string()), + Err(e) => Err(e.to_string()), + }; + match written { + Ok(descriptor_output) => outcome.write_cost += descriptor_output.meta.cost, + Err(e) => { + OTLP_RESOURCE_INFO_WRITE_ERRORS.inc(); + warn!("Failed to write the OTLP resource descriptor table: {e}"); + outcome.error_message.get_or_insert(format!( + "metric data was accepted, but writing the resource \ + descriptor table `{}` failed: {e}", + otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME + )); + } + } + } + Ok(outcome) } diff --git a/src/frontend/src/metrics.rs b/src/frontend/src/metrics.rs index aba33637cfc..f649db0b46f 100644 --- a/src/frontend/src/metrics.rs +++ b/src/frontend/src/metrics.rs @@ -45,6 +45,14 @@ lazy_static! { ) .unwrap(); + /// Failed writes of the synthesized OTLP resource descriptor table; these + /// surface as OTLP partial-success warnings, not request failures. + pub static ref OTLP_RESOURCE_INFO_WRITE_ERRORS: IntCounter = register_int_counter!( + "greptime_frontend_otlp_resource_info_write_errors", + "frontend otlp resource descriptor write errors" + ) + .unwrap(); + /// The number of OpenTelemetry traces send by frontend node. pub static ref OTLP_TRACES_ROWS: IntCounter = register_int_counter!( "greptime_frontend_otlp_traces_rows", diff --git a/src/frontend/src/service_config/otlp.rs b/src/frontend/src/service_config/otlp.rs index 5dcab2c8ce2..d6a36f702e4 100644 --- a/src/frontend/src/service_config/otlp.rs +++ b/src/frontend/src/service_config/otlp.rs @@ -23,6 +23,11 @@ pub struct OtlpOptions { pub experimental_enable_exponential_histogram: bool, /// Maximum spans per trace ingest chunk. Set to 0 to disable splitting. pub trace_ingest_chunk_size: usize, + /// Whether to synthesize the `greptime_otel_resource_info` descriptor + /// table from the resource attributes of OTLP metrics, so metrics-only + /// services reach the semantic graph. Off by default: it creates and + /// writes a table the user did not send. + pub experimental_enable_resource_info: bool, } impl Default for OtlpOptions { @@ -31,6 +36,7 @@ impl Default for OtlpOptions { enable: true, experimental_enable_exponential_histogram: false, trace_ingest_chunk_size: DEFAULT_TRACE_INGEST_CHUNK_SIZE, + experimental_enable_resource_info: false, } } } @@ -45,6 +51,7 @@ mod tests { assert!(default.enable); assert!(!default.experimental_enable_exponential_histogram); assert_eq!(default.trace_ingest_chunk_size, 512); + assert!(!default.experimental_enable_resource_info); let options: OtlpOptions = toml::from_str("enable = false").unwrap(); assert!(!options.enable); diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index 7014a94f8ac..86827cf8c7a 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -41,8 +41,9 @@ use common_catalog::consts::{ ENTITY_DESCRIPTIVE_COLUMN, ENTITY_ID_ATTRS_COLUMN, ENTITY_ID_COLUMN, ENTITY_SCOPE_COLUMN, ENTITY_TYPE_COLUMN, ERROR_COUNT_COLUMN, FRESH_UNTIL_COLUMN, GENERATION_ID_COLUMN, OBSERVED_AT_COLUMN, PROVENANCE_COLUMN, REL_TYPE_COLUMN, REQUEST_COUNT_COLUMN, - SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN, SRC_ID_COLUMN, - SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, WINDOW_START_COLUMN, + SEMANTIC_GRAPH_WINDOW_NANOS, SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN, + SRC_ID_COLUMN, SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, + WINDOW_START_COLUMN, }; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; @@ -60,7 +61,7 @@ use datafusion::dataframe::DataFrame; use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns}; use datafusion::functions_nested::expr_fn::make_array; use datafusion_common::{Column, Result as DfResult, ScalarValue}; -use datafusion_expr::{Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; +use datafusion_expr::{Case, Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; pub use relationships::{ CallsSource, CoDeclaredSource, DeclaredSource, RelationshipSources, build_relationships_plan, }; @@ -122,9 +123,10 @@ pub fn declared_relationships_schema_matches(table_info: &table::metadata::Table }) } -/// Bin width for the temporal window of derived rows: 60s buckets, matching the -/// service-graph convention. -const BIN_NANOS: i64 = 60 * 1_000_000_000; +/// Bin width for the temporal window of derived rows, matching the +/// service-graph convention. Shared so ingestion-synthesized observations +/// land in the same buckets. +const BIN_NANOS: i64 = SEMANTIC_GRAPH_WINDOW_NANOS; /// Default retention for the declared-edge table; expiry slides the topology window. const DEFAULT_DECLARED_RELATIONSHIPS_TTL: &str = "90d"; @@ -284,8 +286,12 @@ pub struct EntityDeclaration { /// The table's time index column, used for the temporal window filter. pub time_index: String, pub entity_type: String, - /// Identifying columns (>= 1). One column → id verbatim; several → composite. + /// Identifying columns (>= 1), ordered broad to narrow. pub id_columns: Vec, + /// Optional column qualifying the first identity component, so a source + /// carrying the parts separately (a trace table's `service.namespace`) + /// yields the same id as one carrying them pre-composed (`job`). + pub id_qualifier: Option, /// Descriptive columns snapshotted into the `descriptive` JSON (may be empty). pub descriptive_columns: Vec, /// Scope columns (namespace/environment). One column → scope verbatim; @@ -417,28 +423,91 @@ fn cast_string_or_empty(column: &str) -> Expr { core_fns::coalesce().call(vec![cast(ident(column), DataType::Utf8), lit("")]) } -/// The canonical entity-id expression for `id_columns`: the value verbatim for -/// a single column, the sorted `k=v,k=v` rendering for a composite. `col` -/// constructs the column reference (unqualified for registry branches, -/// join-side-qualified for the calls derivation). /// A row identifies an entity only when every identity component is present /// and non-empty: kube-state-metrics descriptors emit empty-string labels (an /// unscheduled pod's `node`, an owner-less pod's `owner_*`), and an empty -/// string is never a meaningful entity id. +/// string is never a meaningful entity id. The qualifier is optional by +/// construction and so is not guarded. pub(crate) fn identifies(column: &str) -> Expr { ident(column) .is_not_null() .and(cast(ident(column), DataType::Utf8).not_eq(lit(""))) } -fn entity_id_expr(id_columns: &[String], col: &dyn Fn(&str) -> Expr) -> Expr { - if let [id] = id_columns { - cast(col(id), DataType::Utf8) - } else { - let mut cols = id_columns.to_vec(); - cols.sort(); - sorted_kv_expr_with(&cols, false, col) +const ID_SEPARATOR: &str = ","; +const ID_ESCAPE: &str = "\\"; +/// Left unescaped: `/` is how Prometheus renders `job`. +const ID_QUALIFIER_SEPARATOR: &str = "/"; + +/// Escapes so a composite id decodes back to its components. The escape +/// character goes first, or it would double the escapes the separator pass +/// introduces. +fn escaped_id_value(value: Expr) -> Expr { + let escape_escape = string_fns::replace().call(vec![ + value, + lit(ID_ESCAPE), + lit(format!("{ID_ESCAPE}{ID_ESCAPE}")), + ]); + string_fns::replace().call(vec![ + escape_escape, + lit(ID_SEPARATOR), + lit(format!("{ID_ESCAPE}{ID_SEPARATOR}")), + ]) +} + +/// The identity values in declared order (broad to narrow), escaped and +/// joined. +/// +/// The identifying *column names* are deliberately absent: the same identity +/// reaches us under different names per source — a trace table's +/// `service_name` against a metric table's `job` — and encoding them would +/// split one entity into one per signal. `entity_id_attrs` keeps the +/// structured form. +/// +/// `col` constructs the column reference (unqualified for registry branches, +/// join-side-qualified for the calls derivation). +fn entity_id_expr( + id_columns: &[String], + qualifier: Option<&str>, + col: &dyn Fn(&str) -> Expr, +) -> Expr { + let mut parts = Vec::with_capacity(id_columns.len() * 2); + for (index, column) in id_columns.iter().enumerate() { + if index > 0 { + parts.push(lit(ID_SEPARATOR)); + } + let value = escaped_id_value(cast(col(column), DataType::Utf8)); + parts.push(match qualifier { + Some(qualifier) if index == 0 => qualified_id_expr(qualifier, value, col), + _ => value, + }); } + if let [single] = parts.as_slice() { + single.clone() + } else { + concat_expr(parts) + } +} + +/// `/`, or bare `value` when the qualifier is empty on this +/// row — the spec's rule composing `job` from `service.namespace` and +/// `service.name`. +fn qualified_id_expr(qualifier: &str, value: Expr, col: &dyn Fn(&str) -> Expr) -> Expr { + let qualifier = escaped_id_value( + core_fns::coalesce().call(vec![cast(col(qualifier), DataType::Utf8), lit("")]), + ); + Expr::Case(Case::new( + None, + vec![( + Box::new(qualifier.clone().eq(lit(""))), + Box::new(value.clone()), + )], + Some(Box::new(concat_expr(vec![ + qualifier, + lit(ID_QUALIFIER_SEPARATOR), + value, + ]))), + )) } /// The `parse_json` UDF, shared by all derivation plans. Resolved from the @@ -583,14 +652,22 @@ fn registry_source( let mut rows = Vec::with_capacity(1 + rest.len()); for decl in std::iter::once(first).chain(rest) { // CAST even a single-column id: id columns need not be strings, and - // the computed table declares entity_id STRING. Composite ids - // additionally carry a JSON object of the id columns in - // entity_id_attrs. - let entity_id = entity_id_expr(&decl.id_columns, &|c| ident(c)); - let entity_id_attrs = if decl.id_columns.len() == 1 { + // the computed table declares entity_id STRING. An id assembled from + // more than one column additionally carries its parts as a JSON + // object in entity_id_attrs, the structured form of the id string. + let entity_id = entity_id_expr(&decl.id_columns, decl.id_qualifier.as_deref(), &|c| { + ident(c) + }); + let id_parts = decl + .id_qualifier + .iter() + .chain(&decl.id_columns) + .cloned() + .collect::>(); + let entity_id_attrs = if id_parts.len() == 1 { null_json() } else { - json_object_expr(&decl.id_columns) + json_object_expr(&id_parts) }; let scope = match decl.scope_columns.as_slice() { @@ -811,11 +888,74 @@ mod tests { time_index: "ts".to_string(), entity_type: entity_type.to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } } + /// The id string must decode back to its components, and a qualifier must + /// reproduce the identity a source that pre-composed it already emits. + #[tokio::test] + async fn registry_id_escaping_and_qualifier() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("service_name", DataType::Utf8, false), + Field::new("instance", DataType::Utf8, false), + Field::new("namespace", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 2_000, 3_000])) as ArrayRef, + Arc::new(StringArray::from(vec!["cart", "a,b", "we\\ird"])), + Arc::new(StringArray::from(vec!["i-1", "c", "i-3"])), + Arc::new(StringArray::from(vec![Some("shop"), None, Some("")])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "svc", + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + + let mut declaration = decl("service.instance", &["service_name", "instance"]); + declaration.id_qualifier = Some("namespace".to_string()); + declaration.time_index = "ts".to_string(); + let plan = build_registry_plan( + vec![RegistrySource { + declarations: vec![declaration], + scan: ctx.table("svc").await.unwrap(), + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let mut ids: Vec = collect(&ctx, plan) + .await + .iter() + .flat_map(|batch| strings(batch, 5)) + .collect(); + ids.sort(); + assert_eq!( + ids, + vec![ + // an absent qualifier leaves the identity bare, and a value + // holding the separator stays distinguishable from two values + "a\\,b,c".to_string(), + "shop/cart,i-1".to_string(), + "we\\\\ird,i-3".to_string(), + ] + ); + } + #[tokio::test] async fn registry_single_column_identity() { let ctx = metric_table_ctx(); @@ -900,24 +1040,24 @@ mod tests { .collect(); rows.sort(); - // Composite id -> sorted `k=v,k=v` plus a JSON object of the id columns; - // descriptive JSON keeps `\`, `"` and control characters intact in - // runtime values, NULL -> "". + // Composite id -> values in declared order plus a JSON object of the + // id columns; descriptive JSON keeps `\`, `"` and control characters + // intact in runtime values, NULL -> "". assert_eq!( rows, vec![ ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":""}"#.to_string()), ), ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":"h2"}"#.to_string()), ), ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":"we\"ird\\\nhost"}"#.to_string()), ), diff --git a/src/operator/src/statement/semantic_graph/conventions.rs b/src/operator/src/statement/semantic_graph/conventions.rs index 36d058921b5..9ec6f52e158 100644 --- a/src/operator/src/statement/semantic_graph/conventions.rs +++ b/src/operator/src/statement/semantic_graph/conventions.rs @@ -19,7 +19,7 @@ //! with the binary, not an operator-editable configuration surface; explicit //! `greptime.semantic.entity.*` declarations always override it. -use std::collections::{BTreeMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::LazyLock; use serde::Deserialize; @@ -43,15 +43,20 @@ pub struct VirtualDstCandidate { pub connection_type: String, } -/// One implicit entity declaration: of a whitelisted Prometheus info metric, -/// or of a trace-v1 table's flattened resource attributes. +/// One implicit entity declaration: of a whitelisted Prometheus or OTel info +/// metric, or of a trace-v1 table's flattened resource attributes. #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] pub struct ImplicitEntity { pub entity: String, - /// Identifying label columns; every one must exist on the table for the - /// declaration to apply. + /// Identifying label columns, ordered broad to narrow; every one must + /// exist on the table for the declaration to apply. pub id: Vec, + /// Column qualifying `id[0]` as `/`, so a source + /// carrying the parts separately matches one carrying them pre-composed. + /// Skipped when the column is absent or empty. + #[serde(default)] + pub qualified_by: Option, /// Descriptive label columns, filtered to those present (kube-state-metrics /// label sets vary across versions). #[serde(default)] @@ -70,7 +75,12 @@ pub struct Conventions { pub trace_co_declared_edges: Vec, pub virtual_dst_candidates: Vec, pub otlp_trace_entities: Vec, + /// Table name -> the entities that table declares, for Prometheus-sourced + /// descriptor metrics (`source = prometheus`). pub prometheus_info_metrics: BTreeMap>, + /// Table name -> the entities that table declares, for OTLP-sourced + /// descriptor tables (`source = opentelemetry`). + pub otel_info_metrics: BTreeMap>, } /// The built-in entity-type vocabulary. User-declared types are open-ended; @@ -78,6 +88,7 @@ pub struct Conventions { pub const ENTITY_TYPE_SERVICE: &str = "service"; pub const ENTITY_TYPE_SERVICE_INSTANCE: &str = "service.instance"; pub const ENTITY_TYPE_HOST: &str = "host"; +pub const ENTITY_TYPE_CONTAINER: &str = "container"; pub const ENTITY_TYPE_PROCESS: &str = "process"; pub const ENTITY_TYPE_K8S_POD: &str = "k8s.pod"; pub const ENTITY_TYPE_K8S_NODE: &str = "k8s.node"; @@ -88,10 +99,11 @@ pub const ENTITY_TYPE_GEN_AI_AGENT: &str = "gen_ai.agent"; pub const ENTITY_TYPE_GEN_AI_MODEL: &str = "gen_ai.model"; pub const ENTITY_TYPE_GEN_AI_TOOL: &str = "gen_ai.tool"; -const ENTITY_TYPES: [&str; 12] = [ +const ENTITY_TYPES: [&str; 13] = [ ENTITY_TYPE_SERVICE, ENTITY_TYPE_SERVICE_INSTANCE, ENTITY_TYPE_HOST, + ENTITY_TYPE_CONTAINER, ENTITY_TYPE_PROCESS, ENTITY_TYPE_K8S_POD, ENTITY_TYPE_K8S_NODE, @@ -200,11 +212,16 @@ fn validate(conventions: &Conventions) -> Result<(), String> { let per_table = conventions .prometheus_info_metrics .iter() + .chain(&conventions.otel_info_metrics) .map(|(table, entities)| (table.as_str(), entities)) .chain(std::iter::once(( "otlp traces", &conventions.otlp_trace_entities, ))); + // An entity type declared with a different number of id columns by two + // sources yields ids that can never match, silently splitting one entity + // in two. + let mut arity = HashMap::new(); for (table, entities) in per_table { let mut seen_types = HashSet::new(); for implicit in entities { @@ -226,6 +243,12 @@ fn validate(conventions: &Conventions) -> Result<(), String> { implicit.entity )); } + if implicit.qualified_by.as_ref().is_some_and(String::is_empty) { + return Err(format!( + "entity `{}` of info metric `{table}` has an empty qualified_by", + implicit.entity + )); + } if implicit.descriptive_rest && !implicit.descriptive.is_empty() { return Err(format!( "entity `{}` of info metric `{table}` sets both descriptive and \ @@ -233,6 +256,17 @@ fn validate(conventions: &Conventions) -> Result<(), String> { implicit.entity )); } + match arity.insert(implicit.entity.as_str(), implicit.id.len()) { + Some(previous) if previous != implicit.id.len() => { + return Err(format!( + "entity `{}` is declared with {previous} id columns elsewhere but \ + {} for `{table}`", + implicit.entity, + implicit.id.len() + )); + } + _ => {} + } } } Ok(()) @@ -249,25 +283,28 @@ mod tests { /// Each case must fail on exactly the rule it names, so the shared /// boilerplate is valid and the mutated part is inside the vocabulary. - fn broken(edges: &str, info_metrics: &str) -> String { + fn broken(edges: &str, info_metrics: &str, otel_metrics: &str) -> String { format!( "co_declared_edges: [{edges}]\ntrace_co_declared_edges: []\n\ virtual_dst_candidates: []\notlp_trace_entities: []\n\ - prometheus_info_metrics: {{{info_metrics}}}" + prometheus_info_metrics: {{{info_metrics}}}\n\ + otel_info_metrics: {{{otel_metrics}}}" ) } #[test] fn validation_rejects_broken_conventions() { - let err = |edges, info| parse(&broken(edges, info)).unwrap_err(); + let err = |edges, info, otel| parse(&broken(edges, info, otel)).unwrap_err(); - assert!(err("{src: host, dst: service, rel: pets}", "").contains("unknown rel_type")); + assert!(err("{src: host, dst: service, rel: pets}", "", "").contains("unknown rel_type")); assert!( - err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "").contains("unknown entity type") + err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "", "") + .contains("unknown entity type") ); assert!( err( "{src: host, dst: service, rel: uses}, {src: host, dst: service, rel: uses}", + "", "" ) .contains("duplicate edge rule") @@ -275,14 +312,27 @@ mod tests { assert!( err( "", - "t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]" + "t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]", + "" ) .contains("descriptive_rest") ); + // the otel map runs through the same per-table validation + assert!(err("", "", "t: [{entity: hosts, id: [x]}]").contains("unknown entity type")); + // ids of a different arity for one type can never match each other + assert!( + err( + "", + "a: [{entity: host, id: [x]}]", + "b: [{entity: host, id: [x, y]}]" + ) + .contains("id columns") + ); // Unknown YAML keys are rejected, catching typos in the embedded file. assert!( parse(&broken( "{src: host, dst: service, rel: uses, direction: down}", + "", "" )) .is_err() diff --git a/src/operator/src/statement/semantic_graph/conventions.yaml b/src/operator/src/statement/semantic_graph/conventions.yaml index 9fc37fe2a02..808fe0f2658 100644 --- a/src/operator/src/statement/semantic_graph/conventions.yaml +++ b/src/operator/src/statement/semantic_graph/conventions.yaml @@ -2,6 +2,10 @@ # versioned with the binary (include_str!), not an operator-editable surface; # non-standard deployments override per table with explicit # greptime.semantic.entity.* declarations, which always win. +# +# Every `id` list is ordered broad to narrow, and that order is part of the +# identity: the id renders as the values joined in this order, so all +# declarations of one entity type must agree on it or they name two entities. co_declared_edges: - { src: service.instance, dst: host, rel: runs_on } @@ -11,6 +15,8 @@ co_declared_edges: - { src: service.instance, dst: service, rel: part_of } - { src: k8s.pod, dst: k8s.workload, rel: part_of } - { src: service.instance, dst: k8s.pod, rel: runs_on } + - { src: service.instance, dst: container, rel: runs_on } + - { src: container, dst: host, rel: runs_on } # Applied only to trace sources; the endpoints are Greptime entity types # derived from GenAI semantic-convention attributes. @@ -30,12 +36,25 @@ virtual_dst_candidates: - { column: span_attributes.server.address, connection_type: virtual_node } # Implicit declarations for greptime_trace_v1 tables; the pod UID bridges -# trace-side pods onto the kube-state-metrics entities below. +# trace-side pods onto the kube-state-metrics entities below. host.name and +# container.name vary by SDK and resource detector, so identity is the id. +# qualified_by folds the namespace into the service identity the same way the +# metric sources' `job` already carries it, so both signals name one entity. otlp_trace_entities: - entity: service id: [service_name] + qualified_by: resource_attributes.service.namespace - entity: service.instance id: [service_name, resource_attributes.service.instance.id] + qualified_by: resource_attributes.service.namespace + - entity: host + id: [resource_attributes.host.id] + descriptive: + - resource_attributes.host.name + - entity: container + id: [resource_attributes.container.id] + descriptive: + - resource_attributes.container.name - entity: k8s.pod id: [resource_attributes.k8s.pod.uid] descriptive: @@ -113,3 +132,32 @@ prometheus_info_metrics: - entity: service.instance id: [job, instance] descriptive_rest: true + +# Implicit declarations for the OTLP resource descriptor synthesized at +# metrics ingestion, gated on source=opentelemetry and metric.type=info. +# job/instance are the derived Prometheus-compatible identity, so services +# line up with target_info. Unlike target_info the attribute set is known, so +# descriptive columns are listed instead of swept up by descriptive_rest, +# which would misfile other entities' ids under service.instance. +otel_info_metrics: + greptime_otel_resource_info: + - entity: service + id: [job] + descriptive: + - service.name + - service.namespace + - entity: service.instance + id: [job, instance] + - entity: host + id: [host.id] + descriptive: + - host.name + - entity: container + id: [container.id] + descriptive: + - container.name + - entity: k8s.pod + id: [k8s.pod.uid] + descriptive: + - k8s.pod.name + - k8s.namespace.name diff --git a/src/operator/src/statement/semantic_graph/relationships.rs b/src/operator/src/statement/semantic_graph/relationships.rs index 3b5c6d37b78..257c844a9a8 100644 --- a/src/operator/src/statement/semantic_graph/relationships.rs +++ b/src/operator/src/statement/semantic_graph/relationships.rs @@ -200,9 +200,9 @@ fn co_declared_branch( bin.clone() + bin_interval(), bin.clone() + bin_interval(), lit(src.entity_type.as_str()), - entity_id_expr(&src.id_columns, &|c| ident(c)), + entity_id_expr(&src.id_columns, src.id_qualifier.as_deref(), &|c| ident(c)), lit(dst.entity_type.as_str()), - entity_id_expr(&dst.id_columns, &|c| ident(c)), + entity_id_expr(&dst.id_columns, dst.id_qualifier.as_deref(), &|c| ident(c)), lit(*rel_type), lit(*provenance), lit(1.0_f64), @@ -501,6 +501,13 @@ fn span_predicate(service: &EntityDeclaration, window: &GraphQueryWindow, strict } /// One trace table's client spans, normalized to +/// A trace table's `duration_nano` is UInt64 on tables created before the +/// signed-integer ingest change and Int64 after it. The per-table selects are +/// unioned, and those two have no common integer type. +fn duration_nano_expr() -> Expr { + cast(ident(DURATION_NANO_COLUMN), DataType::Int64).alias(DURATION_NANO_COLUMN) +} + /// `(timestamp, trace_id, span_id, src_id, status_code, duration_nano, /// virtual_dst, virtual_conn)`. `src_id` is built from the table's `service` /// declaration, so edges land on exactly the entity ids the registry emits (a @@ -553,9 +560,12 @@ fn client_spans( ident(SPAN_ID_COLUMN), // The cast inside entity_id_expr also normalizes tag columns, which // come out of the storage engine dictionary-encoded. - entity_id_expr(&service.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), + entity_id_expr(&service.id_columns, service.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(SRC_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), virtual_dst.alias("virtual_dst"), virtual_conn.alias("virtual_conn"), ]) @@ -578,9 +588,12 @@ fn server_spans( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&service.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), + entity_id_expr(&service.id_columns, service.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), ]) } @@ -609,7 +622,10 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), + entity_id_expr(&agent.id_columns, agent.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(SRC_ID_COLUMN), ])?; let child = trace .scan @@ -619,9 +635,12 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), + entity_id_expr(&agent.id_columns, agent.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), ])?; parents = union_all(parents, parent)?; children = union_all(children, child)?; @@ -732,6 +751,16 @@ mod tests { name: &str, extra: &[(&str, &[Option<&str>])], spans: &[Span<'_>], + ) { + register_typed_trace_table(ctx, name, extra, spans, DataType::UInt64) + } + + fn register_typed_trace_table( + ctx: &SessionContext, + name: &str, + extra: &[(&str, &[Option<&str>])], + spans: &[Span<'_>], + duration_type: DataType, ) { let mut fields = vec![ Field::new( @@ -745,7 +774,7 @@ mod tests { Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false), Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false), Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false), - Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false), + Field::new(DURATION_NANO_COLUMN, duration_type.clone(), false), ]; for (column, _) in extra { fields.push(Field::new(*column, DataType::Utf8, true)); @@ -774,9 +803,14 @@ mod tests { Arc::new(StringArray::from( spans.iter().map(|s| s.6).collect::>(), )), - Arc::new(UInt64Array::from( - spans.iter().map(|s| s.7).collect::>(), - )), + match duration_type { + DataType::Int64 => Arc::new(Int64Array::from( + spans.iter().map(|s| s.7 as i64).collect::>(), + )) as ArrayRef, + _ => Arc::new(UInt64Array::from( + spans.iter().map(|s| s.7).collect::>(), + )), + }, ]; for (_, values) in extra { columns.push(Arc::new(StringArray::from(values.to_vec()))); @@ -843,6 +877,7 @@ mod tests { time_index: TRACE_TIMESTAMP_COLUMN.to_string(), entity_type: "service".to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } @@ -972,7 +1007,7 @@ mod tests { ), ], ); - register_trace_table( + register_typed_trace_table( &ctx, "trace_b", &[], @@ -989,6 +1024,7 @@ mod tests { 1_500_000_000, ), ], + DataType::Int64, ); let a = ctx.table("trace_a").await.unwrap(); let b = ctx.table("trace_b").await.unwrap(); @@ -1294,6 +1330,7 @@ mod tests { time_index: "ts".to_string(), entity_type: entity_type.to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } @@ -1544,16 +1581,10 @@ mod tests { let total: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(total, 1); let batch = &batches[0]; - // Composite service identity renders the same sorted `k=v` form the - // registry emits, so edges land on registry entity ids. - assert_eq!( - strings(batch, 5), - vec!["service_name=frontend,service_namespace=ns1"] - ); - assert_eq!( - strings(batch, 7), - vec!["service_name=cart,service_namespace=ns1"] - ); + // Composite service identity renders exactly as the registry emits it, + // so edges land on registry entity ids. + assert_eq!(strings(batch, 5), vec!["frontend,ns1"]); + assert_eq!(strings(batch, 7), vec!["cart,ns1"]); } fn build_relationships_plan_for_test( diff --git a/src/servers/src/http/otlp.rs b/src/servers/src/http/otlp.rs index d04022cdf88..3ecaba74a4e 100644 --- a/src/servers/src/http/otlp.rs +++ b/src/servers/src/http/otlp.rs @@ -118,8 +118,9 @@ pub async fn metrics( promote_scope_attrs: http_opts.promote_scope_attrs, with_metric_engine, experimental_enable_exponential_histogram, - // set is_legacy later + // set by the frontend from its config is_legacy: false, + resource_info: false, metric_type: MetricType::Init, metric_translation_strategy: http_opts.metric_translation_strategy, })); @@ -128,7 +129,7 @@ pub async fn metrics( handler.metrics(request, query_ctx).await.map(|outcome| { if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 { OtlpMetricsResponse::Failure(outcome) - } else if outcome.rejected_data_points > 0 { + } else if outcome.rejected_data_points > 0 || outcome.error_message.is_some() { OtlpMetricsResponse::PartialSuccess(outcome) } else { OtlpMetricsResponse::FullSuccess(outcome) diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index bd0bfb2ce77..d65e117e07c 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -25,6 +25,7 @@ use common_query::native_histogram::{ }; use common_query::prelude::{GREPTIME_COUNT, greptime_timestamp, greptime_value}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; +use common_telemetry::warn; use lazy_static::lazy_static; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; @@ -36,7 +37,7 @@ use table::requests::{ }; use crate::error::{self, Result}; -use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME}; +use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME, KEY_SERVICE_NAMESPACE}; use crate::query_handler::MetricsIngestOutcome; use crate::row_writer::{self, MultiTableData, TableData}; pub use crate::semantic::SemanticIndex; @@ -45,8 +46,11 @@ use crate::semantic::{ METRIC_TYPE_UPDOWN_COUNTER, }; +mod resource_info; mod translator; +pub use resource_info::OTEL_RESOURCE_INFO_TABLE_NAME; +use resource_info::ResourceInfoData; pub use translator::legacy_normalize_otlp_name; pub(crate) use translator::ucum_to_openmetrics_unit; use translator::{translate_label_name, translate_metric_name}; @@ -96,28 +100,43 @@ const MIN_EXPONENTIAL_HISTOGRAM_SCALE: i32 = -4; const MAX_EXPONENTIAL_HISTOGRAM_SCALE: i32 = 8; const MAX_REJECTION_MESSAGE_BYTES: usize = 512; +/// Result of converting one OTLP metrics request. +#[derive(Debug)] +pub struct MetricsConversion { + pub requests: RowInsertRequests, + /// Row count of `requests`; the resource descriptor is not counted. + pub rows: usize, + /// Per-table semantic index for the auto-create path to stamp as table + /// options; covers the descriptor table too. + pub semantic_index: SemanticIndex, + /// The synthesized resource descriptor, written separately from the + /// metric data. See [`resource_info`]. + pub resource_info: Option, + pub outcome: MetricsIngestOutcome, +} + /// Convert OpenTelemetry metrics to GreptimeDB insert requests /// /// See /// /// for data structure of OTLP metrics. -/// -/// Returns `InsertRequests`, total number of rows to ingest, and the per-table -/// semantic index for the auto-create path to stamp as table options. pub fn to_grpc_insert_requests( request: ExportMetricsServiceRequest, metric_ctx: &mut OtlpMetricCtx, -) -> Result<( - RowInsertRequests, - usize, - SemanticIndex, - MetricsIngestOutcome, -)> { +) -> Result { let mut table_writer = MultiTableData::default(); let mut semantic_index = SemanticIndex::default(); let mut outcome = MetricsIngestOutcome::default(); + let mut resource_info = ResourceInfoData::default(); for resource in &request.resource_metrics { + if metric_ctx.resource_info + && !metric_ctx.is_legacy + && let Some(r) = resource.resource.as_ref() + { + resource_info.observe(&r.attributes, resource, metric_ctx); + } + let resource_attrs = resource.resource.as_ref().map(|r| { let mut attrs = r.attributes.clone(); process_resource_attrs(&mut attrs, metric_ctx); @@ -149,8 +168,46 @@ pub fn to_grpc_insert_requests( } let (requests, rows) = table_writer.into_row_insert_requests(); + validate_sample_kinds(&requests)?; - Ok((requests, rows, semantic_index, outcome)) + + // The metric and the descriptor would fight over the same table across + // two engines. The user's metric wins. + let resource_info = if !metric_ctx.resource_info { + None + } else if requests + .inserts + .iter() + .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) + { + warn!( + "Skipping OTLP resource descriptor synthesis: the request writes \ + a metric named `{OTEL_RESOURCE_INFO_TABLE_NAME}`" + ); + None + } else { + resource_info.into_row_insert_requests()? + }; + if resource_info.is_some() { + semantic_index.record_scalar( + OTEL_RESOURCE_INFO_TABLE_NAME, + SEMANTIC_METRIC_TYPE, + crate::semantic::METRIC_TYPE_INFO, + ); + semantic_index.record_scalar( + OTEL_RESOURCE_INFO_TABLE_NAME, + SEMANTIC_METRIC_METADATA_QUALITY, + METADATA_QUALITY_DECLARED, + ); + } + + Ok(MetricsConversion { + requests, + rows, + semantic_index, + outcome, + resource_info, + }) } fn validate_sample_kinds(requests: &RowInsertRequests) -> Result<()> { @@ -287,30 +344,61 @@ fn from_metric_type(data: &metric::Data) -> MetricType { } } +/// Non-scalar values (bool, arrays, maps, bytes) are not representable as tags. +fn scalar_value_string(value: Option<&AnyValue>) -> Option { + match value.and_then(|v| v.value.as_ref())? { + any_value::Value::StringValue(s) => Some(s.clone()), + any_value::Value::IntValue(v) => Some(v.to_string()), + any_value::Value::DoubleValue(v) => Some(v.to_string()), + _ => None, + } +} + +/// Prometheus-style `(job, instance)` identity. Per the OTel Prometheus +/// compatibility spec `job` folds in `service.namespace` when present; a +/// resource without `service.name` gets no job rather than a fabricated one. +/// Both fields are optional strings, so a named type keeps a swapped +/// destructuring from type-checking at the call sites. +pub(crate) struct ServiceIdentity { + pub job: Option, + pub instance: Option, +} + +pub(crate) fn service_identity(attrs: &[KeyValue]) -> ServiceIdentity { + let mut name = None; + let mut namespace = None; + let mut instance = None; + for kv in attrs { + match kv.key.as_str() { + KEY_SERVICE_NAME => name = scalar_value_string(kv.value.as_ref()), + KEY_SERVICE_NAMESPACE => namespace = scalar_value_string(kv.value.as_ref()), + KEY_SERVICE_INSTANCE_ID => instance = scalar_value_string(kv.value.as_ref()), + _ => {} + } + } + let job = name.map(|name| match namespace { + Some(ns) if !ns.is_empty() => format!("{ns}/{name}"), + _ => name, + }); + ServiceIdentity { job, instance } +} + +fn string_key_value(key: &str, value: String) -> KeyValue { + KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value)), + }), + } +} + fn process_resource_attrs(attrs: &mut Vec, metric_ctx: &OtlpMetricCtx) { if metric_ctx.is_legacy { return; } - // remap service.name and service.instance.id to job and instance - let mut tmp = Vec::with_capacity(2); - for kv in attrs.iter() { - match &kv.key as &str { - KEY_SERVICE_NAME => { - tmp.push(KeyValue { - key: JOB_KEY.to_string(), - value: kv.value.clone(), - }); - } - KEY_SERVICE_INSTANCE_ID => { - tmp.push(KeyValue { - key: INSTANCE_KEY.to_string(), - value: kv.value.clone(), - }); - } - _ => {} - } - } + // remap the service identity attributes to job and instance + let ServiceIdentity { job, instance } = service_identity(attrs); // if promote all, then exclude the list, else, include the list if metric_ctx.promote_all_resource_attrs { @@ -322,7 +410,12 @@ fn process_resource_attrs(attrs: &mut Vec, metric_ctx: &OtlpMetricCtx) }); } - attrs.extend(tmp); + if let Some(job) = job { + attrs.push(string_key_value(JOB_KEY, job)); + } + if let Some(instance) = instance { + attrs.push(string_key_value(INSTANCE_KEY, instance)); + } } fn process_scope_attrs(scope: &ScopeMetrics, metric_ctx: &OtlpMetricCtx) -> Option> { @@ -536,35 +629,13 @@ fn encode_exponential_histogram( metric_ctx: &OtlpMetricCtx, outcome: &mut MetricsIngestOutcome, ) -> Result { - if !metric_ctx.experimental_enable_exponential_histogram { + if let Err(rejection) = exponential_histogram_gate(histogram, metric_ctx) { reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" - ) + rejection.message(name) })?; return Ok(false); } - match AggregationTemporality::try_from(histogram.aggregation_temporality) { - Ok(AggregationTemporality::Cumulative) => {} - Ok(AggregationTemporality::Delta) => { - reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" - ) - })?; - return Ok(false); - } - _ => { - reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` has unspecified OTLP exponential histogram temporality; cumulative temporality is required" - ) - })?; - return Ok(false); - } - } - let column_schema = native_histogram_column_schema().map_err(|error| { error::InternalSnafu { err_msg: error.to_string(), @@ -611,7 +682,46 @@ fn encode_exponential_histogram( Ok(emitted) } -fn exponential_histogram_value( +pub(crate) enum ExponentialHistogramRejection { + Disabled, + DeltaTemporality, + UnspecifiedTemporality, +} + +impl ExponentialHistogramRejection { + fn message(&self, name: &str) -> String { + match self { + Self::Disabled => format!( + "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" + ), + Self::DeltaTemporality => format!( + "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" + ), + Self::UnspecifiedTemporality => format!( + "metric `{name}` has unspecified OTLP exponential histogram temporality; cumulative temporality is required" + ), + } + } +} + +/// Whole-metric acceptance, decided once for the encoder and for the resource +/// descriptor, which must not describe a resource whose data was rejected. +/// Individual points can still fail [`exponential_histogram_value`]. +pub(crate) fn exponential_histogram_gate( + histogram: &ExponentialHistogram, + metric_ctx: &OtlpMetricCtx, +) -> std::result::Result<(), ExponentialHistogramRejection> { + if !metric_ctx.experimental_enable_exponential_histogram { + return Err(ExponentialHistogramRejection::Disabled); + } + match AggregationTemporality::try_from(histogram.aggregation_temporality) { + Ok(AggregationTemporality::Cumulative) => Ok(()), + Ok(AggregationTemporality::Delta) => Err(ExponentialHistogramRejection::DeltaTemporality), + _ => Err(ExponentialHistogramRejection::UnspecifiedTemporality), + } +} + +pub(crate) fn exponential_histogram_value( data_point: &ExponentialHistogramDataPoint, ) -> std::result::Result<(ValueData, i64), String> { if data_point.start_time_unix_nano > data_point.time_unix_nano { @@ -823,29 +933,21 @@ fn write_attributes( }; let tags = attrs.iter().filter_map(|attr| { - attr.value - .as_ref() - .and_then(|v| v.value.as_ref()) - .and_then(|val| { - let key = match attribute_type { - AttributeType::Resource | AttributeType::DataPoint => { - translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) - } - AttributeType::Scope => { - format!( - "otel_scope_{}", - translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) - ) - } - AttributeType::Legacy => legacy_normalize_otlp_name(&attr.key), - }; - match val { - any_value::Value::StringValue(s) => Some((key, s.clone())), - any_value::Value::IntValue(v) => Some((key, v.to_string())), - any_value::Value::DoubleValue(v) => Some((key, v.to_string())), - _ => None, // TODO(sunng87): allow different type of values - } - }) + // TODO(sunng87): allow different type of values + let value = scalar_value_string(attr.value.as_ref())?; + let key = match attribute_type { + AttributeType::Resource | AttributeType::DataPoint => { + translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) + } + AttributeType::Scope => { + format!( + "otel_scope_{}", + translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) + ) + } + AttributeType::Legacy => legacy_normalize_otlp_name(&attr.key), + }; + Some((key, value)) }); row_writer::write_tags(writer, tags, row)?; @@ -1295,6 +1397,169 @@ mod tests { } } + fn descriptor_ctx() -> OtlpMetricCtx { + OtlpMetricCtx { + resource_info: true, + ..Default::default() + } + } + + fn attr_value(attrs: &[KeyValue], key: &str) -> Option { + attrs + .iter() + .find(|kv| kv.key == key) + .and_then(|kv| scalar_value_string(kv.value.as_ref())) + } + + fn gauge_request( + resource_attrs: Vec, + metric_name: &str, + ) -> ExportMetricsServiceRequest { + use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: Some(Resource { + attributes: resource_attrs, + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: metric_name.to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano: 1_000_000, + value: Some(Value::AsInt(1)), + ..Default::default() + }], + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + } + } + + fn column_names(request: &RowInsertRequests, table: &str) -> Vec { + request + .inserts + .iter() + .find(|r| r.table_name == table) + .unwrap_or_else(|| panic!("missing table {table}")) + .rows + .as_ref() + .unwrap() + .schema + .iter() + .map(|c| c.column_name.clone()) + .collect() + } + + #[test] + fn test_conversion_synthesizes_resource_descriptor() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api"), keyvalue("host.id", "h-1")], + "my_gauge", + ); + let conversion = to_grpc_insert_requests(request, &mut descriptor_ctx()).unwrap(); + + // descriptor keeps raw OTel keys while the metric table's labels went + // through the (default underscore-escaping) translation strategy + let resource_info = conversion.resource_info.expect("descriptor synthesized"); + let descriptor_cols = column_names(&resource_info, OTEL_RESOURCE_INFO_TABLE_NAME); + assert!(descriptor_cols.contains(&"host.id".to_string())); + assert!(descriptor_cols.contains(&"service.name".to_string())); + assert!(descriptor_cols.contains(&"job".to_string())); + let metric_cols = column_names(&conversion.requests, "my_gauge"); + assert!(metric_cols.contains(&"service_name".to_string())); + assert!(!metric_cols.contains(&"service.name".to_string())); + // host.id is not in the promote list: only the descriptor keeps it + assert!(!metric_cols.contains(&"host_id".to_string())); + + let decoded = decode(&conversion.semantic_index); + let t = &decoded[OTEL_RESOURCE_INFO_TABLE_NAME]; + assert_eq!( + t.get(SEMANTIC_METRIC_TYPE).map(String::as_str), + Some("info") + ); + assert_eq!( + t.get(SEMANTIC_METRIC_METADATA_QUALITY).map(String::as_str), + Some("declared") + ); + } + + #[test] + fn test_conversion_skips_descriptor_for_legacy_mode() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api"), keyvalue("host.id", "h-1")], + "my_gauge", + ); + let mut ctx = OtlpMetricCtx { + is_legacy: true, + ..descriptor_ctx() + }; + let conversion = to_grpc_insert_requests(request, &mut ctx).unwrap(); + assert!(conversion.resource_info.is_none()); + + // legacy tables predate job/instance and the promote filter; adding + // either would alter the schema of tables already in use + let cols = column_names(&conversion.requests, "my_gauge"); + assert!(!cols.contains(&"job".to_string())); + assert!(cols.contains(&"host_id".to_string())); + } + + #[test] + fn test_conversion_skips_descriptor_on_metric_name_collision() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api")], + OTEL_RESOURCE_INFO_TABLE_NAME, + ); + let conversion = to_grpc_insert_requests(request, &mut descriptor_ctx()).unwrap(); + assert!(conversion.resource_info.is_none()); + // the metric itself still goes through the main path + assert!( + conversion + .requests + .inserts + .iter() + .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) + ); + } + + #[test] + fn test_job_composition_follows_service_namespace() { + let mut attrs = vec![ + keyvalue("service.name", "api"), + keyvalue("service.namespace", "shop"), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("shop/api")); + + let mut attrs = vec![keyvalue("service.name", "api")]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("api")); + + let mut attrs = vec![ + keyvalue("service.name", "api"), + keyvalue("service.namespace", ""), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("api")); + + // no service.name: no job is fabricated + let mut attrs = vec![ + keyvalue("service.namespace", "shop"), + keyvalue("service.instance.id", "inst-1"), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job"), None); + assert_eq!(attr_value(&attrs, "instance").as_deref(), Some("inst-1")); + } + #[test] fn test_encode_gauge() { let mut tables = MultiTableData::default(); @@ -2104,8 +2369,12 @@ mod tests { AggregationTemporality::Cumulative, ), ]); - let (requests, _, semantic_index, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let MetricsConversion { + requests, + semantic_index, + outcome, + .. + } = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); assert_eq!(outcome.accepted_data_points, 1); assert_eq!(outcome.rejected_data_points, 1); @@ -2127,8 +2396,9 @@ mod tests { vec![], AggregationTemporality::Cumulative, )]); - let (_, _, _, outcome) = - to_grpc_insert_requests(empty, &mut OtlpMetricCtx::default()).unwrap(); + let outcome = to_grpc_insert_requests(empty, &mut OtlpMetricCtx::default()) + .unwrap() + .outcome; assert_eq!(outcome.rejected_data_points, 0); assert_eq!(outcome.error_message, None); } @@ -2208,8 +2478,12 @@ mod tests { .collect(), }; - let (requests, rows, _, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let MetricsConversion { + requests, + rows, + outcome, + .. + } = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); assert_eq!(outcome.accepted_data_points, 2); assert_eq!(rows, 6); @@ -2243,8 +2517,13 @@ mod tests { experimental_enable_exponential_histogram: true, ..Default::default() }; - let (requests, rows, semantic_index, outcome) = - to_grpc_insert_requests(request, &mut ctx).unwrap(); + let MetricsConversion { + requests, + rows, + semantic_index, + outcome, + .. + } = to_grpc_insert_requests(request, &mut ctx).unwrap(); assert_eq!(outcome.accepted_data_points, 0); assert_eq!(outcome.rejected_data_points, 1); @@ -2269,14 +2548,17 @@ mod tests { experimental_enable_exponential_histogram: true, ..Default::default() }; - let (new_requests, _, _, _) = - to_grpc_insert_requests(request.clone(), &mut new_ctx).unwrap(); + let new_requests = to_grpc_insert_requests(request.clone(), &mut new_ctx) + .unwrap() + .requests; let mut legacy_ctx = OtlpMetricCtx { experimental_enable_exponential_histogram: true, is_legacy: true, ..Default::default() }; - let (legacy_requests, _, _, _) = to_grpc_insert_requests(request, &mut legacy_ctx).unwrap(); + let legacy_requests = to_grpc_insert_requests(request, &mut legacy_ctx) + .unwrap() + .requests; let new_insert = &new_requests.inserts[0]; let legacy_insert = &legacy_requests.inserts[0]; @@ -2325,8 +2607,9 @@ mod tests { vec![ExponentialHistogramDataPoint::default()], AggregationTemporality::Cumulative, )]); - let (_, _, _, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let outcome = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()) + .unwrap() + .outcome; assert_eq!(outcome.rejected_data_points, 1); assert!(outcome.error_message.unwrap().len() <= MAX_REJECTION_MESSAGE_BYTES); diff --git a/src/servers/src/otlp/metrics/resource_info.rs b/src/servers/src/otlp/metrics/resource_info.rs new file mode 100644 index 00000000000..cd3427f7163 --- /dev/null +++ b/src/servers/src/otlp/metrics/resource_info.rs @@ -0,0 +1,393 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Resource descriptor synthesized from OTLP metrics requests, so the entity +//! graph reads one table instead of scanning every logical metric table for +//! attributes the promote filter may have dropped. +//! +//! Columns are a fixed allowlist under the raw OTel attribute names: the +//! conventions whitelist matches fixed names, so they must not follow the +//! per-request label translation strategy or the promote/ignore headers. + +use std::collections::BTreeMap; + +use api::v1::RowInsertRequests; +use common_catalog::consts::SEMANTIC_GRAPH_WINDOW_NANOS; +use common_grpc::precision::Precision; +use common_query::prelude::{greptime_timestamp, greptime_value}; +use otel_arrow_rust::proto::opentelemetry::common::v1::KeyValue; +use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ResourceMetrics, metric}; +use session::protocol_ctx::OtlpMetricCtx; + +use crate::error::Result; +use crate::otlp::metrics::{ + INSTANCE_KEY, JOB_KEY, ServiceIdentity, exponential_histogram_gate, + exponential_histogram_value, scalar_value_string, service_identity, +}; +use crate::otlp::trace::{ + KEY_CONTAINER_ID, KEY_CONTAINER_NAME, KEY_HOST_ID, KEY_HOST_NAME, KEY_K8S_NAMESPACE_NAME, + KEY_K8S_POD_NAME, KEY_K8S_POD_UID, KEY_SERVICE_NAME, KEY_SERVICE_NAMESPACE, +}; +use crate::row_writer::{self, MultiTableData}; + +/// Prefixed like the other engine-managed tables, so a user metric is +/// unlikely to claim the name. +pub const OTEL_RESOURCE_INFO_TABLE_NAME: &str = "greptime_otel_resource_info"; + +/// Attributes projected under their raw OTel keys. `service.instance.id` is +/// absent on purpose: it lands in `instance`. +/// +/// Matched instead of scanned: this runs for every attribute of every +/// resource, and the compiler turns it into a length-and-prefix dispatch. +fn is_projected_attr(key: &str) -> bool { + matches!( + key, + KEY_SERVICE_NAME + | KEY_SERVICE_NAMESPACE + | KEY_HOST_ID + | KEY_HOST_NAME + | KEY_CONTAINER_ID + | KEY_CONTAINER_NAME + | KEY_K8S_POD_UID + | KEY_K8S_POD_NAME + | KEY_K8S_NAMESPACE_NAME + ) +} + +/// Upper bound of [`is_projected_attr`] plus the derived `job`/`instance`, +/// used to size the per-row buffers. +const MAX_PROJECTED_TAGS: usize = 11; + +/// Projected attributes (sorted `(name, value)` pairs) -> graph window -> +/// the newest data-point time seen in that window, which is what the row for +/// that window is stamped with. +/// +/// Windows are an inner map so the attributes are stored, and moved, once per +/// resource. They are keyed separately because one request may carry data for +/// several of them: a single row per resource would describe only the newest +/// window, leaving the earlier ones with metric rows but no entities. +#[derive(Debug, Default)] +pub struct ResourceInfoData { + rows: BTreeMap, BTreeMap>, +} + +impl ResourceInfoData { + /// Takes the raw attributes, before the promote filter runs on them. + pub fn observe( + &mut self, + raw_attrs: &[KeyValue], + resource: &ResourceMetrics, + metric_ctx: &OtlpMetricCtx, + ) { + let mut tags = Vec::with_capacity(MAX_PROJECTED_TAGS); + let ServiceIdentity { job, instance } = service_identity(raw_attrs); + if let Some(job) = job { + tags.push((JOB_KEY.to_string(), job)); + } + if let Some(instance) = instance { + tags.push((INSTANCE_KEY.to_string(), instance)); + } + for kv in raw_attrs { + if is_projected_attr(&kv.key) + && let Some(value) = scalar_value_string(kv.value.as_ref()) + { + tags.push((kv.key.clone(), value)); + } + } + if tags.is_empty() { + return; + } + // Sorted so equal attribute sets share a key, and so the emitted + // columns keep a stable order. + tags.sort_unstable(); + + let mut observed: BTreeMap = BTreeMap::new(); + for_each_encoded_time(resource, metric_ctx, |ts| { + let window = ts - ts.rem_euclid(SEMANTIC_GRAPH_WINDOW_NANOS); + observed + .entry(window) + .and_modify(|newest| *newest = (*newest).max(ts)) + .or_insert(ts); + }); + if observed.is_empty() { + return; + } + + let windows = self.rows.entry(tags).or_default(); + for (window, newest) in observed { + windows + .entry(window) + .and_modify(|seen| *seen = (*seen).max(newest)) + .or_insert(newest); + } + } + + /// Every projected attribute becomes a tag, so auto-create puts it in the + /// primary key where the conventions expect it. + pub fn into_row_insert_requests(self) -> Result> { + if self.rows.is_empty() { + return Ok(None); + } + + let mut writer = MultiTableData::default(); + let table = writer.get_or_default_table_data( + OTEL_RESOURCE_INFO_TABLE_NAME, + MAX_PROJECTED_TAGS + 2, + self.rows.values().map(BTreeMap::len).sum(), + ); + for (tags, windows) in &self.rows { + for ts_nanos in windows.values().copied() { + let mut row = table.alloc_one_row(); + row_writer::write_tags(table, tags.iter().cloned(), &mut row)?; + row_writer::write_f64(table, greptime_value(), 1.0, &mut row)?; + row_writer::write_ts_to_millis( + table, + greptime_timestamp(), + Some(ts_nanos), + Precision::Nanosecond, + &mut row, + )?; + table.add_row(row); + } + } + + let (requests, _) = writer.into_row_insert_requests(); + Ok(Some(requests)) + } +} + +/// Visits the times of the data points the encoder writes rows for, so a +/// resource is described exactly where it is measured rather than wherever +/// its request happens to reach. +fn for_each_encoded_time( + resource: &ResourceMetrics, + metric_ctx: &OtlpMetricCtx, + mut visit: impl FnMut(i64), +) { + fn visit_all(points: impl Iterator, visit: &mut impl FnMut(i64)) { + for ts in points { + visit(ts as i64); + } + } + for scope in &resource.scope_metrics { + for m in &scope.metrics { + match &m.data { + Some(metric::Data::Gauge(g)) => { + visit_all(g.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Sum(s)) => { + visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Histogram(h)) => { + visit_all(h.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Summary(s)) => { + visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::ExponentialHistogram(h)) + if exponential_histogram_gate(h, metric_ctx).is_ok() => + { + for point in &h.data_points { + if let Ok((_, ts)) = exponential_histogram_value(point) { + visit(ts); + } + } + } + Some(metric::Data::ExponentialHistogram(_)) => {} + None => {} + } + } + } +} + +#[cfg(test)] +mod tests { + use api::v1::SemanticType; + use api::v1::value::ValueData; + use common_query::prelude::set_default_prefix; + use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, any_value}; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Gauge, Metric, + NumberDataPoint, ScopeMetrics, + }; + + use super::*; + + fn kv(key: &str, value: &str) -> KeyValue { + KeyValue { + key: key.into(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.into())), + }), + } + } + + fn gauge_at(times: &[i64]) -> ResourceMetrics { + ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + data: Some(metric::Data::Gauge(Gauge { + data_points: times + .iter() + .map(|ts| NumberDataPoint { + time_unix_nano: *ts as u64, + ..Default::default() + }) + .collect(), + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + } + } + + #[test] + fn observe_projects_allowlist_and_dedups_per_request() { + let mut data = ResourceInfoData::default(); + let attrs = vec![ + kv("service.name", "api"), + kv("service.namespace", "shop"), + kv("service.instance.id", "inst-1"), + kv("host.id", "h-1"), + kv("os.type", "linux"), + ]; + data.observe(&attrs, &gauge_at(&[100, 50]), &OtlpMetricCtx::default()); + assert_eq!(data.rows.len(), 1); + let (tags, windows) = data.rows.iter().next().unwrap(); + assert_eq!(windows.values().copied().collect::>(), vec![100]); + assert!(tags.contains(&("job".to_string(), "shop/api".to_string()))); + assert!(tags.contains(&("instance".to_string(), "inst-1".to_string()))); + assert!(tags.contains(&("service.name".to_string(), "api".to_string()))); + assert!( + tags.iter() + .all(|(k, _)| k != "os.type" && k != "service.instance.id") + ); + + data.observe( + &[kv("host.id", "h-2")], + &gauge_at(&[10]), + &OtlpMetricCtx::default(), + ); + assert_eq!(data.rows.len(), 2); + + let mut empty = ResourceInfoData::default(); + empty.observe( + &[kv("os.type", "linux")], + &gauge_at(&[100]), + &OtlpMetricCtx::default(), + ); + assert!(empty.into_row_insert_requests().unwrap().is_none()); + } + + /// Earlier windows would keep their metric rows but lose their entities. + #[test] + fn observe_keeps_one_row_per_graph_window() { + let window = SEMANTIC_GRAPH_WINDOW_NANOS; + let mut data = ResourceInfoData::default(); + data.observe( + &[kv("service.name", "api")], + &gauge_at(&[window + 1, window + 2, 3 * window + 7]), + &OtlpMetricCtx::default(), + ); + + let windows = data.rows.values().next().unwrap(); + assert_eq!( + windows.iter().collect::>(), + vec![(&window, &(window + 2)), (&(3 * window), &(3 * window + 7))] + ); + } + + /// Describing a resource whose only data the encoder drops invents an + /// entity with no measurements. + #[test] + fn observe_ignores_data_the_encoder_drops() { + let exponential = |temporality: AggregationTemporality| ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points: vec![ExponentialHistogramDataPoint { + time_unix_nano: 100, + ..Default::default() + }], + aggregation_temporality: temporality as i32, + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }; + let enabled = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + + for (resource, ctx) in [ + ( + exponential(AggregationTemporality::Cumulative), + OtlpMetricCtx::default(), + ), + (exponential(AggregationTemporality::Delta), enabled), + ] { + let mut data = ResourceInfoData::default(); + data.observe(&[kv("service.name", "api")], &resource, &ctx); + assert!(data.into_row_insert_requests().unwrap().is_none()); + } + } + + #[test] + fn rows_carry_raw_key_tags_value_and_millis_timestamp() { + set_default_prefix(None).unwrap(); + let mut data = ResourceInfoData::default(); + data.observe( + &[kv("service.name", "api"), kv("host.id", "h-1")], + &gauge_at(&[1_700_000_000_123_456_789]), + &OtlpMetricCtx::default(), + ); + let requests = data.into_row_insert_requests().unwrap().unwrap(); + assert_eq!(requests.inserts.len(), 1); + let insert = &requests.inserts[0]; + assert_eq!(insert.table_name, OTEL_RESOURCE_INFO_TABLE_NAME); + + let rows = insert.rows.as_ref().unwrap(); + let names = rows + .schema + .iter() + .map(|c| c.column_name.as_str()) + .collect::>(); + assert_eq!( + names, + vec![ + "host.id", + "job", + "service.name", + greptime_value(), + greptime_timestamp() + ] + ); + for column in &rows.schema[..3] { + assert_eq!(column.semantic_type, SemanticType::Tag as i32); + } + + assert_eq!(rows.rows.len(), 1); + let values = &rows.rows[0].values; + assert_eq!(values[3].value_data, Some(ValueData::F64Value(1.0))); + assert_eq!( + values[4].value_data, + Some(ValueData::TimestampMillisecondValue(1_700_000_000_123)) + ); + } +} diff --git a/src/servers/src/otlp/trace.rs b/src/servers/src/otlp/trace.rs index a8439cdb2c4..9ee16507bda 100644 --- a/src/servers/src/otlp/trace.rs +++ b/src/servers/src/otlp/trace.rs @@ -46,7 +46,15 @@ pub const TRACE_STATE_COLUMN: &str = "trace_state"; // const keys pub const KEY_SERVICE_NAME: &str = "service.name"; +pub const KEY_SERVICE_NAMESPACE: &str = "service.namespace"; pub const KEY_SERVICE_INSTANCE_ID: &str = "service.instance.id"; +pub const KEY_HOST_ID: &str = "host.id"; +pub const KEY_HOST_NAME: &str = "host.name"; +pub const KEY_CONTAINER_ID: &str = "container.id"; +pub const KEY_CONTAINER_NAME: &str = "container.name"; +pub const KEY_K8S_POD_UID: &str = "k8s.pod.uid"; +pub const KEY_K8S_POD_NAME: &str = "k8s.pod.name"; +pub const KEY_K8S_NAMESPACE_NAME: &str = "k8s.namespace.name"; pub const KEY_SPAN_KIND: &str = "span.kind"; // jaeger const keys, not sure if they are general diff --git a/src/session/src/protocol_ctx.rs b/src/session/src/protocol_ctx.rs index 537dec5c414..d9a4a386f91 100644 --- a/src/session/src/protocol_ctx.rs +++ b/src/session/src/protocol_ctx.rs @@ -56,6 +56,9 @@ pub struct OtlpMetricCtx { pub with_metric_engine: bool, pub experimental_enable_exponential_histogram: bool, pub is_legacy: bool, + /// Set from the server's `otlp.experimental_enable_resource_info`; off + /// means the resource descriptor is not synthesized at all. + pub resource_info: bool, pub metric_type: MetricType, pub metric_translation_strategy: OtlpMetricTranslationStrategy, } diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index 2daa440c11d..471be8b891c 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -98,6 +98,10 @@ pub const SEMANTIC_ENTITY_SERVICE_ID: &str = "greptime.semantic.entity.service.i /// ([`has_stable_string_form`]). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum EntityRole { + /// Identifying columns. Their **order is part of the identity**: the + /// entity id is their values joined in the declared order, so two tables + /// naming the same entity must list them the same way (broad to narrow) + /// or they name two entities. Id, Descriptive, Scope, diff --git a/tests-integration/src/instance.rs b/tests-integration/src/instance.rs index 08d2a1aff15..9f48d01a467 100644 --- a/tests-integration/src/instance.rs +++ b/tests-integration/src/instance.rs @@ -691,4 +691,119 @@ mod tests { "denied source leaked into entities:\n{pretty_print}" ); } + + /// Denying the descriptor's table must not fail the request: the metric + /// data is already committed, so the client would retry what the server + /// accepted. + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_descriptor_permission_denial_degrades() { + use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; + use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::number_data_point::Value; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + Gauge, Metric, NumberDataPoint, ResourceMetrics, ScopeMetrics, metric, + }; + use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; + use servers::query_handler::OpenTelemetryProtocolHandler; + + common_telemetry::init_default_ut_logging(); + + /// Denies writes to the descriptor table only. + struct DenyResourceInfo; + + impl auth::PermissionChecker for DenyResourceInfo { + fn check_permission( + &self, + _user_info: auth::UserInfoRef, + _req: auth::PermissionReq, + ) -> auth::error::Result { + Ok(auth::PermissionResp::Allow) + } + + fn check_permission_with_table_targets( + &self, + _user_info: auth::UserInfoRef, + _req: auth::PermissionReq, + targets: auth::PermissionTableTargets, + ) -> auth::error::Result { + if let auth::PermissionTableTargets::Resolved(targets) = targets + && targets + .iter() + .any(|target| target.table == "greptime_otel_resource_info") + { + return Ok(auth::PermissionResp::Reject); + } + Ok(auth::PermissionResp::Allow) + } + } + + let plugins = Plugins::new(); + plugins.insert::(Arc::new(DenyResourceInfo)); + + let standalone = GreptimeDbStandaloneBuilder::new("test_otlp_descriptor_denied") + .with_plugin(plugins) + .build() + .await; + let instance = standalone.fe_instance().clone(); + + let attr = |key: &str, value: &str| KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.to_string())), + }), + }; + let request = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: Some(Resource { + attributes: vec![attr("service.name", "api"), attr("host.id", "h-1")], + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: "denied_descriptor_gauge".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano: 1_700_000_000_000_000_000, + value: Some(Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + + let outcome = + OpenTelemetryProtocolHandler::metrics(&*instance, request, QueryContext::arc()) + .await + .unwrap(); + let warning = outcome + .error_message + .expect("denial must surface as a warning"); + assert!( + warning.contains("greptime_otel_resource_info"), + "unexpected warning: {warning}" + ); + + // the metric data is committed, the descriptor table was never created + let sql = "select table_name from information_schema.tables \ + where table_schema = 'public' order by table_name"; + let output = query(&instance, sql).await; + let OutputData::Stream(s) = output.data else { + unreachable!() + }; + let batches = common_recordbatch::util::collect_batches(s).await.unwrap(); + let pretty_print = batches.pretty_print().unwrap(); + assert!( + pretty_print.contains("denied_descriptor_gauge"), + "metric data was not committed:\n{pretty_print}" + ); + assert!( + !pretty_print.contains("greptime_otel_resource_info"), + "denied descriptor table was created:\n{pretty_print}" + ); + } } diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 1085c21983e..16fc2e65634 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -369,6 +369,11 @@ impl GreptimeDbStandaloneBuilder { grpc: GrpcOptions::default().with_server_addr("127.0.0.1:4001"), slow_query: self.slow_query_options.clone(), auto_create_table: self.auto_create_table, + // Tests cover the descriptor, so they run with it enabled. + otlp: frontend::service_config::OtlpOptions { + experimental_enable_resource_info: true, + ..Default::default() + }, ..StandaloneOptions::default() }; diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 7a787833376..4a8f96260e8 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -164,6 +164,7 @@ macro_rules! http_tests { test_otlp_metrics_new, test_otlp_exponential_histogram, test_otlp_metric_translation_strategies, + test_otlp_metrics_resource_info_conflicts, test_otlp_traces_v0, test_otlp_traces_v1, test_otlp_traces_v1_entity_graph, @@ -2208,6 +2209,7 @@ enable = true enable = true experimental_enable_exponential_histogram = false trace_ingest_chunk_size = 512 +experimental_enable_resource_info = true [prom_store] enable = true @@ -6500,7 +6502,7 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .await; assert_eq!(StatusCode::OK, res.status()); - let expected = "[[\"claude_code_cost_usage_USD_total\"],[\"claude_code_token_usage_tokens_total\"],[\"demo\"],[\"greptime_physical_table\"],[\"numbers\"]]"; + let expected = "[[\"claude_code_cost_usage_USD_total\"],[\"claude_code_token_usage_tokens_total\"],[\"demo\"],[\"greptime_otel_resource_info\"],[\"greptime_physical_table\"],[\"numbers\"]]"; validate_data("otlp_metrics_all_tables", &client, "show tables;", expected).await; // Metric-engine logical table carries the semantic identity. Match substrings @@ -6567,6 +6569,33 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { ) .await; + // The synthesized resource descriptor: a plain mito info table keyed by + // the raw OTel attribute keys (service.name is the only allowlisted + // resource attribute in this payload), independent of the promote/ignore + // headers. The timestamp is the newest data-point time. + validate_data( + "otlp_metrics_resource_info", + &client, + "select job, \"service.name\", greptime_value, greptime_timestamp from greptime_otel_resource_info;", + "[[\"claude-code\",\"claude-code\",1.0,1753780559836]]", + ) + .await; + // All four ownership markers must land on auto-create: the descriptor + // write path refuses tables that miss any of them, so a missing stamp + // here would silently stop every follow-up descriptor write. + validate_data( + "otlp_metrics_resource_info_options", + &client, + "select count(*) from information_schema.tables where table_name = 'greptime_otel_resource_info' \ + and engine = 'mito' \ + and create_options like '%greptime.semantic.metric.type=info%' \ + and create_options like '%greptime.semantic.metric.metadata_quality=declared%' \ + and create_options like '%greptime.semantic.signal_type=metric%' \ + and create_options like '%greptime.semantic.source=opentelemetry%';", + "[[1]]", + ) + .await; + // drop table let res = client .get("/v1/sql?sql=drop table `claude_code_cost_usage_USD_total`;") @@ -6578,6 +6607,11 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .send() .await; assert_eq!(res.status(), StatusCode::OK); + let res = client + .get("/v1/sql?sql=drop table greptime_otel_resource_info;") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); // write metrics data // with scope attrs @@ -6704,6 +6738,16 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { ) .await; + // the descriptor was auto-created again after the earlier drop + validate_data( + "otlp_metrics_resource_info_recreated", + &client, + "select count(*) from information_schema.tables \ + where table_name = 'greptime_otel_resource_info' and engine = 'mito';", + "[[1]]", + ) + .await; + // drop table let res = client .get("/v1/sql?sql=drop table `claude_code_cost_usage_USD_total`;") @@ -6715,6 +6759,44 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .send() .await; assert_eq!(res.status(), StatusCode::OK); + let res = client + .get("/v1/sql?sql=drop table greptime_otel_resource_info;") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + // job composition end-to-end: service.namespace folds into "/" + // on both the descriptor and the metric tables + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}},{"key":"service.namespace","value":{"stringValue":"shop"}},{"key":"host.id","value":{"stringValue":"h-1"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"ns_gauge","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":1.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + )], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + validate_data( + "otlp_metrics_namespace_job_descriptor", + &client, + "select job, \"service.name\", \"service.namespace\", \"host.id\" from greptime_otel_resource_info;", + "[[\"shop/api\",\"api\",\"shop\",\"h-1\"]]", + ) + .await; + validate_data( + "otlp_metrics_namespace_job_metric_table", + &client, + "select job from ns_gauge;", + "[[\"shop/api\"]]", + ) + .await; guard.remove_all().await; } @@ -6912,6 +6994,24 @@ pub async fn test_otlp_metric_translation_strategies(store_type: StorageType) { ) .await; + // the descriptor's columns are the fixed raw keys regardless of the + // translation strategy, and only allowlisted attributes are projected + validate_data( + "otlp_metric_translation_strategy_resource_info", + &client, + "select job, \"service.name\" from greptime_otel_resource_info;", + "[[\"strategy-service\",\"strategy-service\"]]", + ) + .await; + validate_data( + "otlp_metric_translation_strategy_resource_info_allowlist", + &client, + "select count(*) from information_schema.columns \ + where table_name = 'greptime_otel_resource_info' and column_name = 'resource.attr';", + "[[0]]", + ) + .await; + let res = send_req( &client, vec![ @@ -6934,6 +7034,122 @@ pub async fn test_otlp_metric_translation_strategies(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_otlp_metrics_resource_info_conflicts(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_otlp_metrics_resource_info_conflicts") + .await; + let client = TestClient::new(app).await; + + let content_type = || { + ( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + ) + }; + + // A metric named like the descriptor wins the table name: it goes through + // the normal metric-engine path and no descriptor is synthesized. + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}},{"key":"host.id","value":{"stringValue":"h-1"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"greptime_otel_resource_info","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":1.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + assert!(body.partial_success.is_none()); + validate_data( + "otlp_metrics_name_collision_engine", + &client, + "select count(*) from information_schema.tables \ + where table_name = 'greptime_otel_resource_info' and engine = 'metric';", + "[[1]]", + ) + .await; + let res = execute_sql(&client, "drop table greptime_otel_resource_info;").await; + assert_eq!(res.status(), StatusCode::OK); + + // A pre-existing table with an incompatible schema fails the descriptor + // write, which degrades to an OTLP partial-success warning while the + // metric data itself is accepted. + let res = execute_sql( + &client, + "create table greptime_otel_resource_info (ts timestamp time index, job double, greptime_value double);", + ) + .await; + assert_eq!(res.status(), StatusCode::OK); + + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"conflict_gauge","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":7.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + let partial_success = body.partial_success.as_ref().unwrap(); + assert_eq!(partial_success.rejected_data_points, 0); + assert!( + partial_success + .error_message + .contains("greptime_otel_resource_info"), + "unexpected warning: {}", + partial_success.error_message + ); + validate_data( + "otlp_metrics_conflict_main_data_accepted", + &client, + "select greptime_value from conflict_gauge;", + "[[7.0]]", + ) + .await; + + // The descriptor is auto-created again once the conflicting table is + // gone, and a follow-up write into it does not degrade. + let res = execute_sql(&client, "drop table greptime_otel_resource_info;").await; + assert_eq!(res.status(), StatusCode::OK); + for round in 0..2 { + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + assert!( + body.partial_success.is_none(), + "round {round} unexpectedly degraded: {:?}", + body.partial_success + ); + } + validate_data( + "otlp_metrics_owned_descriptor_keeps_accepting", + &client, + "select count(*) from greptime_otel_resource_info;", + "[[1]]", + ) + .await; + + guard.remove_all().await; +} + pub async fn test_otlp_traces_v0(store_type: StorageType) { // init common_telemetry::init_default_ut_logging(); @@ -7136,7 +7352,7 @@ pub async fn test_otlp_traces_v1_entity_graph(store_type: StorageType) { let res = send_trace_v1_req(&client, "graph_traces", req, false).await; assert_eq!(StatusCode::OK, res.status()); - let expected = r#"[["service","frontend","service","cart","calls"],["service.instance","resource_attributes.service.instance.id=binst-1,service_name=batch","service","batch","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","service","frontend","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","k8s.pod","uid-9","runs_on"]]"#; + let expected = r#"[["service","frontend","service","cart","calls"],["service.instance","batch,binst-1","service","batch","part_of"],["service.instance","frontend,inst-1","service","frontend","part_of"],["service.instance","frontend,inst-1","k8s.pod","uid-9","runs_on"]]"#; validate_data( "otlp_traces_entity_graph", &client, diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index 6af28cfa2b5..f8894319842 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -94,26 +94,26 @@ select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_table from greptime_private.semantic_entities order by entity_type, entity_id; -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| host | h1 | | | | ["public.graph_app_metrics"] | -| process | host=h1,service_name=cart | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_metrics"] | -| service | cart | | us-east | | ["public.graph_app_metrics"] | -| service.instance | cart-0 | | | | ["public.graph_app_metrics"] | -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ +| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ +| host | h1 | | | | ["public.graph_app_metrics"] | +| process | cart,h1 | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_metrics"] | +| service | cart | | us-east | | ["public.graph_app_metrics"] | +| service.instance | cart-0 | | | | ["public.graph_app_metrics"] | ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence from greptime_private.semantic_relationships order by rel_type, src_id; -+------------------+---------------------------+----------+--------+----------+------------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | -+------------------+---------------------------+----------+--------+----------+------------+------------+ -| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | -| service.instance | cart-0 | host | h1 | runs_on | attribute | 1.0 | -| process | host=h1,service_name=cart | host | h1 | runs_on | attribute | 1.0 | -+------------------+---------------------------+----------+--------+----------+------------+------------+ ++------------------+---------+----------+--------+----------+------------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | ++------------------+---------+----------+--------+----------+------------+------------+ +| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | +| process | cart,h1 | host | h1 | runs_on | attribute | 1.0 | +| service.instance | cart-0 | host | h1 | runs_on | attribute | 1.0 | ++------------------+---------+----------+--------+----------+------------+------------+ drop table graph_app_metrics; @@ -570,38 +570,38 @@ select entity_type, entity_id, source_tables from greptime_private.semantic_entities order by entity_type, entity_id, source_tables; -+------------------+-----------------------------------------------------------+------------------------------------+ -| entity_type | entity_id | source_tables | -+------------------+-----------------------------------------------------------+------------------------------------+ -| k8s.container | container=main,uid=uid-1 | ["public.kube_pod_container_info"] | -| k8s.node | node-a | ["public.kube_pod_info"] | -| k8s.pod | uid-1 | ["public.kube_pod_info"] | -| k8s.pod | uid-1 | ["public.kube_pod_owner"] | -| k8s.pod | uid-1 | ["public.kube_pod_container_info"] | -| k8s.pod | uid-2 | ["public.kube_pod_info"] | -| k8s.pod | uid-2 | ["public.kube_pod_owner"] | -| k8s.pod | uid-3 | ["public.kube_pod_info"] | -| k8s.service | svc-uid-1 | ["public.kube_service_info"] | -| k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | ["public.kube_pod_owner"] | -| service | shop/api | ["public.target_info"] | -| service.instance | instance=inst-1,job=shop/api | ["public.target_info"] | -+------------------+-----------------------------------------------------------+------------------------------------+ ++------------------+---------------------------+------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+---------------------------+------------------------------------+ +| k8s.container | uid-1,main | ["public.kube_pod_container_info"] | +| k8s.node | node-a | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_owner"] | +| k8s.pod | uid-1 | ["public.kube_pod_container_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_owner"] | +| k8s.pod | uid-3 | ["public.kube_pod_info"] | +| k8s.service | svc-uid-1 | ["public.kube_service_info"] | +| k8s.workload | default,ReplicaSet,api-rs | ["public.kube_pod_owner"] | +| service | shop/api | ["public.target_info"] | +| service.instance | shop/api,inst-1 | ["public.target_info"] | ++------------------+---------------------------+------------------------------------+ -- SQLNESS PROTOCOL MYSQL select src_type, src_id, dst_type, dst_id, rel_type, provenance from greptime_private.semantic_relationships order by rel_type, src_id, dst_id; -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ -| k8s.pod | uid-1 | k8s.container | container=main,uid=uid-1 | contains | attribute | -| service.instance | instance=inst-1,job=shop/api | service | shop/api | part_of | attribute | -| k8s.pod | uid-1 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | -| k8s.pod | uid-2 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | -| k8s.pod | uid-1 | k8s.node | node-a | runs_on | attribute | -| k8s.pod | uid-2 | k8s.node | node-a | runs_on | attribute | -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ ++------------------+-----------------+---------------+---------------------------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++------------------+-----------------+---------------+---------------------------+----------+------------+ +| k8s.pod | uid-1 | k8s.container | uid-1,main | contains | attribute | +| service.instance | shop/api,inst-1 | service | shop/api | part_of | attribute | +| k8s.pod | uid-1 | k8s.workload | default,ReplicaSet,api-rs | part_of | attribute | +| k8s.pod | uid-2 | k8s.workload | default,ReplicaSet,api-rs | part_of | attribute | +| k8s.pod | uid-1 | k8s.node | node-a | runs_on | attribute | +| k8s.pod | uid-2 | k8s.node | node-a | runs_on | attribute | ++------------------+-----------------+---------------+---------------------------+----------+------------+ drop table kube_pod_info; @@ -627,3 +627,182 @@ drop table http_requests_total; Affected Rows: 0 +-- OTel conventions: the ingestion-synthesized greptime_otel_resource_info descriptor +-- (stamped signal_type=metric + source=opentelemetry + metric.type=info) gets implicit +-- declarations under the raw OTel column names. host/container identities are +-- the stable ids only, so rows with an empty host.id / container.id link no +-- infrastructure, and a row without an instance value declares no +-- service.instance. The same table name stamped source=prometheus is not +-- whitelisted. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + "host.id" string, + "host.name" string, + "container.id" string, + "container.name" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace", + "host.id", "host.name", "container.id", "container.name") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +Affected Rows: 0 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 'h-1', 'node-a', 'c-1', 'api-ctr', 1), + (now(), 'shop/api', 'inst-2', 'api', 'shop', '', 'laptop', '', '', 1), + (now(), 'worker', '', 'worker', '', 'h-1', 'node-a', '', '', 1); + +Affected Rows: 3 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + ++------------------+-----------------+----------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+-----------------+----------------------------------------+ +| container | c-1 | ["public.greptime_otel_resource_info"] | +| host | h-1 | ["public.greptime_otel_resource_info"] | +| service | shop/api | ["public.greptime_otel_resource_info"] | +| service | worker | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-1 | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-2 | ["public.greptime_otel_resource_info"] | ++------------------+-----------------+----------------------------------------+ + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + ++------------------+-----------------+-----------+----------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++------------------+-----------------+-----------+----------+----------+------------+ +| service.instance | shop/api,inst-1 | service | shop/api | part_of | attribute | +| service.instance | shop/api,inst-2 | service | shop/api | part_of | attribute | +| container | c-1 | host | h-1 | runs_on | attribute | +| service.instance | shop/api,inst-1 | container | c-1 | runs_on | attribute | +| service.instance | shop/api,inst-1 | host | h-1 | runs_on | attribute | ++------------------+-----------------+-----------+----------+----------+------------+ + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + +-- The descriptor whitelist is gated on source=opentelemetry: the same table +-- shape stamped as a prometheus source contributes nothing. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "host.id" string, + greptime_value double, + primary key (job, instance, "host.id") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'h-1', 1); + +Affected Rows: 1 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id +from greptime_private.semantic_entities +order by entity_type, entity_id; + +affected_rows: 0 + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + +-- Cross-signal identity: the same service reaches the graph as trace spans and +-- as an ingestion-synthesized descriptor. The two sources name their identity +-- columns differently and only the metric side pre-composes the namespace into +-- job, so this asserts one entity per service and per instance, sourced from +-- both tables. +create table graph_xsignal_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "resource_attributes.service.namespace" string, + "resource_attributes.service.instance.id" string, + primary key (service_name, "resource_attributes.service.namespace", + "resource_attributes.service.instance.id") +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +Affected Rows: 0 + +insert into graph_xsignal_traces values + (now(), 't1', 's1', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'api', 100, 'shop', 'inst-1'), + (now(), 't2', 's2', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'worker', 100, '', 'inst-2'); + +Affected Rows: 2 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 1), + (now(), 'worker', 'inst-2', 'worker', '', 1); + +Affected Rows: 2 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +where entity_type in ('service', 'service.instance') +order by entity_type, entity_id, source_tables; + ++------------------+-----------------+----------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+-----------------+----------------------------------------+ +| service | shop/api | ["public.graph_xsignal_traces"] | +| service | shop/api | ["public.greptime_otel_resource_info"] | +| service | worker | ["public.graph_xsignal_traces"] | +| service | worker | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-1 | ["public.graph_xsignal_traces"] | +| service.instance | shop/api,inst-1 | ["public.greptime_otel_resource_info"] | +| service.instance | worker,inst-2 | ["public.graph_xsignal_traces"] | +| service.instance | worker,inst-2 | ["public.greptime_otel_resource_info"] | ++------------------+-----------------+----------------------------------------+ + +drop table graph_xsignal_traces; + +Affected Rows: 0 + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index bfc9ea8fd7b..60d95630d91 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -387,3 +387,122 @@ drop table kube_service_info; drop table target_info; drop table http_requests_total; + +-- OTel conventions: the ingestion-synthesized greptime_otel_resource_info descriptor +-- (stamped signal_type=metric + source=opentelemetry + metric.type=info) gets implicit +-- declarations under the raw OTel column names. host/container identities are +-- the stable ids only, so rows with an empty host.id / container.id link no +-- infrastructure, and a row without an instance value declares no +-- service.instance. The same table name stamped source=prometheus is not +-- whitelisted. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + "host.id" string, + "host.name" string, + "container.id" string, + "container.name" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace", + "host.id", "host.name", "container.id", "container.name") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 'h-1', 'node-a', 'c-1', 'api-ctr', 1), + (now(), 'shop/api', 'inst-2', 'api', 'shop', '', 'laptop', '', '', 1), + (now(), 'worker', '', 'worker', '', 'h-1', 'node-a', '', '', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + +drop table greptime_otel_resource_info; + +-- The descriptor whitelist is gated on source=opentelemetry: the same table +-- shape stamped as a prometheus source contributes nothing. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "host.id" string, + greptime_value double, + primary key (job, instance, "host.id") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'h-1', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id +from greptime_private.semantic_entities +order by entity_type, entity_id; + +drop table greptime_otel_resource_info; + +-- Cross-signal identity: the same service reaches the graph as trace spans and +-- as an ingestion-synthesized descriptor. The two sources name their identity +-- columns differently and only the metric side pre-composes the namespace into +-- job, so this asserts one entity per service and per instance, sourced from +-- both tables. +create table graph_xsignal_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "resource_attributes.service.namespace" string, + "resource_attributes.service.instance.id" string, + primary key (service_name, "resource_attributes.service.namespace", + "resource_attributes.service.instance.id") +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +insert into graph_xsignal_traces values + (now(), 't1', 's1', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'api', 100, 'shop', 'inst-1'), + (now(), 't2', 's2', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'worker', 100, '', 'inst-2'); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 1), + (now(), 'worker', 'inst-2', 'worker', '', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +where entity_type in ('service', 'service.instance') +order by entity_type, entity_id, source_tables; + +drop table graph_xsignal_traces; + +drop table greptime_otel_resource_info; diff --git a/tests/conf/frontend-test.toml.template b/tests/conf/frontend-test.toml.template index 632947a3cde..5a879e04882 100644 --- a/tests/conf/frontend-test.toml.template +++ b/tests/conf/frontend-test.toml.template @@ -7,3 +7,6 @@ parallelism = 4 [event_recorder] event_types = [] + +[otlp] +experimental_enable_resource_info = true diff --git a/tests/conf/standalone-test.toml.template b/tests/conf/standalone-test.toml.template index ebe25f82487..8d7e307d511 100644 --- a/tests/conf/standalone-test.toml.template +++ b/tests/conf/standalone-test.toml.template @@ -53,3 +53,6 @@ retry_delay = "500ms" [event_recorder] event_types = [] + +[otlp] +experimental_enable_resource_info = true From 144f83528db8648bc5607fdc65c249bbfa3c7ffc Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 26 Aug 2026 06:15:37 +0000 Subject: [PATCH 109/126] test: renew etcd TLS certificates (#8956) Signed-off-by: WenyXu --- .../fixtures/etcd-tls-certs/ca-key.pem | 52 +++++++++---------- .../fixtures/etcd-tls-certs/ca.crt | 34 ++++++------ .../fixtures/etcd-tls-certs/ca.srl | 2 +- .../fixtures/etcd-tls-certs/client-key.pem | 52 +++++++++---------- .../fixtures/etcd-tls-certs/client.crt | 34 ++++++------ .../fixtures/etcd-tls-certs/server-key.pem | 52 +++++++++---------- .../fixtures/etcd-tls-certs/server.crt | 32 ++++++------ 7 files changed, 129 insertions(+), 129 deletions(-) diff --git a/tests-integration/fixtures/etcd-tls-certs/ca-key.pem b/tests-integration/fixtures/etcd-tls-certs/ca-key.pem index 2066dfb7b47..07a0a86ee39 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/ca-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCfIi0iTllUQyNY -ka+d9Zt30ILA0QBzLKHwez9QJj8bQLpfKR26d794CQn/OO6XO7MWsJ6nvqwRNYM0 -WX+UwUtuT9MrttYWpS5yEAVO79dMtF67gCNafXKjQyK8L1vz6Avq4XJN4VqZEifT -Fdk3s2MmS0twv4sgrK//jh7aZzBQJtw+ROe9RN1dqlKoa3TtpB7iwNTDxGiR0EYj -wHj+GdiNHtDMWvrnTNNX/IKeEwLP/Mjirzr7GnaOjR25ruqcLuwWrufu+EMhqFTU -r+LIQpZD0TFuPX6DPlQZOs/4UpjteoC84m5AnQXNlh0vzjbW1efEYWYQoKDesx9e -oE5jcABpAgMBAAECggEAA3Rw/r6CYaTG1tdBilA47LF/XTihu15mh6Y4BNadESDn -IEYav0p2mDW4sgH7lWyg4ecPaBHo124jfUFM1p8Zsvk9sERwbR5vSIpWKyq5hbtM -Fo2yIXakGgJM9ev19vDR7WqxRMVAkR6GnyZpcwmh0e+ENvGZpNR1n7qSAJOJRibW -Ye8qMF/zMLF9VekBRPFJFIZBvKgHxSGCLbXRzsi5LKMbwfBoJmxEZSOC9zODl43o -Z2i+LvlS0R/Rsf2KG+Hwgj1P7HH1znJeA5bVJFDewljjSwu1mk0JpQL7Jgp2xpqs -F/mRZmtoQSH3EsTm2Kjz1sLGXwUb9JaQ274aV4t7pwKBgQDXYF+qFYip7k+b9FTF -sFbfp3CjOqxXiAxnoE5hlfV/h0w4V1Su42AdejnbIVYFX8bvmezkNWPXhB6P5oHU -4nO+oAuhsRT4317mfRM5NWO6mv2eglSBnS04GLsX142A7jMPr9e+Cqh4dYfw/A1L -bSdsdILemTFQrak27pC11mHZOwKBgQC9JhDhQf5uRh6k8GvgKGu9ZEZmkYCVu2sr -7kpD60UhklnyQo6DKYsa6OcChRQMv9ss3Hqu3Y1G0/sTe1Tjz5zwrWWAZoLlbdOt -vipesvwNTxvJoUWPtrjOoFZtmyg0Q0TqxaxfEWAssDU9IfnIy+zx7sAg4zZIPUyy -6kudaY1SqwKBgQCAQ/71Fjn7qddzc4GA8lHqhJeKPokg3/8zP78uUtaQCo2UCD6A -oR0+sOn/3MyUCsQ5MZxpFHrPgPmKjabIl8yCvGHw+7sXtD+aWOa37VnlaiSc39Vg -E7E4dVIHEvJM1I9ISlrb7REEHErHc/Se9PTDnGfMFcPO3n2mH1HDWVeQvQKBgGGn -BHH3a08tXmbTRS5uT+lwmrQbjKJBJ3x/wtG75m4Fq/BaEk9/JDUZZyKy5/4JEzPf -BGvBME4P5QFS3CndJu5O5ydaRVwDzpRVqHRJvb11SShY3Zvrvw/WUai2wRPyYuM+ -eNaAFwIbWvEb2GSle8gP9htEkuLK2w1HzxAOzYqPAoGAKl5l/WsoLfQP5UIUjSQe -JtVTDVmoz8m8Mg2nn56UdUetG246EVhDx5ngj84kZDM0GDXhfVWV6MWnIZr5NQsr -bORIDNP7Czat8FnlXeAQ+R6A/87o5g5p6ydH9jYqrrLxYjaZyXJptE+ja4zW7Lbl -ogWAc76iZvXLXqKRfepG0zQ= +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQC05De2Tk8TMZ0J +ME+7Rp+zvvLkEBPPXCCFcXfo91gYVHMMr1TGQJuVtbx7rIqC1tFhc2luAGfsBmW3 +IfTXZjzY6834ZUzmPL7uuBfoRpd0Paj5X4owpn0eHIgcsqx4juHSTYc3fdnrTLpq +lNMulZEemkNULLCE7pyxg3hhYIx9JEe9nv26fkyr1vW5z1cPlEpyOK+SUIeqHgXi +VtWPU9WGv5sdQeXnD4E1JTX2URihb7C8xali3YVbI3SUxy/gAhM6swQgZ4SLgW8k +f87a+PcP9C1N5Gujfw4fT90dJoutyiTSbcDrjZLjF03TgjcDaACvntb0qN85XD96 +dIvhg5jfAgMBAAECggEAAd3vHe6d+KTMDZ1HwPAAKS151ipnpX3W3aRgi/FC6OMm +r0VUpWC4jCJlif3kYvaID66yFeucJveD75CJw6VaF4bSnTa9GmKM5F2BQhsrhKj+ +69jY+MU8cztGfJMdgczNfi79qDc8zSW6TUFpoMhmD6Og8icRCODbkbRWDDnnQi1d +1CVe/lZ4yIapCGgzHfQxdl9I8LEZWyeAfqMCvOTI2eE4Z4ZYto/EmuRoJN3oqOx1 +MAtpS5mw6RTFVitO4b0j5gwAN2u6Cgo8jDc5UTzHGAiGc/dAn7S+mNGv+IXP1/aL +UA/ZsBwQhZbqtGCkSmCo5gGDL+anDTiIy6nRCwIWWQKBgQDtE7V4pNj5J2MciK+h +KAR/kT36NuGoZSnf585HJVB/xAs/d8F4kU5pCofmrtm2Bl5Lx4sCE3wPmv/HMYcI +QH2zt/YxPA3F/9zoBowfTNB8/iCDlzNb/pUry+qkBhKG2HgWaUkkGEtiO16PeCF2 +bidYA+h+XZBnHKexAKL74ee71wKBgQDDVHJ6HRo1uTAT/8Uz3AF5NRj+qj9QC3UX +On6dtzbtDGsJFkcqW5CV+LwUPN6zQ7MW7U6AEeH99q1UAUDkd5Q8/4bqs+0y2Ss7 +0UUWr1LDJVYo6FIs/f2ON0e9LzcJcG9TgZ6T95TL1kGEG1Ng96Je3hAZWolLmr+x +yfjTSYKqOQKBgEwRWtTO799hx+dL5C5tTKQx0hUKrvT9IKZ7FjC1xFJ6cLF4l1c8 +KFCD1H8r8yb4fCEMcYnE/rVzIkajmZQIPU0A5bl+b1zsb9Dy6NrSJsM0NvKB/TSz +RuG6mBrw59jkdAOc3J78PJPUQM7/2JzLU0xmVJ7XHpI3G4crkSAIp/YZAoGATX6m +gF4ldOUI9xZFleKWTxFK3laLEeXJybJyY34582g22v8UsvBq96UccWcI79RPLCxw +NY1ivNBuSeLJbRsoG99BFsLVu5O/fFb1cx+R1Uxt14L8f08xlofGFX+y9TK/aEEH +uirCxPA3RANXXCRDLiIp/vUVfYJixVWdO65xgbkCgYAVQhvJ3jPXZLXPl9fVC3Ro +qweEZHs0knPc4mqL/7bE+BCE1HcccU7O7+ZKbhW2q11EMrPJ7cHUqweHE3GsTgRC +C5FIvRxInZUgtPNx2UIZAVNjxRdCZJHBc/59eu8wVp9f5ETiSFmdz5xlsJECSxdO +mBakpZMqnESl1/6FMtcqeQ== -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/ca.crt b/tests-integration/fixtures/etcd-tls-certs/ca.crt index 8723f7ffedd..ecd404c8918 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca.crt +++ b/tests-integration/fixtures/etcd-tls-certs/ca.crt @@ -1,21 +1,21 @@ -----BEGIN CERTIFICATE----- -MIIDeTCCAmGgAwIBAgIUeRD4bMVd0XEzS3UHo2ptnknsWzEwDQYJKoZIhvcNAQEL +MIIDeTCCAmGgAwIBAgIUM4xko7n+KubCgGJnX4EaUzP7ZicwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjBMMQswCQYDVQQGEwJVUzELMAkGA1UECAwCQ0ExCzAJ +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjBMMQswCQYDVQQGEwJVUzELMAkGA1UECAwCQ0ExCzAJ BgNVBAcMAlNGMREwDwYDVQQKDAhHcmVwdGltZTEQMA4GA1UEAwwHZXRjZC1jYTCC -ASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAJ8iLSJOWVRDI1iRr531m3fQ -gsDRAHMsofB7P1AmPxtAul8pHbp3v3gJCf847pc7sxawnqe+rBE1gzRZf5TBS25P -0yu21halLnIQBU7v10y0XruAI1p9cqNDIrwvW/PoC+rhck3hWpkSJ9MV2TezYyZL -S3C/iyCsr/+OHtpnMFAm3D5E571E3V2qUqhrdO2kHuLA1MPEaJHQRiPAeP4Z2I0e -0Mxa+udM01f8gp4TAs/8yOKvOvsado6NHbmu6pwu7Bau5+74QyGoVNSv4shClkPR -MW49foM+VBk6z/hSmO16gLzibkCdBc2WHS/ONtbV58RhZhCgoN6zH16gTmNwAGkC -AwEAAaNTMFEwHQYDVR0OBBYEFOCkI3Uyx7F38LtXtSlg3ORE9ur6MB8GA1UdIwQY -MBaAFOCkI3Uyx7F38LtXtSlg3ORE9ur6MA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZI -hvcNAQELBQADggEBAJXbWJJ9b+5nXjRiFXOg+wQnzn7kzMf6sWKIFk+AuKJWWt78 -O00t6vyAz6zel5Cj3ho9yAaMFNy8vYEnJCYngy5pT/2hOncnz/w7IKTeoEhzqAnf -MCEmCgHbTKDoFfMfrrRwtyoePVfx4xbiGVWMQFTPG+WNlE/ivFMRvFwsgPJ+7SUK -mR2FscH6DVo4sqF6s729lTmr6/U7bOD5l2HYGpKJ2cjCI8+HDv55aAT43tsTwBoJ -BgX+RT4w8ryGhYV+hrRTjPlMHWiHsNbzJeAi5bNA0f6vUP6k6zHn/Ur3MNb1avcf -cSsdYU2dhs7PeB6IpoJ0QCeJw9MIFK3XTnmbeVY= +ASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBALTkN7ZOTxMxnQkwT7tGn7O+ +8uQQE89cIIVxd+j3WBhUcwyvVMZAm5W1vHusioLW0WFzaW4AZ+wGZbch9NdmPNjr +zfhlTOY8vu64F+hGl3Q9qPlfijCmfR4ciByyrHiO4dJNhzd92etMumqU0y6VkR6a +Q1QssITunLGDeGFgjH0kR72e/bp+TKvW9bnPVw+USnI4r5JQh6oeBeJW1Y9T1Ya/ +mx1B5ecPgTUlNfZRGKFvsLzFqWLdhVsjdJTHL+ACEzqzBCBnhIuBbyR/ztr49w/0 +LU3ka6N/Dh9P3R0mi63KJNJtwOuNkuMXTdOCNwNoAK+e1vSo3zlcP3p0i+GDmN8C +AwEAAaNTMFEwHQYDVR0OBBYEFDjvuYaWEoViLOKyCse6mAwH66mIMB8GA1UdIwQY +MBaAFDjvuYaWEoViLOKyCse6mAwH66mIMA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZI +hvcNAQELBQADggEBAJuGxtprq/GF/fwUJ5hk0cG80VY24G+iKKi7Jw7PWPEMvnlQ +/ZIKgNDupE8mSSyokicKZMrtwT8QF6wlDLVcz4LTPCtKUmYU3FiuNE4V57+yw7zt +XZn6rwiWGxeyfY0ZktLXiSYA5FydOJ9GQQWgxkpbd5QG1L9aoFW49cbz8ob6efxG +p+YRgr0Fr4hXn7ImJWWqeawf2tg4iLwpV4sem5atpI7y9qamxrGGUp2h/CnADqUz +YWYIYP1RCPp2bXzxO40Ni1NYmFsGUU/SL1vtJw5n2Zmx/Phe7u/CLe05dH1O54Qf +vCbk6X0CthFI0HzSepbEVAuniVHkrpWwDLHcQ6s= -----END CERTIFICATE----- diff --git a/tests-integration/fixtures/etcd-tls-certs/ca.srl b/tests-integration/fixtures/etcd-tls-certs/ca.srl index 69ac026da25..ad32d06e16e 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca.srl +++ b/tests-integration/fixtures/etcd-tls-certs/ca.srl @@ -1 +1 @@ -4484A9BBD4F25F32994F3C03D80D294105B119E6 +4484A9BBD4F25F32994F3C03D80D294105B119E8 diff --git a/tests-integration/fixtures/etcd-tls-certs/client-key.pem b/tests-integration/fixtures/etcd-tls-certs/client-key.pem index 13a23c53478..1d64fc66a3d 100644 --- a/tests-integration/fixtures/etcd-tls-certs/client-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/client-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQCOcrzRmE3sEGYo -pazDrpkkzbiaBhHClfmBiZAAGV2cvSXY8e9pc21GOgrGBNeyPPieS305rAE02EZl -oIAt97XRFxkZniUy8C5+PT3eicAKPRcrgGhIrYxUs39rdRocqVVpZsiZGzhHhboB -bXNW1l37SDQpuxAQJuhHu2wu9J7c5/TBXWN9ZOxzWTlrIdEQOA46VclE/N1TlaKs -CSWw8qH87YLwkLEXcFLOJfmWNZxOOhdL5MSzKKNcQKbpmkd65B/PHNwBIBZHpI+K -xphK1SybA+ESkN09qZaa1Kj3e0I62UjHNXUCuAod6RhRRmpiFDFgggXhD9T9Ttuf -1Lq38XxDAgMBAAECggEAANNhqQldwHzi9WCxXObLoufBe9lM73+1Df1IQaDHkZR3 -jHliixxF18XQW7qZ5tAJqjsDWLPm4BxSKXozsjsvBdgIUjbqkyx1dsGlO9TjeGu5 -H2/8jlH14boIMWJ9Y0JBWwvU7EnVU5IkUOEiwsxvFlzM5Q0H+841CrSD+Xd7iSZE -ZrpyZgpHNmvc/u7AmseZVccjX21dxwWtcd7cGUmFWC9o2a/hkgG6+baA8zH2ze3J -Smeowa/oOyOVkTM1EHLCBvddE+/W+BeMGPY7A6SztVw0gLqlSOL3cr/g2llUTmFg -T6MOS6FZZ7MD5c2HBL0V3OXIhxNABzD+c0O3CaRbSQKBgQC/+QuhMpDCzxBPaL/K -sqRp4sQgURDCXPA1nAfI+NoHos/CbjwcMGp6LUSiBW0tLTJ+LZoVgZXj6j/uxji9 -DWxQfULvQ7oanm/q/D+40NnupQQTzmMT04z/mV3XjFtzEhbO30YZfqwXy+tqB8vR -g1KajXVqLmCxKDtsFedAlIVxRQKBgQC99TKGcBwsKlYJJThU2eIIy8G0GkjjXleK -ROMO5i0keJtLac1oRvXo4JN5UBLJM6ZpSy0TEa2bxfzy6Uea+41aZmdCpPg0Kcsq -Pb1S3rsPpVdsfxSxsdKvwZmHXOzAdF+fom1+ZRAdo/8ALIPaNQSSaw+KkmmpHPEV -qdzzWJ4b5wKBgFubEtKcF3nuZwENoh+ueUhRvncRV+b3hGSAjTJ4lUn5hhxoj+R/ -sf+VJGAQKNXa8HJHfnRuvsDgYhulmSOViS8rZspXzjGvkwZV0m51stjvA3AUFzE5 -zNmXLLGTt3vEkP+siX3W9XXxh+ezyq2ydbNsdy/w65D9+sUL+qrVdIvlAoGAOOQy -2ajCB0g2tE59bIxE8jV0MiidI9uhhDvVdSTi6EVm3VM2vcBi7fg0suSUe8YIVQi6 -2zc0M688btQHKhek4ipBSuh1ncnWmzQae7NRewIeCNSWshF79D+bZ7sg/RLdgMX4 -3R4PkZEIUlkCtFuknuWJpgrrskaEveQ91HP6BokCgYAX42nKPANIMrTXuJ2vACWu -qKQ9atAC1ngVzYe2PI8rHjA4gCDraGfCaILewrd89rK9ukBeSDIm5Mn448q3ystw -p/QMxmMdHADCD6UwCjlt9N0Xigv+IKbDzDzhartRVIPpZdTViQXW3cSHt9id8zvO -HTu5C6W+N9Vb7j8Ak/mvzw== +MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQC5F3hVsri2FBzP +dtRe+aJFh9yqhXnJ7zN+xbUmNWsTZVnPwrbVo9H6PtGWzkbZ6YMBzMlQSYP2txaS +ROjE8wwEs38eKbyJrHO2Ccq3iR9NfDxSbHTClngsOjf9fEoYExif84VKWLO6+EKL +KqqOfj3+iuI3hbabr0MqGu4e2Nr6MlQtViBxTs5aSnahc8hmZbujBcWRTsw/GNx5 +Am0a8BrrJRVmVR+BzlYtzbf1ltfLxIhnc8dAGyT9B6GTwcMNST6m4WtTzR/uOr6S +2YEbIqwCHB1VQ+oJUGJvExqaNLoM/x5a8sUNIWx3LBN0/XC4D+MMmh/bbyb2cAOE +gyUYu+hLAgMBAAECggEAEkt+wksVP9toEkecFpEqgss65VimdkPK/WkHDpzxawGL +FDfvq4Py9tRrvLIXRbseObjSCI9jrai2Ne9Sv1jeDhQ8kzDWk8MfItX/8A6Whf8A +KrjRUlneIfz/2Mcxaa9wMSs4YxC2dIYidFqwxhSLWyluJ8Ud4hs7kndb+beUmW3i +DgkLsQ9ytQ0HDQrAERnrbRf19fmVaqQapfxpYT2gg80BMIq+zlzoiXareDz3r/1g +aoponxArWV6yCvQOYYjWcDr4WZPvIA2APEX5isxyvMBvb72Fpd2z+0JX+4LVe6N1 +FqZLNpSeJn6O+Q1vJSScSJcoJYkLWCfZbFLHPttZGQKBgQDdVr5PtgPIVM+67cvv +p3x7NNC10Gp387PgQS/FlG+QtbP87RTBnANT1uFd/FMFqL35or7Y46YKDUdk00OZ +LinpM4ARFRYFqkthM6bcUWatpF6Fg2dfsKuD4H6ARp8YHhlmCnDCQ0swZp1HCmAB +6vltUNTZjQjRdd44GlQ0icX0mQKBgQDWE50TXRszea75S+zUfXeyvxMyptIOS/rY +L50XVuhGVknH4OuBQbGF6UZulIGSpN7gmmbO29Yx1HKmqiJVaWlvB71baK4hWLDd +2grIUy0zuq1B7I1IOLmOpMcL3xau76g7ZWRUprX26hxAuxYyPb9V8gNo1P7HN8si +VLJVE8ZugwKBgQCaxoGmM91JRSVNzeOB3ljJvxEDUo5g+uWZt3u0aivpwWXvQ8nz +6Sjag7RsiHl1x52w5wEVoXsGJGr8Mk9e2k0saXrwdxJDO+YiPoA8KB/o5LvEGTM8 +UspdGarcAIZX0xRnqn1XGr+FRPxOJQ8lyC5LJu7wghLchdOy35Zqdr0aYQKBgQC/ +38yFsonS1VnS8A5RVjOW7lPSrlrPnaIzalmutaJyiJyQnjP3Il5u2+rY6hpIyaVK +QpmrBrcw6m3om80yKMzrS1CZQXXxRYEhF3Fao9J77vGjiNYIyW7nPyF4rneyS/PJ +aNNIXDP0H1k7W3RFi7qW2dfceivxezyChM9iGdtc6QKBgAIfb63+1q98MNlAB7eD +x7K8vuSxHrJ9gX9KcLrG1ScLOxnAC4Op+8xf06ZHt828CajchAOqee2nh52MsFOr +AdPNdhS/8mTq8KosTfgVuawv4yTHpZViMR2KSOcTKhKaBmhufY+k8mJ0YAAOTnzO +JJqjDSJdjWsDG+7UlNgFcvuU -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/client.crt b/tests-integration/fixtures/etcd-tls-certs/client.crt index 935351a48b1..0dee5915ffa 100644 --- a/tests-integration/fixtures/etcd-tls-certs/client.crt +++ b/tests-integration/fixtures/etcd-tls-certs/client.crt @@ -1,20 +1,20 @@ -----BEGIN CERTIFICATE----- -MIIDMjCCAhqgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGeYwDQYJKoZIhvcNAQEL +MIIDMjCCAhqgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGegwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjAWMRQwEgYDVQQDDAtldGNkLWNsaWVudDCCASIwDQYJ -KoZIhvcNAQEBBQADggEPADCCAQoCggEBAI5yvNGYTewQZiilrMOumSTNuJoGEcKV -+YGJkAAZXZy9Jdjx72lzbUY6CsYE17I8+J5LfTmsATTYRmWggC33tdEXGRmeJTLw -Ln49Pd6JwAo9FyuAaEitjFSzf2t1GhypVWlmyJkbOEeFugFtc1bWXftINCm7EBAm -6Ee7bC70ntzn9MFdY31k7HNZOWsh0RA4DjpVyUT83VOVoqwJJbDyofztgvCQsRdw -Us4l+ZY1nE46F0vkxLMoo1xApumaR3rkH88c3AEgFkekj4rGmErVLJsD4RKQ3T2p -lprUqPd7QjrZSMc1dQK4Ch3pGFFGamIUMWCCBeEP1P1O25/UurfxfEMCAwEAAaNC -MEAwHQYDVR0OBBYEFINcPBk3iqCw1B2sZzcDpSF6ZNGnMB8GA1UdIwQYMBaAFOCk -I3Uyx7F38LtXtSlg3ORE9ur6MA0GCSqGSIb3DQEBCwUAA4IBAQBgyF1XSpAyiQNA -sFCqXrBfwPnv6V+R5jhO6Glmn+Nhj0gbzNBlU8EKsK/WZllDrJjmxubPBqb563Zt -J+QcNLqugZdioXPbxmmi9xj3oK25pUCRBj03nlHpGqIwuCJyn/ZxlfXMpNYE1KjM -nprFNUKgnYPk2AGffAngQdm0yrnmSpxXhNTuhQE8avgRhd5nEUEDyPIvHWwo4iU6 -pw2aP2FpWorESv9LVna+JortThYDQrIFgCtPM6BP1N7Aeqn0H846+JsXlq2ch1Iy -aBRqKNQRc1Q2Qb9/QCDEJ3dEixgFjxztIaqrIcr18ZJ797MhTFS3aVEwE8clSRR5 -9ow9jK8B +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjAWMRQwEgYDVQQDDAtldGNkLWNsaWVudDCCASIwDQYJ +KoZIhvcNAQEBBQADggEPADCCAQoCggEBALkXeFWyuLYUHM921F75okWH3KqFecnv +M37FtSY1axNlWc/CttWj0fo+0ZbORtnpgwHMyVBJg/a3FpJE6MTzDASzfx4pvIms +c7YJyreJH018PFJsdMKWeCw6N/18ShgTGJ/zhUpYs7r4Qosqqo5+Pf6K4jeFtpuv +Qyoa7h7Y2voyVC1WIHFOzlpKdqFzyGZlu6MFxZFOzD8Y3HkCbRrwGuslFWZVH4HO +Vi3Nt/WW18vEiGdzx0AbJP0HoZPBww1JPqbha1PNH+46vpLZgRsirAIcHVVD6glQ +Ym8TGpo0ugz/HlryxQ0hbHcsE3T9cLgP4wyaH9tvJvZwA4SDJRi76EsCAwEAAaNC +MEAwHQYDVR0OBBYEFD1aTqJO87ovA2usg/v8UckSsnwhMB8GA1UdIwQYMBaAFDjv +uYaWEoViLOKyCse6mAwH66mIMA0GCSqGSIb3DQEBCwUAA4IBAQArP3mtjdj5a9uT +XTys4DUjp+CJaSmsZ/4ICf/KeddaEyijx6uoiqbNsAtfwlD5A5BtyjJrxg3SpjqD +Fm5mTV0HlwvYHxcFPTUlR56//sOyeRhm+kVPp4xdVq2p6IZXFF/hDkeT/AETWOjj +l4+90gAsJsUHrKUZmzlqHBeOXoAhzo2PwSh8KLcGK0705FGql4QnS731PHo9a9qp +e8uxk1KKoc+RHm2OjJvZ7hvkCXuCm3TJeOHQEtjZ5buQg40BRrt11GVnsPePpnQR +8OtV0DvPT3UcOss88kLDQuRuZeJf8UyaevGGvItv3+Yk5cwUBSaJ0mY6sgrmkXn0 +TstUzYIT -----END CERTIFICATE----- diff --git a/tests-integration/fixtures/etcd-tls-certs/server-key.pem b/tests-integration/fixtures/etcd-tls-certs/server-key.pem index fca8a0122f9..a73ac974434 100644 --- a/tests-integration/fixtures/etcd-tls-certs/server-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/server-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCtYFxH+oet6jTD -0eF55ISO+bwm1NPTluP9zLJrQFNDkqlddyiIasK5q7qN1pPXOw9mJF/EN8Hf1lWj -42an6tr7PgE0NHLcOrqDv4dRjegRi87IWiPkfep8peUE3glI4IOMXfZmAUNp+nRO -JP/ymmR0pzHYJlx+D082OYmeBny7uFQxo+vxgVNo2MVVwy6oTv+jtjLgZ4iOrIvh -iILoXtXcdg/tjB1rS2hSABGPJMNHwKswcrt6xXv+7Llsxroq0cx0E2PrAbc5qB40 -v1Qp7PjFTJ0drgf0g51LkOb7znpB1viPD2BLMxpELTZqa4OruarcDNxShrK8WWMP -1mygLHLxAgMBAAECggEAC4TgtLaQP62VImKGLs1QQliU1+ahiUgX70Ojog0UyyNK -JefWFVQ0ilX+z9AvI+hsYj6t7zE+LBNHPtuLtUHdGT66IUAP1pJ/VGQMB07cmZfW -ngihJFv6UZxLDkr7RnCGRPP0PDw+wKKPiiaaq8F2xapbHS+VSxnUyzdA7bMkI+uj -+gLAFj+kEU0FNbIp/mBzz8sksDzcB29oy3VnZkR1211KkL1BJDvr9oKiUg8NFKaN -NF2WXRYI/m0CQMlX6zlTyv7+h/Ay7y8oNG7cfHavmCrWvYIMfaDggYxbIpNdGxb5 -MAdbrtXJ6daUBa2sezxgQUyJX8+8GGF2zpGvCfhkfQKBgQDiAwA8rlxrBnXclwRK -9s/apILKCxPEGwU6XZi/8na3ugIO14y2WD0zZTjUK/3xwoKN6j6+ObLkUBX2RRYZ -ewwxpqos2cyxi5EBKPhsydqbJQlat68I5+NOG4KOZQkqO92tBtgY3gGFWVntaqtE -UVE/Clh2gFcc/5nR/5C5Mil8LQKBgQDEYXtrl2kl4qcyXuuXKDA/bLSnDuIEMhLO -4oIiqmywXQgjPjGAWy0eztjMq2a353pqTFCqto/cfOif+rjkB7oFfi9gMm0plIKI -+F8QsBB7HpGqWos3zb1yt8X9hbxKuFFwZS/un3Li4lwPlk8EZ1YIvhC44BHAvNkD -VzxzBwoYVQKBgArRe/RroC7bS074x5LTB5X+o+gJ6bNMW860Zjhh4b7fn3OYa7ra -tGs+YB7/0BL/bYJfgQtX9bEqCDMWkX08v5Os1553+m1RMeqtTF7gtp8Qgcce3bj+ -aIn3lSM9wNeNsAm1NyjRj58TbNOJdJM7lTkARMW/VOwla/Z6VjIXLZctAoGBAJKf -xira7eMfi36MaJJ/qyZv36Ir9ozzZh+Z91gyrtwvWfgWY5dWfCXYgv6tqw/8gOYE -/OW5UUhq6rUn2gxHyJh5Up4ciGzXOW9TIoevLV7/v/rVh8SulJimpelYhPG1FPk6 -U8NywbCtGdd5fp3nGdGFN68Rfa/OUKmx5KxtwRfRAoGAWpfzmoI8EmCBZNXYk03k -hnDKmICsKi9seDclcQq2PbyNvE7LDEd8O4pDnwV6RwKVEDTer9Y4IGcjkEBZcJVm -za3l4En9+M27Xlf1I8XeB8o8RYbFjoDQyScFFIgQpvTDsOezXeDR55FGPZN7V6Si -lKmd+8BuzJ7d31KJtRjoS1U= +MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDQLtVfdk9l1rdH +/ssuy/BK9LpmNPx7gatp/E3EqpgekBjizMXIxnCHDFNM4Dtq1v3afaxeijXnnW/b +0CVKSeGkOBmb0dO0K21xe86Y5VfNHkGotQyCZ8mPOzep/EQOGJY5+dprErL0YKM+ +2drgiRNGwidmNk/jD3OtzXSMR+8UvDWJSxUoZ1ncvgmqf3zTTiaTiABHrfLdXaGI +23+JftyLgoeU4vEI+x+RQJMAa6Gzjo9vALxX3kTtR/uH34hfm8vb3yo68UHGKVMp +i3yi6YDP4YEao5lOpm7fIyEes10fPkD224++S34KHykAk4q+eTPetWkDYFCYzdSp +AR2Ob3cxAgMBAAECggEAMDgAxO0mw7xBVGQgFJU48WuQtvaj2kl09gXxz1ECDeYr +VXC/iNrpmmYQ7zfqmzrzrkU4hOc3SA/PplamJHhLUpmJ2Oz3P35liYj3F7PbK8/L +vnM81AGNDmdVY8Jh0u//76q+29kHaRHvDbIw/5vQQq3aqVKAG2PrU8DIM2u/5Qmt +0GHJTMmyVADAOLe7j2agN9hrOrjeoWobcELwtwgKXxmU0l+g0XX1OdJpbjA3qDPa +S5a4MzcucGINeB3bohcqbcx2f1HwVUoC19mti3Cz+KvxF2YZV9bS9izW4hSEsAwU +0upNHF3MwTUQg7ftLPAZcfNoXIZMFXwYMS1HW2v7HwKBgQDuHHkFRN2lvor/Ycjr +gJD6LGb9RlQ7NMYBQHB+5g49Y/q8oDZJgyl4CfgYWdVjjOfEX8oMst1WU/Stbva5 +HlHR7yYFyebScQYOPHZaPyVKGIbcl2p7Qy9H4pxJc1ZMNxO8SxoSC68LZEan08FD +/Qsish51fZ2Sqf37b+8Uz/CGdwKBgQDf0sIp88MPGbhtzdYtF9gA4AoIOPo9UxeI +F6wTb05uCqZyK+hcdK37av1xcNpczd/yvlALxRdM9xGedJkcS2cgeExtMRcfhAzL +as0dYPzuhRkpEIfN/y9cqtHj2wigrbu35eDgRRnJDgXgL9dYl6r+Swqn3jBaasP7 +KtQdQb/RlwKBgQC7qe0n3fLi4p4iUStNkPKyebRiAb/5Ocqkyejf2ul2MQo5B/xB +TAKu/QxwBL1NzIwOFYDlKUOQ+nJpDn+dvuu1jcpl3Y7yZOnk5npQ/luhXltMGHpv +06+79DpBGYn2X6JKUNanSlYXoFyfgSFdOF5CZifjabF7Gkd2l+3SdWCYWQKBgQDB +3k4sBGZKaB7ljUscl/CTIXvPD3tBLv3M9aQo2Vp32mW9suZ7Xt1sTonkfrnFdNWr +7shqyXabRc5PD/OnHHDhIRIh6kl7FOf4MjQkZGPxPfxDI3xeI9EkVRmkYY6hjppw +eX9FAtWI3sqcGxROOmD0Do/WQ5BiYOQMZFaCWPcLVQKBgGVT+iUYkYTs+UJeD4RC +QeL6QEJDszvAI/gavUp7sD/LWn7Ig+hJiAl2O+HFYMurh7eda0JnE2KF40UyPiYA +DDl19xRAXrLPOlOj9YKDA9HkROcaVZCQCqyDalKsbw1IqMrAkRmLCIZW3EHaMrUf +UNjg+GaxwvrRj11gbd4IPgai -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/server.crt b/tests-integration/fixtures/etcd-tls-certs/server.crt index 6ed8b1054e7..8dcb7b13b6d 100644 --- a/tests-integration/fixtures/etcd-tls-certs/server.crt +++ b/tests-integration/fixtures/etcd-tls-certs/server.crt @@ -1,21 +1,21 @@ -----BEGIN CERTIFICATE----- -MIIDjDCCAnSgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGeUwDQYJKoZIhvcNAQEL +MIIDjDCCAnSgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGecwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjATMREwDwYDVQQDDAhldGNkLXRsczCCASIwDQYJKoZI -hvcNAQEBBQADggEPADCCAQoCggEBAK1gXEf6h63qNMPR4XnkhI75vCbU09OW4/3M -smtAU0OSqV13KIhqwrmruo3Wk9c7D2YkX8Q3wd/WVaPjZqfq2vs+ATQ0ctw6uoO/ -h1GN6BGLzshaI+R96nyl5QTeCUjgg4xd9mYBQ2n6dE4k//KaZHSnMdgmXH4PTzY5 -iZ4GfLu4VDGj6/GBU2jYxVXDLqhO/6O2MuBniI6si+GIguhe1dx2D+2MHWtLaFIA -EY8kw0fAqzByu3rFe/7suWzGuirRzHQTY+sBtzmoHjS/VCns+MVMnR2uB/SDnUuQ -5vvOekHW+I8PYEszGkQtNmprg6u5qtwM3FKGsrxZYw/WbKAscvECAwEAAaOBnjCB +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjATMREwDwYDVQQDDAhldGNkLXRsczCCASIwDQYJKoZI +hvcNAQEBBQADggEPADCCAQoCggEBANAu1V92T2XWt0f+yy7L8Er0umY0/HuBq2n8 +TcSqmB6QGOLMxcjGcIcMU0zgO2rW/dp9rF6KNeedb9vQJUpJ4aQ4GZvR07QrbXF7 +zpjlV80eQai1DIJnyY87N6n8RA4Yljn52msSsvRgoz7Z2uCJE0bCJ2Y2T+MPc63N +dIxH7xS8NYlLFShnWdy+Cap/fNNOJpOIAEet8t1doYjbf4l+3IuCh5Ti8Qj7H5FA +kwBrobOOj28AvFfeRO1H+4ffiF+by9vfKjrxQcYpUymLfKLpgM/hgRqjmU6mbt8j +IR6zXR8+QPbbj75LfgofKQCTir55M961aQNgUJjN1KkBHY5vdzECAwEAAaOBnjCB mzAJBgNVHRMEAjAAMAsGA1UdDwQEAwIEMDBBBgNVHREEOjA4gglsb2NhbGhvc3SC CGV0Y2QtdGxzggkxMjcuMC4wLjGHBH8AAAGHEAAAAAAAAAAAAAAAAAAAAAEwHQYD -VR0OBBYEFK8y9gWv8Si9AusP280BrgDF6y3jMB8GA1UdIwQYMBaAFOCkI3Uyx7F3 -8LtXtSlg3ORE9ur6MA0GCSqGSIb3DQEBCwUAA4IBAQBMBsAPIlk17I2ioN9xwxwl -yYbSjZlTs+18wSZAoCNLfxwWIYAQ08dPoEUdALfMUPEEe1Ol2IAp/qx2JoDrZWje -maA43hppBUpFFCkTKWUMYxsetN7d7BVxWL43GDoMwoD/k36nhoUKBUjlbF0+nkem -dOMr8SA4GZbEg7qk2cL93g0UHv/Z/dZgKf3epZkR9hEN4/R2jSP6OPmY9XIvQCoZ -8D9jgbQKavAAXmUcP5a81alQMRroGaBCzI1f5OlS3EuVE7ZTEBgxKK1idbouSrPt -4UHEbOGz9zXDjmfut6CTm247+lJm9jzYe2Xx+XGw29l0pzd/8tGFN9zIvW8mCv7a +VR0OBBYEFNqGaSuuVHwEIvQvYxt9H0kdWFPSMB8GA1UdIwQYMBaAFDjvuYaWEoVi +LOKyCse6mAwH66mIMA0GCSqGSIb3DQEBCwUAA4IBAQAHttiuyEZieeTsggAKzAIC +UXsjWCQMAmjRxery6mzM8P33CdCKA4cwGyy6tzoFPiZC/wh+gPq+LiUO9G+vkqNA +NQ/DSJS+fsqsOBWUt5DSYekD25o80mZVE/4H/ynFOgENBqgkVl1ikpkXu9BK8wTm +Wij+p9RWyNup8lr9OFYZRpy7wrxPp7opwp3qG2Bfu/BGY1YEr04CIrLMKnk1zRMm +73X1lcrlrM9ycCmYcfXxyU69G0iFKaPcbS7T0fngm1NciA1qxYYLuMyzDtF7/SyC +J2ynvDBBnc9+/X8piu8p9D0I9mbcbEySv6zZfYMT0kqXbsn2GijZTH5mCJ3P8cwc -----END CERTIFICATE----- From 35ea88a4ef4b06c8c5d10e57cc8bad44f4f38569 Mon Sep 17 00:00:00 2001 From: localhost Date: Wed, 26 Aug 2026 12:11:14 +0000 Subject: [PATCH 110/126] feat(ci): run query regression on ephemeral Aliyun ECS runners (#8937) * feat(ci): add aliyun ecs ephemeral runner path for query regression Signed-off-by: paomian * fix: improve condition for query-regression job execution in workflow * feat: update Docker installation to use official repository and add GPG key handling * Refactor query regression runner setup and configuration - Removed deprecated PersistentVolumeClaim for build cache. - Introduced a new bootstrap script for setting up the ECS runner host. - Deleted obsolete Helm values files for runner configuration. - Updated the Aliyun ECS runner provisioning script to reflect new cache paths. - Modified GitHub workflows to use the new Aliyun ECS runner setup. - Adjusted documentation to clarify the new runner lifecycle and provisioning process. * fix: enhance runner service management during bootstrap process * fix: update alibabacloud_tea_openapi dependency version in metadata * feat: enhance ECS runner scripts with region_id and resource_group_id support * fix: move containerd content store to data root for improved storage management * feat: rename query-regression runner to ephemeral-github runner and update related scripts * fix: update sentinel polling method to use serial console output for improved reliability * fix: add environment variable checks for Alibaba Cloud access keys in ECS client * fix: improve error handling in GitHub API requests for better diagnostics * fix: improve cache disk detection logic for Aliyun ECS instances * fix: enhance cache disk waiting logic with detailed output and error handling * fix: update dependency version for alibabacloud_tea_openapi in teardown script * fix: enhance cache disk waiting logic for better compatibility and clarity * fix: enhance console output handling and add incremental logging during instance provisioning * fix: add PATH environment variable for runner jobs in service and provision script * fix: add machine telemetry sampling and logging during query regression jobs * fix: update query regression documentation and provision script for cache disk handling * fix: update SCCACHE_CACHE_SIZE validation to 10G for improved caching efficiency * fix: remove outdated cache size checks and cleanup logic for fresh system disk runs * fix: enhance instance deletion logic with region handling and console output export * fix: add swap file setup and OOM handling for ECS runner to improve stability * fix: update OOM handling and service restart logic for ECS runner to enhance stability * fix: increase system disk size to 100 GiB for cold double nightly builds to prevent ENOSPC errors * fix: increase system disk size to 150 GiB for ECS runner to prevent ENOSPC errors * fix: add keep_instance option to preserve ECS instance for post-mortem debugging * fix: disable unattended upgrades to prevent job cancellations during library updates * fix: reduce system disk size to 40 GiB for ECS runner to prevent ENOSPC errors * feat: Refactor Aliyun ECS runner provisioning and introduce nightly regression comparison - Update `aliyun-ecs-runner-provision.py` to remove cache disk handling, simplifying the provisioning process. - Introduce `query-regression-nightly-refs.py` to resolve and compare SHAs from successful nightly builds. - Create `query-regression-nightly.yml` workflow to trigger nightly comparisons based on successful builds. - Enhance `query-regression.yml` to include a `test-tooling` job for validating Python scripts before provisioning. - Update tests for the new nightly reference selection logic and refactor existing tests to align with the new caching strategy. - Modify documentation to reflect changes in caching and nightly comparison workflows. * fix: enhance runner image tool verification with detailed checks * fix: improve error handling in runner image tool verification * fix: update tool versions in ECS image and workflow for consistency * fix: correct typo in error message for unparseable ECS creation time * fix: update README and workflow files for query regression tests and image hygiene --------- Signed-off-by: paomian --- .../query-regression/README.md | 453 ++++++---------- .../query-regression/cache-pvc.yaml | 15 - .../ecs-image/bootstrap-runner-host.sh | 195 +++++++ .../ecs-image/build-ecs-image.py | 392 ++++++++++++++ .../ecs-image/ephemeral-github-runner.service | 45 ++ .../ecs-image/start-runner.sh | 37 ++ .../query-regression/values-8-cores.yaml | 146 ------ .../query-regression/values-paused.yaml | 2 - .../scripts/aliyun-ecs-runner-provision.py | 485 ++++++++++++++++++ .github/scripts/aliyun-ecs-runner-teardown.py | 237 +++++++++ .../scripts/query-regression-nightly-refs.py | 270 ++++++++++ .github/workflows/checks.yml | 8 + .../workflows/query-regression-janitor.yml | 36 ++ .../workflows/query-regression-nightly.yml | 104 ++++ .github/workflows/query-regression.yml | 433 +++++++++++----- .github/workflows/release.yml | 2 +- tests/perf/AGENTS.md | 20 +- tests/perf/README.md | 7 +- tests/perf/test_aliyun_ecs_runner_scripts.py | 122 +++++ .../test_query_regression_nightly_refs.py | 120 +++++ 20 files changed, 2538 insertions(+), 591 deletions(-) delete mode 100644 .github/runner-scale-sets/query-regression/cache-pvc.yaml create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/start-runner.sh delete mode 100644 .github/runner-scale-sets/query-regression/values-8-cores.yaml delete mode 100644 .github/runner-scale-sets/query-regression/values-paused.yaml create mode 100644 .github/scripts/aliyun-ecs-runner-provision.py create mode 100644 .github/scripts/aliyun-ecs-runner-teardown.py create mode 100644 .github/scripts/query-regression-nightly-refs.py create mode 100644 .github/workflows/query-regression-janitor.yml create mode 100644 .github/workflows/query-regression-nightly.yml create mode 100644 tests/perf/test_aliyun_ecs_runner_scripts.py create mode 100644 tests/perf/test_query_regression_nightly_refs.py diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index a1d58de7166..ce447dbdd0a 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -1,321 +1,184 @@ # Query regression self-hosted runners -The `Query Regression` workflow uses the dedicated ARC runner scale set -`perf-regression-8-cores`. ARC runner Pods run in the target Kubernetes cluster -and connect outbound to GitHub. The live scale set is currently **paused**: -`minRunners=0`, `maxRunners=0`, and no runner Pods. Do not resume it without -explicit approval. +The `Query Regression` workflow runs on **Aliyun ECS ephemeral runners** +(`aliyun-ecs`, the default and only automated path): a `provision` job on +`ubuntu-latest` creates one pay-as-you-go ECS instance per run, and the +instance registers itself as an ephemeral GitHub runner with a per-run +label. A `teardown` job (`if: always()`) deletes the instance; +`query-regression-janitor.yml` sweeps tagged leftovers older than 4 hours +daily. -## Prerequisites and trust admission +Build caches live on the instance's system disk, so **every run compiles +cold**; only the within-run reuse (base build warms the candidate build +through the shared target dir and sccache) applies. There is no retained +data disk. -Install the ARC scale set controller if it is not already installed: +Dispatching with any other `runner` value uses it as a literal self-hosted +runner label, which is how a manually prepared host (see +`ecs-image/bootstrap-runner-host.sh`) runs the workflow. For PR labels, set +the repository variable `QUERY_REGRESSION_PR_RUNNER` to such a label to +redirect PR runs away from ECS. + +The office ARC scale set `perf-regression-8-cores` that previously ran this +workflow is retired; see git history for its values files and pause/deploy +procedures. Tearing down the office cluster (helm release, runner namespace, +and the `query-regression-build-cache` PVC) is a manual operator action +outside this repository. This directory keeps its historical +`runner-scale-sets` name for path stability. + +## Nightly vs previous nightly + +`query-regression-nightly.yml` runs after a successful `GreptimeDB Nightly +Build` (`workflow_run`). It resolves that run's `head_sha` as the candidate +and the previous successful nightly (same branch, typically Friday when +Monday's nightly fires) as the base, then calls `query-regression.yml` with +those immutable SHAs. Both binaries are still compiled on the ECS runner; +this is not an artifact-download path. `workflow_dispatch` can pass explicit +`base_ref` / `candidate_ref` or a nightly run id. If there is no previous +nightly, or both nightlies built the same commit, the comparison is skipped. + +## Aliyun ECS path + +Configuration lives in repository variables/secrets: + +| Kind | Name | Purpose | +| --- | --- | --- | +| secret | `ALICLOUD_ECS_ACCESS_KEY_ID` / `ALICLOUD_ECS_ACCESS_KEY_SECRET` | RAM user scoped to ECS RunInstances/DeleteInstances/Describe*/CreateImage/RunCommand. Used only by provision/teardown jobs on `ubuntu-latest`; never reaches the ECS instance. | +| secret | `GH_PERSONAL_ACCESS_TOKEN` | Creates the short-lived runner registration token (shared with the jsonbench EC2 path). | +| vars | `ALIYUN_ECS_REGION_ID` / `ALIYUN_ECS_VSWITCH_ID` / `ALIYUN_ECS_SECURITY_GROUP_ID` | Network placement. The security group should allow egress only; no inbound rules are needed. The vSwitch pins the zone. | +| vars | `ALIYUN_ECS_INSTANCE_TYPE` | Dedicated (non-burstable, non-shared) instance family. Prefer 32 GiB (e.g. `ecs.g8i.2xlarge`); `ecs.c9i.2xlarge` is 8c16g and nightly thin-LTO of greptime peaks above that. Both base and candidate clusters run on the same machine, so noisy neighbors break thresholds. | +| vars | `QUERY_REGRESSION_ECS_IMAGE_ID` | Custom image built by `ecs-image/build-ecs-image.py`. | + +The system disk is 40 GiB, which covers the image, a 16 GiB swapfile, the +checkout, and cold build caches (target dir, cargo registry, sccache). ENOSPC +stops the runner itself from writing logs, which GitHub reports as `The +operation was canceled` with no telemetry, indistinguishable from a +platform-side cancellation. +cloud-init masks `systemd-oomd`, disables `unattended-upgrades` / +`apt-daily-upgrade`, creates `/swapfile`, and sets `OOMPolicy=continue` +on the runner unit. Ubuntu 24.04 defaults to `DefaultOOMPolicy=stop`, +which SIGTERM-s the whole unit when rustc is OOM-killed and GitHub +reports `The operation was canceled` with no telemetry. Unattended +upgrades can do the same via `systemctl restart` of the runner after a +library update; GitHub then records `UserCancelled` even though the job +was still valid. Swap is a safety net for 16 GiB types, not a substitute +for 32 GiB; linking on swap is slow. + +When a run dies with an unexplained `The operation was canceled` (no +"Canceled by" banner, healthy machine), re-dispatch with `keep_instance` +checked: the teardown job is skipped and the instance survives for +post-mortem inspection. The security group is egress-only, so inspect via +Cloud Assistant (`RunCommand`) or VNC: the runner's `_diag` logs under the +runner home record reconnects and worker crashes, `journalctl -u +ephemeral-github-runner.service` mirrors the console stream, `dmesg -T` +shows kernel OOM kills, and `machine-telemetry.log` in the job workspace +has the 30 s sampler history. The janitor sweep still deletes the instance +after its TTL, so finish the inspection within that window. + +Trust model on the ECS path: the instance receives only the one-hour runner +registration token via user data and holds no cloud credentials; the Aliyun +AK/SK exist only in the control-plane jobs. Instance tags +(`managed-by=query-regression-ci`, `query-regression-run-id`) feed the janitor +sweep. + +### Building and updating the ECS image + +The runner `Dockerfile` in the parent directory stays the single source of the +tool contract. Build a new ECS image from it: ```bash -helm upgrade --install arc \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set-controller \ - --namespace arc-systems \ - --create-namespace \ - --version 0.14.2 +ALIBABA_CLOUD_ACCESS_KEY_ID=... ALIBABA_CLOUD_ACCESS_KEY_SECRET=... \ +uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ + --region-id --vswitch-id --security-group-id \ + --base-image-id ``` -Create the GitHub App secret in the runner namespace. Prefer an App limited to -`GreptimeTeam/greptimedb`: +The script boots a temporary builder instance, `docker build`s the runner +image, materializes `/opt/rustup`, `/opt/cargo`, `/usr/local/bin` tools, and +`/home/runner` (actions-runner) onto the host, installs the ephemeral-runner +systemd unit from `ecs-image/`, and snapshots a custom image. It prints the +image id; set it as `QUERY_REGRESSION_ECS_IMAGE_ID`, and bump +`RUNNER_IMAGE_EPOCH` in `query-regression.yml` at the same time so the target +cache invalidates. Builder sentinel polling requires the Cloud Assistant +agent, which Aliyun public Ubuntu images include. -```bash -kubectl -n arc-runners create secret generic greptimedb-arc-github-app \ - --from-literal=github_app_id= \ - --from-literal=github_app_installation_id= \ - --from-file=github_app_private_key= -``` - -The values files here reference that secret by name. +## Trust admission for PR runs A maintainer applying the `query-regression` or `heavy-regression` label is -**trust admission for that exact PR revision**. `query-regression` runs the six -routine default cases; `heavy-regression` runs only the high-cardinality -`prom_remote_write_7913` remote-write case. The admitted job may use this scale -set's dedicated, writable persistent cache. `pull_request: labeled` is the only -PR trigger: the label event snapshots its merge, head, and base SHAs. A queued -job fetches that immutable event merge SHA directly, verifies it is a two-parent -merge whose parents include the snapshotted head exactly once, and uses its -other parent as the actual base build revision. The snapshotted event base is -retained for audit only, so a difference from the merge's non-head parent is -not a failure. The job never follows a newer mutable PR merge ref. An -unavailable event merge, or one that does not contain exactly one snapshotted -head parent, fails closed. A later PR head change does not retarget an already -queued run: it may execute only its previously trusted event revision if that -revision remains fetchable. To run the new revision, the maintainer must review -it, remove the label, and re-add the desired regression label; cancel the old -run if it is no longer wanted. An existing label does not automatically rerun -the benchmark. +**trust admission for that exact PR revision**. `query-regression` runs the +six routine default cases; `heavy-regression` runs only the high-cardinality +`prom_remote_write_7913` remote-write case. `pull_request: labeled` is the only PR +trigger: the label event snapshots its merge, head, and base SHAs. A queued +job fetches that immutable event merge SHA directly, verifies it is a +two-parent merge whose parents include the snapshotted head exactly once, and +uses its other parent as the actual base build revision. The snapshotted +event base is retained for audit only, so a difference from the merge's +non-head parent is not a failure. The job never follows a newer mutable PR +merge ref. An unavailable event merge, or one that does not contain exactly +one snapshotted head parent, fails closed. A later PR head change does not +retarget an already queued run: it may execute only its previously trusted +event revision if that revision remains fetchable. To run the new revision, +the maintainer must review it, remove the label, and re-add the desired +regression label; cancel the old run if it is no longer wanted. An existing +label does not automatically rerun the benchmark. -Admission does not relax runner hardening or GitHub permissions. Keep -service-account token mounting disabled; do not mount host paths, the Docker -socket, kubeconfig, or long-lived credentials. The runner and cache initializer -use UID/GID 1001, disallow privilege escalation, drop all capabilities, and use -the RuntimeDefault seccomp profile. Keep GitHub tokens least-privilege and -review workflow changes before admission. Where the CNI supports it, restrict -egress to required GitHub Actions, artifact/cache, Rust/crate/toolchain, DNS, -and image-registry endpoints; block unrelated cluster services, private ranges, -and metadata endpoints unless a case requires them. - -### Network routing prerequisite - -Required split routing is an **external environment-specific prerequisite**. The -responsible network operator must route GitHub Actions, GitHub content, -artifact/cache, crates.io, Rust toolchain, and image-registry traffic through -the approved path rather than the VPN where applicable. Neither this repository -nor Kubernetes configures that route. Verify it with the responsible network -operator before any canary. +Admission does not relax runner hardening or GitHub permissions. Keep the ECS +instance free of cloud credentials and long-lived tokens, keep the security +group egress-only, keep GitHub tokens least-privilege, and review workflow +changes before admission. ## Runner image and workflow tools -Build and push the derived runner image; it preserves the official -`/home/runner/run.sh` entrypoint and supplies CI tools needed at runtime. The -image builds `otelgen` from +The runner `Dockerfile` builds `otelgen` from [`WenyXu/otelgen`](https://github.com/WenyXu/otelgen) commit [`863a3f395d062c7322cc1de08a38774b7fdaa6c8`](https://github.com/WenyXu/otelgen/commit/863a3f395d062c7322cc1de08a38774b7fdaa6c8) -so trace cases do not download or compile tools during a benchmark run: +so trace cases do not download or compile tools during a benchmark run. It is +built only as a toolchain factory: `build-ecs-image.py` and +`bootstrap-runner-host.sh` both materialize its contents onto a host, and the +benchmark itself runs host-native. -```bash -docker build \ - --platform linux/amd64 \ - -f .github/runner-scale-sets/query-regression/Dockerfile \ - -t greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest \ - .github/runner-scale-sets/query-regression - -docker push greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest -``` - -Deploy by digest, not mutable tag, by updating both image references in -`values-8-cores.yaml` after a rebuild. Update `RUNNER_IMAGE_DIGEST` and bump -`RUNNER_IMAGE_EPOCH` in `query-regression.yml` at the same time. If the registry -is private, use a dedicated read-only pull secret only as `imagePullSecrets`; -never expose registry credentials to runner containers. -Both digest-pinned init and runner containers use `IfNotPresent`: the immutable -digest makes a cached image safe and avoids adding a registry dependency to every -runner startup. - -The runner optionally imports only the non-sensitive `HTTP_PROXY`, -`HTTPS_PROXY`, and `NO_PROXY` variables from the -`query-regression-runner-local-env` ConfigMap. Manage that ConfigMap locally in -the target namespace; private endpoint configuration must not be committed, and -credentials or secrets must never be placed in a ConfigMap. - -Before builds, the workflow asserts UID/GID 1001 and exact image tool versions: -`libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.30.0`, `Python 3.12.3`, `sccache -0.16.0`, `otelgen` commit `863a3f395d062c7322cc1de08a38774b7fdaa6c8`, -root-owned `rustup 1.29.0`, and the image-baked -`nightly-2026-03-21` Rust toolchain. Rustup, Cargo, and Rustc must resolve from -`/opt/cargo/bin`; the runner cannot write `/opt/rustup` or `/opt/cargo/bin`. -Protobuf well-known includes, including `google/protobuf/any.proto` and -`google/protobuf/empty.proto`, are an image contract and must compile with -`protoc`. +Before builds, the workflow asserts the runner UID/GID (1001 in the ECS image; +overridable via `QUERY_REGRESSION_RUNNER_UID`/`QUERY_REGRESSION_RUNNER_GID`) +and exact tool versions: `libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.40.4`, +`Python 3.14.4`, `sccache 0.16.0`, `otelgen` commit +`863a3f395d062c7322cc1de08a38774b7fdaa6c8`, root-owned `rustup 1.29.0`, and +the image-baked `nightly-2026-03-21` Rust toolchain. `mold` and `python3` +come from apt at image-build time (not Ubuntu 24.04's default 3.12); bump +the Verify pins together with `QUERY_REGRESSION_ECS_IMAGE_ID` when the +image is rebuilt. Rustup, Cargo, and Rustc +must resolve from `/opt/cargo/bin`; the runner cannot write `/opt/rustup` or +`/opt/cargo/bin`. Protobuf well-known includes, including +`google/protobuf/any.proto` and `google/protobuf/empty.proto`, are an image +contract and must compile with `protoc`. `actions-rust-lang/setup-rust-toolchain@v1` is intentionally removed. The -workflow sets its warning-denying mold `RUSTFLAGS` directly, disables automatic -Rustup installation, and performs no runtime toolchain downloads. +workflow sets its warning-denying mold `RUSTFLAGS` directly, disables +automatic Rustup installation, and performs no runtime toolchain downloads. -The workflow no longer uses GitHub `rust-cache`, `setup-protoc`, `setup-uv`, or -runtime Rust setup: the image establishes immutable executable state and the PVC -supplies only reusable Cargo data. Do not reintroduce those actions unless the -corresponding cache or image contract changes. +The workflow no longer uses GitHub `rust-cache`, `setup-protoc`, `setup-uv`, +or runtime Rust setup: the image establishes immutable executable state and +each instance's system disk holds only that run's Cargo data. Do not +reintroduce those actions unless the image contract changes. -## Capacity and persistent cache +## Capacity -`values-8-cores.yaml` is normal operation: `minRunners=0`, `maxRunners=1`. -`values-paused.yaml` is the mandatory pause overlay: `minRunners=0`, -`maxRunners=0`. The job uses group `query-regression-persistent-cache-v1`, -`queue: max`, and `cancel-in-progress: false`; admitted jobs queue rather than -replacing older pending jobs. During maintenance, cancel admitted queued runs as -well as pausing ARC. Runner Pods have `activeDeadlineSeconds=12600`. -The runner requests 6 CPU and limits at 8 CPU to preserve `minipc-3` -allocatable-capacity scheduling headroom; do not reset the request to 8 CPU -without revalidating scheduling capacity. +Each run provisions its own ECS instance, so overlapping dispatches proceed +in parallel. There is no workflow `concurrency` group. All Cargo state +(`CARGO_HOME` including registry/git, `CARGO_TARGET_DIR`, sccache, cache +metadata) lives on the 40 GiB system disk and is discarded with the VM. +`RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are image-owned. The runner +sets `RUSTC_WRAPPER=/usr/local/bin/sccache`, +`SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=10G`, and +`CARGO_INCREMENTAL=0`. sccache uses its local disk backend and self-evicts +at 10G. Base and candidate builds share the target on that disk; Cargo +fingerprints invalidate source and dependency changes. -The cache claim `query-regression-build-cache` is a nominal 600Gi `local-path` -PVC in `arc-runners`. It is `ReadWriteOnce`; `local-path` uses -WaitForFirstConsumer binding and Delete reclaim behavior, produces a -node-affine local PV, is non-expandable, and the 600Gi request is not a hard -storage quota. The runner's `minipc-3` selector is its only consumer candidate. - -The initializer mounts the PVC root at `/cache`, creates and write-tests these -versioned subpaths as non-root UID/GID 1001, and the runner mounts them as: - -| Persistent state | PVC subpath | Runner mount | -| --- | --- | --- | -| Ephemeral Cargo home | `emptyDir` | `/home/runner/.cargo` | -| Cargo registry data | `cargo-registry-v1` | `/home/runner/.cargo/registry` | -| Cargo Git data | `cargo-git-v1` | `/home/runner/.cargo/git` | -| Cargo target | `query-regression-target-v1` | `/home/runner/query-regression-target` | -| Cache metadata | `meta-v1` | `/home/runner/query-regression-cache-meta` | -| sccache local disk cache | `sccache-v1` | `/home/runner/.cache/sccache` | -| Immutable Rust toolchain | image-owned | `/opt/rustup`, `/opt/cargo/bin` | - -The Pod security context uses UID/GID and `fsGroup` 1001 with -`fsGroupChangePolicy: OnRootMismatch`; no privileged `chown` or raw `hostPath` -is used. `CARGO_HOME` is a per-Pod `emptyDir`; only its nested `registry` and -`git` mounts are persistent. `RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are -image-owned immutable paths, while `CARGO_TARGET_DIR`, cache metadata, and -`SCCACHE_DIR` are persistent absolute paths. The runner sets -`RUSTC_WRAPPER=/usr/local/bin/sccache`, -`SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=40G`, and -`CARGO_INCREMENTAL=0`. sccache uses its local PVC disk backend and self-evicts -at 40G; do not add runtime downloads, object storage, or a shared backend. - -The repository's `.cargo/config.toml` remains a trusted per-revision build input. -In contrast, `$CARGO_HOME/config*`, credentials, installed bins, and Cargo -metadata outside the persistent `registry` and `git` data mounts are ephemeral -and cannot survive to another Pod. - -The local disk backend has a one-server constraint. `maxRunners=1` and the -unchanged `query-regression-persistent-cache-v1` workflow concurrency group -serialize runs; do not increase runner capacity or relax that serialization -while this backend is in use. Base and candidate builds share the target; Cargo -fingerprints invalidate source and dependency changes. The workflow records the -sccache version and relevant environment in the target ABI marker, starts and -zeros sccache after cache and toolchain checks, shows initial/base/candidate -statistics, and resets statistics between base and candidate builds. - -### Disk preflight and cleanup contract - -Before applying or unpausing, verify the backing filesystem on `minipc-3` has -at least 900GiB free. The current local-path provisioner source is -`/opt/local-path-provisioner`; measure the filesystem containing it: - -```bash -df -PB1G /opt/local-path-provisioner -``` - -The workflow reports `du`, `df -P`, human-readable free space, and inode -availability before builds and in an always-run report. Its cleanup is narrow -and non-destructive: - -- warn at target size 400GiB; at 450GiB clear only the complete target root; -- warn at Cargo registry-plus-Git data size 60GiB; at 80GiB remove only - `registry/src` and `git/checkouts`, then abort if that persistent data remains - at least 80GiB; -- below 300GiB backing free space, clear the complete target root first, - remeasure, then remove only those Cargo extracted trees and checkouts; abort - if free space is still below 300GiB; -- never automatically remove Cargo registry cache/index, Git database, the - image-owned Cargo bin or Rustup toolchain, cache metadata, the self-evicting - sccache directory, or the PVC. - -The target clear uses fixed absolute roots and removes all entries, including -dotfiles. After migration validation, remove obsolete `cargo-home-v1` and -`rustup-home-v1` only in explicit maintenance while ARC is 0/0 and no runner Pod -exists; they are not mounted by the current configuration. - -## Deploy and pause safely - -First verify the configured external network route and the disk preflight. Apply the PVC; while -the scale set is paused, it is expected to remain `Pending` because -WaitForFirstConsumer has no scheduled runner: - -```bash -kubectl apply --dry-run=server \ - -f .github/runner-scale-sets/query-regression/cache-pvc.yaml -kubectl apply -f .github/runner-scale-sets/query-regression/cache-pvc.yaml -``` - -Render normal and paused configurations. Normal values are always first; the -pause overlay is always last: - -```bash -helm template perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --version 0.14.2 \ - --set controllerServiceAccount.name=arc-gha-rs-controller \ - --set controllerServiceAccount.namespace=arc-systems \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml - -helm template perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --version 0.14.2 \ - --set controllerServiceAccount.name=arc-gha-rs-controller \ - --set controllerServiceAccount.namespace=arc-systems \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ - -f .github/runner-scale-sets/query-regression/values-paused.yaml -``` - -The **first post-merge Helm deployment must reconcile the release in paused -mode**. Keep the pause overlay last: - -```bash -# First post-merge deployment and every return to paused mode: 0/0. -helm upgrade --install perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --create-namespace --version 0.14.2 \ - --reset-values --wait \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ - -f .github/runner-scale-sets/query-regression/values-paused.yaml - -# Expect 0/0 and no runner resources before considering normal mode. -kubectl -n arc-runners get autoscalingrunnerset perf-regression-8-cores \ - -o jsonpath='{.spec.minRunners}{"/"}{.spec.maxRunners}{"\n"}' -kubectl -n arc-runners get ephemeralrunners,pods \ - -l actions.github.com/scale-set-name=perf-regression-8-cores -``` - -Only after that verification and separate explicit approval, apply normal 0/1 -operation without the pause overlay: - -```bash -helm upgrade --install perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --create-namespace --version 0.14.2 \ - --reset-values --wait \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml -``` - -Do not use bare `helm rollback`, `--atomic`, or `--reuse-values`: a stored -revision can restore nonzero runner capacity. Inspect rendered manifests for -capacity, the `minipc-3` selector, cache claim and mounts, initializer, -security context, and resources. After approved normal mode receives its first -canary, the PVC binds to `minipc-3`. - -For local-PV node loss, cache recovery is intentionally disposable: return to -0/0, recreate the PVC on a healthy node, cold-fill it, and run a new canary. - -## Canary and rollback - -With explicit approval, run two identical `workflow_dispatch` canaries on -`perf-regression-8-cores`, using immutable full base and candidate commit SHAs -and `cargo_profile=nightly`. The first is the cold fill; the second verifies warm -reuse. Record the workflow's base/candidate build elapsed logs and cache -ABI-marker output, initial/base/candidate sccache statistics, and cache report. -Confirm the image tool contract (root-owned Rustup/Cargo paths, baked nightly, -and non-writable `/opt` roots) and that the ephemeral Cargo home contains only -the mounted registry/Git data before Cargo creates per-Pod state. -Obtain dependency and tool network byte counters from the configured -environment counter source, filtered to `minipc-3` and the dependency/tool -destinations. - -Accept the canary only when all of the following hold: - -- exactly one runner Pod runs on `minipc-3`, and both jobs use the same bound PV; -- UID/GID 1001 cache mounts are writable; the warm run does not invalidate the - target or bulk-redownload crates or toolchains; sccache reports separate base - and candidate build statistics without server or cache-path errors; immutable - Rustup/Cargo roots remain non-writable and only registry/Git data persists; -- warm base build time is at most 50% of cold base build time; -- warm dependency/tool network bytes are at most 10% of cold fill bytes; -- cache sizes remain below soft watermarks, node free space remains at least - 300GiB, and the benchmark is correct without TLS EOFs or timeouts; -- the configured environment counter source confirms this traffic is outside VPN - accounting. - -Immediately return to 0/0 after either canary unless ongoing normal operation -has been explicitly approved; return immediately on any traffic, cache, disk, -TLS, or correctness failure. To roll back, use the paused Helm upgrade above, -or another explicit `helm upgrade --install` with known-good values followed by -`values-paused.yaml`, `--reset-values`, and `--wait`. Do not delete the PVC -automatically; preserve it for diagnosis unless intentionally discarding cache. +The workflow reports `du` / `df` in telemetry. It does not try to reclaim +space across runs: a cold 40 GiB disk that fills up fails the build. ## Future optional phases -The current phase uses a digest-pinned image with sccache 0.16.0 and no shared -cache service. Optional follow-ups are an image additionally seeded with the -exact Rust toolchain and `cargo fetch --locked`; or an internal read/write -sccache backend or Cargo/Git mirror. Evaluate them only if persistent PVC reuse -is insufficient. +The current phase uses a materialized runner toolchain with sccache 0.16.0 and +no shared cache service. Optional follow-ups are an image additionally seeded +with `cargo fetch --locked` results, or an internal read/write sccache +backend or Cargo/Git mirror. Evaluate them only if cold compile time on the +system disk becomes the bottleneck. diff --git a/.github/runner-scale-sets/query-regression/cache-pvc.yaml b/.github/runner-scale-sets/query-regression/cache-pvc.yaml deleted file mode 100644 index 692f3a2d6e2..00000000000 --- a/.github/runner-scale-sets/query-regression/cache-pvc.yaml +++ /dev/null @@ -1,15 +0,0 @@ -apiVersion: v1 -kind: PersistentVolumeClaim -metadata: - name: query-regression-build-cache - namespace: arc-runners - labels: - app.kubernetes.io/part-of: perf-regression-8-cores -spec: - volumeMode: Filesystem - accessModes: - - ReadWriteOnce - storageClassName: local-path - resources: - requests: - storage: 600Gi diff --git a/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh b/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh new file mode 100644 index 00000000000..883bbede092 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh @@ -0,0 +1,195 @@ +#!/usr/bin/env bash +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Bootstrap an Ubuntu 24.04 host as a query-regression self-hosted runner. +# +# This is the host-native equivalent of the ECS image contract: it +# installs docker-ce from Docker's official repository, builds the runner +# image from the sibling Dockerfile (the single source of the tool contract), +# materializes the toolchain onto the host, and optionally registers the +# runner and installs its systemd service. +# +# Large state lives under --data-root (a data disk): docker's data-root, the +# runner home, and the rustup/cargo toolchain roots, symlinked back to the +# hard-coded contract paths (/home/runner, /opt/rustup, /opt/cargo). +# +# The script is idempotent: re-running refreshes the toolchain in place, and +# the runner registration (.runner/.credentials) and _work survive. A live +# runner service is stopped first and restarted at the end. +# +# Usage: +# sudo bash bootstrap-runner-host.sh --repo-dir /path/to/greptimedb +# sudo RUNNER_TOKEN= bash bootstrap-runner-host.sh --repo-dir . \ +# --register --repo-url https://github.com// +# +# Options: +# --repo-dir PATH greptimedb checkout containing the runner Dockerfile (required) +# --data-root PATH data disk mount point (default: /data) +# --uid / --gid N runner user id (default: 3141; the ECS image contract is 1001) +# --register also register the runner and install the systemd service +# --repo-url URL repository URL for registration (required with --register) +# --runner-name N runner name (default: qreg-host) +# --labels L runner labels (default: perf-regression-8-cores) +# +# With --register, provide a fresh registration token via RUNNER_TOKEN. + +set -euo pipefail + +REPO_DIR="" +DATA_ROOT="/data" +RUNNER_UID="3141" +RUNNER_GID="3141" +REGISTER="false" +REPO_URL="" +RUNNER_NAME="qreg-host" +RUNNER_LABELS="perf-regression-8-cores" +IMAGE_TAG="qreg-runner:manual" + +while [[ $# -gt 0 ]]; do + case "$1" in + --repo-dir) REPO_DIR="$2"; shift 2 ;; + --data-root) DATA_ROOT="$2"; shift 2 ;; + --uid) RUNNER_UID="$2"; shift 2 ;; + --gid) RUNNER_GID="$2"; shift 2 ;; + --register) REGISTER="true"; shift ;; + --repo-url) REPO_URL="$2"; shift 2 ;; + --runner-name) RUNNER_NAME="$2"; shift 2 ;; + --labels) RUNNER_LABELS="$2"; shift 2 ;; + -h | --help) sed -n '17,40p' "$0"; exit 0 ;; + *) echo "Unknown argument: $1" >&2; exit 2 ;; + esac +done + +[[ $EUID -eq 0 ]] || { echo "Run as root, e.g. sudo bash $0 ..." >&2; exit 1; } +[[ -n "${REPO_DIR}" ]] || { echo "--repo-dir is required" >&2; exit 2; } +DOCKERFILE="${REPO_DIR}/.github/runner-scale-sets/query-regression/Dockerfile" +[[ -f "${DOCKERFILE}" ]] || { echo "Dockerfile not found at ${DOCKERFILE}" >&2; exit 1; } +if [[ "${REGISTER}" == "true" ]]; then + [[ -n "${REPO_URL}" ]] || { echo "--repo-url is required with --register" >&2; exit 2; } + [[ -n "${RUNNER_TOKEN:-}" ]] || { echo "RUNNER_TOKEN env is required with --register" >&2; exit 2; } +fi + +step() { printf '\n==> %s\n' "$*"; } + +# A re-run refreshes files under the (possibly live) runner service, so stop +# it first and restart it at the end (unless --register re-creates it). +RUNNER_SERVICE="$(systemctl list-units --type=service --all --no-legend 'actions.runner.*' 2>/dev/null | awk '{print $1}' | head -n 1 || true)" +RUNNER_WAS_ACTIVE="false" +if [[ -n "${RUNNER_SERVICE}" ]] && systemctl is-active --quiet "${RUNNER_SERVICE}"; then + step "Stop running runner service ${RUNNER_SERVICE}" + systemctl stop "${RUNNER_SERVICE}" + RUNNER_WAS_ACTIVE="true" +fi + +step "Install docker-ce from Docker's official repository" +for pkg in docker.io docker-doc docker-compose docker-compose-v2 podman-docker containerd runc; do + apt-get remove -y "${pkg}" 2>/dev/null || true +done +apt-get update +apt-get install -y ca-certificates curl +install -m 0755 -d /etc/apt/keyrings +curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc +chmod a+r /etc/apt/keyrings/docker.asc +# shellcheck disable=SC1091 +. /etc/os-release +tee /etc/apt/sources.list.d/docker.sources > /dev/null <&2 + exit 1 +fi +echo "{ \"data-root\": \"${DATA_ROOT}/docker\" }" > /etc/docker/daemon.json +# containerd keeps its own content store (where build-time image data lands); +# move it off the system disk too. +mkdir -p /etc/containerd +containerd config default | sed "s|root = \"/var/lib/containerd\"|root = \"${DATA_ROOT}/containerd\"|" > /etc/containerd/config.toml +systemctl start containerd docker +docker info | grep "Docker Root Dir" + +step "Build the runner image" +docker build --platform linux/amd64 -f "${DOCKERFILE}" -t "${IMAGE_TAG}" "${REPO_DIR}" + +step "Materialize the toolchain onto the host" +mkdir -p "${DATA_ROOT}/opt" "${DATA_ROOT}/runner" /opt /home +container="$(docker create "${IMAGE_TAG}")" +trap 'docker rm -f "${container}" >/dev/null 2>&1 || true' EXIT +# The `/.` suffix copies directory *contents*, so re-running over an existing +# materialization refreshes it in place instead of nesting runner/runner. +# Files absent from the image (runner registration, _work) are preserved. +docker cp "${container}:/home/runner/." "${DATA_ROOT}/runner" +docker cp "${container}:/opt/rustup/." "${DATA_ROOT}/opt/rustup" +docker cp "${container}:/opt/cargo/." "${DATA_ROOT}/opt/cargo" +for tool in uv uvx otelgen sccache; do + docker cp "${container}:/usr/local/bin/${tool}" "/usr/local/bin/${tool}" +done +docker rm "${container}" > /dev/null +trap - EXIT +ln -sfn "${DATA_ROOT}/runner" /home/runner +ln -sfn "${DATA_ROOT}/opt/rustup" /opt/rustup +ln -sfn "${DATA_ROOT}/opt/cargo" /opt/cargo +docker image rm "${IMAGE_TAG}" > /dev/null + +step "Install system packages" +apt-get install -y --no-install-recommends \ + build-essential clang cmake git gzip jq libprotobuf-dev libssl-dev mold \ + openssh-client pkg-config protobuf-compiler python3 sudo tar unzip wget xz-utils zip zstd + +step "Create runner user (${RUNNER_UID}:${RUNNER_GID}) and environment" +getent group "${RUNNER_GID}" > /dev/null || groupadd -g "${RUNNER_GID}" runner +id -u runner > /dev/null 2>&1 || useradd -u "${RUNNER_UID}" -g "${RUNNER_GID}" -d /home/runner -s /bin/bash runner +chown -R "${RUNNER_UID}:${RUNNER_GID}" "${DATA_ROOT}/runner" +echo 'PATH=/opt/cargo/bin:/usr/local/bin:/usr/bin:/bin' > "${DATA_ROOT}/runner/.env" +chown "${RUNNER_UID}:${RUNNER_GID}" "${DATA_ROOT}/runner/.env" + +if [[ "${REGISTER}" == "true" ]]; then + step "Register runner ${RUNNER_NAME} (labels: ${RUNNER_LABELS})" + if [[ -n "${RUNNER_SERVICE}" ]]; then + (cd /home/runner && ./svc.sh uninstall) || true + fi + runuser -u runner -- bash -c "cd /home/runner && HOME=/home/runner ./config.sh \ + --url '${REPO_URL}' --token '${RUNNER_TOKEN}' --name '${RUNNER_NAME}' \ + --labels '${RUNNER_LABELS}' --unattended --replace --disableupdate" + + step "Install and start the runner service" + cd /home/runner + ./svc.sh install runner + ./svc.sh start + ./svc.sh status +else + step "Skipping registration (pass --register --repo-url ... with RUNNER_TOKEN to enable)" + if [[ "${RUNNER_WAS_ACTIVE}" == "true" ]]; then + step "Restart runner service ${RUNNER_SERVICE}" + systemctl start "${RUNNER_SERVICE}" + systemctl --no-pager status "${RUNNER_SERVICE}" || true + fi +fi + +step "Done" +echo "Runner home: /home/runner -> ${DATA_ROOT}/runner" +echo "Before each workflow run on this persistent host, clean transient cargo state:" +echo " sudo rm -f /home/runner/.cargo/.package-cache" +echo " sudo rm -rf /home/runner/.cargo/.global-cache" diff --git a/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py new file mode 100644 index 00000000000..7031917f910 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# PEP 723 inline metadata (see .github/scripts/aliyun-ecs-runner-provision.py +# for the convention). +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Build the query-regression ECS custom image (manual ops tool). + +Boots a temporary pay-as-you-go ECS instance from a public Ubuntu 24.04 image, +builds the existing runner container image (the Dockerfile in the parent +directory remains the single source of the tool contract), materializes the +tool directories onto the host filesystem so the workflow's "Verify runner +image tools" step holds unchanged, installs the ephemeral-runner systemd unit, +and snapshots the result as a custom image. The temporary instance is deleted +afterwards. + +Sentinel polling reads the instance's serial console output +(GetInstanceConsoleOutput) and looks for marker lines the user data writes to +/dev/console. This has no in-guest agent dependency. + +Usage: + uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ + --region-id cn-hangzhou --vswitch-id vsw-... --security-group-id sg-... \ + --base-image-id ubuntu_24_04_x64_20G_alibase_*.vhd +""" + +from __future__ import annotations + +import argparse +import base64 +import os +import time +from pathlib import Path + +ASSETS_DIR = Path(__file__).resolve().parent +DONE_MARKER = "QREG_IMAGE_BUILD_DONE" +FAILED_MARKER = "QREG_IMAGE_BUILD_FAILED" +POLL_INTERVAL_SECONDS = 15 +CONSOLE_POLL_INTERVAL_SECONDS = 30 +BUILD_TIMEOUT_SECONDS = 60 * 60 + +# Same apt package contract as the runner Dockerfile; the base +# actions-runner image is Ubuntu 24.04, so an Ubuntu 24.04 host resolves the +# same tool versions (protoc 3.21.12, mold 2.40.4, Python 3.14.4). +# Docker itself comes from Docker's official repository (docker-ce), not the +# distribution-packaged docker.io. +APT_PACKAGES = [ + "build-essential", + "ca-certificates", + "clang", + "cmake", + "curl", + "git", + "gpg", + "gzip", + "jq", + "libprotobuf-dev", + "libssl-dev", + "mold", + "openssh-client", + "pkg-config", + "protobuf-compiler", + "python3", + "sudo", + "tar", + "unzip", + "wget", + "xz-utils", + "zip", + "zstd", +] + +DOCKER_CE_PACKAGES = "docker-ce docker-ce-cli containerd.io docker-buildx-plugin" + + +def render_user_data(dockerfile: str, start_runner: str, unit: str) -> str: + dockerfile_b64 = base64.b64encode(dockerfile.encode()).decode() + start_runner_b64 = base64.b64encode(start_runner.encode()).decode() + unit_b64 = base64.b64encode(unit.encode()).decode() + packages = " ".join(APT_PACKAGES) + return f"""#!/bin/bash +set -euo pipefail +trap 'echo "{FAILED_MARKER} at line $LINENO" > /dev/console' ERR +# Stream the full build log to the serial console so the poller (and anyone +# watching GetInstanceConsoleOutput) sees real progress, not a silent login +# prompt for the whole build. +exec > >(tee -a /dev/console) 2>&1 + +apt-get update +DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends {packages} + +# Docker from the official repository, not the distribution-packaged docker.io. +install -m 0755 -d /etc/apt/keyrings +curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg +chmod a+r /etc/apt/keyrings/docker.gpg +. /etc/os-release +echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu ${{VERSION_CODENAME}} stable" \ + > /etc/apt/sources.list.d/docker.list +apt-get update +DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends {DOCKER_CE_PACKAGES} + +base64 -d > /tmp/Dockerfile <<'EOF' +{dockerfile_b64} +EOF +mkdir -p /tmp/image-context +docker build --platform linux/amd64 -f /tmp/Dockerfile -t qreg-runner:local /tmp/image-context + +# Materialize the tool contract onto the host filesystem. +container="$(docker create qreg-runner:local)" +trap 'docker rm -f "${{container}}" >/dev/null 2>&1 || true' EXIT +docker cp "${{container}}:/opt/rustup" /opt/rustup +docker cp "${{container}}:/opt/cargo" /opt/cargo +for tool in uv uvx otelgen sccache; do + docker cp "${{container}}:/usr/local/bin/${{tool}}" "/usr/local/bin/${{tool}}" +done +docker cp "${{container}}:/home/runner" /home/runner + +# Runner identity mirrors the container image (UID/GID 1001). +groupadd --gid 1001 runner 2>/dev/null || true +useradd --uid 1001 --gid 1001 --home-dir /home/runner --shell /bin/bash runner 2>/dev/null || true +chown -R 1001:1001 /home/runner + +mkdir -p /opt/ephemeral-github-runner +base64 -d > /opt/ephemeral-github-runner/start-runner.sh <<'EOF' +{start_runner_b64} +EOF +chmod 0755 /opt/ephemeral-github-runner/start-runner.sh +base64 -d > /etc/systemd/system/ephemeral-github-runner.service <<'EOF' +{unit_b64} +EOF +systemctl daemon-reload +systemctl enable ephemeral-github-runner.service + +# Keep the image free of the build-time docker state (also shrinks the +# snapshot: buildkit cache is several GB). +docker rm -f "${{container}}" >/dev/null +docker system prune -af >/dev/null +trap - EXIT + +echo "{DONE_MARKER}" > /dev/console +""" + + +def make_ecs_client(region_id: str): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + return EcsClient( + OpenApiConfig( + access_key_id=os.environ["ALIBABA_CLOUD_ACCESS_KEY_ID"], + access_key_secret=os.environ["ALIBABA_CLOUD_ACCESS_KEY_SECRET"], + region_id=region_id, + endpoint=f"ecs.{region_id}.aliyuncs.com", + ) + ) + + +def wait_for_instance_status(client, region_id: str, instance_id: str, wanted: str, deadline: float) -> None: + from alibabacloud_ecs20140526 import models as ecs_models + + import json + + while time.monotonic() < deadline: + response = call_api_with_retry( + lambda: client.describe_instances( + ecs_models.DescribeInstancesRequest( + region_id=region_id, instance_ids=json.dumps([instance_id]) + ) + ), + "DescribeInstances", + ) + instances = response.body.instances.instance + if instances and instances[0].status == wanted: + return + time.sleep(POLL_INTERVAL_SECONDS) + raise TimeoutError(f"Instance {instance_id} did not reach status {wanted} in time") + + +# Aliyun error codes worth retrying: throttling and server-side faults. Client +# errors (4xx: permissions, bad parameters) are configuration problems and must +# fail fast instead of being retried. +TRANSIENT_ERROR_CODES = {"Throttling", "Throttling.User", "InternalError", "ServiceUnavailable"} + + +def is_transient(error: Exception) -> bool: + # UnretryableException comes from the darabonba network layer: connection + # reset, read timeout, DNS blip. + if type(error).__name__ == "UnretryableException": + return True + code = getattr(error, "code", "") or "" + status = getattr(error, "statusCode", None) + return code in TRANSIENT_ERROR_CODES or (isinstance(status, int) and status >= 500) + + +def call_api_with_retry(fn, description: str, attempts: int = 5): + """Retry transient API/network failures; the build is too long to die on a blip. + + Only call this with idempotent or safely-repeatable requests (reads, + StopInstance, CreateImage). Never with RunInstances: if the request + succeeded but the response was lost, a retry double-creates instances. + """ + for attempt in range(1, attempts + 1): + try: + return fn() + except Exception as error: # noqa: BLE001 + if attempt == attempts or not is_transient(error): + raise + print(f"{description} failed (attempt {attempt}/{attempts}): {error}", flush=True) + time.sleep(POLL_INTERVAL_SECONDS) + raise RuntimeError("unreachable: retry loop exited without returning") + + +def read_console_output(client, region_id: str, instance_id: str) -> str: + """Fetch the instance's serial console output; no in-guest agent needed.""" + from alibabacloud_ecs20140526 import models as ecs_models + + response = call_api_with_retry( + lambda: client.get_instance_console_output( + ecs_models.GetInstanceConsoleOutputRequest(region_id=region_id, instance_id=instance_id) + ), + "GetInstanceConsoleOutput", + ) + return base64.b64decode(response.body.console_output or "").decode("utf-8", "replace") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--vswitch-id", default=os.environ.get("ALIYUN_ECS_VSWITCH_ID")) + parser.add_argument("--security-group-id", default=os.environ.get("ALIYUN_ECS_SECURITY_GROUP_ID")) + parser.add_argument("--base-image-id", default=os.environ.get("ALIYUN_ECS_BASE_IMAGE_ID")) + parser.add_argument("--resource-group-id", default=os.environ.get("ALIYUN_ECS_RESOURCE_GROUP_ID")) + parser.add_argument("--instance-type", default="ecs.g7.xlarge") + parser.add_argument("--image-name", default=None, help="Defaults to a timestamped name.") + args = parser.parse_args() + + for name in ("region_id", "vswitch_id", "security_group_id", "base_image_id"): + if not getattr(args, name): + raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}") + + from alibabacloud_ecs20140526 import models as ecs_models + + client = make_ecs_client(args.region_id) + image_name = args.image_name or time.strftime( + "greptimedb-query-regression-runner-%Y%m%d%H%M%S", time.gmtime() + ) + + user_data = base64.b64encode( + render_user_data( + (ASSETS_DIR.parent / "Dockerfile").read_text(), + (ASSETS_DIR / "start-runner.sh").read_text(), + (ASSETS_DIR / "ephemeral-github-runner.service").read_text(), + ).encode() + ).decode() + + instance_id = None + try: + response = client.run_instances( + ecs_models.RunInstancesRequest( + region_id=args.region_id, + image_id=args.base_image_id, + resource_group_id=args.resource_group_id, + instance_type=args.instance_type, + v_switch_id=args.vswitch_id, + security_group_id=args.security_group_id, + instance_name=f"build-{image_name}", + description="Temporary builder for the query-regression ECS image", + amount=1, + instance_charge_type="PostPaid", + spot_strategy="NoSpot", + internet_charge_type="PayByTraffic", + internet_max_bandwidth_out=100, + # Peak usage is ~18G (OS+apt, docker image, and the materialized + # toolchain coexist briefly): deliberately tight, and a smaller + # disk makes the image snapshot faster. Instances created from + # the image get a larger system disk from the provision side. + system_disk=ecs_models.RunInstancesRequestSystemDisk( + category="cloud_essd", size="20" + ), + user_data=user_data, + tag=[ + ecs_models.RunInstancesRequestTag(key="managed-by", value="query-regression-ci"), + ecs_models.RunInstancesRequestTag(key="role", value="image-builder"), + ], + ) + ) + instance_id = response.body.instance_id_sets.instance_id_set[0] + print(f"Builder instance: {instance_id}", flush=True) + wait_for_instance_status(client, args.region_id, instance_id, "Running", time.monotonic() + 10 * 60) + + deadline = time.monotonic() + BUILD_TIMEOUT_SECONDS + while time.monotonic() < deadline: + console = read_console_output(client, args.region_id, instance_id) + if DONE_MARKER in console: + break + if FAILED_MARKER in console: + tail = "\n".join(console.splitlines()[-20:]) + raise RuntimeError(f"Image build failed on the builder; console tail:\n{tail}") + lines = console.splitlines() + print(f"Build in progress; last console line: {lines[-1] if lines else '(none yet)'}", flush=True) + time.sleep(CONSOLE_POLL_INTERVAL_SECONDS) + else: + raise TimeoutError("Image build did not finish in time") + + print("Stopping builder before image creation", flush=True) + try: + call_api_with_retry( + lambda: client.stop_instance( + ecs_models.StopInstanceRequest( + instance_id=instance_id, + # The instance is deleted right after the snapshot, but + # there is no reason to keep billing vCPU during the stop. + stopped_mode="StopCharging", + ) + ), + "StopInstance", + ) + except Exception as error: # noqa: BLE001 + # Instance families with local disks do not support StopCharging. + print(f"StopCharging unavailable ({error}); stopping with default mode", flush=True) + call_api_with_retry( + lambda: client.stop_instance(ecs_models.StopInstanceRequest(instance_id=instance_id)), + "StopInstance", + ) + wait_for_instance_status(client, args.region_id, instance_id, "Stopped", time.monotonic() + 10 * 60) + + image = call_api_with_retry( + lambda: client.create_image( + ecs_models.CreateImageRequest( + region_id=args.region_id, instance_id=instance_id, image_name=image_name + ) + ), + "CreateImage", + ) + image_id = image.body.image_id + deadline = time.monotonic() + 60 * 60 + while time.monotonic() < deadline: + description = call_api_with_retry( + lambda: client.describe_images( + ecs_models.DescribeImagesRequest(region_id=args.region_id, image_id=image_id) + ), + "DescribeImages", + ) + images = description.body.images.image + if images: + status = images[0].status + print(f"Image {image_id} status: {status} (progress {images[0].progress})", flush=True) + if status == "Available": + break + time.sleep(POLL_INTERVAL_SECONDS) + else: + raise TimeoutError( + f"Image {image_id} did not become Available in time. The snapshot " + f"continues server-side: check its status in the console and reuse it " + f"once Available instead of rebuilding." + ) + + print(f"Custom image ready: {image_id} ({image_name})", flush=True) + print(f"Set the repo variable QUERY_REGRESSION_ECS_IMAGE_ID={image_id}", flush=True) + return 0 + finally: + if instance_id: + try: + client.delete_instance( + ecs_models.DeleteInstanceRequest(instance_id=instance_id, force=True) + ) + print(f"Deleted builder instance {instance_id}", flush=True) + except Exception as error: # noqa: BLE001 + print(f"Failed to delete builder instance {instance_id}: {error}", flush=True) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service b/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service new file mode 100644 index 00000000000..a1cbe6a4503 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service @@ -0,0 +1,45 @@ +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +[Unit] +Description=Ephemeral GitHub Actions runner +After=network-online.target +Wants=network-online.target +# cloud-final runs the provision user data, which writes +# /etc/ephemeral-github-runner.env and the drop-ins below. Ordering after it +# makes the runner come up configured on first boot, no restart needed. +# (cloud-init's `systemctl restart` in user data remains as the bridge for +# images built before this ordering existed.) +After=cloud-final.service + +[Service] +Type=simple +# The toolchain lives in the image at /opt/cargo/bin (the Dockerfile's ENV +# PATH does not survive materialization onto the host). Publish it here so +# every runner job inherits it. /etc/ephemeral-github-runner.env currently +# carries the same PATH line as a bridge for images built before this +# directive existed; EnvironmentFile is applied after Environment=, and the +# two values are kept identical, so there is no conflict. +Environment=PATH=/opt/cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +EnvironmentFile=/etc/ephemeral-github-runner.env +ExecStart=/opt/ephemeral-github-runner/start-runner.sh +Restart=no +# If the kernel OOM killer kills a job child (rustc/cargo), keep the runner +# service alive so the job is reported as a normal failure with logs instead +# of the runner vanishing and GitHub reporting a cancellation. cloud-init's +# oom.conf drop-in carries the same directive for images built before this. +OOMPolicy=continue + +[Install] +WantedBy=multi-user.target diff --git a/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh b/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh new file mode 100644 index 00000000000..a6aff87ca22 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Configure and start the ephemeral GitHub Actions runner on an ECS instance. +# Baked into the custom image at /opt/ephemeral-github-runner/ by +# build-ecs-image.py; the per-run environment is written by cloud-init to +# /etc/ephemeral-github-runner.env. +set -euo pipefail + +# shellcheck disable=SC1091 +source /etc/ephemeral-github-runner.env + +export HOME=/home/runner +cd /home/runner + +if [[ ! -f /home/runner/.runner ]]; then + runuser -u runner -- ./config.sh \ + --url "${REPO_URL}" \ + --token "${RUNNER_TOKEN}" \ + --name "${RUNNER_NAME}" \ + --labels "${RUNNER_LABELS}" \ + --ephemeral --unattended --replace --disableupdate +fi + +exec runuser -u runner -- ./run.sh diff --git a/.github/runner-scale-sets/query-regression/values-8-cores.yaml b/.github/runner-scale-sets/query-regression/values-8-cores.yaml deleted file mode 100644 index cff1322680b..00000000000 --- a/.github/runner-scale-sets/query-regression/values-8-cores.yaml +++ /dev/null @@ -1,146 +0,0 @@ -githubConfigUrl: "https://github.com/GreptimeTeam/greptimedb" -githubConfigSecret: greptimedb-arc-github-app - -runnerScaleSetName: "perf-regression-8-cores" -minRunners: 0 -maxRunners: 1 - -template: - spec: - automountServiceAccountToken: false - activeDeadlineSeconds: 12600 - nodeSelector: - kubernetes.io/hostname: minipc-3 - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - fsGroup: 1001 - fsGroupChangePolicy: OnRootMismatch - seccompProfile: - type: RuntimeDefault - volumes: - - name: cargo-home - emptyDir: {} - - name: build-cache - persistentVolumeClaim: - claimName: query-regression-build-cache - initContainers: - - name: initialize-build-cache - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11 - imagePullPolicy: IfNotPresent - command: - - /bin/sh - - -ec - - | - umask 0002 - for directory in cargo-registry-v1 cargo-git-v1 query-regression-target-v1 meta-v1 sccache-v1; do - cache_directory="/cache/${directory}" - mkdir -p "${cache_directory}" - test -w "${cache_directory}" - probe_file="${cache_directory}/.write-probe" - : > "${probe_file}" - test -f "${probe_file}" - rm "${probe_file}" - done - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - allowPrivilegeEscalation: false - capabilities: - drop: - - ALL - resources: - requests: - cpu: 100m - memory: 128Mi - ephemeral-storage: 1Gi - limits: - cpu: 500m - memory: 512Mi - ephemeral-storage: 1Gi - volumeMounts: - - name: build-cache - mountPath: /cache - containers: - - name: runner - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11 - imagePullPolicy: IfNotPresent - command: ["/home/runner/run.sh"] - env: - - name: CARGO_HOME - value: /home/runner/.cargo - - name: HTTP_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: HTTP_PROXY - optional: true - - name: HTTPS_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: HTTPS_PROXY - optional: true - - name: NO_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: NO_PROXY - optional: true - - name: UV_CACHE_DIR - value: /home/runner/.cargo/uv-cache - - name: RUSTUP_HOME - value: /opt/rustup - - name: RUSTUP_TOOLCHAIN - value: nightly-2026-03-21 - - name: RUSTUP_AUTO_INSTALL - value: "0" - - name: CARGO_TARGET_DIR - value: /home/runner/query-regression-target - - name: QUERY_REGRESSION_CACHE_META - value: /home/runner/query-regression-cache-meta - - name: RUSTC_WRAPPER - value: /usr/local/bin/sccache - - name: SCCACHE_DIR - value: /home/runner/.cache/sccache - - name: SCCACHE_CACHE_SIZE - value: 40G - - name: CARGO_INCREMENTAL - value: "0" - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - allowPrivilegeEscalation: false - capabilities: - drop: - - ALL - resources: - requests: - cpu: "6" - memory: 24Gi - ephemeral-storage: 80Gi - limits: - cpu: "8" - memory: 24Gi - ephemeral-storage: 80Gi - volumeMounts: - - name: cargo-home - mountPath: /home/runner/.cargo - - name: build-cache - mountPath: /home/runner/.cargo/registry - subPath: cargo-registry-v1 - - name: build-cache - mountPath: /home/runner/.cargo/git - subPath: cargo-git-v1 - - name: build-cache - mountPath: /home/runner/query-regression-target - subPath: query-regression-target-v1 - - name: build-cache - mountPath: /home/runner/query-regression-cache-meta - subPath: meta-v1 - - name: build-cache - mountPath: /home/runner/.cache/sccache - subPath: sccache-v1 diff --git a/.github/runner-scale-sets/query-regression/values-paused.yaml b/.github/runner-scale-sets/query-regression/values-paused.yaml deleted file mode 100644 index fe734ef7615..00000000000 --- a/.github/runner-scale-sets/query-regression/values-paused.yaml +++ /dev/null @@ -1,2 +0,0 @@ -minRunners: 0 -maxRunners: 0 diff --git a/.github/scripts/aliyun-ecs-runner-provision.py b/.github/scripts/aliyun-ecs-runner-provision.py new file mode 100644 index 00000000000..74a69681f3d --- /dev/null +++ b/.github/scripts/aliyun-ecs-runner-provision.py @@ -0,0 +1,485 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This is the first `.github/scripts/` user of PEP 723 inline script metadata: +# `uv run` installs the pinned Aliyun SDK before executing. The SDK import stays +# lazy (inside `make_ecs_client`) so unit tests can import this module with a +# plain stdlib interpreter. +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# # Pin exactly after the first live run against the Aliyun API. +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Provision an ephemeral Aliyun ECS query-regression runner. + +Creates one pay-as-you-go ECS instance from the query-regression custom image +and waits until the instance registers itself as an ephemeral GitHub Actions +runner with a per-run label. Build caches live on the instance system disk +and disappear with the VM. The instance id is written to job outputs as soon +as the VM exists so the teardown job is the single DeleteInstance caller (a +timeout here used to delete and then teardown deleted again). On runner +online timeout the console is dumped and the job fails; teardown releases +the instance. + +The Aliyun credentials are used only in this control-plane job (ubuntu-latest) +and are never passed to the ECS instance; the instance receives only a +short-lived GitHub runner registration token via user data. +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import time +import urllib.error +import urllib.request +from dataclasses import dataclass + +RUNNER_NAME_PREFIX = "qreg-ecs" +RUNNER_LABEL_PREFIX = "query-regression-ecs" +MANAGED_BY_TAG_KEY = "managed-by" +MANAGED_BY_TAG_VALUE = "query-regression-ci" +RUN_TAG_KEY = "query-regression-run-id" + +# Runner cache paths on the instance system disk. They are created empty +# every provision and deleted with the VM. +CACHE_PATHS = ( + "/home/runner/.cargo/registry", + "/home/runner/.cargo/git", + "/home/runner/query-regression-target", + "/home/runner/query-regression-cache-meta", + "/home/runner/.cache/sccache", +) + +RUNNER_ONLINE_TIMEOUT_SECONDS = 10 * 60 +POLL_INTERVAL_SECONDS = 5 + +# Nightly thin-LTO of greptime peaks above ecs.c9i.2xlarge's 16 GiB. A swap +# file plus masking systemd-oomd lets the kernel reclaim rustc pages instead +# of SIGTERM-ing the runner service cgroup (GitHub then reports the step as +# "The operation was canceled"). Sized to match that 16 GiB instance; a +# 32 GiB type still benefits as a safety net. Lives on the system disk next +# to the cold build caches. +SWAP_FILE = "/swapfile" +SWAP_SIZE_GIB = 16 + + +@dataclass(frozen=True) +class ProvisionConfig: + region_id: str + vswitch_id: str + security_group_id: str + image_id: str + instance_type: str + repo: str + run_id: str + github_token: str + # Optional resource group; required when the RAM grant is scoped to one. + resource_group_id: str | None = None + # Runner identity inside the image; the workflow asserts the same values. + runner_uid: str = "1001" + runner_gid: str = "1001" + + +def runner_name_for_run(run_id: str) -> str: + return f"{RUNNER_NAME_PREFIX}-{run_id}" + + +def runner_label_for_run(run_id: str) -> str: + return f"{RUNNER_LABEL_PREFIX}-{run_id}" + + +def render_user_data( + runner_name: str, + runner_label: str, + runner_token: str, + repo: str, + runner_uid: str = "1001", + runner_gid: str = "1001", +) -> str: + """Render the cloud-init shell script for the runner instance.""" + destinations = " ".join(f'"{dst}"' for dst in CACHE_PATHS) + cache_setup = f"""# Caches live on the system disk, are deleted with the instance, and every +# run compiles cold. Within-run reuse (base warming candidate via the shared +# target dir and sccache) still applies. +mkdir -p {destinations} +chown -R {runner_uid}:{runner_gid} /home/runner""" + swap_setup = f"""# Mask systemd-oomd before enabling swap: Ubuntu 24.04 kills the whole +# service cgroup on PSI pressure, and swap thrashing looks like pressure. +systemctl disable --now systemd-oomd.socket systemd-oomd.service || true +systemctl mask systemd-oomd.socket systemd-oomd.service || true +# apt-daily-upgrade / unattended-upgrades can `systemctl restart` services +# whose libraries were updated. That SIGTERM-s the runner mid-job; GitHub +# records UserCancelled (job still valid, child exit 143). These VMs live +# ~1h; security updates belong in the image, not at job time. +systemctl disable --now unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl stop unattended-upgrades.service || true +killall -9 unattended-upgr || true +cat > /etc/apt/apt.conf.d/99disable-auto-updates <<'APTEOF' +APT::Periodic::Update-Package-Lists "0"; +APT::Periodic::Unattended-Upgrade "0"; +APT::Periodic::Download-Upgradeable-Packages "0"; +APTEOF +if [[ ! -f "{SWAP_FILE}" ]]; then + fallocate --length {SWAP_SIZE_GIB}G "{SWAP_FILE}" + chmod 600 "{SWAP_FILE}" + mkswap "{SWAP_FILE}" +fi +swapon "{SWAP_FILE}" +sysctl --write vm.swappiness=10 +swapon --show +free --human""" + return f"""#!/bin/bash +set -euo pipefail + +{cache_setup} + +{swap_setup} + +cat > /etc/ephemeral-github-runner.env <<'ENVEOF' +RUNNER_NAME={runner_name} +RUNNER_LABELS={runner_label} +RUNNER_TOKEN={runner_token} +REPO_URL=https://github.com/{repo} +# The container image baked /opt/cargo/bin into PATH via Dockerfile ENV; on +# the host nothing inherits that, so publish it through the unit's +# EnvironmentFile: runner job processes inherit the runner's environment. +PATH=/opt/cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +ENVEOF + +# Stream the runner service's own logs (config output, job lifecycle, crash +# messages) to the serial console, next to the kernel's OOM-killer records. +# Teardown dumps the console tail before deleting the instance, so the +# witness of what killed the runner survives the machine. +mkdir -p /etc/systemd/system/ephemeral-github-runner.service.d +cat > /etc/systemd/system/ephemeral-github-runner.service.d/console.conf <<'CONFEOF' +[Service] +StandardOutput=journal+console +StandardError=journal+console +CONFEOF +# systemd's DefaultOOMPolicy=stop SIGTERM-s the whole unit when *any* +# process in the cgroup is OOM-killed (typically rustc during thin LTO). +# That is the "Session terminated, killing shell..." / job Canceled path: +# runuser dies, GitHub deletes the ephemeral registration, always() steps +# never run. continue keeps the runner up so cargo can report signal 9. +cat > /etc/systemd/system/ephemeral-github-runner.service.d/oom.conf <<'CONFEOF' +[Service] +OOMPolicy=continue +CONFEOF +systemctl daemon-reload +# restart (not start): the image enables this unit, so it may already +# have come up without the drop-ins if it raced cloud-init. restart +# applies OOMPolicy=continue to the running cgroup. --no-block matters: +# newer image units carry After=cloud-final.service, and this script IS +# cloud-final — a blocking restart would deadlock until TimeoutStartSec. +systemctl restart --no-block ephemeral-github-runner.service +""" + + +def encode_user_data(script: str) -> str: + return base64.b64encode(script.encode("utf-8")).decode("ascii") + + +def github_api(token: str, method: str, path: str, body: dict | None = None) -> dict: + data = json.dumps(body).encode("utf-8") if body is not None else None + request = urllib.request.Request( + f"https://api.github.com{path}", + data=data, + method=method, + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as error: + # GitHub's error body says exactly why (e.g. "Must have admin rights to + # Repository" for a PAT without the required scope); surface it instead + # of a bare "HTTP Error 403". + body = error.read().decode("utf-8", "replace") + raise SystemExit( + f"GitHub API {method} {path} failed: HTTP {error.code}: {body}\n" + "The token comes from the GH_PERSONAL_ACCESS_TOKEN secret; it needs " + "'repo' scope (classic PAT) or 'Administration: write' on the " + "repository (fine-grained PAT)." + ) from error + + +def create_registration_token(github_token: str, repo: str) -> str: + response = github_api( + github_token, + "POST", + f"/repos/{repo}/actions/runners/registration-token", + body={}, + ) + return response["token"] + + +def find_runner_by_name(github_token: str, repo: str, name: str) -> dict | None: + page = 1 + while True: + response = github_api( + github_token, + "GET", + f"/repos/{repo}/actions/runners?per_page=100&page={page}", + ) + runners = response.get("runners", []) + for runner in runners: + if runner.get("name") == name: + return runner + if len(runners) < 100: + return None + page += 1 + + +def make_ecs_client(config: ProvisionConfig): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + access_key_id = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID", "") + access_key_secret = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "") + if not access_key_id or not access_key_secret: + raise SystemExit( + "ALIBABA_CLOUD_ACCESS_KEY_ID/SECRET are empty; in CI they come from the " + "ALICLOUD_ECS_ACCESS_KEY_ID/SECRET repository secrets (a missing or " + "misnamed secret expands to an empty string)." + ) + return EcsClient( + OpenApiConfig( + access_key_id=access_key_id, + access_key_secret=access_key_secret, + region_id=config.region_id, + endpoint=f"ecs.{config.region_id}.aliyuncs.com", + ) + ) + + +def wait_for_instance_status(client, region_id: str, instance_id: str, wanted: str, deadline: float) -> None: + from alibabacloud_ecs20140526 import models as ecs_models + + while time.monotonic() < deadline: + response = client.describe_instances( + ecs_models.DescribeInstancesRequest( + region_id=region_id, instance_ids=json.dumps([instance_id]) + ) + ) + instances = response.body.instances.instance + if instances and instances[0].status == wanted: + return + time.sleep(POLL_INTERVAL_SECONDS) + raise TimeoutError(f"Instance {instance_id} did not reach status {wanted} in time") + + +def fetch_console_output(client, region_id: str, instance_id: str) -> str: + from alibabacloud_ecs20140526 import models as ecs_models + + response = client.get_instance_console_output( + ecs_models.GetInstanceConsoleOutputRequest(region_id=region_id, instance_id=instance_id) + ) + return base64.b64decode(response.body.console_output or "").decode("utf-8", "replace") + + +class ConsoleTailer: + """Incrementally prints new serial-console lines so cloud-init progress is + visible in the CI log while waiting, not only after a failure dump.""" + + def __init__(self, client, region_id: str, instance_id: str) -> None: + self.client = client + self.region_id = region_id + self.instance_id = instance_id + self.printed_lines = 0 + + def poll(self) -> None: + try: + lines = fetch_console_output(self.client, self.region_id, self.instance_id).splitlines() + except Exception as error: # noqa: BLE001 + print(f"[console] unable to fetch console output: {error}", flush=True) + return + # The API returns a bounded tail; if it ever shrinks, restart from the + # beginning of the new buffer rather than skipping lines. + if len(lines) < self.printed_lines: + self.printed_lines = 0 + for line in lines[self.printed_lines :]: + print(f"[console] {line}", flush=True) + self.printed_lines = len(lines) + + +def dump_console_output(client, region_id: str, instance_id: str) -> None: + try: + output = fetch_console_output(client, region_id, instance_id) + except Exception as error: # noqa: BLE001 + print(f"Unable to fetch console output for {instance_id}: {error}", flush=True) + return + tail = "\n".join(output.splitlines()[-80:]) + print(f"::group::Console output tail for {instance_id}\n{tail}\n::endgroup::", flush=True) + summary_path = os.environ.get("GITHUB_STEP_SUMMARY") + if summary_path: + with open(summary_path, "a", encoding="utf-8") as summary: + summary.write(f"
ECS console output ({instance_id})\n\n```\n") + summary.write(tail) + summary.write("\n```\n
\n") + + +def append_github_output(name: str, value: str) -> None: + output_path = os.environ.get("GITHUB_OUTPUT") + if output_path: + with open(output_path, "a", encoding="utf-8") as output: + output.write(f"{name}={value}\n") + + +def run_instance(client, config: ProvisionConfig, user_data: str) -> str: + from alibabacloud_ecs20140526 import models as ecs_models + + request = ecs_models.RunInstancesRequest( + region_id=config.region_id, + image_id=config.image_id, + resource_group_id=config.resource_group_id, + instance_type=config.instance_type, + v_switch_id=config.vswitch_id, + security_group_id=config.security_group_id, + instance_name=runner_name_for_run(config.run_id), + host_name=runner_name_for_run(config.run_id), + description=f"Ephemeral query-regression runner for run {config.run_id}", + amount=1, + instance_charge_type="PostPaid", + spot_strategy="NoSpot", + internet_charge_type="PayByTraffic", + internet_max_bandwidth_out=100, + # Holds the image (~12G), the 16G swapfile, the source checkout, the + # base+candidate cluster data homes, and cold build caches. Deleted + # with the instance. + system_disk=ecs_models.RunInstancesRequestSystemDisk(category="cloud_essd", size="40"), + user_data=user_data, + tag=[ + ecs_models.RunInstancesRequestTag(key=MANAGED_BY_TAG_KEY, value=MANAGED_BY_TAG_VALUE), + ecs_models.RunInstancesRequestTag(key=RUN_TAG_KEY, value=config.run_id), + ], + ) + response = client.run_instances(request) + instance_id = response.body.instance_id_sets.instance_id_set[0] + return instance_id + + +def provision(config: ProvisionConfig) -> int: + client = make_ecs_client(config) + runner_name = runner_name_for_run(config.run_id) + runner_label = runner_label_for_run(config.run_id) + + registration_token = create_registration_token(config.github_token, config.repo) + user_data = encode_user_data( + render_user_data( + runner_name, + runner_label, + registration_token, + config.repo, + config.runner_uid, + config.runner_gid, + ) + ) + + print(f"::group::Run instance {runner_name}", flush=True) + instance_id = run_instance(client, config, user_data) + print(f"Instance id: {instance_id}", flush=True) + # Teardown is the only DeleteInstance caller. Write outputs immediately + # so a later status/online failure still releases this VM. + append_github_output("label", runner_label) + append_github_output("instance_id", instance_id) + append_github_output("runner_name", runner_name) + wait_for_instance_status(client, config.region_id, instance_id, "Running", time.monotonic() + 5 * 60) + print("::endgroup::", flush=True) + + print("::group::Wait for runner online", flush=True) + deadline = time.monotonic() + RUNNER_ONLINE_TIMEOUT_SECONDS + tailer = ConsoleTailer(client, config.region_id, instance_id) + last_console_poll = 0.0 + while time.monotonic() < deadline: + runner = find_runner_by_name(config.github_token, config.repo, runner_name) + if runner is not None: + status = runner.get("status") + print(f"Runner {runner_name} status: {status}", flush=True) + if status == "online": + print("::endgroup::", flush=True) + print(f"Runner {runner_name} is online with label {runner_label}", flush=True) + return 0 + # The serial console lags the guest by a minute or so; 30s polling is + # enough to follow cloud-init without tripping API throttling. + if time.monotonic() - last_console_poll >= 30: + tailer.poll() + last_console_poll = time.monotonic() + time.sleep(POLL_INTERVAL_SECONDS) + + print("::endgroup::", flush=True) + print( + f"Runner {runner_name} did not come online in time; " + "leaving instance for the teardown job to delete", + flush=True, + ) + dump_console_output(client, config.region_id, instance_id) + return 1 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--vswitch-id", default=os.environ.get("ALIYUN_ECS_VSWITCH_ID")) + parser.add_argument("--security-group-id", default=os.environ.get("ALIYUN_ECS_SECURITY_GROUP_ID")) + parser.add_argument("--image-id", default=os.environ.get("QUERY_REGRESSION_ECS_IMAGE_ID")) + parser.add_argument("--instance-type", default=os.environ.get("ALIYUN_ECS_INSTANCE_TYPE")) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument("--run-id", default=os.environ.get("GITHUB_RUN_ID")) + parser.add_argument("--github-token", default=os.environ.get("GH_PERSONAL_ACCESS_TOKEN")) + parser.add_argument("--runner-uid", default=os.environ.get("QUERY_REGRESSION_RUNNER_UID", "1001")) + parser.add_argument("--runner-gid", default=os.environ.get("QUERY_REGRESSION_RUNNER_GID", "1001")) + parser.add_argument("--resource-group-id", default=os.environ.get("ALIYUN_ECS_RESOURCE_GROUP_ID")) + args = parser.parse_args() + + missing = [ + name + for name, value in vars(args).items() + if name not in ("resource_group_id",) + and (value is None or (isinstance(value, str) and not value)) + ] + if missing: + flags = ", ".join(f"--{name.replace('_', '-')}" for name in missing) + raise SystemExit(f"Missing required configuration: {flags} (or their env defaults)") + + config = ProvisionConfig( + region_id=args.region_id, + vswitch_id=args.vswitch_id, + security_group_id=args.security_group_id, + image_id=args.image_id, + instance_type=args.instance_type, + repo=args.repo, + run_id=args.run_id, + github_token=args.github_token, + resource_group_id=args.resource_group_id or None, + runner_uid=args.runner_uid, + runner_gid=args.runner_gid, + ) + return provision(config) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/aliyun-ecs-runner-teardown.py b/.github/scripts/aliyun-ecs-runner-teardown.py new file mode 100644 index 00000000000..36c80253b95 --- /dev/null +++ b/.github/scripts/aliyun-ecs-runner-teardown.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# PEP 723 inline metadata (see aliyun-ecs-runner-provision.py for the +# convention). The SDK import stays lazy so unit tests run on a plain +# stdlib interpreter. +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Tear down ephemeral Aliyun ECS query-regression runners. + +Two modes: + +- Targeted (per workflow run): delete one instance by id and deregister its + runner by name. Both steps are idempotent and best-effort; a missing + instance or runner is not an error. +- Sweep (scheduled janitor): delete every instance tagged as managed by + query-regression CI whose creation time is older than the given TTL, and + deregister the matching runners. This is the safety net for runs whose + teardown job never executed. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import os +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path + +# Reuse the GitHub API client, runner lookup, and tag constants from the +# provision script. +_PROVISION_SPEC = importlib.util.spec_from_file_location( + "aliyun_ecs_runner_provision", + Path(__file__).resolve().parent / "aliyun-ecs-runner-provision.py", +) +assert _PROVISION_SPEC is not None and _PROVISION_SPEC.loader is not None +provision = importlib.util.module_from_spec(_PROVISION_SPEC) +sys.modules[_PROVISION_SPEC.name] = provision +_PROVISION_SPEC.loader.exec_module(provision) + +SWEEP_TTL = timedelta(hours=4) +# ECS DescribeInstances creation_time is ISO8601 UTC, e.g. "2026-08-17T02:13Z". +CREATION_TIME_FORMATS = ("%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%MZ") + + +def parse_creation_time(value: str) -> datetime: + for fmt in CREATION_TIME_FORMATS: + try: + return datetime.strptime(value, fmt).replace(tzinfo=timezone.utc) + except ValueError: + continue + raise ValueError(f"Unparsable ECS creation time: {value}") + + +def expired_instance_names( + instances: list[tuple[str, str, str]], now: datetime, ttl: timedelta +) -> list[tuple[str, str]]: + """Pick (instance_id, instance_name) pairs whose creation is older than ttl. + + `instances` items are (instance_id, instance_name, creation_time). + """ + expired = [] + for instance_id, instance_name, creation_time in instances: + age = now - parse_creation_time(creation_time) + if age >= ttl: + expired.append((instance_id, instance_name)) + return expired + + +def make_ecs_client(region_id: str): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + access_key_id = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID", "") + access_key_secret = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "") + if not access_key_id or not access_key_secret: + raise SystemExit( + "ALIBABA_CLOUD_ACCESS_KEY_ID/SECRET are empty; in CI they come from the " + "ALICLOUD_ECS_ACCESS_KEY_ID/SECRET repository secrets (a missing or " + "misnamed secret expands to an empty string)." + ) + return EcsClient( + OpenApiConfig( + access_key_id=access_key_id, + access_key_secret=access_key_secret, + region_id=region_id, + endpoint=f"ecs.{region_id}.aliyuncs.com", + ) + ) + + +def delete_instance(client, instance_id: str, region_id: str | None = None) -> bool: + import time + + from alibabacloud_ecs20140526 import models as ecs_models + + # Export the serial console before destroying the evidence: a cancelled + # run still runs this teardown job, and the console tail (cloud-init, + # runner service logs, kernel OOM records) is the only witness of what + # happened on the machine. + if region_id is not None: + provision.dump_console_output(client, region_id, instance_id) + + # DeleteInstance rejects instances that are still Initializing (or in + # another transitional status). Teardown runs right after a cancelled + # run, when the instance may be only seconds old, so retry the transient + # status errors for a few minutes instead of leaking the instance. + deadline = time.monotonic() + 5 * 60 + while True: + try: + client.delete_instance(ecs_models.DeleteInstanceRequest(instance_id=instance_id, force=True)) + print(f"Deleted instance {instance_id}", flush=True) + return True + except Exception as error: # noqa: BLE001 + message = str(error) + if "InvalidInstanceId.NotFound" in message: + print(f"Instance {instance_id} already gone", flush=True) + return True + if "IncorrectInstanceStatus" in message and time.monotonic() < deadline: + print(f"Instance {instance_id} is in a transitional status; retrying delete", flush=True) + time.sleep(15) + continue + print(f"Failed to delete instance {instance_id}: {error}", flush=True) + return False + + +def deregister_runner(token: str, repo: str, runner_name: str) -> bool: + runner = provision.find_runner_by_name(token, repo, runner_name) + if runner is None: + print(f"Runner {runner_name} is not registered", flush=True) + return True + try: + provision.github_api(token, "DELETE", f"/repos/{repo}/actions/runners/{runner['id']}") + print(f"Deregistered runner {runner_name} (id {runner['id']})", flush=True) + return True + except Exception as error: # noqa: BLE001 + print(f"Failed to deregister runner {runner_name}: {error}", flush=True) + return False + + +def list_managed_instances(client, region_id: str) -> list[tuple[str, str, str]]: + from alibabacloud_ecs20140526 import models as ecs_models + + result: list[tuple[str, str, str]] = [] + next_token = None + while True: + request = ecs_models.DescribeInstancesRequest( + region_id=region_id, + tag=[ + ecs_models.DescribeInstancesRequestTag( + key=provision.MANAGED_BY_TAG_KEY, value=provision.MANAGED_BY_TAG_VALUE + ) + ], + max_results=100, + next_token=next_token, + ) + response = client.describe_instances(request) + for instance in response.body.instances.instance: + result.append((instance.instance_id, instance.instance_name, instance.creation_time)) + next_token = response.body.next_token + if not next_token: + return result + + +def sweep(client, region_id: str, repo: str, github_token: str, ttl: timedelta) -> int: + instances = list_managed_instances(client, region_id) + print(f"Found {len(instances)} managed instance(s) in {region_id}", flush=True) + expired = expired_instance_names(instances, datetime.now(timezone.utc), ttl) + ok = True + for instance_id, instance_name in expired: + print(f"Instance {instance_id} ({instance_name}) exceeds TTL {ttl}; deleting", flush=True) + # Runner names mirror instance names by construction in the provision + # script (both are qreg-ecs-). + ok &= delete_instance(client, instance_id, region_id) + ok &= deregister_runner(github_token, repo, instance_name) + return 0 if ok else 1 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument("--github-token", default=os.environ.get("GH_PERSONAL_ACCESS_TOKEN")) + parser.add_argument("--instance-id", default=os.environ.get("QUERY_REGRESSION_ECS_INSTANCE_ID")) + parser.add_argument("--runner-name", default=os.environ.get("QUERY_REGRESSION_ECS_RUNNER_NAME")) + parser.add_argument( + "--sweep", + action="store_true", + help="Janitor mode: delete all managed instances older than --sweep-ttl-hours.", + ) + parser.add_argument("--sweep-ttl-hours", type=float, default=SWEEP_TTL.total_seconds() / 3600) + args = parser.parse_args() + + for name in ("region_id", "repo", "github_token"): + if not getattr(args, name): + raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}") + + client = make_ecs_client(args.region_id) + + if args.sweep: + return sweep( + client, args.region_id, args.repo, args.github_token, timedelta(hours=args.sweep_ttl_hours) + ) + + ok = True + if args.instance_id: + ok &= delete_instance(client, args.instance_id, args.region_id) + else: + print("No instance id given; skipping instance deletion", flush=True) + if args.runner_name: + ok &= deregister_runner(args.github_token, args.repo, args.runner_name) + else: + print("No runner name given; skipping runner deregistration", flush=True) + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/query-regression-nightly-refs.py b/.github/scripts/query-regression-nightly-refs.py new file mode 100644 index 00000000000..6e9f54682d4 --- /dev/null +++ b/.github/scripts/query-regression-nightly-refs.py @@ -0,0 +1,270 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resolve consecutive successful Nightly Build SHAs for query-regression. + +Query regression builds both binaries from git; this script only picks the +head SHAs of two Nightly Build workflow runs (today vs the previous success). +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import urllib.error +import urllib.parse +import urllib.request +from dataclasses import dataclass +from typing import Any + + +NIGHTLY_WORKFLOW = "nightly-build.yml" +NIGHTLY_EVENTS = frozenset({"schedule", "workflow_dispatch"}) + + +@dataclass(frozen=True) +class WorkflowRun: + id: int + head_sha: str + head_branch: str + html_url: str + created_at: str + conclusion: str + event: str + + +@dataclass(frozen=True) +class RefPair: + base: WorkflowRun | None + candidate: WorkflowRun | None + skip: bool + reason: str + + +def parse_run(payload: dict[str, Any]) -> WorkflowRun: + return WorkflowRun( + id=int(payload["id"]), + head_sha=str(payload["head_sha"]), + head_branch=str(payload.get("head_branch") or ""), + html_url=str(payload.get("html_url") or ""), + created_at=str(payload.get("created_at") or ""), + conclusion=str(payload.get("conclusion") or ""), + event=str(payload.get("event") or ""), + ) + + +def github_get(token: str, path: str, query: dict[str, str] | None = None) -> dict[str, Any]: + encoded = f"?{urllib.parse.urlencode(query)}" if query else "" + request = urllib.request.Request( + f"https://api.github.com{path}{encoded}", + method="GET", + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as error: + body = error.read().decode("utf-8", "replace") + raise SystemExit( + f"GitHub API GET {path} failed: HTTP {error.code}: {body}" + ) from error + + +def list_successful_nightly_runs( + token: str, + repo: str, + *, + branch: str | None = None, + per_page: int = 30, +) -> list[WorkflowRun]: + query = { + "status": "success", + "per_page": str(per_page), + } + if branch: + query["branch"] = branch + payload = github_get( + token, + f"/repos/{repo}/actions/workflows/{NIGHTLY_WORKFLOW}/runs", + query, + ) + runs = [parse_run(item) for item in payload.get("workflow_runs") or []] + return [ + run + for run in runs + if run.conclusion == "success" and run.event in NIGHTLY_EVENTS and run.head_sha + ] + + +def fetch_run(token: str, repo: str, run_id: int) -> WorkflowRun: + payload = github_get(token, f"/repos/{repo}/actions/runs/{run_id}") + return parse_run(payload) + + +def select_base_and_candidate( + runs: list[WorkflowRun], + *, + candidate_run_id: int | None = None, + candidate: WorkflowRun | None = None, +) -> RefPair: + """Pick candidate (newest / requested) and the previous successful run. + + `runs` must be newest-first. Same-SHA consecutive nightlies are skipped + (no commits that day). + """ + if candidate is None: + if candidate_run_id is not None: + candidate = next((run for run in runs if run.id == candidate_run_id), None) + if candidate is None: + return RefPair( + None, + None, + True, + f"candidate nightly run {candidate_run_id} was not a successful " + f"{NIGHTLY_WORKFLOW} run", + ) + elif runs: + candidate = runs[0] + else: + return RefPair(None, None, True, "no successful Nightly Build runs") + + branch = candidate.head_branch + older = [ + run + for run in runs + if run.id != candidate.id + and (not branch or not run.head_branch or run.head_branch == branch) + and (not candidate.created_at or run.created_at <= candidate.created_at) + ] + if not older: + return RefPair( + None, + candidate, + True, + "no previous successful Nightly Build to compare against", + ) + base = older[0] + if base.head_sha.lower() == candidate.head_sha.lower(): + return RefPair( + base, + candidate, + True, + f"previous nightly SHA {base.head_sha} matches candidate; nothing to compare", + ) + return RefPair(base, candidate, False, "") + + +def _override_run(sha: str) -> WorkflowRun: + return WorkflowRun( + id=0, + head_sha=sha, + head_branch="", + html_url="", + created_at="", + conclusion="success", + event="workflow_dispatch", + ) + + +def override_pair(base_ref: str, candidate_ref: str) -> RefPair: + return RefPair(_override_run(base_ref), _override_run(candidate_ref), False, "explicit refs") + + +def write_outputs(pair: RefPair) -> None: + values = { + "skip": "true" if pair.skip else "false", + "reason": pair.reason, + "base_sha": pair.base.head_sha if pair.base else "", + "candidate_sha": pair.candidate.head_sha if pair.candidate else "", + "base_run_url": pair.base.html_url if pair.base else "", + "candidate_run_url": pair.candidate.html_url if pair.candidate else "", + "base_run_id": str(pair.base.id) if pair.base and pair.base.id else "", + "candidate_run_id": str(pair.candidate.id) if pair.candidate and pair.candidate.id else "", + } + output_path = os.environ.get("GITHUB_OUTPUT") + if output_path: + with open(output_path, "a", encoding="utf-8") as handle: + for key, value in values.items(): + handle.write(f"{key}={value}\n") + for key, value in values.items(): + print(f"{key}={value}") + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument( + "--token", + default=os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN"), + ) + parser.add_argument("--branch", default=os.environ.get("NIGHTLY_BRANCH") or "") + parser.add_argument( + "--candidate-run-id", + default=os.environ.get("CANDIDATE_RUN_ID") or "", + ) + parser.add_argument("--base-ref", default=os.environ.get("BASE_REF") or "") + parser.add_argument("--candidate-ref", default=os.environ.get("CANDIDATE_REF") or "") + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + base_ref = args.base_ref.strip() + candidate_ref = args.candidate_ref.strip() + if base_ref and candidate_ref: + write_outputs(override_pair(base_ref, candidate_ref)) + return 0 + if bool(base_ref) != bool(candidate_ref): + print("Provide both --base-ref and --candidate-ref, or neither.", file=sys.stderr) + return 2 + if not args.repo or not args.token: + print("--repo and --token (or GITHUB_REPOSITORY / GITHUB_TOKEN) are required.", file=sys.stderr) + return 2 + + candidate_run_id = int(args.candidate_run_id) if str(args.candidate_run_id).strip() else None + candidate: WorkflowRun | None = None + branch = args.branch.strip() or None + if candidate_run_id is not None: + candidate = fetch_run(args.token, args.repo, candidate_run_id) + if candidate.conclusion != "success": + write_outputs( + RefPair( + None, + candidate, + True, + f"nightly run {candidate_run_id} conclusion is {candidate.conclusion}", + ) + ) + return 0 + branch = branch or candidate.head_branch or None + + runs = list_successful_nightly_runs(args.token, args.repo, branch=branch) + pair = select_base_and_candidate( + runs, + candidate_run_id=candidate_run_id if candidate is None else None, + candidate=candidate, + ) + write_outputs(pair) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 7f064f52eb5..8d7a37080b3 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -80,6 +80,14 @@ jobs: run: python3 scripts/check-enterprise-license-test.py - name: Run fuzz orchestration script tests run: .github/scripts/run-fuzz-targets-test.sh + - name: Run query-regression Python tooling tests + run: | + python3 tests/perf/test_query_regression_runner_compaction_toctou.py + python3 tests/perf/test_query_regression_runner_otlp_trace_load.py + python3 tests/perf/test_query_regression_summary_otlp.py + python3 tests/perf/test_query_regression_case_selection.py + python3 tests/perf/test_query_regression_nightly_refs.py + python3 tests/perf/test_aliyun_ecs_runner_scripts.py check: if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} diff --git a/.github/workflows/query-regression-janitor.yml b/.github/workflows/query-regression-janitor.yml new file mode 100644 index 00000000000..ce277ce2f02 --- /dev/null +++ b/.github/workflows/query-regression-janitor.yml @@ -0,0 +1,36 @@ +name: Query Regression ECS Janitor + +# Safety net for the aliyun-ecs query-regression path: deletes managed ECS +# instances (and their runner registrations) that outlived their workflow run, +# so a failed teardown never leaks pay-as-you-go spend. + +on: + schedule: + - cron: "23 3 * * *" + workflow_dispatch: + +permissions: + contents: read + +jobs: + sweep: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Checkout teardown script + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Sweep expired ECS runners + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-teardown.py + --sweep --sweep-ttl-hours 4 diff --git a/.github/workflows/query-regression-nightly.yml b/.github/workflows/query-regression-nightly.yml new file mode 100644 index 00000000000..027fcd16987 --- /dev/null +++ b/.github/workflows/query-regression-nightly.yml @@ -0,0 +1,104 @@ +name: Query Regression Nightly + +# After a successful GreptimeDB Nightly Build, compare that commit against +# the previous successful nightly. The reusable Query Regression workflow +# still compiles both SHAs; this wrapper only resolves which two SHAs. + +on: + workflow_run: + workflows: + - GreptimeDB Nightly Build + types: + - completed + workflow_dispatch: + inputs: + base_ref: + description: Base ref/SHA (empty = previous successful nightly) + required: false + type: string + default: "" + candidate_ref: + description: Candidate ref/SHA (empty = latest successful nightly) + required: false + type: string + default: "" + candidate_run_id: + description: Nightly Build run id to treat as candidate (empty = latest) + required: false + type: string + default: "" + case: + description: Query perf case path(s), all, or heavy + required: false + type: string + default: all + +permissions: + contents: read + actions: read + +jobs: + resolve-refs: + name: Resolve previous vs current nightly SHAs + if: ${{ github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success' }} + runs-on: ubuntu-latest + timeout-minutes: 5 + outputs: + skip: ${{ steps.refs.outputs.skip }} + reason: ${{ steps.refs.outputs.reason }} + base_sha: ${{ steps.refs.outputs.base_sha }} + candidate_sha: ${{ steps.refs.outputs.candidate_sha }} + steps: + - name: Checkout ref resolver + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Resolve nightly SHAs + id: refs + env: + GITHUB_TOKEN: ${{ github.token }} + GITHUB_REPOSITORY: ${{ github.repository }} + CANDIDATE_RUN_ID: ${{ github.event_name == 'workflow_run' && github.event.workflow_run.id || inputs.candidate_run_id || '' }} + BASE_REF: ${{ github.event_name == 'workflow_dispatch' && inputs.base_ref || '' }} + CANDIDATE_REF: ${{ github.event_name == 'workflow_dispatch' && inputs.candidate_ref || '' }} + run: python3 .github/scripts/query-regression-nightly-refs.py + + - name: Summarize comparison + if: always() + env: + SKIP: ${{ steps.refs.outputs.skip }} + REASON: ${{ steps.refs.outputs.reason }} + BASE_SHA: ${{ steps.refs.outputs.base_sha }} + CANDIDATE_SHA: ${{ steps.refs.outputs.candidate_sha }} + BASE_RUN_URL: ${{ steps.refs.outputs.base_run_url }} + CANDIDATE_RUN_URL: ${{ steps.refs.outputs.candidate_run_url }} + run: | + set -euo pipefail + { + if [[ "${SKIP}" == "true" ]]; then + printf 'Skipping query-regression nightly: %s\n' "${REASON}" + else + printf 'Comparing previous nightly `%s` -> current nightly `%s`\n' \ + "${BASE_SHA}" "${CANDIDATE_SHA}" + if [[ -n "${BASE_RUN_URL}" ]]; then + printf -- '- Previous Nightly Build: %s\n' "${BASE_RUN_URL}" + fi + if [[ -n "${CANDIDATE_RUN_URL}" ]]; then + printf -- '- Current Nightly Build: %s\n' "${CANDIDATE_RUN_URL}" + fi + fi + } | tee -a "${GITHUB_STEP_SUMMARY}" + + query-regression: + name: Query regression nightly + needs: [resolve-refs] + if: ${{ needs.resolve-refs.outputs.skip != 'true' && needs.resolve-refs.outputs.base_sha != '' && needs.resolve-refs.outputs.candidate_sha != '' }} + uses: ./.github/workflows/query-regression.yml + secrets: inherit + with: + case: ${{ github.event_name == 'workflow_dispatch' && inputs.case || 'all' }} + base_ref: ${{ needs.resolve-refs.outputs.base_sha }} + candidate_ref: ${{ needs.resolve-refs.outputs.candidate_sha }} + cargo_profile: nightly + runner: aliyun-ecs diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index 0806d4d3e38..209c8d76fff 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -32,10 +32,10 @@ on: type: string default: nightly runner: - description: Self-hosted runner label or ARC runner scale set for this query regression run + description: Self-hosted runner label; aliyun-ecs provisions a fresh ECS instance per run required: false type: string - default: perf-regression-8-cores + default: aliyun-ecs workflow_dispatch: inputs: case: @@ -68,12 +68,20 @@ on: - release - dev runner: - description: Self-hosted runner label or ARC runner scale set for this query regression run + description: >- + Self-hosted runner label; aliyun-ecs provisions a fresh ECS instance + per run, any other value is used as a literal runner label required: true - type: choice - default: perf-regression-8-cores - options: - - perf-regression-8-cores + type: string + default: aliyun-ecs + keep_instance: + description: >- + Debug: keep the ECS instance after the run (skip teardown) so its + runner _diag logs, journal, and telemetry can be inspected; the + janitor still sweeps it after the TTL + required: false + type: boolean + default: false pull_request: types: [labeled] @@ -81,19 +89,99 @@ permissions: contents: read jobs: - query-regression: + test-tooling: + # Stdlib unittests for the workflow Python (case selection, report + # helpers, nightly SHA picking, and rendered ECS user-data). They do + # not talk to Aliyun or the Actions runner process; running them on + # ubuntu-latest fails fast before any ECS spend. Not gated on the + # regression labels: those labels boot a VM, these tests should not. + # Ordinary PRs also run the same tests from checks.yml, because this + # workflow only starts on `labeled` (or dispatch / workflow_call). + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Test query regression tooling + run: | + python3 tests/perf/test_query_regression_runner_compaction_toctou.py + python3 tests/perf/test_query_regression_runner_otlp_trace_load.py + python3 tests/perf/test_query_regression_summary_otlp.py + python3 tests/perf/test_query_regression_case_selection.py + python3 tests/perf/test_query_regression_nightly_refs.py + python3 tests/perf/test_aliyun_ecs_runner_scripts.py + + provision: + # Runs when the aliyun-ecs path is selected: explicitly via the runner + # input, or for PR labels by default (a QUERY_REGRESSION_PR_RUNNER + # repository variable set to another value redirects PRs to that literal + # runner label instead). Uses trusted scripts from the PR base (or the + # dispatched ref), never from candidate code. Waits for test-tooling so + # a broken user-data template does not still create a VM. + needs: [test-tooling] if: >- - ${{ github.event_name != 'pull_request' || - (github.event_name == 'pull_request' && - !github.event.pull_request.draft && - (github.event.label.name == 'query-regression' || - github.event.label.name == 'heavy-regression')) }} - runs-on: ${{ github.event_name != 'pull_request' && inputs.runner || 'perf-regression-8-cores' }} + ${{ !failure() && !cancelled() && + ((github.event_name != 'pull_request' && inputs.runner == 'aliyun-ecs') || + (github.event_name == 'pull_request' && + !github.event.pull_request.draft && + (github.event.label.name == 'query-regression' || + github.event.label.name == 'heavy-regression') && + (vars.QUERY_REGRESSION_PR_RUNNER || 'aliyun-ecs') == 'aliyun-ecs')) }} + runs-on: ubuntu-latest + timeout-minutes: 45 + outputs: + label: ${{ steps.provision.outputs.label }} + instance_id: ${{ steps.provision.outputs.instance_id }} + runner_name: ${{ steps.provision.outputs.runner_name }} + steps: + - name: Checkout trusted provisioning scripts + uses: actions/checkout@v4 + with: + repository: ${{ github.repository }} + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.sha }} + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Provision ECS runner + id: provision + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + ALIYUN_ECS_VSWITCH_ID: ${{ vars.ALIYUN_ECS_VSWITCH_ID }} + ALIYUN_ECS_SECURITY_GROUP_ID: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }} + ALIYUN_ECS_INSTANCE_TYPE: ${{ vars.ALIYUN_ECS_INSTANCE_TYPE }} + ALIYUN_ECS_RESOURCE_GROUP_ID: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} + QUERY_REGRESSION_ECS_IMAGE_ID: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} + QUERY_REGRESSION_RUNNER_UID: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} + QUERY_REGRESSION_RUNNER_GID: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-provision.py + + query-regression: + needs: [provision, test-tooling] + # `!failure() && !cancelled()` both suppresses the implicit success() and + # lets the job run when provision was intentionally skipped because a + # literal runner label was selected; a failed or cancelled provision still + # blocks the run because no ECS runner would be waiting. + if: >- + ${{ !failure() && !cancelled() && + (github.event_name != 'pull_request' || + (github.event_name == 'pull_request' && + !github.event.pull_request.draft && + (github.event.label.name == 'query-regression' || + github.event.label.name == 'heavy-regression'))) }} + runs-on: >- + ${{ needs.provision.outputs.label || + (github.event_name != 'pull_request' && inputs.runner || + (vars.QUERY_REGRESSION_PR_RUNNER || 'aliyun-ecs')) }} timeout-minutes: 180 - concurrency: - group: query-regression-persistent-cache-v1 - queue: max - cancel-in-progress: false env: CARGO_PROFILE: ${{ github.event_name == 'pull_request' && 'nightly' || inputs.cargo_profile }} CARGO_HOME: /home/runner/.cargo @@ -105,7 +193,8 @@ jobs: QUERY_REGRESSION_CACHE_META: /home/runner/query-regression-cache-meta RUSTC_WRAPPER: /usr/local/bin/sccache SCCACHE_DIR: /home/runner/.cache/sccache - SCCACHE_CACHE_SIZE: 40G + # Caps the local sccache on the system disk next to target and cargo. + SCCACHE_CACHE_SIZE: 10G CARGO_INCREMENTAL: "0" RUSTFLAGS: -D warnings -C link-arg=-fuse-ld=mold QUERY_REGRESSION_CACHE_EPOCH: "1" @@ -114,7 +203,55 @@ jobs: EVENT_MERGE_SHA: ${{ github.event_name == 'pull_request' && github.sha || '' }} EVENT_HEAD_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || '' }} EVENT_BASE_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || '' }} + # Runner identity contract. The ECS image uses 1001; a manually prepared + # host may override via repo variables when 1001 is already taken. + EXPECTED_RUNNER_UID: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} + EXPECTED_RUNNER_GID: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} steps: + - name: Report provisioned ECS runner + if: ${{ needs.provision.outputs.instance_id != '' }} + shell: bash + run: | + set -euo pipefail + { + printf -- '- ECS instance: `%s` (type `%s`, image `%s`, region `%s`)\n' \ + "${{ needs.provision.outputs.instance_id }}" \ + "${{ vars.ALIYUN_ECS_INSTANCE_TYPE }}" \ + "${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }}" \ + "${{ vars.ALIYUN_ECS_REGION_ID }}" + } >> "${GITHUB_STEP_SUMMARY}" + + - name: Start machine telemetry sampler + shell: bash + # The ECS instance is deleted when the run ends, so machine state must + # be captured while the job runs. A background sampler appends load, + # memory, disk, and top-memory processes to a log every 30s; the + # "Dump machine telemetry" step (if: always()) prints and uploads it. + # Note: when the job is *cancelled*, even always() steps are killed — + # telemetry covers failures, not cancellations. + run: | + set -euo pipefail + log="${GITHUB_WORKSPACE}/machine-telemetry.log" + { + echo "== baseline $(date -u +%Y-%m-%dT%H:%M:%SZ) ==" + nproc + free -h + df -h / /home/runner + } >> "${log}" + nohup bash -c " + while true; do + { + date -u '+== %Y-%m-%dT%H:%M:%SZ ==' + uptime + free -m + df -h / /home/runner + ps -eo pid,comm,%mem,%cpu --sort=-%mem | head -6 + } >> '${log}' 2>&1 + sleep 30 + done + " >/dev/null 2>&1 & + echo "Telemetry sampler started (pid $!), logging to ${log}" + - name: Checkout base source uses: actions/checkout@v4 with: @@ -263,14 +400,64 @@ jobs: shell: bash run: | set -euo pipefail - [[ "$(id -u)" == "1001" ]] - [[ "$(id -g)" == "1001" ]] - [[ "${UV_CACHE_DIR}" == "/home/runner/.cargo/uv-cache" ]] - [[ "$(protoc --version)" == "libprotoc 3.21.12" ]] + errors=() + record() { + printf 'FAIL %s\n' "$*" + errors+=("$*") + } + require_eq() { + local name="$1" actual="$2" expected="$3" + printf 'check %s: %s\n' "${name}" "${actual}" + [[ "${actual}" == "${expected}" ]] || record "${name}: expected '${expected}', got '${actual}'" + } + require_match() { + local name="$1" actual="$2" pattern="$3" + printf 'check %s: %s\n' "${name}" "${actual}" + [[ "${actual}" =~ $pattern ]] || record "${name}: expected to match ${pattern}, got '${actual}'" + } + require() { + local name="$1" + shift + printf 'check %s\n' "${name}" + "$@" || record "${name}" + } + capture() { + local out + if out="$("$@" 2>&1)"; then + printf '%s' "${out}" + else + printf '' "$*" + fi + } + + printf 'uid=%s gid=%s PATH=%s CARGO_HOME=%s\n' "$(id -u)" "$(id -g)" "${PATH}" "${CARGO_HOME}" + ls -la "${CARGO_HOME}" 2>&1 || printf '(CARGO_HOME missing)\n' + + # Image/host hygiene: CARGO_HOME must be empty of config and + # install-state before cargo/rustup run. Those tools create + # .package-cache / bin on first use, so this cannot come after. + cargo_home_dirty=() + for entry in config config.toml credentials credentials.toml bin .crates.toml .crates2.json .global-cache .package-cache; do + if [[ -e "${CARGO_HOME}/${entry}" ]]; then + cargo_home_dirty+=("${entry}") + record "CARGO_HOME must not contain ${entry}" + fi + done + if (( ${#cargo_home_dirty[@]} > 0 )); then + printf 'CARGO_HOME contents:\n' >&2 + ls -la "${CARGO_HOME}" >&2 || true + fi + + require_eq uid "$(id -u)" "${EXPECTED_RUNNER_UID}" + require_eq gid "$(id -g)" "${EXPECTED_RUNNER_GID}" + require_eq UV_CACHE_DIR "${UV_CACHE_DIR}" "/home/runner/.cargo/uv-cache" + require_eq protoc "$(capture protoc --version)" "libprotoc 3.21.12" temporary_proto_dir="$(mktemp --directory)" trap 'rm -rf "${temporary_proto_dir}"' EXIT - test -r /usr/include/google/protobuf/any.proto - test -r /usr/include/google/protobuf/empty.proto + require "readable /usr/include/google/protobuf/any.proto" \ + test -r /usr/include/google/protobuf/any.proto + require "readable /usr/include/google/protobuf/empty.proto" \ + test -r /usr/include/google/protobuf/empty.proto printf '%s\n' \ 'syntax = "proto3";' \ 'package smoke;' \ @@ -278,47 +465,51 @@ jobs: 'import "google/protobuf/empty.proto";' \ 'message Smoke { google.protobuf.Any any = 1; google.protobuf.Empty empty = 2; }' \ > "${temporary_proto_dir}/smoke.proto" - protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ + if protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ --descriptor_set_out="${temporary_proto_dir}/smoke.pb" \ - "${temporary_proto_dir}/smoke.proto" - test -s "${temporary_proto_dir}/smoke.pb" - [[ "$(uv --version)" =~ ^uv[[:space:]]0\.11\.26([[:space:]]|$) ]] - mold_version="$(mold --version)" - [[ "${mold_version}" =~ ^mold[[:space:]]2\.30\.0([[:space:]]|$) ]] - [[ "$(python3 --version)" == "Python 3.12.3" ]] + "${temporary_proto_dir}/smoke.proto"; then + require "protoc smoke descriptor is non-empty" test -s "${temporary_proto_dir}/smoke.pb" + else + record "protoc smoke compile failed" + fi + require_match uv "$(capture uv --version)" '^uv[[:space:]]0\.11\.26([[:space:]]|$)' + require_match mold "$(capture mold --version)" '^mold[[:space:]]2\.40\.4([[:space:]]|$)' + require_eq python3 "$(capture python3 --version)" "Python 3.14.4" otelgen_path="$(command -v otelgen || true)" otelgen_version="" if [[ -n "${otelgen_path}" ]]; then - otelgen_version="$("${otelgen_path}" --version 2>&1 || true)" + otelgen_version="$(capture "${otelgen_path}" --version)" fi - printf 'otelgen path: %s\n' "${otelgen_path:-}" - printf 'otelgen version: %s\n' "${otelgen_version}" - [[ "${otelgen_path}" == "/usr/local/bin/otelgen" ]] - [[ "${otelgen_version}" == *"863a3f395d062c7322cc1de08a38774b7fdaa6c8"* ]] - sccache_version="$(sccache --version)" - [[ "${sccache_version}" =~ ^sccache[[:space:]]0\.16\.0([[:space:]]|$) ]] - [[ "$(command -v rustup)" == "/opt/cargo/bin/rustup" ]] - [[ "$(command -v cargo)" == "/opt/cargo/bin/cargo" ]] - [[ "$(command -v rustc)" == "/opt/cargo/bin/rustc" ]] - [[ "$(rustup --version)" =~ ^rustup[[:space:]]1\.29\.0([[:space:]]|$) ]] - cargo_version="$(cargo --version)" - [[ "${cargo_version}" =~ ^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] - rustc_version="$(rustc --version)" - [[ "${rustc_version}" =~ ^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] - active_toolchain="$(rustup show active-toolchain)" - [[ "${active_toolchain}" =~ ^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$) ]] - [[ "${RUSTUP_HOME}" == "/opt/rustup" ]] - [[ "${RUSTUP_TOOLCHAIN}" == "nightly-2026-03-21" ]] - [[ "${RUSTUP_AUTO_INSTALL}" == "0" ]] - test -r /opt/rustup && test -x /opt/rustup - test ! -w /opt/rustup - test ! -w /opt/cargo/bin - for entry in config config.toml credentials credentials.toml bin .crates.toml .crates2.json .global-cache .package-cache; do - test ! -e "${CARGO_HOME}/${entry}" - done + require_eq otelgen_path "${otelgen_path}" "/usr/local/bin/otelgen" + printf 'check otelgen_version: %s\n' "${otelgen_version}" + [[ "${otelgen_version}" == *"863a3f395d062c7322cc1de08a38774b7fdaa6c8"* ]] \ + || record "otelgen_version: expected commit 863a3f395d062c7322cc1de08a38774b7fdaa6c8, got '${otelgen_version}'" + require_match sccache "$(capture sccache --version)" '^sccache[[:space:]]0\.16\.0([[:space:]]|$)' + require_eq rustup_path "$(command -v rustup || true)" "/opt/cargo/bin/rustup" + require_eq cargo_path "$(command -v cargo || true)" "/opt/cargo/bin/cargo" + require_eq rustc_path "$(command -v rustc || true)" "/opt/cargo/bin/rustc" + require_match rustup "$(capture rustup --version)" '^rustup[[:space:]]1\.29\.0([[:space:]]|$)' + require_match cargo "$(capture cargo --version)" \ + '^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' + require_match rustc "$(capture rustc --version)" \ + '^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' + require_match active_toolchain "$(capture rustup show active-toolchain)" \ + '^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$)' + require_eq RUSTUP_HOME "${RUSTUP_HOME}" "/opt/rustup" + require_eq RUSTUP_TOOLCHAIN "${RUSTUP_TOOLCHAIN}" "nightly-2026-03-21" + require_eq RUSTUP_AUTO_INSTALL "${RUSTUP_AUTO_INSTALL}" "0" + require "readable /opt/rustup" test -r /opt/rustup + require "executable /opt/rustup" test -x /opt/rustup + require "runner cannot write /opt/rustup" test ! -w /opt/rustup + require "runner cannot write /opt/cargo/bin" test ! -w /opt/cargo/bin mkdir -p "${CARGO_HOME}/registry" "${CARGO_HOME}/git" + if (( ${#errors[@]} > 0 )); then + printf '\nVerify runner image tools failed (%d checks):\n' "${#errors[@]}" >&2 + printf ' - %s\n' "${errors[@]}" >&2 + exit 1 + fi - - name: Prepare persistent query regression cache + - name: Prepare query regression cache shell: bash working-directory: src run: | @@ -362,24 +553,6 @@ jobs: find "${root}" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + } - clear_cargo_extracted_trees() { - require_expected_root CARGO_REGISTRY "${CARGO_HOME}/registry" "${EXPECTED_CARGO_REGISTRY}" - require_expected_root CARGO_GIT "${CARGO_HOME}/git" "${EXPECTED_CARGO_GIT}" - rm -rf -- "${CARGO_HOME}/registry/src" "${CARGO_HOME}/git/checkouts" - } - - cargo_size_kib() { - du -sk -- "${CARGO_HOME}/registry" "${CARGO_HOME}/git" | awk '{ total += $1 } END { print total }' - } - - target_size_kib() { - du -sk -- "${CARGO_TARGET_DIR}" | cut -f1 - } - - free_kib() { - df -Pk "${CARGO_TARGET_DIR}" | awk 'NR == 2 { print $4 }' - } - report_cache_usage() { du -sh -- "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ "${RUSTUP_HOME}" "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}" @@ -400,7 +573,7 @@ jobs: printf 'Refusing unexpected RUSTC_WRAPPER: %s\n' "${RUSTC_WRAPPER}" >&2 exit 1 } - [[ "${SCCACHE_CACHE_SIZE}" == "40G" ]] || { + [[ "${SCCACHE_CACHE_SIZE}" == "10G" ]] || { printf 'Refusing unexpected SCCACHE_CACHE_SIZE: %s\n' "${SCCACHE_CACHE_SIZE}" >&2 exit 1 } @@ -487,46 +660,10 @@ jobs: report_cache_usage - target_size="$(target_size_kib)" - if (( target_size >= 400 * 1024 * 1024 )); then - printf 'Warning: target cache is at least 400 GiB (%s KiB)\n' "${target_size}" >&2 - fi - if (( target_size >= 450 * 1024 * 1024 )); then - printf 'Target cache is at least 450 GiB; clearing complete target root\n' >&2 - clear_directory "${CARGO_TARGET_DIR}" - fi - - cargo_size="$(cargo_size_kib)" - if (( cargo_size >= 60 * 1024 * 1024 )); then - printf 'Warning: Cargo cache is at least 60 GiB (%s KiB)\n' "${cargo_size}" >&2 - fi - if (( cargo_size >= 80 * 1024 * 1024 )); then - printf 'Cargo cache is at least 80 GiB; removing extracted sources and checkouts\n' >&2 - clear_cargo_extracted_trees - cargo_size="$(cargo_size_kib)" - if (( cargo_size >= 80 * 1024 * 1024 )); then - printf 'Cargo cache remains at least 80 GiB after cleanup (%s KiB)\n' "${cargo_size}" >&2 - exit 1 - fi - fi - - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space is below 300 GiB; clearing complete target root\n' >&2 - clear_directory "${CARGO_TARGET_DIR}" - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space remains below 300 GiB; removing Cargo extracted sources and checkouts\n' >&2 - clear_cargo_extracted_trees - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space remains below 300 GiB after cleanup (%s KiB)\n' "${free_space}" >&2 - exit 1 - fi - fi - fi - - report_cache_usage + # Every run starts from a fresh system disk, so size/free-space + # watermarks from the retired retained-disk era are not restored. + # A run that overflows the disk fails the build outright, which the + # telemetry step makes diagnosable. sccache --start-server sccache --zero-stats @@ -567,14 +704,6 @@ jobs: git reset --hard "${VERIFIED_CANDIDATE_SHA}" git clean -ffdx - - name: Test query regression tooling - working-directory: src - run: | - uv run --no-project python tests/perf/test_query_regression_runner_compaction_toctou.py - uv run --no-project python tests/perf/test_query_regression_runner_otlp_trace_load.py - uv run --no-project python tests/perf/test_query_regression_summary_otlp.py - uv run --no-project python tests/perf/test_query_regression_case_selection.py - - name: Build candidate greptime and query regression helpers working-directory: src run: | @@ -643,6 +772,7 @@ jobs: query-regression-work/**/logs/** query-regression-work/**/otelgen/** query-regression-summary.md + machine-telemetry.log if-no-files-found: warn retention-days: 7 @@ -657,7 +787,7 @@ jobs: if-no-files-found: warn retention-days: 7 - - name: Report persistent cache usage + - name: Report cache usage if: ${{ always() }} shell: bash run: | @@ -697,6 +827,55 @@ jobs: printf 'sccache is unavailable (report only)\n' >&2 fi + - name: Dump machine telemetry + if: ${{ always() }} + shell: bash + run: | + log="${GITHUB_WORKSPACE}/machine-telemetry.log" + if [[ -f "${log}" ]]; then + echo "::group::Machine telemetry (last 200 lines)" + tail -n 200 "${log}" + echo "::endgroup::" + else + echo "No telemetry log found (sampler never started?)" + fi + echo "::group::dmesg tail (OOM killer records)" + sudo dmesg -T 2>/dev/null | tail -n 50 || dmesg -T 2>/dev/null | tail -n 50 || \ + echo "dmesg unavailable without root" + echo "::endgroup::" + - name: Fail on regression failure if: ${{ steps.run.outputs.status != '0' }} run: exit 1 + + teardown: + # Releases the dynamically provisioned ECS runner. Runs even when the + # benchmark job fails or is cancelled; skipped when a literal runner label + # was selected because the provision outputs are empty, or when the + # dispatch set keep_instance to preserve the machine for post-mortem + # debugging (the janitor sweep still reclaims it after the TTL). + if: ${{ always() && needs.provision.outputs.instance_id != '' && !inputs.keep_instance }} + needs: [provision, query-regression] + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Checkout trusted teardown scripts + uses: actions/checkout@v4 + with: + repository: ${{ github.repository }} + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.sha }} + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Teardown ECS runner + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + QUERY_REGRESSION_ECS_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + QUERY_REGRESSION_ECS_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-teardown.py diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 52d62299a07..0a9c39c9034 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -510,7 +510,7 @@ jobs: base_ref: ${{ needs.prepare-release-validation.outputs.previous-release-tag }} candidate_ref: ${{ needs.prepare-release-validation.outputs.candidate-ref }} cargo_profile: nightly - runner: perf-regression-8-cores + runner: aliyun-ecs release-images-to-dockerhub: name: Build and push images to DockerHub diff --git a/tests/perf/AGENTS.md b/tests/perf/AGENTS.md index 1f13dea58dc..822b99080f9 100644 --- a/tests/perf/AGENTS.md +++ b/tests/perf/AGENTS.md @@ -3,9 +3,23 @@ - Keep GitHub Actions YAML thin. Put non-trivial control flow, case expansion, report generation, and metadata writing in scripts under `.github/scripts/`; workflow steps should mostly invoke those scripts. +- Runner lifecycle: the default path provisions one ephemeral Aliyun ECS + instance per run via `.github/scripts/aliyun-ecs-runner-provision.py` and + always releases it via `aliyun-ecs-runner-teardown.py`; a scheduled janitor + workflow sweeps leftovers. Build caches live on that instance's system disk + and are discarded with the VM. Runs do not share a workflow concurrency + group. The ECS custom image is built from the runner Dockerfile by + `.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py`; + keep the Dockerfile the single source of the tool contract. Dispatching with + any other `runner` value treats it as a literal self-hosted runner label + (see `ecs-image/bootstrap-runner-host.sh` for preparing such a host). - Query regression PR runs should build base/candidate binaries once, then run the default case set. Do not hard-code a single case such as `promql_pushdown_7913` into the workflow path. +- Scheduled nightly comparison lives in `query-regression-nightly.yml`: it + waits for a successful Nightly Build, then calls `query-regression.yml` + with the previous vs current nightly SHAs. Keep SHA selection in + `.github/scripts/query-regression-nightly-refs.py`. - The case DSL is not required to keep compatibility inside this PR. When the DSL changes, update TOML cases, the outer lifecycle script, Rust helpers, and docs together. @@ -19,8 +33,10 @@ - Keep the direct-SST generator generic. Issue-specific behavior belongs in case files and thresholds, not in Rust generator logic. - Before pushing perf harness changes, run at least: - - the Python tests in the `Test query regression tooling` step of - `.github/workflows/query-regression.yml` + - the Python tests in the `test-tooling` job of + `.github/workflows/query-regression.yml` (ubuntu-latest, not the ECS runner). + The Checks workflow runs the same tests on ordinary PRs so they are not + gated on the `query-regression` / `heavy-regression` labels. - `cargo fmt --all -- --check` - `cargo build -p cmd --bin query_perf_fixture --features dev-tools` - `cargo build -p cmd --bin query_regression_runner --features dev-tools` diff --git a/tests/perf/README.md b/tests/perf/README.md index cbe03b51794..47495ca531f 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -255,11 +255,12 @@ gh workflow run query-regression.yml \ -f candidate_ref= \ -f cargo_profile=nightly \ -f http_timeout=300 \ - -f runner=perf-regression-8-cores + -f runner=aliyun-ecs ``` -The selected ARC scale set must already be deployed with the runner-image -digest built from the current query-regression Dockerfile. +The `aliyun-ecs` path provisions a fresh ECS instance per run from the custom +image built from the current query-regression Dockerfile; see +`.github/runner-scale-sets/query-regression/README.md` for its configuration. ## Generator contract diff --git a/tests/perf/test_aliyun_ecs_runner_scripts.py b/tests/perf/test_aliyun_ecs_runner_scripts.py new file mode 100644 index 00000000000..581fc88abd0 --- /dev/null +++ b/tests/perf/test_aliyun_ecs_runner_scripts.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Coverage for the pure parts of the Aliyun ECS runner provision/teardown scripts.""" + +import base64 +import importlib.util +import sys +import unittest +from datetime import datetime, timedelta, timezone +from pathlib import Path + + +SCRIPTS_DIR = Path(__file__).parents[2] / ".github/scripts" + + +def load_module(name: str, filename: str): + spec = importlib.util.spec_from_file_location(name, SCRIPTS_DIR / filename) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +provision = load_module("aliyun_ecs_runner_provision_under_test", "aliyun-ecs-runner-provision.py") +teardown = load_module("aliyun_ecs_runner_teardown_under_test", "aliyun-ecs-runner-teardown.py") + + +class ProvisionNamingTest(unittest.TestCase): + def test_runner_name_and_label_derive_from_run_id(self) -> None: + self.assertEqual(provision.runner_name_for_run("12345"), "qreg-ecs-12345") + self.assertEqual(provision.runner_label_for_run("12345"), "query-regression-ecs-12345") + + +class ProvisionUserDataTest(unittest.TestCase): + def render(self) -> str: + return provision.render_user_data( + runner_name="qreg-ecs-12345", + runner_label="query-regression-ecs-12345", + runner_token="TOKEN", + repo="GreptimeTeam/greptimedb", + ) + + def test_user_data_creates_cache_paths_on_the_system_disk(self) -> None: + script = self.render() + self.assertNotIn("DISK_SERIAL", script) + self.assertNotIn("mount --bind", script) + self.assertNotIn("mkfs.ext4", script) + for destination in provision.CACHE_PATHS: + self.assertIn(f'"{destination}"', script) + + def test_user_data_wires_runner_registration(self) -> None: + script = self.render() + self.assertIn("RUNNER_NAME=qreg-ecs-12345", script) + self.assertIn("RUNNER_LABELS=query-regression-ecs-12345", script) + self.assertIn("RUNNER_TOKEN=TOKEN", script) + self.assertIn("REPO_URL=https://github.com/GreptimeTeam/greptimedb", script) + self.assertIn("PATH=/opt/cargo/bin:", script) + self.assertIn("systemctl restart --no-block ephemeral-github-runner.service", script) + + def test_encode_user_data_round_trips(self) -> None: + script = self.render() + self.assertEqual( + base64.b64decode(provision.encode_user_data(script)).decode("utf-8"), script + ) + + def test_user_data_enables_swap_and_masks_oomd(self) -> None: + script = self.render() + self.assertIn("systemctl mask systemd-oomd.socket systemd-oomd.service", script) + self.assertIn("systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer", script) + self.assertIn('APT::Periodic::Unattended-Upgrade "0"', script) + self.assertIn(f'fallocate --length {provision.SWAP_SIZE_GIB}G "{provision.SWAP_FILE}"', script) + self.assertIn(f'swapon "{provision.SWAP_FILE}"', script) + self.assertIn("sysctl --write vm.swappiness=10", script) + self.assertIn("OOMPolicy=continue", script) + self.assertNotIn("OOMScoreAdjust", script) + self.assertLess(script.index("swapon"), script.index("systemctl restart --no-block ephemeral-github-runner.service")) + + +class TeardownExpiryTest(unittest.TestCase): + NOW = datetime(2026, 8, 17, 6, 0, tzinfo=timezone.utc) + TTL = timedelta(hours=4) + + def test_parse_creation_time_formats(self) -> None: + self.assertEqual( + teardown.parse_creation_time("2026-08-17T01:02:03Z"), + datetime(2026, 8, 17, 1, 2, 3, tzinfo=timezone.utc), + ) + self.assertEqual( + teardown.parse_creation_time("2026-08-17T01:02Z"), + datetime(2026, 8, 17, 1, 2, tzinfo=timezone.utc), + ) + with self.assertRaises(ValueError): + teardown.parse_creation_time("not-a-time") + + def test_expired_instance_names_selects_only_old_instances(self) -> None: + instances = [ + ("i-old", "qreg-ecs-1", "2026-08-17T01:00Z"), # 5h old: expired + ("i-edge", "qreg-ecs-2", "2026-08-17T02:00Z"), # exactly TTL: expired + ("i-fresh", "qreg-ecs-3", "2026-08-17T05:30Z"), # 30m old: kept + ] + self.assertEqual( + teardown.expired_instance_names(instances, self.NOW, self.TTL), + [("i-old", "qreg-ecs-1"), ("i-edge", "qreg-ecs-2")], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/perf/test_query_regression_nightly_refs.py b/tests/perf/test_query_regression_nightly_refs.py new file mode 100644 index 00000000000..74b168bb464 --- /dev/null +++ b/tests/perf/test_query_regression_nightly_refs.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for consecutive Nightly Build SHA selection.""" + +import importlib.util +import sys +import unittest +from pathlib import Path + + +SCRIPTS_DIR = Path(__file__).parents[2] / ".github/scripts" + + +def load_module(): + spec = importlib.util.spec_from_file_location( + "query_regression_nightly_refs_under_test", + SCRIPTS_DIR / "query-regression-nightly-refs.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +refs = load_module() + + +def run(*, run_id: int, sha: str, branch: str = "main", created_at: str = "2026-08-24T00:00:00Z"): + return refs.WorkflowRun( + id=run_id, + head_sha=sha, + head_branch=branch, + html_url=f"https://github.com/example/run/{run_id}", + created_at=created_at, + conclusion="success", + event="schedule", + ) + + +class SelectNightlyRefsTest(unittest.TestCase): + def test_picks_newest_and_previous_on_same_branch(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="bbb", created_at="2026-08-21T00:00:00Z"), + run(run_id=1, sha="aaa", created_at="2026-08-20T00:00:00Z"), + ] + ) + self.assertFalse(pair.skip) + assert pair.candidate is not None and pair.base is not None + self.assertEqual(pair.candidate.head_sha, "ccc") + self.assertEqual(pair.base.head_sha, "bbb") + + def test_candidate_run_id_uses_that_run_and_the_next_older(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="bbb", created_at="2026-08-21T00:00:00Z"), + run(run_id=1, sha="aaa", created_at="2026-08-20T00:00:00Z"), + ], + candidate_run_id=2, + ) + self.assertFalse(pair.skip) + assert pair.candidate is not None and pair.base is not None + self.assertEqual(pair.candidate.id, 2) + self.assertEqual(pair.base.id, 1) + + def test_skips_when_consecutive_nightlies_share_a_sha(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=2, sha="same", created_at="2026-08-22T00:00:00Z"), + run(run_id=1, sha="same", created_at="2026-08-21T00:00:00Z"), + ] + ) + self.assertTrue(pair.skip) + self.assertIn("matches candidate", pair.reason) + + def test_skips_other_branches_when_picking_previous(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", branch="main", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="other", branch="feat", created_at="2026-08-21T12:00:00Z"), + run(run_id=1, sha="aaa", branch="main", created_at="2026-08-21T00:00:00Z"), + ] + ) + self.assertFalse(pair.skip) + assert pair.base is not None + self.assertEqual(pair.base.head_sha, "aaa") + + def test_skips_when_only_one_nightly_exists(self) -> None: + pair = refs.select_base_and_candidate( + [run(run_id=1, sha="aaa")], + ) + self.assertTrue(pair.skip) + self.assertIn("no previous", pair.reason) + + def test_explicit_refs_bypass_github(self) -> None: + pair = refs.override_pair("base-sha", "cand-sha") + self.assertFalse(pair.skip) + assert pair.base is not None and pair.candidate is not None + self.assertEqual(pair.base.head_sha, "base-sha") + self.assertEqual(pair.candidate.head_sha, "cand-sha") + + +if __name__ == "__main__": + unittest.main() From 1409e66837d89ff7f9d7d0604dfdbba3ca9fc7a7 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 26 Aug 2026 14:00:39 +0000 Subject: [PATCH 111/126] refactor(flight): add request builder and defer DoGet execution (#8953) * refactor: add flight request builder Signed-off-by: WenyXu * refactor: use flight request builder in flow Signed-off-by: WenyXu * refactor: defer frontend flight query execution Signed-off-by: WenyXu * fix(grpc): avoid cloning requests during auth Signed-off-by: WenyXu * test(grpc): cover flight request timeout Signed-off-by: WenyXu * refactor(client): share Flight message reader Signed-off-by: WenyXu * feat(flow): add Flight DoGet timeout Signed-off-by: WenyXu * refactor(client): gate Flight DDL helpers for testing Signed-off-by: WenyXu * fix(client): restore Flight stream error semantics Signed-off-by: WenyXu * docs(grpc): document Flight stream input constructors Signed-off-by: WenyXu * fix(flight): preserve deferred stream context Signed-off-by: WenyXu * fix(client): use Flight stream SNAFU context Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/client/src/database.rs | 490 ++++++++++++------ src/client/src/error.rs | 18 + src/client/src/flight.rs | 74 ++- src/client/src/region.rs | 189 +++++-- src/datanode/src/region_server.rs | 9 +- src/flow/src/batching_mode.rs | 3 + src/flow/src/batching_mode/frontend_client.rs | 61 +-- src/servers/src/grpc/flight.rs | 75 +-- src/servers/src/grpc/flight/stream.rs | 150 ++++-- src/servers/src/grpc/greptime_handler.rs | 72 +-- tests-integration/src/grpc/flight.rs | 180 ++++++- tests-integration/tests/http.rs | 1 + 12 files changed, 923 insertions(+), 399 deletions(-) diff --git a/src/client/src/database.rs b/src/client/src/database.rs index 4fff3636b63..7c33794f477 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -12,20 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::pin::Pin; use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, RwLock}; use std::task::{Context, Poll}; +use std::time::Duration; use api::v1::auth_header::AuthScheme; +#[cfg(feature = "testing")] use api::v1::ddl_request::Expr as DdlExpr; use api::v1::greptime_database_client::GreptimeDatabaseClient; use api::v1::greptime_request::Request; use api::v1::query_request::Query; +#[cfg(feature = "testing")] +use api::v1::{AlterTableExpr, CreateTableExpr, DdlRequest}; use api::v1::{ - AlterTableExpr, AuthHeader, Basic, CreateTableExpr, DdlRequest, GreptimeRequest, - InsertRequests, QueryRequest, RequestHeader, RowInsertRequests, + AuthHeader, Basic, GreptimeRequest, InsertRequests, QueryRequest, RequestHeader, + RowInsertRequests, }; use arc_swap::ArcSwapOption; use arrow_flight::{FlightData, Ticket}; @@ -34,7 +39,7 @@ use base64::Engine; use base64::prelude::BASE64_STANDARD; use common_catalog::build_db_string; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; -use common_error::ext::BoxedError; +use common_error::ext::{BoxedError, ErrorExt}; use common_grpc::flight::do_put::DoPutResponse; use common_grpc::flight::{ FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY, @@ -49,21 +54,23 @@ use common_telemetry::{error, warn}; use futures::future; use futures_util::{Stream, StreamExt, TryStreamExt}; use prost::Message; -use snafu::ResultExt; +use snafu::{IntoError, ResultExt}; use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue, MetadataMap, MetadataValue}; use tonic::transport::Channel; use crate::error::{ - ConvertFlightDataSnafu, Error, FlightGetSnafu, IllegalFlightMessagesSnafu, + ConvertFlightDataSnafu, Error, FlightGetSnafu, FlightStreamSnafu, IllegalFlightMessagesSnafu, InvalidTonicMetadataValueSnafu, }; -use crate::flight::decode_flight_data; +use crate::flight::{FlightMessageReader, decode_flight_data}; use crate::{Client, Result, error, from_grpc_response}; type FlightDataStream = Pin + Send>>; type DoPutResponseStream = Pin>>>; +const HINTS_METADATA_KEY: &str = "x-greptime-hints"; + /// Terminal metrics associated with a query output. /// /// For streaming outputs, metrics are only final after the stream is fully @@ -257,19 +264,17 @@ fn attach_terminal_metrics(output: Output, terminal_metrics: &OutputMetrics) -> } async fn output_from_flight_message_stream( - mut flight_message_stream: S, + remote_addr: String, + flight_message_stream: S, ) -> Result where S: Stream> + Send + Unpin + 'static, { - let Some(first_flight_message) = flight_message_stream.next().await else { - return IllegalFlightMessagesSnafu { - reason: "Expect the response not to be empty", - } - .fail(); - }; - - let first_flight_message = first_flight_message?; + let mut reader = FlightMessageReader::new(remote_addr, flight_message_stream); + let first_flight_message = reader + .read_first() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; match first_flight_message { FlightMessage::AffectedRows { rows, metrics } => { @@ -277,7 +282,10 @@ where if let Some(metrics) = metrics { terminal_metrics.update(Some(parse_terminal_metrics(&metrics)?)); } - let next_message = flight_message_stream.next().await.transpose()?; + let next_message = reader + .read_next() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; match next_message { None => terminal_metrics.mark_ready(), Some(FlightMessage::Metrics(s)) if terminal_metrics.get().is_none() => { @@ -316,15 +324,19 @@ where ); let schema_cloned = schema.clone(); let stream = Box::pin(stream!({ - while let Some(flight_message_item) = flight_message_stream.next().await { - let flight_message = match flight_message_item { - Ok(message) => message, - Err(e) => { - yield Err(BoxedError::new(e)).context(ExternalSnafu); + loop { + let flight_message = match reader.read_next().await { + Ok(Some(message)) => message, + Ok(None) => break, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error( + reader.remote_addr(), + error, + ))) + .context(ExternalSnafu); break; } }; - match flight_message { FlightMessage::RecordBatch(arrow_batch) => { yield Ok(RecordBatch::from_df_record_batch( @@ -371,6 +383,25 @@ where } } +fn flight_stream_error(addr: &str, error: Error) -> Error { + let tonic_code = error.tonic_code().unwrap_or(tonic::Code::Unknown); + let message = error.to_string(); + if error.status_code().should_log_error() { + error!( + error; "Failed to receive Flight data, addr: {}, code: {}", + addr, + tonic_code + ); + } + + FlightStreamSnafu { + addr: addr.to_string(), + tonic_code, + message, + } + .into_error(BoxedError::new(error)) +} + #[derive(Clone, Debug, Default)] pub struct Database { // The "catalog" and "schema" to be used in processing the requests at the server side. @@ -389,6 +420,42 @@ pub struct Database { ctx: FlightContext, } +#[derive(Default)] +struct FlightRequestOptions { + hints: Option, + flow_extensions: Option, + snapshot_seqs: Option, + timeout: Option, +} + +impl FlightRequestOptions { + fn apply_to(self, request: &mut tonic::Request) -> Result<()> { + let metadata = request.metadata_mut(); + if let Some(hints) = self.hints { + Database::put_metadata_value(metadata, HINTS_METADATA_KEY, hints)?; + } + if let Some(flow_extensions) = self.flow_extensions { + Database::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, flow_extensions)?; + } + if let Some(snapshot_seqs) = self.snapshot_seqs { + Database::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, snapshot_seqs)?; + } + if let Some(timeout) = self.timeout { + request.set_timeout(timeout); + } + Ok(()) + } +} + +/// A single Flight DoGet request to a [`Database`]. +/// +/// The builder carries request-scoped metadata and options. It does not modify +/// the underlying [`Database`], so its configuration cannot affect later RPCs. +pub struct DatabaseFlightRequest<'a> { + database: &'a Database, + options: FlightRequestOptions, +} + pub struct DatabaseClient { pub addr: String, pub inner: GreptimeDatabaseClient, @@ -483,6 +550,14 @@ impl Database { }); } + /// Creates a builder for a single Flight DoGet request. + pub fn flight_request(&self) -> DatabaseFlightRequest<'_> { + DatabaseFlightRequest { + database: self, + options: FlightRequestOptions::default(), + } + } + /// Make an InsertRequests request to the database. pub async fn insert(&self, requests: InsertRequests) -> Result { self.handle(Request::Inserts(requests)).await @@ -538,44 +613,31 @@ impl Database { } fn put_hints(metadata: &mut MetadataMap, hints: &[(&str, &str)]) -> Result<()> { - let Some(value) = hints - .iter() - .map(|(k, v)| format!("{}={}", k, v)) - .reduce(|a, b| format!("{},{}", a, b)) - else { + let Some(value) = Self::encode_hints(hints) else { return Ok(()); }; - let key = AsciiMetadataKey::from_static("x-greptime-hints"); - let value = AsciiMetadataValue::from_str(&value).context(InvalidTonicMetadataValueSnafu)?; - metadata.insert(key, value); - Ok(()) + Self::put_metadata_value(metadata, HINTS_METADATA_KEY, value) } - fn put_flow_extensions( - metadata: &mut MetadataMap, - flow_extensions: &[(&str, &str)], - ) -> Result<()> { - if flow_extensions.is_empty() { - return Ok(()); - } - - let value = serde_json::to_string(&flow_extensions.to_vec()) - .expect("flow extension pairs should serialize"); - Self::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, value) + fn encode_hints(hints: &[(&str, &str)]) -> Option { + hints + .iter() + .map(|(k, v)| format!("{}={}", k, v)) + .reduce(|a, b| format!("{},{}", a, b)) } - fn put_snapshot_seqs( - metadata: &mut MetadataMap, - snapshot_seqs: &std::collections::HashMap, - ) -> Result<()> { - if snapshot_seqs.is_empty() { - return Ok(()); - } + fn encode_flow_extensions(flow_extensions: &[(&str, &str)]) -> Option { + (!flow_extensions.is_empty()).then(|| { + serde_json::to_string(&flow_extensions.to_vec()) + .expect("flow extension pairs should serialize") + }) + } - let value = - serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize"); - Self::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, value) + fn encode_snapshot_seqs(snapshot_seqs: &HashMap) -> Option { + (!snapshot_seqs.is_empty()).then(|| { + serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize") + }) } fn put_metadata_value( @@ -673,7 +735,7 @@ impl Database { where S: AsRef, { - self.sql_with_hint(sql, &[]).await + self.flight_request().sql(sql).await } /// Executes a SQL query with optional hints for query optimization. @@ -681,12 +743,7 @@ impl Database { where S: AsRef, { - let request = Request::Query(QueryRequest { - query: Some(Query::Sql(sql.as_ref().to_string())), - }); - self.do_get(request, hints, &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().with_hints(hints).sql(sql).await } /// Executes a SQL query and returns the output with terminal metrics. @@ -701,77 +758,33 @@ impl Database { where S: AsRef, { - self.query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::Sql(sql.as_ref().to_string())), - }, - hints, - &[], - &Default::default(), - ) - .await + self.flight_request() + .with_hints(hints) + .sql_with_terminal_metrics(sql) + .await } /// Executes a logical plan directly without SQL parsing. pub async fn logical_plan(&self, logical_plan: Vec) -> Result { - self.query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::LogicalPlan(logical_plan)), - }, - &[], - &[], - &Default::default(), - ) - .await - .map(OutputWithMetrics::into_output) - } - - /// Executes a query and carries flow extensions through Flight metadata. - /// - /// This is the lower-level terminal-metrics API for Flow callers that need - /// to pass JSON-bearing flow extensions without going through hint metadata. - pub async fn query_with_terminal_metrics_and_flow_extensions( - &self, - request: QueryRequest, - hints: &[(&str, &str)], - flow_extensions: &[(&str, &str)], - snapshot_seqs: &std::collections::HashMap, - ) -> Result { - self.do_get( - Request::Query(request), - hints, - flow_extensions, - snapshot_seqs, - ) - .await + self.flight_request().logical_plan(logical_plan).await } /// Creates a new table using the provided table expression. + #[cfg(feature = "testing")] pub async fn create(&self, expr: CreateTableExpr) -> Result { - let request = Request::Ddl(DdlRequest { - expr: Some(DdlExpr::CreateTable(expr)), - }); - self.do_get(request, &[], &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().create(expr).await } /// Alters an existing table using the provided alter expression. + #[cfg(feature = "testing")] pub async fn alter(&self, expr: AlterTableExpr) -> Result { - let request = Request::Ddl(DdlRequest { - expr: Some(DdlExpr::AlterTable(expr)), - }); - self.do_get(request, &[], &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().alter(expr).await } async fn do_get( &self, request: Request, - hints: &[(&str, &str)], - flow_extensions: &[(&str, &str)], - snapshot_seqs: &std::collections::HashMap, + options: FlightRequestOptions, ) -> Result { let request = self.to_rpc_request(request); let request = Ticket { @@ -779,12 +792,10 @@ impl Database { }; let mut request = tonic::Request::new(request); - let metadata = request.metadata_mut(); - Self::put_hints(metadata, hints)?; - Self::put_flow_extensions(metadata, flow_extensions)?; - Self::put_snapshot_seqs(metadata, snapshot_seqs)?; + options.apply_to(&mut request)?; let mut client = self.client.make_flight_client(false, false)?; + let remote_addr = client.addr().to_string(); let response = client.mut_inner().do_get(request).await.or_else(|e| { let tonic_code = e.code(); @@ -796,7 +807,7 @@ impl Database { e ); Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu { - addr: client.addr().to_string(), + addr: remote_addr.clone(), tonic_code, }) })?; @@ -808,7 +819,7 @@ impl Database { future::ready(decode_flight_data(&mut decoder, flight_data)) }); - output_from_flight_message_stream(flight_message_stream).await + output_from_flight_message_stream(remote_addr, flight_message_stream).await } /// Ingest a stream of [RecordBatch]es that belong to a table, using Arrow Flight's "`DoPut`" @@ -847,6 +858,98 @@ impl Database { } } +impl<'a> DatabaseFlightRequest<'a> { + /// Adds query optimization hints to this Flight request. + pub fn with_hints(mut self, hints: &[(&str, &str)]) -> Self { + self.options.hints = Database::encode_hints(hints); + self + } + + /// Adds Flow extensions to this Flight request. + pub fn with_flow_extensions(mut self, flow_extensions: &[(&str, &str)]) -> Self { + self.options.flow_extensions = Database::encode_flow_extensions(flow_extensions); + self + } + + /// Adds snapshot sequence fences to this Flight request. + pub fn with_snapshot_seqs(mut self, snapshot_seqs: &HashMap) -> Self { + self.options.snapshot_seqs = Database::encode_snapshot_seqs(snapshot_seqs); + self + } + + /// Sets a timeout for this Flight request only. + pub fn with_timeout(mut self, timeout: Duration) -> Self { + self.options.timeout = Some(timeout); + self + } + + /// Executes a SQL query. + pub async fn sql(self, sql: S) -> Result + where + S: AsRef, + { + let request = Request::Query(QueryRequest { + query: Some(Query::Sql(sql.as_ref().to_string())), + }); + self.do_get(request) + .await + .map(OutputWithMetrics::into_output) + } + + /// Executes a SQL query and returns terminal metrics. + pub async fn sql_with_terminal_metrics(self, sql: S) -> Result + where + S: AsRef, + { + self.query_with_terminal_metrics(QueryRequest { + query: Some(Query::Sql(sql.as_ref().to_string())), + }) + .await + } + + /// Executes a logical plan directly without SQL parsing. + pub async fn logical_plan(self, logical_plan: Vec) -> Result { + self.query_with_terminal_metrics(QueryRequest { + query: Some(Query::LogicalPlan(logical_plan)), + }) + .await + .map(OutputWithMetrics::into_output) + } + + /// Executes a query and returns terminal metrics. + pub async fn query_with_terminal_metrics( + self, + request: QueryRequest, + ) -> Result { + self.do_get(Request::Query(request)).await + } + + /// Creates a new table using the provided table expression. + #[cfg(feature = "testing")] + pub async fn create(self, expr: CreateTableExpr) -> Result { + self.do_get(Request::Ddl(DdlRequest { + expr: Some(DdlExpr::CreateTable(expr)), + })) + .await + .map(OutputWithMetrics::into_output) + } + + /// Alters an existing table using the provided alter expression. + #[cfg(feature = "testing")] + pub async fn alter(self, expr: AlterTableExpr) -> Result { + self.do_get(Request::Ddl(DdlRequest { + expr: Some(DdlExpr::AlterTable(expr)), + })) + .await + .map(OutputWithMetrics::into_output) + } + + async fn do_get(self, request: Request) -> Result { + let Self { database, options } = self; + database.do_get(request, options).await + } +} + /// by grpc standard, only `Unavailable` is retryable, see: https://github.com/grpc/grpc/blob/master/doc/statuscodes.md#status-codes-and-their-use-in-grpc pub fn is_grpc_retryable(err: &tonic::Status) -> bool { matches!(err.code(), tonic::Code::Unavailable) @@ -934,12 +1037,14 @@ mod tests { #[test] fn test_put_flow_extensions_preserves_comma_bearing_values() { let mut metadata = MetadataMap::new(); - Database::put_flow_extensions( + Database::put_metadata_value( &mut metadata, - &[ + FLOW_EXTENSIONS_METADATA_KEY, + Database::encode_flow_extensions(&[ ("flow.return_region_seq", "true"), ("flow.incremental_after_seqs", r#"{"1":10,"2":20}"#), - ], + ]) + .unwrap(), ) .unwrap(); @@ -969,7 +1074,12 @@ mod tests { (9_007_199_254_740_993_u64, 9_007_199_254_740_995_u64), ]); - Database::put_snapshot_seqs(&mut metadata, &snapshot_seqs).unwrap(); + Database::put_metadata_value( + &mut metadata, + SNAPSHOT_SEQS_METADATA_KEY, + Database::encode_snapshot_seqs(&snapshot_seqs).unwrap(), + ) + .unwrap(); let value = metadata .get(SNAPSHOT_SEQS_METADATA_KEY) @@ -980,6 +1090,50 @@ mod tests { assert_eq!(decoded, snapshot_seqs); } + #[test] + fn test_flight_request_builder_applies_request_options() { + let database = Database::new("greptime", "public", Client::default()); + let snapshot_seqs = HashMap::from([(42, 99)]); + let request = database + .flight_request() + .with_hints(&[("query_parallelism", "1")]) + .with_flow_extensions(&[("flow.return_region_seq", "true")]) + .with_snapshot_seqs(&snapshot_seqs) + .with_timeout(Duration::from_millis(50)); + let mut tonic_request = tonic::Request::new(()); + + request.options.apply_to(&mut tonic_request).unwrap(); + + let metadata = tonic_request.metadata(); + assert_eq!( + metadata.get(HINTS_METADATA_KEY).unwrap(), + "query_parallelism=1" + ); + assert_eq!( + serde_json::from_str::>( + metadata + .get(FLOW_EXTENSIONS_METADATA_KEY) + .unwrap() + .to_str() + .unwrap(), + ) + .unwrap(), + vec![("flow.return_region_seq".to_string(), "true".to_string())] + ); + assert_eq!( + serde_json::from_str::>( + metadata + .get(SNAPSHOT_SEQS_METADATA_KEY) + .unwrap() + .to_str() + .unwrap(), + ) + .unwrap(), + snapshot_seqs + ); + assert!(metadata.get("grpc-timeout").is_some()); + } + #[test] fn test_flight_ctx() { let mut ctx = FlightContext::default(); @@ -1020,6 +1174,28 @@ mod tests { assert_eq!(expected.should_retry(), actual.should_retry()); } + #[test] + fn test_flight_stream_error_preserves_addr_and_message() { + let error = flight_stream_error( + "127.0.0.1:4001", + Status::out_of_range("message length too large").into(), + ); + + assert!(matches!( + &error, + Error::FlightStream { + addr, + tonic_code: Code::OutOfRange, + message, + .. + } if addr == "127.0.0.1:4001" && message == "message length too large" + )); + assert_eq!( + "Failed to receive Flight data from 127.0.0.1:4001, code: Operation was attempted past the valid range: message length too large", + error.to_string(), + ); + } + #[test] fn test_from_tonic_status_with_retry_hint() { let mut headers = HeaderMap::new(); @@ -1117,13 +1293,13 @@ mod tests { #[tokio::test] async fn test_affected_rows_inline_metrics_are_parsed() { - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![Ok( - FlightMessage::AffectedRows { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![Ok(FlightMessage::AffectedRows { rows: 3, metrics: Some(terminal_metrics_json()), - }, - )] - as Vec>)) + })] as Vec>), + ) .await .unwrap(); @@ -1138,14 +1314,16 @@ mod tests { #[tokio::test] async fn test_affected_rows_inline_metrics_rejects_trailing_metrics() { let metrics_json = terminal_metrics_json(); - let err = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::AffectedRows { - rows: 3, - metrics: Some(metrics_json.clone()), - }), - Ok(FlightMessage::Metrics(metrics_json)), - ] - as Vec>)) + let err = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: Some(metrics_json.clone()), + }), + Ok(FlightMessage::Metrics(metrics_json)), + ] as Vec>), + ) .await .unwrap_err(); @@ -1167,12 +1345,14 @@ mod tests { vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], ) .unwrap(); - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), - Ok(FlightMessage::Metrics("{not-json}".to_string())), - ] - as Vec>)) + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + Ok(FlightMessage::Metrics("{not-json}".to_string())), + ] as Vec>), + ) .await .unwrap(); let terminal_metrics = output.metrics.clone(); @@ -1211,18 +1391,20 @@ mod tests { vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef], ) .unwrap(); - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::RecordBatch( - first_batch.into_df_record_batch(), - )), - Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(1))), - Ok(FlightMessage::RecordBatch( - second_batch.into_df_record_batch(), - )), - Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(2))), - ] - as Vec>)) + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch( + first_batch.into_df_record_batch(), + )), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(1))), + Ok(FlightMessage::RecordBatch( + second_batch.into_df_record_batch(), + )), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(2))), + ] as Vec>), + ) .await .unwrap(); let terminal_metrics = output.metrics.clone(); diff --git a/src/client/src/error.rs b/src/client/src/error.rs index 1db0858f087..168ba422715 100644 --- a/src/client/src/error.rs +++ b/src/client/src/error.rs @@ -40,6 +40,21 @@ pub enum Error { source: BoxedError, }, + #[snafu(display( + "Failed to receive Flight data from {}, code: {}: {}", + addr, + tonic_code, + message + ))] + FlightStream { + addr: String, + tonic_code: Code, + message: String, + source: BoxedError, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to convert FlightData"))] ConvertFlightData { #[snafu(implicit)] @@ -154,6 +169,7 @@ impl ErrorExt for Error { Error::Server { code, .. } | Error::Tonic { code, .. } => *code, Error::FlightGet { source, .. } + | Error::FlightStream { source, .. } | Error::RegionServer { source, .. } | Error::FlowServer { source, .. } => source.status_code(), Error::CreateChannel { source, .. } @@ -174,6 +190,7 @@ impl ErrorExt for Error { match self { Error::Tonic { retry_hint, .. } => *retry_hint, Error::FlightGet { source, .. } + | Error::FlightStream { source, .. } | Error::RegionServer { source, .. } | Error::FlowServer { source, .. } | Error::External { source, .. } => source.retry_hint(), @@ -193,6 +210,7 @@ impl Error { pub fn tonic_code(&self) -> Option { match self { Self::FlightGet { tonic_code, .. } + | Self::FlightStream { tonic_code, .. } | Self::RegionServer { code: tonic_code, .. } diff --git a/src/client/src/flight.rs b/src/client/src/flight.rs index 4d00cba962e..1753fdd5246 100644 --- a/src/client/src/flight.rs +++ b/src/client/src/flight.rs @@ -12,12 +12,82 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::pin::Pin; + use arrow_flight::FlightData; use common_grpc::flight::{FlightDecoder, FlightMessage}; -use snafu::ResultExt; +use futures_util::stream::Peekable; +use futures_util::{Stream, StreamExt}; +use snafu::{OptionExt, ResultExt}; use crate::Result; -use crate::error::{ConvertFlightDataSnafu, Error}; +use crate::error::{ConvertFlightDataSnafu, Error, IllegalFlightMessagesSnafu}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FlightMessageKind { + Schema, + RecordBatch, + AffectedRows, + Metrics, +} + +impl From<&FlightMessage> for FlightMessageKind { + fn from(message: &FlightMessage) -> Self { + match message { + FlightMessage::Schema(_) => Self::Schema, + FlightMessage::RecordBatch(_) => Self::RecordBatch, + FlightMessage::AffectedRows { .. } => Self::AffectedRows, + FlightMessage::Metrics(_) => Self::Metrics, + } + } +} + +pub(crate) struct FlightMessageReader { + /// Remote Flight peer associated with this response stream. + remote_addr: String, + messages: Peekable, +} + +impl FlightMessageReader +where + S: Stream> + Unpin, +{ + pub(crate) fn new(remote_addr: impl Into, messages: S) -> Self { + Self { + remote_addr: remote_addr.into(), + messages: messages.peekable(), + } + } + + pub(crate) fn remote_addr(&self) -> &str { + &self.remote_addr + } + + pub(crate) async fn read_first(&mut self) -> Result { + self.read_next().await?.context(IllegalFlightMessagesSnafu { + reason: "Expect the response not to be empty", + }) + } + + pub(crate) async fn read_next(&mut self) -> Result> { + self.messages.next().await.transpose() + } + + pub(crate) async fn peek_next_message_kind(&mut self) -> Result> { + match Pin::new(&mut self.messages).peek().await { + Some(Ok(message)) => Ok(Some(message.into())), + None => Ok(None), + Some(Err(_)) => match self.read_next().await { + // `peek` only borrows the error; consume it to preserve the source error. + Err(error) => Err(error), + Ok(_) => IllegalFlightMessagesSnafu { + reason: "Flight stream changed after peek".to_string(), + } + .fail(), + }, + } + } +} pub(crate) fn decode_flight_data( decoder: &mut FlightDecoder, diff --git a/src/client/src/region.rs b/src/client/src/region.rs index 2bb5f7dd0a6..cd155810ea9 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -47,7 +47,7 @@ use crate::error::{ self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu, MissingFieldSnafu, Result, ServerSnafu, }; -use crate::flight::decode_flight_data; +use crate::flight::{FlightMessageKind, FlightMessageReader, decode_flight_data}; use crate::{Client, metrics}; const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); @@ -203,26 +203,22 @@ impl RegionRequester { async fn recordbatches_from_flight_message_stream( addr: String, - mut flight_message_stream: S, + flight_message_stream: S, ) -> Result where S: Stream> + Send + Unpin + 'static, { - let Some(first_flight_message) = flight_message_stream.next().await else { - return IllegalFlightMessagesSnafu { - reason: "Expect the response not to be empty", - } - .fail() - .map_err(|error| flight_stream_error(&addr, error)); - }; - let FlightMessage::Schema(schema) = - first_flight_message.map_err(|e| flight_stream_error(&addr, e))? + let mut reader = FlightMessageReader::new(addr.clone(), flight_message_stream); + let FlightMessage::Schema(schema) = reader + .read_first() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))? else { return IllegalFlightMessagesSnafu { reason: "Expect schema to be the first flight message", } .fail() - .map_err(|error| flight_stream_error(&addr, error)); + .map_err(|error| flight_stream_error(reader.remote_addr(), error)); }; let metrics = Arc::new(ArcSwapOption::from(None)); @@ -233,31 +229,23 @@ where let schema = Arc::new(datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?); let schema_cloned = schema.clone(); - let stream_addr = addr.clone(); + let stream_addr = addr; let stream = Box::pin(stream!({ let _span = tracing_context.attach(common_telemetry::tracing::info_span!( "poll_flight_data_stream" )); - let mut buffered_message: Option = None; let mut stream_ended = false; while !stream_ended { - // get the next message from the buffered message or read from the flight message stream - let flight_message_item = if let Some(msg) = buffered_message.take() { - Some(Ok(msg)) - } else { - flight_message_stream.next().await - }; - - let flight_message = match flight_message_item { - Some(Ok(message)) => message, - Some(Err(e)) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + let flight_message = match reader.read_next().await { + Ok(Some(message)) => message, + Ok(None) => break, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) .context(ExternalSnafu); break; } - None => break, }; match flight_message { @@ -265,42 +253,54 @@ where let result_to_yield = RecordBatch::from_df_record_batch(schema_cloned.clone(), record_batch); - // get the next message from the stream. normally it should be a metrics message. - if let Some(next_flight_message_result) = flight_message_stream.next().await { - match next_flight_message_result { - Ok(FlightMessage::Metrics(s)) => { - let m = serde_json::from_str(&s).ok().map(Arc::new); - metrics_ref.swap(m); - } - Ok(FlightMessage::RecordBatch(rb)) => { - // for some reason it's not a metrics message, so we need to buffer this record batch - // and yield it in the next iteration. - buffered_message = Some(FlightMessage::RecordBatch(rb)); - } - Ok(_) => { - yield IllegalFlightMessagesSnafu { - reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Err(e) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + // Metrics follow a batch so MergeScan can observe them before yielding it. + match reader.peek_next_message_kind().await { + Ok(Some(FlightMessageKind::Metrics)) => { + let metrics_message = match reader.read_next().await { + Ok(Some(FlightMessage::Metrics(metrics))) => metrics, + Ok(Some(_) | None) => { + yield IllegalFlightMessagesSnafu { + reason: "Flight stream changed after peek", + } + .fail() + .map_err(BoxedError::new) .context(ExternalSnafu); - break; + break; + } + Err(error) => { + yield Err(BoxedError::new(flight_stream_error( + &stream_addr, + error, + ))) + .context(ExternalSnafu); + break; + } + }; + let metrics = serde_json::from_str(&metrics_message).ok().map(Arc::new); + metrics_ref.swap(metrics); + } + Ok(Some(FlightMessageKind::RecordBatch)) => {} + Ok(Some(FlightMessageKind::Schema | FlightMessageKind::AffectedRows)) => { + yield IllegalFlightMessagesSnafu { + reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" } + .fail() + .map_err(BoxedError::new) + .context(ExternalSnafu); + break; + } + Ok(None) => stream_ended = true, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) + .context(ExternalSnafu); + break; } - } else { - // the stream has ended - stream_ended = true; } yield Ok(result_to_yield); } FlightMessage::Metrics(s) => { - // just a branch in case of some metrics message comes after other things. + // Metrics may arrive before the next RecordBatch. let m = serde_json::from_str(&s).ok().map(Arc::new); metrics_ref.swap(m); continue; @@ -412,6 +412,7 @@ mod test { RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, remote_dyn_filter_request, }; use common_recordbatch::adapter::RecordBatchMetrics; + use datatypes::arrow::array::Int32Array; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; @@ -598,6 +599,86 @@ mod test { assert_eq!(metrics.elapsed_compute, 7); } + #[tokio::test] + async fn test_record_batch_stream_updates_following_metrics_before_yielding_batch() { + let schema = test_schema(); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + Ok(FlightMessage::Metrics(test_metrics_json())), + ]), + ) + .await + .unwrap(); + + let batch = recordbatches.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), 1); + + let metrics = recordbatches.metrics().unwrap(); + assert_eq!(metrics.elapsed_compute, 7); + assert!(recordbatches.next().await.is_none()); + } + + #[tokio::test] + async fn test_record_batch_stream_preserves_peeked_record_batch() { + let schema = test_schema(); + let first_batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + let second_batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef], + ) + .unwrap(); + + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch( + first_batch.into_df_record_batch(), + )), + Ok(FlightMessage::RecordBatch( + second_batch.into_df_record_batch(), + )), + ]), + ) + .await + .unwrap(); + + let first_batch = recordbatches.next().await.unwrap().unwrap(); + let second_batch = recordbatches.next().await.unwrap().unwrap(); + assert_eq!( + first_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 1 + ); + assert_eq!( + second_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2 + ); + assert!(recordbatches.next().await.is_none()); + } + #[tokio::test] async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() { let schema = test_schema(); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 0d9398b0889..79d0dd19152 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -67,7 +67,8 @@ use servers::error::{ }; use servers::grpc::FlightCompression; use servers::grpc::flight::{ - FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, TonicStream, + FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, FlightRecordBatchStreamInput, + TonicStream, }; use servers::grpc::region_server::RegionServerHandler; use session::context::{ @@ -990,7 +991,11 @@ impl FlightCraft for RegionServer { }; let stream = Box::pin(FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(initializer), + FlightRecordBatchStreamInput::initializer(async move { + initializer + .await + .map(FlightRecordBatchSource::RecordBatches) + }), tracing_context, self.flight_compression, query_ctx, diff --git a/src/flow/src/batching_mode.rs b/src/flow/src/batching_mode.rs index bb946e75740..3981446721b 100644 --- a/src/flow/src/batching_mode.rs +++ b/src/flow/src/batching_mode.rs @@ -44,6 +44,8 @@ pub struct BatchingModeOptions { /// The gRPC connection timeout #[serde(with = "humantime_serde")] pub grpc_conn_timeout: Duration, + #[serde(with = "humantime_serde")] + pub experimental_flight_do_get_timeout: Duration, /// The gRPC max retry number pub experimental_grpc_max_retries: u32, /// Flow wait for available frontend timeout, @@ -72,6 +74,7 @@ impl Default for BatchingModeOptions { slow_query_threshold: Duration::from_secs(60), experimental_min_refresh_duration: Duration::new(5, 0), grpc_conn_timeout: Duration::from_secs(5), + experimental_flight_do_get_timeout: Duration::from_secs(10), experimental_grpc_max_retries: 3, experimental_frontend_scan_timeout: Duration::from_secs(30), experimental_max_filter_num_per_query: 20, diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index aa94b91b2ae..27db96748a6 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -18,7 +18,6 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex, RwLock, Weak}; use api::v1::greptime_request::Request; -use api::v1::query_request::Query; use api::v1::{CreateTableExpr, QueryRequest}; use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputWithMetrics}; use common_error::ext::BoxedError; @@ -340,53 +339,6 @@ impl FrontendClient { }) } - /// Execute a SQL statement on the frontend. - pub async fn sql(&self, catalog: &str, schema: &str, sql: &str) -> Result { - match self { - FrontendClient::Distributed { .. } => { - let db = self.get_random_active_frontend(catalog, schema).await?; - db.database - .sql(sql) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu) - } - FrontendClient::Standalone { - database_client, .. - } => { - let ctx = QueryContextBuilder::default() - .current_catalog(catalog.to_string()) - .current_schema(schema.to_string()) - .build(); - let ctx = Arc::new(ctx); - { - let database_client = { - database_client - .handler - .lock() - .unwrap() - .as_ref() - .context(UnexpectedSnafu { - reason: "Standalone's frontend instance is not set", - })? - .upgrade() - .context(UnexpectedSnafu { - reason: "Failed to upgrade database client", - })? - }; - let req = Request::Query(QueryRequest { - query: Some(Query::Sql(sql.to_string())), - }); - database_client - .do_query(req, ctx) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu) - } - } - } - } - /// Execute a flow query and return terminal metrics. `snapshot_seqs` are /// optional read upper bounds used only by snapshot-fenced repair chunks. pub(crate) async fn query_with_terminal_metrics( @@ -413,12 +365,12 @@ impl FrontendClient { peer: db.peer.clone(), }); db.database - .query_with_terminal_metrics_and_flow_extensions( - request, - &hints, - extensions, - snapshot_seqs, - ) + .flight_request() + .with_hints(&hints) + .with_flow_extensions(extensions) + .with_snapshot_seqs(snapshot_seqs) + .with_timeout(batch_opts.experimental_flight_do_get_timeout) + .query_with_terminal_metrics(request) .await .map_err(BoxedError::new) .context(ExternalSnafu) @@ -623,6 +575,7 @@ mod tests { use std::task::{Context, Poll}; use std::time::Duration; + use api::v1::query_request::Query; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, Ticket}; use common_query::{Output, OutputData}; diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index e6f35e17b29..0b7a9dd5704 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -30,8 +30,7 @@ use bytes::{self, Bytes}; use common_error::ext::ErrorExt; use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse}; use common_grpc::flight::{ - FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightEncoder, FlightMessage, - SNAPSHOT_SEQS_METADATA_KEY, + FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY, }; use common_memory_manager::MemoryGuard; use common_query::{Output, OutputData}; @@ -46,14 +45,16 @@ use prost::Message; use query::metrics::terminal_recordbatch_metrics_from_plan_if_requested; use query::options::FlowQueryExtensions; use session::context::{Channel, QueryContextRef}; -use snafu::{IntoError, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use table::table_name::TableName; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status, Streaming}; -use crate::error::{InvalidParameterSnafu, Result, ToJsonSnafu}; -pub use crate::grpc::flight::stream::{FlightRecordBatchSource, FlightRecordBatchStream}; +use crate::error::{InvalidParameterSnafu, InvalidQuerySnafu, Result, ToJsonSnafu}; +pub use crate::grpc::flight::stream::{ + FlightRecordBatchSource, FlightRecordBatchStream, FlightRecordBatchStreamInput, +}; use crate::grpc::greptime_handler::{ GreptimeRequestHandler, create_query_context, get_request_type, }; @@ -221,9 +222,12 @@ impl FlightCraft for GreptimeRequestHandler { ); let flight_compression = self.flight_compression; async { - let output = self - .handle_request_with_query_ctx(request, query_ctx.clone()) + let query = request.request.context(InvalidQuerySnafu { + reason: "Expecting non-empty GreptimeRequest.", + })?; + self.authenticate_request_with_query_ctx(request.header.as_ref(), &query_ctx) .await?; + let output = self.handle_request_with_query_ctx(query, query_ctx.clone()); let stream = to_flight_data_stream( output, TracingContext::from_current_span(), @@ -573,32 +577,36 @@ fn extract_json_metadata( Ok(Some(parsed)) } -fn to_flight_data_stream( - output: Output, +fn to_flight_data_stream( + output: F, tracing_context: TracingContext, flight_compression: FlightCompression, query_ctx: QueryContextRef, should_emit_terminal_metrics: bool, -) -> TonicStream { +) -> TonicStream +where + F: std::future::Future> + Send + 'static, +{ + let initializer = async move { + let output = output.await.map_err(Status::from)?; + output_to_flight_record_batch_source(output, should_emit_terminal_metrics) + }; + let stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::initializer(initializer), + tracing_context, + flight_compression, + query_ctx, + ); + Box::pin(stream) as _ +} + +fn output_to_flight_record_batch_source( + output: Output, + should_emit_terminal_metrics: bool, +) -> TonicResult { match output.data { - OutputData::Stream(stream) => { - let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(stream), - tracing_context, - flight_compression, - query_ctx, - ); - Box::pin(stream) as _ - } - OutputData::RecordBatches(x) => { - let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(x.as_stream()), - tracing_context, - flight_compression, - query_ctx, - ); - Box::pin(stream) as _ - } + OutputData::Stream(stream) => Ok(FlightRecordBatchSource::RecordBatches(stream)), + OutputData::RecordBatches(x) => Ok(FlightRecordBatchSource::RecordBatches(x.as_stream())), OutputData::AffectedRows(rows) => { let terminal_metrics = match terminal_recordbatch_metrics_from_plan_if_requested( output.meta.plan, @@ -607,20 +615,17 @@ fn to_flight_data_stream( Some(metrics) => match serde_json::to_string(&metrics) { Ok(metrics) => Some(metrics), Err(e) => { - let stream = tokio_stream::once(Err(Status::internal(format!( + return Err(Status::internal(format!( "Failed to serialize terminal metrics: {e}" - )))); - return Box::pin(stream) as _; + ))); } }, None => None, }; - let affected_rows = FlightEncoder::default().encode(FlightMessage::AffectedRows { + Ok(FlightRecordBatchSource::AffectedRows { rows, metrics: terminal_metrics, - }); - let stream = tokio_stream::iter(affected_rows.into_iter().map(Ok)); - Box::pin(stream) as _ + }) } } } diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index 319a797d9ca..b740f12d02e 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/src/servers/src/grpc/flight/stream.rs @@ -41,27 +41,34 @@ use crate::error; use crate::grpc::FlightCompression; use crate::grpc::flight::TonicResult; -pub type FlightRecordBatchStreamInitializer = Pin< - Box< - dyn Future> - + Send - + 'static, - >, ->; - pub enum FlightRecordBatchSource { RecordBatches(SendableRecordBatchStream), - Initializer(FlightRecordBatchStreamInitializer), + AffectedRows { + rows: usize, + metrics: Option, + }, } -impl FlightRecordBatchSource { - pub fn initializer(initializer: F) -> Self - where - F: Future> - + Send - + 'static, - { - Self::Initializer(Box::pin(initializer)) +/// Determines whether a Flight result is ready now or initialized asynchronously. +pub enum FlightRecordBatchStreamInput>> +{ + Ready(FlightRecordBatchSource), + Initializer(F), +} + +impl FlightRecordBatchStreamInput { + /// Creates an input from a source that is already available. + pub fn ready(source: FlightRecordBatchSource) -> Self { + Self::Ready(source) + } +} + +impl FlightRecordBatchStreamInput { + /// Creates an input that obtains its source asynchronously. + /// + /// Errors from the initializer are returned through the Flight response stream. + pub fn initializer(initializer: F) -> Self { + Self::Initializer(initializer) } } @@ -160,31 +167,32 @@ impl FlightRecordBatchStream { Self::send_metrics(tx, metrics, metrics_str).await } - pub fn new( - source: FlightRecordBatchSource, + pub fn new( + input: FlightRecordBatchStreamInput, tracing_context: TracingContext, compression: FlightCompression, query_ctx: QueryContextRef, - ) -> Self { - let should_send_partial_metrics = query_ctx.explain_verbose(); - let can_send_metrics_before_batch = query_ctx.explain_verbose() - && query_ctx.live_analyze_metrics_enabled() - && query_ctx - .remote_query_id() - .zip(query_ctx.extension(SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY)) - .is_some_and(|(remote_query_id, capability)| capability == remote_query_id); + ) -> Self + where + F: Future> + Send + 'static, + { let (mut tx, rx) = mpsc::channel::>(1); - let source_type = match &source { - FlightRecordBatchSource::RecordBatches(_) => "record_batches", - FlightRecordBatchSource::Initializer(_) => "initializer", + let source_type = match &input { + FlightRecordBatchStreamInput::Ready(FlightRecordBatchSource::RecordBatches(_)) => { + "record_batches" + } + FlightRecordBatchStreamInput::Ready(FlightRecordBatchSource::AffectedRows { + .. + }) => "affected_rows", + FlightRecordBatchStreamInput::Initializer(_) => "initializer", }; let initializer_tracing_context = tracing_context.clone(); let join_handle = common_runtime::spawn_global( async move { - let recordbatches = async move { - match source { - FlightRecordBatchSource::RecordBatches(recordbatches) => Ok(recordbatches), - FlightRecordBatchSource::Initializer(initializer) => initializer.await, + let source = async move { + match input { + FlightRecordBatchStreamInput::Ready(source) => Ok(source), + FlightRecordBatchStreamInput::Initializer(initializer) => initializer.await, } } .trace( @@ -193,8 +201,20 @@ impl FlightRecordBatchStream { ) .await; - match recordbatches { - Ok(recordbatches) => { + match source { + Ok(FlightRecordBatchSource::RecordBatches(recordbatches)) => { + let should_send_partial_metrics = query_ctx.explain_verbose(); + let can_send_metrics_before_batch = + query_ctx.explain_verbose() + && query_ctx.live_analyze_metrics_enabled() + && query_ctx + .remote_query_id() + .zip(query_ctx.extension( + SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, + )) + .is_some_and(|(remote_query_id, capability)| { + capability == remote_query_id + }); Self::flight_data_stream( recordbatches, tx, @@ -203,6 +223,11 @@ impl FlightRecordBatchStream { ) .await; } + Ok(FlightRecordBatchSource::AffectedRows { rows, metrics }) => { + let _ = tx + .send(Ok(FlightMessage::AffectedRows { rows, metrics })) + .await; + } Err(status) => { let _ = tx.send(Err(status)).await; } @@ -485,7 +510,9 @@ mod test { .unwrap() .as_stream(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), QueryContext::arc(), @@ -519,10 +546,34 @@ mod test { } } + #[tokio::test] + async fn test_flight_record_batch_stream_encodes_affected_rows() { + let mut stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::AffectedRows { + rows: 42, + metrics: Some(r#"{"region_watermarks":[]}"#.to_string()), + }), + TracingContext::default(), + FlightCompression::default(), + QueryContext::arc(), + ); + + let data = stream.next().await.unwrap().unwrap(); + let message = FlightDecoder::default().try_decode(&data).unwrap().unwrap(); + assert!(matches!( + message, + FlightMessage::AffectedRows { + rows: 42, + metrics: Some(_), + } + )); + assert!(stream.next().await.is_none()); + } + #[tokio::test] async fn test_flight_record_batch_stream_forwards_initializer_error() { let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(async { + FlightRecordBatchStreamInput::initializer(async { Err(tonic::Status::unavailable( "remote read initialization failed", )) @@ -552,9 +603,12 @@ mod test { metrics, }); let query_ctx = query_context_with_live_metrics_and_matching_capability(); - query_ctx.set_explain_verbose(true); + let initializer_query_ctx = query_ctx.clone(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::initializer(async move { + initializer_query_ctx.set_explain_verbose(true); + Ok(FlightRecordBatchSource::RecordBatches(recordbatches)) + }), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -609,7 +663,9 @@ mod test { let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -663,7 +719,9 @@ mod test { let query_ctx = query_context_with_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -706,7 +764,9 @@ mod test { let query_ctx = Arc::new(query_ctx); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -745,7 +805,9 @@ mod test { }); let query_ctx = query_context_with_matching_capability(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index ad7779435b8..98786d30adc 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -15,11 +15,13 @@ //! Handler for Greptime Database service. It's implemented by frontend. use std::collections::HashMap; +use std::future::Future; use std::str::FromStr; use std::sync::{Arc, RwLock}; use std::time::Instant; use api::helper::request_type; +use api::v1::greptime_request::Request as QueryRequest; use api::v1::{GreptimeRequest, RequestHeader}; use auth::UserProviderRef; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; @@ -79,30 +81,39 @@ impl GreptimeRequestHandler { ) -> Result { let header = request.header.as_ref(); let query_ctx = create_query_context(Channel::Grpc, header, hints, HashMap::new())?; - self.handle_request_with_query_ctx(request, query_ctx).await - } - - pub(crate) async fn handle_request_with_query_ctx( - &self, - request: GreptimeRequest, - query_ctx: QueryContextRef, - ) -> Result { let query = request.request.context(InvalidQuerySnafu { reason: "Expecting non-empty GreptimeRequest.", })?; + self.authenticate_request_with_query_ctx(request.header.as_ref(), &query_ctx) + .await?; + self.handle_request_with_query_ctx(query, query_ctx).await + } - let header = request.header.as_ref(); - let user_info = context_auth::auth(self.user_provider.clone(), header, &query_ctx).await?; + pub(crate) async fn authenticate_request_with_query_ctx( + &self, + header: Option<&RequestHeader>, + query_ctx: &QueryContextRef, + ) -> Result<()> { + let user_info = context_auth::auth(self.user_provider.clone(), header, query_ctx).await?; query_ctx.set_current_user(user_info); + Ok(()) + } + pub(crate) fn handle_request_with_query_ctx( + &self, + query: QueryRequest, + query_ctx: QueryContextRef, + ) -> impl Future> + Send + 'static { let handler = self.handler.clone(); + let runtime = self.runtime.clone(); let request_type = request_type(&query).to_string(); let db = query_ctx.get_db_string(); let timer = RequestTimer::new(db.clone(), request_type); let tracing_context = TracingContext::from_current_span(); - let result_future = async move { - handler + async move { + let result_future = async move { + handler .do_query(query, query_ctx) .trace(tracing_context.attach(tracing::info_span!( "GreptimeRequestHandler::handle_request_runtime" @@ -118,26 +129,27 @@ impl GreptimeRequestHandler { } e }) - }; + }; - match &self.runtime { - Some(runtime) => { - // Executes requests in another runtime to - // 1. prevent the execution from being cancelled unexpected by Tonic runtime; - // - Refer to our blog for the rational behind it: - // https://www.greptime.com/blogs/2023-01-12-hidden-control-flow.html - // - Obtaining a `JoinHandle` to get the panic message (if there's any). - // From its docs, `JoinHandle` is cancel safe. The task keeps running even it's handle been dropped. - // 2. avoid the handler blocks the gRPC runtime incidentally. - runtime - .spawn(result_future) - .await - .context(JoinTaskSnafu) - .inspect_err(|e| { - timer.record(e.status_code()); - })? + match runtime { + Some(runtime) => { + // Executes requests in another runtime to + // 1. prevent the execution from being cancelled unexpected by Tonic runtime; + // - Refer to our blog for the rational behind it: + // https://www.greptime.com/blogs/2023-01-12-hidden-control-flow.html + // - Obtaining a `JoinHandle` to get the panic message (if there's any). + // From its docs, `JoinHandle` is cancel safe. The task keeps running even it's handle been dropped. + // 2. avoid the handler blocks the gRPC runtime incidentally. + runtime + .spawn(result_future) + .await + .context(JoinTaskSnafu) + .inspect_err(|e| { + timer.record(e.status_code()); + })? + } + None => result_future.await, } - None => result_future.await, } } diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index cf75ae7fff1..1a2d841600f 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -16,40 +16,45 @@ mod test { use std::collections::HashMap; use std::net::SocketAddr; + use std::pin::Pin; use std::sync::Arc; use std::time::{Duration, Instant}; use api::v1::auth_header::AuthScheme; + use api::v1::greptime_request::Request as GreptimeQueryRequest; use api::v1::query_request::Query; use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, QueryRequest, SemanticType}; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, FlightDescriptor, Ticket}; use auth::user_provider_from_option; + use client::region::RegionRequester; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; - use common_grpc::flight::do_put::DoPutMetadata; + use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse}; use common_grpc::flight::{FlightEncoder, FlightMessage}; - use common_query::OutputData; + use common_query::{Output, OutputData}; use common_recordbatch::adapter::RegionWatermarkEntry; use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use common_telemetry::tracing_context::TracingContext; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector}; - use futures_util::StreamExt; + use futures_util::{Stream, StreamExt}; use hyper_util::rt::TokioIo; use itertools::Itertools; use servers::grpc::builder::GrpcServerBuilder; use servers::grpc::flight::{ FlightCraft, FlightCraftWrapper, FlightRecordBatchSource, FlightRecordBatchStream, - TonicStream, + FlightRecordBatchStreamInput, PutRecordBatchRequestStream, TonicStream, }; use servers::grpc::greptime_handler::GreptimeRequestHandler; use servers::grpc::{FlightCompression, GrpcServerConfig}; + use servers::query_handler::grpc::GrpcQueryHandler; use servers::server::Server; - use tonic::Response; + use session::context::QueryContextRef; use tonic::transport::Server as TonicServer; + use tonic::{Response, Status}; use tower::service_fn; use crate::cluster::GreptimeDbClusterBuilder; @@ -59,6 +64,12 @@ mod test { struct SlowFlightCraft; + struct ErrorFlightCraft; + + struct SlowRemoteQueryHandler { + region_requester: RegionRequester, + } + fn slow_recordbatch_stream() -> SendableRecordBatchStream { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "value", @@ -83,9 +94,11 @@ mod test { _: tonic::Request, ) -> std::result::Result>, tonic::Status> { let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(async { + FlightRecordBatchStreamInput::initializer(async { tokio::time::sleep(Duration::from_secs(2)).await; - Ok(slow_recordbatch_stream()) + Ok(FlightRecordBatchSource::RecordBatches( + slow_recordbatch_stream(), + )) }), TracingContext::default(), FlightCompression::default(), @@ -96,14 +109,72 @@ mod test { } } - #[tokio::test(flavor = "multi_thread")] - async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { + #[async_trait::async_trait] + impl FlightCraft for ErrorFlightCraft { + async fn do_get( + &self, + _: tonic::Request, + ) -> std::result::Result>, tonic::Status> { + let stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::initializer(async { + Err(Status::internal("deferred initializer detail")) + }), + TracingContext::default(), + FlightCompression::default(), + session::context::QueryContext::arc(), + ); + Ok(Response::new(Box::pin(stream))) + } + } + + #[async_trait::async_trait] + impl GrpcQueryHandler for SlowRemoteQueryHandler { + async fn do_query( + &self, + _: GreptimeQueryRequest, + _: QueryContextRef, + ) -> servers::error::Result { + let stream = self + .region_requester + .do_get_inner(Ticket::default()) + .await + .unwrap(); + Ok(Output::new_with_stream(stream)) + } + + fn handle_put_record_batch_stream( + &self, + _: PutRecordBatchRequestStream, + _: QueryContextRef, + ) -> Pin> + Send>> { + Box::pin(futures::stream::empty()) + } + } + + fn client_for_flight_craft(addr: &'static str, craft: T) -> Client + where + T: FlightCraft, + { + client_for_flight_craft_with_max_encoding(addr, craft, None) + } + + fn client_for_flight_craft_with_max_encoding( + addr: &'static str, + craft: T, + max_encoding_message_size: Option, + ) -> Client + where + T: FlightCraft, + { let (client_io, server_io) = tokio::io::duplex(1024); tokio::spawn(async move { + let flight_service = FlightServiceServer::new(FlightCraftWrapper(craft)); + let flight_service = match max_encoding_message_size { + Some(size) => flight_service.max_encoding_message_size(size), + None => flight_service, + }; TonicServer::builder() - .add_service(FlightServiceServer::new(FlightCraftWrapper( - SlowFlightCraft, - ))) + .add_service(flight_service) .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>( server_io, )])) @@ -115,7 +186,7 @@ mod test { let mut client_io = Some(client_io); channel_manager .reset_with_connector( - "slow-flight", + addr, service_fn(move |_| { let client_io = client_io.take(); @@ -127,7 +198,12 @@ mod test { }), ) .unwrap(); - let client = Client::with_manager_and_urls(channel_manager, ["slow-flight"]); + Client::with_manager_and_urls(channel_manager, [addr]) + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { + let client = client_for_flight_craft("slow-flight", SlowFlightCraft); let mut flight_client = client.make_flight_client(false, false).unwrap(); let start = Instant::now(); @@ -143,6 +219,64 @@ mod test { assert!(stream.message().await.unwrap().is_some()); } + + #[tokio::test(flavor = "multi_thread")] + async fn test_deferred_flight_initializer_error_preserves_message() { + let client = client_for_flight_craft("error-flight", ErrorFlightCraft); + let mut flight_client = client.make_flight_client(false, false).unwrap(); + let mut stream = flight_client + .mut_inner() + .do_get(tonic::Request::new(Ticket::default())) + .await + .unwrap() + .into_inner(); + let error = stream.message().await.unwrap_err(); + assert_eq!(tonic::Code::Internal, error.code()); + assert_eq!("deferred initializer detail", error.message()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_deferred_flight_encode_error_preserves_message() { + let client = + client_for_flight_craft_with_max_encoding("limited-flight", SlowFlightCraft, Some(1)); + let mut flight_client = client.make_flight_client(false, false).unwrap(); + let mut stream = flight_client + .mut_inner() + .do_get(tonic::Request::new(Ticket::default())) + .await + .unwrap() + .into_inner(); + let error = stream.message().await.unwrap_err(); + assert_eq!(tonic::Code::OutOfRange, error.code()); + assert!(error.message().contains("message length too large")); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_request_timeout_does_not_cancel_slow_datanode_stream() { + let datanode_client = client_for_flight_craft("slow-datanode", SlowFlightCraft); + let frontend_handler = GreptimeRequestHandler::new( + Arc::new(SlowRemoteQueryHandler { + region_requester: RegionRequester::new(datanode_client, false, false), + }), + None, + None, + FlightCompression::default(), + ); + let frontend_client = client_for_flight_craft("slow-frontend", frontend_handler); + let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, frontend_client); + + let start = Instant::now(); + let output = database + .flight_request() + .with_timeout(Duration::from_secs(1)) + .sql("select 1") + .await + .unwrap(); + // `sql()` waits for the first Flight message. Succeeding after this delay + // proves the request timeout applied only to the response header. + assert!(start.elapsed() >= Duration::from_secs(1)); + assert!(matches!(output.data, OutputData::Stream(_))); + } #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put() { common_telemetry::init_default_ut_logging(); @@ -446,16 +580,14 @@ mod test { ); let result = client - .query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::Sql( - "select ts, a, `B` from foo order by ts".to_string(), - )), - }, - &[], - &[("flow.return_region_seq", "true")], - &HashMap::from([(region_id.as_u64(), stale_snapshot_seq)]), - ) + .flight_request() + .with_flow_extensions(&[("flow.return_region_seq", "true")]) + .with_snapshot_seqs(&HashMap::from([(region_id.as_u64(), stale_snapshot_seq)])) + .query_with_terminal_metrics(QueryRequest { + query: Some(Query::Sql( + "select ts, a, `B` from foo order by ts".to_string(), + )), + }) .await .unwrap(); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 4a8f96260e8..43ba8701d76 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2252,6 +2252,7 @@ query_timeout = "10m" slow_query_threshold = "1m" experimental_min_refresh_duration = "5s" grpc_conn_timeout = "5s" +experimental_flight_do_get_timeout = "10s" experimental_grpc_max_retries = 3 experimental_frontend_scan_timeout = "30s" experimental_max_filter_num_per_query = 20 From b31f05eb5964c17342e7606d80dd04e6b93a3e03 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 27 Aug 2026 03:37:59 +0000 Subject: [PATCH 112/126] fix: update tokio-postgres and correct explain/fetch cursor output schema (#8955) * chore(deps): update tokio-postgres * fix: describing fetch cursor and analyze --- Cargo.lock | 180 ++++++++++++++++++++------- Cargo.toml | 2 +- src/common/meta/Cargo.toml | 2 +- src/common/recordbatch/src/cursor.rs | 9 ++ src/query/src/analyze.rs | 21 +++- src/query/src/lib.rs | 2 +- src/servers/Cargo.toml | 4 +- src/servers/src/postgres/handler.rs | 29 +++++ src/session/src/lib.rs | 9 ++ 9 files changed, 202 insertions(+), 56 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ec477e5c9a3..d77ec340a8d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4628,10 +4628,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" dependencies = [ "const-oid 0.9.6", + "der_derive", + "flagset", "pem-rfc7468", "zeroize", ] +[[package]] +name = "der_derive" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8034092389675178f570469e6c3b0465d3d30b4505c294a6550db47f3c17ad18" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "deranged" version = "0.5.8" @@ -5334,6 +5347,12 @@ version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" +[[package]] +name = "flagset" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7ac824320a75a52197e8f2d787f6a38b6718bb6897a35142d749af3c0e8f4fe" + [[package]] name = "flatbuffers" version = "25.2.10" @@ -5973,7 +5992,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "wasm-bindgen", ] @@ -7866,13 +7885,14 @@ checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" [[package]] name = "libredox" -version = "0.1.4" +version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1580801010e535496706ba011c15f8532df6b42297d2e471fec38ceadd8c0638" +checksum = "28d0a00925a9f930d679b6789b721e3a7f9ed110f41b86d2497caa780c3a070a" dependencies = [ "bitflags 2.12.1", "libc", - "redox_syscall 0.5.13", + "plain", + "redox_syscall 0.9.3", ] [[package]] @@ -8543,7 +8563,7 @@ checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" dependencies = [ "libc", "log", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] @@ -9358,6 +9378,15 @@ dependencies = [ "objc2-encode", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags 2.12.1", +] + [[package]] name = "objc2-encode" version = "4.1.0" @@ -9374,6 +9403,15 @@ dependencies = [ "objc2", ] +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "objc2-core-foundation", +] + [[package]] name = "object" version = "0.36.7" @@ -10552,16 +10590,7 @@ dependencies = [ "tokio", "tokio-rustls", "tokio-util", - "x509-certificate 0.25.0", -] - -[[package]] -name = "phf" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" -dependencies = [ - "phf_shared 0.11.3", + "x509-certificate", ] [[package]] @@ -10814,6 +10843,12 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +[[package]] +name = "plain" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6" + [[package]] name = "plist" version = "1.7.2" @@ -10924,16 +10959,16 @@ dependencies = [ [[package]] name = "postgres-types" -version = "0.2.9" +version = "0.2.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613283563cd90e1dfc3518d548caee47e0e725455ed619881f5cf21f36de4b48" +checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be" dependencies = [ "array-init", "bytes", "chrono", "fallible-iterator", "postgres-protocol", - "serde", + "serde_core", "serde_json", ] @@ -11299,7 +11334,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.4.1", + "heck 0.5.0", "itertools 0.12.1", "log", "multimap", @@ -11319,7 +11354,7 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "itertools 0.14.0", "log", "multimap", @@ -12087,6 +12122,15 @@ dependencies = [ "bitflags 2.12.1", ] +[[package]] +name = "redox_syscall" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d678d17679829e73d371e96880897e98fee2ded7acc0a50bdf8af2affa4b2fe5" +dependencies = [ + "bitflags 2.12.1", +] + [[package]] name = "redox_users" version = "0.4.6" @@ -13785,7 +13829,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -14117,7 +14161,7 @@ dependencies = [ "stringprep", "thiserror 2.0.17", "tracing", - "whoami", + "whoami 1.6.0", ] [[package]] @@ -14156,7 +14200,7 @@ dependencies = [ "stringprep", "thiserror 2.0.17", "tracing", - "whoami", + "whoami 1.6.0", ] [[package]] @@ -15175,6 +15219,27 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tls_codec" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0de2e01245e2bb89d6f05801c564fa27624dbd7b1846859876c7dad82e90bf6b" +dependencies = [ + "tls_codec_derive", + "zeroize", +] + +[[package]] +name = "tls_codec_derive" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d2e76690929402faae40aebdda620a2c0e25dd6d3b9afe48867dfd95991f4bd" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "tokio" version = "1.52.3" @@ -15241,9 +15306,9 @@ dependencies = [ [[package]] name = "tokio-postgres" -version = "0.7.13" +version = "0.7.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c95d533c83082bb6490e0189acaa0bbeef9084e60471b696ca6988cd0541fb0" +checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3" dependencies = [ "async-trait", "byteorder", @@ -15254,29 +15319,29 @@ dependencies = [ "log", "parking_lot 0.12.4", "percent-encoding", - "phf 0.11.3", + "phf 0.13.1", "pin-project-lite", "postgres-protocol", "postgres-types", - "rand 0.9.4", - "socket2 0.5.10", + "rand 0.10.1", + "socket2 0.6.4", "tokio", "tokio-util", - "whoami", + "whoami 2.1.3", ] [[package]] name = "tokio-postgres-rustls" -version = "0.12.0" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04fb792ccd6bbcd4bba408eb8a292f70fc4a3589e5d793626f45190e6454b6ab" +checksum = "4c2ad44aa0ae96db89c4742212ed41645b2f597311ff6e1945542a4d9fadc2fb" dependencies = [ - "ring", "rustls", + "sha2 0.11.0", "tokio", "tokio-postgres", "tokio-rustls", - "x509-certificate 0.23.1", + "x509-cert", ] [[package]] @@ -16284,6 +16349,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasi" +version = "0.14.7+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" +dependencies = [ + "wasip2", +] + [[package]] name = "wasip2" version = "1.0.2+wasi-0.2.9" @@ -16308,6 +16382,15 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8dad83b4f25e74f184f64c43b150b91efe7647395b42289f38e50566d82855b" +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + [[package]] name = "wasm-bindgen" version = "0.2.122" @@ -16530,7 +16613,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6994d13118ab492c3c80c1f81928718159254c53c472bf9ce36f8dae4add02a7" dependencies = [ "redox_syscall 0.5.13", - "wasite", + "wasite 0.1.0", +] + +[[package]] +name = "whoami" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "626c4bac6755d76ffc12cb01b2eac751db1996b9e0041de9aa02c8c211ddc82c" +dependencies = [ + "libc", + "libredox", + "objc2-system-configuration", + "wasite 1.0.2", "web-sys", ] @@ -17139,22 +17234,15 @@ dependencies = [ ] [[package]] -name = "x509-certificate" -version = "0.23.1" +name = "x509-cert" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "66534846dec7a11d7c50a74b7cdb208b9a581cad890b7866430d438455847c85" +checksum = "1301e935010a701ae5f8655edc0ad17c44bad3ac5ce8c39185f75453b720ae94" dependencies = [ - "bcder", - "bytes", - "chrono", + "const-oid 0.9.6", "der", - "hex", - "pem", - "ring", - "signature", "spki", - "thiserror 1.0.69", - "zeroize", + "tls_codec", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 7c3fccfbbf8..fab7b2aa736 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -255,7 +255,7 @@ strum = { version = "0.27", features = ["derive"] } sysinfo = "0.33" tempfile = "3" tokio = { version = "1.47", features = ["full"] } -tokio-postgres = "0.7" +tokio-postgres = "0.7.18" tokio-rustls = { version = "0.26.2", default-features = false } tokio-stream = "0.1" tokio-util = { version = "0.7", features = ["io-util", "compat"] } diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 8e0cbbdf354..8deb3d00814 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -89,7 +89,7 @@ strum.workspace = true table = { workspace = true, features = ["testing"] } tokio.workspace = true tokio-postgres = { workspace = true, optional = true } -tokio-postgres-rustls = { version = "0.12", optional = true } +tokio-postgres-rustls = { version = "0.14", optional = true } tonic.workspace = true tracing.workspace = true typetag.workspace = true diff --git a/src/common/recordbatch/src/cursor.rs b/src/common/recordbatch/src/cursor.rs index a741953ccc2..60ae9d54fa2 100644 --- a/src/common/recordbatch/src/cursor.rs +++ b/src/common/recordbatch/src/cursor.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use datatypes::schema::SchemaRef; use futures::StreamExt; use tokio::sync::Mutex; @@ -28,12 +29,15 @@ struct Inner { /// A cursor on RecordBatchStream that fetches data batch by batch pub struct RecordBatchStreamCursor { + schema: SchemaRef, inner: Mutex, } impl RecordBatchStreamCursor { pub fn new(stream: SendableRecordBatchStream) -> RecordBatchStreamCursor { + let schema = stream.schema(); Self { + schema, inner: Mutex::new(Inner { stream, current_row_index: 0, @@ -43,6 +47,11 @@ impl RecordBatchStreamCursor { } } + /// Returns the schema of the underlying record batch stream. + pub fn schema(&self) -> SchemaRef { + self.schema.clone() + } + /// Take `size` of row from the `RecordBatchStream` and create a new /// `RecordBatch` for these rows. pub async fn take(&self, size: usize) -> Result { diff --git a/src/query/src/analyze.rs b/src/query/src/analyze.rs index ec0abf25ba6..b1522fc05d5 100644 --- a/src/query/src/analyze.rs +++ b/src/query/src/analyze.rs @@ -46,6 +46,21 @@ const STAGE: &str = "stage"; const NODE: &str = "node"; const PLAN: &str = "plan"; +/// Returns the fixed output schema of [`DistAnalyzeExec`]: +/// (`stage`: UInt32, `node`: UInt32, `plan`: Utf8). +/// +/// Exposed so protocol-level `Describe` handlers can advertise the schema +/// clients will actually receive (the physical plan is rewritten in +/// `optimize_physical_plan` and its schema differs from the DataFusion +/// `Analyze` logical plan schema). +pub fn dist_analyze_output_schema() -> SchemaRef { + SchemaRef::new(Schema::new(vec![ + Field::new(STAGE, DataType::UInt32, true), + Field::new(NODE, DataType::UInt32, true), + Field::new(PLAN, DataType::Utf8, true), + ])) +} + #[derive(Debug)] pub struct DistAnalyzeExec { input: Arc, @@ -58,11 +73,7 @@ pub struct DistAnalyzeExec { impl DistAnalyzeExec { /// Create a new DistAnalyzeExec pub fn new(input: Arc, verbose: bool, format: AnalyzeFormat) -> Self { - let schema = SchemaRef::new(Schema::new(vec![ - Field::new(STAGE, DataType::UInt32, true), - Field::new(NODE, DataType::UInt32, true), - Field::new(PLAN, DataType::Utf8, true), - ])); + let schema = dist_analyze_output_schema(); let properties = Arc::new(Self::compute_properties(&input, schema.clone())); Self { input, diff --git a/src/query/src/lib.rs b/src/query/src/lib.rs index 68d8ff3a9eb..4b2a2d5264d 100644 --- a/src/query/src/lib.rs +++ b/src/query/src/lib.rs @@ -45,7 +45,7 @@ pub(crate) mod test_util; #[cfg(test)] mod tests; -pub use crate::analyze::analyze_plan_metrics_to_json_value; +pub use crate::analyze::{analyze_plan_metrics_to_json_value, dist_analyze_output_schema}; pub use crate::datafusion::DfContextProviderAdapter; pub use crate::query_engine::{ QueryEngine, QueryEngineContext, QueryEngineFactory, QueryEngineRef, diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 7936c96fbda..bb358ea76a6 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -167,8 +167,8 @@ servers = { workspace = true, features = ["testing"] } session = { workspace = true, features = ["testing"] } table.workspace = true tempfile = "3.0.0" -tokio-postgres = "0.7" -tokio-postgres-rustls = "0.12" +tokio-postgres.workspace = true +tokio-postgres-rustls = "0.14" [target.'cfg(unix)'.dev-dependencies] pprof = { version = "0.14", features = ["criterion", "flamegraph"] } diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 484bb6a1f11..433eb11cdad 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -40,6 +40,7 @@ use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::PgWireBackendMessage; use pgwire::messages::copy::CopyData; use pgwire::messages::data::DataRow; +use query::dist_analyze_output_schema; use query::planner::DfLogicalPlanner; use query::query_engine::DescribeResult; use session::Session; @@ -554,6 +555,16 @@ fn describe_fields( session: &Arc, ) -> PgWireResult> { match sql_plan { + // EXPLAIN ANALYZE: at execution time the physical plan is replaced with + // GreptimeDB's `DistAnalyzeExec` (see `optimize_physical_plan`), whose + // output schema (stage/node/plan) differs from the DataFusion `Analyze` + // logical plan schema (plan_type/plan). Describe with the schema the + // client will actually receive so the DataRow field count matches. + SqlPlan::Plan(LogicalPlan::Analyze(_), _) => { + let schema: Schema = + Schema::try_from(dist_analyze_output_schema()).map_err(convert_err)?; + schema_to_pg(&schema, format, None).map_err(convert_err) + } // query SqlPlan::Plan(plan, _) if !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_)) => { let schema: Schema = plan.schema().clone().try_into().map_err(convert_err)?; @@ -679,6 +690,24 @@ fn describe_fields( Ok(vec![]) } } + // FETCH cursor: return the cursor's schema so the RowDescription + // matches the DataRow field count sent during Execute. + SqlPlan::Statement(Statement::FetchCursor(fetch), _) => { + let cursor_name = fetch.cursor_name.to_string(); + match session.get_cursor(&cursor_name) { + Some(cursor) => { + // `cursor.schema()` is the GreptimeDB `SchemaRef` captured + // when the cursor was declared; `schema_to_pg` accepts it + // directly (no DataFusion -> GreptimeDB conversion needed). + schema_to_pg(&cursor.schema(), format, None).map_err(convert_err) + } + None => { + // Cursor not found (e.g. DECLARE hasn't executed yet in + // an extended-protocol batch). Return NoData as a fallback. + Ok(vec![]) + } + } + } _ => { // NoData Ok(vec![]) diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index 110560b3243..44f3ee66891 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -111,6 +111,15 @@ impl Session { .into() } + /// Returns the cursor with the given name, if it exists. + /// + /// Cursors are stored in the session's mutable inner data and shared across + /// all query contexts created from this session. + pub fn get_cursor(&self, name: &str) -> Option> { + let guard = self.mutable_inner.read().unwrap(); + guard.cursors.get(name).cloned() + } + pub fn conn_info(&self) -> &ConnInfo { &self.conn_info } From a7ba8b01e1c8ced20a19e39ac68a4637dc64e266 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 27 Aug 2026 04:24:07 +0000 Subject: [PATCH 113/126] fix(flow): restore FrontendClient::sql API (#8963) fix(flow): restore batching frontend SQL Signed-off-by: WenyXu --- src/flow/src/batching_mode/frontend_client.rs | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index 27db96748a6..48314dfa21b 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -18,6 +18,7 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex, RwLock, Weak}; use api::v1::greptime_request::Request; +use api::v1::query_request::Query; use api::v1::{CreateTableExpr, QueryRequest}; use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputWithMetrics}; use common_error::ext::BoxedError; @@ -339,6 +340,53 @@ impl FrontendClient { }) } + /// Execute a SQL statement on the frontend. + pub async fn sql(&self, catalog: &str, schema: &str, sql: &str) -> Result { + match self { + FrontendClient::Distributed { .. } => { + let db = self.get_random_active_frontend(catalog, schema).await?; + db.database + .sql(sql) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu) + } + FrontendClient::Standalone { + database_client, .. + } => { + let ctx = QueryContextBuilder::default() + .current_catalog(catalog.to_string()) + .current_schema(schema.to_string()) + .build(); + let ctx = Arc::new(ctx); + { + let database_client = { + database_client + .handler + .lock() + .unwrap() + .as_ref() + .context(UnexpectedSnafu { + reason: "Standalone's frontend instance is not set", + })? + .upgrade() + .context(UnexpectedSnafu { + reason: "Failed to upgrade database client", + })? + }; + let req = Request::Query(QueryRequest { + query: Some(Query::Sql(sql.to_string())), + }); + database_client + .do_query(req, ctx) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu) + } + } + } + } + /// Execute a flow query and return terminal metrics. `snapshot_seqs` are /// optional read upper bounds used only by snapshot-fenced repair chunks. pub(crate) async fn query_with_terminal_metrics( From beded6e232918147b93f91677e5ec8b8768d1be5 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 27 Aug 2026 04:51:39 +0000 Subject: [PATCH 114/126] fix(flow): avoid insert select HTTP/2 stalls (#8962) * fix(flow): avoid insert select HTTP/2 stalls Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): share record batch forwarding for DML Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/cmd/src/frontend.rs | 2 + src/query/src/datafusion.rs | 126 +++++++++++++++++++++++++++++------- 2 files changed, 103 insertions(+), 25 deletions(-) diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index e3737088e8d..a303041b5a2 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -434,6 +434,8 @@ impl StartCommand { // Some queries are expected to take long time. let mut channel_config = opts.datanode.client.channel_config(); channel_config.timeout = None; + // Source Flight streams and sink unary responses share pooled connections. + channel_config.http2_adaptive_window = Some(true); if opts.grpc.flight_compression.transport_compression() { channel_config.accept_compression = true; channel_config.send_compression = true; diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 754199a78fc..f9d87782a03 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -30,7 +30,7 @@ use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; use common_query::{Output, OutputData, OutputMeta}; use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters}; -use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream}; +use common_recordbatch::{EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream}; use common_telemetry::tracing; use datafusion::catalog::TableFunction; use datafusion::dataframe::DataFrame; @@ -44,6 +44,7 @@ use datafusion_expr::{ use datatypes::prelude::VectorRef; use datatypes::schema::Schema; use futures_util::StreamExt; +use futures_util::future::try_join; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::AnalyzeFormat; @@ -80,6 +81,29 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism"; /// Whether to fallback to the original plan when failed to push down. pub const QUERY_FALLBACK_HINT: &str = "query_fallback"; +// An unbounded queue keeps draining source RPCs while mutation RPCs on a shared +// HTTP/2 connection are pending, trading bounded memory for request liveness. +async fn forward_record_batches( + mut stream: SendableRecordBatchStream, + batch_tx: tokio::sync::mpsc::UnboundedSender>, +) -> Result<()> { + while let Some(batch) = stream.next().await { + match batch.context(CreateRecordBatchSnafu) { + Ok(batch) => { + if batch_tx.send(Ok(batch)).is_err() { + break; + } + tokio::task::yield_now().await; + } + Err(error) => { + let _ = batch_tx.send(Err(error)); + break; + } + } + } + Ok(()) +} + fn query_load_region_id(plan: &Arc) -> Option { let mut region_id = None; let mut stack = vec![plan.clone()]; @@ -205,7 +229,7 @@ impl DatafusionQueryEngine { let Output { data, meta } = self .exec_query_plan((*dml.input).clone(), query_ctx.clone()) .await?; - let mut stream = match data { + let stream = match data { OutputData::RecordBatches(batches) => batches.as_stream(), OutputData::Stream(stream) => stream, _ => unreachable!(), @@ -214,30 +238,35 @@ impl DatafusionQueryEngine { let mut affected_rows = 0; let mut insert_cost = 0; - while let Some(batch) = stream.next().await { - let batch = batch.context(CreateRecordBatchSnafu)?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table.schema()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - - match dml.op { - WriteOp::Insert(_) => { - // We ignore the insert op. - let output = self - .insert(&table_name, column_vectors, query_ctx.clone()) - .await?; - let (rows, cost) = output.extract_rows_and_cost(); - affected_rows += rows; - insert_cost += cost; - } - WriteOp::Delete => { - affected_rows += self - .delete(&table_name, &table, column_vectors, query_ctx.clone()) - .await?; - } - _ => unreachable!("guarded by the 'ensure!' at the beginning"), + match dml.op { + WriteOp::Insert(_) => { + let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); + let producer = forward_record_batches(stream, batch_tx); + let consumer = self.consume_insert_record_batches( + batch_rx, + &table_name, + table.schema(), + query_ctx.clone(), + ); + let ((), (rows, cost)) = try_join(producer, consumer).await?; + affected_rows += rows; + insert_cost += cost; } + WriteOp::Delete => { + // Keep DELETE on the same producer/consumer schedule as INSERT so the source + // stream can continue draining while mutation RPCs are pending. + let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); + let producer = forward_record_batches(stream, batch_tx); + let consumer = self.consume_delete_record_batches( + batch_rx, + &table_name, + &table, + query_ctx.clone(), + ); + let (rows, ()) = try_join(consumer, producer).await?; + affected_rows += rows; + } + _ => unreachable!("guarded by the 'ensure!' at the beginning"), } Ok(Output::new( OutputData::AffectedRows(affected_rows), @@ -245,6 +274,53 @@ impl DatafusionQueryEngine { )) } + async fn consume_insert_record_batches( + &self, + mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, + table_name: &ResolvedTableReference, + table_schema: Arc, + query_ctx: QueryContextRef, + ) -> Result<(usize, usize)> { + let mut affected_rows = 0; + let mut insert_cost = 0; + while let Some(batch) = batch_rx.recv().await { + let batch = batch?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table_schema.clone()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + // We ignore the insert op. + let output = self + .insert(table_name, column_vectors, query_ctx.clone()) + .await?; + let (rows, cost) = output.extract_rows_and_cost(); + affected_rows += rows; + insert_cost += cost; + } + Ok((affected_rows, insert_cost)) + } + + async fn consume_delete_record_batches( + &self, + mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, + table_name: &ResolvedTableReference, + table: &TableRef, + query_ctx: QueryContextRef, + ) -> Result { + let mut affected_rows = 0; + while let Some(batch) = batch_rx.recv().await { + let batch = batch?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + affected_rows += self + .delete(table_name, table, column_vectors, query_ctx.clone()) + .await?; + } + Ok(affected_rows) + } + #[tracing::instrument(skip_all)] async fn delete( &self, From 9befb7e64cdccf25e96b36fa5ca8fa98f1901dd7 Mon Sep 17 00:00:00 2001 From: sun Date: Thu, 27 Aug 2026 09:57:18 +0000 Subject: [PATCH 115/126] feat: update dashboard to v0.13.14 (#8968) --- src/servers/dashboard/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index b68b26e6f78..0822cdb1e68 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.13.13 +v0.13.14 From a2f39ecf7b22da2a99502ff2838d51d90a502e66 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 27 Aug 2026 17:27:52 +0000 Subject: [PATCH 116/126] feat: harden frontend heartbeat extensions (#8803) Add a typed `Command::build_with_heartbeat_extensions` seam in `src/cmd/src/frontend.rs`. Freeze and harden `FrontendHeartbeatExtensions` in `src/frontend/src/heartbeat.rs`, with lifecycle and race coverage in `src/frontend/src/heartbeat/tests.rs`. Signed-off-by: Lei, HUANG --- src/cmd/src/frontend.rs | 209 +++++++++++++++++++++++++--- src/frontend/src/heartbeat.rs | 69 ++++++++- src/frontend/src/heartbeat/tests.rs | 167 +++++++++++++++++++++- 3 files changed, 418 insertions(+), 27 deletions(-) diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index a303041b5a2..1f95a3c54b4 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -13,7 +13,9 @@ // limitations under the License. use std::fmt::Debug; +use std::future::Future; use std::path::Path; +use std::pin::Pin; use std::sync::Arc; use std::time::Duration; @@ -61,6 +63,7 @@ use crate::options::{GlobalOptions, GreptimeOptions}; use crate::{App, create_resource_limit_metrics, log_versions, maybe_activate_heap_profile}; type FrontendOptions = GreptimeOptions; +type HeartbeatExtensionSetupFuture<'a> = Pin> + Send + 'a>>; pub struct Instance { frontend: Frontend, @@ -117,7 +120,35 @@ pub struct Command { impl Command { pub async fn build(&self, opts: FrontendOptions) -> Result { - self.subcmd.build(opts).await + self.build_with_heartbeat_extensions(opts, FrontendHeartbeatExtensions::default()) + .await + } + + /// Builds a frontend with pre-registered heartbeat extensions. + /// + /// Register extensions before calling this method. The supplied registry is installed before + /// normal plugin setup. After plugin heartbeat setup completes, the registry is frozen before + /// heartbeat handlers and the heartbeat task are built, so later registrations are rejected + /// and every heartbeat consumer observes the same extension membership. + pub async fn build_with_heartbeat_extensions( + &self, + opts: FrontendOptions, + heartbeat_extensions: FrontendHeartbeatExtensions, + ) -> Result { + let plugins = plugins_with_heartbeat_extensions(heartbeat_extensions); + self.forward_build_with_plugins(opts, plugins, |command, opts, plugins| { + command.build_with_plugins(opts, plugins) + }) + .await + } + + fn forward_build_with_plugins<'a, T>( + &'a self, + opts: FrontendOptions, + plugins: Plugins, + build: impl FnOnce(&'a StartCommand, FrontendOptions, Plugins) -> T, + ) -> T { + self.subcmd.forward_build_with_plugins(opts, plugins, build) } pub fn load_options(&self, global_options: &GlobalOptions) -> Result { @@ -131,9 +162,14 @@ pub enum SubCommand { } impl SubCommand { - async fn build(&self, opts: FrontendOptions) -> Result { + fn forward_build_with_plugins<'a, T>( + &'a self, + opts: FrontendOptions, + plugins: Plugins, + build: impl FnOnce(&'a StartCommand, FrontendOptions, Plugins) -> T, + ) -> T { match self { - SubCommand::Start(cmd) => cmd.build(opts).await, + SubCommand::Start(cmd) => build(cmd, opts, plugins), } } @@ -329,7 +365,11 @@ impl StartCommand { Ok(()) } - async fn build(&self, opts: FrontendOptions) -> Result { + async fn build_with_plugins( + &self, + opts: FrontendOptions, + plugins: Plugins, + ) -> Result { let guard = common_telemetry::init_global_logging( APP_NAME, &opts.component.logging, @@ -380,15 +420,8 @@ impl StartCommand { .await .context(error::MetaClientInitSnafu)?; - let mut plugins = Plugins::new(); - plugins::setup_frontend_plugins_pre_build( - &mut plugins, - &plugin_opts, - &opts, - Some(&meta_config), - ) - .await - .context(error::StartFrontendSnafu)?; + let mut plugins = + prepare_frontend_plugins(plugins, &plugin_opts, &opts, Some(&meta_config)).await?; // now initialize the meta_client with plugins let meta_client = meta_client::create_meta_client( @@ -499,12 +532,20 @@ impl StartCommand { let instance = Arc::new(instance); - plugins::setup_frontend_heartbeat_extensions(&mut plugins, &plugin_opts, &instance) - .await - .context(error::StartFrontendSnafu)?; - let heartbeat_extensions = plugins - .get::() - .unwrap_or_default(); + let heartbeat_instance = instance.clone(); + let heartbeat_extensions = + setup_and_freeze_frontend_heartbeat_extensions(&mut plugins, move |plugins| { + Box::pin(async move { + plugins::setup_frontend_heartbeat_extensions( + plugins, + &plugin_opts, + &heartbeat_instance, + ) + .await + .context(error::StartFrontendSnafu) + }) + }) + .await?; let heartbeat_task = Some(create_heartbeat_task_with_extensions( &opts, meta_client, @@ -526,6 +567,36 @@ impl StartCommand { } } +async fn prepare_frontend_plugins( + mut plugins: Plugins, + plugin_opts: &[PluginOptions], + opts: &frontend::frontend::FrontendOptions, + meta_config: Option<&[PluginOptions]>, +) -> Result { + plugins::setup_frontend_plugins_pre_build(&mut plugins, plugin_opts, opts, meta_config) + .await + .context(error::StartFrontendSnafu)?; + Ok(plugins) +} + +fn plugins_with_heartbeat_extensions(heartbeat_extensions: FrontendHeartbeatExtensions) -> Plugins { + let plugins = Plugins::new(); + plugins.insert(heartbeat_extensions); + plugins +} + +async fn setup_and_freeze_frontend_heartbeat_extensions( + plugins: &mut Plugins, + setup: impl for<'a> FnOnce(&'a mut Plugins) -> HeartbeatExtensionSetupFuture<'a>, +) -> Result { + setup(plugins).await?; + let extensions = plugins + .get::() + .unwrap_or_default(); + extensions.freeze(); + Ok(extensions) +} + pub fn create_heartbeat_task( options: &frontend::frontend::FrontendOptions, meta_client: MetaClientRef, @@ -565,6 +636,7 @@ fn create_heartbeat_task_with_extensions( #[cfg(test)] mod tests { use std::io::Write; + use std::sync::Arc; use std::time::Duration; use auth::{Identity, Password, UserProviderRef}; @@ -578,6 +650,105 @@ mod tests { use super::*; use crate::options::GlobalOptions; + #[derive(Debug)] + struct TestExtension(&'static str); + + #[async_trait] + impl frontend::heartbeat::FrontendHeartbeatExtension for TestExtension { + fn name(&self) -> &'static str { + self.0 + } + } + + #[test] + fn test_command_forwards_heartbeat_extensions_to_start_build() { + let command = Command { + subcmd: SubCommand::Start(StartCommand::default()), + }; + let extensions = FrontendHeartbeatExtensions::default(); + let shared_extensions = extensions.clone(); + + let forwarded_plugins = command.forward_build_with_plugins( + FrontendOptions::default(), + plugins_with_heartbeat_extensions(extensions), + |_, _, plugins| plugins, + ); + let forwarded_extensions = forwarded_plugins + .get::() + .unwrap(); + + assert_registry_identity(&shared_extensions, &forwarded_extensions); + } + + #[tokio::test] + async fn test_prefilled_heartbeat_extensions_survive_pre_build_setup() { + let extensions = FrontendHeartbeatExtensions::default(); + let shared_extensions = extensions.clone(); + + let plugins = prepare_frontend_plugins( + plugins_with_heartbeat_extensions(extensions), + &[], + &frontend::frontend::FrontendOptions::default(), + None, + ) + .await + .unwrap(); + + let setup_extensions = plugins.get_or_insert(FrontendHeartbeatExtensions::default); + assert_registry_identity(&shared_extensions, &setup_extensions); + } + + #[tokio::test] + async fn test_heartbeat_extension_setup_completes_before_freeze() { + let extensions = FrontendHeartbeatExtensions::default(); + assert_eq!( + extensions.try_register(Arc::new(TestExtension("caller"))), + Ok(()) + ); + let mut plugins = plugins_with_heartbeat_extensions(extensions.clone()); + + let finalized = setup_and_freeze_frontend_heartbeat_extensions(&mut plugins, |plugins| { + Box::pin(async move { + let setup_extensions = plugins.get_or_insert(FrontendHeartbeatExtensions::default); + assert_eq!( + setup_extensions.try_register(Arc::new(TestExtension("plugin"))), + Ok(()) + ); + Ok(()) + }) + }) + .await + .unwrap(); + + assert_eq!( + finalized + .extensions() + .iter() + .map(|extension| extension.name()) + .collect::>(), + ["caller", "plugin"] + ); + assert_eq!( + extensions.try_register(Arc::new(TestExtension("late"))), + Err(frontend::heartbeat::RegistrationError::Frozen) + ); + assert_eq!(finalized.len(), 2); + } + + fn assert_registry_identity( + expected: &FrontendHeartbeatExtensions, + actual: &FrontendHeartbeatExtensions, + ) { + let extension = Arc::new(TestExtension("cmd-test-extension")); + assert!(expected.register(extension.clone())); + let registered = actual.extensions(); + assert_eq!(registered.len(), 1); + assert!(Arc::ptr_eq( + &(extension as Arc), + ®istered[0] + )); + } + #[test] fn test_try_from_start_command() { let command = StartCommand { diff --git a/src/frontend/src/heartbeat.rs b/src/frontend/src/heartbeat.rs index 0d9cde74fbe..b7e3a7aff48 100644 --- a/src/frontend/src/heartbeat.rs +++ b/src/frontend/src/heartbeat.rs @@ -16,8 +16,9 @@ mod tests; use std::collections::{HashMap, HashSet}; +use std::fmt::{Display, Formatter}; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::sync::{Arc, Mutex as StdMutex}; +use std::sync::{Arc, Mutex as StdMutex, MutexGuard, PoisonError}; use api::v1::meta::heartbeat_request::NodeWorkloads; use api::v1::meta::{FrontendWorkloads, HeartbeatRequest, HeartbeatResponse, NodeInfo, Peer}; @@ -93,34 +94,88 @@ pub trait FrontendHeartbeatExtension: Send + Sync { } /// A shareable, ordered registry of frontend heartbeat extensions. +/// +/// Registration is available until [`FrontendHeartbeatExtensions::freeze`] is called. After that, +/// membership is immutable so every heartbeat consumer observes the same extensions. #[derive(Clone, Default)] pub struct FrontendHeartbeatExtensions { inner: Arc>, } +/// The reason a frontend heartbeat extension was not registered. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RegistrationError { + /// An extension with the same name is already registered. + Duplicate, + /// The registry is frozen and no longer accepts registrations. + Frozen, +} + +impl Display for RegistrationError { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + match self { + Self::Duplicate => formatter.write_str("heartbeat extension is already registered"), + Self::Frozen => formatter.write_str("heartbeat extension registry is frozen"), + } + } +} + +impl std::error::Error for RegistrationError {} + #[derive(Default)] struct FrontendHeartbeatExtensionsInner { names: HashSet, extensions: Vec>, + frozen: bool, } impl FrontendHeartbeatExtensions { + fn lock(&self) -> MutexGuard<'_, FrontendHeartbeatExtensionsInner> { + self.inner.lock().unwrap_or_else(PoisonError::into_inner) + } + /// Registers an extension without replacing an existing extension of the same name. /// - /// Returns `true` for a new registration and `false` for an idempotent duplicate. + /// Returns `true` for a new registration and `false` for an idempotent duplicate or when the + /// registry is frozen. Use [`FrontendHeartbeatExtensions::try_register`] when the rejection + /// reason must not be discarded. pub fn register(&self, extension: Arc) -> bool { - let mut inner = self.inner.lock().unwrap(); + self.try_register(extension).is_ok() + } + + /// Registers an extension and returns the reason when registration is rejected. + /// + /// If the registry is already frozen, this method returns [`RegistrationError::Frozen`] + /// without calling [`FrontendHeartbeatExtension::name`]. If freezing races with name + /// evaluation, the registry state is checked again before insertion. + pub fn try_register( + &self, + extension: Arc, + ) -> std::result::Result<(), RegistrationError> { + if self.lock().frozen { + return Err(RegistrationError::Frozen); + } + let name = extension.name().to_string(); + let mut inner = self.lock(); + if inner.frozen { + return Err(RegistrationError::Frozen); + } if !inner.names.insert(name) { - return false; + return Err(RegistrationError::Duplicate); } inner.extensions.push(extension); - true + Ok(()) + } + + /// Prevents further registrations while preserving the current registration order. + pub fn freeze(&self) { + self.lock().frozen = true; } /// Returns the registered extensions in registration order. pub fn extensions(&self) -> Vec> { - self.inner.lock().unwrap().extensions.clone() + self.lock().extensions.clone() } /// Returns response handlers in extension registration order. @@ -133,7 +188,7 @@ impl FrontendHeartbeatExtensions { /// Returns the number of registered extensions. pub fn len(&self) -> usize { - self.inner.lock().unwrap().extensions.len() + self.lock().extensions.len() } /// Returns whether no extension is registered. diff --git a/src/frontend/src/heartbeat/tests.rs b/src/frontend/src/heartbeat/tests.rs index d3c651de326..8f5bfa614b2 100644 --- a/src/frontend/src/heartbeat/tests.rs +++ b/src/frontend/src/heartbeat/tests.rs @@ -14,8 +14,9 @@ use std::collections::{HashMap, VecDeque}; use std::future::pending; +use std::panic::{AssertUnwindSafe, catch_unwind}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Barrier, Mutex}; use std::time::Duration; use api::v1::meta::heartbeat_request::NodeWorkloads; @@ -494,6 +495,170 @@ impl FrontendHeartbeatExtension for TestExtension { } } +struct PanickingNameExtension; + +#[async_trait] +impl FrontendHeartbeatExtension for PanickingNameExtension { + fn name(&self) -> &str { + panic!("extension name panic") + } +} + +#[test] +fn test_panicking_extension_name_does_not_corrupt_registry() { + let extensions = FrontendHeartbeatExtensions::default(); + + let result = catch_unwind(AssertUnwindSafe(|| { + extensions.register(Arc::new(PanickingNameExtension)); + })); + assert!(result.is_err()); + + let valid = TestExtension::new("valid"); + assert!(extensions.register(valid.clone())); + let registered = extensions.extensions(); + assert_eq!(registered.len(), 1); + assert!(Arc::ptr_eq( + &(valid as Arc), + ®istered[0] + )); + assert_eq!(extensions.len(), 1); +} + +#[test] +fn test_registry_recovers_from_poisoned_lock() { + let extensions = FrontendHeartbeatExtensions::default(); + let poisoned = extensions.clone(); + + let result = catch_unwind(AssertUnwindSafe(move || { + let _guard = poisoned.lock(); + panic!("poison registry lock") + })); + assert!(result.is_err()); + + let valid = TestExtension::new("valid"); + assert!(extensions.register(valid.clone())); + let registered = extensions.extensions(); + assert_eq!(registered.len(), 1); + assert!(Arc::ptr_eq( + &(valid as Arc), + ®istered[0] + )); + assert_eq!(extensions.len(), 1); +} + +#[test] +fn test_registry_freeze_rejects_late_registration_without_changing_membership() { + let extensions = FrontendHeartbeatExtensions::default(); + let calls = Arc::new(AtomicUsize::new(0)); + let handler: HeartbeatResponseHandlerRef = Arc::new(ShortCircuitHandler { + result: ShortCircuitResult::Done, + calls, + }); + let registered = TestExtension::with_handler("registered", handler.clone()); + + assert!(extensions.register(registered.clone())); + assert!(!extensions.register(TestExtension::new("registered"))); + extensions.freeze(); + + assert!(!extensions.register(TestExtension::new("late"))); + let frozen_extensions = extensions.extensions(); + assert_eq!(frozen_extensions.len(), 1); + assert!(Arc::ptr_eq( + &(registered as Arc), + &frozen_extensions[0] + )); + let frozen_handlers = extensions.response_handlers(); + assert_eq!(frozen_handlers.len(), 1); + assert!(Arc::ptr_eq(&handler, &frozen_handlers[0])); + assert_eq!(extensions.len(), 1); +} + +#[test] +fn test_try_register_distinguishes_duplicate_and_frozen() { + let extensions = FrontendHeartbeatExtensions::default(); + + assert_eq!( + extensions.try_register(TestExtension::new("registered")), + Ok(()) + ); + assert_eq!( + extensions.try_register(TestExtension::new("registered")), + Err(RegistrationError::Duplicate) + ); + + extensions.freeze(); + assert_eq!( + extensions.try_register(TestExtension::new("late")), + Err(RegistrationError::Frozen) + ); +} + +#[test] +fn test_frozen_registration_does_not_call_extension_name() { + let extensions = FrontendHeartbeatExtensions::default(); + extensions.freeze(); + + let compatible_result = catch_unwind(AssertUnwindSafe(|| { + extensions.register(Arc::new(PanickingNameExtension)) + })); + assert!(!compatible_result.unwrap()); + + let typed_result = catch_unwind(AssertUnwindSafe(|| { + extensions.try_register(Arc::new(PanickingNameExtension)) + })); + + assert_eq!(typed_result.unwrap(), Err(RegistrationError::Frozen)); + assert!(extensions.is_empty()); +} + +struct BlockingNameExtension { + entered_name: Arc, + release_name: Arc, +} + +#[async_trait] +impl FrontendHeartbeatExtension for BlockingNameExtension { + fn name(&self) -> &str { + self.entered_name.wait(); + self.release_name.wait(); + "blocking" + } +} + +#[test] +fn test_freeze_wins_while_registration_name_is_running() { + let extensions = FrontendHeartbeatExtensions::default(); + let entered_name = Arc::new(Barrier::new(2)); + let release_name = Arc::new(Barrier::new(2)); + let registering = extensions.clone(); + let extension = Arc::new(BlockingNameExtension { + entered_name: entered_name.clone(), + release_name: release_name.clone(), + }); + + let registration = std::thread::spawn(move || registering.try_register(extension)); + entered_name.wait(); + let freezing = extensions.clone(); + let (freeze_done_tx, freeze_done_rx) = std::sync::mpsc::sync_channel(1); + let freeze = std::thread::spawn(move || { + freezing.freeze(); + let _ = freeze_done_tx.send(()); + }); + if let Err(error) = freeze_done_rx.recv_timeout(Duration::from_secs(2)) { + release_name.wait(); + let registration = registration.join(); + let freeze = freeze.join(); + panic!( + "freeze did not complete while extension name was blocked: {error:?}; registration: {registration:?}; freeze: {freeze:?}" + ); + } + release_name.wait(); + + freeze.join().unwrap(); + assert_eq!(registration.join().unwrap(), Err(RegistrationError::Frozen)); + assert!(extensions.is_empty()); +} + fn test_task( connector: Arc, extensions: FrontendHeartbeatExtensions, From aaa843104bb2dba1d774b53dbff3adcf1d18eba7 Mon Sep 17 00:00:00 2001 From: wy471x Date: Fri, 28 Aug 2026 02:40:25 +0000 Subject: [PATCH 117/126] fix(mysql): interpret prepared statement datetime params in session timezone (#8923) * fix(mysql): interpret prepared statement datetime params in session timezone Binary DATETIME parameters of server-side prepared statements were converted as if UTC, ignoring the session timezone set via SET time_zone. Convert them with the session timezone and add an integration test covering prepared inserts and predicates under Asia/Shanghai. Signed-off-by: wy471x * refactor: share naive datetime timezone policy via common-time Address review feedback on the prepared-statement timezone fix: - Expose Timestamp::from_naive_datetime in common-time so the DST policy (gap -> error, ambiguous -> earlier instant) lives in one place, shared by the text protocol (Timestamp::from_str) and the MySQL binary protocol. - Route the MySQL prepared-statement datetime conversion through it. - Match the target type before converting datetime params so PreparedStmtTypeMismatch fails fast without wasted conversion. - Use the short Timezone import form for consistency with the rest of servers. Signed-off-by: wy471x --------- Signed-off-by: wy471x Co-authored-by: Ning Sun --- src/common/time/src/timestamp.rs | 75 ++++++++++++-- src/servers/src/mysql/handler.rs | 12 ++- src/servers/src/mysql/helper.rs | 171 +++++++++++++++++++++++++------ tests-integration/tests/sql.rs | 72 +++++++++++++ 4 files changed, 286 insertions(+), 44 deletions(-) diff --git a/src/common/time/src/timestamp.rs b/src/common/time/src/timestamp.rs index 7daa793bafe..88cac7f4fcc 100644 --- a/src/common/time/src/timestamp.rs +++ b/src/common/time/src/timestamp.rs @@ -483,6 +483,31 @@ impl Timestamp { ParseTimestampSnafu { raw: s }.fail() } + /// Interprets a timezone-less [`NaiveDateTime`] in the given timezone and + /// returns the corresponding timestamp. + /// + /// Datetimes that fall into a DST gap (a local time that does not exist) + /// are rejected, while ambiguous datetimes (from a repeated local time) + /// are resolved to the earlier instant. This policy is shared with + /// [`Timestamp::from_str`] so that the text and binary protocols interpret + /// datetimes consistently. + pub fn from_naive_datetime( + datetime: NaiveDateTime, + timezone: &Timezone, + ) -> crate::error::Result { + match datetime_to_utc(&datetime, timezone) { + LocalResult::Single(utc) | LocalResult::Ambiguous(utc, _) => { + Timestamp::from_chrono_datetime(utc).context(ParseTimestampSnafu { + raw: format!("{datetime} (timezone {timezone})"), + }) + } + LocalResult::None => ParseTimestampSnafu { + raw: format!("{datetime} (timezone {timezone})"), + } + .fail(), + } + } + pub fn negative(mut self) -> Self { self.value = -self.value; self @@ -531,12 +556,8 @@ fn naive_datetime_to_timestamp( .context(ParseTimestampSnafu { raw: s }); }; - match datetime_to_utc(&datetime, timezone) { - LocalResult::None => ParseTimestampSnafu { raw: s }.fail(), - LocalResult::Single(utc) | LocalResult::Ambiguous(utc, _) => { - Timestamp::from_chrono_datetime(utc).context(ParseTimestampSnafu { raw: s }) - } - } + Timestamp::from_naive_datetime(datetime, timezone) + .map_err(|_| ParseTimestampSnafu { raw: s }.build()) } impl From for Timestamp { @@ -919,6 +940,48 @@ mod tests { ); } + #[test] + fn test_from_naive_datetime() { + let datetime = NaiveDate::from_ymd_opt(2026, 8, 13) + .unwrap() + .and_hms_opt(8, 0, 0) + .unwrap(); + + // A fixed-offset timezone shifts the datetime by a constant amount. + let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + assert_eq!( + "2026-08-13 00:00:00", + Timestamp::from_naive_datetime(datetime, &shanghai) + .unwrap() + .to_chrono_datetime() + .unwrap() + .to_string() + ); + + // 2026-03-08 02:30 does not exist in America/New_York (DST gap). + let new_york = Timezone::from_tz_string("America/New_York").unwrap(); + let gap = NaiveDate::from_ymd_opt(2026, 3, 8) + .unwrap() + .and_hms_opt(2, 30, 0) + .unwrap(); + assert!(Timestamp::from_naive_datetime(gap, &new_york).is_err()); + + // 2026-11-01 01:30 is ambiguous in America/New_York; picks the first + // instant (EDT, UTC-4). + let ambiguous = NaiveDate::from_ymd_opt(2026, 11, 1) + .unwrap() + .and_hms_opt(1, 30, 0) + .unwrap(); + assert_eq!( + "2026-11-01 05:30:00", + Timestamp::from_naive_datetime(ambiguous, &new_york) + .unwrap() + .to_chrono_datetime() + .unwrap() + .to_string() + ); + } + #[test] fn test_to_iso8601_string() { set_default_timezone(Some("Asia/Shanghai")).unwrap(); diff --git a/src/servers/src/mysql/handler.rs b/src/servers/src/mysql/handler.rs index b0eb33f8fd5..30579d8846a 100644 --- a/src/servers/src/mysql/handler.rs +++ b/src/servers/src/mysql/handler.rs @@ -25,6 +25,7 @@ use common_catalog::parse_optional_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; use common_query::Output; use common_telemetry::{debug, error, tracing, warn}; +use common_time::Timezone; use datafusion_common::ParamValues; use datafusion_expr::LogicalPlan; use datatypes::prelude::ConcreteDataType; @@ -289,12 +290,13 @@ impl MysqlInstanceShim { .fail(); } + let timezone = query_ctx.timezone(); let replaced_plan = match params { Params::ProtocolParams(params) => { - replace_params_with_values(&plan, param_types, ¶ms) + replace_params_with_values(&plan, param_types, ¶ms, &timezone) } Params::CliParams(params) => { - replace_params_with_exprs(&plan, param_types, ¶ms) + replace_params_with_exprs(&plan, param_types, ¶ms, &timezone) } }?; @@ -807,6 +809,7 @@ fn replace_params_with_values( plan: &LogicalPlan, param_types: HashMap>, params: &[ParamValue], + timezone: &Timezone, ) -> Result { debug_assert_eq!(param_types.len(), params.len()); @@ -824,7 +827,7 @@ fn replace_params_with_values( for (i, param) in params.iter().enumerate() { if let Some(Some(t)) = param_types.get(&format_placeholder(i + 1)) { - let value = helper::convert_value(param, t)?; + let value = helper::convert_value(param, t, timezone)?; values.push(value.into()); } @@ -839,6 +842,7 @@ fn replace_params_with_exprs( plan: &LogicalPlan, param_types: HashMap>, params: &[sql::ast::Expr], + timezone: &Timezone, ) -> Result { debug_assert_eq!(param_types.len(), params.len()); @@ -853,7 +857,7 @@ fn replace_params_with_exprs( for (i, param) in params.iter().enumerate() { if let Some(Some(t)) = param_types.get(&format_placeholder(i + 1)) { - let value = helper::convert_expr_to_scalar_value(param, t)?; + let value = helper::convert_expr_to_scalar_value(param, t, timezone)?; values.push(value.into()); } diff --git a/src/servers/src/mysql/helper.rs b/src/servers/src/mysql/helper.rs index a4f289cdfb1..f9ab5325925 100644 --- a/src/servers/src/mysql/helper.rs +++ b/src/servers/src/mysql/helper.rs @@ -15,10 +15,11 @@ use std::ops::ControlFlow; use std::time::Duration; -use chrono::NaiveDate; +use chrono::{NaiveDate, NaiveDateTime}; use common_query::prelude::ScalarValue; use common_sql::convert::sql_value_to_value; -use common_time::{Date, Timestamp}; +use common_time::timestamp::TimeUnit; +use common_time::{Date, Timestamp, Timezone}; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::schema::ColumnSchema; use datatypes::types::TimestampType; @@ -137,11 +138,15 @@ where /// Convert [`ParamValue`] into [`Value`] according to param type. /// It will try it's best to do type conversions if possible -pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result { +pub fn convert_value( + param: &ParamValue, + t: &ConcreteDataType, + timezone: &Timezone, +) -> Result { if let ConcreteDataType::Dictionary(dictionary) = t { return Ok(ScalarValue::Dictionary( Box::new(dictionary.key_type().as_arrow_type()), - Box::new(convert_value(param, dictionary.value_type())?), + Box::new(convert_value(param, dictionary.value_type(), timezone)?), )); } @@ -221,7 +226,9 @@ pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result Ok(ScalarValue::Binary(Some(b.to_vec()))), - ConcreteDataType::Timestamp(ts_type) => convert_bytes_to_timestamp(b, ts_type), + ConcreteDataType::Timestamp(ts_type) => { + convert_bytes_to_timestamp(b, ts_type, timezone) + } ConcreteDataType::Date(_) => convert_bytes_to_date(b), _ => error::PreparedStmtTypeMismatchSnafu { expected: t, @@ -233,29 +240,26 @@ pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result { - let timestamp_millis = to_naive_datetime(param.value) - .map_err(|e| { + ValueInner::Datetime(_) => match t { + ConcreteDataType::Timestamp(_) => { + let datetime = to_naive_datetime(param.value).map_err(|e| { error::MysqlValueConversionSnafu { err_msg: e.to_string(), } .build() - })? - .and_utc() - .timestamp_millis(); - - match t { - ConcreteDataType::Timestamp(_) => Ok(ScalarValue::TimestampMillisecond( + })?; + let timestamp_millis = datetime_to_timestamp_millis(datetime, timezone)?; + Ok(ScalarValue::TimestampMillisecond( Some(timestamp_millis), None, - )), - _ => error::PreparedStmtTypeMismatchSnafu { - expected: t, - actual: param.coltype, - } - .fail(), + )) } - } + _ => error::PreparedStmtTypeMismatchSnafu { + expected: t, + actual: param.coltype, + } + .fail(), + }, ValueInner::Time(_) => Ok(ScalarValue::Time64Nanosecond(Some( Duration::from(param.value).as_millis() as i64, ))), @@ -264,13 +268,18 @@ pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result Result { +pub fn convert_expr_to_scalar_value( + param: &Expr, + t: &ConcreteDataType, + timezone: &Timezone, +) -> Result { if let ConcreteDataType::Dictionary(dictionary) = t { return Ok(ScalarValue::Dictionary( Box::new(dictionary.key_type().as_arrow_type()), Box::new(convert_expr_to_scalar_value( param, dictionary.value_type(), + timezone, )?), )); } @@ -278,7 +287,7 @@ pub fn convert_expr_to_scalar_value(param: &Expr, t: &ConcreteDataType) -> Resul let column_schema = ColumnSchema::new("", t.clone(), true); match param { Expr::Value(v) => { - let v = sql_value_to_value(&column_schema, &v.value, None, None, true); + let v = sql_value_to_value(&column_schema, &v.value, Some(timezone), None, true); match v { Ok(v) => v .try_to_scalar_value(t) @@ -290,7 +299,7 @@ pub fn convert_expr_to_scalar_value(param: &Expr, t: &ConcreteDataType) -> Resul } } Expr::UnaryOp { op, expr } if let Expr::Value(v) = &**expr => { - let v = sql_value_to_value(&column_schema, &v.value, None, Some(*op), true); + let v = sql_value_to_value(&column_schema, &v.value, Some(timezone), Some(*op), true); match v { Ok(v) => v .try_to_scalar_value(t) @@ -308,8 +317,32 @@ pub fn convert_expr_to_scalar_value(param: &Expr, t: &ConcreteDataType) -> Resul } } -fn convert_bytes_to_timestamp(bytes: &[u8], ts_type: &TimestampType) -> Result { - let ts = Timestamp::from_str_utc(&String::from_utf8_lossy(bytes)) +/// Interprets a timezone-less datetime in the given timezone and returns the +/// corresponding epoch timestamp in milliseconds. +fn datetime_to_timestamp_millis(datetime: NaiveDateTime, timezone: &Timezone) -> Result { + let ts = Timestamp::from_naive_datetime(datetime, timezone) + .map_err(|e| { + error::MysqlValueConversionSnafu { + err_msg: e.to_string(), + } + .build() + })? + .convert_to(TimeUnit::Millisecond) + .ok_or_else(|| { + error::MysqlValueConversionSnafu { + err_msg: "Overflow when converting datetime to milliseconds".to_string(), + } + .build() + })?; + Ok(ts.value()) +} + +fn convert_bytes_to_timestamp( + bytes: &[u8], + ts_type: &TimestampType, + timezone: &Timezone, +) -> Result { + let ts = Timestamp::from_str(&String::from_utf8_lossy(bytes), Some(timezone)) .map_err(|e| { error::MysqlValueConversionSnafu { err_msg: e.to_string(), @@ -446,19 +479,20 @@ mod tests { #[test] fn test_convert_expr_to_scalar_value() { + let utc = Timezone::from_tz_string("UTC").unwrap(); let expr = Expr::Value(ValueExpr::Number("123".to_string(), false).into()); let t = ConcreteDataType::int32_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Int32(Some(123)), v); let expr = Expr::Value(ValueExpr::Number("123.456789".to_string(), false).into()); let t = ConcreteDataType::float64_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Float64(Some(123.456789)), v); let expr = Expr::Value(ValueExpr::SingleQuotedString("2001-01-02".to_string()).into()); let t = ConcreteDataType::date_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); let scalar_v = ScalarValue::Utf8(Some("2001-01-02".to_string())) .cast_to(&arrow_schema::DataType::Date32) .unwrap(); @@ -467,7 +501,7 @@ mod tests { let expr = Expr::Value(ValueExpr::SingleQuotedString("2001-01-02 03:04:05".to_string()).into()); let t = ConcreteDataType::timestamp_microsecond_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); let scalar_v = ScalarValue::Utf8(Some("2001-01-02 03:04:05".to_string())) .cast_to(&arrow_schema::DataType::Timestamp( arrow_schema::TimeUnit::Microsecond, @@ -478,14 +512,14 @@ mod tests { let expr = Expr::Value(ValueExpr::SingleQuotedString("hello".to_string()).into()); let t = ConcreteDataType::string_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Utf8(Some("hello".to_string())), v); let t = ConcreteDataType::dictionary_datatype( ConcreteDataType::uint32_datatype(), ConcreteDataType::string_datatype(), ); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!( ScalarValue::Dictionary( Box::new(arrow_schema::DataType::UInt32), @@ -496,7 +530,7 @@ mod tests { let expr = Expr::Value(ValueExpr::Null.into()); let t = ConcreteDataType::time_microsecond_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Time64Microsecond(None), v); } @@ -577,12 +611,81 @@ mod tests { ), ]; + let utc = Timezone::from_tz_string("UTC").unwrap(); for (input, ts_type, expected) in test_cases { - let result = convert_bytes_to_timestamp(input.as_bytes(), &ts_type).unwrap(); + let result = convert_bytes_to_timestamp(input.as_bytes(), &ts_type, &utc).unwrap(); assert_eq!(result, expected); } } + fn utc_millis(year: i32, month: u32, day: u32, hour: u32, minute: u32, second: u32) -> i64 { + NaiveDate::from_ymd_opt(year, month, day) + .unwrap() + .and_hms_opt(hour, minute, second) + .unwrap() + .and_utc() + .timestamp_millis() + } + + #[test] + fn test_datetime_to_timestamp_millis() { + let datetime = NaiveDate::from_ymd_opt(2026, 8, 13) + .unwrap() + .and_hms_opt(8, 0, 0) + .unwrap(); + + let utc = Timezone::from_tz_string("UTC").unwrap(); + assert_eq!( + datetime_to_timestamp_millis(datetime, &utc).unwrap(), + utc_millis(2026, 8, 13, 8, 0, 0) + ); + + let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + assert_eq!( + datetime_to_timestamp_millis(datetime, &shanghai).unwrap(), + utc_millis(2026, 8, 13, 0, 0, 0) + ); + + let minus_7 = Timezone::from_tz_string("-07:00").unwrap(); + assert_eq!( + datetime_to_timestamp_millis(datetime, &minus_7).unwrap(), + utc_millis(2026, 8, 13, 15, 0, 0) + ); + + // 2026-03-08 02:30 does not exist in America/New_York (DST gap). + let new_york = Timezone::from_tz_string("America/New_York").unwrap(); + let gap = NaiveDate::from_ymd_opt(2026, 3, 8) + .unwrap() + .and_hms_opt(2, 30, 0) + .unwrap(); + assert!(datetime_to_timestamp_millis(gap, &new_york).is_err()); + + // 2026-11-01 01:30 is ambiguous in America/New_York; picks the first (EDT, UTC-4). + let ambiguous = NaiveDate::from_ymd_opt(2026, 11, 1) + .unwrap() + .and_hms_opt(1, 30, 0) + .unwrap(); + assert_eq!( + datetime_to_timestamp_millis(ambiguous, &new_york).unwrap(), + utc_millis(2026, 11, 1, 5, 30, 0) + ); + } + + #[test] + fn test_convert_bytes_to_timestamp_with_timezone() { + let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + let result = convert_bytes_to_timestamp( + "2026-08-13 08:00:00".as_bytes(), + &TimestampType::Millisecond(TimestampMillisecondType), + &shanghai, + ) + .unwrap(); + assert_eq!( + result, + ScalarValue::TimestampMillisecond(Some(utc_millis(2026, 8, 13, 0, 0, 0)), None) + ); + } + #[test] fn test_convert_bytes_to_date() { let test_cases = vec![ diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index e46a8f58d1b..32a2881dd78 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -89,6 +89,7 @@ macro_rules! sql_tests { test_postgres_explain_bind_parameter, test_postgres_array_types, test_mysql_prepare_stmt_insert_timestamp, + test_mysql_prepare_stmt_timezone, test_mysql_federated_prepare_stmt, test_declare_fetch_close_cursor, test_alter_update_on, @@ -1779,6 +1780,77 @@ pub async fn test_mysql_prepare_stmt_insert_timestamp(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_mysql_prepare_stmt_timezone(store_type: StorageType) { + let (mut guard, server) = + setup_mysql_server(store_type, "test_mysql_prepare_stmt_timezone").await; + let addr = server.bind_addr().unwrap().to_string(); + + let mut conn = MySqlConnection::connect(&format!("mysql://{addr}/public")) + .await + .unwrap(); + + conn.execute("create table demo(i bigint, ts timestamp time index)") + .await + .unwrap(); + conn.execute("SET time_zone = 'Asia/Shanghai'") + .await + .unwrap(); + + // Server-side prepared statement: the binary DATETIME parameter must be + // interpreted in the session timezone. + sqlx::query("insert into demo values(?, ?)") + .bind(1) + .bind( + NaiveDate::from_ymd_opt(2026, 8, 13) + .and_then(|x| x.and_hms_opt(8, 0, 0)) + .unwrap(), + ) + .execute(&mut conn) + .await + .unwrap(); + + // Text protocol with an equivalent timezone-less literal one hour later. + sqlx::query("insert into demo values(2, '2026-08-13 09:00:00')") + .execute(&mut conn) + .await + .unwrap(); + + // Timestamps are read back in the session timezone, and sqlx parses the + // timezone-less wire representation as UTC. + let rows = sqlx::query("select i, ts from demo order by i") + .fetch_all(&mut conn) + .await + .unwrap(); + assert_eq!(rows.len(), 2); + let ts: DateTime = rows[0].get("ts"); + assert_eq!(ts.to_string(), "2026-08-13 08:00:00 UTC"); + let ts: DateTime = rows[1].get("ts"); + assert_eq!(ts.to_string(), "2026-08-13 09:00:00 UTC"); + + // The prepared predicate compares the same instant as the text literal. + let rows = sqlx::query("select i from demo where ts = ? order by i") + .bind( + NaiveDate::from_ymd_opt(2026, 8, 13) + .and_then(|x| x.and_hms_opt(8, 0, 0)) + .unwrap(), + ) + .fetch_all(&mut conn) + .await + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::("i"), 1); + + let rows = sqlx::query("select i from demo where ts = '2026-08-13 08:00:00'") + .fetch_all(&mut conn) + .await + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::("i"), 1); + + let _ = server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_mysql_federated_prepare_stmt(store_type: StorageType) { common_telemetry::init_default_ut_logging(); From bd7d2c1dfa166f579bc33c496d50d41eb29b775f Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 28 Aug 2026 03:45:41 +0000 Subject: [PATCH 118/126] fix(mito2): use target sequence for foreign SSTs (#8946) * fix(mito2): use target sequence for foreign SSTs Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore(mito2): address foreign SST review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/engine/apply_staging_manifest_test.rs | 215 ++++++++++++++++- src/mito2/src/sst/index.rs | 86 ++++++- src/mito2/src/sst/parquet.rs | 227 ++++++++++++------ src/mito2/src/sst/parquet/reader.rs | 21 +- src/mito2/src/worker/handle_rebuild_index.rs | 1 + 5 files changed, 466 insertions(+), 84 deletions(-) diff --git a/src/mito2/src/engine/apply_staging_manifest_test.rs b/src/mito2/src/engine/apply_staging_manifest_test.rs index e10ed0cbf00..49cc97c7d51 100644 --- a/src/mito2/src/engine/apply_staging_manifest_test.rs +++ b/src/mito2/src/engine/apply_staging_manifest_test.rs @@ -12,20 +12,24 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashSet; use std::sync::Arc; use std::{assert_matches, fs}; use api::v1::Rows; +use api::v1::region::{StrictWindow, compact_request}; use common_function::utils::partition_expr_version; use common_recordbatch::RecordBatches; +use datatypes::arrow::array::AsArray; +use datatypes::arrow::datatypes::Float64Type; use datatypes::value::Value; use partition::expr::{PartitionExpr, col}; use store_api::region_engine::{ RegionEngine, RegionRole, RemapManifestsRequest, SettableRegionRoleState, }; use store_api::region_request::{ - ApplyStagingManifestRequest, EnterStagingRequest, RegionFlushRequest, RegionPutRequest, - RegionRequest, StagingPartitionDirective, + ApplyStagingManifestRequest, EnterStagingRequest, RegionCompactRequest, RegionFlushRequest, + RegionPutRequest, RegionRequest, StagingPartitionDirective, }; use store_api::storage::{FileId, RegionId}; @@ -37,7 +41,9 @@ use crate::manifest::action::{ }; use crate::sst::FormatType; use crate::sst::file::FileMeta; -use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema}; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows, build_rows_for_key, put_rows, rows_schema, +}; fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr { col(col_name) @@ -45,6 +51,209 @@ fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr { .and(col(col_name).lt(Value::Int64(end))) } +#[tokio::test] +async fn test_apply_staging_manifest_sequence_domain() { + common_telemetry::init_default_ut_logging(); + test_apply_staging_manifest_sequence_domain_with_format(false).await; + test_apply_staging_manifest_sequence_domain_with_format(true).await; +} + +async fn test_apply_staging_manifest_sequence_domain_with_format(flat_format: bool) { + let mut env = TestEnv::with_prefix("apply-staging-sequence-domain").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let source = RegionId::new(1, 1); + let target = RegionId::new(1, 2); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + + engine + .handle_request(source, RegionRequest::Create(request.clone())) + .await + .unwrap(); + for value in 0..3 { + put_rows( + &engine, + source, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("0", 0, 1, value), + }, + ) + .await; + } + engine + .handle_request(source, RegionRequest::Flush(RegionFlushRequest::default())) + .await + .unwrap(); + let source_manifest = engine + .get_region(source) + .unwrap() + .manifest_ctx + .manifest() + .await; + assert_eq!(source_manifest.files.len(), 1); + assert_eq!( + source_manifest.files.values().next().unwrap().sequence, + Some(std::num::NonZeroU64::new(3).unwrap()) + ); + + engine + .set_region_role_state_gracefully(source, SettableRegionRoleState::StagingLeader) + .await + .unwrap(); + let partition_expr = float_range_expr("field_0", 0.1, 100.1) + .as_json_str() + .unwrap(); + let result = engine + .remap_manifests(RemapManifestsRequest { + region_id: source, + input_regions: vec![source], + region_mapping: [(source, vec![target])].into_iter().collect(), + new_partition_exprs: [(target, partition_expr.clone())].into_iter().collect(), + }) + .await + .unwrap(); + engine + .handle_request(target, RegionRequest::Create(request.clone())) + .await + .unwrap(); + engine + .handle_request( + target, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::UpdatePartitionExpr( + partition_expr.clone(), + ), + }), + ) + .await + .unwrap(); + engine + .handle_request( + target, + RegionRequest::ApplyStagingManifest(ApplyStagingManifestRequest { + partition_expr, + central_region_id: source, + manifest_path: result.manifest_paths[&target].clone(), + }), + ) + .await + .unwrap(); + + let manifest = engine + .get_region(target) + .unwrap() + .manifest_ctx + .manifest() + .await; + assert_eq!(manifest.files.len(), 1); + assert_eq!(manifest.committed_sequence, Some(1)); + let imported_file = manifest.files.values().next().unwrap(); + assert_eq!(imported_file.region_id, source); + assert_eq!( + imported_file.sequence, + Some(std::num::NonZeroU64::new(1).unwrap()) + ); + + put_rows( + &engine, + target, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("0", 0, 1, 99), + }, + ) + .await; + assert_target_value(&engine, target, 99.0).await; + + engine + .handle_request(target, RegionRequest::Flush(RegionFlushRequest::default())) + .await + .unwrap(); + let input_ids = current_file_ids(&engine, target); + assert_eq!( + input_ids.len(), + 2, + "imported and target-write SSTs must both exist" + ); + + engine + .handle_request( + target, + RegionRequest::Compact(RegionCompactRequest { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 60, + }), + parallelism: None, + time_range: None, + }), + ) + .await + .unwrap(); + let output_version = engine.get_region(target).unwrap().version(); + let output_files = output_version + .ssts + .levels() + .iter() + .flat_map(|level| level.files.values()) + .collect::>(); + let output_ids = output_files + .iter() + .map(|file| file.meta_ref().file_id) + .collect::>(); + assert_eq!(output_ids.len(), 1); + assert!( + output_files + .iter() + .all(|file| file.meta_ref().region_id == target) + ); + assert!( + input_ids.iter().all(|id| !output_ids.contains(id)), + "real compaction must replace both input SSTs" + ); + assert_target_value(&engine, target, 99.0).await; +} + +fn current_file_ids(engine: &crate::engine::MitoEngine, region_id: RegionId) -> HashSet { + engine + .get_region(region_id) + .unwrap() + .version() + .ssts + .levels() + .iter() + .flat_map(|level| level.files.values()) + .map(|file| file.meta_ref().file_id) + .collect() +} + +async fn assert_target_value( + engine: &crate::engine::MitoEngine, + region_id: RegionId, + expected: f64, +) { + let scan = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(scan).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 1 + ); + let batch = batches.iter().next().unwrap(); + let values = batch + .column_by_name("field_0") + .unwrap() + .as_primitive::(); + assert_eq!(values.value(0), expected); +} + fn float_range_expr(col_name: &str, start: f64, end: f64) -> PartitionExpr { col(col_name) .gt_eq(Value::Float64(start.into())) diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 581d1edebdf..7954af2b920 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -706,6 +706,14 @@ pub struct IndexBuildTask { pub region_id: RegionId, /// The SST file handle to build index for. pub file: FileHandle, + /// The target region metadata used to decode rows from the SST. + /// + /// An SST may originate in another region while being visible in the target + /// manifest. This metadata defines the target schema and sequence domain; + /// applying the staging manifest only makes imported files visible. Index + /// rebuild happens later when a flush, compaction, schema change, or manual + /// index build request schedules it. + pub(crate) target_region_metadata: RegionMetadataRef, /// The manifest state this build is based on. pub(crate) source: IndexBuildSource, pub reason: IndexBuildType, @@ -846,6 +854,7 @@ impl IndexBuildTask { let mut parquet_reader = self .access_layer .read_sst(self.file.clone()) // use the latest file handle instead of creating a new one + .expected_metadata(Some(self.target_region_metadata.clone())) .build() .await?; @@ -1494,8 +1503,12 @@ mod tests { use datatypes::schema::{ ColumnSchema, FulltextOptions, SkippingIndexOptions, SkippingIndexType, }; + use datatypes::value::Value; + use index::inverted_index::format::reader::InvertedIndexReader; use object_store::ObjectStore; use object_store::services::Memory; + use partition::expr::col; + use puffin::puffin_manager::{PuffinManager, PuffinReader}; use puffin_manager::PuffinManagerFactory; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use tokio::sync::mpsc; @@ -2088,6 +2101,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta, version_control.current().version.metadata.schema_version, @@ -2120,16 +2134,34 @@ mod tests { } #[tokio::test] - async fn test_index_build_task_increments_legacy_index_version() { + async fn test_index_build_task_foreign_file_uses_target_metadata() { let env = SchedulerEnv::new().await; let mut scheduler = env.mock_index_build_scheduler(4); - let metadata = Arc::new(sst_region_metadata()); - let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; - let region_id = metadata.region_id; + let source_metadata = Arc::new(sst_region_metadata()); + let mut target_metadata = (*source_metadata).clone(); + target_metadata.region_id = RegionId::new(1, 3); + let mut target_builder = RegionMetadataBuilder::new(target_metadata.region_id); + for mut column_metadata in target_metadata.column_metadatas.clone() { + if column_metadata.column_id == 2 { + column_metadata.column_schema = + column_metadata.column_schema.with_inverted_index(true); + } + target_builder.push_column_metadata(column_metadata); + } + let partition_expr = col("field_0") + .gt_eq(Value::UInt64(100)) + .and(col("field_0").lt(Value::UInt64(200))); + target_builder + .primary_key(target_metadata.primary_key.clone()) + .partition_expr_json(Some(partition_expr.as_json_str().unwrap())) + .bump_version(); + let target_metadata = Arc::new(target_builder.build().unwrap()); + let manifest_ctx = env.mock_manifest_context(target_metadata.clone()).await; + let region_id = target_metadata.region_id; let file_purger = Arc::new(NoopFilePurger {}); - let sst_info = mock_sst_file(metadata.clone(), &env, IndexBuildMode::Async).await; + let sst_info = mock_sst_file(source_metadata.clone(), &env, IndexBuildMode::Async).await; let file_meta = FileMeta { - region_id, + region_id: source_metadata.region_id, file_id: sst_info.file_id, file_size: sst_info.file_size, max_row_group_uncompressed_size: sst_info.max_row_group_uncompressed_size, @@ -2144,8 +2176,8 @@ mod tests { seed_manifest_file(&manifest_ctx, &file_meta).await; let files = HashMap::from([(file_meta.file_id, file_meta.clone())]); let version_control = - mock_version_control(metadata.clone(), file_purger.clone(), files).await; - let indexer_builder = mock_indexer_builder(metadata.clone(), &env).await; + mock_version_control(target_metadata.clone(), file_purger.clone(), files).await; + let indexer_builder = mock_indexer_builder(target_metadata.clone(), &env).await; let file = FileHandle::new(file_meta.clone(), file_purger.clone()); @@ -2155,6 +2187,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2199,9 +2232,40 @@ mod tests { assert!(updated_meta.index_file_size > 0); assert_eq!(updated_meta.file_id, file_meta.file_id); assert_eq!(updated_meta.index_version, 1); + let field_0_index = updated_meta + .indexes + .iter() + .find(|index| index.column_id == 2) + .expect("field_0 should have an inverted index"); + assert_eq!( + field_0_index.created_indexes.as_slice(), + [IndexType::InvertedIndex] + ); } _ => panic!("Unexpected worker request: {:?}", worker_req), } + + let puffin_reader = env + .access_layer + .build_puffin_manager() + .reader(&RegionIndexId::new( + RegionFileId::new(source_metadata.region_id, file_meta.file_id), + 1, + )) + .await + .unwrap(); + let blob = puffin_reader + .blob(inverted_index::INDEX_BLOB_TYPE) + .await + .unwrap(); + let blob_reader = blob.reader().await.unwrap(); + let index_metadata = + index::inverted_index::format::reader::InvertedIndexBlobReader::new(blob_reader) + .metadata(None) + .await + .unwrap(); + assert!(index_metadata.metas.contains_key("2")); + assert_eq!(index_metadata.total_row_count, 100); } async fn schedule_index_build_task_with_mode(build_mode: IndexBuildMode) { @@ -2236,6 +2300,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2346,6 +2411,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2444,6 +2510,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2505,7 +2572,7 @@ mod tests { let schema_version = metadata.schema_version; let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; let file_purger = Arc::new(NoopFilePurger {}); - let indexer_builder = mock_indexer_builder(metadata, env).await; + let indexer_builder = mock_indexer_builder(metadata.clone(), env).await; let (tx, _rx) = mpsc::channel(4); let (result_tx, result_rx) = mpsc::channel::>(4); @@ -2521,6 +2588,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: metadata, source: IndexBuildSource::new(file_meta, schema_version), reason, access_layer: env.access_layer.clone(), diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index e67dcf2bee5..8a05c099e98 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -118,8 +118,8 @@ mod tests { use datafusion_expr::{BinaryExpr, Expr, Literal, Operator, col, lit}; use datatypes::arrow; use datatypes::arrow::array::{ - ArrayRef, BinaryDictionaryBuilder, RecordBatch, StringArray, StringDictionaryBuilder, - TimestampMillisecondArray, UInt8Array, UInt64Array, + ArrayRef, AsArray, BinaryDictionaryBuilder, RecordBatch, StringArray, + StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array, }; use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type}; use datatypes::arrow::util::pretty::pretty_format_batches; @@ -1625,91 +1625,178 @@ mod tests { #[tokio::test] async fn test_read_with_override_sequence() { + test_read_with_override_sequence_with_format(false).await; + test_read_with_override_sequence_with_format(true).await; + } + + async fn test_read_with_override_sequence_with_format(flat_format: bool) { let mut env = TestEnv::new().await; let object_store = env.init_object_store_manager(); - let handle = sst_file_handle(0, 1000); - let file_path = FixedPathProvider { - region_file_id: handle.file_id(), - }; let metadata = Arc::new(sst_region_metadata()); - // Create batches with sequence 0 to trigger override functionality. - let source = new_flat_source_from_record_batches(vec![ - new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, 0), - new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, 0), - ]); + async fn read_sequences(builder: ParquetReaderBuilder) -> Vec { + let mut reader = builder.build().await.unwrap().unwrap(); + let mut sequences = Vec::new(); + while let Some(batch) = reader.next_record_batch().await.unwrap() { + let sequence = batch + .column(batch.num_columns() - 2) + .as_primitive::(); + sequences.extend((0..sequence.len()).map(|idx| sequence.value(idx))); + } + sequences + } - let write_opts = WriteOptions { - row_group_size: 50, - ..Default::default() - }; + async fn write_sst( + object_store: ObjectStore, + metadata: Arc, + handle: FileHandle, + flat_format: bool, + sequence: u64, + ) { + let file_path = FixedPathProvider { + region_file_id: handle.file_id(), + }; + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, sequence), + new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, sequence), + ]); + let write_opts = WriteOptions { + row_group_size: 50, + ..Default::default() + }; + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + object_store, + metadata, + IndexConfig::default(), + NoopIndexBuilder, + file_path, + &mut metrics, + ) + .await; + if flat_format { + writer + .write_all_flat(source, None, &write_opts) + .await + .unwrap(); + } else { + writer + .write_all_flat_as_primary_key(source, None, &write_opts) + .await + .unwrap(); + } + } - let mut metrics = Metrics::new(WriteType::Flush); - let mut writer = ParquetWriter::new_with_object_store( + let custom_sequence = 12345u64; + let local_zero_handle = sst_file_handle(0, 1000); + write_sst( object_store.clone(), metadata.clone(), - IndexConfig::default(), - NoopIndexBuilder, - file_path, - &mut metrics, + local_zero_handle.clone(), + flat_format, + 0, ) .await; - writer - .write_all_flat_as_primary_key(source, None, &write_opts) - .await - .unwrap() - .remove(0); + // Local all-zero SSTs retain the compatibility override. + let local_zero_none = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_zero_handle.clone(), + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!(local_zero_none.iter().all(|sequence| *sequence == 0)); - // Read without override sequence (should read sequence 0) - let builder = ParquetReaderBuilder::new( - FILE_DIR.to_string(), - PathType::Bare, - handle.clone(), - object_store.clone(), - ); - let mut reader = builder.build().await.unwrap().unwrap(); - let mut normal_batches = Vec::new(); - while let Some(batch) = reader.next_record_batch().await.unwrap() { - normal_batches.push(batch); - } - - // Read with override sequence using FileMeta.sequence - let custom_sequence = 12345u64; - let file_meta = handle.meta_ref(); - let mut override_file_meta = file_meta.clone(); - override_file_meta.sequence = Some(std::num::NonZero::new(custom_sequence).unwrap()); - let override_handle = FileHandle::new( - override_file_meta, + let mut local_zero_meta = local_zero_handle.meta_ref().clone(); + local_zero_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap()); + let local_zero_override_handle = FileHandle::new( + local_zero_meta, Arc::new(crate::sst::file_purger::NoopFilePurger), ); - - let builder = ParquetReaderBuilder::new( - FILE_DIR.to_string(), - PathType::Bare, - override_handle, - object_store.clone(), + let local_zero_override = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_zero_override_handle, + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!( + local_zero_override + .iter() + .all(|sequence| *sequence == custom_sequence) ); - let mut reader = builder.build().await.unwrap().unwrap(); - let mut override_batches = Vec::new(); - while let Some(batch) = reader.next_record_batch().await.unwrap() { - override_batches.push(batch); - } - // Compare the results - assert_eq!(normal_batches.len(), override_batches.len()); - for (normal, override_batch) in normal_batches.into_iter().zip(override_batches.iter()) { - let expected_batch = { - let mut columns = normal.columns().to_vec(); - let num_cols = columns.len(); - columns[num_cols - 2] = - Arc::new(UInt64Array::from_value(custom_sequence, normal.num_rows())); - RecordBatch::try_new(normal.schema(), columns).unwrap() - }; + let local_nonzero_handle = sst_file_handle(0, 1000); + write_sst( + object_store.clone(), + metadata.clone(), + local_nonzero_handle.clone(), + flat_format, + 7, + ) + .await; - // Override batch should match expected batch - assert_eq!(*override_batch, expected_batch); - } + // Local nonzero SSTs retain physical per-row sequences, even with FileMeta.sequence. + let mut local_nonzero_meta = local_nonzero_handle.meta_ref().clone(); + local_nonzero_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap()); + let local_nonzero_override_handle = FileHandle::new( + local_nonzero_meta, + Arc::new(crate::sst::file_purger::NoopFilePurger), + ); + let local_nonzero_override = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_nonzero_override_handle, + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!(local_nonzero_override.iter().all(|sequence| *sequence == 7)); + + // None never overrides a local nonzero physical sequence. + let local_nonzero_none = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_nonzero_handle.clone(), + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!(local_nonzero_none.iter().all(|sequence| *sequence == 7)); + + // A source-owned handle is foreign when read against target metadata, so the + // target-local manifest barrier is applied even for nonzero physical sequences. + let mut target_metadata = (*metadata).clone(); + target_metadata.region_id = RegionId::new(0, 1); + let target_metadata = Arc::new(target_metadata); + let mut foreign_meta = local_nonzero_handle.meta_ref().clone(); + foreign_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap()); + let foreign_handle = FileHandle::new( + foreign_meta, + Arc::new(crate::sst::file_purger::NoopFilePurger), + ); + let foreign = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + foreign_handle, + object_store, + ) + .expected_metadata(Some(target_metadata)), + ) + .await; + assert!(foreign.iter().all(|sequence| *sequence == custom_sequence)); } #[tokio::test] diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 4f890cb83e8..3adb335564b 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -22,7 +22,7 @@ use std::time::{Duration, Instant}; use api::v1::SemanticType; use common_recordbatch::filter::SimpleFilterEvaluator; -use common_telemetry::{error, tracing, warn}; +use common_telemetry::{debug, error, tracing, warn}; use datafusion::physical_plan::PhysicalExpr; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_expr::utils::expr_to_columns; @@ -471,7 +471,24 @@ impl ParquetReaderBuilder { &file_path, skip_auto_convert, )?; - if need_override_sequence(&parquet_meta) { + // `region_meta` comes from the Parquet/source file and must not be used as the + // target identity. When the caller has no current metadata, the handle is the + // only local identity available and therefore denotes a local read. + let expected_region_id = self + .expected_metadata + .as_ref() + .map(|metadata| metadata.region_id) + .unwrap_or(self.file_handle.region_id()); + let is_foreign = self.file_handle.region_id() != expected_region_id; + if is_foreign { + debug!( + "Reading foreign SST, file_id: {}, source_region_id: {}, expected_region_id: {}", + self.file_handle.file_id().file_id(), + self.file_handle.region_id(), + expected_region_id, + ); + } + if is_foreign || need_override_sequence(&parquet_meta) { read_format .set_override_sequence(self.file_handle.meta_ref().sequence.map(|x| x.get())); } diff --git a/src/mito2/src/worker/handle_rebuild_index.rs b/src/mito2/src/worker/handle_rebuild_index.rs index ad0a09e1d32..3eb45b1ab21 100644 --- a/src/mito2/src/worker/handle_rebuild_index.rs +++ b/src/mito2/src/worker/handle_rebuild_index.rs @@ -79,6 +79,7 @@ impl RegionWorkerLoop { IndexBuildTask { region_id: region.region_id, file: file.clone(), + target_region_metadata: version.metadata.clone(), source: IndexBuildSource::new(file_meta, version.metadata.schema_version), reason: build_type, access_layer: access_layer.clone(), From ba3c5a939e3319b1973bfc57db77c011f1b511f0 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 28 Aug 2026 05:56:44 +0000 Subject: [PATCH 119/126] chore(mito2): reduce default auto flush interval (#8971) Signed-off-by: evenyag --- config/config.md | 4 ++-- config/datanode.example.toml | 2 +- config/standalone.example.toml | 2 +- src/cmd/tests/load_config_test.rs | 4 ++-- src/mito2/src/config.rs | 12 ++++++++++-- tests-integration/tests/http.rs | 2 +- 6 files changed, 17 insertions(+), 9 deletions(-) diff --git a/config/config.md b/config/config.md index 992296d4827..1343f60dbdd 100644 --- a/config/config.md +++ b/config/config.md @@ -161,7 +161,7 @@ | `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). | | `region_engine.mito.experimental_compaction_memory_limit` | String | 0 | Memory budget for compaction tasks.
Supports absolute size (e.g., "2GiB", "512MB") or percentage of system memory (e.g., "50%").
Setting it to 0 or "unlimited" disables the limit. | | `region_engine.mito.experimental_compaction_on_exhausted` | String | wait | Behavior when compaction cannot acquire memory from the budget.
Options: "wait" (default, 10s), "wait()", "fail" | -| `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | +| `region_engine.mito.auto_flush_interval` | String | `10m` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size`. | | `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. | @@ -585,7 +585,7 @@ | `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). | | `region_engine.mito.experimental_compaction_memory_limit` | String | 0 | Memory budget for compaction tasks.
Supports absolute size (e.g., "2GiB", "512MB") or percentage of system memory (e.g., "50%").
Setting it to 0 or "unlimited" disables the limit. | | `region_engine.mito.experimental_compaction_on_exhausted` | String | wait | Behavior when compaction cannot acquire memory from the budget.
Options: "wait" (default, 10s), "wait()", "fail" | -| `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | +| `region_engine.mito.auto_flush_interval` | String | `10m` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size` | | `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 5966405452a..20c84fd9ee4 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -462,7 +462,7 @@ compress_manifest = false #+ experimental_compaction_on_exhausted = "wait" ## Interval to auto flush a region if it has not flushed yet. -auto_flush_interval = "1h" +auto_flush_interval = "10m" ## Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. ## @toml2docs:none-default="Auto" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 70f74c7f946..6ca13566d7b 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -625,7 +625,7 @@ compress_manifest = false #+ experimental_compaction_on_exhausted = "wait" ## Interval to auto flush a region if it has not flushed yet. -auto_flush_interval = "1h" +auto_flush_interval = "10m" ## Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. ## @toml2docs:none-default="Auto" diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 70594d0653f..8af216256f9 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -113,7 +113,7 @@ fn test_load_datanode_example_config() { }, region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { - auto_flush_interval: Duration::from_secs(3600), + auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), scan_memory_limit: MemoryLimit::Unlimited, @@ -354,7 +354,7 @@ fn test_load_standalone_example_config() { }), region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { - auto_flush_interval: Duration::from_secs(3600), + auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), scan_memory_limit: MemoryLimit::Unlimited, diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 54445549a05..be74fa67c60 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -102,7 +102,7 @@ pub struct MitoConfig { pub experimental_compaction_on_exhausted: OnExhaustedPolicy, // Flush configs: - /// Interval to auto flush a region if it has not flushed yet (default 30 min). + /// Interval to auto flush a region if it has not flushed yet (default 10 min). #[serde(with = "humantime_serde")] pub auto_flush_interval: Duration, /// Global write buffer size threshold to trigger flush. @@ -210,7 +210,7 @@ impl Default for MitoConfig { max_background_purges: get_total_cpu_cores(), experimental_compaction_memory_limit: MemoryLimit::Unlimited, experimental_compaction_on_exhausted: OnExhaustedPolicy::default(), - auto_flush_interval: Duration::from_secs(30 * 60), + auto_flush_interval: Duration::from_secs(10 * 60), global_write_buffer_size: ReadableSize::gb(1), global_write_buffer_reject_size: ReadableSize::gb(2), default_region_write_buffer_size: ReadableSize::mb(0), @@ -381,6 +381,14 @@ impl MitoConfig { mod tests { use super::*; + #[test] + fn test_default_auto_flush_interval() { + assert_eq!( + Duration::from_secs(10 * 60), + MitoConfig::default().auto_flush_interval + ); + } + #[test] fn test_adjust_sst_metadata_and_prefilter_cache_caps_independently() { let mut config = MitoConfig::default(); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 43ba8701d76..f4cfda24dc5 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2278,7 +2278,7 @@ experimental_manifest_keep_removed_file_ttl = "1h" compress_manifest = false experimental_compaction_memory_limit = "unlimited" experimental_compaction_on_exhausted = "wait" -auto_flush_interval = "30m" +auto_flush_interval = "10m" default_region_write_buffer_size = "0KiB" enable_write_cache = false write_cache_path = "" From 919846286964c34356487cb4cc437edc166c35c8 Mon Sep 17 00:00:00 2001 From: Dhruv Vaishnav Date: Fri, 28 Aug 2026 06:17:55 +0000 Subject: [PATCH 120/126] feat(meta): record physical table reconciliation events (#8935) * feat(meta): record physical table reconciliation events Signed-off-by: dhruvxvaishnav * docs(config): add reconciliation table event Signed-off-by: dhruvxvaishnav * fix(meta): address reconciliation event review feedback Signed-off-by: dhruvxvaishnav * fix(meta): keep reconciliation event summary volatile Signed-off-by: dhruvxvaishnav * refactor(meta): remove unused table state downcasting Signed-off-by: dhruvxvaishnav --------- Signed-off-by: dhruvxvaishnav --- config/config.md | 4 +- config/metasrv.example.toml | 2 +- config/standalone.example.toml | 2 +- src/common/meta/src/reconciliation.rs | 1 + src/common/meta/src/reconciliation/event.rs | 351 +++++++++++ .../src/reconciliation/reconcile_table.rs | 544 +++++++++++++++++- .../reconcile_table/reconcile_regions.rs | 25 +- .../reconcile_table/reconciliation_end.rs | 6 - .../reconcile_table/reconciliation_start.rs | 25 +- .../resolve_column_metadata.rs | 34 +- .../reconcile_table/update_table_info.rs | 13 +- tests-integration/tests/main.rs | 1 + .../tests/reconciliation_event.rs | 144 +++++ 13 files changed, 1102 insertions(+), 50 deletions(-) create mode 100644 src/common/meta/src/reconciliation/event.rs create mode 100644 tests-integration/tests/reconciliation_event.rs diff --git a/config/config.md b/config/config.md index 1343f60dbdd..e8a89e94e1f 100644 --- a/config/config.md +++ b/config/config.md @@ -237,7 +237,7 @@ | `tracing.tokio_console_addr` | String | Unset | The tokio console address. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`, `reconcile_table`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `memory` | -- | -- | The memory options. | | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | @@ -456,7 +456,7 @@ | `wal.create_topic_timeout` | String | `30s` | The timeout for creating a Kafka topic.
**It's only used when the provider is `kafka`**. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `region_migration`,
`create_database`, `alter_database`, `drop_database`, `create_flow`,
`drop_flow`, `create_table`, `create_logical_tables`, `alter_table`,
`alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`,
`truncate_table`, `create_view`, `drop_view`, `repartition`,
`repartition_group`, `wal_prune`, `batch_gc`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `region_migration`,
`create_database`, `alter_database`, `drop_database`, `create_flow`,
`drop_flow`, `create_table`, `create_logical_tables`, `alter_table`,
`alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`,
`truncate_table`, `create_view`, `drop_view`, `repartition`,
`repartition_group`, `wal_prune`, `batch_gc`, `reconcile_table`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `stats_persistence` | -- | -- | Configuration options for the stats persistence. | | `stats_persistence.ttl` | String | `0s` | TTL for the stats table that will be used to store the stats.
Set to `0s` to disable stats persistence.
Default is `0s`.
If you want to enable stats persistence, set the TTL to a value greater than 0.
It is recommended to set a small value, e.g., `3h`. | | `stats_persistence.interval` | String | `10m` | The interval to persist the stats. Default is `10m`.
The minimum value is `10m`, if the value is less than `10m`, it will be overridden to `10m`. | diff --git a/config/metasrv.example.toml b/config/metasrv.example.toml index 6017bfcc49a..ae373d07482 100644 --- a/config/metasrv.example.toml +++ b/config/metasrv.example.toml @@ -320,7 +320,7 @@ ttl = "90d" ## `drop_flow`, `create_table`, `create_logical_tables`, `alter_table`, ## `alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`, ## `truncate_table`, `create_view`, `drop_view`, `repartition`, -## `repartition_group`, `wal_prune`, `batch_gc`. +## `repartition_group`, `wal_prune`, `batch_gc`, `reconcile_table`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. #+ event_types = ["region_migration"] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 6ca13566d7b..1fd8fe5a522 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -930,7 +930,7 @@ ttl = "90d" ## `alter_database`, `drop_database`, `create_flow`, `drop_flow`, ## `create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`, ## `drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`, -## `create_view`, `drop_view`, `admin_function`. +## `create_view`, `drop_view`, `admin_function`, `reconcile_table`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. #+ event_types = ["create_database"] diff --git a/src/common/meta/src/reconciliation.rs b/src/common/meta/src/reconciliation.rs index 351a415566a..06ff9e704af 100644 --- a/src/common/meta/src/reconciliation.rs +++ b/src/common/meta/src/reconciliation.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub(crate) mod event; pub mod manager; pub(crate) mod reconcile_catalog; pub(crate) mod reconcile_database; diff --git a/src/common/meta/src/reconciliation/event.rs b/src/common/meta/src/reconciliation/event.rs new file mode 100644 index 00000000000..9430fa575f0 --- /dev/null +++ b/src/common/meta/src/reconciliation/event.rs @@ -0,0 +1,351 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::any::Any; + +use api::v1::value::ValueData; +use api::v1::{ColumnSchema, Row}; +use common_event_recorder::Event; +use common_event_recorder::error::{Result, SerializeEventSnafu}; +use common_event_recorder::event_table::{ + CATALOG_NAME_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN, + TABLE_NAME_COLUMN, column_schemas, nullable_string, nullable_value, +}; +use serde::Serialize; +use snafu::ResultExt; +use store_api::storage::TableId; + +use crate::reconciliation::ResolveStrategy; + +/// Stable event type stored for physical table reconciliation procedures. +pub(crate) const RECONCILE_TABLE_EVENT_TYPE: &str = "reconcile_table"; +const PAYLOAD_VERSION: u8 = 1; + +#[derive(Debug, Clone, Copy)] +enum ReconciliationEventType { + Table, +} + +impl ReconciliationEventType { + const fn as_str(self) -> &'static str { + match self { + Self::Table => RECONCILE_TABLE_EVENT_TYPE, + } + } +} + +/// Nullable object locators shared by all reconciliation event types. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct ReconciliationLocator { + catalog_name: Option, + schema_name: Option, + table_name: Option, + table_id: Option, + physical_table_id: Option, +} + +impl ReconciliationLocator { + /// Creates a locator for a physical table with its fully qualified name and ID. + pub(crate) fn physical_table( + catalog_name: &str, + schema_name: &str, + table_name: &str, + table_id: TableId, + ) -> Self { + Self { + catalog_name: Some(catalog_name.to_string()), + schema_name: Some(schema_name.to_string()), + table_name: Some(table_name.to_string()), + table_id: Some(table_id), + ..Default::default() + } + } +} + +#[derive(Debug, Serialize)] +#[serde(untagged)] +enum ReconciliationPayload { + TableSubmitted(TableSubmittedPayload), + TableResult(TableResultPayload), +} + +#[derive(Debug, Serialize)] +struct TableSubmittedPayload { + version: u8, + resolve_strategy: &'static str, + is_subprocedure: bool, +} + +#[derive(Debug, Serialize)] +struct TableResultPayload { + version: u8, + complete: bool, + metadata_state: Option<&'static str>, + resolution_strategy_applied: Option<&'static str>, + resolved_column_count: Option, + scanned_region_count: usize, + updated_region_count: usize, + table_info_updated: bool, + last_completed_phase: Option<&'static str>, +} + +/// Event representation shared by reconciliation procedures. +#[derive(Debug)] +pub(crate) struct ReconciliationEvent { + event_type: ReconciliationEventType, + locator: ReconciliationLocator, + payload: Option, +} + +impl ReconciliationEvent { + /// Builds the bounded intent event emitted when table reconciliation is submitted. + pub(crate) fn table_submitted( + locator: ReconciliationLocator, + resolve_strategy: ResolveStrategy, + is_subprocedure: bool, + ) -> Self { + Self { + event_type: ReconciliationEventType::Table, + locator, + payload: Some(ReconciliationPayload::TableSubmitted( + TableSubmittedPayload { + version: PAYLOAD_VERSION, + resolve_strategy: resolve_strategy_name(resolve_strategy), + is_subprocedure, + }, + )), + } + } + + /// Builds a terminal event from the bounded reconciliation result summary. + #[allow(clippy::too_many_arguments)] + pub(crate) fn table_result( + locator: ReconciliationLocator, + complete: bool, + metadata_state: Option<&'static str>, + resolution_strategy_applied: Option, + resolved_column_count: Option, + scanned_region_count: usize, + updated_region_count: usize, + table_info_updated: bool, + last_completed_phase: Option<&'static str>, + ) -> Self { + Self { + event_type: ReconciliationEventType::Table, + locator, + payload: Some(ReconciliationPayload::TableResult(TableResultPayload { + version: PAYLOAD_VERSION, + complete, + metadata_state, + resolution_strategy_applied: resolution_strategy_applied.map(resolve_strategy_name), + resolved_column_count, + scanned_region_count, + updated_region_count, + table_info_updated, + last_completed_phase, + })), + } + } + + /// Builds a table lifecycle event whose reconciliation payload is null. + pub(crate) fn table_lifecycle(locator: ReconciliationLocator) -> Self { + Self::lifecycle(ReconciliationEventType::Table, locator) + } + + fn lifecycle(event_type: ReconciliationEventType, locator: ReconciliationLocator) -> Self { + Self { + event_type, + locator, + payload: None, + } + } + + fn schema() -> Vec { + column_schemas([ + &CATALOG_NAME_COLUMN, + &SCHEMA_NAME_COLUMN, + &TABLE_NAME_COLUMN, + &TABLE_ID_COLUMN, + &PHYSICAL_TABLE_ID_COLUMN, + ]) + } +} + +impl Event for ReconciliationEvent { + fn event_type(&self) -> &str { + self.event_type.as_str() + } + + fn json_payload(&self) -> Result { + match &self.payload { + Some(payload) => serde_json::to_value(payload).context(SerializeEventSnafu), + None => Ok(serde_json::Value::Null), + } + } + + fn extra_schema(&self) -> Vec { + Self::schema() + } + + fn extra_rows(&self) -> Result> { + Ok(vec![Row { + values: vec![ + nullable_string(self.locator.catalog_name.as_deref()), + nullable_string(self.locator.schema_name.as_deref()), + nullable_string(self.locator.table_name.as_deref()), + nullable_table_id(self.locator.table_id), + nullable_table_id(self.locator.physical_table_id), + ], + }]) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +fn resolve_strategy_name(strategy: ResolveStrategy) -> &'static str { + match strategy { + ResolveStrategy::UseLatest => "use_latest", + ResolveStrategy::UseMetasrv => "use_metasrv", + ResolveStrategy::AbortOnConflict => "abort_on_conflict", + } +} + +fn nullable_table_id(value: Option) -> api::v1::Value { + nullable_value(value.map(ValueData::U32Value)) +} + +#[cfg(test)] +mod tests { + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row, SemanticType, Value}; + use common_event_recorder::Event; + use serde_json::json; + + use super::*; + + #[test] + fn reconciliation_events_use_the_shared_locator_contract() { + let table = ReconciliationEvent::table_lifecycle(ReconciliationLocator::physical_table( + "greptime", "public", "metrics", 42, + )); + assert_eq!(table.event_type(), RECONCILE_TABLE_EVENT_TYPE); + assert_eq!( + table + .extra_schema() + .into_iter() + .map(|column| { + ( + column.column_name, + ColumnDataType::try_from(column.datatype).unwrap(), + SemanticType::try_from(column.semantic_type).unwrap(), + ) + }) + .collect::>(), + vec![ + ( + "catalog_name".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + ( + "schema_name".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + ( + "table_name".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + ( + "table_id".to_string(), + ColumnDataType::Uint32, + SemanticType::Field, + ), + ( + "physical_table_id".to_string(), + ColumnDataType::Uint32, + SemanticType::Field, + ), + ] + ); + assert_eq!( + table.extra_rows().unwrap(), + vec![Row { + values: vec![ + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue("public".to_string()).into(), + ValueData::StringValue("metrics".to_string()).into(), + ValueData::U32Value(42).into(), + Value::default(), + ], + }] + ); + assert_eq!(table.json_payload().unwrap(), serde_json::Value::Null); + } + + #[test] + fn submitted_payloads_are_versioned_and_use_stable_strategy_names() { + for (strategy, expected) in [ + (ResolveStrategy::UseLatest, "use_latest"), + (ResolveStrategy::UseMetasrv, "use_metasrv"), + (ResolveStrategy::AbortOnConflict, "abort_on_conflict"), + ] { + let table = ReconciliationEvent::table_submitted( + ReconciliationLocator::physical_table("greptime", "public", "metrics", 42), + strategy, + true, + ); + assert_eq!( + table.json_payload().unwrap(), + json!({ + "version": 1, + "resolve_strategy": expected, + "is_subprocedure": true, + }) + ); + } + } + + #[test] + fn terminal_payloads_distinguish_complete_and_partial_results() { + let table = ReconciliationEvent::table_result( + ReconciliationLocator::physical_table("greptime", "public", "metrics", 42), + false, + Some("inconsistent"), + Some(ResolveStrategy::UseMetasrv), + Some(4), + 3, + 2, + true, + Some("update_table_info"), + ); + assert_eq!( + table.json_payload().unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_metasrv", + "resolved_column_count": 4, + "scanned_region_count": 3, + "updated_region_count": 2, + "table_info_updated": true, + "last_completed_phase": "update_table_info", + }) + ); + } +} diff --git a/src/common/meta/src/reconciliation/reconcile_table.rs b/src/common/meta/src/reconciliation/reconcile_table.rs index 004b93d38c5..562a2061b0a 100644 --- a/src/common/meta/src/reconciliation/reconcile_table.rs +++ b/src/common/meta/src/reconciliation/reconcile_table.rs @@ -18,13 +18,12 @@ pub(crate) mod reconciliation_start; pub(crate) mod resolve_column_metadata; pub(crate) mod update_table_info; -use std::any::Any; use std::fmt::Debug; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure, - Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, Result as ProcedureResult, Status, }; use serde::{Deserialize, Serialize}; use snafu::ResultExt; @@ -42,6 +41,9 @@ use crate::key::{DeserializedValueWithBytes, TableMetadataManagerRef}; use crate::lock_key::{CatalogLock, SchemaLock, TableNameLock}; use crate::metrics; use crate::node_manager::NodeManagerRef; +use crate::reconciliation::event::{ + RECONCILE_TABLE_EVENT_TYPE, ReconciliationEvent, ReconciliationLocator, +}; use crate::reconciliation::reconcile_table::reconciliation_start::ReconciliationStart; use crate::reconciliation::reconcile_table::resolve_column_metadata::ResolveStrategy; use crate::reconciliation::utils::{ @@ -56,6 +58,105 @@ pub struct ReconcileTableContext { pub volatile_ctx: VolatileContext, } +#[derive(Debug, Clone, Copy)] +enum TablePhase { + Start, + ResolveColumnMetadata, + ReconcileRegions, + UpdateTableInfo, +} + +impl TablePhase { + const fn as_event_value(self) -> &'static str { + match self { + Self::Start => "start", + Self::ResolveColumnMetadata => "resolve_column_metadata", + Self::ReconcileRegions => "reconcile_regions", + Self::UpdateTableInfo => "update_table_info", + } + } +} + +#[derive(Debug, Clone, Copy)] +enum TableMetadataState { + Consistent, + Inconsistent, +} + +impl TableMetadataState { + const fn as_event_value(self) -> &'static str { + match self { + Self::Consistent => "consistent", + Self::Inconsistent => "inconsistent", + } + } +} + +#[derive(Debug, Default)] +struct ReconcileTableResultSummary { + metadata_state: Option, + resolution_strategy_applied: Option, + resolved_column_count: Option, + scanned_region_count: usize, + updated_region_count: usize, + table_info_updated: bool, + last_completed_phase: Option, +} + +impl ReconcileTableResultSummary { + fn record_scanned_regions(&mut self, scanned_region_count: usize) { + self.scanned_region_count = scanned_region_count; + } + + fn mark_start_completed(&mut self) { + self.last_completed_phase = Some(TablePhase::Start); + } + + fn record_metadata_state(&mut self, metadata_state: TableMetadataState) { + self.metadata_state = Some(metadata_state); + } + + fn record_resolution_strategy(&mut self, resolve_strategy: ResolveStrategy) { + self.resolution_strategy_applied = Some(resolve_strategy); + } + + fn record_resolved_columns( + &mut self, + metadata_state: TableMetadataState, + resolution_strategy_applied: Option, + resolved_column_count: Option, + ) { + self.metadata_state = Some(metadata_state); + self.resolution_strategy_applied = resolution_strategy_applied; + self.resolved_column_count = resolved_column_count; + self.last_completed_phase = Some(TablePhase::ResolveColumnMetadata); + } + + fn record_updated_regions(&mut self, updated_region_count: usize) { + self.updated_region_count = updated_region_count; + } + + fn mark_region_phase_completed(&mut self) { + self.last_completed_phase = Some(TablePhase::ReconcileRegions); + } + + fn mark_table_info_updated(&mut self) { + self.table_info_updated = true; + } + + fn mark_table_info_phase_completed(&mut self) { + self.last_completed_phase = Some(TablePhase::UpdateTableInfo); + } + + fn metadata_state(&self) -> Option<&'static str> { + self.metadata_state.map(TableMetadataState::as_event_value) + } + + fn last_completed_phase(&self) -> Option<&'static str> { + self.last_completed_phase.map(TablePhase::as_event_value) + } +} + impl ReconcileTableContext { /// Creates a new [`ReconcileTableContext`] with the given [`Context`] and [`PersistentContext`]. pub fn new(ctx: Context, persistent_ctx: PersistentContext) -> Self { @@ -147,6 +248,7 @@ impl PersistentContext { pub(crate) struct VolatileContext { pub(crate) table_meta: Option, pub(crate) metrics: ReconcileTableMetrics, + result_summary: ReconcileTableResultSummary, } pub struct ReconcileTableProcedure { @@ -256,6 +358,56 @@ impl Procedure for ReconcileTableProcedure { TableNameLock::new(table_ref.catalog, table_ref.schema, table_ref.table).into(), ]) } + + fn event(&self, ctx: &EventContext<'_>) -> Option> { + if !ctx.event_type_filter.allows(RECONCILE_TABLE_EVENT_TYPE) { + return None; + } + + let persistent_ctx = &self.context.persistent_ctx; + let table_name = &persistent_ctx.table_name; + let locator = ReconciliationLocator::physical_table( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + persistent_ctx.table_id, + ); + let event = match ctx.trigger { + EventTrigger::Submitted => ReconciliationEvent::table_submitted( + locator, + persistent_ctx.resolve_strategy, + persistent_ctx.is_subprocedure, + ), + EventTrigger::Succeeded => { + Self::result_event(locator, &self.context.volatile_ctx.result_summary, true) + } + EventTrigger::Failed | EventTrigger::Poisoned => { + Self::result_event(locator, &self.context.volatile_ctx.result_summary, false) + } + _ => ReconciliationEvent::table_lifecycle(locator), + }; + Some(Box::new(event)) + } +} + +impl ReconcileTableProcedure { + fn result_event( + locator: ReconciliationLocator, + summary: &ReconcileTableResultSummary, + complete: bool, + ) -> ReconciliationEvent { + ReconciliationEvent::table_result( + locator, + complete, + summary.metadata_state(), + summary.resolution_strategy_applied, + summary.resolved_column_count, + summary.scanned_region_count, + summary.updated_region_count, + summary.table_info_updated, + summary.last_completed_phase(), + ) + } } #[async_trait::async_trait] @@ -272,6 +424,390 @@ pub(crate) trait State: Sync + Send + Debug { ctx: &mut ReconcileTableContext, procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)>; +} - fn as_any(&self) -> &dyn Any; +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_event_recorder::{EventTypeFilter, EventTypeFilterRef}; + use common_procedure::{ + ChildSubmissionOutcome, EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState, + RetryPhase, + }; + use common_procedure_test::MockContextProvider; + use serde_json::{Value, json}; + use store_api::storage::RegionId; + + use super::*; + use crate::ddl::test_util::datanode_handler::PartialSuccessDatanodeHandler; + use crate::key::DeserializedValueWithBytes; + use crate::key::table_info::TableInfoValue; + use crate::key::table_route::PhysicalTableRouteValue; + use crate::key::test_utils::new_test_table_info_with_name; + use crate::peer::Peer; + use crate::reconciliation::reconcile_table::reconcile_regions::ReconcileRegions; + use crate::reconciliation::utils::build_column_metadata_from_table_info; + use crate::rpc::router::{Region, RegionRoute}; + use crate::test_util::{MockDatanodeManager, new_ddl_context}; + + struct TableEventHarness { + procedure_id: ProcedureId, + lifecycle_state: ProcedureState, + event_type_filter: EventTypeFilterRef, + } + + impl TableEventHarness { + fn all() -> Self { + Self { + procedure_id: ProcedureId::random(), + lifecycle_state: ProcedureState::Running, + event_type_filter: Arc::new(EventTypeFilter::All), + } + } + + fn selected(event_types: impl IntoIterator) -> Self { + Self { + event_type_filter: Arc::new(EventTypeFilter::Only( + event_types.into_iter().map(str::to_string).collect(), + )), + ..Self::all() + } + } + + fn event( + &self, + procedure: &dyn Procedure, + trigger: EventTrigger, + ) -> Option> { + procedure.event(&EventContext { + procedure_id: self.procedure_id, + lifecycle_state: &self.lifecycle_state, + trigger, + event_type_filter: self.event_type_filter.clone(), + event_context: None, + }) + } + } + + #[test] + fn table_submitted_events_cover_root_and_child_intent() { + let events = TableEventHarness::all(); + let root = test_procedure(false); + let child = test_procedure(true); + for (procedure, is_subprocedure) in [(&root, false), (&child, true)] { + let submitted = events.event(procedure, EventTrigger::Submitted).unwrap(); + assert_eq!(submitted.event_type(), RECONCILE_TABLE_EVENT_TYPE); + assert_eq!( + submitted.json_payload().unwrap(), + json!({ + "version": 1, + "resolve_strategy": "use_latest", + "is_subprocedure": is_subprocedure, + }) + ); + } + } + + #[test] + fn table_non_terminal_lifecycle_events_have_null_payloads() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(true); + procedure.context.volatile_ctx.result_summary = populated_summary(); + + for trigger in [ + EventTrigger::Recovered, + EventTrigger::ChildSubmitted { + procedure_id: ProcedureId::random(), + outcome: ChildSubmissionOutcome::Accepted, + }, + EventTrigger::Retrying { + phase: RetryPhase::Execute, + attempt: 2, + }, + EventTrigger::RollingBack, + ] { + assert_eq!( + events + .event(&procedure, trigger) + .unwrap() + .json_payload() + .unwrap(), + Value::Null + ); + } + } + + #[test] + fn table_terminal_events_report_bounded_results() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(true); + procedure.context.volatile_ctx.result_summary = populated_summary(); + + for (trigger, complete) in [ + (EventTrigger::Succeeded, true), + (EventTrigger::Failed, false), + (EventTrigger::Poisoned, false), + ] { + assert_eq!( + events + .event(&procedure, trigger) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": complete, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_latest", + "resolved_column_count": 3, + "scanned_region_count": 2, + "updated_region_count": 1, + "table_info_updated": true, + "last_completed_phase": "update_table_info", + }) + ); + } + } + + #[test] + fn table_event_filtering_uses_the_reconciliation_event_type() { + let procedure = test_procedure(false); + assert!( + TableEventHarness::selected([RECONCILE_TABLE_EVENT_TYPE]) + .event(&procedure, EventTrigger::Submitted) + .is_some() + ); + assert!( + TableEventHarness::selected(["create_table"]) + .event(&procedure, EventTrigger::Submitted) + .is_none() + ); + assert!( + TableEventHarness::selected([]) + .event(&procedure, EventTrigger::Submitted) + .is_none() + ); + } + + #[test] + fn table_result_summary_is_not_persisted() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(false); + procedure.context.volatile_ctx.result_summary = populated_summary(); + let value: Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + assert!(value["persistent_ctx"].get("result_summary").is_none()); + + let loaded = ReconcileTableProcedure::from_json( + test_context(), + &serde_json::to_string(&value).unwrap(), + ) + .unwrap(); + assert_eq!( + events + .event(&loaded, EventTrigger::Succeeded) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": true, + "metadata_state": null, + "resolution_strategy_applied": null, + "resolved_column_count": null, + "scanned_region_count": 0, + "updated_region_count": 0, + "table_info_updated": false, + "last_completed_phase": null, + }) + ); + } + + #[test] + fn table_partial_summary_preserves_completed_work_without_advancing_phases() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(false); + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.record_scanned_regions(2); + + assert_eq!( + events + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": null, + "resolution_strategy_applied": null, + "resolved_column_count": null, + "scanned_region_count": 2, + "updated_region_count": 0, + "table_info_updated": false, + "last_completed_phase": null, + }) + ); + + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.mark_start_completed(); + summary.record_metadata_state(TableMetadataState::Inconsistent); + assert_eq!( + events + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": null, + "resolved_column_count": null, + "scanned_region_count": 2, + "updated_region_count": 0, + "table_info_updated": false, + "last_completed_phase": "start", + }) + ); + + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(ResolveStrategy::UseLatest), + Some(3), + ); + summary.record_updated_regions(1); + assert_eq!( + events + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_latest", + "resolved_column_count": 3, + "scanned_region_count": 2, + "updated_region_count": 1, + "table_info_updated": false, + "last_completed_phase": "resolve_column_metadata", + }) + ); + } + + #[tokio::test] + async fn table_mixed_region_outcome_preserves_partial_result() { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new( + PartialSuccessDatanodeHandler { retryable: false }, + ))); + let context = Context { + node_manager: ddl_context.node_manager, + table_metadata_manager: ddl_context.table_metadata_manager, + cache_invalidator: ddl_context.cache_invalidator, + }; + let mut procedure = ReconcileTableProcedure::new( + context, + 42, + TableName::new("greptime", "public", "metrics"), + ResolveStrategy::UseLatest, + false, + ); + + let mut table_info = new_test_table_info_with_name(42, "metrics"); + table_info.meta.column_ids = vec![0, 1, 2]; + let name_to_ids = table_info.name_to_ids().unwrap(); + let column_metadatas = build_column_metadata_from_table_info( + table_info.meta.schema.column_schemas(), + &table_info.meta.primary_key_indices, + &name_to_ids, + ) + .unwrap(); + let region_ids = [RegionId::new(42, 1), RegionId::new(42, 2)]; + procedure.context.persistent_ctx.table_info_value = Some( + DeserializedValueWithBytes::from_inner(TableInfoValue::new(table_info)), + ); + procedure.context.persistent_ctx.physical_table_route = + Some(PhysicalTableRouteValue::new(vec![ + RegionRoute { + region: Region::new_test(region_ids[0]), + leader_peer: Some(Peer::empty(1)), + ..Default::default() + }, + RegionRoute { + region: Region::new_test(region_ids[1]), + leader_peer: Some(Peer::empty(2)), + ..Default::default() + }, + ])); + + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.record_scanned_regions(region_ids.len()); + summary.mark_start_completed(); + summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(ResolveStrategy::UseLatest), + Some(column_metadatas.len()), + ); + procedure.state = Box::new(ReconcileRegions::new(column_metadatas, region_ids.to_vec())); + + let procedure_ctx = ProcedureContext { + procedure_id: ProcedureId::random(), + provider: Arc::new(MockContextProvider::default()), + event_context: None, + }; + assert!(procedure.execute(&procedure_ctx).await.is_err()); + + assert_eq!( + TableEventHarness::all() + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_latest", + "resolved_column_count": 3, + "scanned_region_count": 2, + "updated_region_count": 1, + "table_info_updated": false, + "last_completed_phase": "resolve_column_metadata", + }) + ); + } + + fn populated_summary() -> ReconcileTableResultSummary { + ReconcileTableResultSummary { + metadata_state: Some(TableMetadataState::Inconsistent), + resolution_strategy_applied: Some(ResolveStrategy::UseLatest), + resolved_column_count: Some(3), + scanned_region_count: 2, + updated_region_count: 1, + table_info_updated: true, + last_completed_phase: Some(TablePhase::UpdateTableInfo), + } + } + + fn test_procedure(is_subprocedure: bool) -> ReconcileTableProcedure { + ReconcileTableProcedure::new( + test_context(), + 42, + TableName::new("greptime", "public", "metrics"), + ResolveStrategy::UseLatest, + is_subprocedure, + ) + } + + fn test_context() -> Context { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + Context { + node_manager: ddl_context.node_manager, + table_metadata_manager: ddl_context.table_metadata_manager, + cache_invalidator: ddl_context.cache_invalidator, + } + } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs b/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs index 6ee3df70f37..fd5d55fb0cb 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; use std::collections::{HashMap, HashSet}; use api::v1::column_def::try_as_column_def; @@ -65,7 +64,6 @@ impl State for ReconcileRegions { let table_meta = ctx.build_table_meta(&self.column_metadatas)?; ctx.volatile_ctx.table_meta = Some(table_meta); let table_id = ctx.table_id(); - let table_name = ctx.table_name(); let primary_keys = self .column_metadatas @@ -126,10 +124,12 @@ impl State for ReconcileRegions { } } - let mut results = future::join_all(sync_column_tsks) - .await - .into_iter() - .collect::>>()?; + let results = future::join_all(sync_column_tsks).await; + let updated_region_count = results.iter().filter(|result| result.is_ok()).count(); + ctx.volatile_ctx + .result_summary + .record_updated_regions(updated_region_count); + let mut results = results.into_iter().collect::>>()?; // Ensures all the column metadatas are the same. let column_metadatas = @@ -137,17 +137,22 @@ impl State for ReconcileRegions { UnexpectedSnafu { err_msg: format!( "The table column metadata schemas from datanodes are not the same, table: {}, table_id: {}", - table_name, + ctx.table_name(), table_id ), }, )?; + ctx.volatile_ctx + .result_summary + .mark_region_phase_completed(); + // Checks all column metadatas are consistent, and updates the table info if needed. if column_metadatas != self.column_metadatas { info!( "Datanode column metadatas are not consistent with metasrv, updating metasrv's column metadatas, table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ); // Safety: fetched in the above. let table_info_value = ctx.persistent_ctx.table_info_value.clone().unwrap(); @@ -159,10 +164,6 @@ impl State for ReconcileRegions { Ok((Box::new(ReconciliationEnd), Status::executing(false))) } - - fn as_any(&self) -> &dyn Any { - self - } } /// Makes an alter region request to sync columns. diff --git a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs index 9f28d7f4be5..b6e695dca46 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -46,8 +44,4 @@ impl State for ReconciliationEnd { Ok((Box::new(ReconciliationEnd), Status::done())) } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs index 2fcc8ded0c0..48a7fff8b96 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -46,7 +44,6 @@ impl State for ReconciliationStart { procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)> { let table_id = ctx.table_id(); - let table_name = ctx.table_name(); let (physical_table_id, physical_table_route) = ctx .table_metadata_manager @@ -58,14 +55,17 @@ impl State for ReconciliationStart { error::UnexpectedSnafu { err_msg: format!( "Reconcile table only works for physical table, but got logical table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ), } ); info!( "Reconciling table: {}, table_id: {}, procedure_id: {}", - table_name, table_id, procedure_ctx.procedure_id + ctx.table_name(), + table_id, + procedure_ctx.procedure_id ); // TODO(weny): Repairs the table route if needed. let region_metadata_lister = RegionMetadataLister::new(ctx.node_manager.clone()); @@ -79,6 +79,9 @@ impl State for ReconciliationStart { .list(physical_table_id, &physical_table_route.region_routes) .await? }; + ctx.volatile_ctx + .result_summary + .record_scanned_regions(region_metadatas.len()); ensure!(!region_metadatas.is_empty(), { metrics::METRIC_META_RECONCILIATION_STATS @@ -92,7 +95,8 @@ impl State for ReconciliationStart { error::UnexpectedSnafu { err_msg: format!( "No region metadata found for table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ), } }); @@ -109,11 +113,14 @@ impl State for ReconciliationStart { error::UnexpectedSnafu { err_msg: format!( "Some regions are not opened, table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ), } }); + ctx.volatile_ctx.result_summary.mark_start_completed(); + // Persist the physical table route. // TODO(weny): refetch the physical table route if repair is needed. ctx.persistent_ctx.physical_table_route = Some(physical_table_route); @@ -127,8 +134,4 @@ impl State for ReconciliationStart { Status::executing(false), )) } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs index d45cfe464ea..27950f0fd5d 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use async_trait::async_trait; use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; @@ -25,7 +23,7 @@ use strum::AsRefStr; use crate::error::{self, MissingColumnIdsSnafu, Result}; use crate::reconciliation::reconcile_table::reconcile_regions::ReconcileRegions; use crate::reconciliation::reconcile_table::update_table_info::UpdateTableInfo; -use crate::reconciliation::reconcile_table::{ReconcileTableContext, State}; +use crate::reconciliation::reconcile_table::{ReconcileTableContext, State, TableMetadataState}; use crate::reconciliation::utils::{ ResolveColumnMetadataResult, build_column_metadata_from_table_info, check_column_metadatas_consistent, resolve_column_metadatas_with_latest, @@ -101,6 +99,12 @@ impl State for ResolveColumnMetadata { table_name, table_id ); + ctx.volatile_ctx.result_summary.record_resolved_columns( + TableMetadataState::Consistent, + None, + Some(column_metadatas.len()), + ); + // Update metrics. ctx.mut_metrics().resolve_column_metadata_result = Some(ResolveColumnMetadataResult::Consistent); @@ -110,6 +114,10 @@ impl State for ResolveColumnMetadata { )); }; + ctx.volatile_ctx + .result_summary + .record_metadata_state(TableMetadataState::Inconsistent); + match self.strategy { ResolveStrategy::UseMetasrv => { let table_info_value = ctx.persistent_ctx.table_info_value.as_ref().unwrap(); @@ -126,6 +134,12 @@ impl State for ResolveColumnMetadata { let region_ids = resolve_column_metadatas_with_metasrv(&column_metadata, &self.region_metadata)?; + ctx.volatile_ctx.result_summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(self.strategy), + Some(column_metadata.len()), + ); + // Update metrics. let metrics = ctx.mut_metrics(); metrics.resolve_column_metadata_result = @@ -139,6 +153,12 @@ impl State for ResolveColumnMetadata { let (column_metadatas, region_ids) = resolve_column_metadatas_with_latest(&self.region_metadata)?; + ctx.volatile_ctx.result_summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(self.strategy), + Some(column_metadatas.len()), + ); + // Update metrics. let metrics = ctx.mut_metrics(); metrics.resolve_column_metadata_result = @@ -151,6 +171,10 @@ impl State for ResolveColumnMetadata { ResolveStrategy::AbortOnConflict => { let table_name = table_name.to_string(); + ctx.volatile_ctx + .result_summary + .record_resolution_strategy(self.strategy); + // Update metrics. let metrics = ctx.mut_metrics(); metrics.resolve_column_metadata_result = @@ -163,8 +187,4 @@ impl State for ResolveColumnMetadata { } } } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs b/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs index 790184c07ab..7ef62265de3 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -81,6 +79,9 @@ impl State for UpdateTableInfo { ctx.table_name(), ctx.table_id() ); + ctx.volatile_ctx + .result_summary + .mark_table_info_phase_completed(); return Ok((Box::new(ReconciliationEnd), Status::executing(true))); } @@ -98,6 +99,7 @@ impl State for UpdateTableInfo { new_table_info, ) .await?; + ctx.volatile_ctx.result_summary.mark_table_info_updated(); let table_ref = ctx.table_name().table_ref(); let table_id = ctx.table_id(); @@ -119,11 +121,10 @@ impl State for UpdateTableInfo { // Update metrics. let metrics = ctx.mut_metrics(); metrics.update_table_info = true; + ctx.volatile_ctx + .result_summary + .mark_table_info_phase_completed(); Ok((Box::new(ReconciliationEnd), Status::executing(true))) } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 965a3a67462..56108d4d3ff 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -37,6 +37,7 @@ mod repartition_event; #[macro_use] mod repartition_expr_version; mod mysql; +mod reconciliation_event; mod view_ddl_event; mod wal_prune_event; diff --git a/tests-integration/tests/reconciliation_event.rs b/tests-integration/tests/reconciliation_event.rs new file mode 100644 index 00000000000..c4b4b88ea2d --- /dev/null +++ b/tests-integration/tests/reconciliation_event.rs @@ -0,0 +1,144 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use client::OutputData; +use common_meta::reconciliation::ResolveStrategy; +use common_procedure::watcher; +use frontend::instance::Instance; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::QueryContext; +use table::table_reference::TableReference; +use tests_integration::cluster::GreptimeDbClusterBuilder; + +use crate::event_recorder_test_util::{assert_eventually_eq, find_eventually_string}; + +const EVENTS_TABLE: &str = "greptime_private.events"; +const CATALOG: &str = "greptime"; +const DATABASE: &str = "reconciliation_event_database"; + +#[tokio::test(flavor = "multi_thread")] +async fn test_table_reconciliation_events() { + common_telemetry::init_default_ut_logging(); + + let cluster = GreptimeDbClusterBuilder::new("table_reconciliation_events") + .await + .with_datanodes(1) + .build(true) + .await; + let frontend = cluster.fe_instance().clone(); + run_sql(frontend.as_ref(), &format!("CREATE DATABASE {DATABASE}")).await; + run_sql( + frontend.as_ref(), + &format!("CREATE TABLE {DATABASE}.metrics (ts TIMESTAMP TIME INDEX)"), + ) + .await; + + let table_procedure_id = cluster + .metasrv + .reconciliation_manager() + .reconcile_table( + TableReference { + catalog: CATALOG, + schema: DATABASE, + table: "metrics", + }, + ResolveStrategy::UseLatest, + ) + .await + .unwrap(); + let mut procedure_watcher = cluster + .metasrv + .procedure_manager() + .procedure_watcher(table_procedure_id) + .unwrap(); + watcher::wait(&mut procedure_watcher).await.unwrap(); + let table_procedure_id = table_procedure_id.to_string(); + + assert_eventually_eq( + &frontend, + &format!( + "SELECT count(*) = 1 AS matches FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_table' AND procedure_id = '{table_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Submitted' \ + AND catalog_name = '{CATALOG}' AND schema_name = '{DATABASE}' \ + AND table_name = 'metrics' AND table_id IS NOT NULL \ + AND physical_table_id IS NULL \ + AND json_get_int(payload, 'version') = 1 \ + AND json_get_string(payload, 'resolve_strategy') = 'use_latest' \ + AND json_get_bool(payload, 'is_subprocedure') = false" + ), + "\ ++---------+ +| matches | ++---------+ +| true | ++---------+", + ) + .await; + + let table_result_payload = find_eventually_string( + &frontend, + &format!( + "SELECT json_to_string(payload) AS payload FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_table' AND procedure_id = '{table_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Succeeded' \ + AND catalog_name = '{CATALOG}' AND schema_name = '{DATABASE}' \ + AND table_name = 'metrics' AND table_id IS NOT NULL \ + AND physical_table_id IS NULL LIMIT 1" + ), + "payload", + ) + .await; + assert_eq!( + serde_json::from_str::(&table_result_payload).unwrap(), + serde_json::json!({ + "version": 1, + "complete": true, + "metadata_state": "consistent", + "resolution_strategy_applied": null, + "resolved_column_count": 1, + "scanned_region_count": 1, + "updated_region_count": 0, + "table_info_updated": true, + "last_completed_phase": "update_table_info", + }) + ); + + assert_eventually_eq( + &frontend, + &format!( + "SELECT count(*) AS terminal_event_count FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_table' AND procedure_id = '{table_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Succeeded' \ + AND catalog_name = '{CATALOG}' AND schema_name = '{DATABASE}' \ + AND table_name = 'metrics' AND table_id IS NOT NULL \ + AND physical_table_id IS NULL" + ), + "\ ++----------------------+ +| terminal_event_count | ++----------------------+ +| 1 | ++----------------------+", + ) + .await; +} + +async fn run_sql(instance: &Instance, sql: &str) { + let output = SqlQueryHandler::do_query(instance, sql, QueryContext::arc()) + .await + .remove(0) + .unwrap(); + assert!(matches!(output.data, OutputData::AffectedRows(_)), "{sql}"); +} From 28e2b5aca26fdaf49d5851e88de80ac787c4a14d Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 28 Aug 2026 07:55:57 +0000 Subject: [PATCH 121/126] feat: expose missing SST manifest fields (#8965) * feat: expose missing SST manifest fields Signed-off-by: evenyag * fix: preserve SST manifest column ordinals Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/metric-engine/src/engine/flush.rs | 13 +-- src/mito2/src/engine/basic_test.rs | 13 +-- src/mito2/src/region.rs | 2 + src/store-api/src/sst_entry.rs | 75 ++++++++++++++++ .../common/information_schema/ssts.result | 90 ++++++++++--------- .../common/system/information_schema.result | 2 + 6 files changed, 139 insertions(+), 56 deletions(-) diff --git a/src/metric-engine/src/engine/flush.rs b/src/metric-engine/src/engine/flush.rs index 62b2266249c..21e642de0b2 100644 --- a/src/metric-engine/src/engine/flush.rs +++ b/src/metric-engine/src/engine/flush.rs @@ -121,6 +121,7 @@ mod tests { .map(|path| path.replace(&e.file_id, "")); e.file_id = "".to_string(); e.index_version = 0; + e.max_row_group_uncompressed_size = 0; // Round down sizes to nearest 1000 to avoid exact size // comparisons that break when the SST format changes. e.file_size = e.file_size / 1000 * 1000; @@ -133,12 +134,12 @@ mod tests { assert_eq!( debug_format, r#" -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/data/.parquet", file_size: 3000, index_file_path: Some("test_metric_region/11_0000000001/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(20), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/data/.parquet", file_size: 3000, index_file_path: Some("test_metric_region/11_0000000002/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417473(11, 16777217), table_id: 11, region_number: 16777217, region_group: 1, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/metadata/.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 8, num_row_groups: 1, num_series: Some(8), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(8), origin_region_id: 47261417473(11, 16777217), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665867_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665868\0\0\x06") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417474(11, 16777218), table_id: 11, region_number: 16777218, region_group: 1, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/metadata/.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), origin_region_id: 47261417474(11, 16777218), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665858_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665858\0\0\x06") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/data/.parquet", file_size: 3000, index_file_path: Some("test_metric_region/22_0000000042/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94506057770(22, 16777258), table_id: 22, region_number: 16777258, region_group: 1, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/metadata/.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), origin_region_id: 94506057770(22, 16777258), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_9543417\t331754_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_9543417\t331754\0\0\x06") }"#, +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/data/.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/11_0000000001/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(20), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/data/.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/11_0000000002/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417473(11, 16777217), table_id: 11, region_number: 16777217, region_group: 1, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/metadata/.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 8, num_row_groups: 1, num_series: Some(8), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(8), partition_expr: None, origin_region_id: 47261417473(11, 16777217), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665867_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665868\0\0\x06") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417474(11, 16777218), table_id: 11, region_number: 16777218, region_group: 1, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/metadata/.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), partition_expr: None, origin_region_id: 47261417474(11, 16777218), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665858_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665858\0\0\x06") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/data/.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/22_0000000042/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94506057770(22, 16777258), table_id: 22, region_number: 16777258, region_group: 1, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/metadata/.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), partition_expr: None, origin_region_id: 94506057770(22, 16777258), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_9543417\t331754_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_9543417\t331754\0\0\x06") }"#, ); // list from storage let storage_entries = mito diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index da36c827665..6d90f90b784 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -882,9 +882,9 @@ async fn test_cache_null_primary_key_with_format(flat_format: bool) { #[tokio::test] async fn test_list_ssts() { test_list_ssts_with_format(false, r#" -ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 2951, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 2951, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 2951, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"# , +ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"# , r#" StorageSstEntry { file_path: "test/11_0000000001/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/11_0000000001/index/.puffin", file_size: None, last_modified_ms: None, node_id: None } @@ -893,9 +893,9 @@ StorageSstEntry { file_path: "test/11_0000000002/index/.puffin", file_s StorageSstEntry { file_path: "test/22_0000000042/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/22_0000000042/index/.puffin", file_size: None, last_modified_ms: None, node_id: None }"#).await; test_list_ssts_with_format(true, r#" -ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 3365, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 3365, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 3365, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"#, +ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"#, r#" StorageSstEntry { file_path: "test/11_0000000001/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/11_0000000001/index/.puffin", file_size: None, last_modified_ms: None, node_id: None } @@ -967,6 +967,7 @@ async fn test_list_ssts_with_format( .map(|p| p.replace(&e.file_id, "")); e.file_id = "".to_string(); e.index_version = 0; + e.max_row_group_uncompressed_size = 0; format!("\n{:?}", e) }) .sorted() diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 76204da9cf6..b68b0d1cccd 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -839,6 +839,7 @@ impl MitoRegion { level: meta.level, file_path: sst_file_path(table_dir, meta.file_id(), path_type), file_size: meta.file_size, + max_row_group_uncompressed_size: meta.max_row_group_uncompressed_size, index_file_path, index_file_size, num_rows: meta.num_rows, @@ -847,6 +848,7 @@ impl MitoRegion { min_ts: meta.time_range.0, max_ts: meta.time_range.1, sequence: meta.sequence.map(|s| s.get()), + partition_expr: meta.partition_expr.as_ref().map(ToString::to_string), origin_region_id, node_id: None, visible, diff --git a/src/store-api/src/sst_entry.rs b/src/store-api/src/sst_entry.rs index f1bb1f9c36f..c70e909ebd6 100644 --- a/src/store-api/src/sst_entry.rs +++ b/src/store-api/src/sst_entry.rs @@ -65,6 +65,8 @@ pub struct ManifestSstEntry { pub file_path: String, /// File size in bytes. pub file_size: u64, + /// Maximum uncompressed row group size in bytes. Zero means unknown. + pub max_row_group_uncompressed_size: u64, /// Full path of the index file in object store. pub index_file_path: Option, /// File size of the index file in object store. @@ -81,6 +83,8 @@ pub struct ManifestSstEntry { pub max_ts: Timestamp, /// The sequence number associated with this file. pub sequence: Option, + /// Human-readable partition expression associated with this file. + pub partition_expr: Option, /// The region id of region that creates the file. pub origin_region_id: RegionId, /// The node id fetched from the manifest. @@ -122,6 +126,12 @@ impl ManifestSstEntry { ColumnSchema::new("visible", Ty::boolean_datatype(), false), ColumnSchema::new("primary_key_min", Ty::binary_datatype(), true), ColumnSchema::new("primary_key_max", Ty::binary_datatype(), true), + ColumnSchema::new( + "max_row_group_uncompressed_size", + Ty::uint64_datatype(), + false, + ), + ColumnSchema::new("partition_expr", Ty::string_datatype(), true), ])) } @@ -139,6 +149,8 @@ impl ManifestSstEntry { let levels = entries.iter().map(|e| e.level); let file_paths = entries.iter().map(|e| e.file_path.as_str()); let file_sizes = entries.iter().map(|e| e.file_size); + let max_row_group_uncompressed_sizes = + entries.iter().map(|e| e.max_row_group_uncompressed_size); let index_file_paths = entries.iter().map(|e| e.index_file_path.as_ref()); let index_file_sizes = entries.iter().map(|e| e.index_file_size); let num_rows = entries.iter().map(|e| e.num_rows); @@ -155,6 +167,7 @@ impl ManifestSstEntry { .map(|ts| ts.value()) }); let sequences = entries.iter().map(|e| e.sequence); + let partition_exprs = entries.iter().map(|e| e.partition_expr.as_ref()); let origin_region_ids = entries.iter().map(|e| e.origin_region_id.as_u64()); let node_ids = entries.iter().map(|e| e.node_id); let visible_flags = entries.iter().map(|e| Some(e.visible)); @@ -186,6 +199,10 @@ impl ManifestSstEntry { Arc::new(BooleanArray::from_iter(visible_flags)), Arc::new(BinaryArray::from_iter(primary_key_min)), Arc::new(BinaryArray::from_iter(primary_key_max)), + Arc::new(UInt64Array::from_iter_values( + max_row_group_uncompressed_sizes, + )), + Arc::new(StringArray::from_iter(partition_exprs)), ]; DfRecordBatch::try_new(schema.arrow_schema().clone(), columns) @@ -458,6 +475,7 @@ mod tests { level: 1, file_path: "/p1".to_string(), file_size: 100, + max_row_group_uncompressed_size: 80, index_file_path: None, index_file_size: None, num_rows: 10, @@ -466,6 +484,7 @@ mod tests { min_ts: Timestamp::new_millisecond(1000), // 1s -> 1_000_000_000ns max_ts: Timestamp::new_second(2), // 2s -> 2_000_000_000ns sequence: None, + partition_expr: Some("a < 10".to_string()), origin_region_id: region_id1, node_id: Some(1), visible: false, @@ -484,6 +503,7 @@ mod tests { level: 3, file_path: "/p2".to_string(), file_size: 200, + max_row_group_uncompressed_size: 160, index_file_path: Some("idx".to_string()), index_file_size: Some(11), num_rows: 20, @@ -492,6 +512,7 @@ mod tests { min_ts: Timestamp::new_nanosecond(5), // 5ns max_ts: Timestamp::new_microsecond(2000), // 2ms -> 2_000_000ns sequence: Some(9), + partition_expr: None, origin_region_id: region_id2, node_id: None, visible: true, @@ -506,6 +527,44 @@ mod tests { // Schema checks assert_eq!(schema.arrow_schema().fields().len(), batch.num_columns()); assert_eq!(2, batch.num_rows()); + let expected_columns = [ + "table_dir", + "region_id", + "table_id", + "region_number", + "region_group", + "region_sequence", + "file_id", + "index_version", + "level", + "file_path", + "file_size", + "index_file_path", + "index_file_size", + "num_rows", + "num_row_groups", + "num_series", + "min_ts", + "max_ts", + "sequence", + "origin_region_id", + "node_id", + "visible", + "primary_key_min", + "primary_key_max", + "max_row_group_uncompressed_size", + "partition_expr", + ]; + assert_eq!( + expected_columns, + schema + .arrow_schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>() + .as_slice() + ); for (i, f) in schema.arrow_schema().fields().iter().enumerate() { assert_eq!(f.name(), batch.schema().field(i).name()); assert_eq!(f.is_nullable(), batch.schema().field(i).is_nullable()); @@ -704,6 +763,22 @@ mod tests { .unwrap(); assert_eq!(b"zzz", primary_key_max.value(0)); assert!(primary_key_max.is_null(1)); + + let max_row_group_uncompressed_sizes = batch + .column(24) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(80, max_row_group_uncompressed_sizes.value(0)); + assert_eq!(160, max_row_group_uncompressed_sizes.value(1)); + + let partition_exprs = batch + .column(25) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!("a < 10", partition_exprs.value(0)); + assert!(partition_exprs.is_null(1)); } #[test] diff --git a/tests/cases/standalone/common/information_schema/ssts.result b/tests/cases/standalone/common/information_schema/ssts.result index caf6a239607..9777e3b12a9 100644 --- a/tests/cases/standalone/common/information_schema/ssts.result +++ b/tests/cases/standalone/common/information_schema/ssts.result @@ -1,33 +1,35 @@ DESC TABLE information_schema.ssts_manifest; -+------------------+---------------------+-----+------+---------+---------------+ -| Column | Type | Key | Null | Default | Semantic Type | -+------------------+---------------------+-----+------+---------+---------------+ -| table_dir | String | | NO | | FIELD | -| region_id | UInt64 | | NO | | FIELD | -| table_id | UInt32 | | NO | | FIELD | -| region_number | UInt32 | | NO | | FIELD | -| region_group | UInt8 | | NO | | FIELD | -| region_sequence | UInt32 | | NO | | FIELD | -| file_id | String | | NO | | FIELD | -| index_version | UInt64 | | NO | | FIELD | -| level | UInt8 | | NO | | FIELD | -| file_path | String | | NO | | FIELD | -| file_size | UInt64 | | NO | | FIELD | -| index_file_path | String | | YES | | FIELD | -| index_file_size | UInt64 | | YES | | FIELD | -| num_rows | UInt64 | | NO | | FIELD | -| num_row_groups | UInt64 | | NO | | FIELD | -| num_series | UInt64 | | YES | | FIELD | -| min_ts | TimestampNanosecond | | YES | | FIELD | -| max_ts | TimestampNanosecond | | YES | | FIELD | -| sequence | UInt64 | | YES | | FIELD | -| origin_region_id | UInt64 | | NO | | FIELD | -| node_id | UInt64 | | YES | | FIELD | -| visible | Boolean | | NO | | FIELD | -| primary_key_min | Binary | | YES | | FIELD | -| primary_key_max | Binary | | YES | | FIELD | -+------------------+---------------------+-----+------+---------+---------------+ ++---------------------------------+---------------------+-----+------+---------+---------------+ +| Column | Type | Key | Null | Default | Semantic Type | ++---------------------------------+---------------------+-----+------+---------+---------------+ +| table_dir | String | | NO | | FIELD | +| region_id | UInt64 | | NO | | FIELD | +| table_id | UInt32 | | NO | | FIELD | +| region_number | UInt32 | | NO | | FIELD | +| region_group | UInt8 | | NO | | FIELD | +| region_sequence | UInt32 | | NO | | FIELD | +| file_id | String | | NO | | FIELD | +| index_version | UInt64 | | NO | | FIELD | +| level | UInt8 | | NO | | FIELD | +| file_path | String | | NO | | FIELD | +| file_size | UInt64 | | NO | | FIELD | +| index_file_path | String | | YES | | FIELD | +| index_file_size | UInt64 | | YES | | FIELD | +| num_rows | UInt64 | | NO | | FIELD | +| num_row_groups | UInt64 | | NO | | FIELD | +| num_series | UInt64 | | YES | | FIELD | +| min_ts | TimestampNanosecond | | YES | | FIELD | +| max_ts | TimestampNanosecond | | YES | | FIELD | +| sequence | UInt64 | | YES | | FIELD | +| origin_region_id | UInt64 | | NO | | FIELD | +| node_id | UInt64 | | YES | | FIELD | +| visible | Boolean | | NO | | FIELD | +| primary_key_min | Binary | | YES | | FIELD | +| primary_key_max | Binary | | YES | | FIELD | +| max_row_group_uncompressed_size | UInt64 | | NO | | FIELD | +| partition_expr | String | | YES | | FIELD | ++---------------------------------+---------------------+-----+------+---------+---------------+ DESC TABLE information_schema.ssts_storage; @@ -99,13 +101,13 @@ ADMIN FLUSH_TABLE('sst_case'); -- SQLNESS REPLACE (/public/\d+) /public/ SELECT * FROM information_schema.ssts_manifest order by file_path; -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | max_row_group_uncompressed_size | partition_expr | ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 || a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc || a >=AND a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 || a >=| ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ -- SQLNESS REPLACE (\s+\d+\s+) -- SQLNESS REPLACE ([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}) @@ -168,15 +170,15 @@ ADMIN FLUSH_TABLE('sst_case'); -- SQLNESS REPLACE (/public/\d+) /public/ SELECT * FROM information_schema.ssts_manifest order by region_id, sequence; -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true || 018000007c | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true ||| -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | max_row_group_uncompressed_size | partition_expr | ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 || a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true || 018000007c || a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc || a >=AND a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true |||| a >=AND a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 || a >=| ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ -- SQLNESS REPLACE (\s+\d+\s+) -- SQLNESS REPLACE ([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}) diff --git a/tests/cases/standalone/common/system/information_schema.result b/tests/cases/standalone/common/system/information_schema.result index e05cd064e1d..17e1f0ef80d 100644 --- a/tests/cases/standalone/common/system/information_schema.result +++ b/tests/cases/standalone/common/system/information_schema.result @@ -436,6 +436,7 @@ order by table_schema, table_name, column_name; | greptime | information_schema | ssts_manifest | index_file_size | 13 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | | greptime | information_schema | ssts_manifest | index_version | 8 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | | greptime | information_schema | ssts_manifest | level | 9 | | | 3 | 0 | | | | | | select,insert | | UInt8 | tinyint unsigned | FIELD | | NO | tinyint unsigned | | | +| greptime | information_schema | ssts_manifest | max_row_group_uncompressed_size | 25 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | | greptime | information_schema | ssts_manifest | max_ts | 18 | | | | | 9 | | | | | select,insert | | TimestampNanosecond | timestamp(9) | FIELD | | YES | timestamp(9) | | | | greptime | information_schema | ssts_manifest | min_ts | 17 | | | | | 9 | | | | | select,insert | | TimestampNanosecond | timestamp(9) | FIELD | | YES | timestamp(9) | | | | greptime | information_schema | ssts_manifest | node_id | 21 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | @@ -443,6 +444,7 @@ order by table_schema, table_name, column_name; | greptime | information_schema | ssts_manifest | num_rows | 14 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | | greptime | information_schema | ssts_manifest | num_series | 16 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | | greptime | information_schema | ssts_manifest | origin_region_id | 20 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | +| greptime | information_schema | ssts_manifest | partition_expr | 26 | 2147483647 | 2147483647 | | | | utf8 | utf8_bin | | | select,insert | | String | string | FIELD | | YES | string | | | | greptime | information_schema | ssts_manifest | primary_key_max | 24 | | | | | | | | | | select,insert | | Binary | varbinary | FIELD | | YES | varbinary | | | | greptime | information_schema | ssts_manifest | primary_key_min | 23 | | | | | | | | | | select,insert | | Binary | varbinary | FIELD | | YES | varbinary | | | | greptime | information_schema | ssts_manifest | region_group | 5 | | | 3 | 0 | | | | | | select,insert | | UInt8 | tinyint unsigned | FIELD | | NO | tinyint unsigned | | | From c01de4afdc7e2ef65edb5efdfaccebfd62b0d1c6 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 28 Aug 2026 08:49:49 +0000 Subject: [PATCH 122/126] fix(operator): whitelist private system table auto create (#8930) Signed-off-by: WenyXu --- src/common/meta/src/datanode.rs | 2 + .../src/handler/persist_stats_handler.rs | 8 +- src/operator/src/insert.rs | 68 ++++++-- tests-integration/src/cluster.rs | 14 +- tests-integration/src/standalone.rs | 14 ++ tests-integration/tests/database_ddl_event.rs | 148 ++++++++++++++++- tests-integration/tests/grpc.rs | 155 +++++++++++++++++- 7 files changed, 385 insertions(+), 24 deletions(-) diff --git a/src/common/meta/src/datanode.rs b/src/common/meta/src/datanode.rs index 8af86462129..d50c37c381b 100644 --- a/src/common/meta/src/datanode.rs +++ b/src/common/meta/src/datanode.rs @@ -31,6 +31,8 @@ use crate::heartbeat::utils::get_datanode_workloads; const DATANODE_STAT_PREFIX: &str = "__meta_datanode_stat"; pub const REGION_STATISTIC_KEY: &str = "__region_statistic"; +/// The private table name for persisted region statistics. +pub const REGION_STATS_HISTORY_TABLE_NAME: &str = "region_statistics_history"; lazy_static! { pub(crate) static ref DATANODE_LEASE_KEY_PATTERN: Regex = diff --git a/src/meta-srv/src/handler/persist_stats_handler.rs b/src/meta-srv/src/handler/persist_stats_handler.rs index b1378c39e6b..fcdb9fa04bc 100644 --- a/src/meta-srv/src/handler/persist_stats_handler.rs +++ b/src/meta-srv/src/handler/persist_stats_handler.rs @@ -22,7 +22,7 @@ use client::inserter::{Context as InserterContext, Inserter}; use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME; use common_macro::{Schema, ToRow}; use common_meta::DatanodeId; -use common_meta::datanode::RegionStat; +use common_meta::datanode::{REGION_STATS_HISTORY_TABLE_NAME, RegionStat}; use common_telemetry::warn; use dashmap::DashMap; use store_api::region_engine::RegionRole; @@ -40,8 +40,6 @@ pub struct PersistStatsHandler { persist_interval: Duration, } -/// The name of the table to persist region stats. -const META_REGION_STATS_HISTORY_TABLE_NAME: &str = "region_statistics_history"; /// The default context to persist region stats. const DEFAULT_CONTEXT: InserterContext = InserterContext { catalog: DEFAULT_CATALOG_NAME, @@ -207,7 +205,7 @@ impl PersistStatsHandler { &DEFAULT_CONTEXT, RowInsertRequests { inserts: vec![RowInsertRequest { - table_name: META_REGION_STATS_HISTORY_TABLE_NAME.to_string(), + table_name: REGION_STATS_HISTORY_TABLE_NAME.to_string(), rows: Some(Rows { schema: PersistRegionStat::schema(), rows, @@ -537,7 +535,7 @@ mod tests { }; assert_eq!( request.table_name, - META_REGION_STATS_HISTORY_TABLE_NAME.to_string() + REGION_STATS_HISTORY_TABLE_NAME.to_string() ); assert_eq!(request.rows.unwrap().rows, vec![expected_row]); diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index ead517ab80f..52a73892dc0 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -29,12 +29,15 @@ use api::v1::{ use catalog::CatalogManagerRef; use client::{OutputData, OutputMeta}; use common_catalog::consts::{ - PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN, TRACE_TABLE_NAME, - TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table, + DEFAULT_PRIVATE_SCHEMA_NAME, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN, + TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table, trace_operations_table_name, trace_services_table_name, }; +use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME; +use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; use common_grpc_expr::util::ColumnExpr; use common_meta::cache::TableFlownodeSetCacheRef; +use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME; use common_meta::node_manager::{AffectedRows, NodeManagerRef}; use common_meta::peer::Peer; use common_meta::rpc::ddl::TriggerReason; @@ -490,9 +493,9 @@ impl Inserter { Ok(inserts) } - /// Returns `None` if auto table creation is allowed, or `Some(reason)` if - /// disabled by either the global config or the request hint. The reason tells - /// which one, for a clearer error. + /// Returns `Some(reason)` if the config or request hint disables automatic + /// table creation. Exempt private system tables are handled by + /// [`Self::is_auto_create_exempt_private_table`]. fn auto_create_disabled_reason(&self, ctx: &QueryContextRef) -> Result> { let auto_create_table_hint = ctx .extension(AUTO_CREATE_TABLE_KEY) @@ -514,6 +517,16 @@ impl Inserter { }) } + /// Returns whether a private system table may infer and reconcile its schema + /// even when automatic table creation is disabled. + fn is_auto_create_exempt_private_table(schema: &str, table: &str) -> bool { + schema == DEFAULT_PRIVATE_SCHEMA_NAME + && matches!( + table, + DEFAULT_EVENTS_TABLE_NAME | SLOW_QUERY_TABLE_NAME | REGION_STATS_HISTORY_TABLE_NAME + ) + } + /// Ensures a trace table has the request-global schema without requiring a /// padded data row to drive on-demand creation or alteration. When /// `alter_existing` is false, a table created after planning is left for the @@ -571,12 +584,21 @@ impl Inserter { let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND .with_label_values(&[auto_create_table_type.as_str()]) .start_timer(); - let catalog = ctx.current_catalog(); let schema = ctx.current_schema(); + let auto_create_disabled_reason = self.auto_create_disabled_reason(ctx)?; + // Enabled batches permit every table, so only disabled batches need a whitelist scan. + let has_auto_create_exempt_table = auto_create_disabled_reason.is_some() + && requests + .inserts + .iter() + .any(|req| Self::is_auto_create_exempt_private_table(&schema, &req.table_name)); let mut table_infos = HashMap::new(); - if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? { + // Without exempt tables, verify existing tables and reject missing ones without inferring schemas. + if let Some(disabled_reason) = auto_create_disabled_reason + && !has_auto_create_exempt_table + { let mut instant_table_ids = HashSet::new(); for req in &requests.inserts { let table = match self.get_table(catalog, &schema, &req.table_name).await? { @@ -618,20 +640,25 @@ impl Inserter { let mut per_table_semantics: Option> = None; for req in &mut requests.inserts { + // Mixed batches need a per-table decision so an exempt table cannot authorize others. + let auto_create_allowed = auto_create_disabled_reason.is_none() + || Self::is_auto_create_exempt_private_table(&schema, &req.table_name); match self.get_table(catalog, &schema, &req.table_name).await? { Some(table) => { let table_info = table.table_info(); if table_info.is_ttl_instant_table() { instant_table_ids.insert(table_info.table_id()); } - if let Some(alter_expr) = self.get_alter_table_expr_on_demand( - req, - &table, - ctx, - accommodate_existing_schema, - is_single_value, - auto_create_table_type.alter_existing(), - )? { + if auto_create_allowed + && let Some(alter_expr) = self.get_alter_table_expr_on_demand( + req, + &table, + ctx, + accommodate_existing_schema, + is_single_value, + auto_create_table_type.alter_existing(), + )? + { alter_tables.push(alter_expr); need_refresh_table_infos.insert(( catalog.to_string(), @@ -655,6 +682,17 @@ impl Inserter { } table_infos.insert(table_info.table_id(), table_info); } + None if !auto_create_allowed + && let Some(disabled_reason) = auto_create_disabled_reason => + { + return InvalidInsertRequestSnafu { + reason: format!( + "Table `{}` does not exist, and {}", + req.table_name, disabled_reason, + ), + } + .fail(); + } None => { let semantic_index = per_table_semantics .get_or_insert_with(|| parse_per_table_semantic_index(ctx)) diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index d0cd2e584c3..e40c707fe98 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -169,6 +169,7 @@ pub struct GreptimeDbClusterBuilder { metasrv_wal_config: MetasrvWalConfig, datanode_gc_config: GcConfig, metasrv_gc_config: GcSchedulerOptions, + frontend_auto_create_table: bool, shared_home_dir: Option>, meta_selector: Option, } @@ -202,6 +203,7 @@ impl GreptimeDbClusterBuilder { metasrv_wal_config: MetasrvWalConfig::default(), datanode_gc_config: GcConfig::default(), metasrv_gc_config: GcSchedulerOptions::default(), + frontend_auto_create_table: true, shared_home_dir: None, meta_selector: None, } @@ -248,6 +250,13 @@ impl GreptimeDbClusterBuilder { self } + /// Sets whether the test frontend automatically creates tables on write. + #[must_use] + pub fn with_frontend_auto_create_table(mut self, auto_create_table: bool) -> Self { + self.frontend_auto_create_table = auto_create_table; + self + } + #[must_use] pub fn with_shared_home_dir(mut self, shared_home_dir: Arc) -> Self { self.shared_home_dir = Some(shared_home_dir); @@ -529,7 +538,10 @@ impl GreptimeDbClusterBuilder { } fn build_frontend_options(&self) -> FrontendOptions { - let mut fe_opts = FrontendOptions::default(); + let mut fe_opts = FrontendOptions { + auto_create_table: self.frontend_auto_create_table, + ..Default::default() + }; // Choose a random unused port between [14000, 24000] for local test to avoid conflicts. let port_range = 14000..=24000; diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 16fc2e65634..0704b184e89 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -27,6 +27,7 @@ use common_base::Plugins; use common_catalog::consts::{MIN_USER_FLOW_ID, MIN_USER_TABLE_ID}; use common_config::KvBackendConfig; use common_datasource::object_store::LocalFileAccess; +use common_event_recorder::EventRecorderOptions; use common_meta::cache::LayeredCacheRegistryBuilder; use common_meta::ddl::flow_meta::FlowMetadataAllocator; use common_meta::ddl::table_meta::TableMetadataAllocator; @@ -84,6 +85,7 @@ pub struct GreptimeDbStandaloneBuilder { default_store: Option, plugin: Option, slow_query_options: SlowQueryOptions, + event_recorder_options: EventRecorderOptions, auto_create_table: bool, } @@ -102,6 +104,7 @@ impl GreptimeDbStandaloneBuilder { threshold: Duration::from_secs(1), ..Default::default() }, + event_recorder_options: EventRecorderOptions::default(), auto_create_table: true, } } @@ -112,6 +115,16 @@ impl GreptimeDbStandaloneBuilder { self } + /// Sets the event recorder options for the standalone instance. + #[must_use] + pub fn with_event_recorder_options( + mut self, + event_recorder_options: EventRecorderOptions, + ) -> Self { + self.event_recorder_options = event_recorder_options; + self + } + #[must_use] pub fn with_default_store_type(self, store_type: StorageType) -> Self { Self { @@ -368,6 +381,7 @@ impl GreptimeDbStandaloneBuilder { wal: self.metasrv_wal_config.clone().into(), grpc: GrpcOptions::default().with_server_addr("127.0.0.1:4001"), slow_query: self.slow_query_options.clone(), + event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, // Tests cover the descriptor, so they run with it enabled. otlp: frontend::service_config::OtlpOptions { diff --git a/tests-integration/tests/database_ddl_event.rs b/tests-integration/tests/database_ddl_event.rs index 4dfda4f32b9..c08b81211de 100644 --- a/tests-integration/tests/database_ddl_event.rs +++ b/tests-integration/tests/database_ddl_event.rs @@ -15,20 +15,163 @@ use std::sync::Arc; use client::{Database, OutputData}; +use common_event_recorder::{EventRecorderOptions, EventTypeFilter}; use common_test_util::temp_dir::create_temp_dir; use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; use tests_integration::test_util::{ - StorageType, get_test_store_config, setup_authenticated_grpc_database, + StorageType, execute_sql, get_test_store_config, setup_authenticated_grpc_database, }; use crate::event_recorder_test_util::{ - assert_procedure_actor, assert_single_event, find_eventually_string, + assert_eventually_eq, assert_procedure_actor, assert_single_event, find_eventually_string, }; const DATABASE_NAME: &str = "database_ddl_events"; const PROCEDURE_ACTOR: &str = "procedure_actor"; const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; +const CREATION_DATABASE_NAME: &str = "event_schema_creation"; +const RECONCILIATION_DATABASE_NAME: &str = "event_schema_reconciliation"; + +#[tokio::test(flavor = "multi_thread")] +async fn test_event_table_auto_creation_with_auto_create_disabled() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_event_table_auto_creation") + .with_auto_create_table(false) + .with_event_recorder_options(EventRecorderOptions { + event_types: Arc::new(EventTypeFilter::Only( + [String::from("create_database")].into_iter().collect(), + )), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; + database + .sql(format!("CREATE DATABASE {CREATION_DATABASE_NAME}")) + .await + .unwrap(); + + assert_eventually_eq( + instance, + "SELECT count(*) AS events_tables \ + FROM information_schema.tables \ + WHERE table_catalog = 'greptime' \ + AND table_schema = 'greptime_private' \ + AND table_name = 'events'", + "+---------------+\n| events_tables |\n+---------------+\n| 1 |\n+---------------+", + ) + .await; + + let procedure_id = find_eventually_string( + instance, + &format!( + "SELECT procedure_id FROM greptime_private.events \ + WHERE type = 'create_database' \ + AND schema_name = '{CREATION_DATABASE_NAME}' \ + AND json_path_match(procedure_trigger, '$.type == \"Submitted\"') \ + ORDER BY timestamp DESC LIMIT 1" + ), + "procedure_id", + ) + .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_event_table_schema_reconciliation_with_auto_create_disabled() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_event_table_schema_reconciliation") + .with_auto_create_table(false) + .with_event_recorder_options(EventRecorderOptions { + event_types: Arc::new(EventTypeFilter::Only( + [String::from("create_database")].into_iter().collect(), + )), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + + // Matches the pre-actor procedure-event schema. The recorder filter keeps + // this setup DDL from creating an event before the regression is exercised. + execute_sql( + instance, + r#" + CREATE TABLE greptime_private.events ( + "type" STRING, + payload JSON, + "timestamp" TIMESTAMP(9) NOT NULL, + procedure_id STRING, + procedure_state STRING, + procedure_error STRING, + procedure_trigger JSON, + catalog_name STRING, + schema_name STRING, + event_context JSON, + TIME INDEX ("timestamp"), + PRIMARY KEY ("type") + ) WITH (append_mode = 'true') + "#, + ) + .await; + + assert_eventually_eq( + instance, + "SELECT count(*) AS actor_columns \ + FROM information_schema.columns \ + WHERE table_catalog = 'greptime' \ + AND table_schema = 'greptime_private' \ + AND table_name = 'events' \ + AND column_name = 'actor'", + "+---------------+\n| actor_columns |\n+---------------+\n| 0 |\n+---------------+", + ) + .await; + + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; + database + .sql(format!("CREATE DATABASE {RECONCILIATION_DATABASE_NAME}")) + .await + .unwrap(); + + assert_eventually_eq( + instance, + "SELECT count(*) AS actor_columns \ + FROM information_schema.columns \ + WHERE table_catalog = 'greptime' \ + AND table_schema = 'greptime_private' \ + AND table_name = 'events' \ + AND column_name = 'actor'", + "+---------------+\n| actor_columns |\n+---------------+\n| 1 |\n+---------------+", + ) + .await; + + let procedure_id = find_eventually_string( + instance, + &format!( + "SELECT procedure_id FROM greptime_private.events \ + WHERE type = 'create_database' \ + AND schema_name = '{RECONCILIATION_DATABASE_NAME}' \ + AND json_path_match(procedure_trigger, '$.type == \"Submitted\"') \ + ORDER BY timestamp DESC LIMIT 1" + ), + "procedure_id", + ) + .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; +} #[tokio::test(flavor = "multi_thread")] async fn test_database_ddl_events() { @@ -43,6 +186,7 @@ async fn test_database_ddl_events() { let cluster = GreptimeDbClusterBuilder::new("test_database_ddl_events") .await .with_datanodes(1) + .with_frontend_auto_create_table(false) .with_store_config(store_config) .with_shared_home_dir(Arc::new(home_dir)) .build(true) diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 737b4fb6c1d..f06ffa65224 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -26,9 +26,12 @@ use api::v1::{ use auth::user_provider_from_option; use base64::prelude::{BASE64_STANDARD, Engine as _}; use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputData}; -use common_catalog::consts::MITO_ENGINE; +use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, MITO_ENGINE}; +use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME; +use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; use common_grpc::channel_manager::ClientTlsOption; use common_memory_manager::OnExhaustedPolicy; +use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME; use common_query::Output; use common_recordbatch::RecordBatches; use common_runtime::Runtime; @@ -98,6 +101,8 @@ macro_rules! grpc_tests { test_auto_create_table, test_auto_create_table_with_hints, test_auto_create_table_disabled_by_config, + test_private_system_tables_auto_create_table_with_global_disabled, + test_private_system_tables_bypass_auto_create_hint, test_otel_arrow_auth, test_otel_arrow_exponential_histogram, test_insert_and_select, @@ -679,6 +684,154 @@ pub async fn test_auto_create_table_disabled_by_config(store_type: StorageType) let _ = fe_grpc_server.shutdown().await; } +pub async fn test_private_system_tables_auto_create_table_with_global_disabled( + store_type: StorageType, +) { + let (_db, fe_grpc_server) = setup_grpc_server_with_auto_create_table_disabled( + store_type, + "test_private_system_tables_auto_create_table_with_global_disabled", + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + + let grpc_client = Client::with_urls(vec![addr]); + let db = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_PRIVATE_SCHEMA_NAME, + grpc_client, + ); + let (host, cpu, mem, ts) = expect_data(); + + for table_name in [ + DEFAULT_EVENTS_TABLE_NAME, + SLOW_QUERY_TABLE_NAME, + REGION_STATS_HISTORY_TABLE_NAME, + ] { + let result = db + .insert(InsertRequests { + inserts: vec![InsertRequest { + table_name: table_name.to_string(), + columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()], + row_count: 4, + }], + }) + .await; + assert_eq!(result.unwrap(), 4); + } + + let load = Column { + column_name: "load".to_string(), + values: Some(column::Values { + f64_values: vec![0.4, 0.5, 0.6, 0.7], + ..Default::default() + }), + semantic_type: SemanticType::Field as i32, + datatype: ColumnDataType::Float64 as i32, + ..Default::default() + }; + let result = db + .insert(InsertRequests { + inserts: vec![InsertRequest { + table_name: DEFAULT_EVENTS_TABLE_NAME.to_string(), + columns: vec![host, cpu, mem, ts, load], + row_count: 4, + }], + }) + .await; + assert_eq!(result.unwrap(), 4); + + let output = db + .sql(format!("SHOW CREATE TABLE {DEFAULT_EVENTS_TABLE_NAME}")) + .await + .unwrap(); + let record_batches = match output.data { + OutputData::RecordBatches(record_batches) => record_batches, + OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(), + OutputData::AffectedRows(_) => unreachable!(), + }; + assert!(record_batches.pretty_print().unwrap().contains("\"load\"")); + + let _ = fe_grpc_server.shutdown().await; +} + +pub async fn test_private_system_tables_bypass_auto_create_hint(store_type: StorageType) { + let (_db, fe_grpc_server) = setup_grpc_server( + store_type, + "test_private_system_tables_bypass_auto_create_hint", + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + + let grpc_client = Client::with_urls(vec![addr]); + let db = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_PRIVATE_SCHEMA_NAME, + grpc_client, + ); + let (host, cpu, mem, ts) = expect_data(); + + for table_name in [ + DEFAULT_EVENTS_TABLE_NAME, + SLOW_QUERY_TABLE_NAME, + REGION_STATS_HISTORY_TABLE_NAME, + ] { + let result = db + .insert_with_hints( + InsertRequests { + inserts: vec![InsertRequest { + table_name: table_name.to_string(), + columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()], + row_count: 4, + }], + }, + &[("auto_create_table", "false")], + ) + .await; + assert_eq!(result.unwrap(), 4); + } + + let ordinary_table = "ordinary_private_table"; + let result = db + .insert_with_hints( + InsertRequests { + inserts: vec![ + InsertRequest { + table_name: DEFAULT_EVENTS_TABLE_NAME.to_string(), + columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()], + row_count: 4, + }, + InsertRequest { + table_name: ordinary_table.to_string(), + columns: vec![host, cpu, mem, ts], + row_count: 4, + }, + ], + }, + &[("auto_create_table", "false")], + ) + .await; + let err = result.unwrap_err().to_string(); + assert!( + err.contains(ordinary_table) && err.contains("auto_create_table"), + "unexpected error: {err}" + ); + + let output = db.sql("SHOW TABLES").await.unwrap(); + let record_batches = match output.data { + OutputData::RecordBatches(record_batches) => record_batches, + OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(), + OutputData::AffectedRows(_) => unreachable!(), + }; + assert!( + !record_batches + .pretty_print() + .unwrap() + .contains(ordinary_table) + ); + + let _ = fe_grpc_server.shutdown().await; +} + fn expect_data() -> (Column, Column, Column, Column) { // testing data: let expected_host_col = Column { From c0b8612c41bf0b2e7e46d76c5251777c0a68057d Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Fri, 28 Aug 2026 13:45:06 +0000 Subject: [PATCH 123/126] fix(cmd): configure meta client in frontend plugin test (#8977) Provide valid distributed frontend options in `src/cmd/src/frontend.rs` so enterprise plugin setup can preserve prefilled heartbeat extensions without weakening production validation. Signed-off-by: Lei, HUANG --- src/cmd/src/frontend.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index 1f95a3c54b4..6cae195f3e6 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -685,10 +685,14 @@ mod tests { let extensions = FrontendHeartbeatExtensions::default(); let shared_extensions = extensions.clone(); + let options = frontend::frontend::FrontendOptions { + meta_client: Some(MetaClientOptions::default()), + ..Default::default() + }; let plugins = prepare_frontend_plugins( plugins_with_heartbeat_extensions(extensions), &[], - &frontend::frontend::FrontendOptions::default(), + &options, None, ) .await From 9cfbc42126704c2cc8c358e85d89cec9482c9a9f Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Sat, 29 Aug 2026 06:51:48 +0000 Subject: [PATCH 124/126] refactor(mito2): prioritize file count in TWCS picker (#8765) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(mito2): prioritize file count in TWCS picker Signed-off-by: Lei, HUANG * test(mito2): migrate TWCS tests to async picker API Signed-off-by: Lei, HUANG * refactor(mito2): remove legacy reduce_runs and merge_seq_files pickers Signed-off-by: Lei, HUANG * feat(mito2): balance TWCS picks by file group Signed-off-by: Lei, HUANG * refactor(mito2): treat multi-SST groups as barriers in TWCS pick_count_first Previously pick_count_first filtered out multi-SST file groups and could pick singleton groups across them in one interval. Now multi-SST groups split the candidates into independent segments, so a single pick never crosses a multi-SST group. Signed-off-by: Lei, HUANG * refactor(mito2): remove redundant overlaps_files_left_behind check The rebase onto main kept both the upstream fix (#8872) and the branch's selected_overlaps_unselected check for the same deletion-marker problem. The upstream check is fully subsumed: files_to_merge only differs from the window files in append mode, where filter_deleted is already false, and selected_overlaps_unselected treats partially-selected groups as unselected, making it strictly stronger. Signed-off-by: Lei, HUANG * perf(mito2): pre-filter by selection time span in selected_overlaps_unselected A window group outside the overall time span of the selected groups cannot overlap any of them, so skip the precise overlap check (and the per-group file id set lookup) for it. In typical TWCS windows the pick is clustered in one segment, so most unselected groups are rejected in O(1). Signed-off-by: Lei, HUANG * refactor(mito2): remove FileGroup abstraction from TWCS picker Signed-off-by: Lei, HUANG * feat(mito2): drain compactable backlog after successful compaction Signed-off-by: Lei, HUANG * refactor(mito2): score TWCS candidates by predicted progress Redefine what a pick candidate is worth: an interval is only eligible if compacting it makes progress on at least one axis — it reduces the physical file count given the max_output_file_size split threshold (predicted output = ceil(input bytes / threshold)), or it resolves at least one overlap between sorted runs. A pure rewrite that achieves neither (e.g. 32 large balanced files whose output would split back into just as many SSTs) is skipped instead of burning I/O. The candidate metrics are accumulated in a Candidate struct as the interval expands, and the score ranks by predicted file reduction, then overlap participants, then smaller input bytes. With no output size limit the behavior degenerates to the previous count-first rule. Signed-off-by: Lei, HUANG * fix(mito2): make TWCS input limit configurable Allow operators to tune the maximum SST inputs through a hidden environment variable while keeping 32 as the validated default. Signed-off-by: Lei, HUANG * fix(mito2): apply TWCS trigger at window level Check the total physical SST count before candidate selection so trigger values above the per-task input limit still allow bounded compactions. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/benches/bench_compaction_picker.rs | 66 -- src/mito2/src/compaction/run.rs | 481 ++--------- src/mito2/src/compaction/scheduler.rs | 27 +- src/mito2/src/compaction/scheduler_test.rs | 176 +++- src/mito2/src/compaction/twcs.rs | 863 +++++++++++++++---- src/mito2/src/engine/compaction_test.rs | 25 +- src/mito2/src/engine/merge_mode_test.rs | 3 +- src/mito2/src/gc/worker_test.rs | 10 +- src/mito2/src/worker/handle_compaction.rs | 6 + 9 files changed, 981 insertions(+), 676 deletions(-) diff --git a/src/mito2/benches/bench_compaction_picker.rs b/src/mito2/benches/bench_compaction_picker.rs index 74142127aee..439b9e6f226 100644 --- a/src/mito2/benches/bench_compaction_picker.rs +++ b/src/mito2/benches/bench_compaction_picker.rs @@ -17,7 +17,6 @@ use std::hint::black_box; use criterion::{Criterion, criterion_group, criterion_main}; use mito2::compaction::run::{ Item, Ranged, SortedRun, find_overlapping_items, find_sorted_runs, find_sorted_runs_original, - merge_seq_files, reduce_runs, }; #[derive(Clone, Debug, Eq, Hash, PartialEq)] @@ -49,20 +48,6 @@ impl Item for MockFile { } } -fn generate_test_files(n: usize) -> Vec { - let mut files = Vec::with_capacity(n); - for _ in 0..n { - // Create slightly overlapping ranges to force multiple sorted runs - files.push(MockFile { - start: 0, - end: 10, - pk: 0, - size: 10, - }); - } - files -} - fn generate_same_timestamp_files(total_files: usize, files_per_timestamp: usize) -> Vec { const TIMESTAMP_INTERVAL_SECS: i64 = 10 * 60; @@ -107,21 +92,6 @@ fn bench_find_sorted_runs(c: &mut Criterion) { group.finish(); } -fn bench_reduce_runs(c: &mut Criterion) { - let mut group = c.benchmark_group("reduce_runs"); - - for size in [10, 100, 1000].iter() { - group.bench_function(format!("size_{}", size), |b| { - let mut files = generate_test_files(*size); - let runs = find_sorted_runs(&mut files); - b.iter(|| { - reduce_runs(black_box(runs.clone())); - }); - }); - } - group.finish(); -} - fn bench_find_overlapping_items(c: &mut Criterion) { let mut group = c.benchmark_group("find_overlapping_items"); @@ -158,45 +128,9 @@ fn bench_find_overlapping_items(c: &mut Criterion) { group.finish(); } -fn bench_merge_seq_files(c: &mut Criterion) { - let mut group = c.benchmark_group("merge_seq_files"); - - for size in [10, 100, 1000].iter() { - group.bench_function(format!("size_{}", size), |b| { - // Create a set of files with varying sizes - let mut files = Vec::with_capacity(*size); - - for i in 0..*size { - // Create files with different sizes to test the scoring algorithm - let file_size = if i % 3 == 0 { - 5 - } else if i % 3 == 1 { - 10 - } else { - 15 - }; - - files.push(MockFile { - start: i as i64, - end: (i + 1) as i64, - pk: 0, - size: file_size, - }); - } - - b.iter(|| { - merge_seq_files(black_box(&files), black_box(Some(50))); - }); - }); - } - group.finish(); -} - criterion_group!( benches, bench_find_sorted_runs, - bench_reduce_runs, bench_find_overlapping_items, - bench_merge_seq_files ); criterion_main!(benches); diff --git a/src/mito2/src/compaction/run.rs b/src/mito2/src/compaction/run.rs index 1ba45690215..5698a6ce667 100644 --- a/src/mito2/src/compaction/run.rs +++ b/src/mito2/src/compaction/run.rs @@ -20,14 +20,9 @@ use std::collections::BinaryHeap; use bytes::{Buf, Bytes}; use common_base::BitVec; -use common_base::readable_size::ReadableSize; use common_time::Timestamp; -use smallvec::{SmallVec, smallvec}; -use crate::sst::file::{FileHandle, RegionFileId}; - -/// Default max compaction output file size when not specified. -const DEFAULT_MAX_OUTPUT_SIZE: u64 = ReadableSize::gb(2).as_bytes(); +use crate::sst::file::FileHandle; /// Trait for any items with specific range (both boundaries are inclusive). pub trait Ranged { @@ -152,71 +147,11 @@ pub trait Item: Ranged + Clone { fn size(&self) -> usize; } -/// A group of files that are created by the same compaction task. -#[derive(Debug, Clone)] -pub struct FileGroup { - files: SmallVec<[FileHandle; 2]>, - size: usize, - num_rows: usize, - min_timestamp: Timestamp, - max_timestamp: Timestamp, - primary_key_range: Option<(Bytes, Bytes)>, -} - -impl FileGroup { - pub(crate) fn new_with_file(file: FileHandle) -> Self { - let size = file.size() as usize; - let (min_timestamp, max_timestamp) = file.time_range(); - let num_rows = file.num_rows(); - let primary_key_range = file.primary_key_range(); - Self { - files: smallvec![file], - size, - num_rows, - min_timestamp, - max_timestamp, - primary_key_range, - } - } - - pub(crate) fn num_rows(&self) -> usize { - self.num_rows - } - - pub(crate) fn add_file(&mut self, file: FileHandle) { - self.size += file.size() as usize; - self.num_rows += file.num_rows(); - let (min_timestamp, max_timestamp) = file.time_range(); - self.min_timestamp = self.min_timestamp.min(min_timestamp); - self.max_timestamp = self.max_timestamp.max(max_timestamp); - self.primary_key_range = - merge_primary_key_ranges(self.primary_key_range.take(), file.primary_key_range()); - self.files.push(file); - } - - pub(crate) fn num_files(&self) -> usize { - self.files.len() - } - - #[cfg(test)] - pub(crate) fn files(&self) -> &[FileHandle] { - &self.files[..] - } - - pub(crate) fn file_ids(&self) -> SmallVec<[RegionFileId; 2]> { - SmallVec::from_iter(self.files.iter().map(|f| f.file_id())) - } - - pub(crate) fn into_files(self) -> impl Iterator { - self.files.into_iter() - } -} - -impl Ranged for FileGroup { +impl Ranged for FileHandle { type BoundType = Timestamp; fn range(&self) -> (Self::BoundType, Self::BoundType) { - (self.min_timestamp, self.max_timestamp) + self.time_range() } fn overlap(&self, other: &Self) -> bool { @@ -226,7 +161,7 @@ impl Ranged for FileGroup { return false; } - match (&self.primary_key_range, &other.primary_key_range) { + match (&self.primary_key_range(), &other.primary_key_range()) { (Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs), _ => true, } @@ -239,16 +174,16 @@ impl Ranged for FileGroup { return false; } - match (&self.primary_key_range, &other.primary_key_range) { + match (&self.primary_key_range(), &other.primary_key_range()) { (Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs), _ => true, } } } -impl Item for FileGroup { +impl Item for FileHandle { fn size(&self) -> usize { - self.size + self.size() as usize } } @@ -553,124 +488,8 @@ where runs.into_iter().map(|x| x.run).collect() } -/// Finds a set of files with minimum penalty to merge that can reduce the total num of runs. -/// The penalty of merging is defined as the size of all overlapping files between two runs. -pub fn reduce_runs(mut runs: Vec>) -> Vec { - assert!(runs.len() > 1); - // sort runs by size - runs.sort_unstable_by_key(|a| a.size); - // limit max probe runs to 100 - let probe_end = runs.len().min(100); - let mut min_penalty = usize::MAX; - let mut files = vec![]; - let mut temp_files = vec![]; - for i in 0..probe_end { - for j in i + 1..probe_end { - let (a, b) = runs.split_at_mut(j); - find_overlapping_items(&mut a[i], &mut b[0], &mut temp_files); - let penalty = temp_files.iter().map(|e| e.size()).sum(); - if penalty < min_penalty { - min_penalty = penalty; - files.clear(); - files.extend_from_slice(&temp_files); - } - } - } - files -} - -/// Finds the optimal set of adjacent files to merge based on a scoring system. -/// -/// This function evaluates all possible contiguous subsets of files to find the best -/// candidates for merging, considering: -/// -/// 1. File reduction - prioritizes merging more files to reduce the total count -/// 2. Write amplification - minimizes the ratio of largest file to total size -/// 3. Size efficiency - prefers merges that utilize available space effectively -/// -/// When multiple merge candidates have the same score, older files (those with lower indices) -/// are preferred. -/// -/// # Arguments -/// * `input_files` - Slice of files to consider for merging -/// * `max_file_size` - Optional maximum size constraint for the merged file. -/// If None, uses 1.5 times the average file size. -/// -/// # Returns -/// A vector containing the best set of adjacent files to merge. -/// Returns an empty vector if input is empty or contains only one file. -pub fn merge_seq_files(input_files: &[T], max_file_size: Option) -> Vec { - if input_files.is_empty() || input_files.len() == 1 { - return vec![]; - } - - // Limit the number of files to process to 100 to control time complexity - let files_to_process = if input_files.len() > 100 { - &input_files[0..100] - } else { - input_files - }; - - // Calculate target size based on max_file_size or average file size - let target_size = match max_file_size { - Some(size) => size as usize, - None => { - // Calculate 1.5*average_file_size if max_file_size is not provided and clamp to 2GB - let total_size: usize = files_to_process.iter().map(|f| f.size()).sum(); - ((((total_size as f64) / (files_to_process.len() as f64)) * 1.5) as usize) - .min(DEFAULT_MAX_OUTPUT_SIZE as usize) - } - }; - - // Find the best group of adjacent files to merge - let mut best_group = Vec::new(); - let mut best_score = f64::NEG_INFINITY; - - // Try different starting positions - iterate from end to start to prefer older files - for start_idx in (0..files_to_process.len()).rev() { - // Try different ending positions - also iterate from end to start - for end_idx in (start_idx + 1..files_to_process.len()).rev() { - let group = &files_to_process[start_idx..=end_idx]; - let total_size: usize = group.iter().map(|f| f.size()).sum(); - - // Skip if total size exceeds target size - if total_size > target_size { - continue; // Use continue instead of break to check smaller ranges - } - - // Calculate amplification factor (largest file size / total size) - let largest_file_size = group.iter().map(|f| f.size()).max().unwrap_or(0); - let amplification_factor = largest_file_size as f64 / total_size as f64; - - // Calculate file reduction (number of files that will be reduced) - let file_reduction = group.len() - 1; - - // Calculate score based on multiple factors: - // 1. File reduction (higher is better) - // 2. Amplification factor (lower is better) - // 3. Size efficiency (how close to target size) - let file_reduction_score = file_reduction as f64 / files_to_process.len() as f64; - let amp_factor_score = (1.0 - amplification_factor) * 1.5; // Lower amplification is better - let size_efficiency = (total_size as f64 / target_size as f64).min(1.0); // Reward using available space - - let score = file_reduction_score + amp_factor_score + size_efficiency; - - // Check if this group is better than our current best - // Use >= instead of > to prefer older files (which we encounter first due to reverse iteration) - if score >= best_score { - best_score = score; - best_group = group.to_vec(); - } - } - } - - best_group -} - #[cfg(test)] mod tests { - use std::collections::HashSet; - use bytes::Bytes; use store_api::storage::FileId; @@ -709,17 +528,6 @@ mod tests { .collect() } - fn build_items_with_size(items: &[(i64, i64, usize)]) -> Vec { - items - .iter() - .map(|(start, end, size)| MockFile { - start: *start, - end: *end, - size: *size, - }) - .collect() - } - fn pk_range(min: &'static [u8], max: &'static [u8]) -> Option<(Bytes, Bytes)> { Some((Bytes::from_static(min), Bytes::from_static(max))) } @@ -844,99 +652,6 @@ mod tests { } } - fn check_reduce_runs( - files: &[(i64, i64)], - expected_runs: &[Vec<(i64, i64)>], - expected: &[(i64, i64)], - ) { - let runs = check_sorted_runs(files, expected_runs); - if runs.len() <= 1 { - assert!(expected.is_empty()); - return; - } - let files_to_merge = reduce_runs(runs); - let file_to_merge_timestamps = files_to_merge - .into_iter() - .map(|f| (f.start, f.end)) - .collect::>(); - - let expected = expected.iter().cloned().collect::>(); - assert_eq!(&expected, &file_to_merge_timestamps); - } - - #[test] - fn test_reduce_runs() { - // [1..3] [5..6] - // [2..4] - check_reduce_runs( - &[(1, 3), (2, 4), (5, 6)], - &[vec![(1, 3), (5, 6)], vec![(2, 4)]], - &[(1, 3), (2, 4)], - ); - - // [1..2][3..5] - // [4..6] - check_reduce_runs( - &[(1, 2), (3, 5), (4, 6)], - &[vec![(1, 2), (3, 5)], vec![(4, 6)]], - &[(3, 5), (4, 6)], - ); - - // [1..2][3..4] [7..8] - // [4..6] - check_reduce_runs( - &[(1, 2), (3, 4), (4, 6), (7, 8)], - &[vec![(1, 2), (3, 4), (4, 6), (7, 8)]], - &[], - ); - - // [1..2][3........6][7..8][8..9] - // [3..4][5..6] - check_reduce_runs( - &[(1, 2), (3, 4), (5, 6), (3, 6), (7, 8), (8, 9)], - &[vec![(1, 2), (3, 6), (7, 8), (8, 9)], vec![(3, 4), (5, 6)]], - &[(5, 6), (3, 4), (3, 6)], // already satisfied - ); - - // [1..2][3........6][7..8][8..9] - // [3..4][5..6] - check_reduce_runs( - &[(1, 2), (3, 4), (5, 6), (3, 6), (7, 8), (8, 9)], - &[vec![(1, 2), (3, 6), (7, 8), (8, 9)], vec![(3, 4), (5, 6)]], - &[(3, 4), (3, 6), (5, 6)], - ); - - // [10..20] [30..40] [50........80][80...100][100..110] - // [50..60] [80..90] - // - check_reduce_runs( - &[ - (10, 20), - (30, 40), - (50, 60), - (50, 80), - (80, 90), - (80, 100), - (100, 110), - ], - &[ - vec![(10, 20), (30, 40), (50, 80), (80, 100), (100, 110)], - vec![(50, 60), (80, 90)], - ], - &[(50, 80), (80, 100), (50, 60), (80, 90)], - ); - - // [0..10] - // [0...11] - // [0....12] - // [0.....13] - check_reduce_runs( - &[(0, 10), (0, 11), (0, 12), (0, 13)], - &[vec![(0, 13)], vec![(0, 12)], vec![(0, 11)], vec![(0, 10)]], - &[(0, 10), (0, 11)], - ); - } - #[test] fn test_find_overlapping_items() { let mut result = Vec::new(); @@ -1038,27 +753,25 @@ mod tests { } #[test] - fn test_file_group_overlap_time_overlap_pk_disjoint() { - let lhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 0, - 100, - 0, - 1, - 10, - pk_range(b"a", b"f"), - )); - let rhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 50, - 150, - 0, - 2, - 10, - pk_range(b"x", b"z"), - )); + fn test_file_overlap_time_overlap_pk_disjoint() { + let lhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 0, + 100, + 0, + 1, + 10, + pk_range(b"a", b"f"), + ); + let rhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 50, + 150, + 0, + 2, + 10, + pk_range(b"x", b"z"), + ); assert!(!lhs.overlap(&rhs)); } @@ -1066,7 +779,7 @@ mod tests { #[test] fn test_find_sorted_runs_collapses_pk_disjoint_files_into_one_run() { let mut files = vec![ - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 0, 100, @@ -1074,8 +787,8 @@ mod tests { 1, 10, pk_range(b"a", b"f"), - )), - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + ), + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, 150, @@ -1083,7 +796,7 @@ mod tests { 2, 10, pk_range(b"x", b"z"), - )), + ), ]; let runs = find_sorted_runs(&mut files); @@ -1095,7 +808,7 @@ mod tests { #[test] fn test_find_sorted_runs_handles_2d_transitivity_break() { let mut files = vec![ - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 0, 100, @@ -1103,8 +816,8 @@ mod tests { 1, 10, pk_range(b"a", b"f"), - )), - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + ), + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, 150, @@ -1112,8 +825,8 @@ mod tests { 2, 10, pk_range(b"x", b"z"), - )), - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + ), + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, 150, @@ -1121,7 +834,7 @@ mod tests { 3, 10, pk_range(b"a", b"f"), - )), + ), ]; let runs = find_sorted_runs(&mut files); @@ -1133,7 +846,7 @@ mod tests { #[test] fn test_find_overlapping_items_skips_pk_disjoint_pairs() { - let mut left = SortedRun::from(vec![FileGroup::new_with_file( + let mut left = SortedRun::from(vec![ new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 0, @@ -1143,8 +856,8 @@ mod tests { 10, pk_range(b"a", b"f"), ), - )]); - let mut right = SortedRun::from(vec![FileGroup::new_with_file( + ]); + let mut right = SortedRun::from(vec![ new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, @@ -1154,7 +867,7 @@ mod tests { 10, pk_range(b"x", b"z"), ), - )]); + ]); let mut result = Vec::new(); find_overlapping_items(&mut left, &mut right, &mut result); @@ -1163,102 +876,26 @@ mod tests { } #[test] - fn test_file_group_touching_time_boundary_with_same_pk_is_not_overlap() { - let lhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 0, - 100, - 0, - 1, - 10, - pk_range(b"a", b"f"), - )); - let rhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 100, - 150, - 0, - 2, - 10, - pk_range(b"a", b"f"), - )); + fn test_file_touching_time_boundary_with_same_pk_is_not_overlap() { + let lhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 0, + 100, + 0, + 1, + 10, + pk_range(b"a", b"f"), + ); + let rhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 100, + 150, + 0, + 2, + 10, + pk_range(b"a", b"f"), + ); assert!(!lhs.overlap(&rhs)); } - - #[test] - fn test_merge_seq_files() { - // Test empty input - let files = Vec::::new(); - assert_eq!(merge_seq_files(&files, None), Vec::::new()); - - // Test single file input (should return empty vec as no merge needed) - let files = build_items(&[(1, 5)]); - assert_eq!(merge_seq_files(&files, None), Vec::::new()); - - // Test the example case: [10, 1, 1, 1] - should merge the last three files - let files = build_items_with_size(&[(1, 2, 10), (3, 4, 1), (5, 6, 1), (7, 8, 1)]); - let result = merge_seq_files(&files, None); - assert_eq!(result.len(), 3); - assert_eq!(result[0].size, 1); - assert_eq!(result[1].size, 1); - assert_eq!(result[2].size, 1); - - // Test with files of equal size - should merge as many as possible - let files = build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 5), (7, 8, 5)]); - let result = merge_seq_files(&files, Some(20)); - assert_eq!(result.len(), 4); // Should merge all 4 files as total size is 20 - - // Test with max_file_size constraint - let files = build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 5), (7, 8, 5)]); - let result = merge_seq_files(&files, Some(10)); - assert_eq!(result.len(), 2); // Should merge only 2 files as max size is 10 - - // Test with uneven file sizes - should prioritize reducing file count - let files = build_items_with_size(&[(1, 2, 2), (3, 4, 3), (5, 6, 4), (7, 8, 10)]); - let result = merge_seq_files(&files, Some(10)); - assert_eq!(result.len(), 3); // Should merge the first 3 files (total size 9) - - // Test amplification factor prioritization - // Two possible merges: [5, 5] (amp factor 0.5) vs [10, 1, 1] (amp factor 0.83) - let files = - build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 10), (7, 8, 1), (9, 10, 1)]); - let result = merge_seq_files(&files, Some(12)); - assert_eq!(result.len(), 2); - assert_eq!(result[0].size, 5); - assert_eq!(result[1].size, 5); - - // Test with large file preventing merges - let files = build_items_with_size(&[(1, 2, 100), (3, 4, 1), (5, 6, 1), (7, 8, 1)]); - let result = merge_seq_files(&files, Some(10)); - assert_eq!(result.len(), 3); // Should merge the last 3 small files - assert_eq!(result[0].size, 1); - assert_eq!(result[1].size, 1); - assert_eq!(result[2].size, 1); - - let files = build_items_with_size(&[(1, 2, 100), (3, 4, 20), (5, 6, 20), (7, 8, 20)]); - let result = merge_seq_files(&files, Some(200)); - assert_eq!(result.len(), 4); - - let files = build_items_with_size(&[(1, 2, 160), (3, 4, 20), (5, 6, 20), (7, 8, 20)]); - let result = merge_seq_files(&files, None); - assert_eq!(result.len(), 3); - assert_eq!(result[0].size, 20); - assert_eq!(result[1].size, 20); - assert_eq!(result[2].size, 20); - - let files = build_items_with_size(&[(1, 2, 100), (3, 4, 1)]); - let result = merge_seq_files(&files, Some(200)); - assert_eq!(result.len(), 2); - assert_eq!(result[0].size, 100); - assert_eq!(result[1].size, 1); - - let files = build_items_with_size(&[(1, 2, 20), (3, 4, 20), (5, 6, 20), (7, 8, 20)]); - let result = merge_seq_files(&files, Some(40)); - assert_eq!(result.len(), 2); - assert_eq!(result[0].start, 1); - assert_eq!(result[1].start, 3); - } } diff --git a/src/mito2/src/compaction/scheduler.rs b/src/mito2/src/compaction/scheduler.rs index 69a7acc5a5a..6aebb25eebc 100644 --- a/src/mito2/src/compaction/scheduler.rs +++ b/src/mito2/src/compaction/scheduler.rs @@ -264,9 +264,13 @@ impl CompactionScheduler { /// /// # Effects /// - /// Notifies waiters, schedules a pending manual or explicit automatic - /// follow-up, or removes the region status. The returned transition reports - /// a dispatched automatic follow-up or DDLs that are now safe to execute. + /// Notifies waiters, schedules a pending manual or automatic follow-up, or + /// removes the region status. A follow-up is scheduled whenever the cycle + /// latched an automatic trigger, or when `made_progress` is true (the + /// execution removed more files than it added): successful compaction keeps + /// draining the region until the picker returns no plan. The returned + /// transition reports a dispatched automatic follow-up or DDLs that are now + /// safe to execute. /// /// # Constraints /// @@ -279,12 +283,18 @@ impl CompactionScheduler { execution: &CompactionExecution, manifest_ctx: &ManifestContextRef, schema_metadata_manager: SchemaMetadataManagerRef, + made_progress: bool, ) -> CompactionTransition { if !self.is_current_execution(region_id, execution) { return CompactionTransition::NoAction; } - self.on_compaction_finished(region_id, manifest_ctx, schema_metadata_manager) - .await + self.on_compaction_finished( + region_id, + manifest_ctx, + schema_metadata_manager, + made_progress, + ) + .await } /// Completes a cooperatively canceled execution. @@ -571,6 +581,7 @@ impl CompactionScheduler { region_id: RegionId, manifest_ctx: &ManifestContextRef, schema_metadata_manager: SchemaMetadataManagerRef, + made_progress: bool, ) -> CompactionTransition { if !self.region_status.contains_key(®ion_id) { return CompactionTransition::NoAction; @@ -605,7 +616,11 @@ impl CompactionScheduler { return CompactionTransition::DdlReady(pending_ddl_requests); } - if status.active.reset_automatic_followup() + // Keep draining when the cycle latched an automatic trigger, or when the + // execution reduced the file count. A no-progress rewrite stops here so a + // split-heavy output cannot loop forever; the next flush trigger resumes. + let should_continue = status.active.reset_automatic_followup() || made_progress; + if should_continue && self.schedule_automatic_followup(region_id, manifest_ctx, schema_metadata_manager) { return CompactionTransition::AutomaticFollowupScheduled; diff --git a/src/mito2/src/compaction/scheduler_test.rs b/src/mito2/src/compaction/scheduler_test.rs index d3548ee8a5f..9545d0eea2a 100644 --- a/src/mito2/src/compaction/scheduler_test.rs +++ b/src/mito2/src/compaction/scheduler_test.rs @@ -826,6 +826,7 @@ async fn test_stale_plan_execution_does_not_affect_replacement_status() { &stale_execution, &manifest_ctx, schema_metadata_manager, + true, ) .await; assert!(pending_ddls.is_empty()); @@ -907,7 +908,7 @@ async fn test_schedule_compaction_skips_task_exceeding_memory_limit() { } #[tokio::test] -async fn test_execution_finished_without_followup_removes_status() { +async fn test_execution_finished_drains_until_no_plan() { common_telemetry::init_default_ut_logging(); let job_scheduler = Arc::new(VecScheduler::default()); let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); @@ -975,23 +976,59 @@ async fn test_execution_finished_without_followup_removes_status() { &file_metas, purger.clone(), ); - // A completed execution without an explicit follow-up removes its status. - scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + // A completed execution that reduced the file count chains another pick even + // without an explicit trigger, because the layout is still compactable. + let transition = scheduler + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; + assert_matches!(transition, CompactionTransition::AutomaticFollowupScheduled); + assert!(scheduler.region_status.contains_key(®ion_id)); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_eq!(2, job_scheduler.num_jobs()); + + // Replace the layout with a single file: the drain must stop once the picker + // finds nothing to do. + let data = version_control.current(); + let file_metas: Vec<_> = data.version.ssts.levels()[0] + .files + .values() + .map(|file| file.meta_ref().clone()) + .collect(); + apply_edit(&version_control, &[(0, end)], &file_metas, purger.clone()); + let transition = scheduler + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) + .await; + assert_matches!(transition, CompactionTransition::AutomaticFollowupScheduled); + let finished = recv_compaction_pick_finished(&mut rx).await; + let transition = scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_matches!(transition, CompactionTransition::NoAction); assert!(scheduler.region_status.is_empty()); - assert_eq!(1, job_scheduler.num_jobs()); + assert_eq!(2, job_scheduler.num_jobs()); } #[tokio::test] -async fn test_time_range_compaction_when_compaction_in_progress() { +async fn test_execution_finished_without_progress_removes_status() { common_telemetry::init_default_ut_logging(); let job_scheduler = Arc::new(VecScheduler::default()); let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); - let purger = builder.file_purger(); let region_id = builder.region_id(); let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); @@ -1017,6 +1054,70 @@ async fn test_time_range_compaction_when_compaction_in_progress() { .push_l0_file(90, end) .build(), ); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let scheduled = scheduler + .schedule_automatic_compaction( + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .unwrap(); + assert!(scheduled); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_eq!(1, job_scheduler.num_jobs()); + + // The execution rewrote files without reducing the file count (e.g. its output + // was split into more files than its input). Chaining would loop without making + // progress, so the lifecycle ends here; the next flush trigger resumes compaction. + let transition = scheduler + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + false, + ) + .await; + assert_matches!(transition, CompactionTransition::NoAction); + assert!(scheduler.region_status.is_empty()); + assert_eq!(1, job_scheduler.num_jobs()); +} + +#[tokio::test] +async fn test_time_range_compaction_when_compaction_in_progress() { + common_telemetry::init_default_ut_logging(); + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let mut builder = VersionControlBuilder::new(); + let purger = builder.file_purger(); + let region_id = builder.region_id(); + + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + schema_metadata_manager + .register_region_table_info( + builder.region_id().table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + kv_backend, + ) + .await; + + // 40 files to compact. The first task picks 32, leaving 8 for the pending request. + let end = 1000 * 1000; + for offset in 0..40 { + builder.push_l0_file(offset * 10, end); + } + let version_control = Arc::new(builder.build()); let manifest_ctx = env .mock_manifest_context(version_control.current().version.metadata.clone()) .await; @@ -1027,13 +1128,9 @@ async fn test_time_range_compaction_when_compaction_in_progress() { .map(|file| file.meta_ref().clone()) .collect(); - // 5 files for next compaction and removes old files. - apply_edit( - &version_control, - &[(0, end), (20, end), (40, end), (60, end), (80, end)], - &file_metas, - purger.clone(), - ); + // Replaces the files before scheduling the first compaction. + let next_files = (0..40).map(|offset| (offset * 20, end)).collect::>(); + apply_edit(&version_control, &next_files, &file_metas, purger.clone()); scheduler .schedule_automatic_compaction( @@ -1094,7 +1191,12 @@ async fn test_time_range_compaction_when_compaction_in_progress() { // On compaction finished and schedule next compaction. scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; assert_eq!(1, scheduler.region_status.len()); assert!( @@ -1285,10 +1387,21 @@ async fn test_ranged_compaction_stops_without_followup() { assert_eq!(1, job_scheduler.num_jobs()); scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; - assert!(!scheduler.region_status.contains_key(®ion_id)); - assert_eq!(1, job_scheduler.num_jobs()); + // The execution reduced the file count and the remaining window is still + // compactable, so the scheduler drains it with an unrestricted follow-up. + assert!(scheduler.region_status.contains_key(®ion_id)); + let followup = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(followup, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_eq!(2, job_scheduler.num_jobs()); } #[tokio::test] @@ -1369,7 +1482,12 @@ async fn test_automatic_trigger_during_execution_clears_continuation_scope() { // Finishing the manual execution must chain an unrestricted follow-up instead of // continuing with the manual request's time range. scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; let continuation = recv_compaction_pick_finished(&mut rx).await; @@ -1936,7 +2054,7 @@ async fn test_on_compaction_finished_returns_pending_ddl_requests() { }); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true) .await; assert_eq!(pending_ddls.len(), 1); @@ -1985,7 +2103,12 @@ async fn test_planning_terminal_prioritizes_pending_ddl_over_automatic_followup( assert!(result.is_ok()); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; assert!(pending_ddls.is_empty()); @@ -2034,7 +2157,7 @@ async fn test_on_compaction_finished_dispatches_pending_ddl_before_chained_regul }); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true) .await; assert_eq!(pending_ddls.len(), 1); @@ -2056,7 +2179,7 @@ async fn test_on_compaction_finished_returns_empty_when_region_absent() { let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true) .await; assert!(pending_ddls.is_empty()); @@ -2111,7 +2234,12 @@ async fn test_on_compaction_finished_manual_schedule_error_cleans_status() { }); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; assert!(pending_ddls.is_empty()); diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index 342426bbefe..cdd613b3424 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -15,8 +15,7 @@ use std::collections::hash_map::Entry; use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt::Debug; -use std::num::NonZeroU64; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use common_base::readable_size::ReadableSize; use common_telemetry::{debug, info}; @@ -32,8 +31,8 @@ use crate::compaction::buckets::infer_time_bucket; use crate::compaction::compactor::CompactionRegion; use crate::compaction::picker::{Picker, PickerOutput, get_expired_ssts}; use crate::compaction::run::{ - FileGroup, Item, Ranged, find_sorted_runs, find_sorted_runs_by_time_range, - merge_primary_key_ranges, merge_seq_files, primary_key_ranges_overlap, reduce_runs, + Ranged, SortedRun, find_sorted_runs, find_sorted_runs_by_time_range, merge_primary_key_ranges, + primary_key_ranges_overlap, }; use crate::error::{JoinSnafu, Result}; use crate::sst::file::{FileHandle, Level, overlaps}; @@ -41,8 +40,24 @@ use crate::sst::version::LevelMeta; const LEVEL_COMPACTED: Level = 1; -/// Default value for max compaction input file num. -const DEFAULT_MAX_INPUT_FILE_NUM: usize = 32; +/// Default maximum number of input SST files in one compaction input. +const DEFAULT_MAX_INPUT_FILES: usize = 32; + +const MAX_INPUT_FILES_ENV: &str = "GREPTIME_TWCS_MAX_INPUT_FILES"; + +/// Maximum number of input SST files in one compaction input. +/// Configurable via [`MAX_INPUT_FILES_ENV`]. +static MAX_INPUT_FILES: LazyLock = LazyLock::new(|| { + let env_value = std::env::var(MAX_INPUT_FILES_ENV).ok(); + parse_max_input_files(env_value.as_deref()) +}); + +fn parse_max_input_files(env_value: Option<&str>) -> usize { + env_value + .and_then(|env_value| env_value.parse().ok()) + .filter(|max_input_files| *max_input_files >= 2) + .unwrap_or(DEFAULT_MAX_INPUT_FILES) +} /// `TwcsPicker` picks files of which the max timestamp are in the same time window as compaction /// candidates. @@ -52,7 +67,8 @@ pub struct TwcsPicker { pub trigger_file_num: usize, /// Compaction time window in seconds. pub time_window_seconds: Option, - /// Max allowed compaction output file size. + /// Max allowed compaction output file size. The picker also uses it to predict + /// output splits when scoring candidates. pub max_output_file_size: Option, /// Whether the target region is in append mode. pub append_mode: bool, @@ -113,7 +129,7 @@ impl TwcsPicker { output.push(CompactionOutput { output_level: LEVEL_COMPACTED, // always compact to l1 - inputs: inputs.into_iter().flat_map(|fg| fg.into_files()).collect(), + inputs, filter_deleted, output_time_range: None, // we do not enforce output time range in twcs compactions. }); @@ -138,7 +154,11 @@ impl TwcsPicker { active_window: Option, files: &Window, windows: &BTreeMap, - ) -> (Vec, bool) { + ) -> (Vec, bool) { + if files.files.len() < self.trigger_file_num { + return (vec![], false); + } + let window = &files.time_window; let mut files_to_merge: Vec<_> = files.files().cloned().collect(); @@ -148,7 +168,7 @@ impl TwcsPicker { { let (kept_files, ignored_files) = files_to_merge .into_iter() - .partition(|fg| fg.size() <= max_size as usize); + .partition(|file| file.size() <= max_size); files_to_merge = kept_files; if !ignored_files.is_empty() { info!( @@ -167,59 +187,13 @@ impl TwcsPicker { find_sorted_runs_by_time_range(&mut files_to_merge) }; let found_runs = sorted_runs.len(); - // We only remove deletion markers if we found less than 2 runs and not in append mode. - // because after compaction there will be no overlapping files. - let mut filter_deleted = - found_runs <= 2 && !self.append_mode && !window_has_overlap(files, windows); - if found_runs == 0 { - return (vec![], filter_deleted); - } - - let mut inputs = if found_runs > 1 { - reduce_runs(sorted_runs) - } else { - let run = sorted_runs.last().unwrap(); - if run.items().len() < self.trigger_file_num { - return (vec![], filter_deleted); - } - // no overlapping files, try merge small files - merge_seq_files(run.items(), self.max_output_file_size) - }; - - // Limits the number of input files. - let total_input_files: usize = inputs.iter().map(|fg| fg.num_files()).sum(); - if total_input_files > DEFAULT_MAX_INPUT_FILE_NUM { - // Sorts file groups by size first. - inputs.sort_unstable_by_key(|fg| fg.size()); - let mut num_picked_files = 0; - inputs = inputs - .into_iter() - .take_while(|fg| { - let current_group_file_num = fg.num_files(); - if current_group_file_num + num_picked_files <= DEFAULT_MAX_INPUT_FILE_NUM { - num_picked_files += current_group_file_num; - true - } else { - false - } - }) - .collect::>(); - info!( - "Compaction for region {} enforces max input file num limit: {}, current total: {}, input: {:?}", - region_id, DEFAULT_MAX_INPUT_FILE_NUM, total_input_files, inputs - ); - } - - // The inputs are only a subset of the window: `reduce_runs` and `merge_seq_files` both - // narrow the selection down and the file num limit above narrows it further. A deletion - // marker may only be dropped when every file that can hold a row it masks is compacted - // along with it, so re-check the final selection against what stays in the window. - if filter_deleted && overlaps_files_left_behind(&inputs, &files_to_merge) { - filter_deleted = false; - } + let inputs = pick_count_first(sorted_runs, self.max_output_file_size); + let filter_deleted = !self.append_mode + && !window_has_overlap(files, windows) + && !selected_overlaps_unselected(&inputs, files); if inputs.len() > 1 { - // If we have more than one group to compact. + // If we have more than one file to compact. log_pick_result( region_id, *window, @@ -235,18 +209,225 @@ impl TwcsPicker { } } -/// Returns whether any file group of `window_files` that is not part of `inputs` overlaps -/// `inputs`. Such a group keeps rows a deletion marker in `inputs` masks, so the compaction -/// must not filter deleted rows out. +#[derive(Debug)] +struct OrderedFile<'a> { + file: &'a FileHandle, + run_id: usize, + position_in_run: usize, +} + +/// Metrics of a candidate compaction input, accumulated incrementally as the +/// candidate interval expands one file at a time. +#[derive(Debug, Default)] +struct Candidate { + /// Number of files in the interval. + num_files: usize, + /// Total input bytes of the interval. + total_size: usize, + /// Size of the largest single file in the interval. + largest_file_size: usize, + /// Files in the interval that overlap another interval file from a different + /// sorted run. Resolving these overlaps is what reduces sorted runs. + overlap_participants: usize, +} + +impl Candidate { + /// Absorbs `file` into the candidate. `preceding` holds the interval files added + /// before it and `participations` tracks which of them already participate in an + /// intra-interval overlap; both exist only for overlap accounting. + fn absorb( + &mut self, + file: &OrderedFile, + preceding: &[OrderedFile], + participations: &mut Vec, + ) { + self.num_files += 1; + let file_size = file.file.size() as usize; + self.total_size += file_size; + self.largest_file_size = self.largest_file_size.max(file_size); + + let mut participates = false; + for (offset, other) in preceding.iter().enumerate() { + if file.run_id != other.run_id && file.file.overlap_inclusive(other.file) { + if !participations[offset] { + participations[offset] = true; + self.overlap_participants += 1; + } + participates = true; + } + } + participations.push(participates); + if participates { + self.overlap_participants += 1; + } + } + + /// Predicted number of output files after compacting this candidate. Compaction + /// output holds at most the input bytes, so a split threshold of + /// `max_output_file_size` yields at most `ceil(total_size / threshold)` files. + /// Without a threshold the output is a single file. + fn predicted_output_files(&self, max_output_file_size: Option) -> usize { + match max_output_file_size { + Some(max) if max > 0 => self.total_size.div_ceil(max as usize).max(1), + _ => 1, + } + } + + /// Net file count reduction compacting this candidate would achieve, i.e. how + /// many fewer physical SSTs the window holds afterwards. Zero when the output + /// would be split back into at least as many files as the input. + fn file_reduction(&self, max_output_file_size: Option) -> usize { + self.num_files + .saturating_sub(self.predicted_output_files(max_output_file_size)) + } + + /// A large historical file is only rewritten once its peers contribute at least + /// the same amount of data, so every rewrite moves it into a larger size tier. + fn is_balanced(&self) -> bool { + self.largest_file_size <= self.total_size - self.largest_file_size + } + + /// A candidate is worth compacting only if it makes progress on at least one + /// axis: it reduces the physical file count, or it resolves at least one + /// overlap (merging sorted runs and reducing read amplification). A pure + /// rewrite that achieves neither only burns I/O. + fn makes_progress(&self, max_output_file_size: Option) -> bool { + self.file_reduction(max_output_file_size) > 0 || self.overlap_participants > 0 + } + + fn score(&self, max_output_file_size: Option) -> CandidateScore { + CandidateScore { + file_reduction: self.file_reduction(max_output_file_size), + overlap_participants: self.overlap_participants, + total_size: self.total_size, + } + } +} + +/// Score of a candidate compaction input. Higher is better: first by the predicted +/// net file count reduction, then by the number of files participating in an +/// overlap, then by smaller input bytes. +#[derive(Debug)] +struct CandidateScore { + file_reduction: usize, + overlap_participants: usize, + total_size: usize, +} + +impl CandidateScore { + fn is_better_than(&self, other: &Self) -> bool { + self.file_reduction + .cmp(&other.file_reduction) + .then_with(|| self.overlap_participants.cmp(&other.overlap_participants)) + .then_with(|| other.total_size.cmp(&self.total_size)) + .is_gt() + } +} + +/// Picks a contiguous (in global time order) interval of files to compact. /// -/// Overlapping is inclusive on both boundaries here: two files that only share a boundary -/// timestamp can still hold the same row. -fn overlaps_files_left_behind(inputs: &[FileGroup], window_files: &[FileGroup]) -> bool { - let picked: HashSet<_> = inputs.iter().flat_map(|fg| fg.file_ids()).collect(); - window_files +/// The picker reorders all files from all sorted runs by `(start asc, end desc)` and +/// enumerates every interval of at most [`MAX_INPUT_FILES`] files as a [`Candidate`]. +/// An interval is eligible when it +/// +/// - holds at least 2 files, +/// - is balanced: no single file dominates it (`largest <= sum of the others`), +/// - makes progress on at least one axis: it reduces the physical file count given +/// the output split threshold `max_output_file_size`, or it resolves at least one +/// overlap between sorted runs. +/// +/// The best interval wins by [`CandidateScore`]; ties keep the earliest interval. +fn pick_count_first( + sorted_runs: Vec>, + max_output_file_size: Option, +) -> Vec { + let files = ordered_files(&sorted_runs); + + let mut best = None; + for left in 0..files.len() { + let mut candidate = Candidate::default(); + let right_bound = left.saturating_add(*MAX_INPUT_FILES).min(files.len()); + let mut participations: Vec = Vec::with_capacity(right_bound - left); + for right in left..right_bound { + candidate.absorb(&files[right], &files[left..right], &mut participations); + if candidate.num_files < 2 + || !candidate.is_balanced() + || !candidate.makes_progress(max_output_file_size) + { + continue; + } + + let score = candidate.score(max_output_file_size); + if best + .as_ref() + .is_none_or(|(best_score, _)| score.is_better_than(best_score)) + { + best = Some((score, &files[left..=right])); + } + } + } + + let Some((_, best)) = best else { + return vec![]; + }; + best.iter().map(|file| file.file.clone()).collect() +} + +/// Flattens sorted runs into files ordered by `(start asc, end desc)`, breaking ties +/// by run and position to keep the order deterministic. +fn ordered_files(sorted_runs: &[SortedRun]) -> Vec> { + let mut files = sorted_runs .iter() - .filter(|fg| !fg.file_ids().iter().any(|id| picked.contains(id))) - .any(|fg| inputs.iter().any(|input| input.overlap_inclusive(fg))) + .enumerate() + .flat_map(|(run_id, run)| { + run.items() + .iter() + .enumerate() + .map(move |(position_in_run, file)| OrderedFile { + file, + run_id, + position_in_run, + }) + }) + .collect::>(); + files.sort_unstable_by(|lhs, rhs| { + let (lhs_start, lhs_end) = lhs.file.range(); + let (rhs_start, rhs_end) = rhs.file.range(); + lhs_start + .cmp(&rhs_start) + .then_with(|| rhs_end.cmp(&lhs_end)) + .then_with(|| lhs.run_id.cmp(&rhs.run_id)) + .then_with(|| lhs.position_in_run.cmp(&rhs.position_in_run)) + }); + files +} + +fn selected_overlaps_unselected(selected: &[FileHandle], window: &Window) -> bool { + // The overall time span of the selection: a file outside it cannot overlap any + // selected file (ranges are inclusive), so it needs no precise overlap check. + let Some((span_start, span_end)) = selected + .iter() + .map(Ranged::range) + .reduce(|(start_a, end_a), (start_b, end_b)| (start_a.min(start_b), end_a.max(end_b))) + else { + return false; + }; + let selected_file_ids = selected + .iter() + .map(FileHandle::file_id) + .collect::>(); + window + .files() + .filter(|file| { + let (start, end) = file.range(); + start <= span_end && span_start <= end + }) + .filter(|file| !selected_file_ids.contains(&file.file_id())) + .any(|unselected| { + selected + .iter() + .any(|selected| selected.overlap_inclusive(unselected)) + }) } #[allow(clippy::too_many_arguments)] @@ -258,7 +439,7 @@ fn log_pick_result( file_num: usize, max_output_file_size: Option, filter_deleted: bool, - inputs: &[FileGroup], + inputs: &[FileHandle], ) { let input_file_str: Vec = inputs .iter() @@ -268,11 +449,11 @@ fn log_pick_result( let end = range.1.to_iso8601_string(); let num_rows = f.num_rows(); format!( - "FileGroup{{id: {:?}, range: ({}, {}), size: {}, num rows: {} }}", - f.file_ids(), + "File{{id: {:?}, range: ({}, {}), size: {}, num rows: {} }}", + f.file_id(), start, end, - ReadableSize(f.size() as u64), + ReadableSize(f.size()), num_rows ) }) @@ -369,9 +550,7 @@ impl Picker for TwcsPicker { struct Window { start: Timestamp, end: Timestamp, - // Mapping from file sequence to file groups. Files with the same sequence is considered - // created from the same compaction task. - files: HashMap, FileGroup>, + files: Vec, time_window: i64, primary_key_range: Option<(bytes::Bytes, bytes::Bytes)>, } @@ -381,11 +560,10 @@ impl Window { fn new_with_file(file: FileHandle) -> Self { let (start, end) = file.time_range(); let primary_key_range = file.primary_key_range(); - let files = HashMap::from([(file.meta_ref().sequence, FileGroup::new_with_file(file))]); Self { start, end, - files, + files: vec![file], time_window: 0, primary_key_range, } @@ -403,19 +581,11 @@ impl Window { self.end = self.end.max(end); self.primary_key_range = merge_primary_key_ranges(self.primary_key_range.take(), file.primary_key_range()); - - match self.files.entry(file.meta_ref().sequence) { - Entry::Occupied(mut o) => { - o.get_mut().add_file(file); - } - Entry::Vacant(v) => { - v.insert(FileGroup::new_with_file(file)); - } - } + self.files.push(file); } - fn files(&self) -> impl Iterator { - self.files.values() + fn files(&self) -> impl Iterator { + self.files.iter() } } @@ -547,6 +717,18 @@ mod tests { use crate::test_util::memtable_util::metadata_for_test; use crate::test_util::scheduler_util::SchedulerEnv; + #[test] + fn test_valid_max_input_files_env_overrides_default() { + assert_eq!(64, parse_max_input_files(Some("64"))); + } + + #[test] + fn test_invalid_max_input_files_env_falls_back_to_default() { + for env_value in [None, Some(""), Some("invalid"), Some("0"), Some("1")] { + assert_eq!(32, parse_max_input_files(env_value)); + } + } + async fn compaction_region_with_expired_sst() -> CompactionRegion { let env = SchedulerEnv::new().await; let metadata = metadata_for_test(); @@ -664,8 +846,7 @@ mod tests { 3, ); let fgs = &windows.get(&0).unwrap().files; - assert_eq!(1, fgs.len()); - assert_eq!(fgs.values().map(|f| f.files().len()).sum::(), 5); + assert_eq!(5, fgs.len()); let files = [FileId::random(); 3]; let windows = assign_to_windows( @@ -679,26 +860,41 @@ mod tests { ); assert_eq!( files[0], - windows.get(&0).unwrap().files().next().unwrap().files()[0] + windows + .get(&0) + .unwrap() + .files() + .next() + .unwrap() .file_id() .file_id() ); assert_eq!( files[1], - windows.get(&3).unwrap().files().next().unwrap().files()[0] + windows + .get(&3) + .unwrap() + .files() + .next() + .unwrap() .file_id() .file_id() ); assert_eq!( files[2], - windows.get(&12).unwrap().files().next().unwrap().files()[0] + windows + .get(&12) + .unwrap() + .files() + .next() + .unwrap() .file_id() .file_id() ); } #[test] - fn test_assign_file_groups_to_windows() { + fn test_assign_files_to_windows() { let files = [ FileId::random(), FileId::random(), @@ -716,25 +912,14 @@ mod tests { 3, ); assert_eq!(windows.len(), 1); - let fgs = &windows.get(&0).unwrap().files; - assert_eq!(2, fgs.len()); + let window_files = &windows.get(&0).unwrap().files; + assert_eq!(4, window_files.len()); assert_eq!( - fgs.get(&NonZeroU64::new(1)) - .unwrap() - .files() + window_files .iter() .map(|f| f.file_id().file_id()) .collect::>(), - [files[0], files[1]].into_iter().collect() - ); - assert_eq!( - fgs.get(&NonZeroU64::new(2)) - .unwrap() - .files() - .iter() - .map(|f| f.file_id().file_id()) - .collect::>(), - [files[2], files[3]].into_iter().collect() + files.into_iter().collect() ); } @@ -751,11 +936,10 @@ mod tests { files[2].set_compacting(true); let mut windows = assign_to_windows(files.iter(), 3); let window0 = windows.remove(&0).unwrap(); - assert_eq!(1, window0.files.len()); + assert_eq!(3, window0.files.len()); let candidates = window0 .files - .into_values() - .flat_map(|fg| fg.into_files()) + .iter() .map(|f| f.file_id().file_id()) .collect::>(); assert_eq!(candidates.len(), 3); @@ -801,12 +985,10 @@ mod tests { assert_eq!(*overlapping, actual_overlapping); let mut file_ranges = actual_window .files - .values() - .flat_map(|f| { - f.files().iter().map(|f| { - let (s, e) = f.time_range(); - (s.value(), e.value()) - }) + .iter() + .map(|f| { + let (s, e) = f.time_range(); + (s.value(), e.value()) }) .collect::>(); file_ranges.sort_unstable_by(|l, r| l.0.cmp(&r.0).then(l.1.cmp(&r.1))); @@ -988,7 +1170,7 @@ mod tests { let active_window = find_latest_window_in_seconds(self.input_files.iter(), self.window_size); let output = TwcsPicker { - trigger_file_num: 4, + trigger_file_num: 2, time_window_seconds: None, max_output_file_size: None, append_mode: false, @@ -1075,7 +1257,7 @@ mod tests { .to_vec(), expected_outputs: vec![ ExpectedOutput { - input_files: vec![2, 4], + input_files: vec![2, 3, 4], output_level: 1, }, ExpectedOutput { @@ -1088,8 +1270,9 @@ mod tests { .await; // Case 3: - // A compaction may split output into several files that have overlapping time ranges and same sequence, - // we should treat these files as one FileGroup. + // A compaction may split output into several files that have overlapping time ranges and + // the same sequence. They are ordinary compaction candidates now: the picker merges the + // overlapping ones in the window that reaches the trigger. let file_ids = (0..6).map(|_| FileId::random()).collect::>(); CompactionPickerTestCase { window_size: 3, @@ -1101,10 +1284,16 @@ mod tests { new_file_handle_with_sequence(file_ids[4], 11, 2990, 0, 3), ] .to_vec(), - expected_outputs: vec![ExpectedOutput { - input_files: vec![0, 1, 4], - output_level: 1, - }], + expected_outputs: vec![ + ExpectedOutput { + input_files: vec![2, 3], + output_level: 1, + }, + ExpectedOutput { + input_files: vec![0, 1, 4], + output_level: 1, + }, + ], } .check() .await; @@ -1160,30 +1349,24 @@ mod tests { let small_file_2 = new_file_handle_with_size_and_sequence(file_ids[2], 0, 999, 0, 3, 800); let large_file_2 = new_file_handle_with_size_and_sequence(file_ids[3], 0, 999, 0, 4, 2000); - // Create file groups (each file is in its own group due to different sequences) - let mut files_to_merge = vec![ - FileGroup::new_with_file(small_file_1), - FileGroup::new_with_file(large_file_1), - FileGroup::new_with_file(small_file_2), - FileGroup::new_with_file(large_file_2), - ]; + let mut files_to_merge = vec![small_file_1, large_file_1, small_file_2, large_file_2]; // Test filtering logic directly let original_count = files_to_merge.len(); // Apply append mode filtering - files_to_merge.retain(|fg| fg.size() <= max_output_file_size as usize); + files_to_merge.retain(|file| file.size() <= max_output_file_size); // Should have filtered out 2 large files, leaving 2 small files assert_eq!(files_to_merge.len(), 2); assert_eq!(original_count, 4); // Verify the remaining files are the small ones - for fg in &files_to_merge { + for file in &files_to_merge { assert!( - fg.size() <= max_output_file_size as usize, + file.size() <= max_output_file_size, "File size {} should be <= {}", - fg.size(), + file.size(), max_output_file_size ); } @@ -1191,7 +1374,7 @@ mod tests { #[tokio::test] async fn test_build_output_multiple_windows_with_zero_runs() { - let file_ids = (0..6).map(|_| FileId::random()).collect::>(); + let file_ids = (0..7).map(|_| FileId::random()).collect::>(); let files = [ // Window 0: Contains 3 files but not forming any runs (not enough files in sequence to reach trigger_file_num) @@ -1202,6 +1385,7 @@ mod tests { new_file_handle_with_sequence(file_ids[3], 3000, 3999, 0, 4), new_file_handle_with_sequence(file_ids[4], 3000, 3999, 0, 5), new_file_handle_with_sequence(file_ids[5], 3000, 3999, 0, 6), + new_file_handle_with_sequence(file_ids[6], 3000, 3999, 0, 7), ]; let windows = assign_to_windows(files.iter(), 3); @@ -1431,11 +1615,11 @@ mod tests { .unwrap(); assert_eq!(1, output.len()); - assert_eq!(output[0].inputs.len(), 2); + assert_eq!(output[0].inputs.len(), num_files); } #[tokio::test] - async fn test_limit_max_input_files() { + async fn test_window_trigger_can_exceed_input_limit() { common_telemetry::init_default_ut_logging(); let num_files = 50; @@ -1460,7 +1644,7 @@ mod tests { let windows = assign_to_windows(files.iter(), 3); let picker = TwcsPicker { - trigger_file_num: 4, + trigger_file_num: num_files, time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, @@ -1475,7 +1659,7 @@ mod tests { .unwrap(); assert_eq!(1, output.len()); - assert_eq!(output[0].inputs.len(), 32); + assert_eq!(output[0].inputs.len(), num_files.min(*MAX_INPUT_FILES)); } #[tokio::test] @@ -1671,5 +1855,394 @@ mod tests { ); } + #[tokio::test] + async fn test_count_first_prefers_more_files_over_smaller_overlap() { + let files = [ + new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 20, 30, 0, 2, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 40, 50, 0, 3, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 5, 15, 0, 4, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 2, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + assert_eq!(4, output[0].inputs.len()); + } + + #[tokio::test] + async fn test_count_first_trigger_counts_physical_ssts() { + let files = [ + new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 20, 30, 0, 2, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 3, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + assert_eq!(3, output[0].inputs.len()); + } + + #[tokio::test] + async fn test_count_first_does_not_compact_overlap_below_trigger() { + let files = [ + new_file_handle_with_size_and_sequence(FileId::random(), 0, 20, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 10, 30, 0, 2, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 3, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert!(output.is_empty()); + } + + #[tokio::test] + async fn test_filter_deleted_is_false_when_selected_files_overlap_unselected_file() { + let mut files = (0..32) + .map(|idx| { + new_file_handle_with_size_and_sequence( + FileId::random(), + idx * 10, + idx * 10 + 9, + 0, + idx as u64 + 1, + 10, + ) + }) + .collect::>(); + files.push(new_file_handle_with_size_and_sequence( + FileId::random(), + 0, + 320, + 0, + 33, + 10, + )); + let windows = assign_to_windows(files.iter(), 1000); + let picker = TwcsPicker { + trigger_file_num: 2, + time_window_seconds: Some(1000), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + assert_eq!(DEFAULT_MAX_INPUT_FILES, output[0].inputs.len()); + assert!(!output[0].filter_deleted); + } + + fn new_file(start: i64, end: i64, sequence: u64, file_size: u64) -> FileHandle { + new_file_handle_with_size_and_sequence(FileId::random(), start, end, 0, sequence, file_size) + } + + fn picked_ranges(files: &[FileHandle]) -> Vec<(i64, i64)> { + files + .iter() + .map(|file| { + let (start, end) = file.range(); + (start.value(), end.value()) + }) + .collect() + } + + #[test] + fn test_count_first_rejects_dominant_historical_file() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 400), + new_file(20, 29, 2, 100), + ])], + None, + ); + + assert!(picked.is_empty()); + } + + #[test] + fn test_count_first_accepts_balanced_historical_file() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 400), + new_file(20, 29, 2, 100), + new_file(40, 49, 3, 100), + new_file(60, 69, 4, 100), + new_file(80, 89, 5, 100), + ])], + None, + ); + + assert_eq!( + vec![(0, 9), (20, 29), (40, 49), (60, 69), (80, 89)], + picked_ranges(&picked) + ); + } + + #[test] + fn test_count_first_finds_smaller_balanced_interval_when_larger_one_is_unbalanced() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 1000), + new_file(20, 29, 2, 100), + new_file(40, 49, 3, 100), + new_file(60, 69, 4, 100), + new_file(80, 89, 5, 100), + ])], + None, + ); + + assert_eq!( + vec![(20, 29), (40, 49), (60, 69), (80, 89)], + picked_ranges(&picked) + ); + } + + #[test] + fn test_count_first_prefers_overlap_participants_when_file_counts_match() { + let first_run = (0..DEFAULT_MAX_INPUT_FILES) + .map(|idx| { + let start = idx as i64 * 20; + let end = if idx + 1 == DEFAULT_MAX_INPUT_FILES { + 700 + } else { + start + 9 + }; + new_file(start, end, idx as u64 + 1, 10) + }) + .collect::>(); + let overlapping = new_file(690, 710, 100, 10); + + let picked = pick_count_first( + vec![ + SortedRun::from(first_run), + SortedRun::from(vec![overlapping]), + ], + None, + ); + let ranges = picked_ranges(&picked); + + assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len()); + assert!(!ranges.contains(&(0, 9))); + assert!(ranges.contains(&(690, 710))); + } + + #[test] + fn test_count_first_prefers_smaller_bytes_when_file_counts_match() { + let files = (0..=DEFAULT_MAX_INPUT_FILES) + .map(|idx| { + let start = idx as i64 * 20; + let size = if idx == 0 { + 100 + } else if idx == DEFAULT_MAX_INPUT_FILES { + 1 + } else { + 10 + }; + new_file(start, start + 9, idx as u64 + 1, size) + }) + .collect::>(); + + let picked = pick_count_first(vec![SortedRun::from(files)], None); + let ranges = picked_ranges(&picked); + + assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len()); + assert_eq!(Some(&(20, 29)), ranges.first()); + assert_eq!(Some(&(640, 649)), ranges.last()); + } + + #[test] + fn test_count_first_skips_pure_rewrite_without_progress() { + // Four non-overlapping files whose combined size splits back into at least + // four outputs: compacting them reduces neither the file count nor any + // overlap, so there is nothing worth picking. + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 600), + new_file(20, 29, 2, 600), + new_file(40, 49, 3, 600), + new_file(60, 69, 4, 600), + ])], + Some(512), + ); + + assert!(picked.is_empty()); + } + + #[test] + fn test_count_first_allows_overlap_resolution_without_file_reduction() { + // Two large overlapping files from different runs: compacting them cannot + // reduce the file count (the output splits back into just as many files), + // but it resolves the overlap and merges two sorted runs into one. + let picked = pick_count_first( + vec![ + SortedRun::from(vec![new_file(0, 19, 1, 600)]), + SortedRun::from(vec![new_file(10, 29, 2, 600)]), + ], + Some(512), + ); + + assert_eq!(vec![(0, 19), (10, 29)], picked_ranges(&picked)); + } + + #[test] + fn test_count_first_prefers_guaranteed_reduction_over_pure_rewrite() { + // A window mixing large and small files: the large triple on its own is a + // pure rewrite, and every interval containing large files rewrites more + // bytes for the same predicted reduction, so the small triple wins. + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 600), + new_file(10, 19, 2, 600), + new_file(20, 29, 3, 600), + new_file(30, 39, 4, 10), + new_file(40, 49, 5, 10), + new_file(50, 59, 6, 10), + ])], + Some(512), + ); + + assert_eq!(vec![(30, 39), (40, 49), (50, 59)], picked_ranges(&picked)); + } + + #[test] + fn test_count_first_prefers_earlier_time_on_exact_tie() { + let files = (0..=DEFAULT_MAX_INPUT_FILES) + .map(|idx| { + let start = idx as i64 * 20; + new_file(start, start + 9, idx as u64 + 1, 10) + }) + .collect::>(); + + let picked = pick_count_first(vec![SortedRun::from(files)], None); + let ranges = picked_ranges(&picked); + + assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len()); + assert_eq!(Some(&(0, 9)), ranges.first()); + assert_eq!(Some(&(620, 629)), ranges.last()); + } + + #[test] + fn test_count_first_keeps_each_interleaved_run_contiguous() { + let picked = pick_count_first( + vec![ + SortedRun::from(vec![ + new_file(0, 9, 1, 1), + new_file(20, 29, 2, 1), + new_file(40, 49, 3, 1), + ]), + SortedRun::from(vec![new_file(10, 19, 4, 1), new_file(30, 39, 5, 1)]), + ], + None, + ); + + assert_eq!( + vec![(0, 9), (10, 19), (20, 29), (30, 39), (40, 49)], + picked_ranges(&picked) + ); + } + // TODO(hl): TTL tester that checks if get_expired_ssts function works as expected. + + // Reproduces the leftover-small-files deadlock observed in the compaction bench: + // an "EngineFull" flush emits several SSTs sharing one sequence. They must be ordinary + // candidates so a contiguous candidate can span them instead of treating them as an + // indivisible block that fragments the timeline. + #[test] + fn test_count_first_candidate_spans_same_sequence_files() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 10), + // One EngineFull flush output: three SSTs with the same sequence. + new_file(10, 19, 2, 10), + new_file(10, 19, 2, 10), + new_file(10, 19, 2, 10), + new_file(20, 29, 3, 10), + ])], + None, + ); + + assert_eq!( + vec![(0, 9), (10, 19), (10, 19), (10, 19), (20, 29)], + picked_ranges(&picked) + ); + } + + // End-to-end repro of the compaction-bench leftover files: a window where + // EngineFull flush groups (multiple SSTs sharing a sequence) sit between + // singleton flushes. The picker must be able to merge the whole contiguous + // range; otherwise these files can never be compacted once the active window + // moves on. + #[tokio::test] + async fn test_count_first_merges_window_with_interleaved_flush_groups() { + let files = [ + // Singleton flush. + new_file_handle_with_size_and_sequence(FileId::random(), 0, 9, 0, 1, 10), + // EngineFull flush: three SSTs sharing sequence 2. + new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10), + // Singleton flush. + new_file_handle_with_size_and_sequence(FileId::random(), 20, 29, 0, 3, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 3, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + // All five SSTs of the window are merged in one compaction. + assert_eq!(5, output[0].inputs.len()); + } } diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index 1fc3a8f785d..04c97bdfbba 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -768,6 +768,7 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() { .await; let create = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "4") .build(); let column_schemas = create .column_metadatas @@ -779,7 +780,9 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() { .await .unwrap(); put_and_flush(&engine, region_id, &column_schemas, 0..10).await; - put_and_flush(&engine, region_id, &column_schemas, 5..20).await; + put_and_flush(&engine, region_id, &column_schemas, 10..20).await; + put_and_flush(&engine, region_id, &column_schemas, 20..30).await; + put_and_flush(&engine, region_id, &column_schemas, 30..40).await; let commit_guard = gate.arm_commit(); let compact_engine = engine.clone(); @@ -803,7 +806,7 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() { RegionRequest::Truncate(RegionTruncateRequest::ByTimeRanges { time_ranges: vec![( Timestamp::new_millisecond(0), - Timestamp::new_millisecond(19_000), + Timestamp::new_millisecond(39_000), )], }), ) @@ -981,12 +984,10 @@ async fn test_compaction_region_with_format(flat_format: bool) { // [20....29] // -[15.........29]- (delete) // [15.....24] - // Output: - // [0..9] - // [10............29] (contains delete) - // [15....24] + // Count-first compaction consumes the first 4 SSTs as soon as the trigger is reached. + // The compacted output and final flush leave 2 SSTs. assert_eq!( - 3, + 2, scanner.num_files(), "unexpected files: {:?}", scanner.file_ids() @@ -1530,6 +1531,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) { let request = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "4") .build(); let column_schemas = request @@ -1542,7 +1544,9 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) { .await .unwrap(); let listener_guard = CompactionListenerGuard::new(listener.clone()); - // Flush 2 SSTs for compaction. + // Flush 4 balanced SSTs for compaction. + put_and_flush(&engine, region_id, &column_schemas, 0..10).await; + put_and_flush(&engine, region_id, &column_schemas, 5..20).await; put_and_flush(&engine, region_id, &column_schemas, 0..10).await; put_and_flush(&engine, region_id, &column_schemas, 5..20).await; @@ -1573,7 +1577,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) { .await .unwrap(); assert_eq!( - 2, + 4, scanner.num_files(), "unexpected files: {:?}", scanner.file_ids() @@ -1614,6 +1618,7 @@ async fn test_local_compaction_cancellation_notifies_before_pending_ddl_dispatch let request = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "4") .build(); let column_schemas = request .column_metadatas @@ -1628,6 +1633,8 @@ async fn test_local_compaction_cancellation_notifies_before_pending_ddl_dispatch let merge_guard = gate.arm_merge(); put_and_flush(&engine, region_id, &column_schemas, 0..10).await; put_and_flush(&engine, region_id, &column_schemas, 5..20).await; + put_and_flush(&engine, region_id, &column_schemas, 0..10).await; + put_and_flush(&engine, region_id, &column_schemas, 5..20).await; tokio::time::timeout(Duration::from_secs(5), gate.wait_until_merge_entered()) .await diff --git a/src/mito2/src/engine/merge_mode_test.rs b/src/mito2/src/engine/merge_mode_test.rs index 40a87642ae4..afb984e77bd 100644 --- a/src/mito2/src/engine/merge_mode_test.rs +++ b/src/mito2/src/engine/merge_mode_test.rs @@ -127,6 +127,7 @@ async fn test_merge_mode_compaction_with_format(flat_format: bool) { let request = CreateRequestBuilder::new() .field_num(2) .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "3") .insert_option("merge_mode", "last_non_null") .build(); let table_dir = request.table_dir.clone(); @@ -208,7 +209,7 @@ async fn test_merge_mode_compaction_with_format(flat_format: bool) { .scanner(region_id, ScanRequest::default()) .await .unwrap(); - assert_eq!(2, scanner.num_files()); + assert_eq!(1, scanner.num_files()); assert_eq!(1, scanner.num_memtables()); scanner.set_target_partitions(2); let stream = scanner.scan().await.unwrap(); diff --git a/src/mito2/src/gc/worker_test.rs b/src/mito2/src/gc/worker_test.rs index cd8b875e9e8..e0b03484996 100644 --- a/src/mito2/src/gc/worker_test.rs +++ b/src/mito2/src/gc/worker_test.rs @@ -399,7 +399,8 @@ async fn test_gc_worker_basic_compact() { let region = engine.get_region(region_id).unwrap(); let manifest = region.manifest_ctx.manifest().await; - assert_eq!(manifest.removed_files.removed_files[0].files.len(), 3); + let removed_file_num = manifest.removed_files.removed_files[0].files.len(); + assert!(removed_file_num > 0); let version = manifest.manifest_version; @@ -413,7 +414,10 @@ async fn test_gc_worker_basic_compact() { let gc_worker = create_gc_worker(&engine, regions, &file_ref_manifest, true).await; let report = gc_worker.run().await.unwrap(); - assert_eq!(report.deleted_files.get(®ion_id).unwrap().len(), 3,); + assert_eq!( + report.deleted_files.get(®ion_id).unwrap().len(), + removed_file_num + ); assert!(report.need_retry_regions.is_empty()); } @@ -476,7 +480,7 @@ async fn test_gc_worker_compact_with_ref() { let region = engine.get_region(region_id).unwrap(); let manifest = region.manifest_ctx.manifest().await; - assert_eq!(manifest.removed_files.removed_files[0].files.len(), 3); + assert!(!manifest.removed_files.removed_files[0].files.is_empty()); let version = manifest.manifest_version; diff --git a/src/mito2/src/worker/handle_compaction.rs b/src/mito2/src/worker/handle_compaction.rs index 15dc884a1e5..1f8a66887cd 100644 --- a/src/mito2/src/worker/handle_compaction.rs +++ b/src/mito2/src/worker/handle_compaction.rs @@ -131,6 +131,11 @@ impl RegionWorkerLoop { return; } let execution = request.execution.clone(); + // Whether this execution reduced the SST file count. The scheduler keeps + // draining while compaction makes progress; a rewrite that did not reduce + // files (e.g. its output was split into more files than its input) must end + // the chain to avoid a no-progress loop. + let made_progress = request.edit.files_to_remove.len() > request.edit.files_to_add.len(); region.version_control.apply_edit( Some(request.edit.clone()), @@ -167,6 +172,7 @@ impl RegionWorkerLoop { &execution, ®ion.manifest_ctx, self.schema_metadata_manager.clone(), + made_progress, ) .await; match transition { From de6a903cb196d1ad01aa314f8b138f53f8c27ff7 Mon Sep 17 00:00:00 2001 From: zhaiAohan <1573612521@qq.com> Date: Sun, 30 Aug 2026 05:26:00 +0000 Subject: [PATCH 125/126] feat: Add built-in daemon mode for standalone service (#8960) * feat: Add built-in daemon mode for standalone service Add an optional --daemon/-d flag to `greptimedb standalone start` to run the server as a background daemon detached from the shell session, similar to Redis `daemonize yes`. - Use the daemonize crate (Unix only) to fork/setsid and detach from the controlling terminal before the Tokio runtime is built. - stdin/stdout/stderr are redirected to /dev/null; logs continue to go to data_home/logs/ via the existing tracing appender. - On non-Unix platforms, --daemon is a no-op with a warning (runs in foreground). Closes #7314 Signed-off-by: tian1220A <1573612521@qq.com> * test(cmd): add integration test for standalone daemon mode Verify that standalone start --daemon detaches from the shell and brings up the HTTP listener, guarding against regressions where the daemon blocks in the foreground or crashes after forking. Signed-off-by: tian1220A <1573612521@qq.com> * fix(cmd): gate --daemon flag with #[cfg(unix)] and fix rustfmt Address review feedback to not expose --daemon/-d on non-Unix platforms. The daemon field is now gated with #[cfg(unix)], and a cross-platform is_daemon() accessor returns false on non-Unix so maybe_daemonize() stays unconditional. Also fix the rustfmt formatting issues that caused fmt-check to fail. Signed-off-by: tian1220A <1573612521@qq.com> * chore(cmd): declare daemonize dependency directly in cmd daemonize is only used by the cmd crate. Move it out of the workspace dependencies and declare it directly in the unix-only dependencies of src/cmd/Cargo.toml. Signed-off-by: tian1220A <1573612521@qq.com> --------- Signed-off-by: tian1220A <1573612521@qq.com> --- Cargo.lock | 10 +++ src/cmd/Cargo.toml | 1 + src/cmd/src/bin/greptime.rs | 58 ++++++++++++++-- src/cmd/src/error.rs | 10 ++- src/cmd/src/standalone.rs | 39 +++++++++++ src/cmd/tests/standalone_daemon.rs | 107 +++++++++++++++++++++++++++++ 6 files changed, 217 insertions(+), 8 deletions(-) create mode 100644 src/cmd/tests/standalone_daemon.rs diff --git a/Cargo.lock b/Cargo.lock index d77ec340a8d..8fe6d74dd2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2167,6 +2167,7 @@ dependencies = [ "common-time", "common-version", "common-wal", + "daemonize", "datafusion", "datafusion-common", "datafusion-physical-plan", @@ -3570,6 +3571,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "daemonize" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab8bfdaacb3c887a54d41bdf48d3af8873b3f5566469f8ba21b92057509f116e" +dependencies = [ + "libc", +] + [[package]] name = "darling" version = "0.14.4" diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 579e5b1454c..8cb8cff4300 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -118,6 +118,7 @@ tonic.workspace = true tracing-appender.workspace = true [target.'cfg(unix)'.dependencies] +daemonize = "0.5" pprof = { version = "0.14", features = [ "flamegraph", ] } diff --git a/src/cmd/src/bin/greptime.rs b/src/cmd/src/bin/greptime.rs index f5f4a590487..2fa75ff4171 100644 --- a/src/cmd/src/bin/greptime.rs +++ b/src/cmd/src/bin/greptime.rs @@ -17,6 +17,8 @@ use clap::{Parser, Subcommand}; use cmd::datanode::builder::InstanceBuilder; +#[cfg(unix)] +use cmd::error::DaemonizeSnafu; use cmd::error::{InitTlsProviderSnafu, Result}; use cmd::options::GlobalOptions; use cmd::{App, cli, datanode, flownode, frontend, metasrv, standalone, user}; @@ -73,20 +75,24 @@ static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; #[cfg(debug_assertions)] fn main() -> Result<()> { use snafu::ResultExt; + + let cli = Command::parse(); + maybe_daemonize(&cli)?; + // Set the stack size to 8MB for the thread so it wouldn't overflow on large stack usage in debug mode // see https://github.com/GreptimeTeam/greptimedb/pull/4317 // and https://github.com/rust-lang/rust/issues/34283 std::thread::Builder::new() .name("main_spawn".to_string()) .stack_size(8 * 1024 * 1024) - .spawn(|| { + .spawn(move || { { tokio::runtime::Builder::new_multi_thread() .thread_stack_size(8 * 1024 * 1024) .enable_all() .build() .expect("Failed building the Runtime") - .block_on(main_body()) + .block_on(main_body(cli)) } }) .context(cmd::error::SpawnThreadSnafu)? @@ -95,15 +101,53 @@ fn main() -> Result<()> { } #[cfg(not(debug_assertions))] -#[tokio::main] -async fn main() -> Result<()> { - main_body().await +fn main() -> Result<()> { + let cli = Command::parse(); + maybe_daemonize(&cli)?; + + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("Failed building the Runtime") + .block_on(main_body(cli)) } -async fn main_body() -> Result<()> { +async fn main_body(cli: Command) -> Result<()> { setup_human_panic(); install_default_crypto_provider().map_err(|msg| InitTlsProviderSnafu { msg }.build())?; - start(Command::parse()).await + start(cli).await +} + +/// If `--daemon` was passed for `standalone start`, fork into the background +/// before any Tokio runtime is built. On non-Unix platforms the `--daemon` +/// flag does not exist, so this is a no-op. +fn maybe_daemonize(cli: &Command) -> Result<()> { + let daemon = matches!(&cli.subcmd, SubCommand::Standalone(cmd) if cmd.is_daemon()); + if !daemon { + return Ok(()); + } + + #[cfg(unix)] + { + let cwd = std::env::current_dir().map_err(|e| { + DaemonizeSnafu { + msg: format!("failed to read current working directory: {e}"), + } + .build() + })?; + + daemonize::Daemonize::new() + .working_directory(cwd) + .start() + .map_err(|e| { + DaemonizeSnafu { + msg: format!("failed to daemonize: {e}"), + } + .build() + })?; + } + + Ok(()) } async fn start(cli: Command) -> Result<()> { diff --git a/src/cmd/src/error.rs b/src/cmd/src/error.rs index b80828d13ac..82273c0b67d 100644 --- a/src/cmd/src/error.rs +++ b/src/cmd/src/error.rs @@ -162,6 +162,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to daemonize: {msg}"))] + Daemonize { + msg: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Unsupported selector type: {}", selector_type))] UnsupportedSelectorType { selector_type: String, @@ -361,7 +368,8 @@ impl ErrorExt for Error { Error::SerdeJson { .. } | Error::FileIo { .. } | Error::SpawnThread { .. } - | Error::InitTlsProvider { .. } => StatusCode::Unexpected, + | Error::InitTlsProvider { .. } + | Error::Daemonize { .. } => StatusCode::Unexpected, Error::Other { source, .. } => source.status_code(), diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 2b5bd5e8cec..d628461ba75 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -152,6 +152,11 @@ impl Command { ) -> Result> { self.subcmd.load_options(global_options) } + + /// Whether the `standalone start` command requested daemonization. + pub fn is_daemon(&self) -> bool { + self.subcmd.is_daemon() + } } #[derive(Parser)] @@ -174,6 +179,12 @@ impl SubCommand { SubCommand::Start(cmd) => cmd.load_options(global_options), } } + + fn is_daemon(&self) -> bool { + match self { + SubCommand::Start(cmd) => cmd.is_daemon(), + } + } } pub struct Instance { @@ -291,9 +302,24 @@ pub struct StartCommand { /// The working home directory of this standalone instance. #[clap(long)] data_home: Option, + /// Run in the background as a daemon. + #[cfg(unix)] + #[clap(short, long)] + daemon: bool, } impl StartCommand { + /// Whether the `standalone start` command requested daemonization. + #[cfg(unix)] + pub(crate) fn is_daemon(&self) -> bool { + self.daemon + } + + #[cfg(not(unix))] + pub(crate) fn is_daemon(&self) -> bool { + false + } + /// Load the GreptimeDB options from various sources (command line, config file or env). pub fn load_options( &self, @@ -1356,6 +1382,19 @@ mod tests { assert_eq!(command.grpc_bind_addr.as_deref(), Some("127.0.0.1:34001")); } + #[cfg(unix)] + #[test] + fn test_parse_daemon_flag() { + let command = StartCommand::try_parse_from(["standalone", "--daemon"]).unwrap(); + assert!(command.is_daemon()); + + let command = StartCommand::try_parse_from(["standalone", "-d"]).unwrap(); + assert!(command.is_daemon()); + + let command = StartCommand::try_parse_from(["standalone"]).unwrap(); + assert!(!command.is_daemon()); + } + #[test] fn test_help_uses_grpc_option_names() { let mut cmd = StartCommand::command(); diff --git a/src/cmd/tests/standalone_daemon.rs b/src/cmd/tests/standalone_daemon.rs new file mode 100644 index 00000000000..2f4d1cd068e --- /dev/null +++ b/src/cmd/tests/standalone_daemon.rs @@ -0,0 +1,107 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::net::TcpStream; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; + +/// Stops the orphaned background daemon when the test finishes (including on +/// panic). The daemon's command line contains the unique `--data-home` path, so +/// we match on it to locate the process. +struct DaemonGuard(String); + +impl Drop for DaemonGuard { + fn drop(&mut self) { + let _ = Command::new("pkill") + .args(["-f", &self.0]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + } +} + +#[test] +#[cfg(unix)] +fn test_standalone_daemon_mode_detaches_from_shell() { + let bin = env!("CARGO_BIN_EXE_greptime"); + let data_home = tempfile::tempdir().unwrap(); + let data_home_path = data_home.path().to_str().unwrap().to_string(); + + let http_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + let grpc_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + let mysql_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + let postgres_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + + let mut child = Command::new(bin) + .args([ + "standalone", + "start", + "--daemon", + "--data-home", + &data_home_path, + "--http-addr", + &http_addr, + "--grpc-bind-addr", + &grpc_addr, + "--mysql-addr", + &mysql_addr, + "--postgres-addr", + &postgres_addr, + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(); + + let _guard = DaemonGuard(data_home_path); + + // When daemonization works, the parent forks and exits right away. When it + // does not, the process stays in the foreground and never returns. + let detach_timeout = Duration::from_secs(30); + let started = Instant::now(); + loop { + match child.try_wait().unwrap() { + Some(status) => { + assert!( + status.success(), + "daemon start should exit successfully, got {status:?}" + ); + break; + } + None => { + assert!( + started.elapsed() < detach_timeout, + "daemon mode should detach and return quickly, but it blocked for over {detach_timeout:?}" + ); + std::thread::sleep(Duration::from_millis(200)); + } + } + } + + // A successful detach is not enough: a broken implementation could fork, + // have the parent exit, then crash the child. Poll the HTTP address to + // confirm the daemon actually came up and is listening. + let ready_timeout = Duration::from_secs(30); + let started = Instant::now(); + loop { + if TcpStream::connect(&http_addr).is_ok() { + break; + } + assert!( + started.elapsed() < ready_timeout, + "daemon should come up and listen on {http_addr}, but it was unreachable for {ready_timeout:?}" + ); + std::thread::sleep(Duration::from_millis(200)); + } +} From d32cd77505d3cb304d66796b515c45ab42e739c6 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Sun, 30 Aug 2026 13:24:17 +0000 Subject: [PATCH 126/126] fix: postgres describe for more statements (#8974) * fix: postgres describe for more statements Signed-off-by: Ning Sun * fix: cover more show statements Signed-off-by: Ning Sun * fix: address review comments - add missing `clippy::too_many_arguments` allow on `query_from_information_schema_dataframe` (CI clippy failure) - take `&ShowKind` in the information-schema dataframe helper so `kind` is no longer cloned at every call site; only the WHERE arm (which needs an owned expression for `sql_to_expr`) clones internally - document why re-applying TQL explain formats never overwrites an existing value (per-query context state) Signed-off-by: Ning Sun * chore: trim comments to essentials Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- src/common/recordbatch/src/cursor.rs | 1 - src/frontend/src/instance.rs | 101 ++++++++ src/frontend/src/lib.rs | 2 + src/operator/src/statement.rs | 2 +- src/operator/src/statement/admin.rs | 171 ++++++++----- src/query/src/analyze.rs | 10 +- src/query/src/sql.rs | 249 +++++++++++++++---- src/servers/src/postgres/handler.rs | 83 ++++--- src/session/src/lib.rs | 5 +- tests-integration/tests/sql.rs | 343 +++++++++++++++++++++++++++ 10 files changed, 825 insertions(+), 142 deletions(-) diff --git a/src/common/recordbatch/src/cursor.rs b/src/common/recordbatch/src/cursor.rs index 60ae9d54fa2..125a68f4f13 100644 --- a/src/common/recordbatch/src/cursor.rs +++ b/src/common/recordbatch/src/cursor.rs @@ -47,7 +47,6 @@ impl RecordBatchStreamCursor { } } - /// Returns the schema of the underlying record batch stream. pub fn schema(&self) -> SchemaRef { self.schema.clone() } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 412ff495f65..90650014e6a 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -59,6 +59,7 @@ use common_recordbatch::error::StreamTimeoutSnafu; use common_telemetry::logging::SlowQueryOptions; use common_telemetry::{debug, error, tracing}; use dashmap::DashMap; +use datafusion::dataframe::DataFrame; use datafusion::physical_plan::ExecutionPlan; use datafusion_expr::LogicalPlan; use futures::{Stream, StreamExt, future}; @@ -847,6 +848,22 @@ impl Instance { query_interceptor.pre_execute(stmt.as_ref(), Some(&plan), query_ctx.clone())?; + // TQL EXPLAIN/ANALYZE formats are consumed from the query context at + // execution time (see `optimize_physical_plan`); re-apply the side + // effect of `plan_tql` that was lost when the plan was built during + // Describe. `explain_format` is per-query state, so this never + // overwrites anything. + if let Some(Statement::Tql(tql)) = &stmt { + let format = match tql { + Tql::Explain(explain) => explain.format.as_ref(), + Tql::Analyze(analyze) => analyze.format.as_ref(), + Tql::Eval(_) => None, + }; + if let Some(format) = format { + query_ctx.set_explain_format(format.to_string()); + } + } + let query = stmt .as_ref() .map(|s| s.to_string()) @@ -928,6 +945,59 @@ impl Instance { vec![result] } + /// Builds the [`DataFrame`] for an information-schema-backed `SHOW` + /// statement; `None` for other statements. The future is boxed to keep + /// `do_describe_inner`'s state machine small. + fn show_statement_dataframe<'a>( + &'a self, + stmt: &'a Statement, + query_ctx: &'a QueryContextRef, + ) -> Pin>> + Send + 'a>> { + Box::pin(async move { + let engine = &self.query_engine; + let catalog_manager = self.catalog_manager(); + let ctx = query_ctx.clone(); + let dataframe = match stmt { + Statement::ShowDatabases(show) => { + query::sql::show_databases_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowTables(show) => { + query::sql::show_tables_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowViews(show) => { + query::sql::show_views_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowFlows(show) => { + query::sql::show_flows_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowColumns(show) => { + query::sql::show_columns_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowTableStatus(show) => { + query::sql::show_table_status_dataframe(show, engine, catalog_manager, ctx) + .await + } + Statement::ShowCharset(kind) => { + query::sql::show_charsets_dataframe(kind, engine, catalog_manager, ctx).await + } + Statement::ShowCollation(kind) => { + query::sql::show_collations_dataframe(kind, engine, catalog_manager, ctx).await + } + Statement::ShowIndex(show) => { + query::sql::show_index_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowRegion(show) => { + query::sql::show_region_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowProcesslist(show) => { + query::sql::show_processlist_dataframe(show, engine, catalog_manager, ctx).await + } + _ => return None, + }; + Some(dataframe) + }) + } + async fn do_describe_inner( &self, stmt: Statement, @@ -947,6 +1017,37 @@ impl Instance { let plannable = is_inner_plannable(&stmt) || matches!(&stmt, Statement::Explain(explain) if is_inner_plannable(explain.statement.as_ref())); + if let Statement::Tql(tql) = stmt { + // TQL produces a logical plan; describe it from the plan so the + // extended-protocol RowDescription matches the executed DataRows. + self.check_sql_permission(&Statement::Tql(tql.clone()), &query_ctx) + .await?; + let plan = self.statement_executor.plan_tql(tql, &query_ctx).await?; + return self + .query_engine + .describe(plan, query_ctx) + .await + .map(Some) + .context(error::DescribeStatementSnafu); + } + + // Describe SHOW statements from the same projection the executor builds. + if let Some(dataframe) = self + .show_statement_dataframe(&stmt, &query_ctx) + .await + .transpose() + .context(PlanStatementSnafu)? + { + self.check_sql_permission(&stmt, &query_ctx).await?; + let plan = dataframe.into_unoptimized_plan(); + return self + .query_engine + .describe(plan, query_ctx) + .await + .map(Some) + .context(error::DescribeStatementSnafu); + } + if plannable { self.check_sql_permission(&stmt, &query_ctx).await?; diff --git a/src/frontend/src/lib.rs b/src/frontend/src/lib.rs index c1702360739..1468fe7c0b2 100644 --- a/src/frontend/src/lib.rs +++ b/src/frontend/src/lib.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#![recursion_limit = "256"] + pub mod error; pub mod events; pub mod frontend; diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 478f06e83f7..21c3ed61464 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -84,7 +84,7 @@ use table::table_reference::TableReference; pub use self::admin::{ AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef, AdminFunctionRecordingLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, - AdminFunctionServiceRef, + AdminFunctionServiceRef, admin_output_schema, }; use self::set::{ set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding, diff --git a/src/operator/src/statement/admin.rs b/src/operator/src/statement/admin.rs index 2b4a909e3cd..ff8b26915cc 100644 --- a/src/operator/src/statement/admin.rs +++ b/src/operator/src/statement/admin.rs @@ -19,6 +19,7 @@ use std::sync::Arc; use common_function::function::FunctionContext; use common_function::function_registry::{FUNCTION_REGISTRY, get_admin_function}; +use common_function::state::FunctionState; use common_query::Output; use common_recordbatch::{RecordBatch, RecordBatches}; use common_sql::convert::sql_value_to_value; @@ -77,6 +78,103 @@ struct CoreAdminFunctionService { query_engine: query::QueryEngineRef, } +/// Parts of an `ADMIN` call needed both for execution and schema derivation. +struct ResolvedAdminFunction { + admin_udf: datafusion_expr::ScalarUDF, + fn_name: String, + args: Vec, + arg_types: Vec, + ret_type: ArrowDataType, +} + +/// Resolves the function, parses its literal arguments and derives its +/// return type, without executing it. +fn resolve_admin_function( + stmt: &Admin, + query_ctx: &QueryContextRef, + state: Arc, +) -> Result { + let Admin::Func(func) = stmt; + // the function name should be in lower case. + let func_name = func.name.to_string().to_lowercase(); + let factory = get_admin_function(&func_name) + .or_else(|| FUNCTION_REGISTRY.get_function(&func_name)) + .context(error::AdminFunctionNotFoundSnafu { + name: func_name.clone(), + })?; + + let func_ctx = FunctionContext { + query_ctx: query_ctx.clone(), + state, + }; + + let admin_udf = factory.provide(func_ctx); + admin_udf + .as_async() + .context(error::AdminFunctionNotFoundSnafu { name: func_name })?; + + let fn_name = admin_udf.name().to_string(); + let signature = admin_udf.signature(); + + // Parse function arguments + let FunctionArguments::List(args) = &func.args else { + return error::BuildAdminFunctionArgsSnafu { + msg: format!("unsupported function args {} for {}", func.args, fn_name), + } + .fail(); + }; + let arg_values = args + .args + .iter() + .map(|arg| { + let FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) = arg else { + return error::BuildAdminFunctionArgsSnafu { + msg: format!("unsupported function arg {arg} for {}", fn_name), + } + .fail(); + }; + Ok(&value.value) + }) + .collect::>>()?; + + let args = args_to_vector(&signature.type_signature, &arg_values, query_ctx)?; + let arg_types = args + .iter() + .map(|arg| arg.data_type().as_arrow_type()) + .collect::>(); + let ret_type = + admin_udf + .return_type(&arg_types) + .map_err(|e| error::Error::BuildAdminFunctionArgs { + msg: format!( + "Failed to get return type of admin function {}: {}", + fn_name, e + ), + })?; + + Ok(ResolvedAdminFunction { + admin_udf, + fn_name, + args, + arg_types, + ret_type, + }) +} + +/// Output schema of an `ADMIN` statement, mirroring what +/// [`CoreAdminFunctionService::execute`] produces. `None` if the function +/// or arguments are unresolvable; execution will then surface the error. +pub fn admin_output_schema(stmt: &Admin, query_ctx: &QueryContextRef) -> Option { + let resolved = + resolve_admin_function(stmt, query_ctx, Arc::new(FunctionState::default())).ok()?; + Some(Schema::new(vec![ColumnSchema::new( + // Use statement as the result column name + stmt.to_string(), + ConcreteDataType::from_arrow_type(&resolved.ret_type), + false, + )])) +} + impl CoreAdminFunctionService { fn new(query_engine: query::QueryEngineRef) -> Self { Self { query_engine } @@ -88,62 +186,23 @@ impl CoreAdminFunctionService { query_ctx, } = request; - let Admin::Func(func) = &stmt; - // the function name should be in lower case. - let func_name = func.name.to_string().to_lowercase(); - let factory = get_admin_function(&func_name) - .or_else(|| FUNCTION_REGISTRY.get_function(&func_name)) - .context(error::AdminFunctionNotFoundSnafu { - name: func_name.clone(), - })?; - - let func_ctx = FunctionContext { - query_ctx: query_ctx.clone(), - state: self.query_engine.engine_state().function_state(), - }; - - let admin_udf = factory.provide(func_ctx); + let resolved = resolve_admin_function( + &stmt, + &query_ctx, + self.query_engine.engine_state().function_state(), + )?; + let ResolvedAdminFunction { + admin_udf, + fn_name, + args, + arg_types, + ret_type, + } = resolved; let admin_async_fn = admin_udf .as_async() - .context(error::AdminFunctionNotFoundSnafu { name: func_name })?; - - let fn_name = admin_udf.name(); - let signature = admin_udf.signature(); - - // Parse function arguments - let FunctionArguments::List(args) = &func.args else { - return error::BuildAdminFunctionArgsSnafu { - msg: format!("unsupported function args {} for {}", func.args, fn_name), - } - .fail(); - }; - let arg_values = args - .args - .iter() - .map(|arg| { - let FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) = arg else { - return error::BuildAdminFunctionArgsSnafu { - msg: format!("unsupported function arg {arg} for {}", fn_name), - } - .fail(); - }; - Ok(&value.value) - }) - .collect::>>()?; - - let args = args_to_vector(&signature.type_signature, &arg_values, &query_ctx)?; - let arg_types = args - .iter() - .map(|arg| arg.data_type().as_arrow_type()) - .collect::>(); - let ret_type = admin_udf.return_type(&arg_types).map_err(|e| { - error::Error::BuildAdminFunctionArgs { - msg: format!( - "Failed to get return type of admin function {}: {}", - fn_name, e - ), - } - })?; + .context(error::AdminFunctionNotFoundSnafu { + name: fn_name.clone(), + })?; // Convert arguments to DataFusion ColumnarValue format let columnar_args: Vec = args @@ -174,9 +233,7 @@ impl CoreAdminFunctionService { let result_columnar = admin_async_fn .invoke_async_with_args(func_args) .await - .with_context(|_| ExecuteAdminFunctionSnafu { - msg: fn_name.to_string(), - })?; + .with_context(|_| ExecuteAdminFunctionSnafu { msg: fn_name })?; // Convert result back to VectorRef let result_columnar: common_query::prelude::ColumnarValue = diff --git a/src/query/src/analyze.rs b/src/query/src/analyze.rs index b1522fc05d5..2ab3ec5176e 100644 --- a/src/query/src/analyze.rs +++ b/src/query/src/analyze.rs @@ -46,13 +46,9 @@ const STAGE: &str = "stage"; const NODE: &str = "node"; const PLAN: &str = "plan"; -/// Returns the fixed output schema of [`DistAnalyzeExec`]: -/// (`stage`: UInt32, `node`: UInt32, `plan`: Utf8). -/// -/// Exposed so protocol-level `Describe` handlers can advertise the schema -/// clients will actually receive (the physical plan is rewritten in -/// `optimize_physical_plan` and its schema differs from the DataFusion -/// `Analyze` logical plan schema). +/// Fixed output schema of [`DistAnalyzeExec`], for `Describe` handlers: +/// execution rewrites the plan in `optimize_physical_plan`, so this schema +/// differs from the logical `Analyze` plan's. pub fn dist_analyze_output_schema() -> SchemaRef { SchemaRef::new(Schema::new(vec![ Field::new(STAGE, DataType::UInt32, true), diff --git a/src/query/src/sql.rs b/src/query/src/sql.rs index b2d979c6664..93e74ca33ec 100644 --- a/src/query/src/sql.rs +++ b/src/query/src/sql.rs @@ -41,6 +41,7 @@ use common_recordbatch::RecordBatches; use common_recordbatch::adapter::RecordBatchStreamAdapter; use common_time::Timestamp; use common_time::timezone::get_timezone; +use datafusion::dataframe::DataFrame; use datafusion::prelude::SessionContext; use datafusion_expr::{Expr, SortExpr, col, lit}; use datatypes::prelude::*; @@ -109,7 +110,7 @@ const INDEX_KEY_NAME_COLUMN: &str = "Key_name"; const INDEX_SEQ_IN_INDEX_COLUMN: &str = "Seq_in_index"; const INDEX_COLUMN_NAME_COLUMN: &str = "Column_name"; -static DESCRIBE_TABLE_OUTPUT_SCHEMA: Lazy> = Lazy::new(|| { +pub static DESCRIBE_TABLE_OUTPUT_SCHEMA: Lazy> = Lazy::new(|| { Arc::new(Schema::new(vec![ ColumnSchema::new( COLUMN_NAME_COLUMN, @@ -178,6 +179,18 @@ pub async fn show_databases( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_databases_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW DATABASES` without executing it. +pub async fn show_databases_dataframe( + stmt: &ShowDatabases, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { let projects = if stmt.full { vec![ (schemata::SCHEMA_NAME, SCHEMAS_COLUMN), @@ -191,7 +204,7 @@ pub async fn show_databases( let like_field = Some(schemata::SCHEMA_NAME); let sort = vec![col(schemata::SCHEMA_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -201,7 +214,7 @@ pub async fn show_databases( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -249,6 +262,45 @@ async fn query_from_information_schema_table( sort: Vec, kind: ShowKind, ) -> Result { + let dataframe = query_from_information_schema_dataframe( + query_engine, + catalog_manager, + query_ctx, + table_name, + select, + projects, + filters, + like_field, + sort, + &kind, + ) + .await?; + dataframe_to_output(dataframe).await +} + +async fn dataframe_to_output(dataframe: DataFrame) -> Result { + let stream = dataframe.execute_stream().await?; + Ok(Output::new_with_stream(Box::pin( + RecordBatchStreamAdapter::try_new(stream).context(error::CreateRecordBatchSnafu)?, + ))) +} + +/// Builds the [`DataFrame`] for a `SHOW` statement without executing it, +/// so `Describe` handlers can derive the output schema from the same +/// projection the executor uses. +#[allow(clippy::too_many_arguments)] +async fn query_from_information_schema_dataframe( + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, + table_name: &str, + select: Vec, + projects: Vec<(&str, &str)>, + filters: Vec, + like_field: Option<&str>, + sort: Vec, + kind: &ShowKind, +) -> Result { let table = catalog_manager .table( query_ctx.current_catalog(), @@ -336,7 +388,7 @@ async fn query_from_information_schema_table( .expect("Must be the datafusion planner"); let filter = planner - .sql_to_expr(filter, dataframe.schema(), false, query_ctx) + .sql_to_expr(filter.clone(), dataframe.schema(), false, query_ctx) .await?; // Apply the `where` clause filters @@ -344,11 +396,7 @@ async fn query_from_information_schema_table( } }; - let stream = dataframe.execute_stream().await?; - - Ok(Output::new_with_stream(Box::pin( - RecordBatchStreamAdapter::try_new(stream).context(error::CreateRecordBatchSnafu)?, - ))) + Ok(dataframe) } /// Execute `SHOW COLUMNS` statement. @@ -358,8 +406,19 @@ pub async fn show_columns( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_columns_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW COLUMNS` without executing it. +pub async fn show_columns_dataframe( + stmt: &ShowColumns, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -397,7 +456,7 @@ pub async fn show_columns( let like_field = Some(columns::COLUMN_NAME); let sort = vec![col(columns::COLUMN_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -407,7 +466,7 @@ pub async fn show_columns( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -419,8 +478,19 @@ pub async fn show_index( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_index_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW INDEX` without executing it. +pub async fn show_index_dataframe( + stmt: &ShowIndex, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -471,7 +541,7 @@ pub async fn show_index( col(statistics::SEQ_IN_INDEX).sort(true, true), ]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -481,7 +551,7 @@ pub async fn show_index( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -493,8 +563,19 @@ pub async fn show_region( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_region_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW REGION` without executing it. +pub async fn show_region_dataframe( + stmt: &ShowRegion, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -517,7 +598,7 @@ pub async fn show_region( col(columns::PEER_ID).sort(true, true), ]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -527,7 +608,7 @@ pub async fn show_region( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -539,8 +620,19 @@ pub async fn show_tables( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_tables_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowTables`] without executing it. +pub async fn show_tables_dataframe( + stmt: &ShowTables, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -564,15 +656,16 @@ pub async fn show_tables( // Transform the WHERE clause for backward compatibility: // Replace "Tables" with "Tables_in_{schema}" to support old queries - let kind = match stmt.kind { - ShowKind::Where(mut filter) => { + let rewritten_kind = match &stmt.kind { + ShowKind::Where(filter) => { + let mut filter = filter.clone(); replace_column_in_expr(&mut filter, "Tables", &tables_column); ShowKind::Where(filter) } - other => other, + kind => kind.clone(), }; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -582,7 +675,7 @@ pub async fn show_tables( filters, like_field, sort, - kind, + &rewritten_kind, ) .await } @@ -594,8 +687,20 @@ pub async fn show_table_status( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = + show_table_status_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowTableStatus`] without executing it. +pub async fn show_table_status_dataframe( + stmt: &ShowTableStatus, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -629,7 +734,7 @@ pub async fn show_table_status( let like_field = Some(tables::TABLE_NAME); let sort = vec![col(tables::TABLE_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -639,7 +744,7 @@ pub async fn show_table_status( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -651,6 +756,18 @@ pub async fn show_collations( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_collations_dataframe(&kind, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW COLLATION` without executing it. +pub async fn show_collations_dataframe( + kind: &ShowKind, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { // Refer to https://dev.mysql.com/doc/refman/8.0/en/show-collation.html let projects = vec![ ("collation_name", "Collation"), @@ -665,7 +782,7 @@ pub async fn show_collations( let like_field = Some("collation_name"); let sort = vec![]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -687,6 +804,18 @@ pub async fn show_charsets( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_charsets_dataframe(&kind, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW CHARSET` without executing it. +pub async fn show_charsets_dataframe( + kind: &ShowKind, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { // Refer to https://dev.mysql.com/doc/refman/8.0/en/show-character-set.html let projects = vec![ ("character_set_name", "Charset"), @@ -699,7 +828,7 @@ pub async fn show_charsets( let like_field = Some("character_set_name"); let sort = vec![]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -916,8 +1045,19 @@ pub async fn show_views( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_views_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowViews`] without executing it. +pub async fn show_views_dataframe( + stmt: &ShowViews, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -930,7 +1070,7 @@ pub async fn show_views( let like_field = Some(tables::TABLE_NAME); let sort = vec![col(tables::TABLE_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -940,7 +1080,7 @@ pub async fn show_views( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -952,12 +1092,23 @@ pub async fn show_flows( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = show_flows_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowFlows`] without executing it. +pub async fn show_flows_dataframe( + stmt: &ShowFlows, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { let projects = vec![(flows::FLOW_NAME, FLOWS_COLUMN)]; let filters = vec![col(flows::TABLE_CATALOG).eq(lit(query_ctx.current_catalog()))]; let like_field = Some(flows::FLOW_NAME); let sort = vec![col(flows::FLOW_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -967,7 +1118,7 @@ pub async fn show_flows( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -1340,6 +1491,18 @@ pub async fn show_processlist( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_processlist_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW PROCESSLIST` without executing it. +pub async fn show_processlist_dataframe( + stmt: &ShowProcessList, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { let projects = if stmt.full { vec![ (process_list::ID, "Id"), @@ -1367,17 +1530,17 @@ pub async fn show_processlist( }; let like_field = None; let sort = vec![col("id").sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx.clone(), "process_list", vec![], - projects.clone(), + projects, filters, like_field, sort, - ShowKind::All, + &ShowKind::All, ) .await } diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 433eb11cdad..aca4362bd34 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -28,6 +28,7 @@ use datafusion_pg_catalog::sql::PostgresCompatibilityParser; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{Schema, SchemaRef}; use futures::{Sink, SinkExt, Stream, StreamExt, future, stream}; +use operator::statement::admin_output_schema; use pgwire::api::portal::{Format, Portal}; use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler}; use pgwire::api::results::{ @@ -43,6 +44,7 @@ use pgwire::messages::data::DataRow; use query::dist_analyze_output_schema; use query::planner::DfLogicalPlanner; use query::query_engine::DescribeResult; +use query::sql::DESCRIBE_TABLE_OUTPUT_SCHEMA; use session::Session; use session::context::QueryContextRef; use snafu::ResultExt; @@ -555,11 +557,8 @@ fn describe_fields( session: &Arc, ) -> PgWireResult> { match sql_plan { - // EXPLAIN ANALYZE: at execution time the physical plan is replaced with - // GreptimeDB's `DistAnalyzeExec` (see `optimize_physical_plan`), whose - // output schema (stage/node/plan) differs from the DataFusion `Analyze` - // logical plan schema (plan_type/plan). Describe with the schema the - // client will actually receive so the DataRow field count matches. + // Execution swaps in DistAnalyzeExec (stage/node/plan), whose schema + // differs from the logical `Analyze` plan's (plan_type/plan). SqlPlan::Plan(LogicalPlan::Analyze(_), _) => { let schema: Schema = Schema::try_from(dist_analyze_output_schema()).map_err(convert_err)?; @@ -658,17 +657,6 @@ fn describe_fields( ), ]), - // single column show statements - SqlPlan::Statement( - Statement::ShowTables(_) | Statement::ShowFlows(_) | Statement::ShowViews(_), - _, - ) => Ok(vec![FieldInfo::new( - "name".to_string(), - None, - None, - Type::TEXT, - format.format_for(0), - )]), #[cfg(feature = "enterprise")] SqlPlan::Statement(Statement::ShowTriggers(_), _) => Ok(vec![FieldInfo::new( "name".to_string(), @@ -690,22 +678,59 @@ fn describe_fields( Ok(vec![]) } } - // FETCH cursor: return the cursor's schema so the RowDescription - // matches the DataRow field count sent during Execute. + // Single column named after the variable (see `query::sql::show_variable`). + SqlPlan::Statement(Statement::ShowVariables(show), _) => Ok(vec![FieldInfo::new( + show.variable.to_string().to_uppercase(), + None, + None, + Type::TEXT, + format.format_for(0), + )]), + // Mirrors `query::sql::show_status` (currently always empty). + SqlPlan::Statement(Statement::ShowStatus(_), _) => Ok(vec![ + FieldInfo::new( + "Variable_name".to_string(), + None, + None, + Type::TEXT, + format.format_for(0), + ), + FieldInfo::new( + "Value".to_string(), + None, + None, + Type::TEXT, + format.format_for(1), + ), + ]), + SqlPlan::Statement(Statement::ShowSearchPath(_), _) => Ok(vec![FieldInfo::new( + "search_path".to_string(), + None, + None, + Type::TEXT, + format.format_for(0), + )]), + // Mirrors `query::sql::describe_table`. + SqlPlan::Statement(Statement::DescribeTable(_), _) => { + schema_to_pg(&DESCRIBE_TABLE_OUTPUT_SCHEMA, format, None).map_err(convert_err) + } + // Single column typed with the function's return type (see + // `operator::statement::admin_output_schema`). + SqlPlan::Statement(Statement::Admin(admin), _) => { + let query_ctx = session.new_query_context(); + match admin_output_schema(admin, &query_ctx) { + Some(schema) => schema_to_pg(&schema, format, None).map_err(convert_err), + // Unresolvable; execution will surface the error. + None => Ok(vec![]), + } + } + // Describe from the declared cursor's schema. SqlPlan::Statement(Statement::FetchCursor(fetch), _) => { let cursor_name = fetch.cursor_name.to_string(); match session.get_cursor(&cursor_name) { - Some(cursor) => { - // `cursor.schema()` is the GreptimeDB `SchemaRef` captured - // when the cursor was declared; `schema_to_pg` accepts it - // directly (no DataFusion -> GreptimeDB conversion needed). - schema_to_pg(&cursor.schema(), format, None).map_err(convert_err) - } - None => { - // Cursor not found (e.g. DECLARE hasn't executed yet in - // an extended-protocol batch). Return NoData as a fallback. - Ok(vec![]) - } + Some(cursor) => schema_to_pg(&cursor.schema(), format, None).map_err(convert_err), + // Cursor not declared yet; execution will error. + None => Ok(vec![]), } } _ => { diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index 44f3ee66891..c1414fc5a3b 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -111,10 +111,7 @@ impl Session { .into() } - /// Returns the cursor with the given name, if it exists. - /// - /// Cursors are stored in the session's mutable inner data and shared across - /// all query contexts created from this session. + /// Cursors are shared across query contexts created from this session. pub fn get_cursor(&self, name: &str) -> Option> { let guard = self.mutable_inner.read().unwrap(); guard.cursors.get(name).cloned() diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 32a2881dd78..c95adb483ec 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -91,6 +91,8 @@ macro_rules! sql_tests { test_mysql_prepare_stmt_insert_timestamp, test_mysql_prepare_stmt_timezone, test_mysql_federated_prepare_stmt, + test_mysql_prepare_tql_and_show, + test_postgres_extended_query_row_returning_statements, test_declare_fetch_close_cursor, test_alter_update_on, ); @@ -1487,6 +1489,280 @@ fn assert_pg_numeric_range_error(error: tokio_postgres::Error) { assert_eq!("numeric_value_out_of_range", error.message()); } +pub async fn test_postgres_extended_query_row_returning_statements(store_type: StorageType) { + // Regression test for the tokio-postgres >= 0.7.14 DataRow/RowDescription + // mismatch: statements answered with NoData at Describe but emitting + // DataRows at Execute must describe their real output schema. + let (mut guard, fe_pg_server) = setup_pg_server(store_type, "test_pg_extended_row_stmts").await; + let addr = fe_pg_server.bind_addr().unwrap().to_string(); + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + connection.await.unwrap(); + tx.send(()).unwrap(); + }); + + client + .execute( + "CREATE TABLE demo_metrics (ts timestamp time index, val double, host string primary key skipping index)", + &[], + ) + .await + .unwrap(); + client + .execute( + "INSERT INTO demo_metrics (ts, host, val) VALUES (1000, 'host-a', 1.0), (2000, 'host-b', 2.0)", + &[], + ) + .await + .unwrap(); + + // ---- SHOW DATABASES: single `Database` column ---- + let rows = client.query("SHOW DATABASES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("Database", rows[0].columns()[0].name()); + assert!(rows.iter().any(|r| r.get::<_, String>(0) == "public")); + + // ---- SHOW FULL DATABASES: `Database` + `Options` columns ---- + let rows = client.query("SHOW FULL DATABASES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + assert_eq!("Database", rows[0].columns()[0].name()); + assert_eq!("Options", rows[0].columns()[1].name()); + assert!(rows.iter().any(|r| r.get::<_, String>(0) == "public")); + + // ---- SHOW TABLES: single `Tables_in_` column ---- + let rows = client.query("SHOW TABLES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("Tables_in_public", rows[0].columns()[0].name()); + assert!(rows.iter().any(|r| r.get::<_, String>(0) == "demo_metrics")); + + // ---- SHOW FULL TABLES: `Tables_in_` + `Table_type` columns ---- + let rows = client.query("SHOW FULL TABLES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + assert_eq!("Tables_in_public", rows[0].columns()[0].name()); + assert_eq!("Table_type", rows[0].columns()[1].name()); + + // ---- SHOW VIEWS / SHOW FLOWS: empty results, still described with one column ---- + let rows = client.query("SHOW VIEWS", &[]).await.unwrap(); + assert!(rows.is_empty()); + let stmt = client.prepare("SHOW VIEWS").await.unwrap(); + assert_eq!(1, stmt.columns().len()); + assert_eq!("Views", stmt.columns()[0].name()); + let rows = client.query("SHOW FLOWS", &[]).await.unwrap(); + assert!(rows.is_empty()); + let stmt = client.prepare("SHOW FLOWS").await.unwrap(); + assert_eq!(1, stmt.columns().len()); + assert_eq!("Flows", stmt.columns()[0].name()); + + // ---- SHOW TABLE STATUS: fixed eighteen-column schema ---- + let rows = client.query("SHOW TABLE STATUS", &[]).await.unwrap(); + assert!(!rows.is_empty()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec![ + "Name", + "Engine", + "Version", + "Row_format", + "Rows", + "Avg_row_length", + "Data_length", + "Max_data_length", + "Index_length", + "Data_free", + "Auto_increment", + "Create_time", + "Update_time", + "Check_time", + "Collation", + "Checksum", + "Create_options", + "Comment", + ], + names + ); + + // ---- SHOW COLUMNS / SHOW FULL COLUMNS ---- + let rows = client + .query("SHOW COLUMNS FROM demo_metrics", &[]) + .await + .unwrap(); + assert_eq!(3, rows.len()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec![ + "Field", + "Type", + "Null", + "Key", + "Default", + "Extra", + "Greptime_type" + ], + names + ); + let rows = client + .query("SHOW FULL COLUMNS FROM demo_metrics", &[]) + .await + .unwrap(); + assert_eq!(10, rows[0].columns().len()); + + // ---- SHOW CHARSET / SHOW COLLATION ---- + let rows = client.query("SHOW CHARSET", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(4, rows[0].columns().len()); + let rows = client.query("SHOW COLLATION", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(6, rows[0].columns().len()); + + // ---- SHOW INDEX: fixed fifteen-column schema ---- + let rows = client + .query("SHOW INDEX IN demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec![ + "Table", + "Non_unique", + "Key_name", + "Seq_in_index", + "Column_name", + "Collation", + "Cardinality", + "Sub_part", + "Packed", + "Null", + "Index_type", + "Comment", + "Index_comment", + "Visible", + "Expression", + ], + names + ); + + // ---- SHOW REGION ---- + let rows = client + .query("SHOW REGION IN demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(4, rows[0].columns().len()); + + // ---- SHOW SEARCH_PATH: single string column ---- + let rows = client.query("SHOW SEARCH_PATH", &[]).await.unwrap(); + assert_eq!(1, rows.len()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("search_path", rows[0].columns()[0].name()); + assert_eq!("public", rows[0].get::<_, String>(0)); + + // ---- SHOW VARIABLES: single column named after the variable ---- + let rows = client.query("SHOW VARIABLES timezone", &[]).await.unwrap(); + assert_eq!(1, rows.len()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("TIMEZONE", rows[0].columns()[0].name()); + let _ = rows[0].get::<_, String>(0); + + // ---- DESCRIBE TABLE: fixed six-column string schema ---- + let rows = client + .query("DESCRIBE TABLE demo_metrics", &[]) + .await + .unwrap(); + assert_eq!(3, rows.len()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec!["Column", "Type", "Key", "Null", "Default", "Semantic Type"], + names + ); + // first column of each row is the column name; ensure values decode as TEXT + let columns: Vec = rows.iter().map(|r| r.get::<_, String>(0)).collect(); + assert!(columns.contains(&"ts".to_string())); + assert!(columns.contains(&"val".to_string())); + assert!(columns.contains(&"host".to_string())); + + // ---- ADMIN function: single column named after the statement ---- + let rows = client + .query("ADMIN flush_table('demo_metrics')", &[]) + .await + .unwrap(); + assert_eq!(1, rows.len()); + assert_eq!(1, rows[0].columns().len()); + assert!(rows[0].columns()[0].name().contains("flush_table")); + + // ---- TQL EVAL / EXPLAIN / ANALYZE: described from the planned query ---- + let rows = client + .query("TQL EVAL (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!(vec!["ts", "val", "host"], names); + + let rows = client + .query("TQL EXPLAIN (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + + let rows = client + .query("TQL ANALYZE (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(3, rows[0].columns().len()); + + // FORMAT JSON variants keep working through the plan-based execution path + let rows = client + .query("TQL EXPLAIN FORMAT JSON (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + let rows = client + .query("TQL ANALYZE FORMAT JSON (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + + // ---- the same statements also work over the simple query protocol ---- + for sql in [ + "SHOW DATABASES", + "SHOW FULL TABLES", + "SHOW TABLE STATUS", + "SHOW COLUMNS FROM demo_metrics", + "SHOW CHARSET", + "SHOW COLLATION", + "SHOW INDEX IN demo_metrics", + "SHOW REGION IN demo_metrics", + "SHOW SEARCH_PATH", + "DESCRIBE TABLE demo_metrics", + "ADMIN flush_table('demo_metrics')", + "TQL EVAL (0, 3000, '1s') demo_metrics", + ] { + let msgs = client.simple_query(sql).await.unwrap(); + assert!( + msgs.iter().any(|m| matches!(m, SimpleQueryMessage::Row(_))), + "simple query {sql} should return rows" + ); + } + + drop(client); + rx.await.unwrap(); + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_postgres_explain_bind_parameter(store_type: StorageType) { // Regression test for #8029: EXPLAIN / EXPLAIN ANALYZE must accept bind // parameters over the Postgres extended query protocol. @@ -1890,6 +2166,73 @@ pub async fn test_mysql_federated_prepare_stmt(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_mysql_prepare_tql_and_show(store_type: StorageType) { + // `do_describe` now plans TQL and information-schema-backed SHOW + // statements, so MySQL prepared statements derive their column metadata + // from the planned query and execute the plan directly. This exercises + // that path end-to-end. + common_telemetry::init_default_ut_logging(); + + let (mut guard, fe_mysql_server) = + setup_mysql_server(store_type, "test_mysql_prepare_tql_and_show").await; + let addr = fe_mysql_server.bind_addr().unwrap().to_string(); + + let pool = MySqlPoolOptions::new() + .max_connections(2) + .connect(&format!("mysql://{addr}/public")) + .await + .unwrap(); + + sqlx::query( + "CREATE TABLE demo_metrics (ts timestamp time index, val double, host string primary key)", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query("INSERT INTO demo_metrics (ts, host, val) VALUES (1000, 'host-a', 1.0), (2000, 'host-b', 2.0)") + .execute(&pool) + .await + .unwrap(); + + // sqlx::query uses the binary prepared statement protocol + // (COM_STMT_PREPARE + COM_STMT_EXECUTE). + let rows = sqlx::query("TQL EVAL (0, 3000, '1s') demo_metrics") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(3, rows[0].columns().len()); + + let rows = sqlx::query("TQL ANALYZE (0, 3000, '1s') demo_metrics") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + + // SHOW statements share the same describe path; prepared SHOW FULL + // TABLES must report both columns. + let rows = sqlx::query("SHOW TABLES").fetch_all(&pool).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + + let rows = sqlx::query("SHOW FULL TABLES") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + + let rows = sqlx::query("SHOW DATABASES") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + + let _ = fe_mysql_server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_postgres_array_types(store_type: StorageType) { let (mut guard, fe_pg_server) = setup_pg_server(store_type, "test_postgres_array_types").await; let addr = fe_pg_server.bind_addr().unwrap().to_string();