diff --git a/src/client/src/database.rs b/src/client/src/database.rs index 13ff9fc651..6616962838 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -36,7 +36,9 @@ use common_catalog::build_db_string; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_error::ext::BoxedError; use common_grpc::flight::do_put::DoPutResponse; -use common_grpc::flight::{FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage}; +use common_grpc::flight::{ + FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY, +}; use common_query::Output; use common_recordbatch::adapter::RecordBatchMetrics; use common_recordbatch::error::ExternalSnafu; @@ -559,7 +561,28 @@ impl Database { let value = serde_json::to_string(&flow_extensions.to_vec()) .expect("flow extension pairs should serialize"); - let key = AsciiMetadataKey::from_static(FLOW_EXTENSIONS_METADATA_KEY); + Self::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, value) + } + + fn put_snapshot_seqs( + metadata: &mut MetadataMap, + snapshot_seqs: &std::collections::HashMap, + ) -> Result<()> { + if snapshot_seqs.is_empty() { + return Ok(()); + } + + 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 put_metadata_value( + metadata: &mut MetadataMap, + key: &'static str, + value: String, + ) -> Result<()> { + let key = AsciiMetadataKey::from_static(key); let value = AsciiMetadataValue::from_str(&value).context(InvalidTonicMetadataValueSnafu)?; metadata.insert(key, value); Ok(()) @@ -660,7 +683,7 @@ impl Database { let request = Request::Query(QueryRequest { query: Some(Query::Sql(sql.as_ref().to_string())), }); - self.do_get(request, hints, &[]) + self.do_get(request, hints, &[], &Default::default()) .await .map(OutputWithMetrics::into_output) } @@ -683,6 +706,7 @@ impl Database { }, hints, &[], + &Default::default(), ) .await } @@ -695,6 +719,7 @@ impl Database { }, &[], &[], + &Default::default(), ) .await .map(OutputWithMetrics::into_output) @@ -709,9 +734,15 @@ impl Database { request: QueryRequest, hints: &[(&str, &str)], flow_extensions: &[(&str, &str)], + snapshot_seqs: &std::collections::HashMap, ) -> Result { - self.do_get(Request::Query(request), hints, flow_extensions) - .await + self.do_get( + Request::Query(request), + hints, + flow_extensions, + snapshot_seqs, + ) + .await } /// Creates a new table using the provided table expression. @@ -719,7 +750,7 @@ impl Database { let request = Request::Ddl(DdlRequest { expr: Some(DdlExpr::CreateTable(expr)), }); - self.do_get(request, &[], &[]) + self.do_get(request, &[], &[], &Default::default()) .await .map(OutputWithMetrics::into_output) } @@ -729,7 +760,7 @@ impl Database { let request = Request::Ddl(DdlRequest { expr: Some(DdlExpr::AlterTable(expr)), }); - self.do_get(request, &[], &[]) + self.do_get(request, &[], &[], &Default::default()) .await .map(OutputWithMetrics::into_output) } @@ -739,6 +770,7 @@ impl Database { request: Request, hints: &[(&str, &str)], flow_extensions: &[(&str, &str)], + snapshot_seqs: &std::collections::HashMap, ) -> Result { let request = self.to_rpc_request(request); let request = Ticket { @@ -749,6 +781,7 @@ impl Database { let metadata = request.metadata_mut(); Self::put_hints(metadata, hints)?; Self::put_flow_extensions(metadata, flow_extensions)?; + Self::put_snapshot_seqs(metadata, snapshot_seqs)?; let mut client = self.client.make_flight_client(false, false)?; @@ -928,6 +961,25 @@ mod tests { ); } + #[test] + fn test_put_snapshot_seqs_preserves_u64_precision() { + let mut metadata = MetadataMap::new(); + let snapshot_seqs = std::collections::HashMap::from([ + (u64::MAX, u64::MAX - 1), + (9_007_199_254_740_993_u64, 9_007_199_254_740_995_u64), + ]); + + Database::put_snapshot_seqs(&mut metadata, &snapshot_seqs).unwrap(); + + let value = metadata + .get(SNAPSHOT_SEQS_METADATA_KEY) + .unwrap() + .to_str() + .unwrap(); + let decoded: std::collections::HashMap = serde_json::from_str(value).unwrap(); + assert_eq!(decoded, snapshot_seqs); + } + #[test] fn test_flight_ctx() { let mut ctx = FlightContext::default(); diff --git a/src/common/grpc/src/flight.rs b/src/common/grpc/src/flight.rs index c7860f5706..235223559a 100644 --- a/src/common/grpc/src/flight.rs +++ b/src/common/grpc/src/flight.rs @@ -39,6 +39,8 @@ use crate::error::{DecodeFlightDataSnafu, InvalidFlightDataSnafu, Result}; /// Flight metadata key used to carry flow query extensions as JSON pairs. pub const FLOW_EXTENSIONS_METADATA_KEY: &str = "x-greptime-flow-extensions"; +/// Flight metadata key used to carry query snapshot read upper bounds as JSON. +pub const SNAPSHOT_SEQS_METADATA_KEY: &str = "x-greptime-snapshot-seqs"; #[derive(Debug, Clone)] pub enum FlightMessage { diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index a4c5cd8ba1..2f1e4b73b4 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -411,7 +411,12 @@ impl FrontendClient { peer: db.peer.clone(), }); db.database - .query_with_terminal_metrics_and_flow_extensions(request, &hints, extensions) + .query_with_terminal_metrics_and_flow_extensions( + request, + &hints, + extensions, + &Default::default(), + ) .await .map_err(BoxedError::new) .context(ExternalSnafu) diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 782c82b54a..b70b808753 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -132,6 +132,7 @@ use crate::engine::puffin_index::{IndexEntryContext, collect_index_entries_from_ use crate::error::{ IncrementalQueryStaleSnafu, InvalidRequestSnafu, JoinSnafu, MitoManifestInfoSnafu, RecvSnafu, RegionNotFoundSnafu, Result, SerdeJsonSnafu, SerializeColumnMetadataSnafu, + SnapshotFenceStaleSnafu, }; #[cfg(feature = "enterprise")] use crate::extension::BoxedExtensionRangeProviderFactory; @@ -1045,6 +1046,26 @@ impl EngineInner { ); } + if let Some(given_seq) = request.memtable_max_sequence + && !request.skip_sst_files + { + // Explicit snapshot fences that include SST reads are enforceable + // only while the requested upper bound is not older than the + // region's flushed frontier. If H has already been flushed into SST, + // mito cannot apply a memtable-only sequence upper bound to that + // SST scan, so fail and let Flow rebind the fenced repair instead + // of reading rows beyond H. + let min_enforceable_seq = version.flushed_sequence; + ensure!( + given_seq >= min_enforceable_seq, + SnapshotFenceStaleSnafu { + region_id, + given_seq, + min_enforceable_seq, + } + ); + } + // Get cache. let cache_manager = self.workers.cache_manager(); diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index a4a1c847e7..5ba5025734 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -247,9 +247,10 @@ async fn test_scan_with_min_sst_sequence() { } #[tokio::test] -async fn test_full_snapshot_upper_bound_does_not_constrain_sst_rows() { +async fn test_full_snapshot_upper_bound_returns_outdated_after_late_flush() { let mut env = - TestEnv::with_prefix("test_full_snapshot_upper_bound_does_not_constrain_sst_rows").await; + TestEnv::with_prefix("test_full_snapshot_upper_bound_returns_outdated_after_late_flush") + .await; let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -277,7 +278,7 @@ async fn test_full_snapshot_upper_bound_does_not_constrain_sst_rows() { test_util::put_rows(&engine, region_id, second_rows).await; test_util::flush_region(&engine, region_id, None).await; - let scanner = engine + let err = engine .scanner( region_id, ScanRequest { @@ -286,14 +287,14 @@ async fn test_full_snapshot_upper_bound_does_not_constrain_sst_rows() { }, ) .await - .unwrap(); + .err() + .expect("expect stale snapshot fence error"); - let stream = scanner.scan().await.unwrap(); - let batches = RecordBatches::try_collect(stream).await.unwrap(); - let pretty = batches.pretty_print().unwrap(); - - assert!(pretty.contains("1970-01-01T00:00:03")); - assert!(pretty.contains("1970-01-01T00:00:04")); + assert_eq!(StatusCode::RequestOutdated, err.status_code()); + let err_msg = err.to_string(); + assert!(err_msg.contains("STALE_SNAPSHOT_FENCE")); + assert!(err_msg.contains(®ion_id.to_string())); + assert!(err_msg.contains(&format!("given_seq: {snapshot_upper_bound}"))); } #[tokio::test] diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 2278b61669..b4d8342d20 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -235,6 +235,20 @@ pub enum Error { location: Location, }, + #[snafu(display( + "STALE_SNAPSHOT_FENCE: snapshot upper bound stale, region: {}, given_seq: {}, min_enforceable_seq: {}, retry_hint: REBIND_SNAPSHOT_FENCE", + region_id, + given_seq, + min_enforceable_seq + ))] + SnapshotFenceStale { + region_id: RegionId, + given_seq: u64, + min_enforceable_seq: u64, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Old manifest missing for region {}", region_id))] MissingOldManifest { region_id: RegionId, @@ -1356,7 +1370,7 @@ impl ErrorExt for Error { | SerializePartitionExpr { .. } | InvalidSourceAndTargetRegion { .. } => StatusCode::InvalidArguments, - IncrementalQueryStale { .. } => StatusCode::RequestOutdated, + IncrementalQueryStale { .. } | SnapshotFenceStale { .. } => StatusCode::RequestOutdated, RegionMetadataNotFound { .. } | Join { .. } diff --git a/src/query/src/dummy_catalog.rs b/src/query/src/dummy_catalog.rs index c79c8d88ea..eb8644c2cf 100644 --- a/src/query/src/dummy_catalog.rs +++ b/src/query/src/dummy_catalog.rs @@ -407,7 +407,15 @@ fn decide_flow_scan(query_ctx: &QueryContext, region_id: RegionId) -> Result TonicResult>> { let mut hints = hint_headers::extract_hints(request.metadata()); hints.extend(extract_flow_extensions(request.metadata())?); + let snapshot_seqs = extract_snapshot_seqs(request.metadata())?; let ticket = request.into_inner().ticket; let request = GreptimeRequest::decode(ticket.as_ref()).context(error::InvalidFlightTicketSnafu)?; let query_ctx = - create_query_context(Channel::Grpc, request.header.as_ref(), hints.clone())?; + create_query_context(Channel::Grpc, request.header.as_ref(), hints, snapshot_seqs)?; // Validate flow hint syntax at the transport boundary before dispatching the request. // This does not authorize or execute anything; `handle_request()` below still performs // the normal frontend handling and auth checks before query execution. @@ -218,7 +221,9 @@ impl FlightCraft for GreptimeRequestHandler { ); let flight_compression = self.flight_compression; async { - let output = self.handle_request(request, hints).await?; + let output = self + .handle_request_with_query_ctx(request, query_ctx.clone()) + .await?; let stream = to_flight_data_stream( output, TracingContext::from_current_span(), @@ -542,21 +547,30 @@ impl Stream for PutRecordBatchRequestStream { fn extract_flow_extensions( metadata: &tonic::metadata::MetadataMap, ) -> TonicResult> { - let Some(value) = metadata.get(FLOW_EXTENSIONS_METADATA_KEY) else { - return Ok(vec![]); + Ok(extract_json_metadata(metadata, FLOW_EXTENSIONS_METADATA_KEY)?.unwrap_or_default()) +} + +fn extract_snapshot_seqs( + metadata: &tonic::metadata::MetadataMap, +) -> TonicResult> { + Ok(extract_json_metadata(metadata, SNAPSHOT_SEQS_METADATA_KEY)?.unwrap_or_default()) +} + +fn extract_json_metadata( + metadata: &tonic::metadata::MetadataMap, + key: &'static str, +) -> TonicResult> { + let Some(value) = metadata.get(key) else { + return Ok(None); }; - let value = value.to_str().map_err(|e| { - Status::invalid_argument(format!( - "Invalid {FLOW_EXTENSIONS_METADATA_KEY} metadata value: {e}" - )) - })?; + let value = value + .to_str() + .map_err(|e| Status::invalid_argument(format!("Invalid {key} metadata value: {e}")))?; - serde_json::from_str::>(value).map_err(|e| { - Status::invalid_argument(format!( - "Invalid {FLOW_EXTENSIONS_METADATA_KEY} metadata JSON: {e}" - )) - }) + let parsed = serde_json::from_str::(value) + .map_err(|e| Status::invalid_argument(format!("Invalid {key} metadata JSON: {e}")))?; + Ok(Some(parsed)) } fn to_flight_data_stream( diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index 4669ade118..6ad982dc6a 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -14,7 +14,9 @@ //! Handler for Greptime Database service. It's implemented by frontend. +use std::collections::HashMap; use std::str::FromStr; +use std::sync::{Arc, RwLock}; use std::time::Instant; use api::helper::request_type; @@ -74,13 +76,22 @@ impl GreptimeRequestHandler { &self, request: GreptimeRequest, hints: Vec<(String, String)>, + ) -> 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.", })?; let header = request.header.as_ref(); - let query_ctx = create_query_context(Channel::Grpc, header, hints)?; let user_info = context_auth::auth(self.user_provider.clone(), header, &query_ctx).await?; query_ctx.set_current_user(user_info); @@ -181,6 +192,7 @@ pub(crate) fn create_query_context( channel: Channel, header: Option<&RequestHeader>, mut extensions: Vec<(String, String)>, + snapshot_seqs: HashMap, ) -> Result { let (catalog, schema) = header .map(|header| { @@ -214,7 +226,8 @@ pub(crate) fn create_query_context( .current_catalog(catalog) .current_schema(schema) .timezone(timezone) - .channel(channel); + .channel(channel) + .snapshot_seqs(Arc::new(RwLock::new(snapshot_seqs))); if let Some(x) = extensions .iter() @@ -314,8 +327,10 @@ mod tests { "spoofed-regs".to_string(), ), ], + HashMap::from([(7, 88)]), ) .unwrap(); + assert_eq!(query_context.get_snapshot(7), Some(88)); assert_eq!(query_context.current_catalog(), "cat-a-log"); assert_eq!(query_context.current_schema(), DEFAULT_SCHEMA_NAME); assert_eq!( @@ -354,6 +369,7 @@ mod tests { REMOTE_QUERY_ID_EXTENSION_KEY.to_string(), "spoofed-query-id".to_string(), )], + HashMap::new(), ) .unwrap(); diff --git a/src/servers/src/grpc/prom_query_gateway.rs b/src/servers/src/grpc/prom_query_gateway.rs index 35e4c0f712..d4ee0b91d3 100644 --- a/src/servers/src/grpc/prom_query_gateway.rs +++ b/src/servers/src/grpc/prom_query_gateway.rs @@ -80,7 +80,12 @@ impl PrometheusGateway for PrometheusGatewayService { }; let header = inner.header.as_ref(); - let query_ctx = create_query_context(Channel::Promql, header, Default::default())?; + let query_ctx = create_query_context( + Channel::Promql, + header, + Default::default(), + Default::default(), + )?; let user_info = auth(self.user_provider.clone(), header, &query_ctx).await?; query_ctx.set_current_user(user_info); diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index d9bf15a64c..d0765a46cf 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -14,11 +14,13 @@ #[cfg(test)] mod test { + use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Arc; use api::v1::auth_header::AuthScheme; - use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, SemanticType}; + use api::v1::query_request::Query; + use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, QueryRequest, SemanticType}; use arrow_flight::FlightDescriptor; use auth::user_provider_from_option; use client::{Client, Database}; @@ -259,6 +261,90 @@ mod test { ); } + #[tokio::test(flavor = "multi_thread")] + async fn test_distributed_flight_snapshot_seqs_rejects_stale_sst_fence() { + common_telemetry::init_default_ut_logging(); + + let db = GreptimeDbClusterBuilder::new( + "test_distributed_flight_snapshot_seqs_rejects_stale_sst_fence", + ) + .await + .build(false) + .await; + + let runtime = common_runtime::global_runtime().clone(); + let greptime_request_handler = GreptimeRequestHandler::new( + db.frontend.instance.clone(), + user_provider_from_option("static_user_provider:cmd:greptime_user=greptime_pwd").ok(), + Some(runtime.clone()), + FlightCompression::default(), + ); + let mut grpc_server = GrpcServerBuilder::new(GrpcServerConfig::default(), runtime) + .flight_handler(Arc::new(greptime_request_handler)) + .build(); + grpc_server + .start("127.0.0.1:0".parse::().unwrap()) + .await + .unwrap(); + + let client = Client::with_urls(vec![grpc_server.bind_addr().unwrap().to_string()]); + let mut client = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); + client.set_auth(AuthScheme::Basic(Basic { + username: "greptime_user".to_string(), + password: "greptime_pwd".to_string(), + })); + + create_table(&client).await; + test_put_record_batches(&client, create_record_batches(1)).await; + client.sql("admin flush_table('foo')").await.unwrap(); + + let regions = db.list_all_regions().await; + assert_eq!(regions.len(), 1); + let (region_id, region) = regions.into_iter().next().unwrap(); + let stale_snapshot_seq = region.find_committed_sequence(); + + test_put_record_batches(&client, create_record_batches(10)).await; + client.sql("admin flush_table('foo')").await.unwrap(); + assert!( + region.find_committed_sequence() > stale_snapshot_seq, + "second flush should make the first snapshot sequence stale" + ); + + 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)]), + ) + .await + .unwrap(); + + let OutputData::Stream(mut stream) = result.output.data else { + panic!("expected stream output"); + }; + + let mut stale_fence_error = None; + while let Some(batch) = stream.next().await { + if let Err(err) = batch { + stale_fence_error = Some(format!("{err:?}")); + break; + } + } + + let err_msg = stale_fence_error.expect("expected stale snapshot fence rejection"); + assert!( + err_msg.contains("STALE_SNAPSHOT_FENCE") + || err_msg.contains("RequestOutdated") + || err_msg.contains("snapshot upper bound stale"), + "expected stale snapshot fence rejection, got: {err_msg}" + ); + } + async fn test_put_record_batches(client: &Database, record_batches: Vec) { let requests_count = record_batches.len(); let schema = record_batches[0].schema.arrow_schema().clone();