mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-11-06 19:06:58 +00:00
Lei, HUANG
557843dd68
* fix(catalog): use advertised peer address as process identity The process manager identified the local frontend by its public gRPC address, while metasrv discovers peers by the address advertised in the heartbeat (the internal gRPC address when configured). With separate public and internal listeners, list_all_processes failed to exclude the local node, queried its own management service as a remote frontend, and appended the local registry again, so every local query was listed twice. The same mismatch broke routing of KILL to other frontends. Identify a frontend by the address it advertises to metasrv: - add FrontendOptions::peer_addr, shared by the heartbeat and the process manager, and FrontendOptions::create_process_manager to wire it once; - use the advertised address as the process identity, for self-exclusion in list_all_processes, and for local/remote routing in kill_process. Process ids now carry the address other frontends use to reach the node. When no internal gRPC listener is configured this is the public gRPC address, so the previous behavior is unchanged. The internal listener is the only one serving the frontend management service, so process ids stay routable. The new distributed regression tests exercise separate public and internal addresses: the local query must be listed exactly once while remote frontends are still listed, and KILL must route to a remote frontend. Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * test: keep local peer address reserved in kill-routing test The test reserved the local frontend address with a temporary listener that was dropped immediately, then reserved the remote address separately. The kernel could assign the same port to the remote frontend, so the remote kill request would take the local branch and fail to find the remote process. Bind the local address once and keep the listener alive for the whole test so the local and remote frontend identities cannot collide. Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> --------- Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
2912 lines
102 KiB
Rust
2912 lines
102 KiB
Rust
// 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.
|
|
|
|
// Retain coverage of the legacy constructors, including TLS configuration.
|
|
#![allow(deprecated)]
|
|
|
|
use std::net::SocketAddr;
|
|
use std::sync::Arc;
|
|
|
|
use api::v1::alter_table_expr::Kind;
|
|
use api::v1::greptime_database_client::GreptimeDatabaseClient;
|
|
use api::v1::greptime_request::Request as RequestBody;
|
|
use api::v1::greptime_response::Response as ResponseBody;
|
|
use api::v1::promql_request::Promql;
|
|
use api::v1::query_request::Query;
|
|
use api::v1::value::ValueData;
|
|
use api::v1::{
|
|
AddColumn, AddColumns, AlterTableExpr, Basic, Column, ColumnDataType, ColumnDef, ColumnSchema,
|
|
CreateTableExpr, GreptimeRequest, InsertRequest, InsertRequests, PromInstantQuery,
|
|
PromRangeQuery, PromqlRequest, QueryRequest, RequestHeader, Row, RowInsertRequest,
|
|
RowInsertRequests, Rows, SemanticType, Value, column,
|
|
};
|
|
use auth::user_provider_from_option;
|
|
use base64::prelude::{BASE64_STANDARD, Engine as _};
|
|
use catalog::process_manager::{ProcessManager, ProcessManagerRef};
|
|
use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputData};
|
|
use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, METRICS_EXEMPLARS_TABLE, MITO_ENGINE};
|
|
use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME;
|
|
use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME;
|
|
use common_grpc::channel_manager::ClientTlsOption;
|
|
use common_memory_manager::OnExhaustedPolicy;
|
|
use common_meta::cluster::{FrontendStatus, NodeInfo, NodeInfoKey, NodeStatus, Role};
|
|
use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME;
|
|
use common_meta::peer::Peer;
|
|
use common_meta::rpc::store::PutRequest;
|
|
use common_query::Output;
|
|
use common_recordbatch::RecordBatches;
|
|
use common_runtime::runtime::{BuilderBuild, RuntimeTrait};
|
|
use common_runtime::{Builder as RuntimeBuilder, Runtime};
|
|
use common_test_util::find_workspace_path;
|
|
use common_time::util;
|
|
use datatypes::arrow::array::{
|
|
Array, ArrayRef, Float64Array, Float64Builder, Int32Array, ListBuilder, StringArray,
|
|
StructArray, TimestampNanosecondArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array,
|
|
UInt64Builder,
|
|
};
|
|
use datatypes::arrow::datatypes::{DataType, Field};
|
|
use datatypes::arrow::ipc::writer::StreamWriter;
|
|
use datatypes::arrow::record_batch::RecordBatch as ArrowRecordBatch;
|
|
use frontend::frontend::FrontendOptions;
|
|
use meta_client::client::{MetaClient, MetaClientBuilder};
|
|
use meta_srv::mocks::{self, MockInfo};
|
|
use otel_arrow_rust::otlp::metrics::MetricType as ArrowMetricType;
|
|
use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_client::ArrowMetricsServiceClient;
|
|
use otel_arrow_rust::proto::opentelemetry::arrow::v1::{
|
|
ArrowPayload, ArrowPayloadType, BatchArrowRecords, StatusCode as ArrowStatusCode,
|
|
};
|
|
use otel_arrow_rust::proto::opentelemetry::metrics::v1::AggregationTemporality;
|
|
use otel_arrow_rust::schema::consts as arrow_consts;
|
|
use rstest_reuse::apply;
|
|
use servers::grpc::builder::GrpcServerBuilder;
|
|
use servers::grpc::frontend_grpc_handler::FrontendGrpcHandler;
|
|
use servers::grpc::{GrpcOptions, GrpcServer, GrpcServerConfig};
|
|
use servers::http::prometheus::{
|
|
PromData, PromQueryResult, PromSampleValue, PromSeriesMatrix, PromSeriesVector,
|
|
PrometheusJsonResponse, PrometheusResponse,
|
|
};
|
|
use servers::request_memory_limiter::ServerMemoryLimiter;
|
|
use servers::server::Server;
|
|
use servers::tls::{TlsMode, TlsOption};
|
|
use session::hints::{HINTS_KEY, INSERT_SKIP_WAL_HINT};
|
|
use tests_integration::test_util::{
|
|
MockInstanceImpl, StorageType, assert_wal_delta, setup_grpc_server,
|
|
setup_grpc_server_for_frontend_instance, setup_grpc_server_with,
|
|
setup_grpc_server_with_auto_create_table_disabled, setup_grpc_server_with_user_provider,
|
|
};
|
|
use tonic::Request;
|
|
use tonic::codec::CompressionEncoding;
|
|
use tonic::metadata::MetadataValue;
|
|
|
|
use crate::both_deployment_cases;
|
|
|
|
#[macro_export]
|
|
macro_rules! grpc_test {
|
|
($service:ident, $($(#[$meta:meta])* $test:ident),*,) => {
|
|
paste::item! {
|
|
mod [<integration_grpc_ $service:lower _test>] {
|
|
$(
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
$(
|
|
#[$meta]
|
|
)*
|
|
async fn [< $test >]() {
|
|
let store_type = tests_integration::test_util::StorageType::$service;
|
|
if store_type.test_on() {
|
|
// Support both unit tests and fallible tests without discarding errors.
|
|
let result = $crate::grpc::$test(store_type).await;
|
|
assert_eq!(
|
|
std::process::Termination::report(result),
|
|
std::process::ExitCode::SUCCESS,
|
|
);
|
|
}
|
|
}
|
|
)*
|
|
}
|
|
}
|
|
};
|
|
}
|
|
|
|
#[macro_export]
|
|
macro_rules! grpc_tests {
|
|
($($service:ident),*) => {
|
|
$(
|
|
grpc_test!(
|
|
$service,
|
|
|
|
test_invalid_dbname,
|
|
test_auto_create_table,
|
|
test_auto_create_table_with_hints,
|
|
test_auto_create_table_disabled_by_config,
|
|
test_private_system_tables_auto_create_table_with_global_disabled,
|
|
test_private_system_tables_bypass_auto_create_hint,
|
|
test_otel_arrow_auth,
|
|
test_otel_arrow_delta_histogram,
|
|
test_otel_arrow_exponential_histogram,
|
|
test_otel_arrow_exemplars_without_metric_engine,
|
|
test_insert_and_select,
|
|
test_dbname,
|
|
test_grpc_message_size_ok,
|
|
test_grpc_zstd_compression,
|
|
test_grpc_message_size_limit_recv,
|
|
test_grpc_message_size_limit_send,
|
|
test_grpc_auth,
|
|
test_health_check,
|
|
test_prom_gateway_query,
|
|
test_grpc_timezone,
|
|
test_grpc_tls_config,
|
|
test_grpc_memory_limit,
|
|
test_grpc_compressed_memory_reservation,
|
|
);
|
|
)*
|
|
};
|
|
}
|
|
|
|
pub async fn test_invalid_dbname(store_type: StorageType) {
|
|
let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_invalid_dbname").await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new_with_dbname("tom", grpc_client);
|
|
|
|
let (expected_host_col, expected_cpu_col, expected_mem_col, expected_ts_col) = expect_data();
|
|
let request = InsertRequest {
|
|
table_name: "demo".to_string(),
|
|
columns: vec![
|
|
expected_host_col.clone(),
|
|
expected_cpu_col.clone(),
|
|
expected_mem_col.clone(),
|
|
expected_ts_col.clone(),
|
|
],
|
|
row_count: 4,
|
|
};
|
|
let result = db
|
|
.insert(InsertRequests {
|
|
inserts: vec![request],
|
|
})
|
|
.await;
|
|
assert!(result.is_err());
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_dbname(store_type: StorageType) {
|
|
let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_dbname").await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new_with_dbname(
|
|
format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
|
|
grpc_client,
|
|
);
|
|
insert_and_assert(&db).await;
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_grpc_message_size_ok(store_type: StorageType) {
|
|
let config = GrpcServerConfig {
|
|
max_recv_message_size: 1024,
|
|
max_send_message_size: 1024,
|
|
..Default::default()
|
|
};
|
|
let (_db, fe_grpc_server) = setup_grpc_server_with(
|
|
store_type,
|
|
"test_grpc_message_size_ok",
|
|
None,
|
|
Some(config),
|
|
None,
|
|
)
|
|
.await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new_with_dbname(
|
|
format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
|
|
grpc_client,
|
|
);
|
|
db.sql("show tables;").await.unwrap();
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
/// Both the server (`servers::grpc::builder`) and the standard client
|
|
/// (`configure_tonic_client!`) enable zstd, so a plain round trip never shows
|
|
/// whether negotiation actually happened. This drives raw tonic clients instead
|
|
/// and asserts on the `grpc-encoding` the server answers with.
|
|
pub async fn test_grpc_zstd_compression(store_type: StorageType) {
|
|
let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_grpc_zstd_compression").await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let ddl = |sql: &str| GreptimeRequest {
|
|
header: Some(RequestHeader {
|
|
catalog: DEFAULT_CATALOG_NAME.to_string(),
|
|
schema: DEFAULT_SCHEMA_NAME.to_string(),
|
|
..Default::default()
|
|
}),
|
|
request: Some(RequestBody::Query(QueryRequest {
|
|
query: Some(Query::Sql(sql.to_string())),
|
|
})),
|
|
};
|
|
|
|
// Sends zstd and accepts zstd: the server has to decode a compressed request
|
|
// body and compress its response.
|
|
let mut zstd_client = GreptimeDatabaseClient::connect(format!("http://{addr}"))
|
|
.await
|
|
.unwrap()
|
|
.send_compressed(CompressionEncoding::Zstd)
|
|
.accept_compressed(CompressionEncoding::Zstd);
|
|
let response = zstd_client
|
|
.handle(Request::new(ddl(
|
|
"CREATE TABLE zstd_compression (ts TIMESTAMP TIME INDEX, payload STRING)",
|
|
)))
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(
|
|
response
|
|
.metadata()
|
|
.get("grpc-encoding")
|
|
.map(|v| v.to_str().unwrap()),
|
|
Some("zstd")
|
|
);
|
|
|
|
// A payload far above the zstd frame overhead, so the request body really is
|
|
// compressed rather than passed through.
|
|
let payload = "compressible-".repeat(4096);
|
|
let insert = |ts: i64| GreptimeRequest {
|
|
header: Some(RequestHeader {
|
|
catalog: DEFAULT_CATALOG_NAME.to_string(),
|
|
schema: DEFAULT_SCHEMA_NAME.to_string(),
|
|
..Default::default()
|
|
}),
|
|
request: Some(RequestBody::RowInserts(RowInsertRequests {
|
|
inserts: vec![RowInsertRequest {
|
|
table_name: "zstd_compression".to_string(),
|
|
rows: Some(Rows {
|
|
schema: vec![
|
|
ColumnSchema {
|
|
column_name: "ts".to_string(),
|
|
semantic_type: SemanticType::Timestamp as i32,
|
|
datatype: ColumnDataType::TimestampMillisecond as i32,
|
|
..Default::default()
|
|
},
|
|
ColumnSchema {
|
|
column_name: "payload".to_string(),
|
|
semantic_type: SemanticType::Field as i32,
|
|
datatype: ColumnDataType::String as i32,
|
|
..Default::default()
|
|
},
|
|
],
|
|
rows: vec![Row {
|
|
values: vec![
|
|
Value {
|
|
value_data: Some(ValueData::TimestampMillisecondValue(ts)),
|
|
},
|
|
Value {
|
|
value_data: Some(ValueData::StringValue(payload.clone())),
|
|
},
|
|
],
|
|
}],
|
|
}),
|
|
}],
|
|
})),
|
|
};
|
|
|
|
let response = zstd_client
|
|
.handle(Request::new(insert(1000)))
|
|
.await
|
|
.unwrap();
|
|
let ResponseBody::AffectedRows(rows) = response.into_inner().response.unwrap();
|
|
assert_eq!(rows.value, 1);
|
|
|
|
// A client that does not advertise zstd gets an uncompressed response.
|
|
let mut plain_client = GreptimeDatabaseClient::connect(format!("http://{addr}"))
|
|
.await
|
|
.unwrap();
|
|
let response = plain_client
|
|
.handle(Request::new(insert(2000)))
|
|
.await
|
|
.unwrap();
|
|
assert!(response.metadata().get("grpc-encoding").is_none());
|
|
|
|
// Both payloads survived their respective paths intact.
|
|
let db = Database::new_with_dbname(
|
|
format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
|
|
Client::with_urls(vec![addr]),
|
|
);
|
|
let sql = format!(
|
|
"SELECT count(*) AS c FROM zstd_compression WHERE length(payload) = {}",
|
|
payload.len()
|
|
);
|
|
let OutputData::Stream(stream) = db.sql(&sql).await.unwrap().data else {
|
|
panic!("expected a stream");
|
|
};
|
|
let recordbatches = RecordBatches::try_collect(stream).await.unwrap();
|
|
assert_eq!(
|
|
recordbatches.pretty_print().unwrap(),
|
|
"\
|
|
+---+
|
|
| c |
|
|
+---+
|
|
| 2 |
|
|
+---+"
|
|
);
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_grpc_message_size_limit_send(store_type: StorageType) {
|
|
let config = GrpcServerConfig {
|
|
max_recv_message_size: 1024,
|
|
max_send_message_size: 50,
|
|
..Default::default()
|
|
};
|
|
let (_db, fe_grpc_server) = setup_grpc_server_with(
|
|
store_type,
|
|
"test_grpc_message_size_limit_send",
|
|
None,
|
|
Some(config),
|
|
None,
|
|
)
|
|
.await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new_with_dbname(
|
|
format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
|
|
grpc_client,
|
|
);
|
|
let err_msg = db.sql("show tables;").await.unwrap_err().to_string();
|
|
assert!(err_msg.contains("message length too large"), "{}", err_msg);
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
/// Transport-compressed gRPC requests reserve the worst-case decoded message
|
|
/// size (`max_recv_message_size`) against the aggregate quota *before* tonic
|
|
/// decompresses them. With a quota smaller than that bound, compressed
|
|
/// requests must fail fast with `RESOURCE_EXHAUSTED` (before decoding), while
|
|
/// uncompressed requests are still admitted through the exact post-decode
|
|
/// charge. With the default unlimited limiter, compression keeps working.
|
|
pub async fn test_grpc_compressed_memory_reservation(store_type: StorageType) {
|
|
use api::v1::query_request::Query;
|
|
use api::v1::{GreptimeRequest, QueryRequest};
|
|
use tonic::codec::CompressionEncoding;
|
|
|
|
let config = GrpcServerConfig {
|
|
max_recv_message_size: 1024 * 1024,
|
|
..Default::default()
|
|
};
|
|
// Quota smaller than the decoding bound: compressed requests cannot be
|
|
// pre-admitted.
|
|
let memory_limiter = ServerMemoryLimiter::new(512 * 1024, OnExhaustedPolicy::Fail);
|
|
let (_db, fe_grpc_server) = setup_grpc_server_with(
|
|
store_type,
|
|
"test_grpc_compressed_memory_reservation",
|
|
None,
|
|
Some(config),
|
|
Some(memory_limiter),
|
|
)
|
|
.await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let db = Database::new_with_dbname(
|
|
format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
|
|
Client::with_urls(vec![addr.clone()]),
|
|
);
|
|
db.sql("CREATE TABLE grpc_quota(ts TIMESTAMP TIME INDEX, v DOUBLE)")
|
|
.await
|
|
.unwrap();
|
|
|
|
let insert_request = || {
|
|
tonic::Request::new(GreptimeRequest {
|
|
header: Some(RequestHeader {
|
|
catalog: DEFAULT_CATALOG_NAME.to_string(),
|
|
schema: DEFAULT_SCHEMA_NAME.to_string(),
|
|
..Default::default()
|
|
}),
|
|
request: Some(RequestBody::Query(QueryRequest {
|
|
query: Some(Query::Sql(
|
|
"INSERT INTO grpc_quota VALUES (1000, 1.0)".to_string(),
|
|
)),
|
|
})),
|
|
})
|
|
};
|
|
|
|
// Compressed request: the pre-decode reservation (1 MiB) exceeds the
|
|
// 512 KiB quota and must be rejected before any decoding happens.
|
|
let mut compressed = GreptimeDatabaseClient::connect(format!("http://{addr}"))
|
|
.await
|
|
.unwrap()
|
|
.send_compressed(CompressionEncoding::Zstd);
|
|
let status = compressed.handle(insert_request()).await.unwrap_err();
|
|
assert_eq!(
|
|
status.code(),
|
|
tonic::Code::ResourceExhausted,
|
|
"compressed request must be pre-admitted against the quota, got: {status}"
|
|
);
|
|
|
|
// Uncompressed request on the same server: still admitted (its exact
|
|
// serialized size fits the quota).
|
|
let mut plain = GreptimeDatabaseClient::connect(format!("http://{addr}"))
|
|
.await
|
|
.unwrap();
|
|
plain.handle(insert_request()).await.unwrap();
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
|
|
// Default (unlimited) limiter: compressed requests keep working.
|
|
let (_db, fe_grpc_server) = setup_grpc_server_with(
|
|
store_type,
|
|
"test_grpc_compressed_memory_reservation_unlimited",
|
|
None,
|
|
Some(GrpcServerConfig::default()),
|
|
None,
|
|
)
|
|
.await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
let db = Database::new_with_dbname(
|
|
format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
|
|
Client::with_urls(vec![addr.clone()]),
|
|
);
|
|
db.sql("CREATE TABLE grpc_quota(ts TIMESTAMP TIME INDEX, v DOUBLE)")
|
|
.await
|
|
.unwrap();
|
|
let mut compressed = GreptimeDatabaseClient::connect(format!("http://{addr}"))
|
|
.await
|
|
.unwrap()
|
|
.send_compressed(CompressionEncoding::Zstd);
|
|
compressed.handle(insert_request()).await.unwrap();
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_grpc_message_size_limit_recv(store_type: StorageType) {
|
|
let config = GrpcServerConfig {
|
|
max_recv_message_size: 10,
|
|
max_send_message_size: 1024,
|
|
..Default::default()
|
|
};
|
|
let (_db, fe_grpc_server) = setup_grpc_server_with(
|
|
store_type,
|
|
"test_grpc_message_size_limit_recv",
|
|
None,
|
|
Some(config),
|
|
None,
|
|
)
|
|
.await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new_with_dbname(
|
|
format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
|
|
grpc_client,
|
|
);
|
|
let err_msg = db.sql("show tables;").await.unwrap_err().to_string();
|
|
assert!(
|
|
err_msg.contains("Operation was attempted past the valid range"),
|
|
"{}",
|
|
err_msg
|
|
);
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_grpc_auth(store_type: StorageType) {
|
|
let user_provider =
|
|
user_provider_from_option("static_user_provider:cmd:greptime_user=greptime_pwd").unwrap();
|
|
let (_db, fe_grpc_server) =
|
|
setup_grpc_server_with_user_provider(store_type, "auto_create_table", Some(user_provider))
|
|
.await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let mut db = Database::new_with_dbname(
|
|
format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
|
|
grpc_client,
|
|
);
|
|
|
|
// 1. test without auth
|
|
let re = db.sql("show tables;").await;
|
|
assert!(re.is_err());
|
|
assert!(matches!(
|
|
re,
|
|
Err(client::Error::FlightGet {
|
|
tonic_code: tonic::Code::Unauthenticated,
|
|
..
|
|
})
|
|
));
|
|
|
|
// 2. test wrong auth
|
|
db.set_auth(api::v1::auth_header::AuthScheme::Basic(Basic {
|
|
username: "greptime_user".to_string(),
|
|
password: "wrong_pwd".to_string(),
|
|
}));
|
|
let re = db.sql("show tables;").await;
|
|
assert!(re.is_err());
|
|
assert!(matches!(
|
|
re,
|
|
Err(client::Error::FlightGet {
|
|
tonic_code: tonic::Code::Unauthenticated,
|
|
..
|
|
})
|
|
));
|
|
|
|
// 3. test right auth
|
|
db.set_auth(api::v1::auth_header::AuthScheme::Basic(Basic {
|
|
username: "greptime_user".to_string(),
|
|
password: "greptime_pwd".to_string(),
|
|
}));
|
|
let re = db.sql("show tables;").await;
|
|
assert!(re.is_ok());
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
async fn start_process_rpc_server(process_manager: Option<ProcessManagerRef>) -> GrpcServer {
|
|
start_process_rpc_server_at(
|
|
process_manager,
|
|
"127.0.0.1:0".parse::<SocketAddr>().unwrap(),
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn start_process_rpc_server_at(
|
|
process_manager: Option<ProcessManagerRef>,
|
|
addr: SocketAddr,
|
|
) -> GrpcServer {
|
|
let runtime: Runtime = RuntimeBuilder::default()
|
|
.worker_threads(2)
|
|
.thread_name("process-rpc-test")
|
|
.build()
|
|
.unwrap();
|
|
let mut builder = GrpcServerBuilder::new(GrpcServerConfig::default(), runtime);
|
|
if let Some(process_manager) = process_manager {
|
|
builder = builder.frontend_grpc_handler(FrontendGrpcHandler::new(process_manager));
|
|
}
|
|
let mut server = builder.build();
|
|
server.start(addr).await.unwrap();
|
|
server
|
|
}
|
|
|
|
/// Reserves a loopback address for a test server.
|
|
///
|
|
/// A process manager must be created with the address it advertises before its
|
|
/// gRPC listener starts, so pick a free port first and bind it right after.
|
|
fn reserve_loopback_addr() -> SocketAddr {
|
|
std::net::TcpListener::bind("127.0.0.1:0")
|
|
.unwrap()
|
|
.local_addr()
|
|
.unwrap()
|
|
}
|
|
|
|
async fn build_meta_client_for_mock(mock: &MockInfo) -> Arc<MetaClient> {
|
|
let mut meta_client = MetaClientBuilder::frontend_default_options()
|
|
.channel_manager(mock.channel_manager.clone())
|
|
.build();
|
|
meta_client.start(&[&mock.server_addr]).await.unwrap();
|
|
Arc::new(meta_client)
|
|
}
|
|
|
|
async fn register_process_rpc_peer(mock: &MockInfo, node_id: u64, addr: String) {
|
|
let node_info = NodeInfo {
|
|
peer: Peer { id: node_id, addr },
|
|
last_activity_ts: util::current_time_millis(),
|
|
status: NodeStatus::Frontend(FrontendStatus::default()),
|
|
version: String::new(),
|
|
git_commit: String::new(),
|
|
start_time_ms: 0,
|
|
total_cpu_millicores: 0,
|
|
total_memory_bytes: 0,
|
|
cpu_usage_millicores: 0,
|
|
memory_usage_bytes: 0,
|
|
hostname: String::new(),
|
|
env_vars: Default::default(),
|
|
schedulable: None,
|
|
};
|
|
mock.in_memory
|
|
.as_ref()
|
|
.unwrap()
|
|
.put(PutRequest {
|
|
key: Vec::from(&NodeInfoKey {
|
|
role: Role::Frontend,
|
|
node_id,
|
|
}),
|
|
value: node_info.try_into().unwrap(),
|
|
prev_kv: false,
|
|
})
|
|
.await
|
|
.unwrap();
|
|
}
|
|
|
|
async fn build_process_manager_with_remote_peer(
|
|
mock: &MockInfo,
|
|
remote_addr: String,
|
|
) -> ProcessManagerRef {
|
|
register_process_rpc_peer(mock, 1, remote_addr).await;
|
|
let meta_client = build_meta_client_for_mock(mock).await;
|
|
Arc::new(ProcessManager::new("local".to_string(), Some(meta_client)))
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
async fn test_list_all_processes_reports_unimplemented_remote_service() {
|
|
let mock = mocks::mock_with_memstore().await;
|
|
let server = start_process_rpc_server(None).await;
|
|
let remote_addr = server.bind_addr().unwrap().to_string();
|
|
let process_manager = build_process_manager_with_remote_peer(&mock, remote_addr).await;
|
|
let _local = process_manager.register_query(
|
|
"greptime".to_string(),
|
|
vec!["public".to_string()],
|
|
"local query".to_string(),
|
|
String::new(),
|
|
None,
|
|
None,
|
|
);
|
|
|
|
let result = process_manager.list_all_processes(Some("greptime")).await;
|
|
server.shutdown().await.unwrap();
|
|
|
|
assert!(
|
|
matches!(
|
|
&result,
|
|
Err(catalog::error::Error::InvokeFrontend {
|
|
source: common_frontend::error::Error::InvokeFrontend { error, .. },
|
|
..
|
|
}) if error.code() == tonic::Code::Unimplemented
|
|
),
|
|
"a missing remote management service must not return an incomplete list: {result:?}"
|
|
);
|
|
assert_eq!(
|
|
process_manager
|
|
.local_processes(Some("greptime"))
|
|
.unwrap()
|
|
.len(),
|
|
1
|
|
);
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
async fn test_list_all_processes_collects_remote_and_local_catalog() {
|
|
let mock = mocks::mock_with_memstore().await;
|
|
let remote_manager = Arc::new(ProcessManager::new("remote".to_string(), None));
|
|
let _selected = remote_manager.register_query(
|
|
"greptime".to_string(),
|
|
vec!["public".to_string()],
|
|
"selected remote query".to_string(),
|
|
String::new(),
|
|
None,
|
|
None,
|
|
);
|
|
let _other = remote_manager.register_query(
|
|
"other".to_string(),
|
|
vec!["public".to_string()],
|
|
"other remote query".to_string(),
|
|
String::new(),
|
|
None,
|
|
None,
|
|
);
|
|
let server = start_process_rpc_server(Some(remote_manager)).await;
|
|
let remote_addr = server.bind_addr().unwrap().to_string();
|
|
let process_manager = build_process_manager_with_remote_peer(&mock, remote_addr).await;
|
|
let _local = process_manager.register_query(
|
|
"greptime".to_string(),
|
|
vec!["public".to_string()],
|
|
"selected local query".to_string(),
|
|
String::new(),
|
|
None,
|
|
None,
|
|
);
|
|
|
|
let processes = process_manager
|
|
.list_all_processes(Some("greptime"))
|
|
.await
|
|
.unwrap();
|
|
server.shutdown().await.unwrap();
|
|
|
|
assert_eq!(processes.len(), 2);
|
|
assert!(
|
|
processes
|
|
.iter()
|
|
.any(|process| process.query == "selected local query")
|
|
);
|
|
assert!(
|
|
processes
|
|
.iter()
|
|
.any(|process| process.query == "selected remote query")
|
|
);
|
|
}
|
|
|
|
/// Regression test for frontends with separate public and internal gRPC
|
|
/// addresses.
|
|
///
|
|
/// The local process manager must recognize itself by the peer address it
|
|
/// advertises to metasrv (the internal gRPC address), not by its public gRPC
|
|
/// address. Otherwise it queries its own management service as if it were a
|
|
/// remote frontend and then appends the local registry a second time, so every
|
|
/// local query is listed twice.
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
async fn test_list_all_processes_excludes_local_peer_with_internal_grpc() {
|
|
let mock = mocks::mock_with_memstore().await;
|
|
let local_peer_addr = reserve_loopback_addr();
|
|
let mut options = FrontendOptions::default();
|
|
options.grpc.bind_addr = "127.0.0.1:4001".to_string();
|
|
options.grpc.server_addr = "127.0.0.1:4001".to_string();
|
|
options.internal_grpc = Some(GrpcOptions {
|
|
bind_addr: local_peer_addr.to_string(),
|
|
server_addr: local_peer_addr.to_string(),
|
|
..Default::default()
|
|
});
|
|
assert_ne!(options.grpc.server_addr, options.peer_addr());
|
|
|
|
let meta_client = build_meta_client_for_mock(&mock).await;
|
|
let local_manager = options.create_process_manager(Some(meta_client));
|
|
let local_peer_addr_str = local_peer_addr.to_string();
|
|
let _local_query = local_manager.register_query(
|
|
"greptime".to_string(),
|
|
vec!["public".to_string()],
|
|
"local query".to_string(),
|
|
String::new(),
|
|
None,
|
|
None,
|
|
);
|
|
// The internal gRPC listener of this frontend serves its process list.
|
|
let local_server =
|
|
start_process_rpc_server_at(Some(local_manager.clone()), local_peer_addr).await;
|
|
|
|
let remote_peer_addr = reserve_loopback_addr();
|
|
let remote_manager = Arc::new(ProcessManager::new(remote_peer_addr.to_string(), None));
|
|
let _remote_query = remote_manager.register_query(
|
|
"greptime".to_string(),
|
|
vec!["public".to_string()],
|
|
"remote query".to_string(),
|
|
String::new(),
|
|
None,
|
|
None,
|
|
);
|
|
let remote_server = start_process_rpc_server_at(Some(remote_manager), remote_peer_addr).await;
|
|
|
|
register_process_rpc_peer(&mock, 1, local_peer_addr_str.clone()).await;
|
|
register_process_rpc_peer(&mock, 2, remote_peer_addr.to_string()).await;
|
|
|
|
let processes = local_manager
|
|
.list_all_processes(Some("greptime"))
|
|
.await
|
|
.unwrap();
|
|
local_server.shutdown().await.unwrap();
|
|
remote_server.shutdown().await.unwrap();
|
|
|
|
assert_eq!(
|
|
processes
|
|
.iter()
|
|
.filter(|process| process.query == "local query")
|
|
.count(),
|
|
1,
|
|
"the local query must be listed exactly once: {processes:?}"
|
|
);
|
|
assert_eq!(
|
|
processes
|
|
.iter()
|
|
.filter(|process| process.query == "remote query")
|
|
.count(),
|
|
1,
|
|
"a remote frontend must still be listed exactly once: {processes:?}"
|
|
);
|
|
assert_eq!(processes.len(), 2, "{processes:?}");
|
|
assert!(
|
|
processes
|
|
.iter()
|
|
.filter(|process| process.query == "local query")
|
|
.all(|process| process.frontend == local_peer_addr_str)
|
|
);
|
|
}
|
|
|
|
/// `KILL` must route by the peer address that the process list displays and
|
|
/// metasrv advertises, for both local and remote processes.
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
async fn test_kill_process_routes_by_advertised_peer_addr() {
|
|
let mock = mocks::mock_with_memstore().await;
|
|
let meta_client = build_meta_client_for_mock(&mock).await;
|
|
|
|
// Keep this listener bound for the whole test so the remote server cannot be
|
|
// assigned the same address; otherwise the remote kill would take the local
|
|
// branch and fail to find the remote process.
|
|
let local_listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
|
let local_peer_addr = local_listener.local_addr().unwrap();
|
|
let local_manager = Arc::new(ProcessManager::new(
|
|
local_peer_addr.to_string(),
|
|
Some(meta_client),
|
|
));
|
|
let local_query_id = 7;
|
|
let _local_query = local_manager.register_query(
|
|
"greptime".to_string(),
|
|
vec!["public".to_string()],
|
|
"local query".to_string(),
|
|
String::new(),
|
|
Some(local_query_id),
|
|
None,
|
|
);
|
|
|
|
let remote_peer_addr = reserve_loopback_addr();
|
|
let remote_manager = Arc::new(ProcessManager::new(remote_peer_addr.to_string(), None));
|
|
let remote_query_id = 8;
|
|
let _remote_query = remote_manager.register_query(
|
|
"greptime".to_string(),
|
|
vec!["public".to_string()],
|
|
"remote query".to_string(),
|
|
String::new(),
|
|
Some(remote_query_id),
|
|
None,
|
|
);
|
|
let remote_server =
|
|
start_process_rpc_server_at(Some(remote_manager.clone()), remote_peer_addr).await;
|
|
register_process_rpc_peer(&mock, 1, remote_peer_addr.to_string()).await;
|
|
|
|
let local_killed = local_manager
|
|
.kill_process(
|
|
local_peer_addr.to_string(),
|
|
"greptime".to_string(),
|
|
local_query_id,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert!(local_killed);
|
|
assert!(
|
|
local_manager
|
|
.local_processes(Some("greptime"))
|
|
.unwrap()
|
|
.is_empty()
|
|
);
|
|
|
|
let remote_killed = local_manager
|
|
.kill_process(
|
|
remote_peer_addr.to_string(),
|
|
"greptime".to_string(),
|
|
remote_query_id,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert!(remote_killed);
|
|
assert!(
|
|
remote_manager
|
|
.local_processes(Some("greptime"))
|
|
.unwrap()
|
|
.is_empty()
|
|
);
|
|
|
|
remote_server.shutdown().await.unwrap();
|
|
}
|
|
|
|
pub async fn test_otel_arrow_auth(store_type: StorageType) {
|
|
let user_provider =
|
|
user_provider_from_option("static_user_provider:cmd:greptime_user=greptime_pwd").unwrap();
|
|
let (_db, fe_grpc_server) = setup_grpc_server_with_user_provider(
|
|
store_type,
|
|
"test_otel_arrow_auth",
|
|
Some(user_provider),
|
|
)
|
|
.await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let mut client = ArrowMetricsServiceClient::connect(format!("http://{}", addr))
|
|
.await
|
|
.unwrap();
|
|
|
|
let batch_arrow_records = BatchArrowRecords {
|
|
batch_id: 1,
|
|
arrow_payloads: vec![],
|
|
headers: vec![],
|
|
};
|
|
|
|
// test without auth
|
|
{
|
|
let records = batch_arrow_records.clone();
|
|
let stream = futures::stream::once(async { records });
|
|
let request = Request::new(stream);
|
|
let response = client.arrow_metrics(request).await;
|
|
assert!(response.is_err());
|
|
let error = response.unwrap_err();
|
|
assert_eq!(error.code(), tonic::Code::Unauthenticated);
|
|
}
|
|
// test auth
|
|
{
|
|
let records = batch_arrow_records.clone();
|
|
let stream = futures::stream::once(async { records });
|
|
let mut request = Request::new(stream);
|
|
request.metadata_mut().insert(
|
|
"authorization",
|
|
MetadataValue::try_from(basic_auth("greptime_user", "greptime_pwd")).unwrap(),
|
|
);
|
|
let response = client.arrow_metrics(request).await;
|
|
assert!(response.is_ok());
|
|
|
|
let mut response_stream = response.unwrap().into_inner();
|
|
let resp = response_stream.message().await;
|
|
assert!(resp.is_err());
|
|
let error = resp.unwrap_err();
|
|
assert_eq!(
|
|
error.message(),
|
|
"Failed to handle otel-arrow request, error message: Batch is empty"
|
|
);
|
|
}
|
|
// test old auth
|
|
{
|
|
let stream = futures::stream::once(async { batch_arrow_records });
|
|
let mut request = Request::new(stream);
|
|
request.metadata_mut().insert(
|
|
"authorization",
|
|
MetadataValue::try_from(basic_auth_credentials("greptime_user", "greptime_pwd"))
|
|
.unwrap(),
|
|
);
|
|
let response = client.arrow_metrics(request).await;
|
|
assert!(response.is_ok());
|
|
|
|
let mut response_stream = response.unwrap().into_inner();
|
|
let resp = response_stream.message().await;
|
|
assert!(resp.is_err());
|
|
let error = resp.unwrap_err();
|
|
assert_eq!(
|
|
error.message(),
|
|
"Failed to handle otel-arrow request, error message: Batch is empty"
|
|
);
|
|
}
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
// The pinned otel-arrow Producer cannot hash List-typed bucket schemas yet, so
|
|
// serialize these test-only record batches directly into the same Arrow stream format.
|
|
fn serialize_arrow_record_batch(record_batch: &ArrowRecordBatch) -> Vec<u8> {
|
|
let mut bytes = Vec::new();
|
|
let mut writer = StreamWriter::try_new(&mut bytes, record_batch.schema_ref()).unwrap();
|
|
writer.write(record_batch).unwrap();
|
|
writer.finish().unwrap();
|
|
drop(writer);
|
|
bytes
|
|
}
|
|
|
|
fn exponential_histogram_arrow_batch(batch_id: i64, scales: &[i32]) -> BatchArrowRecords {
|
|
let resource = StructArray::from(vec![(
|
|
Arc::new(Field::new(arrow_consts::ID, DataType::UInt16, true)),
|
|
Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef,
|
|
)]);
|
|
let scope = StructArray::from(vec![(
|
|
Arc::new(Field::new(arrow_consts::ID, DataType::UInt16, true)),
|
|
Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef,
|
|
)]);
|
|
let metrics = ArrowRecordBatch::try_from_iter(vec![
|
|
(
|
|
arrow_consts::ID,
|
|
Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef,
|
|
),
|
|
(arrow_consts::RESOURCE, Arc::new(resource) as ArrayRef),
|
|
(arrow_consts::SCOPE, Arc::new(scope) as ArrayRef),
|
|
(
|
|
arrow_consts::METRIC_TYPE,
|
|
Arc::new(UInt8Array::from(vec![
|
|
ArrowMetricType::ExponentialHistogram as u8,
|
|
])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::NAME,
|
|
Arc::new(StringArray::from(vec!["otel.arrow.exponential.latency"])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::AGGREGATION_TEMPORALITY,
|
|
Arc::new(Int32Array::from(vec![
|
|
AggregationTemporality::Cumulative as i32,
|
|
])) as ArrayRef,
|
|
),
|
|
])
|
|
.unwrap();
|
|
|
|
let point_count = scales.len();
|
|
let mut positive_counts = ListBuilder::new(UInt64Builder::new());
|
|
let mut negative_counts = ListBuilder::new(UInt64Builder::new());
|
|
for _ in scales {
|
|
positive_counts.values().append_slice(&[1, 2]);
|
|
positive_counts.append(true);
|
|
negative_counts.append(true);
|
|
}
|
|
let positive_counts = positive_counts.finish();
|
|
let negative_counts = negative_counts.finish();
|
|
let positive = StructArray::from(vec![
|
|
(
|
|
Arc::new(Field::new(
|
|
arrow_consts::EXP_HISTOGRAM_OFFSET,
|
|
DataType::Int32,
|
|
true,
|
|
)),
|
|
Arc::new(Int32Array::from(vec![-1; point_count])) as ArrayRef,
|
|
),
|
|
(
|
|
Arc::new(Field::new(
|
|
arrow_consts::EXP_HISTOGRAM_BUCKET_COUNTS,
|
|
positive_counts.data_type().clone(),
|
|
true,
|
|
)),
|
|
Arc::new(positive_counts) as ArrayRef,
|
|
),
|
|
]);
|
|
let negative = StructArray::from(vec![
|
|
(
|
|
Arc::new(Field::new(
|
|
arrow_consts::EXP_HISTOGRAM_OFFSET,
|
|
DataType::Int32,
|
|
true,
|
|
)),
|
|
Arc::new(Int32Array::from(vec![0; point_count])) as ArrayRef,
|
|
),
|
|
(
|
|
Arc::new(Field::new(
|
|
arrow_consts::EXP_HISTOGRAM_BUCKET_COUNTS,
|
|
negative_counts.data_type().clone(),
|
|
true,
|
|
)),
|
|
Arc::new(negative_counts) as ArrayRef,
|
|
),
|
|
]);
|
|
let data_points = ArrowRecordBatch::try_from_iter(vec![
|
|
(
|
|
arrow_consts::ID,
|
|
Arc::new(UInt32Array::from_iter_values(
|
|
(0..point_count).map(|id| u32::try_from(id).unwrap()),
|
|
)) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::PARENT_ID,
|
|
Arc::new(UInt16Array::from(vec![0_u16; point_count])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::START_TIME_UNIX_NANO,
|
|
Arc::new(TimestampNanosecondArray::from(vec![
|
|
1_000_000_000;
|
|
point_count
|
|
])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::TIME_UNIX_NANO,
|
|
Arc::new(TimestampNanosecondArray::from(vec![
|
|
3_000_000_000;
|
|
point_count
|
|
])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::HISTOGRAM_COUNT,
|
|
Arc::new(UInt64Array::from(vec![4_u64; point_count])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::HISTOGRAM_SUM,
|
|
Arc::new(Float64Array::from(vec![8.0; point_count])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::EXP_HISTOGRAM_SCALE,
|
|
Arc::new(Int32Array::from(scales.to_vec())) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::EXP_HISTOGRAM_ZERO_COUNT,
|
|
Arc::new(UInt64Array::from(vec![1_u64; point_count])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::EXP_HISTOGRAM_POSITIVE,
|
|
Arc::new(positive) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::EXP_HISTOGRAM_NEGATIVE,
|
|
Arc::new(negative) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::FLAGS,
|
|
Arc::new(UInt32Array::from(vec![0_u32; point_count])) as ArrayRef,
|
|
),
|
|
])
|
|
.unwrap();
|
|
BatchArrowRecords {
|
|
batch_id,
|
|
arrow_payloads: vec![
|
|
ArrowPayload {
|
|
schema_id: format!("metrics-{batch_id}"),
|
|
r#type: ArrowPayloadType::UnivariateMetrics as i32,
|
|
record: serialize_arrow_record_batch(&metrics),
|
|
},
|
|
ArrowPayload {
|
|
schema_id: format!("exp-histogram-{batch_id}"),
|
|
r#type: ArrowPayloadType::ExpHistogramDataPoints as i32,
|
|
record: serialize_arrow_record_batch(&data_points),
|
|
},
|
|
],
|
|
headers: vec![],
|
|
}
|
|
}
|
|
|
|
fn delta_histogram_arrow_batch(
|
|
batch_id: i64,
|
|
points: &[(u64, &[u64], &[f64])],
|
|
) -> BatchArrowRecords {
|
|
let resource = StructArray::from(vec![(
|
|
Arc::new(Field::new(arrow_consts::ID, DataType::UInt16, true)),
|
|
Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef,
|
|
)]);
|
|
let scope = StructArray::from(vec![(
|
|
Arc::new(Field::new(arrow_consts::ID, DataType::UInt16, true)),
|
|
Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef,
|
|
)]);
|
|
let metrics = ArrowRecordBatch::try_from_iter(vec![
|
|
(
|
|
arrow_consts::ID,
|
|
Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef,
|
|
),
|
|
(arrow_consts::RESOURCE, Arc::new(resource) as ArrayRef),
|
|
(arrow_consts::SCOPE, Arc::new(scope) as ArrayRef),
|
|
(
|
|
arrow_consts::METRIC_TYPE,
|
|
Arc::new(UInt8Array::from(vec![ArrowMetricType::Histogram as u8])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::NAME,
|
|
Arc::new(StringArray::from(vec!["otel.arrow.delta.histogram"])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::AGGREGATION_TEMPORALITY,
|
|
Arc::new(Int32Array::from(vec![AggregationTemporality::Delta as i32])) as ArrayRef,
|
|
),
|
|
])
|
|
.unwrap();
|
|
|
|
let mut bucket_counts = ListBuilder::new(UInt64Builder::new());
|
|
let mut explicit_bounds = ListBuilder::new(Float64Builder::new());
|
|
for (_, counts, bounds) in points {
|
|
bucket_counts.values().append_slice(counts);
|
|
bucket_counts.append(true);
|
|
explicit_bounds.values().append_slice(bounds);
|
|
explicit_bounds.append(true);
|
|
}
|
|
let bucket_counts = bucket_counts.finish();
|
|
let explicit_bounds = explicit_bounds.finish();
|
|
let data_points = ArrowRecordBatch::try_from_iter(vec![
|
|
(
|
|
arrow_consts::ID,
|
|
Arc::new(UInt32Array::from_iter_values(
|
|
(0..points.len()).map(|id| u32::try_from(id).unwrap()),
|
|
)) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::PARENT_ID,
|
|
Arc::new(UInt16Array::from(vec![0_u16; points.len()])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::START_TIME_UNIX_NANO,
|
|
Arc::new(TimestampNanosecondArray::from(vec![
|
|
1_000_000_000;
|
|
points.len()
|
|
])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::TIME_UNIX_NANO,
|
|
Arc::new(TimestampNanosecondArray::from_iter_values(
|
|
(1..=points.len()).map(|second| i64::try_from(second).unwrap() * 1_000_000_000),
|
|
)) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::HISTOGRAM_COUNT,
|
|
Arc::new(UInt64Array::from(
|
|
points
|
|
.iter()
|
|
.map(|(count, _, _)| *count)
|
|
.collect::<Vec<_>>(),
|
|
)) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::HISTOGRAM_SUM,
|
|
Arc::new(Float64Array::from(vec![1.0; points.len()])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::HISTOGRAM_BUCKET_COUNTS,
|
|
Arc::new(bucket_counts) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::HISTOGRAM_EXPLICIT_BOUNDS,
|
|
Arc::new(explicit_bounds) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::FLAGS,
|
|
Arc::new(UInt32Array::from(vec![0_u32; points.len()])) as ArrayRef,
|
|
),
|
|
])
|
|
.unwrap();
|
|
BatchArrowRecords {
|
|
batch_id,
|
|
arrow_payloads: vec![
|
|
ArrowPayload {
|
|
schema_id: format!("metrics-{batch_id}"),
|
|
r#type: ArrowPayloadType::UnivariateMetrics as i32,
|
|
record: serialize_arrow_record_batch(&metrics),
|
|
},
|
|
ArrowPayload {
|
|
schema_id: format!("histogram-{batch_id}"),
|
|
r#type: ArrowPayloadType::HistogramDataPoints as i32,
|
|
record: serialize_arrow_record_batch(&data_points),
|
|
},
|
|
],
|
|
headers: vec![],
|
|
}
|
|
}
|
|
|
|
fn gauge_arrow_batch(batch_id: i64, reserved_attr: bool) -> BatchArrowRecords {
|
|
let resource = StructArray::from(vec![(
|
|
Arc::new(Field::new(arrow_consts::ID, DataType::UInt16, true)),
|
|
Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef,
|
|
)]);
|
|
let scope = StructArray::from(vec![(
|
|
Arc::new(Field::new(arrow_consts::ID, DataType::UInt16, true)),
|
|
Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef,
|
|
)]);
|
|
let metrics = ArrowRecordBatch::try_from_iter(vec![
|
|
(
|
|
arrow_consts::ID,
|
|
Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef,
|
|
),
|
|
(arrow_consts::RESOURCE, Arc::new(resource) as ArrayRef),
|
|
(arrow_consts::SCOPE, Arc::new(scope) as ArrayRef),
|
|
(
|
|
arrow_consts::METRIC_TYPE,
|
|
Arc::new(UInt8Array::from(vec![ArrowMetricType::Gauge as u8])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::NAME,
|
|
Arc::new(StringArray::from(vec!["otel.arrow.gauge"])) as ArrayRef,
|
|
),
|
|
])
|
|
.unwrap();
|
|
let data_points = ArrowRecordBatch::try_from_iter(vec![
|
|
(
|
|
arrow_consts::ID,
|
|
Arc::new(UInt32Array::from(vec![0_u32])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::PARENT_ID,
|
|
Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::START_TIME_UNIX_NANO,
|
|
Arc::new(TimestampNanosecondArray::from(vec![1_000_000_000])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::TIME_UNIX_NANO,
|
|
Arc::new(TimestampNanosecondArray::from(vec![2_000_000_000])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::DOUBLE_VALUE,
|
|
Arc::new(Float64Array::from(vec![1.0])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::FLAGS,
|
|
Arc::new(UInt32Array::from(vec![0_u32])) as ArrayRef,
|
|
),
|
|
])
|
|
.unwrap();
|
|
let mut arrow_payloads = vec![
|
|
ArrowPayload {
|
|
schema_id: format!("metrics-{batch_id}"),
|
|
r#type: ArrowPayloadType::UnivariateMetrics as i32,
|
|
record: serialize_arrow_record_batch(&metrics),
|
|
},
|
|
ArrowPayload {
|
|
schema_id: format!("number-{batch_id}"),
|
|
r#type: ArrowPayloadType::NumberDataPoints as i32,
|
|
record: serialize_arrow_record_batch(&data_points),
|
|
},
|
|
];
|
|
if reserved_attr {
|
|
let attributes = ArrowRecordBatch::try_from_iter(vec![
|
|
(
|
|
arrow_consts::PARENT_ID,
|
|
Arc::new(UInt32Array::from(vec![0_u32])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::ATTRIBUTE_KEY,
|
|
Arc::new(StringArray::from(vec!["otlp_aggregation_temporality"])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::ATTRIBUTE_TYPE,
|
|
Arc::new(UInt8Array::from(vec![1_u8])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::ATTRIBUTE_STR,
|
|
Arc::new(StringArray::from(vec!["user"])) as ArrayRef,
|
|
),
|
|
])
|
|
.unwrap();
|
|
arrow_payloads.push(ArrowPayload {
|
|
schema_id: format!("number-attrs-{batch_id}"),
|
|
r#type: ArrowPayloadType::NumberDpAttrs as i32,
|
|
record: serialize_arrow_record_batch(&attributes),
|
|
});
|
|
}
|
|
BatchArrowRecords {
|
|
batch_id,
|
|
arrow_payloads,
|
|
headers: vec![],
|
|
}
|
|
}
|
|
|
|
#[apply(both_deployment_cases)]
|
|
async fn test_skip_wal_otel_arrow_metrics(distributed: bool) {
|
|
// OTEL Arrow metrics decode into ordinary metric inserts, unlike Flight DoPut bulk inserts.
|
|
let mut env = MockInstanceImpl::new("skip_wal_otel_arrow", distributed).await;
|
|
let server = setup_grpc_server_for_frontend_instance(env.frontend(), None).await;
|
|
let mut client =
|
|
ArrowMetricsServiceClient::connect(format!("http://{}", server.bind_addr().unwrap()))
|
|
.await
|
|
.unwrap();
|
|
// Warm up auto-created logical/physical tables and metadata before comparing data-region WAL.
|
|
for (batch_id, hint) in [None, Some("true"), Some("false"), None]
|
|
.into_iter()
|
|
.enumerate()
|
|
{
|
|
let before = env.flush_and_snapshot_wal().await;
|
|
let batch = gauge_arrow_batch(batch_id as i64, false);
|
|
let request = with_skip_wal_hint(futures::stream::iter([batch]), hint);
|
|
let mut response = client.arrow_metrics(request).await.unwrap().into_inner();
|
|
let status = response.message().await.unwrap().unwrap();
|
|
assert_eq!(
|
|
status.status_code,
|
|
ArrowStatusCode::Ok as i32,
|
|
"{}",
|
|
status.status_message
|
|
);
|
|
assert!(response.message().await.unwrap().is_none());
|
|
if batch_id != 0 {
|
|
assert_wal_delta(
|
|
&before,
|
|
&env.flush_and_snapshot_wal().await,
|
|
hint == Some("true"),
|
|
);
|
|
}
|
|
}
|
|
server.shutdown().await.unwrap();
|
|
env.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_otel_arrow_delta_histogram(store_type: StorageType) {
|
|
let (_instance, server) =
|
|
setup_grpc_server(store_type, "test_otel_arrow_delta_histogram").await;
|
|
let addr = server.bind_addr().unwrap().to_string();
|
|
let mut client = ArrowMetricsServiceClient::connect(format!("http://{addr}"))
|
|
.await
|
|
.unwrap();
|
|
let valid = (3, &[1, 2][..], &[1.0][..]);
|
|
let malformed = (1, &[1][..], &[1.0][..]);
|
|
let request = Request::new(futures::stream::iter([
|
|
delta_histogram_arrow_batch(10, &[valid, malformed]),
|
|
delta_histogram_arrow_batch(11, &[malformed]),
|
|
delta_histogram_arrow_batch(12, &[valid]),
|
|
gauge_arrow_batch(13, true),
|
|
gauge_arrow_batch(14, false),
|
|
]));
|
|
let mut response = client.arrow_metrics(request).await.unwrap().into_inner();
|
|
|
|
let mixed = response.message().await.unwrap().unwrap();
|
|
assert_eq!(10, mixed.batch_id);
|
|
assert_eq!(ArrowStatusCode::Ok as i32, mixed.status_code);
|
|
assert!(mixed.status_message.contains("bucket_counts length"));
|
|
|
|
let rejected = response.message().await.unwrap().unwrap();
|
|
assert_eq!(11, rejected.batch_id);
|
|
assert_eq!(
|
|
ArrowStatusCode::InvalidArgument as i32,
|
|
rejected.status_code
|
|
);
|
|
assert!(rejected.status_message.contains("bucket_counts length"));
|
|
|
|
let later_valid = response.message().await.unwrap().unwrap();
|
|
assert_eq!(12, later_valid.batch_id);
|
|
assert_eq!(ArrowStatusCode::Ok as i32, later_valid.status_code);
|
|
assert!(later_valid.status_message.is_empty());
|
|
|
|
let collision = response.message().await.unwrap().unwrap();
|
|
assert_eq!(13, collision.batch_id);
|
|
assert_eq!(
|
|
ArrowStatusCode::InvalidArgument as i32,
|
|
collision.status_code
|
|
);
|
|
assert!(collision.status_message.contains("reserved label"));
|
|
|
|
let after_collision = response.message().await.unwrap().unwrap();
|
|
assert_eq!(14, after_collision.batch_id);
|
|
assert_eq!(ArrowStatusCode::Ok as i32, after_collision.status_code);
|
|
assert!(after_collision.status_message.is_empty());
|
|
let _ = server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_otel_arrow_exemplars_without_metric_engine(store_type: StorageType) {
|
|
use datatypes::arrow::array::BinaryArray;
|
|
use otel_arrow_rust::Consumer;
|
|
use otel_arrow_rust::proto::opentelemetry::metrics::v1::metric;
|
|
|
|
let (_instance, server) = setup_grpc_server(
|
|
store_type,
|
|
"test_otel_arrow_exemplars_without_metric_engine",
|
|
)
|
|
.await;
|
|
let addr = server.bind_addr().unwrap().to_string();
|
|
let db = Database::new_with_dbname(
|
|
format!("{DEFAULT_CATALOG_NAME}-{DEFAULT_SCHEMA_NAME}"),
|
|
Client::with_urls(vec![addr.clone()]),
|
|
);
|
|
let mut client = ArrowMetricsServiceClient::connect(format!("http://{addr}"))
|
|
.await
|
|
.unwrap();
|
|
let mut batch = gauge_arrow_batch(20, false);
|
|
let exemplars = ArrowRecordBatch::try_from_iter(vec![
|
|
(
|
|
arrow_consts::PARENT_ID,
|
|
Arc::new(UInt32Array::from(vec![0_u32])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::TIME_UNIX_NANO,
|
|
Arc::new(TimestampNanosecondArray::from(vec![2_000_000_000])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::DOUBLE_VALUE,
|
|
Arc::new(Float64Array::from(vec![0.5])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::SPAN_ID,
|
|
Arc::new(BinaryArray::from_vec(vec![&[1_u8; 8]])) as ArrayRef,
|
|
),
|
|
(
|
|
arrow_consts::TRACE_ID,
|
|
Arc::new(BinaryArray::from_vec(vec![&[2_u8; 16]])) as ArrayRef,
|
|
),
|
|
])
|
|
.unwrap();
|
|
batch.arrow_payloads.push(ArrowPayload {
|
|
schema_id: "exemplars-20".into(),
|
|
r#type: ArrowPayloadType::NumberDpExemplars as i32,
|
|
record: serialize_arrow_record_batch(&exemplars),
|
|
});
|
|
let decoded = Consumer::default()
|
|
.consume_metrics_batches(&mut batch.clone())
|
|
.unwrap();
|
|
let Some(metric::Data::Gauge(gauge)) =
|
|
&decoded.resource_metrics[0].scope_metrics[0].metrics[0].data
|
|
else {
|
|
panic!("expected a gauge")
|
|
};
|
|
assert_eq!(gauge.data_points[0].exemplars.len(), 1);
|
|
|
|
let request = Request::new(futures::stream::iter([
|
|
batch,
|
|
delta_histogram_arrow_batch(21, &[(3, &[1, 2], &[1.0])]),
|
|
]));
|
|
let mut response = client.arrow_metrics(request).await.unwrap().into_inner();
|
|
for batch_id in [20, 21] {
|
|
let status = response.message().await.unwrap().unwrap();
|
|
assert_eq!(status.batch_id, batch_id);
|
|
assert_eq!(status.status_code, ArrowStatusCode::Ok as i32);
|
|
assert!(status.status_message.is_empty());
|
|
}
|
|
assert!(response.message().await.unwrap().is_none());
|
|
for (sql, count) in [
|
|
("SELECT COUNT(*) AS c FROM otel_arrow_gauge", 1),
|
|
(
|
|
"SELECT COUNT(*) AS c FROM otel_arrow_delta_histogram_bucket",
|
|
2,
|
|
),
|
|
(
|
|
&format!(
|
|
"SELECT COUNT(*) AS c FROM information_schema.tables WHERE table_name='{METRICS_EXEMPLARS_TABLE}'"
|
|
),
|
|
0,
|
|
),
|
|
] {
|
|
let OutputData::Stream(stream) = db.sql(sql).await.unwrap().data else {
|
|
panic!("expected a stream")
|
|
};
|
|
assert_eq!(
|
|
RecordBatches::try_collect(stream)
|
|
.await
|
|
.unwrap()
|
|
.pretty_print()
|
|
.unwrap(),
|
|
format!("+---+\n| c |\n+---+\n| {count} |\n+---+"),
|
|
);
|
|
}
|
|
server.shutdown().await.unwrap();
|
|
}
|
|
|
|
pub async fn test_otel_arrow_exponential_histogram(store_type: StorageType) {
|
|
let (_instance, server) =
|
|
setup_grpc_server(store_type, "test_otel_arrow_exponential_histogram").await;
|
|
let addr = server.bind_addr().unwrap().to_string();
|
|
let mut client = ArrowMetricsServiceClient::connect(format!("http://{addr}"))
|
|
.await
|
|
.unwrap();
|
|
let batch = exponential_histogram_arrow_batch(0, &[0]);
|
|
let request = Request::new(futures::stream::once(async { batch }));
|
|
let mut response = client.arrow_metrics(request).await.unwrap().into_inner();
|
|
let status = response.message().await.unwrap().unwrap();
|
|
assert_eq!(0, status.batch_id);
|
|
assert_eq!(ArrowStatusCode::InvalidArgument as i32, status.status_code);
|
|
assert!(status.status_message.contains("omits zero_threshold"));
|
|
let _ = server.shutdown().await;
|
|
}
|
|
|
|
fn basic_auth(username: &str, password: &str) -> String {
|
|
format!("Basic {}", basic_auth_credentials(username, password))
|
|
}
|
|
|
|
fn basic_auth_credentials(username: &str, password: &str) -> String {
|
|
BASE64_STANDARD.encode(format!("{username}:{password}"))
|
|
}
|
|
|
|
pub async fn test_auto_create_table(store_type: StorageType) {
|
|
let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_auto_create_table").await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, grpc_client);
|
|
insert_and_assert(&db).await;
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_auto_create_table_with_hints(store_type: StorageType) {
|
|
let (_db, fe_grpc_server) =
|
|
setup_grpc_server(store_type, "test_auto_create_table_with_hints").await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, grpc_client);
|
|
insert_with_hints_and_assert(&db).await;
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
/// When the frontend global switch disables auto table creation, a write to a
|
|
/// missing table must fail even if the request sets `auto_create_table=true`,
|
|
/// proving the global config is an upper bound that hints cannot bypass.
|
|
pub async fn test_auto_create_table_disabled_by_config(store_type: StorageType) {
|
|
let (_db, fe_grpc_server) = setup_grpc_server_with_auto_create_table_disabled(
|
|
store_type,
|
|
"test_auto_create_table_disabled_by_config",
|
|
)
|
|
.await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, grpc_client);
|
|
|
|
// Plain row insert to a missing table: must fail even with `auto_create_table=true`.
|
|
let (host, cpu, mem, ts) = expect_data();
|
|
let request = InsertRequest {
|
|
table_name: "demo".to_string(),
|
|
columns: vec![host, cpu, mem, ts],
|
|
row_count: 4,
|
|
};
|
|
let result = db
|
|
.insert_with_hints(
|
|
InsertRequests {
|
|
inserts: vec![request],
|
|
},
|
|
&[("auto_create_table", "true")],
|
|
)
|
|
.await;
|
|
let err = result.unwrap_err().to_string();
|
|
assert!(
|
|
err.contains("does not exist") && err.contains("disabled by frontend config"),
|
|
"unexpected error: {err}"
|
|
);
|
|
|
|
// Metric path (via `physical_table` hint): must also fail without leaking the physical table.
|
|
let (host, cpu, mem, ts) = expect_data();
|
|
let request = InsertRequest {
|
|
table_name: "demo_metric".to_string(),
|
|
columns: vec![host, cpu, mem, ts],
|
|
row_count: 4,
|
|
};
|
|
let result = db
|
|
.insert_with_hints(
|
|
InsertRequests {
|
|
inserts: vec![request],
|
|
},
|
|
&[
|
|
("auto_create_table", "true"),
|
|
("physical_table", "greptime_physical_table"),
|
|
],
|
|
)
|
|
.await;
|
|
let err = result.unwrap_err().to_string();
|
|
assert!(
|
|
err.contains("does not exist") && err.contains("disabled by frontend config"),
|
|
"unexpected error: {err}"
|
|
);
|
|
|
|
// The physical table must not have been created before the failure.
|
|
let output = db.sql("SHOW TABLES").await.unwrap();
|
|
let record_batches = match output.data {
|
|
OutputData::RecordBatches(record_batches) => record_batches,
|
|
OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(),
|
|
OutputData::AffectedRows(_) => unreachable!(),
|
|
};
|
|
let tables = record_batches.pretty_print().unwrap();
|
|
assert!(
|
|
!tables.contains("greptime_physical_table"),
|
|
"physical table leaked despite disabled auto-create:\n{tables}"
|
|
);
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_private_system_tables_auto_create_table_with_global_disabled(
|
|
store_type: StorageType,
|
|
) {
|
|
let (_db, fe_grpc_server) = setup_grpc_server_with_auto_create_table_disabled(
|
|
store_type,
|
|
"test_private_system_tables_auto_create_table_with_global_disabled",
|
|
)
|
|
.await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new(
|
|
DEFAULT_CATALOG_NAME,
|
|
DEFAULT_PRIVATE_SCHEMA_NAME,
|
|
grpc_client,
|
|
);
|
|
let (host, cpu, mem, ts) = expect_data();
|
|
|
|
for table_name in [
|
|
DEFAULT_EVENTS_TABLE_NAME,
|
|
SLOW_QUERY_TABLE_NAME,
|
|
REGION_STATS_HISTORY_TABLE_NAME,
|
|
] {
|
|
let result = db
|
|
.insert(InsertRequests {
|
|
inserts: vec![InsertRequest {
|
|
table_name: table_name.to_string(),
|
|
columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()],
|
|
row_count: 4,
|
|
}],
|
|
})
|
|
.await;
|
|
assert_eq!(result.unwrap(), 4);
|
|
}
|
|
|
|
let load = Column {
|
|
column_name: "load".to_string(),
|
|
values: Some(column::Values {
|
|
f64_values: vec![0.4, 0.5, 0.6, 0.7],
|
|
..Default::default()
|
|
}),
|
|
semantic_type: SemanticType::Field as i32,
|
|
datatype: ColumnDataType::Float64 as i32,
|
|
..Default::default()
|
|
};
|
|
let result = db
|
|
.insert(InsertRequests {
|
|
inserts: vec![InsertRequest {
|
|
table_name: DEFAULT_EVENTS_TABLE_NAME.to_string(),
|
|
columns: vec![host, cpu, mem, ts, load],
|
|
row_count: 4,
|
|
}],
|
|
})
|
|
.await;
|
|
assert_eq!(result.unwrap(), 4);
|
|
|
|
let output = db
|
|
.sql(format!("SHOW CREATE TABLE {DEFAULT_EVENTS_TABLE_NAME}"))
|
|
.await
|
|
.unwrap();
|
|
let record_batches = match output.data {
|
|
OutputData::RecordBatches(record_batches) => record_batches,
|
|
OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(),
|
|
OutputData::AffectedRows(_) => unreachable!(),
|
|
};
|
|
assert!(record_batches.pretty_print().unwrap().contains("\"load\""));
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_private_system_tables_bypass_auto_create_hint(store_type: StorageType) {
|
|
let (_db, fe_grpc_server) = setup_grpc_server(
|
|
store_type,
|
|
"test_private_system_tables_bypass_auto_create_hint",
|
|
)
|
|
.await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new(
|
|
DEFAULT_CATALOG_NAME,
|
|
DEFAULT_PRIVATE_SCHEMA_NAME,
|
|
grpc_client,
|
|
);
|
|
let (host, cpu, mem, ts) = expect_data();
|
|
|
|
for table_name in [
|
|
DEFAULT_EVENTS_TABLE_NAME,
|
|
SLOW_QUERY_TABLE_NAME,
|
|
REGION_STATS_HISTORY_TABLE_NAME,
|
|
] {
|
|
let result = db
|
|
.insert_with_hints(
|
|
InsertRequests {
|
|
inserts: vec![InsertRequest {
|
|
table_name: table_name.to_string(),
|
|
columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()],
|
|
row_count: 4,
|
|
}],
|
|
},
|
|
&[("auto_create_table", "false")],
|
|
)
|
|
.await;
|
|
assert_eq!(result.unwrap(), 4);
|
|
}
|
|
|
|
let ordinary_table = "ordinary_private_table";
|
|
let result = db
|
|
.insert_with_hints(
|
|
InsertRequests {
|
|
inserts: vec![
|
|
InsertRequest {
|
|
table_name: DEFAULT_EVENTS_TABLE_NAME.to_string(),
|
|
columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()],
|
|
row_count: 4,
|
|
},
|
|
InsertRequest {
|
|
table_name: ordinary_table.to_string(),
|
|
columns: vec![host, cpu, mem, ts],
|
|
row_count: 4,
|
|
},
|
|
],
|
|
},
|
|
&[("auto_create_table", "false")],
|
|
)
|
|
.await;
|
|
let err = result.unwrap_err().to_string();
|
|
assert!(
|
|
err.contains(ordinary_table) && err.contains("auto_create_table"),
|
|
"unexpected error: {err}"
|
|
);
|
|
|
|
let output = db.sql("SHOW TABLES").await.unwrap();
|
|
let record_batches = match output.data {
|
|
OutputData::RecordBatches(record_batches) => record_batches,
|
|
OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(),
|
|
OutputData::AffectedRows(_) => unreachable!(),
|
|
};
|
|
assert!(
|
|
!record_batches
|
|
.pretty_print()
|
|
.unwrap()
|
|
.contains(ordinary_table)
|
|
);
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
fn expect_data() -> (Column, Column, Column, Column) {
|
|
// testing data:
|
|
let expected_host_col = Column {
|
|
column_name: "host".to_string(),
|
|
values: Some(column::Values {
|
|
string_values: vec!["host1", "host2", "host3", "host4"]
|
|
.into_iter()
|
|
.map(|s| s.to_string())
|
|
.collect(),
|
|
..Default::default()
|
|
}),
|
|
semantic_type: SemanticType::Tag as i32,
|
|
datatype: ColumnDataType::String as i32,
|
|
..Default::default()
|
|
};
|
|
let expected_cpu_col = Column {
|
|
column_name: "cpu".to_string(),
|
|
values: Some(column::Values {
|
|
f64_values: vec![0.31, 0.41, 0.2],
|
|
..Default::default()
|
|
}),
|
|
null_mask: vec![2],
|
|
semantic_type: SemanticType::Field as i32,
|
|
datatype: ColumnDataType::Float64 as i32,
|
|
..Default::default()
|
|
};
|
|
let expected_mem_col = Column {
|
|
column_name: "memory".to_string(),
|
|
values: Some(column::Values {
|
|
f64_values: vec![0.1, 0.2, 0.3],
|
|
..Default::default()
|
|
}),
|
|
null_mask: vec![4],
|
|
semantic_type: SemanticType::Field as i32,
|
|
datatype: ColumnDataType::Float64 as i32,
|
|
..Default::default()
|
|
};
|
|
let expected_ts_col = Column {
|
|
column_name: "ts".to_string(),
|
|
values: Some(column::Values {
|
|
timestamp_millisecond_values: vec![100, 101, 102, 103],
|
|
..Default::default()
|
|
}),
|
|
semantic_type: SemanticType::Timestamp as i32,
|
|
datatype: ColumnDataType::TimestampMillisecond as i32,
|
|
..Default::default()
|
|
};
|
|
|
|
(
|
|
expected_host_col,
|
|
expected_cpu_col,
|
|
expected_mem_col,
|
|
expected_ts_col,
|
|
)
|
|
}
|
|
|
|
fn with_skip_wal_hint<T>(body: T, hint: Option<&str>) -> Request<T> {
|
|
let mut request = Request::new(body);
|
|
if let Some(hint) = hint {
|
|
request.metadata_mut().insert(
|
|
HINTS_KEY,
|
|
format!("{INSERT_SKIP_WAL_HINT}={hint}").parse().unwrap(),
|
|
);
|
|
}
|
|
request
|
|
}
|
|
|
|
fn skip_wal_insert_body(columnar: bool) -> RequestBody {
|
|
if columnar {
|
|
RequestBody::Inserts(InsertRequests {
|
|
inserts: vec![InsertRequest {
|
|
table_name: "skip_wal_grpc".to_string(),
|
|
row_count: 1,
|
|
columns: vec![Column {
|
|
column_name: "ts".to_string(),
|
|
semantic_type: SemanticType::Timestamp as i32,
|
|
datatype: ColumnDataType::TimestampMillisecond as i32,
|
|
values: Some(column::Values {
|
|
timestamp_millisecond_values: vec![1000],
|
|
..Default::default()
|
|
}),
|
|
..Default::default()
|
|
}],
|
|
}],
|
|
})
|
|
} else {
|
|
RequestBody::RowInserts(RowInsertRequests {
|
|
inserts: vec![RowInsertRequest {
|
|
table_name: "skip_wal_grpc".to_string(),
|
|
rows: Some(Rows {
|
|
schema: vec![ColumnSchema {
|
|
column_name: "ts".to_string(),
|
|
semantic_type: SemanticType::Timestamp as i32,
|
|
datatype: ColumnDataType::TimestampMillisecond as i32,
|
|
..Default::default()
|
|
}],
|
|
rows: vec![Row {
|
|
values: vec![Value {
|
|
value_data: Some(ValueData::TimestampMillisecondValue(1000)),
|
|
}],
|
|
}],
|
|
}),
|
|
}],
|
|
})
|
|
}
|
|
}
|
|
|
|
#[apply(both_deployment_cases)]
|
|
async fn test_skip_wal_grpc_unary_stream_and_flight_sql(distributed: bool) {
|
|
let mut env = MockInstanceImpl::new("skip_wal_grpc_protocols", distributed).await;
|
|
let server = setup_grpc_server_for_frontend_instance(env.frontend(), None).await;
|
|
let addr = server.bind_addr().unwrap().to_string();
|
|
let mut grpc = GreptimeDatabaseClient::connect(format!("http://{addr}"))
|
|
.await
|
|
.unwrap();
|
|
let database = Database::new_with_dbname("greptime-public", Client::with_urls(vec![addr]));
|
|
database
|
|
.sql("CREATE TABLE skip_wal_grpc (ts TIMESTAMP TIME INDEX)")
|
|
.await
|
|
.unwrap();
|
|
|
|
for streaming in [false, true] {
|
|
for columnar in [false, true] {
|
|
for hint in [Some("true"), Some("false"), None] {
|
|
let before = env.flush_and_snapshot_wal().await;
|
|
let request = GreptimeRequest {
|
|
header: Some(RequestHeader {
|
|
catalog: "greptime".to_string(),
|
|
schema: "public".to_string(),
|
|
..Default::default()
|
|
}),
|
|
request: Some(skip_wal_insert_body(columnar)),
|
|
};
|
|
let response = if streaming {
|
|
grpc.handle_requests(with_skip_wal_hint(
|
|
futures::stream::iter([request.clone(), request]),
|
|
hint,
|
|
))
|
|
.await
|
|
.unwrap()
|
|
} else {
|
|
grpc.handle(with_skip_wal_hint(request, hint))
|
|
.await
|
|
.unwrap()
|
|
}
|
|
.into_inner();
|
|
let expected_rows = if streaming { 2 } else { 1 };
|
|
assert!(
|
|
matches!(response.response, Some(ResponseBody::AffectedRows(rows)) if rows.value == expected_rows)
|
|
);
|
|
assert_wal_delta(
|
|
&before,
|
|
&env.flush_and_snapshot_wal().await,
|
|
hint == Some("true"),
|
|
);
|
|
}
|
|
}
|
|
}
|
|
for hint in [Some("true"), Some("false"), None] {
|
|
let before = env.flush_and_snapshot_wal().await;
|
|
let hints = hint
|
|
.map(|value| vec![(INSERT_SKIP_WAL_HINT, value)])
|
|
.unwrap_or_default();
|
|
database
|
|
.sql_with_hint("INSERT INTO skip_wal_grpc VALUES (1000)", &hints)
|
|
.await
|
|
.unwrap();
|
|
assert_wal_delta(
|
|
&before,
|
|
&env.flush_and_snapshot_wal().await,
|
|
hint == Some("true"),
|
|
);
|
|
}
|
|
// Strict validation must reject the request before any data write.
|
|
let before = env.flush_and_snapshot_wal().await;
|
|
assert!(
|
|
database
|
|
.sql_with_hint(
|
|
"INSERT INTO skip_wal_grpc VALUES (1000)",
|
|
&[(INSERT_SKIP_WAL_HINT, "yes")]
|
|
)
|
|
.await
|
|
.is_err()
|
|
);
|
|
assert_eq!(before, env.flush_and_snapshot_wal().await);
|
|
server.shutdown().await.unwrap();
|
|
env.shutdown().await;
|
|
}
|
|
|
|
#[apply(both_deployment_cases)]
|
|
async fn test_grpc_json2_row_inserts(distributed: bool) {
|
|
let mut env = MockInstanceImpl::new("grpc_json2_row_inserts", distributed).await;
|
|
let server = setup_grpc_server_for_frontend_instance(env.frontend(), None).await;
|
|
let addr = server.bind_addr().unwrap().to_string();
|
|
let database = Database::new_with_dbname("greptime-public", Client::with_urls(vec![addr]));
|
|
database
|
|
.sql(
|
|
"CREATE TABLE grpc_json2 (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, j JSON2) \
|
|
WITH (append_mode='true', 'memtable.type'='bulk')",
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
let (datatype, extension) =
|
|
api::helper::ColumnDataTypeWrapper::try_from(datatypes::prelude::ConcreteDataType::json2(
|
|
datatypes::types::json_type::JsonNativeType::Null,
|
|
))
|
|
.unwrap()
|
|
.into_parts();
|
|
|
|
// Send the all-NULL JSON2 batch separately to check its type marker.
|
|
for (index, (host, payload)) in [
|
|
(
|
|
"host-a",
|
|
Some(serde_json::json!({
|
|
"active": true,
|
|
"nested": {"items": [1, "two", null, {"ok": false}], "ratio": 1.5},
|
|
"tags": ["api", "prod"]
|
|
})),
|
|
),
|
|
(
|
|
"host-b",
|
|
Some(serde_json::json!({
|
|
"active": false,
|
|
"nested": {"items": [-2, "three", null, {"ok": true}], "ratio": -0.25},
|
|
"tags": ["worker"]
|
|
})),
|
|
),
|
|
("host-a", None),
|
|
]
|
|
.into_iter()
|
|
.enumerate()
|
|
{
|
|
let output = database
|
|
.row_inserts(RowInsertRequests {
|
|
inserts: vec![RowInsertRequest {
|
|
table_name: "grpc_json2".into(),
|
|
rows: Some(Rows {
|
|
schema: vec![
|
|
ColumnSchema {
|
|
column_name: "ts".into(),
|
|
datatype: ColumnDataType::TimestampMillisecond as i32,
|
|
semantic_type: SemanticType::Timestamp as i32,
|
|
..Default::default()
|
|
},
|
|
ColumnSchema {
|
|
column_name: "host".into(),
|
|
datatype: ColumnDataType::String as i32,
|
|
semantic_type: SemanticType::Tag as i32,
|
|
..Default::default()
|
|
},
|
|
ColumnSchema {
|
|
column_name: "j".into(),
|
|
datatype: datatype as i32,
|
|
datatype_extension: extension.clone(),
|
|
semantic_type: SemanticType::Field as i32,
|
|
..Default::default()
|
|
},
|
|
],
|
|
rows: vec![Row {
|
|
values: vec![
|
|
Value {
|
|
value_data: Some(ValueData::TimestampMillisecondValue(
|
|
index as i64,
|
|
)),
|
|
},
|
|
Value {
|
|
value_data: Some(ValueData::StringValue(host.into())),
|
|
},
|
|
Value {
|
|
value_data: payload.map(|x| {
|
|
ValueData::JsonValue(api::helper::encode_json_value(
|
|
x.into(),
|
|
))
|
|
}),
|
|
},
|
|
],
|
|
}],
|
|
}),
|
|
}],
|
|
})
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(output, 1);
|
|
}
|
|
|
|
let output = database
|
|
.sql("SELECT host, j FROM grpc_json2 ORDER BY ts")
|
|
.await
|
|
.unwrap();
|
|
let batches = match output.data {
|
|
OutputData::RecordBatches(batches) => batches,
|
|
OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(),
|
|
OutputData::AffectedRows(_) => unreachable!(),
|
|
};
|
|
let pretty = batches.pretty_print().unwrap();
|
|
let expected = r#"+--------+---------------------------------------------------------------------------------------------------+
|
|
| host | j |
|
|
+--------+---------------------------------------------------------------------------------------------------+
|
|
| host-a | {"active":true,"nested":{"items":[1,"two",null,{"ok":false}],"ratio":1.5},"tags":["api","prod"]} |
|
|
| host-b | {"active":false,"nested":{"items":[-2,"three",null,{"ok":true}],"ratio":-0.25},"tags":["worker"]} |
|
|
| host-a | |
|
|
+--------+---------------------------------------------------------------------------------------------------+"#;
|
|
assert_eq!(pretty, expected);
|
|
server.shutdown().await.unwrap();
|
|
env.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_insert_and_select(store_type: StorageType) {
|
|
common_telemetry::init_default_ut_logging();
|
|
let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_insert_and_select").await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, grpc_client);
|
|
|
|
// create
|
|
let expr = testing_create_expr();
|
|
let result = db.create(expr).await.unwrap();
|
|
assert!(matches!(result.data, OutputData::AffectedRows(0)));
|
|
|
|
//alter
|
|
let add_column = ColumnDef {
|
|
name: "test_column".to_string(),
|
|
data_type: ColumnDataType::Int64.into(),
|
|
is_nullable: true,
|
|
default_constraint: vec![],
|
|
semantic_type: SemanticType::Field as i32,
|
|
..Default::default()
|
|
};
|
|
let kind = Kind::AddColumns(AddColumns {
|
|
add_columns: vec![AddColumn {
|
|
column_def: Some(add_column),
|
|
location: None,
|
|
add_if_not_exists: false,
|
|
}],
|
|
});
|
|
let expr = AlterTableExpr {
|
|
catalog_name: DEFAULT_CATALOG_NAME.to_string(),
|
|
schema_name: DEFAULT_SCHEMA_NAME.to_string(),
|
|
table_name: "demo".to_string(),
|
|
kind: Some(kind),
|
|
};
|
|
let result = db.alter(expr).await.unwrap();
|
|
assert!(matches!(result.data, OutputData::AffectedRows(0)));
|
|
|
|
// insert
|
|
insert_and_assert(&db).await;
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
async fn insert_with_hints_and_assert(db: &Database) {
|
|
// testing data:
|
|
let (expected_host_col, expected_cpu_col, expected_mem_col, expected_ts_col) = expect_data();
|
|
|
|
let request = InsertRequest {
|
|
table_name: "demo".to_string(),
|
|
columns: vec![
|
|
expected_host_col.clone(),
|
|
expected_cpu_col.clone(),
|
|
expected_mem_col.clone(),
|
|
expected_ts_col.clone(),
|
|
],
|
|
row_count: 4,
|
|
};
|
|
let result = db
|
|
.insert_with_hints(
|
|
InsertRequests {
|
|
inserts: vec![request],
|
|
},
|
|
&[("append_mode", "true")],
|
|
)
|
|
.await;
|
|
assert_eq!(result.unwrap(), 4);
|
|
|
|
// show table
|
|
let output = db.sql("SHOW CREATE TABLE demo;").await.unwrap();
|
|
|
|
let record_batches = match output.data {
|
|
OutputData::RecordBatches(record_batches) => record_batches,
|
|
OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(),
|
|
OutputData::AffectedRows(_) => unreachable!(),
|
|
};
|
|
|
|
let pretty = record_batches.pretty_print().unwrap();
|
|
let expected = "\
|
|
+-------+---------------------------------------+
|
|
| Table | Create Table |
|
|
+-------+---------------------------------------+
|
|
| demo | CREATE TABLE IF NOT EXISTS \"demo\" ( |
|
|
| | \"host\" STRING NULL, |
|
|
| | \"cpu\" DOUBLE NULL, |
|
|
| | \"memory\" DOUBLE NULL, |
|
|
| | \"ts\" TIMESTAMP(3) NOT NULL, |
|
|
| | TIME INDEX (\"ts\"), |
|
|
| | PRIMARY KEY (\"host\") |
|
|
| | ) |
|
|
| | |
|
|
| | ENGINE=mito |
|
|
| | WITH( |
|
|
| | 'comment' = 'Created on insertion', |
|
|
| | append_mode = 'true' |
|
|
| | ) |
|
|
+-------+---------------------------------------+\
|
|
";
|
|
assert_eq!(pretty, expected);
|
|
|
|
// testing data with ttl=instant and auto_create_table = true can be handled correctly
|
|
let (expected_host_col, expected_cpu_col, expected_mem_col, expected_ts_col) = expect_data();
|
|
|
|
let request = InsertRequest {
|
|
table_name: "demo1".to_string(),
|
|
columns: vec![
|
|
expected_host_col.clone(),
|
|
expected_cpu_col.clone(),
|
|
expected_mem_col.clone(),
|
|
expected_ts_col.clone(),
|
|
],
|
|
row_count: 4,
|
|
};
|
|
let result = db
|
|
.insert_with_hints(
|
|
InsertRequests {
|
|
inserts: vec![request],
|
|
},
|
|
&[("auto_create_table", "true"), ("ttl", "instant")],
|
|
)
|
|
.await;
|
|
assert_eq!(result.unwrap(), 0);
|
|
|
|
// check table is empty
|
|
let output = db.sql("SELECT * FROM demo1").await.unwrap();
|
|
|
|
let record_batches = match output.data {
|
|
OutputData::RecordBatches(record_batches) => record_batches,
|
|
OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(),
|
|
OutputData::AffectedRows(_) => unreachable!(),
|
|
};
|
|
|
|
assert!(record_batches.iter().all(|r| r.num_rows() == 0));
|
|
}
|
|
|
|
async fn insert_and_assert(db: &Database) {
|
|
// testing data:
|
|
let (expected_host_col, expected_cpu_col, expected_mem_col, expected_ts_col) = expect_data();
|
|
|
|
let request = InsertRequest {
|
|
table_name: "demo".to_string(),
|
|
columns: vec![
|
|
expected_host_col.clone(),
|
|
expected_cpu_col.clone(),
|
|
expected_mem_col.clone(),
|
|
expected_ts_col.clone(),
|
|
],
|
|
row_count: 4,
|
|
};
|
|
let result = db
|
|
.insert(InsertRequests {
|
|
inserts: vec![request],
|
|
})
|
|
.await;
|
|
assert_eq!(result.unwrap(), 4);
|
|
|
|
let result = db
|
|
.sql(
|
|
"INSERT INTO demo(host, cpu, memory, ts) VALUES \
|
|
('host5', 66.6, 1024, 1672201027000),\
|
|
('host6', 88.8, 333.3, 1672201028000)",
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert!(matches!(result.data, OutputData::AffectedRows(2)));
|
|
|
|
// select
|
|
let output = db
|
|
.sql("SELECT host, cpu, memory, ts FROM demo order by host")
|
|
.await
|
|
.unwrap();
|
|
|
|
let record_batches = match output.data {
|
|
OutputData::RecordBatches(record_batches) => record_batches,
|
|
OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(),
|
|
OutputData::AffectedRows(_) => unreachable!(),
|
|
};
|
|
|
|
let pretty = record_batches.pretty_print().unwrap();
|
|
let expected = "\
|
|
+-------+------+--------+-------------------------+
|
|
| host | cpu | memory | ts |
|
|
+-------+------+--------+-------------------------+
|
|
| host1 | 0.31 | 0.1 | 1970-01-01T00:00:00.100 |
|
|
| host2 | | 0.2 | 1970-01-01T00:00:00.101 |
|
|
| host3 | 0.41 | | 1970-01-01T00:00:00.102 |
|
|
| host4 | 0.2 | 0.3 | 1970-01-01T00:00:00.103 |
|
|
| host5 | 66.6 | 1024.0 | 2022-12-28T04:17:07 |
|
|
| host6 | 88.8 | 333.3 | 2022-12-28T04:17:08 |
|
|
+-------+------+--------+-------------------------+\
|
|
";
|
|
assert_eq!(pretty, expected);
|
|
}
|
|
|
|
fn testing_create_expr() -> CreateTableExpr {
|
|
let column_defs = vec![
|
|
ColumnDef {
|
|
name: "host".to_string(),
|
|
data_type: ColumnDataType::String as i32,
|
|
is_nullable: false,
|
|
default_constraint: vec![],
|
|
semantic_type: SemanticType::Tag as i32,
|
|
..Default::default()
|
|
},
|
|
ColumnDef {
|
|
name: "cpu".to_string(),
|
|
data_type: ColumnDataType::Float64 as i32,
|
|
is_nullable: true,
|
|
default_constraint: vec![],
|
|
semantic_type: SemanticType::Field as i32,
|
|
..Default::default()
|
|
},
|
|
ColumnDef {
|
|
name: "memory".to_string(),
|
|
data_type: ColumnDataType::Float64 as i32,
|
|
is_nullable: true,
|
|
default_constraint: vec![],
|
|
semantic_type: SemanticType::Field as i32,
|
|
..Default::default()
|
|
},
|
|
ColumnDef {
|
|
name: "ts".to_string(),
|
|
data_type: ColumnDataType::TimestampMillisecond as i32, // timestamp
|
|
is_nullable: false,
|
|
default_constraint: vec![],
|
|
semantic_type: SemanticType::Timestamp as i32,
|
|
..Default::default()
|
|
},
|
|
];
|
|
CreateTableExpr {
|
|
catalog_name: "greptime".to_string(),
|
|
schema_name: "public".to_string(),
|
|
table_name: "demo".to_string(),
|
|
desc: "blabla little magic fairy".to_string(),
|
|
column_defs,
|
|
time_index: "ts".to_string(),
|
|
primary_keys: vec!["host".to_string()],
|
|
create_if_not_exists: true,
|
|
table_options: Default::default(),
|
|
table_id: None,
|
|
engine: MITO_ENGINE.to_string(),
|
|
}
|
|
}
|
|
|
|
pub async fn test_health_check(store_type: StorageType) {
|
|
let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_health_check").await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
grpc_client.health_check().await.unwrap();
|
|
|
|
let db = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, grpc_client);
|
|
assert!(db.sql("SHOW TABLES").await.is_ok());
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_prom_gateway_query(store_type: StorageType) {
|
|
common_telemetry::init_default_ut_logging();
|
|
|
|
// prepare connection
|
|
let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_prom_gateway_query").await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let db = Database::new(
|
|
DEFAULT_CATALOG_NAME,
|
|
DEFAULT_SCHEMA_NAME,
|
|
grpc_client.clone(),
|
|
);
|
|
let mut gateway_client = grpc_client.make_prometheus_gateway_client().unwrap();
|
|
|
|
// create table and insert data
|
|
assert!(matches!(
|
|
db.sql("CREATE TABLE test(i DOUBLE, j TIMESTAMP TIME INDEX, k STRING PRIMARY KEY);")
|
|
.await
|
|
.unwrap()
|
|
.data,
|
|
OutputData::AffectedRows(0)
|
|
));
|
|
assert!(matches!(
|
|
db.sql(r#"INSERT INTO test VALUES (1, 1, "a"), (1, 1, "b"), (2, 2, "a");"#)
|
|
.await
|
|
.unwrap()
|
|
.data,
|
|
OutputData::AffectedRows(3)
|
|
));
|
|
|
|
// Instant query using prometheus gateway service
|
|
let header = RequestHeader {
|
|
dbname: "public".to_string(),
|
|
..Default::default()
|
|
};
|
|
let instant_query = PromInstantQuery {
|
|
query: "test".to_string(),
|
|
time: "5".to_string(),
|
|
lookback: "5m".to_string(),
|
|
};
|
|
let instant_query_request = PromqlRequest {
|
|
header: Some(header.clone()),
|
|
promql: Some(Promql::InstantQuery(instant_query)),
|
|
};
|
|
let json_bytes = gateway_client
|
|
.handle(instant_query_request)
|
|
.await
|
|
.unwrap()
|
|
.into_inner()
|
|
.body;
|
|
let instant_query_result =
|
|
serde_json::from_slice::<PrometheusJsonResponse>(&json_bytes).unwrap();
|
|
assert_eq!(&instant_query_result.status, "success");
|
|
assert!(instant_query_result.error.is_none());
|
|
assert!(instant_query_result.error_type.is_none());
|
|
assert!(instant_query_result.warnings.is_none());
|
|
assert!(instant_query_result.resp_metrics.is_empty());
|
|
assert!(instant_query_result.status_code.is_none());
|
|
let PrometheusResponse::PromData(data) = instant_query_result.data else {
|
|
panic!("unexpected result data type")
|
|
};
|
|
assert_eq!(&data.result_type, "vector");
|
|
let PromQueryResult::Vector(mut vector) = data.result else {
|
|
panic!("unexpected result type")
|
|
};
|
|
|
|
vector.sort_unstable_by_key(|v| v.value.as_ref().map(|f| f.1.clone()));
|
|
|
|
assert_eq!(
|
|
vector,
|
|
vec![
|
|
PromSeriesVector {
|
|
metric: [
|
|
("__name__".to_string(), "test".to_string()),
|
|
("k".to_string(), "b".to_string()),
|
|
]
|
|
.into_iter()
|
|
.collect(),
|
|
value: Some((5.0, "1.0".to_string())),
|
|
..Default::default()
|
|
},
|
|
PromSeriesVector {
|
|
metric: [
|
|
("k".to_string(), "a".to_string()),
|
|
("__name__".to_string(), "test".to_string()),
|
|
]
|
|
.into_iter()
|
|
.collect(),
|
|
value: Some((5.0, "2.0".to_string())),
|
|
..Default::default()
|
|
},
|
|
]
|
|
);
|
|
|
|
// Range query using prometheus gateway service
|
|
let range_query = PromRangeQuery {
|
|
query: "test".to_string(),
|
|
start: "0".to_string(),
|
|
end: "10".to_string(),
|
|
step: "5s".to_string(),
|
|
lookback: "5m".to_string(),
|
|
};
|
|
let range_query_request: PromqlRequest = PromqlRequest {
|
|
header: Some(header.clone()),
|
|
promql: Some(Promql::RangeQuery(range_query)),
|
|
};
|
|
let json_bytes = gateway_client
|
|
.handle(range_query_request)
|
|
.await
|
|
.unwrap()
|
|
.into_inner()
|
|
.body;
|
|
let range_query_result = serde_json::from_slice::<PrometheusJsonResponse>(&json_bytes).unwrap();
|
|
|
|
assert_eq!(&range_query_result.status, "success");
|
|
assert!(range_query_result.error.is_none());
|
|
assert!(range_query_result.error_type.is_none());
|
|
assert!(range_query_result.warnings.is_none());
|
|
assert!(range_query_result.resp_metrics.is_empty());
|
|
assert!(range_query_result.status_code.is_none());
|
|
let PrometheusResponse::PromData(data) = range_query_result.data else {
|
|
panic!("unexpected result data type")
|
|
};
|
|
assert_eq!(&data.result_type, "matrix");
|
|
let PromQueryResult::Matrix(mut mat) = data.result else {
|
|
panic!("unexpected result type")
|
|
};
|
|
|
|
mat.sort_unstable_by_key(|v| serde_json::to_string(&v.values[0].1).unwrap());
|
|
|
|
assert_eq!(
|
|
mat,
|
|
vec![
|
|
PromSeriesMatrix {
|
|
metric: [
|
|
("__name__".to_string(), "test".to_string()),
|
|
("k".to_string(), "b".to_string()),
|
|
]
|
|
.into_iter()
|
|
.collect(),
|
|
values: vec![
|
|
(5.0, PromSampleValue::Text("1.0".to_string())),
|
|
(10.0, PromSampleValue::Text("1.0".to_string())),
|
|
],
|
|
..Default::default()
|
|
},
|
|
PromSeriesMatrix {
|
|
metric: [
|
|
("__name__".to_string(), "test".to_string()),
|
|
("k".to_string(), "a".to_string()),
|
|
]
|
|
.into_iter()
|
|
.collect(),
|
|
values: vec![
|
|
(5.0, PromSampleValue::Text("2.0".to_string())),
|
|
(10.0, PromSampleValue::Text("2.0".to_string())),
|
|
],
|
|
..Default::default()
|
|
},
|
|
]
|
|
);
|
|
|
|
// query nonexistent data
|
|
let range_query = PromRangeQuery {
|
|
query: "test".to_string(),
|
|
start: "1000000000".to_string(),
|
|
end: "1000001000".to_string(),
|
|
step: "5s".to_string(),
|
|
lookback: "5m".to_string(),
|
|
};
|
|
let range_query_request: PromqlRequest = PromqlRequest {
|
|
header: Some(header),
|
|
promql: Some(Promql::RangeQuery(range_query)),
|
|
};
|
|
let json_bytes = gateway_client
|
|
.handle(range_query_request)
|
|
.await
|
|
.unwrap()
|
|
.into_inner()
|
|
.body;
|
|
let range_query_result = serde_json::from_slice::<PrometheusJsonResponse>(&json_bytes).unwrap();
|
|
let expected = PrometheusJsonResponse {
|
|
status: "success".to_string(),
|
|
data: PrometheusResponse::PromData(PromData {
|
|
result_type: "matrix".to_string(),
|
|
result: PromQueryResult::Matrix(vec![]),
|
|
}),
|
|
error: None,
|
|
error_type: None,
|
|
warnings: None,
|
|
infos: None,
|
|
resp_metrics: Default::default(),
|
|
status_code: None,
|
|
};
|
|
assert_eq!(range_query_result, expected);
|
|
|
|
// clean up
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_grpc_timezone(store_type: StorageType) {
|
|
let config = GrpcServerConfig {
|
|
max_recv_message_size: 1024,
|
|
max_send_message_size: 1024,
|
|
..Default::default()
|
|
};
|
|
let (_db, fe_grpc_server) =
|
|
setup_grpc_server_with(store_type, "auto_create_table", None, Some(config), None).await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls(vec![addr]);
|
|
let mut db = Database::new_with_dbname(
|
|
format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
|
|
grpc_client,
|
|
);
|
|
db.set_timezone("Asia/Shanghai");
|
|
let sys1 = to_batch(db.sql("show variables system_time_zone;").await.unwrap()).await;
|
|
let user1 = to_batch(db.sql("show variables time_zone;").await.unwrap()).await;
|
|
db.set_timezone("");
|
|
let sys2 = to_batch(db.sql("show variables system_time_zone;").await.unwrap()).await;
|
|
let user2 = to_batch(db.sql("show variables time_zone;").await.unwrap()).await;
|
|
assert_eq!(sys1, sys2);
|
|
assert_eq!(
|
|
sys2,
|
|
"\
|
|
+------------------+
|
|
| SYSTEM_TIME_ZONE |
|
|
+------------------+
|
|
| UTC |
|
|
+------------------+"
|
|
);
|
|
assert_eq!(
|
|
user1,
|
|
"\
|
|
+---------------+
|
|
| TIME_ZONE |
|
|
+---------------+
|
|
| Asia/Shanghai |
|
|
+---------------+"
|
|
);
|
|
assert_eq!(
|
|
user2,
|
|
"\
|
|
+-----------+
|
|
| TIME_ZONE |
|
|
+-----------+
|
|
| UTC |
|
|
+-----------+"
|
|
);
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
async fn to_batch(output: Output) -> String {
|
|
match output.data {
|
|
OutputData::RecordBatches(batch) => batch,
|
|
OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(),
|
|
OutputData::AffectedRows(_) => unreachable!(),
|
|
}
|
|
.pretty_print()
|
|
.unwrap()
|
|
}
|
|
|
|
pub async fn test_grpc_tls_config(store_type: StorageType) {
|
|
let comm_dir = find_workspace_path("/src/common/grpc/tests/tls");
|
|
let ca_path = comm_dir.join("ca.pem").to_str().unwrap().to_string();
|
|
let server_cert_path = comm_dir.join("server.pem").to_str().unwrap().to_string();
|
|
let server_key_path = comm_dir.join("server.key").to_str().unwrap().to_string();
|
|
let client_cert_path = comm_dir.join("client.pem").to_str().unwrap().to_string();
|
|
let client_key_path = comm_dir.join("client.key").to_str().unwrap().to_string();
|
|
let client_corrupted = comm_dir.join("corrupted").to_str().unwrap().to_string();
|
|
|
|
let tls = TlsOption::new(
|
|
Some(TlsMode::Require),
|
|
Some(server_cert_path),
|
|
Some(server_key_path),
|
|
false,
|
|
);
|
|
let config = GrpcServerConfig {
|
|
max_recv_message_size: 1024,
|
|
max_send_message_size: 1024,
|
|
tls,
|
|
max_connection_age: None,
|
|
};
|
|
let (_db, fe_grpc_server) =
|
|
setup_grpc_server_with(store_type, "tls_create_table", None, Some(config), None).await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let mut client_tls = ClientTlsOption {
|
|
enabled: true,
|
|
server_ca_cert_path: Some(ca_path),
|
|
client_cert_path: Some(client_cert_path),
|
|
client_key_path: Some(client_key_path),
|
|
watch: false,
|
|
};
|
|
{
|
|
let grpc_client =
|
|
Client::with_tls_and_urls(vec![addr.clone()], client_tls.clone()).unwrap();
|
|
let db = Database::new_with_dbname(
|
|
format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
|
|
grpc_client,
|
|
);
|
|
db.sql("show tables;").await.unwrap();
|
|
}
|
|
// test corrupted client key
|
|
{
|
|
client_tls.client_key_path = Some(client_corrupted);
|
|
let grpc_client = Client::with_tls_and_urls(vec![addr], client_tls.clone()).unwrap();
|
|
let db = Database::new_with_dbname(
|
|
format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME),
|
|
grpc_client,
|
|
);
|
|
let re = db.sql("show tables;").await;
|
|
assert!(re.is_err());
|
|
}
|
|
// test grpc unsupported tls watch
|
|
{
|
|
let tls = TlsOption {
|
|
watch: true,
|
|
..Default::default()
|
|
};
|
|
let config = GrpcServerConfig {
|
|
max_recv_message_size: 1024,
|
|
max_send_message_size: 1024,
|
|
tls,
|
|
max_connection_age: None,
|
|
};
|
|
let runtime = Runtime::builder().build().unwrap();
|
|
let grpc_builder =
|
|
GrpcServerBuilder::new(config.clone(), runtime).with_tls_config(config.tls);
|
|
// ok but print warning
|
|
assert!(grpc_builder.is_ok());
|
|
}
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|
|
|
|
pub async fn test_grpc_memory_limit(store_type: StorageType) {
|
|
let config = GrpcServerConfig {
|
|
max_recv_message_size: 1024 * 1024,
|
|
max_send_message_size: 1024 * 1024,
|
|
tls: Default::default(),
|
|
max_connection_age: None,
|
|
};
|
|
|
|
// Create memory limiter with 2KB limit and fail-fast policy.
|
|
// Note: MemoryManager uses 1KB granularity (PermitGranularity::Kilobyte),
|
|
// so 2KB = 2 permits. Small/medium requests should fit, large should fail.
|
|
let memory_limiter = ServerMemoryLimiter::new(2048, OnExhaustedPolicy::Fail);
|
|
|
|
let (_db, fe_grpc_server) = setup_grpc_server_with(
|
|
store_type,
|
|
"test_grpc_memory_limit",
|
|
None,
|
|
Some(config),
|
|
Some(memory_limiter),
|
|
)
|
|
.await;
|
|
let addr = fe_grpc_server.bind_addr().unwrap().to_string();
|
|
|
|
let grpc_client = Client::with_urls([&addr]);
|
|
let db = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, grpc_client);
|
|
|
|
let table_name = "demo";
|
|
|
|
let column_schemas = vec![
|
|
ColumnDef {
|
|
name: "host".to_string(),
|
|
data_type: ColumnDataType::String as i32,
|
|
is_nullable: false,
|
|
default_constraint: vec![],
|
|
semantic_type: SemanticType::Tag as i32,
|
|
comment: String::new(),
|
|
datatype_extension: None,
|
|
options: None,
|
|
},
|
|
ColumnDef {
|
|
name: "ts".to_string(),
|
|
data_type: ColumnDataType::TimestampMillisecond as i32,
|
|
is_nullable: false,
|
|
default_constraint: vec![],
|
|
semantic_type: SemanticType::Timestamp as i32,
|
|
comment: String::new(),
|
|
datatype_extension: None,
|
|
options: None,
|
|
},
|
|
ColumnDef {
|
|
name: "cpu".to_string(),
|
|
data_type: ColumnDataType::Float64 as i32,
|
|
is_nullable: true,
|
|
default_constraint: vec![],
|
|
semantic_type: SemanticType::Field as i32,
|
|
comment: String::new(),
|
|
datatype_extension: None,
|
|
options: None,
|
|
},
|
|
];
|
|
|
|
let expr = CreateTableExpr {
|
|
catalog_name: DEFAULT_CATALOG_NAME.to_string(),
|
|
schema_name: DEFAULT_SCHEMA_NAME.to_string(),
|
|
table_name: table_name.to_string(),
|
|
desc: String::new(),
|
|
column_defs: column_schemas.clone(),
|
|
time_index: "ts".to_string(),
|
|
primary_keys: vec!["host".to_string()],
|
|
create_if_not_exists: true,
|
|
table_options: Default::default(),
|
|
table_id: None,
|
|
engine: MITO_ENGINE.to_string(),
|
|
};
|
|
|
|
db.create(expr).await.unwrap();
|
|
|
|
// Test that small request succeeds
|
|
let small_row_insert = RowInsertRequest {
|
|
table_name: table_name.to_owned(),
|
|
rows: Some(api::v1::Rows {
|
|
schema: column_schemas
|
|
.iter()
|
|
.map(|c| api::v1::ColumnSchema {
|
|
column_name: c.name.clone(),
|
|
datatype: c.data_type,
|
|
semantic_type: c.semantic_type,
|
|
datatype_extension: None,
|
|
options: None,
|
|
})
|
|
.collect(),
|
|
rows: vec![Row {
|
|
values: vec![
|
|
Value {
|
|
value_data: Some(ValueData::StringValue("host1".to_string())),
|
|
},
|
|
Value {
|
|
value_data: Some(ValueData::TimestampMillisecondValue(1000)),
|
|
},
|
|
Value {
|
|
value_data: Some(ValueData::F64Value(1.2)),
|
|
},
|
|
],
|
|
}],
|
|
}),
|
|
};
|
|
|
|
let result = db
|
|
.row_inserts(RowInsertRequests {
|
|
inserts: vec![small_row_insert],
|
|
})
|
|
.await;
|
|
assert!(result.is_ok());
|
|
|
|
// Test that medium request in the 200-1024 byte range should also succeed
|
|
// (due to 1KB granularity alignment)
|
|
let medium_rows: Vec<Row> = (0..5)
|
|
.map(|i| Row {
|
|
values: vec![
|
|
Value {
|
|
value_data: Some(ValueData::StringValue(format!("host{}", i))),
|
|
},
|
|
Value {
|
|
value_data: Some(ValueData::TimestampMillisecondValue(2000 + i)),
|
|
},
|
|
Value {
|
|
value_data: Some(ValueData::F64Value(i as f64 * 2.5)),
|
|
},
|
|
],
|
|
})
|
|
.collect();
|
|
|
|
let medium_row_insert = RowInsertRequest {
|
|
table_name: table_name.to_owned(),
|
|
rows: Some(api::v1::Rows {
|
|
schema: column_schemas
|
|
.iter()
|
|
.map(|c| api::v1::ColumnSchema {
|
|
column_name: c.name.clone(),
|
|
datatype: c.data_type,
|
|
semantic_type: c.semantic_type,
|
|
datatype_extension: None,
|
|
options: None,
|
|
})
|
|
.collect(),
|
|
rows: medium_rows,
|
|
}),
|
|
};
|
|
|
|
let result = db
|
|
.row_inserts(RowInsertRequests {
|
|
inserts: vec![medium_row_insert],
|
|
})
|
|
.await;
|
|
assert!(
|
|
result.is_ok(),
|
|
"Medium request (~500 bytes) should succeed within aligned 1KB limit"
|
|
);
|
|
|
|
// Test that large request exceeds limit (> 1KB aligned limit)
|
|
// Create a very large string to ensure we definitely exceed 1KB
|
|
// Use 100 rows with very long strings (>50 chars each) = definitely >5KB total
|
|
let large_rows: Vec<Row> = (0..100)
|
|
.map(|i| Row {
|
|
values: vec![
|
|
Value {
|
|
value_data: Some(ValueData::StringValue(format!(
|
|
"this_is_a_very_long_hostname_string_designed_to_make_the_request_exceed_memory_limit_row_number_{}",
|
|
i
|
|
))),
|
|
},
|
|
Value {
|
|
value_data: Some(ValueData::TimestampMillisecondValue(1000 + i)),
|
|
},
|
|
Value {
|
|
value_data: Some(ValueData::F64Value(i as f64 * 1.2)),
|
|
},
|
|
],
|
|
})
|
|
.collect();
|
|
|
|
let large_row_insert = RowInsertRequest {
|
|
table_name: table_name.to_owned(),
|
|
rows: Some(api::v1::Rows {
|
|
schema: column_schemas
|
|
.iter()
|
|
.map(|c| api::v1::ColumnSchema {
|
|
column_name: c.name.clone(),
|
|
datatype: c.data_type,
|
|
semantic_type: c.semantic_type,
|
|
datatype_extension: None,
|
|
options: None,
|
|
})
|
|
.collect(),
|
|
rows: large_rows,
|
|
}),
|
|
};
|
|
|
|
let result = db
|
|
.row_inserts(RowInsertRequests {
|
|
inserts: vec![large_row_insert],
|
|
})
|
|
.await;
|
|
assert!(
|
|
result.is_err(),
|
|
"Large request should exceed 1KB limit and fail"
|
|
);
|
|
let err = result.unwrap_err();
|
|
let err_msg = err.to_string();
|
|
assert!(
|
|
err_msg.contains("Memory limit exceeded"),
|
|
"Expected 'Memory limit exceeded' error, got: {}",
|
|
err_msg
|
|
);
|
|
|
|
let _ = fe_grpc_server.shutdown().await;
|
|
}
|