// Copyright 2023 Greptime Team // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. use std::sync::Arc; use 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 client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputData}; use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, MITO_ENGINE}; use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME; use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; use common_grpc::channel_manager::ClientTlsOption; use common_memory_manager::OnExhaustedPolicy; use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME; use common_query::Output; use common_recordbatch::RecordBatches; use common_runtime::Runtime; use common_runtime::runtime::{BuilderBuild, RuntimeTrait}; use common_test_util::find_workspace_path; 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 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::GrpcServerConfig; use servers::grpc::builder::GrpcServerBuilder; 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 [] { $( #[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_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, ); )* }; } 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; } 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; } 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 { 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::>(), )) 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_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(body: T, hint: Option<&str>) -> Request { 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::(&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::(&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::(&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 = (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 = (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; }