mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-12 06:52:34 +00:00
feat(datanode): execute nested MergeScan plans on the datanode (#9206)
* feat(datanode): execute nested MergeScan plans on the datanode Give a datanode the ability to decode and execute a plan whose MergeScan carries a nested plan (e.g. an inner MergeScan over a build table), so a distributed sub-plan pushed to a datanode can itself fan out to other datanodes. This is the execution/codec capability only; it is not wired to any user-facing switch and produces no plan by itself. Codec (src/query/src/query_engine/default_serializer.rs): - MergeScanAwareSerializer keeps the request's SessionState and the engine catalog manager, and decodes a MergeScan payload's tables by (catalog, schema, table) name through the engine catalog rather than the region-aware request catalog list (which binds every table to the request's region and breaks on mismatched columns). Datanode (src/datanode): - New DatanodeRegionQueryHandler (region_query.rs) serves region queries for the MergeScan nodes of a received plan, resolving region leaders via PartitionRuleManager and dispatching to peer datanodes. - Wire KvBackendCatalogManager + PartitionRuleManager + NodeClients into the datanode's catalog/cache registry (with_dist_planner stays false: a datanode never rewrites its own plan). Bug fixes uncovered while wiring this up: - cache registry: register the table/table_info/table_name/table_route/ partition_info caches the datanode's catalog manager and partition rule manager require (previously missing -> panic on lookup). - default_serializer: decode_sub_plan bridge no longer drives a future on a current-thread runtime from within another executor (spawn a dedicated thread + runtime when not on a multi-thread runtime). Tested by query_engine::default_serializer unit tests: nested MergeScan round-trip and engine-catalog resolution with mismatched column names. No user-facing behavior change. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): keep GreptimeDB function bindings after rebuilding decode state SessionStateBuilder::new_from_existing collects the function registry (HashMap) into a Vec and build() re-registers each entry; register_udf registers aliases before the function's own name and the Vec order is unstable. Rebuilding the session state to install MergeScanAwareSerializer after the GreptimeDB functions were registered could therefore rebind a name/alias (e.g. date_format, an alias of DataFusion's to_char) back to the built-in function. This path runs for every decoded plan, not only nested MergeScan ones. Extract the GreptimeDB function registration into register_greptime_functions and re-apply it after the final build of both the top-level decode state and the MergeScan payload state. Add regression tests asserting date_format stays bound to the GreptimeDB implementation on both paths. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(datanode): address review on nested MergeScan capability Three follow-ups from review of the nested MergeScan execution capability: 1. Layer the datanode caches (P2b). table_cache derives from table_info_cache + table_name_cache, and partition_info_cache derives from table_route_cache, but they were all registered in a single flat CacheRegistry whose invalidate() runs concurrently with no ordering. A derived cache could be invalidated and refilled from a base cache that had not been invalidated yet, leaving stale metadata. Build the base caches and the derived caches as separate layers of a LayeredCacheRegistry (reusing the existing layering mechanism) so base caches finish invalidating before derived ones. Both the production path (cmd) and the test scaffolding now build the same layered registry; the flat build_datanode_cache_registry is removed. 2. Don't let a nested MergeScan's internal region query compete for the concurrency permit (P2a). An outer query holds its permit until its stream is fully consumed, and a nested MergeScan makes the datanode issue region queries to peer datanodes as execution stages of that same query. If such an inner stage also took a permit, the outer query waiting on it would deadlock against its own held permit (or time out when max_concurrent_queries > 0). Mark inner DN->DN region queries as internal stages: QueryRequest gains an flag carried over the wire in the request header's query_context extensions (reserved key, not injectable via hints), and the region server skips acquiring a concurrency permit for internal-stage requests on both acquire sites. Outer queries still acquire normally. 3. Add an end-to-end capability test (tests-integration/.../nested_merge_scan_capability_test.rs): hand-build Join(<local probe region>, MergeScan(build regions on both datanodes)) with mismatched column names and duplicate build keys, ship it to a datanode over a real TCP Flight endpoint, consume the stream fully and compare the row multiset against the frontend's reference join result, assert a real cross-datanode DoGet via an RPC counter, and assert an inner region failure fails the whole query. Also migrate the with_real_datanode_grpc_addr + DatanodeRpcStats test scaffolding. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): adapt date_format regression test to the updated DataFusion fork ScalarUDFImpl's trait bounds changed so its impl can no longer be downcast via as_any(); assert the function equals the GreptimeDB implementation instead, which to_char can never satisfy. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(datanode): cover the internal stage under the concurrency limit Add an end-to-end regression test for the nested MergeScan capability under max_concurrent_queries = 1: both datanodes' only permit is held by the outer queries' streams while their inner MergeScans fetch the build regions from the same datanodes over real Flight connections. The result consumption is bounded by a timeout, so a broken link anywhere in the chain (marking the inner request internal, carrying the marker in the request header, honoring it at the receiving acquire sites) fails the test instead of deadlocking. Also add a minimal with_datanode_options_override hook to the cluster builder, and fix the test module doc: the outer request is handed to RegionServer::handle_remote_read in process; only the inner region queries go over the wire. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(datanode): cover cold route loads under unrelated invalidations The table route cache uses InitStrategy::VersionChecked, whose version counter belongs to the whole CacheContainer rather than to a single table: any TableId invalidation bumps the shared version, so a cold load of one table's route retries when an unrelated table is invalidated concurrently. Document this as a known limitation on the region query handler, and pin the behavior with an integration test: run the nested plans while a background task invalidates an unrelated table id on every datanode, assert the results still match the expected multiset (retries only add work, they never return a stale route), and report the observed load amplification via the cache miss counter (2 loads without invalidations, 8 over three runs with 24 unrelated invalidations in this run). Also expose the datanode cache registries on the test cluster so a test can drive the local invalidation path that the heartbeat handler drives after a metasrv broadcast. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(query): trim nested plan decoder commentary Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor: remove redundant nested scan support plumbing Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): decode nested MergeScan payloads asynchronously Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): avoid unnecessary extension name allocation Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): retain extension decoding for manual session states Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(datanode): keep nested scan query admission unchanged Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(datanode): honor query handler factories and trim nested scan fixtures Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(datanode): retain required cache error import Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(query): reuse registered state for nested scan decoding Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): verify payload information schema table functions Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): retain nested scan runtime and cold-route guards Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(query): address nested scan review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(session): preserve read preference across region query hops Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(session): verify strict context rejection at Flow ingress Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): remove redundant serializer error case Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
Generated
+4
@@ -4473,6 +4473,8 @@ dependencies = [
|
||||
"async-trait",
|
||||
"bytes",
|
||||
"cache",
|
||||
"catalog",
|
||||
"client",
|
||||
"common-base",
|
||||
"common-catalog",
|
||||
"common-config",
|
||||
@@ -4480,6 +4482,7 @@ dependencies = [
|
||||
"common-error",
|
||||
"common-function",
|
||||
"common-greptimedb-telemetry",
|
||||
"common-grpc",
|
||||
"common-macro",
|
||||
"common-meta",
|
||||
"common-options",
|
||||
@@ -13663,6 +13666,7 @@ dependencies = [
|
||||
"datafusion-common",
|
||||
"derive_builder 0.20.2",
|
||||
"derive_more",
|
||||
"prost 0.14.1",
|
||||
"snafu 0.8.6",
|
||||
"sql",
|
||||
"uuid",
|
||||
|
||||
Vendored
+56
-9
@@ -20,8 +20,8 @@ use std::time::Duration;
|
||||
|
||||
use catalog::kvbackend::new_table_cache;
|
||||
use common_meta::cache::{
|
||||
CacheRegistry, CacheRegistryBuilder, LayeredCacheRegistryBuilder, new_schema_cache,
|
||||
new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache,
|
||||
CacheRegistry, CacheRegistryBuilder, LayeredCacheRegistry, LayeredCacheRegistryBuilder,
|
||||
new_schema_cache, new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache,
|
||||
new_table_route_cache, new_table_schema_cache, new_view_info_cache,
|
||||
};
|
||||
use common_meta::kv_backend::KvBackendRef;
|
||||
@@ -53,11 +53,11 @@ pub const TABLE_FLOWNODE_SET_CACHE_NAME: &str = "table_flownode_set_cache";
|
||||
pub const TABLE_ROUTE_CACHE_NAME: &str = "table_route_cache";
|
||||
pub const PARTITION_INFO_CACHE_NAME: &str = "partition_info_cache";
|
||||
|
||||
/// Builds cache registry for datanode, including:
|
||||
/// - Schema cache.
|
||||
/// - Table id to schema name cache.
|
||||
pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry {
|
||||
// Builds table id schema name cache that never expires.
|
||||
/// Builds the datanode cache registry with base metadata caches before derived caches.
|
||||
///
|
||||
/// Invalidating the base layer first prevents a concurrent query from refilling a derived cache
|
||||
/// from stale metadata after that derived cache has been invalidated.
|
||||
pub fn build_datanode_layered_cache_registry(kv_backend: KvBackendRef) -> LayeredCacheRegistry {
|
||||
let cache = CacheBuilder::new(DEFAULT_CACHE_MAX_CAPACITY).build();
|
||||
let table_id_schema_cache = Arc::new(new_table_schema_cache(
|
||||
TABLE_SCHEMA_NAME_CACHE_NAME.to_string(),
|
||||
@@ -65,7 +65,6 @@ pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry
|
||||
kv_backend.clone(),
|
||||
));
|
||||
|
||||
// Builds schema cache
|
||||
let cache = default_cache();
|
||||
let schema_cache = Arc::new(new_schema_cache(
|
||||
SCHEMA_CACHE_NAME.to_string(),
|
||||
@@ -73,9 +72,57 @@ pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry
|
||||
kv_backend.clone(),
|
||||
));
|
||||
|
||||
CacheRegistryBuilder::default()
|
||||
let cache = default_cache();
|
||||
let table_info_cache = Arc::new(new_table_info_cache(
|
||||
TABLE_INFO_CACHE_NAME.to_string(),
|
||||
cache,
|
||||
kv_backend.clone(),
|
||||
));
|
||||
|
||||
let cache = default_cache();
|
||||
let table_name_cache = Arc::new(new_table_name_cache(
|
||||
TABLE_NAME_CACHE_NAME.to_string(),
|
||||
cache,
|
||||
kv_backend.clone(),
|
||||
));
|
||||
|
||||
let cache = default_cache();
|
||||
let table_route_cache = Arc::new(new_table_route_cache(
|
||||
TABLE_ROUTE_CACHE_NAME.to_string(),
|
||||
cache,
|
||||
kv_backend.clone(),
|
||||
));
|
||||
|
||||
let cache = default_cache();
|
||||
let table_cache = Arc::new(new_table_cache(
|
||||
TABLE_CACHE_NAME.to_string(),
|
||||
cache,
|
||||
table_info_cache.clone(),
|
||||
table_name_cache.clone(),
|
||||
));
|
||||
|
||||
let cache = default_cache();
|
||||
let partition_info_cache = Arc::new(new_partition_info_cache(
|
||||
PARTITION_INFO_CACHE_NAME.to_string(),
|
||||
cache,
|
||||
table_route_cache.clone(),
|
||||
));
|
||||
|
||||
let base_registry = CacheRegistryBuilder::default()
|
||||
.add_cache(table_id_schema_cache)
|
||||
.add_cache(schema_cache)
|
||||
.add_cache(table_info_cache)
|
||||
.add_cache(table_name_cache)
|
||||
.add_cache(table_route_cache)
|
||||
.build();
|
||||
let derived_registry = CacheRegistryBuilder::default()
|
||||
.add_cache(table_cache)
|
||||
.add_cache(partition_info_cache)
|
||||
.build();
|
||||
|
||||
LayeredCacheRegistryBuilder::default()
|
||||
.add_cache_registry(base_registry)
|
||||
.add_cache_registry(derived_registry)
|
||||
.build()
|
||||
}
|
||||
|
||||
|
||||
@@ -14,10 +14,9 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use cache::build_datanode_cache_registry;
|
||||
use cache::build_datanode_layered_cache_registry;
|
||||
use catalog::kvbackend::new_read_only_meta_kv_backend;
|
||||
use common_base::Plugins;
|
||||
use common_meta::cache::LayeredCacheRegistryBuilder;
|
||||
use common_telemetry::info;
|
||||
use common_version::{short_version, verbose_version};
|
||||
use datanode::datanode::DatanodeBuilder;
|
||||
@@ -111,11 +110,9 @@ impl InstanceBuilder {
|
||||
let backend = new_read_only_meta_kv_backend(client.clone());
|
||||
let mut builder = DatanodeBuilder::new(dn_opts.clone(), plugins.clone(), backend.clone());
|
||||
|
||||
let registry = Arc::new(
|
||||
LayeredCacheRegistryBuilder::default()
|
||||
.add_cache_registry(build_datanode_cache_registry(backend))
|
||||
.build(),
|
||||
);
|
||||
// The derived caches (table/partition caches) are put into the layer after the caches
|
||||
// they are derived from, so they are invalidated only after them.
|
||||
let registry = Arc::new(build_datanode_layered_cache_registry(backend));
|
||||
builder
|
||||
.with_cache_registry(registry)
|
||||
.with_meta_client(client.clone());
|
||||
|
||||
@@ -25,7 +25,9 @@ use datafusion::execution::runtime_env::RuntimeEnv;
|
||||
use datafusion::optimizer::AnalyzerRule;
|
||||
use datafusion::prelude::SessionConfig;
|
||||
use datafusion_expr::LogicalPlan;
|
||||
use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
|
||||
use datafusion_substrait::logical_plan::consumer::{
|
||||
SubstraitConsumer, from_substrait_plan, from_substrait_plan_with_consumer,
|
||||
};
|
||||
use datafusion_substrait::logical_plan::producer::to_substrait_plan;
|
||||
use datafusion_substrait::substrait::proto::Plan;
|
||||
use prost::Message;
|
||||
@@ -74,6 +76,21 @@ impl SubstraitPlan for DFLogicalSubstraitConvertor {
|
||||
}
|
||||
|
||||
impl DFLogicalSubstraitConvertor {
|
||||
/// Decode with a custom consumer, retaining Greptime's state-aggregate ordering correction.
|
||||
pub async fn decode_with_consumer(
|
||||
&self,
|
||||
plan: &Plan,
|
||||
state: &SessionState,
|
||||
consumer: &impl SubstraitConsumer,
|
||||
) -> Result<LogicalPlan, Error> {
|
||||
let df_plan = from_substrait_plan_with_consumer(consumer, plan)
|
||||
.await
|
||||
.context(DecodeDfPlanSnafu)?;
|
||||
FixStateUdafOrderingAnalyzer {}
|
||||
.analyze(df_plan, state.config_options())
|
||||
.context(DecodeDfPlanSnafu)
|
||||
}
|
||||
|
||||
pub fn to_sub_plan(
|
||||
&self,
|
||||
plan: &LogicalPlan,
|
||||
|
||||
@@ -20,6 +20,7 @@ use async_trait::async_trait;
|
||||
use bytes::{Buf, Bytes};
|
||||
use datafusion::execution::context::SessionState;
|
||||
pub use datafusion::execution::registry::SerializerRegistry;
|
||||
pub use datafusion_substrait::extensions::Extensions;
|
||||
/// Re-export the Substrait module of datafusion,
|
||||
pub use datafusion_substrait::substrait as substrait_proto_df;
|
||||
pub use datafusion_substrait::{logical_plan as df_logical_plan, variation_const};
|
||||
|
||||
@@ -16,6 +16,8 @@ api.workspace = true
|
||||
arrow-flight.workspace = true
|
||||
async-trait.workspace = true
|
||||
bytes.workspace = true
|
||||
catalog.workspace = true
|
||||
client.workspace = true
|
||||
common-base.workspace = true
|
||||
common-catalog.workspace = true
|
||||
common-config.workspace = true
|
||||
@@ -23,6 +25,7 @@ common-datasource.workspace = true
|
||||
common-error.workspace = true
|
||||
common-function.workspace = true
|
||||
common-greptimedb-telemetry.workspace = true
|
||||
common-grpc.workspace = true
|
||||
common-macro.workspace = true
|
||||
common-meta.workspace = true
|
||||
common-options.workspace = true
|
||||
@@ -54,6 +57,7 @@ metric-engine.workspace = true
|
||||
mito2.workspace = true
|
||||
num_cpus.workspace = true
|
||||
object-store.workspace = true
|
||||
partition.workspace = true
|
||||
prometheus.workspace = true
|
||||
prost.workspace = true
|
||||
query.workspace = true
|
||||
@@ -78,5 +82,4 @@ common-query.workspace = true
|
||||
common-test-util.workspace = true
|
||||
datafusion-common.workspace = true
|
||||
mito2 = { workspace = true, features = ["test"] }
|
||||
partition.workspace = true
|
||||
session.workspace = true
|
||||
|
||||
@@ -19,17 +19,22 @@ use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use catalog::information_schema::NoopInformationExtension;
|
||||
use catalog::kvbackend::KvBackendCatalogManagerBuilder;
|
||||
use client::client_manager::NodeClients;
|
||||
use common_base::Plugins;
|
||||
use common_catalog::consts::{FILE_ENGINE, METRIC_ENGINE, MITO_ENGINE};
|
||||
use common_datasource::object_store::LocalFileAccess;
|
||||
use common_error::ext::BoxedError;
|
||||
use common_greptimedb_telemetry::GreptimeDBTelemetryTask;
|
||||
use common_grpc::channel_manager::ChannelConfig;
|
||||
use common_meta::cache::{LayeredCacheRegistry, SchemaCacheRef, TableSchemaCacheRef};
|
||||
use common_meta::cache_invalidator::CacheInvalidatorRef;
|
||||
use common_meta::datanode::TopicStatsReporter;
|
||||
use common_meta::key::runtime_switch::RuntimeSwitchManager;
|
||||
use common_meta::key::{SchemaMetadataManager, SchemaMetadataManagerRef};
|
||||
use common_meta::kv_backend::KvBackendRef;
|
||||
use common_meta::node_manager::NodeManagerRef;
|
||||
pub use common_procedure::options::ProcedureConfig;
|
||||
use common_query::prelude::set_default_prefix;
|
||||
use common_stat::ResourceStatImpl;
|
||||
@@ -54,7 +59,8 @@ use object_store::ObjectStore;
|
||||
use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef};
|
||||
use object_store::util::normalize_dir;
|
||||
use query::QueryEngineFactory;
|
||||
use query::dummy_catalog::{DummyCatalogManager, TableProviderFactoryRef};
|
||||
use query::dummy_catalog::TableProviderFactoryRef;
|
||||
use query::region_query::RegionQueryHandlerFactoryRef;
|
||||
use servers::server::ServerHandlers;
|
||||
use snafu::{OptionExt, ResultExt, ensure};
|
||||
use store_api::logstore::LogStore;
|
||||
@@ -80,6 +86,7 @@ use crate::event_listener::{
|
||||
use crate::greptimedb_telemetry::get_greptimedb_telemetry_task;
|
||||
use crate::heartbeat::HeartbeatTask;
|
||||
use crate::partition_expr_fetcher::MetaPartitionExprFetcher;
|
||||
use crate::region_query::DatanodeRegionQueryHandler;
|
||||
use crate::region_server::{DummyTableProviderFactory, RegionServer};
|
||||
use crate::store::{self, new_object_store_without_cache};
|
||||
use crate::utils::{RegionOpenRequests, build_region_open_requests};
|
||||
@@ -357,6 +364,7 @@ impl DatanodeBuilder {
|
||||
schema_metadata_manager,
|
||||
region_event_listener,
|
||||
file_ref_manager,
|
||||
cache_registry.clone(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -497,14 +505,41 @@ impl DatanodeBuilder {
|
||||
schema_metadata_manager: SchemaMetadataManagerRef,
|
||||
event_listener: RegionServerEventListenerRef,
|
||||
file_ref_manager: FileReferenceManagerRef,
|
||||
cache_registry: Arc<LayeredCacheRegistry>,
|
||||
) -> Result<RegionServer> {
|
||||
let opts: &DatanodeOptions = &self.opts;
|
||||
|
||||
// Keep `with_dist_planner` false: datanodes execute received MergeScan plans.
|
||||
let catalog_manager = KvBackendCatalogManagerBuilder::new(
|
||||
Arc::new(NoopInformationExtension),
|
||||
self.kv_backend.clone(),
|
||||
cache_registry.clone(),
|
||||
)
|
||||
.build();
|
||||
let partition_manager = catalog_manager.partition_manager();
|
||||
|
||||
// MergeScan region requests stream until completion, so they have no client timeout.
|
||||
let mut channel_config = ChannelConfig {
|
||||
timeout: None,
|
||||
..Default::default()
|
||||
};
|
||||
if opts.grpc.flight_compression.transport_compression() {
|
||||
channel_config.accept_compression = true;
|
||||
channel_config.send_compression = true;
|
||||
}
|
||||
let node_manager: NodeManagerRef = Arc::new(NodeClients::new(channel_config));
|
||||
let region_query_handler =
|
||||
if let Some(factory) = self.plugins.get::<RegionQueryHandlerFactoryRef>() {
|
||||
factory.build(partition_manager.clone(), node_manager)
|
||||
} else {
|
||||
DatanodeRegionQueryHandler::arc(partition_manager.clone(), node_manager)
|
||||
};
|
||||
|
||||
let query_engine_factory = QueryEngineFactory::try_new_with_plugins(
|
||||
// query engine in datanode only executes plan with resolved table source.
|
||||
DummyCatalogManager::arc(),
|
||||
None,
|
||||
None,
|
||||
catalog_manager,
|
||||
Some(partition_manager),
|
||||
Some(region_query_handler),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
@@ -1060,17 +1095,17 @@ mod tests {
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use cache::build_datanode_cache_registry;
|
||||
use cache::build_datanode_layered_cache_registry;
|
||||
use common_base::Plugins;
|
||||
use common_base::readable_size::ReadableSize;
|
||||
use common_config::Configurable;
|
||||
use common_error::ext::ErrorExt;
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_meta::cache::LayeredCacheRegistryBuilder;
|
||||
use common_meta::key::RegionRoleSet;
|
||||
use common_meta::key::datanode_table::DatanodeTableManager;
|
||||
use common_meta::kv_backend::KvBackendRef;
|
||||
use common_meta::kv_backend::memory::MemoryKvBackend;
|
||||
use common_meta::node_manager::NodeManagerRef;
|
||||
use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir};
|
||||
use common_wal::config::DatanodeWalConfig;
|
||||
use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION};
|
||||
@@ -1080,6 +1115,8 @@ mod tests {
|
||||
use mito2::engine::MITO_ENGINE_NAME;
|
||||
use object_store::ObjectStore;
|
||||
use object_store::manager::ObjectStoreManager;
|
||||
use partition::manager::PartitionRuleManagerRef;
|
||||
use query::region_query::{RegionQueryHandlerFactory, RegionQueryHandlerFactoryRef};
|
||||
use store_api::logstore::LogStore;
|
||||
use store_api::region_request::RegionRequest;
|
||||
use store_api::storage::RegionId;
|
||||
@@ -1090,6 +1127,7 @@ mod tests {
|
||||
wal_object_store,
|
||||
};
|
||||
use crate::error::{self, Error};
|
||||
use crate::region_query::DatanodeRegionQueryHandler;
|
||||
use crate::tests::{MockRegionEngine, mock_region_server};
|
||||
|
||||
async fn setup_table_datanode(kv: &KvBackendRef) {
|
||||
@@ -1118,11 +1156,8 @@ mod tests {
|
||||
mock_region_server.register_engine(mock_region.clone());
|
||||
|
||||
let kv_backend = Arc::new(MemoryKvBackend::new());
|
||||
let layered_cache_registry = Arc::new(
|
||||
LayeredCacheRegistryBuilder::default()
|
||||
.add_cache_registry(build_datanode_cache_registry(kv_backend.clone()))
|
||||
.build(),
|
||||
);
|
||||
let layered_cache_registry =
|
||||
Arc::new(build_datanode_layered_cache_registry(kv_backend.clone()));
|
||||
|
||||
let mut builder = DatanodeBuilder::new(
|
||||
DatanodeOptions {
|
||||
@@ -1164,11 +1199,8 @@ mod tests {
|
||||
wal: DatanodeWalConfig,
|
||||
kv_backend: KvBackendRef,
|
||||
) -> DatanodeBuilder {
|
||||
let layered_cache_registry = Arc::new(
|
||||
LayeredCacheRegistryBuilder::default()
|
||||
.add_cache_registry(build_datanode_cache_registry(kv_backend.clone()))
|
||||
.build(),
|
||||
);
|
||||
let layered_cache_registry =
|
||||
Arc::new(build_datanode_layered_cache_registry(kv_backend.clone()));
|
||||
let mut opts = DatanodeOptions {
|
||||
node_id: Some(0),
|
||||
wal,
|
||||
@@ -1446,6 +1478,38 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
struct CountingRegionQueryHandlerFactory(Arc<AtomicUsize>);
|
||||
|
||||
impl RegionQueryHandlerFactory for CountingRegionQueryHandlerFactory {
|
||||
fn build(
|
||||
&self,
|
||||
partition_manager: PartitionRuleManagerRef,
|
||||
node_manager: NodeManagerRef,
|
||||
) -> query::region_query::RegionQueryHandlerRef {
|
||||
self.0.fetch_add(1, Ordering::SeqCst);
|
||||
DatanodeRegionQueryHandler::arc(partition_manager, node_manager)
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_build_uses_installed_region_query_handler_factory() {
|
||||
let data_home = create_temp_dir("region-query-handler-factory");
|
||||
let builds = Arc::new(AtomicUsize::new(0));
|
||||
let plugins = Plugins::new();
|
||||
plugins.insert(Arc::new(CountingRegionQueryHandlerFactory(builds.clone()))
|
||||
as RegionQueryHandlerFactoryRef);
|
||||
let mut builder = datanode_builder(
|
||||
data_home.path().to_str().unwrap(),
|
||||
DatanodeWalConfig::default(),
|
||||
Arc::new(MemoryKvBackend::new()),
|
||||
);
|
||||
builder.set_plugins(plugins);
|
||||
|
||||
let mut datanode = builder.build().await.unwrap();
|
||||
assert_eq!(builds.load(Ordering::SeqCst), 1);
|
||||
datanode.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_build_and_shutdown_with_raft_engine_wal() {
|
||||
let data_home = create_temp_dir("raft-engine-wal");
|
||||
|
||||
@@ -145,6 +145,13 @@ pub enum Error {
|
||||
source: Box<log_store::error::Error>,
|
||||
},
|
||||
|
||||
#[snafu(display("Invalid query context: {}", reason))]
|
||||
InvalidQueryContext {
|
||||
reason: String,
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Invalid SQL, error: {}", msg))]
|
||||
InvalidSql { msg: String },
|
||||
|
||||
@@ -485,7 +492,8 @@ impl ErrorExt for Error {
|
||||
|
||||
Delete { source, .. } => source.status_code(),
|
||||
|
||||
InvalidSql { .. }
|
||||
InvalidQueryContext { .. }
|
||||
| InvalidSql { .. }
|
||||
| IllegalPrimaryKeysDef { .. }
|
||||
| MissingTimestampColumn { .. }
|
||||
| SchemaNotFound { .. }
|
||||
|
||||
@@ -22,6 +22,7 @@ pub mod heartbeat;
|
||||
pub mod metrics;
|
||||
mod partition_expr_fetcher;
|
||||
pub mod query_stream;
|
||||
pub mod region_query;
|
||||
pub mod region_server;
|
||||
pub mod service;
|
||||
pub mod store;
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
// 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 datanode's [`RegionQueryHandler`] implementation for querying remote regions in a
|
||||
//! `MergeScan` plan.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate};
|
||||
use async_trait::async_trait;
|
||||
use common_error::ext::BoxedError;
|
||||
use common_meta::node_manager::NodeManagerRef;
|
||||
use common_query::request::QueryRequest;
|
||||
use common_recordbatch::SendableRecordBatchStream;
|
||||
use partition::manager::PartitionRuleManagerRef;
|
||||
use query::error::{RegionQuerySnafu, Result as QueryResult, UnimplementedSnafu};
|
||||
use query::region_query::{RegionQueryHandler, RegionQueryTarget};
|
||||
use session::ReadPreference;
|
||||
use snafu::ResultExt;
|
||||
use store_api::storage::RegionId;
|
||||
|
||||
/// Serves region queries by resolving their leader and forwarding requests to its datanode.
|
||||
pub(crate) struct DatanodeRegionQueryHandler {
|
||||
partition_manager: PartitionRuleManagerRef,
|
||||
node_manager: NodeManagerRef,
|
||||
}
|
||||
|
||||
impl DatanodeRegionQueryHandler {
|
||||
pub fn arc(
|
||||
partition_manager: PartitionRuleManagerRef,
|
||||
node_manager: NodeManagerRef,
|
||||
) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
partition_manager,
|
||||
node_manager,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RegionQueryHandler for DatanodeRegionQueryHandler {
|
||||
async fn select_target(
|
||||
&self,
|
||||
_read_preference: ReadPreference,
|
||||
region_id: RegionId,
|
||||
) -> QueryResult<RegionQueryTarget> {
|
||||
// The leader of a region is the only peer that can serve a region query.
|
||||
//
|
||||
let peer = self
|
||||
.partition_manager
|
||||
.find_region_leader(region_id)
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.context(RegionQuerySnafu)?;
|
||||
|
||||
Ok(RegionQueryTarget::new(peer))
|
||||
}
|
||||
|
||||
async fn do_get(
|
||||
&self,
|
||||
target: &RegionQueryTarget,
|
||||
request: QueryRequest,
|
||||
) -> QueryResult<SendableRecordBatchStream> {
|
||||
self.node_manager
|
||||
.datanode(target.peer())
|
||||
.await
|
||||
.handle_query(request)
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.context(RegionQuerySnafu)
|
||||
}
|
||||
|
||||
/// Datanodes forward received region queries and do not produce remote dynamic filters.
|
||||
async fn handle_remote_dyn_filter_update(
|
||||
&self,
|
||||
_target: &RegionQueryTarget,
|
||||
_query_id: String,
|
||||
_update: RemoteDynFilterUpdate,
|
||||
) -> QueryResult<()> {
|
||||
UnimplementedSnafu {
|
||||
operation: "remote dynamic filter update from a datanode",
|
||||
}
|
||||
.fail()
|
||||
}
|
||||
|
||||
/// See [`RegionQueryHandler::handle_remote_dyn_filter_update`].
|
||||
async fn handle_remote_dyn_filter_unregister(
|
||||
&self,
|
||||
_target: &RegionQueryTarget,
|
||||
_query_id: String,
|
||||
_unregister: RemoteDynFilterUnregister,
|
||||
) -> QueryResult<()> {
|
||||
UnimplementedSnafu {
|
||||
operation: "remote dynamic filter unregister from a datanode",
|
||||
}
|
||||
.fail()
|
||||
}
|
||||
}
|
||||
@@ -354,7 +354,17 @@ impl RegionServer {
|
||||
None
|
||||
};
|
||||
|
||||
let ctx = request.header.as_ref().map(|h| h.into());
|
||||
let ctx = request
|
||||
.header
|
||||
.as_ref()
|
||||
.map(QueryContext::try_from)
|
||||
.transpose()
|
||||
.map_err(|err| {
|
||||
crate::error::InvalidQueryContextSnafu {
|
||||
reason: err.to_string(),
|
||||
}
|
||||
.build()
|
||||
})?;
|
||||
let query_ctx = Arc::new(ctx.unwrap_or_else(|| QueryContextBuilder::default().build()));
|
||||
|
||||
let region_id = request.region_id;
|
||||
@@ -977,7 +987,15 @@ impl FlightCraft for RegionServer {
|
||||
let query_ctx = request
|
||||
.header
|
||||
.as_ref()
|
||||
.map(|h| Arc::new(QueryContext::from(h)))
|
||||
.map(QueryContext::try_from)
|
||||
.transpose()
|
||||
.map_err(|err| {
|
||||
servers_error::InvalidParameterSnafu {
|
||||
reason: err.to_string(),
|
||||
}
|
||||
.build()
|
||||
})?
|
||||
.map(Arc::new)
|
||||
.unwrap_or(QueryContext::arc());
|
||||
|
||||
let region_server = self.clone();
|
||||
|
||||
@@ -826,7 +826,16 @@ impl common_meta::node_manager::Flownode for FlowDualEngine {
|
||||
let query_ctx = request
|
||||
.header
|
||||
.and_then(|h| h.query_context)
|
||||
.map(|ctx| ctx.into());
|
||||
.map(session::context::QueryContext::try_from)
|
||||
.transpose()
|
||||
.map_err(|err| {
|
||||
to_meta_err(snafu::location!())(
|
||||
crate::error::InvalidQuerySnafu {
|
||||
reason: err.to_string(),
|
||||
}
|
||||
.build(),
|
||||
)
|
||||
})?;
|
||||
match request.body {
|
||||
Some(flow_request::Body::Create(CreateRequest {
|
||||
flow_id: Some(task_id),
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -322,7 +322,8 @@ mod tests {
|
||||
use common_time::Timezone;
|
||||
use query::options::FLOW_SCHEDULED_TIME_MILLIS;
|
||||
use session::hints::{
|
||||
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY,
|
||||
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY,
|
||||
REMOTE_QUERY_ID_EXTENSION_KEY,
|
||||
};
|
||||
use snafu::ResultExt;
|
||||
use tonic::Code;
|
||||
@@ -418,6 +419,10 @@ mod tests {
|
||||
vec![
|
||||
("auto_create_table".to_string(), "true".to_string()),
|
||||
("read_preference".to_string(), "leader".to_string()),
|
||||
(
|
||||
READ_PREFERENCE_EXTENSION_KEY.to_string(),
|
||||
"follower".to_string(),
|
||||
),
|
||||
(
|
||||
REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
|
||||
"spoofed".to_string(),
|
||||
@@ -446,6 +451,8 @@ mod tests {
|
||||
ReadPreference::Leader
|
||||
));
|
||||
assert_eq!(query_context.extension("auto_create_table"), Some("true"));
|
||||
assert_eq!(query_context.extension(READ_PREFERENCE_HINT), None);
|
||||
assert_eq!(query_context.extension(READ_PREFERENCE_EXTENSION_KEY), None);
|
||||
assert_ne!(query_context.remote_query_id(), Some("spoofed"));
|
||||
assert!(
|
||||
query_context
|
||||
|
||||
@@ -49,7 +49,8 @@ mod tests {
|
||||
use query::options::FLOW_SCHEDULED_TIME_MILLIS;
|
||||
use session::context::{QueryContextBuilder, generate_remote_query_id};
|
||||
use session::hints::{
|
||||
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY,
|
||||
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY,
|
||||
REMOTE_QUERY_ID_EXTENSION_KEY,
|
||||
};
|
||||
|
||||
use super::apply_hints;
|
||||
@@ -75,6 +76,10 @@ mod tests {
|
||||
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY.to_string(),
|
||||
"spoofed-regs".to_string(),
|
||||
),
|
||||
(
|
||||
READ_PREFERENCE_EXTENSION_KEY.to_string(),
|
||||
"follower".to_string(),
|
||||
),
|
||||
(
|
||||
FLOW_SCHEDULED_TIME_MILLIS.to_string(),
|
||||
"1700000000000".to_string(),
|
||||
@@ -92,6 +97,8 @@ mod tests {
|
||||
.extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)
|
||||
.is_none()
|
||||
);
|
||||
assert_eq!(query_ctx.read_preference(), session::ReadPreference::Leader);
|
||||
assert!(query_ctx.extension(READ_PREFERENCE_EXTENSION_KEY).is_none());
|
||||
assert_eq!(
|
||||
query_ctx.extension(FLOW_SCHEDULED_TIME_MILLIS),
|
||||
Some("1700000000000")
|
||||
|
||||
@@ -27,6 +27,9 @@ snafu.workspace = true
|
||||
sql.workspace = true
|
||||
uuid.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
prost.workspace = true
|
||||
|
||||
[package.metadata.cargo-shear]
|
||||
# `common-error` is required by `common-macro`'s `stack_trace_debug`: the
|
||||
# generated code references `::common_error` paths, so it must stay a direct
|
||||
|
||||
+159
-17
@@ -15,6 +15,7 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::net::SocketAddr;
|
||||
use std::str::FromStr;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -34,12 +35,14 @@ use derive_builder::Builder;
|
||||
use sql::dialect::{Dialect, GenericDialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect};
|
||||
|
||||
pub use crate::hints::{
|
||||
LIVE_ANALYZE_METRICS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY,
|
||||
SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
|
||||
LIVE_ANALYZE_METRICS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY,
|
||||
REMOTE_QUERY_ID_EXTENSION_KEY, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
|
||||
};
|
||||
use crate::protocol_ctx::ProtocolCtx;
|
||||
use crate::query_id::QueryId;
|
||||
use crate::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle};
|
||||
use crate::session_config::{
|
||||
InvalidConfigValueSnafu, PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle,
|
||||
};
|
||||
use crate::{MutableInner, ReadPreference};
|
||||
|
||||
pub type QueryContextRef = Arc<QueryContext>;
|
||||
@@ -180,25 +183,41 @@ impl QueryContextBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&RegionRequestHeader> for QueryContext {
|
||||
fn from(value: &RegionRequestHeader) -> Self {
|
||||
impl TryFrom<&RegionRequestHeader> for QueryContext {
|
||||
type Error = crate::session_config::Error;
|
||||
|
||||
fn try_from(value: &RegionRequestHeader) -> Result<Self, Self::Error> {
|
||||
if let Some(ctx) = &value.query_context {
|
||||
ctx.clone().into()
|
||||
ctx.clone().try_into()
|
||||
} else {
|
||||
QueryContextBuilder::default()
|
||||
Ok(QueryContextBuilder::default()
|
||||
.set_extension(
|
||||
REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
|
||||
generate_remote_query_id(),
|
||||
)
|
||||
.build()
|
||||
.build())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<api::v1::QueryContext> for QueryContext {
|
||||
fn from(ctx: api::v1::QueryContext) -> Self {
|
||||
impl TryFrom<api::v1::QueryContext> for QueryContext {
|
||||
type Error = crate::session_config::Error;
|
||||
|
||||
fn try_from(mut ctx: api::v1::QueryContext) -> Result<Self, Self::Error> {
|
||||
let read_preference = ctx.extensions.remove(READ_PREFERENCE_EXTENSION_KEY);
|
||||
let read_preference = match read_preference {
|
||||
Some(value) => ReadPreference::from_str(&value).map_err(|_| {
|
||||
InvalidConfigValueSnafu {
|
||||
name: "read_preference",
|
||||
value,
|
||||
hint: "Expected a supported read preference",
|
||||
}
|
||||
.build()
|
||||
})?,
|
||||
None => ReadPreference::default(),
|
||||
};
|
||||
let sequences = ctx.snapshot_seqs.as_ref();
|
||||
QueryContextBuilder::default()
|
||||
Ok(QueryContextBuilder::default()
|
||||
.current_catalog(ctx.current_catalog)
|
||||
.current_schema(ctx.current_schema)
|
||||
.timezone(parse_timezone(Some(&ctx.timezone)))
|
||||
@@ -215,7 +234,8 @@ impl From<api::v1::QueryContext> for QueryContext {
|
||||
.unwrap_or_default(),
|
||||
)))
|
||||
.explain_options(ctx.explain)
|
||||
.build()
|
||||
.read_preference(read_preference)
|
||||
.build())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -234,6 +254,17 @@ impl From<QueryContext> for api::v1::QueryContext {
|
||||
) -> Self {
|
||||
let explain = mutable_query_context_data.read().unwrap().explain_options;
|
||||
let mutable_inner = mutable_inner.read().unwrap();
|
||||
let mut extensions = extensions;
|
||||
extensions.remove(READ_PREFERENCE_EXTENSION_KEY);
|
||||
if mutable_inner.read_preference != ReadPreference::Leader {
|
||||
extensions.insert(
|
||||
READ_PREFERENCE_EXTENSION_KEY.to_string(),
|
||||
mutable_inner
|
||||
.read_preference
|
||||
.to_string()
|
||||
.to_ascii_lowercase(),
|
||||
);
|
||||
}
|
||||
api::v1::QueryContext {
|
||||
current_catalog,
|
||||
current_schema: mutable_inner.schema.clone(),
|
||||
@@ -843,7 +874,7 @@ mod test {
|
||||
0
|
||||
);
|
||||
let wire: api::v1::QueryContext = admitted.into();
|
||||
let restored = QueryContext::from(wire);
|
||||
let restored = QueryContext::try_from(wire).unwrap();
|
||||
assert_eq!(restored.write_rows_to_admit("greptime", "public", 100), 100);
|
||||
}
|
||||
|
||||
@@ -856,7 +887,7 @@ mod test {
|
||||
.extensions
|
||||
.contains_key(crate::hints::INSERT_SKIP_WAL_HINT)
|
||||
);
|
||||
let restored: QueryContext = api_context.into();
|
||||
let restored: QueryContext = api_context.try_into().unwrap();
|
||||
assert!(!restored.skip_wal());
|
||||
}
|
||||
|
||||
@@ -875,7 +906,7 @@ mod test {
|
||||
explain: None,
|
||||
};
|
||||
|
||||
let session_ctx: QueryContext = api_ctx.clone().into();
|
||||
let session_ctx: QueryContext = api_ctx.clone().try_into().unwrap();
|
||||
let roundtrip_api: api::v1::QueryContext = session_ctx.into();
|
||||
|
||||
assert_eq!(roundtrip_api.current_catalog, api_ctx.current_catalog);
|
||||
@@ -910,7 +941,7 @@ mod test {
|
||||
assert_eq!(ctx.remote_query_id_value().unwrap().to_string(), query_id);
|
||||
|
||||
let proto: api::v1::QueryContext = (&ctx).into();
|
||||
let restored = QueryContext::from(proto);
|
||||
let restored = QueryContext::try_from(proto).unwrap();
|
||||
assert_eq!(restored.remote_query_id(), Some(query_id));
|
||||
assert_eq!(
|
||||
restored.remote_query_id_value().unwrap().to_string(),
|
||||
@@ -918,6 +949,117 @@ mod test {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_query_context_read_preference_wire_values() {
|
||||
let wire = |value: Option<&str>| api::v1::QueryContext {
|
||||
current_catalog: DEFAULT_CATALOG_NAME.to_string(),
|
||||
current_schema: DEFAULT_SCHEMA_NAME.to_string(),
|
||||
timezone: "UTC".to_string(),
|
||||
extensions: value
|
||||
.map(|value| {
|
||||
HashMap::from([(READ_PREFERENCE_EXTENSION_KEY.to_string(), value.to_string())])
|
||||
})
|
||||
.unwrap_or_default(),
|
||||
channel: Channel::Internal as u32,
|
||||
snapshot_seqs: None,
|
||||
explain: None,
|
||||
};
|
||||
|
||||
let header = RegionRequestHeader::default();
|
||||
let header_context = QueryContext::try_from(&header).unwrap();
|
||||
assert_eq!(header_context.read_preference(), ReadPreference::Leader);
|
||||
|
||||
for value in [None, Some("leader"), Some("LEADER")] {
|
||||
let ctx = QueryContext::try_from(wire(value)).unwrap();
|
||||
assert_eq!(ctx.read_preference(), ReadPreference::Leader);
|
||||
assert!(!ctx.extensions().contains_key(READ_PREFERENCE_EXTENSION_KEY));
|
||||
}
|
||||
for value in ["", "unknown", "follower", "follower_preferred", "random"] {
|
||||
let error = QueryContext::try_from(wire(Some(value))).unwrap_err();
|
||||
assert!(matches!(
|
||||
error,
|
||||
crate::session_config::Error::InvalidConfigValue {
|
||||
name,
|
||||
value: actual,
|
||||
..
|
||||
} if name == "read_preference" && actual == value
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_legacy_query_context_wire_defaults_to_leader() {
|
||||
use prost::Message;
|
||||
|
||||
let legacy = api::v1::QueryContext::decode(
|
||||
include_bytes!(
|
||||
"../../../tests/compatibility/wire/read_preference/legacy_query_context.pb"
|
||||
)
|
||||
.as_slice(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(legacy.current_catalog, "c1");
|
||||
assert_eq!(legacy.current_schema, "s1");
|
||||
assert_eq!(legacy.timezone, "UTC");
|
||||
assert_eq!(legacy.channel, Channel::Internal as u32);
|
||||
assert_eq!(
|
||||
legacy
|
||||
.extensions
|
||||
.get("flow.return_region_seq")
|
||||
.map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
assert!(
|
||||
!legacy
|
||||
.extensions
|
||||
.contains_key(READ_PREFERENCE_EXTENSION_KEY)
|
||||
);
|
||||
|
||||
let restored = QueryContext::try_from(legacy).unwrap();
|
||||
assert_eq!(restored.read_preference(), ReadPreference::Leader);
|
||||
assert_eq!(restored.current_catalog(), "c1");
|
||||
assert_eq!(restored.current_schema(), "s1");
|
||||
assert_eq!(restored.timezone().to_string(), "UTC");
|
||||
assert_eq!(restored.channel(), Channel::Internal);
|
||||
assert_eq!(restored.extension("flow.return_region_seq"), Some("true"));
|
||||
|
||||
let outbound: api::v1::QueryContext = restored.into();
|
||||
let roundtrip = api::v1::QueryContext::decode(outbound.encode_to_vec().as_slice()).unwrap();
|
||||
assert_eq!(roundtrip.current_catalog, "c1");
|
||||
assert_eq!(roundtrip.current_schema, "s1");
|
||||
assert_eq!(roundtrip.timezone, "UTC");
|
||||
assert_eq!(roundtrip.channel, Channel::Internal as u32);
|
||||
assert_eq!(
|
||||
roundtrip
|
||||
.extensions
|
||||
.get("flow.return_region_seq")
|
||||
.map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
assert!(
|
||||
!roundtrip
|
||||
.extensions
|
||||
.contains_key(READ_PREFERENCE_EXTENSION_KEY)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_outbound_read_preference_clears_stale_extension() {
|
||||
let context = QueryContextBuilder::default()
|
||||
.set_extension(
|
||||
READ_PREFERENCE_EXTENSION_KEY.to_string(),
|
||||
"stale".to_string(),
|
||||
)
|
||||
.set_extension("unrelated".to_string(), "preserved".to_string())
|
||||
.build();
|
||||
let wire: api::v1::QueryContext = context.into();
|
||||
assert!(!wire.extensions.contains_key(READ_PREFERENCE_EXTENSION_KEY));
|
||||
assert_eq!(
|
||||
wire.extensions.get("unrelated").map(String::as_str),
|
||||
Some("preserved")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_live_analyze_metrics_requires_matching_remote_query_id() {
|
||||
let mut ctx = QueryContext::arc().as_ref().clone();
|
||||
@@ -944,7 +1086,7 @@ mod test {
|
||||
assert!(ctx.clone().batching_enabled());
|
||||
assert!(ctx.fork().batching_enabled());
|
||||
let wire: api::v1::QueryContext = ctx.into();
|
||||
let restored = QueryContext::from(wire);
|
||||
let restored = QueryContext::try_from(wire).unwrap();
|
||||
assert!(!restored.batching_enabled());
|
||||
assert!(!restored.logical_batching_enabled());
|
||||
let ctx = QueryContextBuilder::default()
|
||||
|
||||
@@ -27,11 +27,13 @@ pub const LIVE_ANALYZE_METRICS_EXTENSION_KEY: &str = "query.live_analyze_metrics
|
||||
pub const INSERT_SKIP_WAL_HINT: &str = "insert_skip_wal";
|
||||
|
||||
pub const READ_PREFERENCE_HINT: &str = "read_preference";
|
||||
pub const RESERVED_EXTENSION_KEYS: [&str; 4] = [
|
||||
pub const READ_PREFERENCE_EXTENSION_KEY: &str = "query.read_preference";
|
||||
pub const RESERVED_EXTENSION_KEYS: [&str; 5] = [
|
||||
REMOTE_QUERY_ID_EXTENSION_KEY,
|
||||
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY,
|
||||
SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
|
||||
LIVE_ANALYZE_METRICS_EXTENSION_KEY,
|
||||
READ_PREFERENCE_EXTENSION_KEY,
|
||||
];
|
||||
|
||||
/// Deprecated, use `HINTS_KEY` instead.
|
||||
@@ -65,6 +67,7 @@ mod tests {
|
||||
assert!(is_reserved_extension_key(
|
||||
LIVE_ANALYZE_METRICS_EXTENSION_KEY
|
||||
));
|
||||
assert!(is_reserved_extension_key(READ_PREFERENCE_EXTENSION_KEY));
|
||||
assert!(!is_reserved_extension_key(READ_PREFERENCE_HINT));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,13 +17,15 @@ use std::env;
|
||||
use std::net::TcpListener;
|
||||
use std::ops::RangeInclusive;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Duration;
|
||||
|
||||
use api::v1::health_check_server::HealthCheckServer;
|
||||
use api::v1::region::region_server::RegionServer;
|
||||
use arrow_flight::flight_service_server::FlightServiceServer;
|
||||
use cache::{
|
||||
build_datanode_cache_registry, build_fundamental_cache_registry,
|
||||
build_datanode_layered_cache_registry, build_fundamental_cache_registry,
|
||||
with_default_composite_cache_registry,
|
||||
};
|
||||
use catalog::information_extension::DistributedInformationExtension;
|
||||
@@ -74,9 +76,11 @@ use servers::grpc::{GrpcOptions, HealthCheckHandler};
|
||||
use servers::server::ServerHandlers;
|
||||
use store_api::storage::RegionId;
|
||||
use tempfile::TempDir;
|
||||
use tokio::io::{AsyncRead, AsyncWrite};
|
||||
use tonic::codec::CompressionEncoding;
|
||||
use tonic::transport::Server;
|
||||
use tower::service_fn;
|
||||
use tonic::transport::server::Connected;
|
||||
use tower::{Layer, Service, service_fn};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::test_util::{
|
||||
@@ -84,11 +88,31 @@ use crate::test_util::{
|
||||
create_tmp_dir_and_datanode_opts,
|
||||
};
|
||||
|
||||
/// Flight `DoGet` requests another node sent to a datanode over its network address (see
|
||||
/// [`GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`]).
|
||||
///
|
||||
/// The frontend and datanodes inside the test process use in-process clients, so counted requests
|
||||
/// come from remote nodes querying regions owned by this datanode.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct DatanodeRpcStats {
|
||||
do_get_requests: AtomicUsize,
|
||||
}
|
||||
|
||||
impl DatanodeRpcStats {
|
||||
/// The number of Flight `DoGet` requests another node sent to the datanode.
|
||||
pub fn requests(&self) -> usize {
|
||||
self.do_get_requests.load(Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct GreptimeDbCluster {
|
||||
pub guards: Vec<TestGuard>,
|
||||
pub datanode_options: Vec<DatanodeOptions>,
|
||||
|
||||
pub datanode_instances: HashMap<DatanodeId, Datanode>,
|
||||
/// The stats of the network address of every datanode, empty unless the cluster serves the
|
||||
/// datanodes at their own addresses.
|
||||
pub datanode_rpc_stats: HashMap<DatanodeId, Arc<DatanodeRpcStats>>,
|
||||
pub kv_backend: KvBackendRef,
|
||||
pub metasrv: Arc<Metasrv>,
|
||||
pub frontend: Arc<Frontend>,
|
||||
@@ -99,6 +123,13 @@ impl GreptimeDbCluster {
|
||||
&self.frontend.instance
|
||||
}
|
||||
|
||||
/// The stats of the network address of the datanode `datanode_id`, if the cluster serves the
|
||||
/// datanodes at their own addresses (see
|
||||
/// [`GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`]).
|
||||
pub fn datanode_rpc_stats(&self, datanode_id: DatanodeId) -> Option<&Arc<DatanodeRpcStats>> {
|
||||
self.datanode_rpc_stats.get(&datanode_id)
|
||||
}
|
||||
|
||||
/// List all SST files from all datanodes.
|
||||
pub async fn list_sst_files_from_all_datanodes(&self) -> BTreeSet<String> {
|
||||
let mut sst_files = BTreeSet::new();
|
||||
@@ -177,6 +208,7 @@ pub struct GreptimeDbClusterBuilder {
|
||||
plugins: Plugins,
|
||||
local_file_access: LocalFileAccess,
|
||||
event_recorder_options: EventRecorderOptions,
|
||||
real_datanode_grpc_addr: bool,
|
||||
}
|
||||
|
||||
impl GreptimeDbClusterBuilder {
|
||||
@@ -214,6 +246,7 @@ impl GreptimeDbClusterBuilder {
|
||||
plugins: Plugins::default(),
|
||||
local_file_access: LocalFileAccess::default(),
|
||||
event_recorder_options: EventRecorderOptions::default(),
|
||||
real_datanode_grpc_addr: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -299,6 +332,21 @@ impl GreptimeDbClusterBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Serves every datanode on its own real localhost TCP port and registers that
|
||||
/// address in metasrv, instead of the default in-process duplex placeholder address
|
||||
/// that all datanodes share.
|
||||
///
|
||||
/// Datanodes serve the region queries of the plans they execute by querying the
|
||||
/// leader of every region of the plan (see `DatanodeRegionQueryHandler`), so a test
|
||||
/// that exercises a datanode-to-datanode region query needs each datanode to be
|
||||
/// reachable at the address recorded in the table route, i.e. the address the
|
||||
/// datanode registers with its heartbeat.
|
||||
#[must_use]
|
||||
pub fn with_real_datanode_grpc_addr(mut self, real_datanode_grpc_addr: bool) -> Self {
|
||||
self.real_datanode_grpc_addr = real_datanode_grpc_addr;
|
||||
self
|
||||
}
|
||||
|
||||
pub async fn build_with(
|
||||
&self,
|
||||
datanode_options: Vec<DatanodeOptions>,
|
||||
@@ -346,7 +394,22 @@ impl GreptimeDbClusterBuilder {
|
||||
.build_datanodes_with_options(&metasrv, &datanode_options)
|
||||
.await;
|
||||
|
||||
build_datanode_clients(datanode_clients.clone(), &datanode_instances).await;
|
||||
let datanode_rpc_stats = if self.real_datanode_grpc_addr {
|
||||
datanode_options
|
||||
.iter()
|
||||
.map(|opts| (opts.node_id.unwrap(), Arc::new(DatanodeRpcStats::default())))
|
||||
.collect::<HashMap<_, _>>()
|
||||
} else {
|
||||
HashMap::new()
|
||||
};
|
||||
|
||||
build_datanode_clients(
|
||||
datanode_clients.clone(),
|
||||
&datanode_instances,
|
||||
&datanode_options,
|
||||
&datanode_rpc_stats,
|
||||
)
|
||||
.await;
|
||||
|
||||
self.wait_datanodes_alive(metasrv.metasrv.meta_peer_client(), datanodes)
|
||||
.await;
|
||||
@@ -365,6 +428,7 @@ impl GreptimeDbClusterBuilder {
|
||||
datanode_options,
|
||||
guards,
|
||||
datanode_instances,
|
||||
datanode_rpc_stats,
|
||||
kv_backend: self.kv_backend.clone(),
|
||||
metasrv: metasrv.metasrv,
|
||||
frontend: Arc::new(frontend),
|
||||
@@ -421,12 +485,28 @@ impl GreptimeDbClusterBuilder {
|
||||
opts
|
||||
};
|
||||
opts.node_id = Some(datanode_id);
|
||||
|
||||
if self.real_datanode_grpc_addr {
|
||||
// Give every datanode its own address, so that the other datanodes can
|
||||
// reach it (see `with_real_datanode_grpc_addr`).
|
||||
let addr = self.choose_datanode_grpc_addr();
|
||||
opts.grpc.bind_addr = addr.clone();
|
||||
opts.grpc.server_addr = addr;
|
||||
}
|
||||
options.push(opts);
|
||||
}
|
||||
(options, guards)
|
||||
}
|
||||
|
||||
/// Chooses a free localhost port for a datanode.
|
||||
///
|
||||
/// Uses a port range disjoint from the one of the frontend servers
|
||||
/// (see [`Self::build_frontend_options`]).
|
||||
fn choose_datanode_grpc_addr(&self) -> String {
|
||||
let localhost = "127.0.0.1";
|
||||
let port = self.choose_random_unused_port(24000..=34000, 20, localhost);
|
||||
format!("{}:{}", localhost, port)
|
||||
}
|
||||
|
||||
async fn build_datanodes_with_options(
|
||||
&self,
|
||||
metasrv: &MockInfo,
|
||||
@@ -436,7 +516,8 @@ impl GreptimeDbClusterBuilder {
|
||||
|
||||
for opts in options {
|
||||
let datanode = self.create_datanode(opts.clone(), metasrv.clone()).await;
|
||||
instances.insert(opts.node_id.unwrap(), datanode);
|
||||
let datanode_id = opts.node_id.unwrap();
|
||||
instances.insert(datanode_id, datanode);
|
||||
}
|
||||
|
||||
instances
|
||||
@@ -474,15 +555,12 @@ impl GreptimeDbClusterBuilder {
|
||||
|
||||
let meta_backend = new_read_only_meta_kv_backend(meta_client.clone());
|
||||
|
||||
let layered_cache_registry = Arc::new(
|
||||
LayeredCacheRegistryBuilder::default()
|
||||
.add_cache_registry(build_datanode_cache_registry(meta_backend.clone()))
|
||||
.build(),
|
||||
);
|
||||
let layered_cache_registry =
|
||||
Arc::new(build_datanode_layered_cache_registry(meta_backend.clone()));
|
||||
|
||||
let mut builder = DatanodeBuilder::new(opts.clone(), Plugins::default(), meta_backend);
|
||||
builder
|
||||
.with_cache_registry(layered_cache_registry)
|
||||
.with_cache_registry(layered_cache_registry.clone())
|
||||
.with_meta_client(meta_client);
|
||||
let mut datanode = builder.build().await.unwrap();
|
||||
|
||||
@@ -630,18 +708,34 @@ impl GreptimeDbClusterBuilder {
|
||||
async fn build_datanode_clients(
|
||||
clients: Arc<NodeClients>,
|
||||
instances: &HashMap<DatanodeId, Datanode>,
|
||||
options: &[DatanodeOptions],
|
||||
rpc_stats: &HashMap<DatanodeId, Arc<DatanodeRpcStats>>,
|
||||
) {
|
||||
for (&datanode_id, instance) in instances {
|
||||
let (addr, client) = create_datanode_client(instance).await;
|
||||
for opts in options {
|
||||
let datanode_id = opts.node_id.unwrap();
|
||||
let instance = &instances[&datanode_id];
|
||||
let (addr, client) = create_datanode_client(
|
||||
instance,
|
||||
&opts.grpc.server_addr,
|
||||
rpc_stats.get(&datanode_id).cloned(),
|
||||
)
|
||||
.await;
|
||||
clients
|
||||
.insert_client(Peer::new(datanode_id, addr), client)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn create_datanode_client(datanode: &Datanode) -> (String, Client) {
|
||||
let (client, server) = tokio::io::duplex(1024);
|
||||
|
||||
/// Runs a gRPC server that serves the region server of `datanode` over `incoming` until the
|
||||
/// cluster is dropped.
|
||||
fn serve_datanode_region_server<I, IO>(
|
||||
datanode: &Datanode,
|
||||
incoming: I,
|
||||
rpc_stats: Option<Arc<DatanodeRpcStats>>,
|
||||
) where
|
||||
I: futures::Stream<Item = Result<IO, std::io::Error>> + Send + 'static,
|
||||
IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
|
||||
{
|
||||
let runtime = RuntimeBuilder::default()
|
||||
.worker_threads(2)
|
||||
.thread_name("grpc-handlers")
|
||||
@@ -652,8 +746,9 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) {
|
||||
let region_server_handler =
|
||||
RegionServerRequestHandler::new(Arc::new(datanode.region_server()), runtime);
|
||||
|
||||
let _handle = tokio::spawn(async move {
|
||||
tokio::spawn(async move {
|
||||
Server::builder()
|
||||
.layer(CountRequestsLayer { rpc_stats })
|
||||
.add_service(
|
||||
FlightServiceServer::new(flight_handler)
|
||||
.accept_compressed(CompressionEncoding::Gzip)
|
||||
@@ -669,12 +764,40 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) {
|
||||
.send_compressed(CompressionEncoding::Zstd),
|
||||
)
|
||||
.add_service(HealthCheckServer::new(HealthCheckHandler))
|
||||
.serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]))
|
||||
.serve_with_incoming(incoming)
|
||||
.await
|
||||
});
|
||||
}
|
||||
|
||||
/// Serves the region server of `datanode` and returns the address and the client of the address.
|
||||
///
|
||||
/// The returned client is always an in-process client: the frontend and the datanodes of the test
|
||||
/// process reach the region server through a duplex stream, no matter which address they use. This
|
||||
/// keeps the in-process traffic out of `rpc_stats`.
|
||||
///
|
||||
/// [`test_util::PEER_PLACEHOLDER_ADDR`] is not a real address, so it cannot be served over the
|
||||
/// network at all. Any other address is a real one: the region server is additionally served on
|
||||
/// its local TCP port (`rpc_stats` counts the requests of the remote nodes that reach it, e.g. the
|
||||
/// region query that a datanode sends for the regions of another datanode).
|
||||
async fn create_datanode_client(
|
||||
datanode: &Datanode,
|
||||
addr: &str,
|
||||
rpc_stats: Option<Arc<DatanodeRpcStats>>,
|
||||
) -> (String, Client) {
|
||||
let (client, server) = tokio::io::duplex(1024);
|
||||
serve_datanode_region_server(
|
||||
datanode,
|
||||
futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]),
|
||||
None,
|
||||
);
|
||||
|
||||
if addr != test_util::PEER_PLACEHOLDER_ADDR {
|
||||
let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
|
||||
let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
|
||||
serve_datanode_region_server(datanode, incoming, rpc_stats);
|
||||
}
|
||||
|
||||
let mut client = Some(client);
|
||||
let addr = test_util::PEER_PLACEHOLDER_ADDR;
|
||||
let channel_manager = ChannelManager::new();
|
||||
channel_manager
|
||||
.reset_with_connector(
|
||||
@@ -699,6 +822,52 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) {
|
||||
)
|
||||
}
|
||||
|
||||
/// Counts Flight `DoGet` requests received by the wrapped server.
|
||||
#[derive(Clone)]
|
||||
struct CountRequests<S> {
|
||||
inner: S,
|
||||
rpc_stats: Option<Arc<DatanodeRpcStats>>,
|
||||
}
|
||||
|
||||
impl<S, ReqBody> Service<http::Request<ReqBody>> for CountRequests<S>
|
||||
where
|
||||
S: Service<http::Request<ReqBody>>,
|
||||
{
|
||||
type Response = S::Response;
|
||||
type Error = S::Error;
|
||||
type Future = S::Future;
|
||||
|
||||
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
self.inner.poll_ready(cx)
|
||||
}
|
||||
|
||||
fn call(&mut self, request: http::Request<ReqBody>) -> Self::Future {
|
||||
if request.uri().path() == "/arrow.flight.protocol.FlightService/DoGet"
|
||||
&& let Some(rpc_stats) = &self.rpc_stats
|
||||
{
|
||||
rpc_stats.do_get_requests.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
self.inner.call(request)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct CountRequestsLayer {
|
||||
rpc_stats: Option<Arc<DatanodeRpcStats>>,
|
||||
}
|
||||
|
||||
impl<S> Layer<S> for CountRequestsLayer {
|
||||
type Service = CountRequests<S>;
|
||||
|
||||
fn layer(&self, inner: S) -> Self::Service {
|
||||
CountRequests {
|
||||
inner,
|
||||
rpc_stats: self.rpc_stats.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Mock connectors are registered on a shared channel manager.
|
||||
#[cfg(test)]
|
||||
#[allow(deprecated)]
|
||||
@@ -993,13 +1162,41 @@ mod tests {
|
||||
|
||||
// Assert: the inactive peer was rejected and the production Flow service handled the RPC.
|
||||
assert_eq!(
|
||||
(vec![active_peer.addr], vec![inactive_peer.addr]),
|
||||
(vec![active_peer.addr.clone()], vec![inactive_peer.addr]),
|
||||
client.peer_addresses_by_state()
|
||||
);
|
||||
assert_eq!(
|
||||
StatusCode::InvalidArguments,
|
||||
response.unwrap_err().status_code()
|
||||
);
|
||||
|
||||
for value in ["", "unknown", "follower"] {
|
||||
let mut query_context =
|
||||
api::v1::QueryContext::from(session::context::QueryContext::arc().as_ref());
|
||||
query_context.extensions.insert(
|
||||
session::hints::READ_PREFERENCE_EXTENSION_KEY.to_string(),
|
||||
value.to_string(),
|
||||
);
|
||||
let error = clients
|
||||
.flownode(&active_peer)
|
||||
.await
|
||||
.handle(FlowRequest {
|
||||
header: Some(api::v1::flow::FlowRequestHeader {
|
||||
tracing_context: Default::default(),
|
||||
query_context: Some(query_context),
|
||||
}),
|
||||
body: None,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(StatusCode::InvalidArguments, error.status_code());
|
||||
let message = error.output_msg();
|
||||
assert!(message.contains("read_preference"), "{message}");
|
||||
if !value.is_empty() {
|
||||
assert!(message.contains(value), "{message}");
|
||||
}
|
||||
}
|
||||
|
||||
flownode.shutdown().await.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,7 +16,7 @@ use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use cache::{
|
||||
build_datanode_cache_registry, build_fundamental_cache_registry,
|
||||
build_datanode_layered_cache_registry, build_fundamental_cache_registry,
|
||||
with_default_composite_cache_registry,
|
||||
};
|
||||
use catalog::information_schema::NoopInformationExtension;
|
||||
@@ -209,11 +209,8 @@ impl GreptimeDbStandaloneBuilder {
|
||||
) -> GreptimeDbStandalone {
|
||||
let plugins = self.plugin.clone().unwrap_or_default();
|
||||
|
||||
let layered_cache_registry = Arc::new(
|
||||
LayeredCacheRegistryBuilder::default()
|
||||
.add_cache_registry(build_datanode_cache_registry(kv_backend.clone()))
|
||||
.build(),
|
||||
);
|
||||
let layered_cache_registry =
|
||||
Arc::new(build_datanode_layered_cache_registry(kv_backend.clone()));
|
||||
|
||||
let mut builder =
|
||||
DatanodeBuilder::new(opts.datanode_options(), plugins.clone(), kv_backend.clone());
|
||||
|
||||
@@ -16,6 +16,7 @@ mod gc;
|
||||
mod instance_kafka_wal_test;
|
||||
mod instance_noop_wal_test;
|
||||
mod instance_test;
|
||||
mod nested_merge_scan_capability_test;
|
||||
mod promql_test;
|
||||
mod reconcile_table;
|
||||
mod remote_dyn_filter_test;
|
||||
|
||||
@@ -0,0 +1,749 @@
|
||||
// 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.
|
||||
|
||||
//! Exercises hand-built nested `MergeScan` plans at datanode execution, not optimizer behavior.
|
||||
//! The outer `handle_remote_read` request is in-process; inner region reads use real TCP Flight.
|
||||
//! The probe scan resolves against the request's region catalog, while the nested payload resolves
|
||||
//! against the engine catalog; distinct column names ensure these scopes cannot be confused.
|
||||
//! Success compares the complete duplicate-sensitive result multiset with frontend SQL and known
|
||||
//! rows. Missing inner tables and closed remote inner regions must each fail the whole query.
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::v1::region::{QueryRequest, RegionRequestHeader};
|
||||
use client::Client;
|
||||
use client::region::RegionRequester;
|
||||
use common_error::ext::ErrorExt;
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_grpc::channel_manager::ChannelManager;
|
||||
use common_query::Output;
|
||||
use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream};
|
||||
use common_telemetry::info;
|
||||
use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder};
|
||||
use datanode::region_server::RegionServer;
|
||||
use frontend::instance::Instance;
|
||||
use query::datafusion::QUERY_PARALLELISM_HINT;
|
||||
use query::dist_plan::MergeScanLogicalPlan;
|
||||
use query::parser::QueryLanguageParser;
|
||||
use query::query_engine::DefaultSerializer;
|
||||
use servers::query_handler::sql::SqlQueryHandler;
|
||||
use session::context::{QueryContext, QueryContextRef};
|
||||
use store_api::region_request::{RegionCloseRequest, RegionRequest};
|
||||
use store_api::storage::RegionId;
|
||||
use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
|
||||
|
||||
use crate::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder};
|
||||
use crate::test_util::execute_sql;
|
||||
|
||||
/// Partitioned probe table.
|
||||
const PROBE_TABLE: &str = "nested_cap_probe";
|
||||
/// Partitioned build table with duplicate join keys.
|
||||
const BUILD_TABLE: &str = "nested_cap_build";
|
||||
/// Missing inner table used to verify plan-resolution failure.
|
||||
const MISSING_INNER_TABLE: &str = "nested_cap_missing_build";
|
||||
|
||||
/// `(a_id, probe_key, probe_v)` fixture rows; `a_id` is the partition key.
|
||||
const PROBE_ROWS: &[(i32, i32, i32)] = &[
|
||||
(1, 10, 100),
|
||||
(2, 20, 200),
|
||||
// The second partition of the probe table.
|
||||
(120, 30, 300),
|
||||
(121, 10, 400),
|
||||
];
|
||||
|
||||
/// `(b_id, build_key, build_v)` fixture rows; key `10` duplicates and `40` is unmatched.
|
||||
const BUILD_ROWS: &[(i32, i32, i32)] = &[
|
||||
(1, 10, 1),
|
||||
(2, 10, 2),
|
||||
(3, 20, 3),
|
||||
// The second partition of the build table.
|
||||
(120, 30, 4),
|
||||
(121, 40, 5),
|
||||
];
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_nested_merge_scan_capability_join_across_datanodes() {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
|
||||
let cluster = build_cluster("test_nested_merge_scan_capability_join_across_datanodes").await;
|
||||
let frontend = cluster.fe_instance().clone();
|
||||
|
||||
prepare_tables(&frontend).await;
|
||||
|
||||
let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await;
|
||||
let build_leaders = region_leaders(&frontend, BUILD_TABLE).await;
|
||||
assert_eq!(
|
||||
2,
|
||||
datanodes(&probe_leaders).len(),
|
||||
"expected the regions of {PROBE_TABLE} to be spread over both datanodes, actual region \
|
||||
leaders: {probe_leaders:?}"
|
||||
);
|
||||
assert_eq!(
|
||||
2,
|
||||
datanodes(&build_leaders).len(),
|
||||
"expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual region \
|
||||
leaders: {build_leaders:?}"
|
||||
);
|
||||
|
||||
let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await;
|
||||
assert_eq!(
|
||||
expected_join_rows(),
|
||||
multiset(table_cells(&reference)),
|
||||
"unexpected result of the join on the frontend:\n{reference}"
|
||||
);
|
||||
|
||||
let query_ctx = query_ctx();
|
||||
let nested_plan = nested_plan(&frontend, &query_ctx).await;
|
||||
let plan = encode(&nested_plan);
|
||||
|
||||
let remote_datanodes = datanodes(&build_leaders);
|
||||
let baseline_requests = remote_datanodes
|
||||
.iter()
|
||||
.map(|datanode| (*datanode, rpc_requests(&cluster, *datanode)))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
|
||||
let rows = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await;
|
||||
|
||||
let actual_rows = multiset(rows);
|
||||
assert_eq!(
|
||||
expected_join_rows(),
|
||||
actual_rows,
|
||||
"the nested plans of the datanodes returned an unexpected multiset of rows"
|
||||
);
|
||||
assert_eq!(
|
||||
multiset(table_cells(&reference)),
|
||||
actual_rows,
|
||||
"frontend result differs:\n{reference}"
|
||||
);
|
||||
assert_cross_datanode_region_query(&cluster, &remote_datanodes, &baseline_requests);
|
||||
}
|
||||
|
||||
/// Internal read-preference wire metadata is rejected before nested execution on both inbound paths.
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_nested_merge_scan_capability_rejects_internal_read_preference() {
|
||||
use common_meta::node_manager::Datanode;
|
||||
use common_query::request::QueryRequest as LogicalQueryRequest;
|
||||
use futures::TryStreamExt;
|
||||
use session::hints::READ_PREFERENCE_EXTENSION_KEY;
|
||||
|
||||
common_telemetry::init_default_ut_logging();
|
||||
|
||||
let cluster = build_cluster("test_nested_merge_scan_internal_read_preference").await;
|
||||
let frontend = cluster.fe_instance().clone();
|
||||
prepare_tables(&frontend).await;
|
||||
|
||||
let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await;
|
||||
let query_ctx = query_ctx();
|
||||
let plan = nested_plan(&frontend, &query_ctx).await;
|
||||
let (region_id, (datanode_id, address)) = probe_leaders.iter().next().unwrap();
|
||||
let target = region_server(&cluster, *datanode_id);
|
||||
// Confirm the same plan completes on the valid default-Leader path before asserting rejection.
|
||||
let valid_rows = run_nested_queries(&cluster, &probe_leaders, &encode(&plan), &query_ctx).await;
|
||||
assert_eq!(expected_join_rows(), multiset(valid_rows));
|
||||
let mut baseline = cluster
|
||||
.datanode_instances
|
||||
.keys()
|
||||
.map(|id| (*id, rpc_requests(&cluster, *id)))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
|
||||
for value in ["", "invalid", "follower"] {
|
||||
let mut header = region_request_header(&query_ctx);
|
||||
header
|
||||
.query_context
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.extensions
|
||||
.insert(READ_PREFERENCE_EXTENSION_KEY.to_string(), value.to_string());
|
||||
let header = Some(header);
|
||||
let error = target
|
||||
.handle_read(LogicalQueryRequest {
|
||||
header: header.clone(),
|
||||
region_id: RegionId::from_u64(*region_id),
|
||||
plan: plan.clone(),
|
||||
})
|
||||
.await
|
||||
.err()
|
||||
.expect("internal read preference must be rejected at logical ingress");
|
||||
assert_eq!(
|
||||
StatusCode::InvalidArguments,
|
||||
error.status_code(),
|
||||
"value={value:?}"
|
||||
);
|
||||
let logical_message = error.output_msg().to_ascii_lowercase();
|
||||
assert!(
|
||||
logical_message.contains("read_preference"),
|
||||
"expected an invalid query-context read-preference error, got: {logical_message}"
|
||||
);
|
||||
if !value.is_empty() {
|
||||
assert!(
|
||||
logical_message.contains(value),
|
||||
"unexpected logical error: {logical_message}"
|
||||
);
|
||||
}
|
||||
for id in cluster.datanode_instances.keys() {
|
||||
assert_eq!(
|
||||
baseline[id],
|
||||
rpc_requests(&cluster, *id),
|
||||
"logical invalid context caused DoGet at datanode {id}"
|
||||
);
|
||||
}
|
||||
|
||||
let client = Client::with_query_and_control_managers(
|
||||
ChannelManager::new(),
|
||||
ChannelManager::new(),
|
||||
[address.as_str()],
|
||||
);
|
||||
let requester = RegionRequester::new(client, false, false);
|
||||
let request = LogicalQueryRequest {
|
||||
header,
|
||||
region_id: RegionId::from_u64(*region_id),
|
||||
plan: plan.clone(),
|
||||
};
|
||||
let flight_message = match Datanode::handle_query(&requester, request).await {
|
||||
Err(error) => {
|
||||
assert_eq!(StatusCode::InvalidArguments, error.status_code());
|
||||
error.output_msg().to_ascii_lowercase()
|
||||
}
|
||||
Ok(stream) => {
|
||||
let error = stream
|
||||
.try_collect::<Vec<RecordBatch>>()
|
||||
.await
|
||||
.expect_err("Flight must reject invalid internal read preference");
|
||||
assert_eq!(StatusCode::InvalidArguments, error.status_code());
|
||||
error.output_msg().to_ascii_lowercase()
|
||||
}
|
||||
};
|
||||
assert!(
|
||||
flight_message.contains("read_preference"),
|
||||
"expected an invalid-parameter read-preference error, got: {flight_message}"
|
||||
);
|
||||
if !value.is_empty() {
|
||||
assert!(
|
||||
flight_message.contains(value),
|
||||
"unexpected Flight error: {flight_message}"
|
||||
);
|
||||
}
|
||||
for id in cluster.datanode_instances.keys() {
|
||||
let expected = baseline[id] + u8::from(*id == *datanode_id) as usize;
|
||||
assert_eq!(
|
||||
expected,
|
||||
rpc_requests(&cluster, *id),
|
||||
"invalid Flight context caused nested DoGet at datanode {id}"
|
||||
);
|
||||
}
|
||||
for id in cluster.datanode_instances.keys() {
|
||||
baseline.insert(*id, rpc_requests(&cluster, *id));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A failing inner region makes the whole query fail: the datanode must not return the rows of the
|
||||
/// inner regions that succeed.
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_nested_merge_scan_capability_inner_failure_fails_query() {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
|
||||
let cluster =
|
||||
build_cluster("test_nested_merge_scan_capability_inner_failure_fails_query").await;
|
||||
let frontend = cluster.fe_instance().clone();
|
||||
|
||||
prepare_tables(&frontend).await;
|
||||
|
||||
let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await;
|
||||
let build_leaders = region_leaders(&frontend, BUILD_TABLE).await;
|
||||
|
||||
let (probe_region, (probe_datanode, _)) = probe_leaders
|
||||
.iter()
|
||||
.find(|(_, (datanode, _))| *datanode == 1)
|
||||
.unwrap_or_else(|| {
|
||||
panic!(
|
||||
"expected a region of {PROBE_TABLE} on the datanode 1, actual region leaders: \
|
||||
{probe_leaders:?}"
|
||||
)
|
||||
});
|
||||
let (build_region, (build_datanode, _)) = build_leaders
|
||||
.iter()
|
||||
.find(|(_, (datanode, _))| *datanode != *probe_datanode)
|
||||
.unwrap_or_else(|| {
|
||||
panic!(
|
||||
"expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual \
|
||||
region leaders: {build_leaders:?}"
|
||||
)
|
||||
});
|
||||
|
||||
// The nested payload resolves against the engine catalog, not the request's region catalog.
|
||||
let query_ctx = query_ctx();
|
||||
let missing_inner_plan =
|
||||
nested_plan_with_inner_table(&frontend, &query_ctx, MISSING_INNER_TABLE).await;
|
||||
let plan = encode(&missing_inner_plan);
|
||||
let error = region_server(&cluster, *probe_datanode)
|
||||
.handle_remote_read(
|
||||
QueryRequest {
|
||||
header: Some(region_request_header(&query_ctx)),
|
||||
region_id: *probe_region,
|
||||
plan,
|
||||
},
|
||||
query_ctx.clone(),
|
||||
)
|
||||
.await
|
||||
.err()
|
||||
.expect("a nested plan whose inner relation does not exist must not be executed");
|
||||
// The top level error of the datanode only names the failing step (`Failed to decode logical
|
||||
// plan`): the reason is the source of it, which names the table that could not be resolved.
|
||||
let chain = error_chain(&error);
|
||||
assert!(
|
||||
chain.contains(MISSING_INNER_TABLE),
|
||||
"expected the datanode to fail while resolving the inner relation {MISSING_INNER_TABLE}, \
|
||||
actual error: {chain}"
|
||||
);
|
||||
|
||||
// A closed remote inner region must fail the whole query, not return a partial result.
|
||||
region_server(&cluster, *build_datanode)
|
||||
.handle_request(
|
||||
RegionId::from_u64(*build_region),
|
||||
RegionRequest::Close(RegionCloseRequest {
|
||||
flush_on_close: false,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.expect("the datanode must close the region of the build table");
|
||||
|
||||
let nested_plan = nested_plan(&frontend, &query_ctx).await;
|
||||
let plan = encode(&nested_plan);
|
||||
let baseline = rpc_requests(&cluster, *build_datanode);
|
||||
let stream = region_server(&cluster, *probe_datanode)
|
||||
.handle_remote_read(
|
||||
QueryRequest {
|
||||
header: Some(region_request_header(&query_ctx)),
|
||||
region_id: *probe_region,
|
||||
plan,
|
||||
},
|
||||
query_ctx.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("the request itself must be accepted by the datanode");
|
||||
let error = collect_error(stream)
|
||||
.await
|
||||
.expect("the query must fail when one of its inner regions fails");
|
||||
info!("the nested plan failed as a whole: {error}");
|
||||
|
||||
// The increased DoGet count proves execution reached the remote region owner.
|
||||
let actual = rpc_requests(&cluster, *build_datanode);
|
||||
assert!(
|
||||
actual > baseline,
|
||||
"expected the failed nested plan to reach datanode {build_datanode} for region {build_region} over Flight DoGet (count {actual}, baseline {baseline})"
|
||||
);
|
||||
}
|
||||
|
||||
async fn build_cluster(test_name: &str) -> GreptimeDbCluster {
|
||||
GreptimeDbClusterBuilder::new(test_name)
|
||||
.await
|
||||
// A datanode serves the inner `MergeScan` of a nested plan by querying the datanodes that
|
||||
// own the regions of the inner table, so every datanode must be reachable at the address it
|
||||
// registers in metasrv, i.e. the address of the table routes. The cluster then also counts
|
||||
// the requests of the remote nodes, see [`assert_cross_datanode_region_query`].
|
||||
.with_real_datanode_grpc_addr(true)
|
||||
.with_datanodes(2)
|
||||
.build(false)
|
||||
.await
|
||||
}
|
||||
|
||||
fn assert_cross_datanode_region_query(
|
||||
cluster: &GreptimeDbCluster,
|
||||
remote_datanodes: &BTreeSet<u64>,
|
||||
baseline_requests: &BTreeMap<u64, usize>,
|
||||
) {
|
||||
assert_eq!(
|
||||
2,
|
||||
remote_datanodes.len(),
|
||||
"build table must span both datanodes"
|
||||
);
|
||||
for datanode in remote_datanodes {
|
||||
let actual = rpc_requests(cluster, *datanode);
|
||||
let baseline = baseline_requests[datanode];
|
||||
assert!(
|
||||
actual > baseline,
|
||||
"expected a nested Flight DoGet at datanode {datanode}, count {actual}, baseline {baseline}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn rpc_requests(cluster: &GreptimeDbCluster, datanode_id: u64) -> usize {
|
||||
cluster
|
||||
.datanode_rpc_stats(datanode_id)
|
||||
.unwrap_or_else(|| panic!("expected the rpc stats of the datanode {datanode_id}"))
|
||||
.requests()
|
||||
}
|
||||
|
||||
fn region_server(cluster: &GreptimeDbCluster, datanode_id: u64) -> RegionServer {
|
||||
cluster
|
||||
.datanode_instances
|
||||
.get(&datanode_id)
|
||||
.unwrap_or_else(|| panic!("expected a datanode {datanode_id}"))
|
||||
.region_server()
|
||||
}
|
||||
|
||||
async fn run_nested_queries(
|
||||
cluster: &GreptimeDbCluster,
|
||||
probe_leaders: &BTreeMap<u64, (u64, String)>,
|
||||
plan: &[u8],
|
||||
query_ctx: &QueryContextRef,
|
||||
) -> Vec<Vec<String>> {
|
||||
let mut rows = Vec::new();
|
||||
for (probe_region, (probe_datanode, _)) in probe_leaders {
|
||||
let stream = region_server(cluster, *probe_datanode)
|
||||
.handle_remote_read(
|
||||
QueryRequest {
|
||||
header: Some(region_request_header(query_ctx)),
|
||||
region_id: *probe_region,
|
||||
plan: plan.to_vec(),
|
||||
},
|
||||
query_ctx.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap_or_else(|e| {
|
||||
panic!(
|
||||
"the datanode {probe_datanode} must execute the nested plan of the region \
|
||||
{probe_region}: {e}"
|
||||
)
|
||||
});
|
||||
let batches = RecordBatches::try_collect(stream)
|
||||
.await
|
||||
.unwrap_or_else(|e| {
|
||||
panic!(
|
||||
"the datanode {probe_datanode} must execute the nested plan of the region \
|
||||
{probe_region} completely: {e}"
|
||||
)
|
||||
});
|
||||
let actual = batches
|
||||
.pretty_print()
|
||||
.expect("the result of the datanode must be printable");
|
||||
rows.extend(table_cells(&actual));
|
||||
}
|
||||
|
||||
rows
|
||||
}
|
||||
|
||||
/// Creates the probe and build tables of the tests and inserts the rows.
|
||||
async fn prepare_tables(frontend: &Arc<Instance>) {
|
||||
execute_sql(
|
||||
frontend,
|
||||
&format!(
|
||||
"CREATE TABLE {PROBE_TABLE} (
|
||||
a_id INT,
|
||||
probe_key INT,
|
||||
probe_v INT,
|
||||
a_ts TIMESTAMP,
|
||||
TIME INDEX (a_ts),
|
||||
PRIMARY KEY(a_id)
|
||||
)
|
||||
PARTITION ON COLUMNS (a_id) (
|
||||
a_id < 100,
|
||||
a_id >= 100
|
||||
)
|
||||
engine=mito"
|
||||
),
|
||||
)
|
||||
.await;
|
||||
execute_sql(
|
||||
frontend,
|
||||
&format!(
|
||||
"CREATE TABLE {BUILD_TABLE} (
|
||||
b_id INT,
|
||||
build_key INT,
|
||||
build_v INT,
|
||||
b_ts TIMESTAMP,
|
||||
TIME INDEX (b_ts),
|
||||
PRIMARY KEY(b_id)
|
||||
)
|
||||
PARTITION ON COLUMNS (b_id) (
|
||||
b_id < 100,
|
||||
b_id >= 100
|
||||
)
|
||||
engine=mito"
|
||||
),
|
||||
)
|
||||
.await;
|
||||
|
||||
execute_sql(
|
||||
frontend,
|
||||
&format!(
|
||||
"INSERT INTO {PROBE_TABLE}(a_id, probe_key, probe_v, a_ts) VALUES {}",
|
||||
PROBE_ROWS
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, (a_id, probe_key, probe_v))| {
|
||||
format!("({a_id}, {probe_key}, {probe_v}, {}000)", i + 1)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
),
|
||||
)
|
||||
.await;
|
||||
execute_sql(
|
||||
frontend,
|
||||
&format!(
|
||||
"INSERT INTO {BUILD_TABLE}(b_id, build_key, build_v, b_ts) VALUES {}",
|
||||
BUILD_ROWS
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, (b_id, build_key, build_v))| {
|
||||
format!("({b_id}, {build_key}, {build_v}, {}000)", i + 1)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
fn encode(plan: &LogicalPlan) -> Vec<u8> {
|
||||
DFLogicalSubstraitConvertor
|
||||
.encode(plan, DefaultSerializer)
|
||||
.expect("the plan must be encodable")
|
||||
.to_vec()
|
||||
}
|
||||
|
||||
/// Plans the hand-built nested join sent in the region request.
|
||||
async fn nested_plan(frontend: &Arc<Instance>, query_ctx: &QueryContextRef) -> LogicalPlan {
|
||||
nested_plan_with_inner_table(frontend, query_ctx, BUILD_TABLE).await
|
||||
}
|
||||
|
||||
async fn nested_plan_with_inner_table(
|
||||
frontend: &Arc<Instance>,
|
||||
query_ctx: &QueryContextRef,
|
||||
inner_table: &str,
|
||||
) -> LogicalPlan {
|
||||
let probe = plan_sql(
|
||||
frontend,
|
||||
&format!("SELECT a_id, probe_key, probe_v FROM {PROBE_TABLE}"),
|
||||
query_ctx,
|
||||
)
|
||||
.await;
|
||||
let build = plan_sql(
|
||||
frontend,
|
||||
&format!("SELECT build_key, build_v FROM {BUILD_TABLE}"),
|
||||
query_ctx,
|
||||
)
|
||||
.await;
|
||||
let build = rename_scans(build, inner_table);
|
||||
|
||||
nested_plan_of(&probe, build)
|
||||
}
|
||||
|
||||
fn nested_plan_of(probe: &LogicalPlan, build: LogicalPlan) -> LogicalPlan {
|
||||
info!("probe side of the nested plan:\n{probe}");
|
||||
info!("build side of the nested plan:\n{build}");
|
||||
|
||||
let inner = MergeScanLogicalPlan::new(build, false, Default::default()).into_logical_plan();
|
||||
let plan = LogicalPlanBuilder::from(probe.clone())
|
||||
.join(
|
||||
inner,
|
||||
JoinType::Inner,
|
||||
(vec!["probe_key"], vec!["build_key"]),
|
||||
None,
|
||||
)
|
||||
.expect("the join keys must be resolvable in the sides of the join")
|
||||
.build()
|
||||
.expect("the nested plan must be a valid logical plan");
|
||||
info!("nested plan sent to the datanode:\n{plan}");
|
||||
|
||||
plan
|
||||
}
|
||||
|
||||
fn rename_scans(plan: LogicalPlan, table_name: &str) -> LogicalPlan {
|
||||
use datafusion::common::TableReference;
|
||||
use datafusion::common::tree_node::{Transformed, TreeNode};
|
||||
|
||||
plan.transform_up(|plan| match plan {
|
||||
LogicalPlan::TableScan(mut scan) => {
|
||||
info!(
|
||||
"renaming the scan of the nested plan from {} to {table_name}",
|
||||
scan.table_name
|
||||
);
|
||||
scan.table_name = TableReference::bare(table_name);
|
||||
Ok(Transformed::yes(LogicalPlan::TableScan(scan)))
|
||||
}
|
||||
other => Ok(Transformed::no(other)),
|
||||
})
|
||||
.expect("the plan must be rewritable")
|
||||
.data
|
||||
}
|
||||
|
||||
async fn plan_sql(frontend: &Arc<Instance>, sql: &str, query_ctx: &QueryContextRef) -> LogicalPlan {
|
||||
let stmt = QueryLanguageParser::parse_sql(sql, query_ctx)
|
||||
.unwrap_or_else(|e| panic!("failed to parse `{sql}`: {e}"));
|
||||
frontend
|
||||
.statement_executor()
|
||||
.plan(&stmt, query_ctx.clone())
|
||||
.await
|
||||
.unwrap_or_else(|e| panic!("failed to plan `{sql}`: {e}"))
|
||||
}
|
||||
|
||||
/// Frontend reference join over the nested plan's output columns.
|
||||
fn join_sql() -> String {
|
||||
format!(
|
||||
"SELECT p.a_id, p.probe_key, p.probe_v, b.build_key, b.build_v
|
||||
FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b ON p.probe_key = b.build_key"
|
||||
)
|
||||
}
|
||||
|
||||
/// Expected complete join result, retaining duplicate rows.
|
||||
fn expected_join_rows() -> Vec<Vec<String>> {
|
||||
let mut rows = Vec::new();
|
||||
for (a_id, probe_key, probe_v) in PROBE_ROWS {
|
||||
for (_, build_key, build_v) in BUILD_ROWS {
|
||||
if probe_key == build_key {
|
||||
rows.push(vec![
|
||||
a_id.to_string(),
|
||||
probe_key.to_string(),
|
||||
probe_v.to_string(),
|
||||
build_key.to_string(),
|
||||
build_v.to_string(),
|
||||
]);
|
||||
}
|
||||
}
|
||||
}
|
||||
multiset(rows)
|
||||
}
|
||||
|
||||
async fn region_leaders(frontend: &Arc<Instance>, table: &str) -> BTreeMap<u64, (u64, String)> {
|
||||
use datatypes::arrow::array::AsArray;
|
||||
use datatypes::arrow::datatypes::UInt64Type;
|
||||
|
||||
let sql = format!(
|
||||
"SELECT region_id, peer_id, peer_addr FROM information_schema.region_peers
|
||||
WHERE table_schema = 'public' AND table_name = '{table}' AND is_leader = 'Yes'
|
||||
ORDER BY region_id"
|
||||
);
|
||||
let batches = output_batches(run_sql(frontend, &sql, query_ctx()).await).await;
|
||||
|
||||
let mut leaders = BTreeMap::new();
|
||||
for batch in batches.iter() {
|
||||
let batch = batch.df_record_batch();
|
||||
let region_ids = batch
|
||||
.column_by_name("region_id")
|
||||
.expect("region_id column")
|
||||
.as_primitive::<UInt64Type>();
|
||||
let peer_ids = batch
|
||||
.column_by_name("peer_id")
|
||||
.expect("peer_id column")
|
||||
.as_primitive::<UInt64Type>();
|
||||
let peer_addrs = batch
|
||||
.column_by_name("peer_addr")
|
||||
.expect("peer_addr column")
|
||||
.as_string::<i32>();
|
||||
for row in 0..batch.num_rows() {
|
||||
leaders.insert(
|
||||
region_ids.value(row),
|
||||
(peer_ids.value(row), peer_addrs.value(row).to_string()),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
leaders
|
||||
}
|
||||
|
||||
fn datanodes(leaders: &BTreeMap<u64, (u64, String)>) -> BTreeSet<u64> {
|
||||
leaders.values().map(|(id, _)| *id).collect()
|
||||
}
|
||||
|
||||
async fn output_batches(output: Output) -> RecordBatches {
|
||||
match output.data {
|
||||
common_query::OutputData::Stream(stream) => RecordBatches::try_collect(stream)
|
||||
.await
|
||||
.expect("the output stream must be collectable"),
|
||||
common_query::OutputData::RecordBatches(batches) => batches,
|
||||
common_query::OutputData::AffectedRows(rows) => {
|
||||
panic!("expected a query output, got {rows} affected rows")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Consumes the full stream to detect errors after any partial rows.
|
||||
async fn collect_error(stream: SendableRecordBatchStream) -> Option<String> {
|
||||
use futures::TryStreamExt;
|
||||
|
||||
let result = stream.try_collect::<Vec<RecordBatch>>().await;
|
||||
match result {
|
||||
Ok(rows) => {
|
||||
info!(
|
||||
"the query returned {} batches instead of failing",
|
||||
rows.len()
|
||||
);
|
||||
None
|
||||
}
|
||||
Err(error) => Some(error_chain(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
fn error_chain(error: &(dyn std::error::Error + 'static)) -> String {
|
||||
let mut messages = vec![error.to_string()];
|
||||
let mut source = error.source();
|
||||
while let Some(inner) = source {
|
||||
messages.push(inner.to_string());
|
||||
source = inner.source();
|
||||
}
|
||||
messages.join(": ")
|
||||
}
|
||||
|
||||
async fn query_pretty(frontend: &Arc<Instance>, sql: &str, query_ctx: QueryContextRef) -> String {
|
||||
run_sql(frontend, sql, query_ctx)
|
||||
.await
|
||||
.data
|
||||
.pretty_print()
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_sql(frontend: &Arc<Instance>, sql: &str, query_ctx: QueryContextRef) -> Output {
|
||||
SqlQueryHandler::do_query(frontend.as_ref(), sql, query_ctx)
|
||||
.await
|
||||
.remove(0)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Uses one target partition to keep the plans stable.
|
||||
fn query_ctx() -> QueryContextRef {
|
||||
let mut query_ctx = QueryContext::with_db_name(None);
|
||||
query_ctx.set_extension(QUERY_PARALLELISM_HINT, "1");
|
||||
Arc::new(query_ctx)
|
||||
}
|
||||
|
||||
fn region_request_header(query_ctx: &QueryContextRef) -> RegionRequestHeader {
|
||||
RegionRequestHeader {
|
||||
query_context: Some(query_ctx.as_ref().into()),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn table_cells(pretty: &str) -> Vec<Vec<String>> {
|
||||
let mut lines = pretty.lines().filter(|line| line.starts_with('|'));
|
||||
// The first `|` line is the header of the table.
|
||||
lines.next();
|
||||
|
||||
lines
|
||||
.map(|line| {
|
||||
line.trim_matches('|')
|
||||
.split('|')
|
||||
.map(|cell| cell.trim().to_string())
|
||||
.collect()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn multiset(mut rows: Vec<Vec<String>>) -> Vec<Vec<String>> {
|
||||
rows.sort();
|
||||
rows
|
||||
}
|
||||
@@ -2,7 +2,18 @@
|
||||
|
||||
Compatibility tests verify that one GreptimeDB version can restart on state written by another version.
|
||||
|
||||
Tests are run via `cargo sqlness compat` and reuse the sqlness-runner infrastructure.
|
||||
SQL restart tests are run via `cargo sqlness compat` and reuse the sqlness-runner infrastructure.
|
||||
|
||||
## Wire fixtures
|
||||
|
||||
Frozen protobuf fixtures are exercised by the Rust tests that own their context
|
||||
conversion, not by SQLness case discovery. The
|
||||
[legacy query-context fixture](wire/read_preference/README.md) checks omitted
|
||||
read-preference defaulting and unrelated metadata preservation:
|
||||
|
||||
```shell
|
||||
cargo nextest run -p session -E 'test(test_legacy_query_context_wire_defaults_to_leader)'
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
# Legacy query-context wire compatibility
|
||||
|
||||
`legacy_query_context.pb` is a generated legacy-shaped protobuf fixture, not a
|
||||
capture from a released binary. It uses the unchanged `greptime.v1.QueryContext`
|
||||
schema at greptime-proto revision
|
||||
`9533078d6fd4bdaddbca0750cb9bf7ad64b2e568` and omits the internal
|
||||
`query.read_preference` extension.
|
||||
|
||||
Logical contents:
|
||||
|
||||
- Catalog `c1`, schema `s1`, timezone `UTC`.
|
||||
- Unrelated extension `flow.return_region_seq = true`.
|
||||
- Channel `Internal` (`255`); no snapshot sequences or explain options.
|
||||
|
||||
Run the real protobuf decode and session conversion regression:
|
||||
|
||||
```shell
|
||||
cargo nextest run -p session -E 'test(test_legacy_query_context_wire_defaults_to_leader)'
|
||||
```
|
||||
|
||||
The test checks that the omitted preference defaults to Leader, that the original
|
||||
fields and unrelated extension survive conversion and re-encoding, and that
|
||||
Leader does not introduce a preference carrier. Existing session tests cover
|
||||
present invalid or unsupported values. This fixture does not demonstrate that
|
||||
an old receiver honors a non-Leader preference sent by an updated node.
|
||||
|
||||
## Reproduce the fixture
|
||||
|
||||
With `PROTO_ROOT` pointing to the `proto` directory of the revision above:
|
||||
|
||||
```shell
|
||||
printf '%s\n' \
|
||||
'current_catalog: "c1"' \
|
||||
'current_schema: "s1"' \
|
||||
'timezone: "UTC"' \
|
||||
'extensions { key: "flow.return_region_seq" value: "true" }' \
|
||||
'channel: 255' |
|
||||
protoc --proto_path="$PROTO_ROOT" --encode=greptime.v1.QueryContext \
|
||||
greptime/v1/common.proto > legacy_query_context.pb
|
||||
```
|
||||
|
||||
Generated with `libprotoc 25.1`. The 48-byte fixture has SHA-256
|
||||
`4127ebbee8cdb1375ef8ad52754be56ed206c12513e5f61c4354c04587cba638`.
|
||||
The test never rewrites it and does not require protobuf map ordering to remain
|
||||
stable.
|
||||
@@ -0,0 +1,3 @@
|
||||
|
||||
c1s1"UTC*
|
||||
flow.return_region_seqtrue0ÿ
|
||||
Reference in New Issue
Block a user