diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index 49a4b8fcd0f..20c7bc323a5 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -51,6 +51,10 @@ remote datanodes via `operator`/`client`. - Internal gRPC listeners mark requests with `Channel::Internal` in middleware (`server.rs`), including requests handled by Enterprise Flight wrappers. +- **Flight bulk insert** (`instance/grpc.rs`): initializes on the first batch after + the lazy schema handshake; checks permissions and reconciles missing columns + through `Inserter` once per stream, then reuses the refreshed table. + - **Logical-table batching** (`instance/logical_batcher.rs`): `Services` initializes one shared batcher for opted-in HTTP Prom and nonlegacy OTLP metric-engine writes. OTLP checks operator eligibility and falls back for incompatible tables. diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index e2435b65689..5cb546485bc 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -392,6 +392,7 @@ impl Instance { let catalog_manager = self.catalog_manager().clone(); let plugins = self.plugins.clone(); let inserter = self.inserter.clone(); + let statement_executor = self.statement_executor.clone(); let ctx = ctx.clone(); let mut table_ref: Option = None; let mut table_checked = false; @@ -422,25 +423,37 @@ impl Instance { .context(PermissionSnafu)?; // Resolve table reference - table_ref = Some( - catalog_manager - .table( - &table_name.catalog_name, - &table_name.schema_name, - &table_name.table_name, - None, - ) - .await - .context(CatalogSnafu)? - .with_context(|| TableNotFoundSnafu { - table_name: table_name.to_string(), - })?, - ); + let table = catalog_manager + .table( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + None, + ) + .await + .context(CatalogSnafu)? + .with_context(|| TableNotFoundSnafu { + table_name: table_name.to_string(), + })?; // Check permissions for the table let interceptor_ref = plugins.get::>(); let interceptor = interceptor_ref.as_ref(); - interceptor.pre_bulk_insert(table_ref.clone().unwrap(), ctx.clone())?; + interceptor.pre_bulk_insert(table.clone(), ctx.clone())?; + + // The client sends the schema lazily after the stream handshake. + // Reconcile it once, before writing the first batch. + table_ref = Some( + inserter + .ensure_bulk_insert_schema( + table, + request.record_batch.schema().as_ref(), + &ctx, + statement_executor.as_ref(), + ) + .await + .context(TableOperationSnafu)?, + ); table_checked = true; } diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index f1f57a64479..f8d58f5d173 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -140,6 +140,24 @@ impl Inserter { return Ok(0); } + // Bulk storage paths may discard unknown columns instead of rejecting them. + for field in record_batch.schema_ref().fields() { + ensure!( + table_info + .meta + .schema + .column_schema_by_name(field.name()) + .is_some(), + error::InvalidInsertRequestSnafu { + reason: format!( + "Column '{}' not found in table '{}'", + field.name(), + table_info.full_table_name() + ), + } + ); + } + // The zero value is WCU, not bytes. Bulk writes have no WCU accounting; // preserve that behavior while admitting their rows before dispatch. write_meter!(MeterRecord::new( diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index a0ebe236767..e1fbe39a9f4 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use api::v1::alter_table_expr::Kind; -use api::v1::column_def::options_from_skipping; +use api::v1::column_def::{options_from_skipping, try_as_column_def}; use api::v1::region::{ InsertRequest as RegionInsertRequest, InsertRequests as RegionInsertRequests, RegionRequestHeader, @@ -27,6 +27,7 @@ use api::v1::{ AlterTableExpr, ColumnDataType, ColumnSchema, CreateTableExpr, InsertRequests, RowInsertRequest, RowInsertRequests, Rows, SemanticType, }; +use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema}; use catalog::CatalogManagerRef; use client::{OutputData, OutputMeta}; use common_catalog::consts::{ @@ -906,6 +907,84 @@ impl Inserter { ) } + /// Adds missing columns from a bulk stream's schema and returns the refreshed table. + /// Call once when initializing the stream, before writing its first batch. + /// Does not infer new nested or dictionary columns. + pub async fn ensure_bulk_insert_schema( + &self, + table: TableRef, + request_schema: &ArrowSchema, + ctx: &QueryContextRef, + statement_executor: &StatementExecutor, + ) -> Result { + let table_info = table.table_info(); + if self.auto_create_disabled_reason(ctx)?.is_some() + && !Self::is_auto_create_exempt_private_table(&table_info.schema_name, &table_info.name) + { + return Ok(table); + } + + let table_schema = table.schema(); + let schema = request_schema + .fields() + .iter() + .filter(|field| table_schema.column_schema_by_name(field.name()).is_none()) + .map(|field| { + let data_type = field.data_type(); + // Dictionary values can reach the same infallible child-type conversion + // as nested types, even when Arrow's is_nested() returns false. + ensure!( + !data_type.is_nested() && !matches!(data_type, ArrowDataType::Dictionary(..)), + crate::error::NotSupportedSnafu { + feat: format!( + "automatically adding bulk insert column '{}' with type {:?}", + field.name(), + data_type + ), + } + ); + let column = datatypes::schema::ColumnSchema::try_from(field.as_ref()) + .context(crate::error::ConvertSchemaSnafu)?; + // Arrow fields do not carry primary-key semantics. New columns are + // fields, unless explicitly marked as a time index. + let column_def = + try_as_column_def(&column, false).context(crate::error::ColumnDataTypeSnafu)?; + Ok(ColumnSchema { + column_name: column_def.name, + datatype: column_def.data_type, + semantic_type: column_def.semantic_type, + datatype_extension: column_def.datatype_extension, + options: column_def.options, + }) + }) + .collect::>>()?; + let mut request = RowInsertRequest { + table_name: table_info.name.clone(), + rows: Some(Rows { + schema, + rows: Vec::new(), + }), + }; + let Some(alter_expr) = + self.get_alter_table_expr_on_demand(&mut request, &table, ctx, false, false, true)? + else { + return Ok(table); + }; + + statement_executor + .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter) + .await?; + self.get_table( + &table_info.catalog_name, + &table_info.schema_name, + &table_info.name, + ) + .await? + .with_context(|| TableNotFoundSnafu { + table_name: table_info.full_table_name(), + }) + } + /// Ensures a trace table has the request-global schema without requiring a /// padded data row to drive on-demand creation or alteration. When /// `alter_existing` is false, a table created after planning is left for the diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index 874e98b7326..ca05ea91844 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -46,6 +46,7 @@ use query::metrics::terminal_recordbatch_metrics_from_plan_if_requested; use query::options::FlowQueryExtensions; use session::context::{Channel, QueryContextRef}; use snafu::{IntoError, OptionExt, ResultExt, ensure}; +use table::requests::AUTO_CREATE_TABLE_KEY; use table::table_name::TableName; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; @@ -254,8 +255,14 @@ impl FlightCraft for GreptimeRequestHandler { let limiter = extensions.get::().cloned(); - let query_ctx = + let mut query_ctx = context_auth::create_query_context_from_grpc_metadata(&headers, &extensions)?; + // Bulk streams use the same schema-on-write switch as row inserts. + for (key, value) in hint_headers::extract_hints(&headers) { + if key == AUTO_CREATE_TABLE_KEY { + Arc::make_mut(&mut query_ctx).set_extension(key, value); + } + } context_auth::check_auth(self.user_provider.clone(), &headers, query_ctx.clone()).await?; const MAX_PENDING_RESPONSES: usize = 32; diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index 2cd0c44d504..ef3df79f558 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -47,6 +47,9 @@ mod test { RecordBatch, RecordBatchStreamWrapper, RecordBatches, SendableRecordBatchStream, }; use common_telemetry::tracing_context::TracingContext; + use datatypes::arrow::datatypes::{ + DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit, + }; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector}; @@ -66,14 +69,17 @@ mod test { use session::context::QueryContextRef; use session::hints::INSERT_SKIP_WAL_HINT; use tokio::net::TcpListener; - use tokio_stream::wrappers::TcpListenerStream; + use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream}; use tonic::transport::Server as TonicServer; use tonic::{Response, Status}; use tower::service_fn; use crate::cluster::GreptimeDbClusterBuilder; use crate::grpc::query_and_expect; - use crate::test_util::{MockInstanceImpl, StorageType, assert_wal_delta, setup_grpc_server}; + use crate::test_util::{ + MockInstanceImpl, StorageType, assert_wal_delta, setup_grpc_server, + setup_grpc_server_with_auto_create_table_disabled, + }; use crate::tests::test_util::MockInstance; struct SlowFlightCraft; @@ -467,6 +473,315 @@ mod test { db.shutdown().await; } + #[rstest] + #[case::standalone_single_region(false, false)] + #[case::standalone_partitioned(false, true)] + #[case::distributed_single_region(true, false)] + #[case::distributed_partitioned(true, true)] + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_auto_add_columns( + #[case] distributed: bool, + #[case] partitioned: bool, + ) { + let mut db = MockInstanceImpl::new( + &format!("flight_bulk_auto_add_columns_{partitioned}"), + distributed, + ) + .await; + let runtime = common_runtime::global_runtime().clone(); + let handler = GreptimeRequestHandler::new( + db.frontend(), + None, + Some(runtime.clone()), + FlightCompression::default(), + ); + let mut server = GrpcServerBuilder::new(GrpcServerConfig::default(), runtime) + .flight_handler(Arc::new(handler)) + .build(); + server.start("127.0.0.1:0".parse().unwrap()).await.unwrap(); + let client = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + let partition = if partitioned { + " PARTITION ON COLUMNS (a) (a < 0, a >= 0)" + } else { + "" + }; + client.sql(&format!( + "CREATE TABLE foo (ts TIMESTAMP TIME INDEX, a INT NOT NULL, PRIMARY KEY (a)){partition}" + )).await.unwrap(); + // Existing rows require the inferred column to be nullable. + client + .sql("INSERT INTO foo VALUES (100, 100)") + .await + .unwrap(); + test_put_record_batches(&client, create_record_batches(-4)).await; + + query_and_expect( + db.frontend().as_ref(), + "SELECT a, \"B\" FROM foo ORDER BY a", + "\ ++-----+-----+ +| a | B | ++-----+-----+ +| -4 | s4 | +| -3 | s3 | +| -2 | s2 | +| -1 | s1 | +| 0 | s0 | +| 1 | s-1 | +| 2 | s-2 | +| 3 | s-3 | +| 4 | s-4 | +| 100 | | ++-----+-----+", + ) + .await; + server.shutdown().await.unwrap(); + db.shutdown().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_auto_add_columns_once_per_stream() { + let (db, server) = setup_grpc_server( + StorageType::File, + "test_flight_bulk_auto_add_columns_once_per_stream", + ) + .await; + let client = Database::new_with_dbname( + "greptime-public", + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + client + .sql("CREATE TABLE foo (ts TIMESTAMP TIME INDEX, a INT NOT NULL, PRIMARY KEY (a))") + .await + .unwrap(); + let batches = create_record_batches(1); + let schema = batches[0].schema.arrow_schema().clone(); + let (tx, rx) = tokio::sync::mpsc::channel(4); + // The client must be able to receive the handshake before sending its schema. + let mut responses = client + .do_put(ReceiverStream::new(rx).boxed()) + .await + .unwrap(); + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + tx.send(encode_put_schema(schema.as_ref())).await.unwrap(); + + let empty_batch = common_recordbatch::DfRecordBatch::new_empty(schema); + for round in 0..2 { + for message in + FlightEncoder::default().encode(FlightMessage::RecordBatch(empty_batch.clone())) + { + tx.send(message).await.unwrap(); + } + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + if round == 0 { + // Even an empty first batch initializes the schema. Dropping the + // column makes a repeated reconciliation on the next batch observable. + client + .sql("ALTER TABLE foo DROP COLUMN \"B\"") + .await + .unwrap(); + } + } + query_and_expect( + db.frontend().as_ref(), + "SELECT column_name FROM information_schema.columns WHERE table_name = 'foo' ORDER BY column_name", + "\ ++-------------+ +| column_name | ++-------------+ +| a | +| ts | ++-------------+", + ) + .await; + drop(tx); + assert!(responses.next().await.is_none()); + + // A new stream reconciles again and persists the new column's values. + test_put_record_batches(&client, batches).await; + query_and_expect( + db.frontend().as_ref(), + "SELECT count(\"B\") AS n FROM foo", + "\ ++---+ +| n | ++---+ +| 9 | ++---+", + ) + .await; + server.shutdown().await.unwrap(); + } + + #[rstest] + #[case::hint(false)] + #[case::config(true)] + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_auto_add_columns_disabled(#[case] disabled_by_config: bool) { + let name = "test_flight_bulk_auto_add_columns_disabled"; + let (db, server) = if disabled_by_config { + setup_grpc_server_with_auto_create_table_disabled(StorageType::File, name).await + } else { + setup_grpc_server(StorageType::File, name).await + }; + let client = Database::new_with_dbname( + "greptime-public", + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + client + .sql("CREATE TABLE foo (ts TIMESTAMP TIME INDEX, a INT NOT NULL, PRIMARY KEY (a))") + .await + .unwrap(); + let batch = create_record_batches(1).remove(0).into_df_record_batch(); + let schema = batch.schema(); + let mut messages = vec![encode_put_schema(schema.as_ref())]; + messages.extend(FlightEncoder::default().encode(FlightMessage::RecordBatch( + common_recordbatch::DfRecordBatch::new_empty(schema), + ))); + messages.extend(FlightEncoder::default().encode(FlightMessage::RecordBatch(batch))); + // A request hint cannot override the server-side setting. + let hint = if disabled_by_config { "true" } else { "false" }; + let mut responses = client + .do_put_with_hints( + tokio_stream::iter(messages).boxed(), + &[("auto_create_table", hint)], + ) + .await + .unwrap(); + // The handshake and empty batch succeed, but the non-empty write must fail. + for _ in 0..2 { + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + } + let Some(Err(err)) = responses.next().await else { + panic!("expected the bulk write with an unknown column to fail"); + }; + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + let message = err.to_string(); + assert!(message.contains("Column 'B' not found"), "{message}"); + assert!(responses.next().await.is_none()); + query_and_expect( + db.frontend().as_ref(), + "SELECT column_name FROM information_schema.columns WHERE table_name = 'foo' ORDER BY column_name", + "\ ++-------------+ +| column_name | ++-------------+ +| a | +| ts | ++-------------+", + ) + .await; + query_and_expect( + db.frontend().as_ref(), + "SELECT count(*) AS n FROM foo", + "\ ++---+ +| n | ++---+ +| 0 | ++---+", + ) + .await; + test_put_record_batches_with_hints( + &client, + create_record_batches_without_nullable_column(1), + &[("auto_create_table", hint)], + ) + .await; + server.shutdown().await.unwrap(); + } + + #[rstest] + #[case::list(ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Int32, true))))] + #[case::list_unsupported_child(ArrowDataType::List(Arc::new(Field::new( + "item", + ArrowDataType::FixedSizeBinary(16), + true, + ))))] + #[case::struct_unsupported_child(ArrowDataType::Struct( + vec![Field::new("item", ArrowDataType::FixedSizeBinary(16), true)].into(), + ))] + #[case::dictionary(ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int32), + Box::new(ArrowDataType::Utf8), + ))] + #[case::dictionary_unsupported_value(ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int32), + Box::new(ArrowDataType::FixedSizeBinary(16)), + ))] + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_rejects_nested_auto_add_without_altering_table( + #[case] data_type: ArrowDataType, + ) { + let (db, server) = setup_grpc_server( + StorageType::File, + "test_flight_bulk_rejects_nested_auto_add_without_altering_table", + ) + .await; + let client = Database::new_with_dbname( + "greptime-public", + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + client + .sql("CREATE TABLE foo (ts TIMESTAMP TIME INDEX)") + .await + .unwrap(); + + // Use raw Arrow fields: Greptime's schema conversion itself used to panic + // on unsupported child types. A preceding scalar must not be partially added. + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("new_scalar", ArrowDataType::Int32, true), + Field::new("new_nested", data_type.clone(), true), + ])); + let mut encoder = FlightEncoder::default(); + let mut schema_data = encoder.encode_schema(schema.as_ref()); + schema_data.flight_descriptor = Some(FlightDescriptor { + r#type: arrow_flight::flight_descriptor::DescriptorType::Path as i32, + path: vec!["foo".to_string()], + ..Default::default() + }); + let mut messages = vec![schema_data]; + // An empty first batch still performs schema reconciliation. + messages.extend(encoder.encode(FlightMessage::RecordBatch( + common_recordbatch::DfRecordBatch::new_empty(schema), + ))); + let mut responses = client + .do_put(tokio_stream::iter(messages).boxed()) + .await + .unwrap(); + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + let Some(Err(err)) = responses.next().await else { + panic!("expected an unsupported bulk schema error"); + }; + assert_eq!(err.status_code(), StatusCode::Unsupported); + let message = err.to_string(); + assert!(message.contains("new_nested"), "{message}"); + assert!(message.contains(&format!("{data_type:?}")), "{message}"); + assert!(responses.next().await.is_none()); + + query_and_expect( + db.frontend().as_ref(), + "SELECT column_name FROM information_schema.columns WHERE table_name = 'foo' ORDER BY column_name", + "\ ++-------------+ +| column_name | ++-------------+ +| ts | ++-------------+", + ) + .await; + server.shutdown().await.unwrap(); + } + #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put() { common_telemetry::init_default_ut_logging(); @@ -816,34 +1131,23 @@ mod test { let requests_count = record_batches.len(); let schema = record_batches[0].schema.arrow_schema().clone(); - let stream = futures::stream::once(async move { - let mut schema_data = FlightEncoder::default().encode_schema(schema.as_ref()); - let metadata = DoPutMetadata::new(0); - schema_data.app_metadata = serde_json::to_vec(&metadata).unwrap().into(); - // first message in "DoPut" stream should carry table name in flight descriptor - schema_data.flight_descriptor = Some(FlightDescriptor { - r#type: arrow_flight::flight_descriptor::DescriptorType::Path as i32, - path: vec!["foo".to_string()], - ..Default::default() - }); - schema_data - }) - .chain( - tokio_stream::iter(record_batches) - .enumerate() - .flat_map(|(i, x)| { - let mut encoder = FlightEncoder::default(); - let message = FlightMessage::RecordBatch(x.into_df_record_batch()); - let mut data = encoder.encode(message); - let metadata = DoPutMetadata::new((i + 1) as i64); - data.iter_mut().for_each(|x| { - x.app_metadata = serde_json::to_vec(&metadata).unwrap().into() - }); - tokio_stream::iter(data) - }) - .boxed(), - ) - .boxed(); + let stream = futures::stream::once(async move { encode_put_schema(schema.as_ref()) }) + .chain( + tokio_stream::iter(record_batches) + .enumerate() + .flat_map(|(i, x)| { + let mut encoder = FlightEncoder::default(); + let message = FlightMessage::RecordBatch(x.into_df_record_batch()); + let mut data = encoder.encode(message); + let metadata = DoPutMetadata::new((i + 1) as i64); + data.iter_mut().for_each(|x| { + x.app_metadata = serde_json::to_vec(&metadata).unwrap().into() + }); + tokio_stream::iter(data) + }) + .boxed(), + ) + .boxed(); let response_stream = client.do_put_with_hints(stream, hints).await.unwrap(); @@ -863,6 +1167,17 @@ mod test { assert_eq!(requests_count + 1, responses_count); } + fn encode_put_schema(schema: &datatypes::arrow::datatypes::Schema) -> FlightData { + let mut schema_data = FlightEncoder::default().encode_schema(schema); + schema_data.app_metadata = serde_json::to_vec(&DoPutMetadata::new(0)).unwrap().into(); + schema_data.flight_descriptor = Some(FlightDescriptor { + r#type: arrow_flight::flight_descriptor::DescriptorType::Path as i32, + path: vec!["foo".to_string()], + ..Default::default() + }); + schema_data + } + fn create_record_batches_without_nullable_column(start: i64) -> Vec { let schema = Arc::new(Schema::new(vec![ ColumnSchema::new(