1#[cfg(feature = "enterprise")]
16mod flow_write;
17
18use std::collections::BTreeMap;
19use std::sync::atomic::{AtomicU64, Ordering};
20use std::sync::{Arc, LazyLock};
21use std::time::{Duration, Instant};
22
23use ahash::{HashMap, HashMapExt, HashSet, HashSetExt};
24use api::v1::alter_table_expr::Kind;
25use api::v1::column_def::{options_from_skipping, try_as_column_def};
26use api::v1::region::{
27 InsertRequest as RegionInsertRequest, InsertRequests as RegionInsertRequests,
28 RegionRequestHeader,
29};
30use api::v1::value::ValueData;
31use api::v1::{
32 AlterTableExpr, ColumnDataType, ColumnSchema, CreateTableExpr, InsertRequests,
33 RowInsertRequest, RowInsertRequests, Rows, SemanticType,
34};
35use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema};
36use catalog::CatalogManagerRef;
37use client::{OutputData, OutputMeta};
38use common_catalog::consts::{
39 DEFAULT_PRIVATE_SCHEMA_NAME, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN,
40 TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table,
41 trace_operations_table_name, trace_services_table_name,
42};
43use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME;
44use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME;
45use common_grpc_expr::util::ColumnExpr;
46use common_meta::cache::TableFlownodeSetCacheRef;
47use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME;
48use common_meta::node_manager::NodeManagerRef;
49use common_meta::peer::Peer;
50use common_meta::rpc::ddl::TriggerReason;
51use common_query::Output;
52use common_query::native_histogram::{is_native_histogram_value_type, native_histogram_value_type};
53use common_query::prelude::{greptime_timestamp, greptime_value};
54use common_telemetry::tracing_context::TracingContext;
55use common_telemetry::{debug, error, warn};
56use common_time::Timestamp;
57use common_time::timestamp::TimeUnit;
58use datatypes::schema::SkippingIndexOptions;
59use futures_util::future;
60use meter_core::data::MeterRecord;
61use meter_macros::write_meter;
62use partition::manager::PartitionRuleManagerRef;
63use session::context::QueryContextRef;
64use snafu::ResultExt;
65use snafu::prelude::*;
66use sql::partition::partition_rule_for_hexstring;
67use sql::statements::create::Partitions;
68use sql::statements::insert::Insert;
69use store_api::metric_engine_consts::{
70 LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
71};
72use store_api::mito_engine_options::{
73 APPEND_MODE_KEY, COMPACTION_OPTION_KEYS, COMPACTION_TWCS_PREFIX, COMPACTION_TYPE,
74 COMPACTION_TYPE_TWCS, MERGE_MODE_KEY, TTL_KEY,
75};
76use store_api::storage::{RegionId, TableId};
77use table::TableRef;
78use table::metadata::{TableInfo, TableInfoRef};
79use table::requests::{
80 AUTO_CREATE_TABLE_KEY, InsertRequest as TableInsertRequest, SEMANTIC_PER_TABLE_INDEX_KEY,
81 SEMANTIC_PIPELINE, TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1, TABLE_DATA_MODEL_TRACE_V2,
82 TRACE_TABLE_PARTITIONS_HINT_KEY, VALID_TABLE_OPTION_KEYS, is_semantic_option_key,
83 validate_semantic_option,
84};
85use table::table_reference::TableReference;
86use tokio_util::sync::CancellationToken;
87
88use crate::batcher::PendingRowsBatcher;
89use crate::error::{
90 CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, InvalidInsertRequestSnafu,
91 JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu, WriteRejectedSnafu,
92};
93use crate::expr_helper;
94use crate::region_req_factory::RegionRequestFactory;
95use crate::req_convert::common::preprocess_row_insert_requests;
96use crate::req_convert::insert::{
97 ColumnToRow, ImpureDefaultFiller, RowToRegion, StatementToRegion, TableToRegion,
98 fill_reqs_with_impure_default, rows_to_record_batch,
99};
100use crate::statement::StatementExecutor;
101use crate::write_dispatcher::{WriteDispatcher, WriteDispatcherRef};
102
103pub struct Inserter {
104 catalog_manager: CatalogManagerRef,
105 pub(crate) partition_manager: PartitionRuleManagerRef,
106 pub(crate) node_manager: NodeManagerRef,
107 pub(crate) write_dispatcher: WriteDispatcherRef,
108 pub(crate) table_flownode_set_cache: TableFlownodeSetCacheRef,
109 auto_create_table: bool,
113 pending_rows_batcher: Option<Arc<dyn PendingRowsBatcher>>,
114 mirror_pending_rows: Arc<AtomicU64>,
118}
119
120pub type InserterRef = Arc<Inserter>;
121
122#[derive(Clone)]
124pub enum AutoCreateTableType {
125 Logical(String),
127 Physical,
129 Log,
131 LastNonNull,
133 Trace { alter_existing: bool },
135}
136
137impl AutoCreateTableType {
138 pub fn as_str(&self) -> &'static str {
139 match self {
140 AutoCreateTableType::Logical(_) => "logical",
141 AutoCreateTableType::Physical => "physical",
142 AutoCreateTableType::Log => "log",
143 AutoCreateTableType::LastNonNull => "last_non_null",
144 AutoCreateTableType::Trace { .. } => "trace",
145 }
146 }
147
148 fn alter_existing(&self) -> bool {
149 !matches!(
150 self,
151 Self::Trace {
152 alter_existing: false
153 }
154 )
155 }
156}
157
158#[derive(Clone)]
165pub struct InstantAndNormalInsertRequests {
166 pub normal_requests: RegionInsertRequests,
168 pub instant_requests: RegionInsertRequests,
171}
172
173impl Inserter {
174 pub async fn can_batch_metric_rows(
178 &self,
179 requests: &RowInsertRequests,
180 ctx: &QueryContextRef,
181 physical_table: &str,
182 ) -> Result<bool> {
183 if self.auto_create_disabled_reason(ctx)?.is_some() || ctx.extension(TTL_KEY).is_some() {
184 return Ok(false);
185 }
186 for request in &requests.inserts {
187 if request.rows.as_ref().is_some_and(|rows| {
192 rows.schema.iter().any(|column| {
193 column.datatype_extension.is_some()
194 || !matches!(
195 ColumnDataType::try_from(column.datatype),
196 Ok(ColumnDataType::TimestampSecond
197 | ColumnDataType::TimestampMillisecond
198 | ColumnDataType::TimestampMicrosecond
199 | ColumnDataType::TimestampNanosecond
200 | ColumnDataType::Float64
201 | ColumnDataType::String)
202 )
203 })
204 }) {
205 return Ok(false);
206 }
207
208 let Some(table) = self
209 .get_table(
210 ctx.current_catalog(),
211 &ctx.current_schema(),
212 &request.table_name,
213 )
214 .await?
215 else {
216 continue;
217 };
218 let info = table.table_info();
219 if info.meta.engine != METRIC_ENGINE_NAME
220 || info.is_ttl_instant_table()
221 || info
222 .meta
223 .options
224 .extra_options
225 .get(LOGICAL_TABLE_METADATA_KEY)
226 .map(String::as_str)
227 != Some(physical_table)
228 || info
229 .meta
230 .schema
231 .column_schemas()
232 .iter()
233 .any(|column| column.default_constraint().is_some())
234 {
235 return Ok(false);
236 }
237 if info
240 .meta
241 .primary_key_indices
242 .iter()
243 .any(|&index| !info.meta.schema.column_schemas()[index].is_nullable())
244 {
245 return Ok(false);
246 }
247 match self.table_flownode_set_cache.get(info.table_id()).await {
250 Ok(None) => {}
251 Ok(Some(flows)) if flows.is_empty() => {}
252 _ => return Ok(false),
253 }
254 }
255 Ok(true)
256 }
257
258 pub async fn meter_row_inserts(
261 requests: &mut RowInsertRequests,
262 ctx: &QueryContextRef,
263 ) -> Result<u64> {
264 let metered = InstantAndNormalInsertRequests {
265 normal_requests: RegionInsertRequests {
266 requests: requests
267 .inserts
268 .iter_mut()
269 .map(|request| RegionInsertRequest {
270 rows: request.rows.take(),
271 ..Default::default()
272 })
273 .collect(),
274 },
275 instant_requests: RegionInsertRequests::default(),
276 };
277 let cost = write_meter!(
278 ctx.current_catalog(),
279 ctx.current_schema(),
280 metered,
281 ctx.write_rows_to_admit(
282 ctx.current_catalog(),
283 &ctx.current_schema(),
284 count_insert_rows(&metered)?
285 ),
286 ctx.channel() as u8
287 )
288 .await
289 .context(WriteRejectedSnafu);
290 for (request, region) in requests
291 .inserts
292 .iter_mut()
293 .zip(metered.normal_requests.requests)
294 {
295 request.rows = region.rows;
296 }
297 cost
298 }
299
300 pub fn new(
302 catalog_manager: CatalogManagerRef,
303 partition_manager: PartitionRuleManagerRef,
304 node_manager: NodeManagerRef,
305 table_flownode_set_cache: TableFlownodeSetCacheRef,
306 auto_create_table: bool,
307 ) -> Self {
308 Self {
309 write_dispatcher: Arc::new(WriteDispatcher::new(
310 node_manager.clone(),
311 partition_manager.clone(),
312 )),
313 catalog_manager,
314 partition_manager,
315 node_manager,
316 table_flownode_set_cache,
317 auto_create_table,
318 pending_rows_batcher: None,
319 mirror_pending_rows: Arc::new(AtomicU64::new(0)),
320 }
321 }
322
323 pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self {
325 self.write_dispatcher = dispatcher;
326 self
327 }
328
329 pub fn with_pending_rows_batcher(
331 mut self,
332 batcher: Option<Arc<dyn PendingRowsBatcher>>,
333 ) -> Self {
334 self.pending_rows_batcher = batcher;
335 self
336 }
337
338 pub async fn handle_column_inserts(
339 &self,
340 requests: InsertRequests,
341 ctx: QueryContextRef,
342 statement_executor: &StatementExecutor,
343 ) -> Result<Output> {
344 let row_inserts = ColumnToRow::convert(requests)?;
345 self.handle_row_inserts(row_inserts, ctx, statement_executor, false, false)
346 .await
347 }
348
349 pub async fn handle_row_inserts(
351 &self,
352 mut requests: RowInsertRequests,
353 ctx: QueryContextRef,
354 statement_executor: &StatementExecutor,
355 accommodate_existing_schema: bool,
356 is_single_value: bool,
357 ) -> Result<Output> {
358 preprocess_row_insert_requests(&mut requests.inserts)?;
359 self.handle_row_inserts_with_create_type(
360 requests,
361 ctx,
362 statement_executor,
363 AutoCreateTableType::Physical,
364 accommodate_existing_schema,
365 is_single_value,
366 )
367 .await
368 }
369
370 pub async fn handle_log_inserts(
372 &self,
373 requests: RowInsertRequests,
374 ctx: QueryContextRef,
375 statement_executor: &StatementExecutor,
376 ) -> Result<Output> {
377 self.handle_row_inserts_with_create_type(
378 requests,
379 ctx,
380 statement_executor,
381 AutoCreateTableType::Log,
382 false,
383 false,
384 )
385 .await
386 }
387
388 pub async fn handle_trace_inserts(
389 &self,
390 requests: RowInsertRequests,
391 ctx: QueryContextRef,
392 statement_executor: &StatementExecutor,
393 ) -> Result<Output> {
394 self.handle_row_inserts_with_create_type(
395 requests,
396 ctx,
397 statement_executor,
398 AutoCreateTableType::Trace {
399 alter_existing: true,
400 },
401 false,
402 false,
403 )
404 .await
405 }
406
407 pub async fn handle_last_non_null_inserts(
409 &self,
410 requests: RowInsertRequests,
411 ctx: QueryContextRef,
412 statement_executor: &StatementExecutor,
413 accommodate_existing_schema: bool,
414 is_single_value: bool,
415 ) -> Result<Output> {
416 self.handle_row_inserts_with_create_type(
417 requests,
418 ctx,
419 statement_executor,
420 AutoCreateTableType::LastNonNull,
421 accommodate_existing_schema,
422 is_single_value,
423 )
424 .await
425 }
426
427 async fn handle_row_inserts_with_create_type(
429 &self,
430 mut requests: RowInsertRequests,
431 ctx: QueryContextRef,
432 statement_executor: &StatementExecutor,
433 create_type: AutoCreateTableType,
434 accommodate_existing_schema: bool,
435 is_single_value: bool,
436 ) -> Result<Output> {
437 let skip_wal = ctx.skip_wal();
438
439 let batcher = self
440 .pending_rows_batcher
441 .as_ref()
442 .filter(|_| ctx.batching_enabled());
443
444 requests.inserts.retain(|req| {
446 req.rows
447 .as_ref()
448 .map(|r| !r.rows.is_empty())
449 .unwrap_or_default()
450 });
451 validate_column_count_match(&requests)?;
452
453 let CreateAlterTableResult {
454 instant_table_ids,
455 table_infos,
456 } = self
457 .create_or_alter_tables_on_demand(
458 &mut requests,
459 &ctx,
460 create_type,
461 statement_executor,
462 accommodate_existing_schema,
463 is_single_value,
464 None,
465 )
466 .await?;
467
468 if let Some(batcher) = batcher
471 && instant_table_ids.is_empty()
472 && table_infos
473 .values()
474 .all(|info| info.meta.engine == default_engine())
475 {
476 return self
477 .submit_pending_rows(requests, table_infos, ctx, batcher)
478 .await;
479 }
480
481 let name_to_info = table_infos
482 .values()
483 .map(|info| (info.name.clone(), info.clone()))
484 .collect::<HashMap<_, _>>();
485 let inserts = RowToRegion::new(
486 name_to_info,
487 instant_table_ids,
488 self.partition_manager.as_ref(),
489 )
490 .convert(requests, skip_wal)
491 .await?;
492
493 self.do_request(inserts, &table_infos, &ctx).await
494 }
495
496 async fn submit_pending_rows(
497 &self,
498 mut requests: RowInsertRequests,
499 table_infos: HashMap<TableId, Arc<TableInfo>>,
500 ctx: QueryContextRef,
501 batcher: &Arc<dyn PendingRowsBatcher>,
502 ) -> Result<Output> {
503 requests.inserts.retain(|request| {
506 request
507 .rows
508 .as_ref()
509 .is_some_and(|rows| !rows.rows.is_empty())
510 });
511 let by_name = table_infos
512 .values()
513 .map(|info| (info.name.as_str(), info))
514 .collect::<HashMap<_, _>>();
515 let mut prepared = Vec::with_capacity(requests.inserts.len());
516 for request in &mut requests.inserts {
517 let table_info =
518 by_name
519 .get(request.table_name.as_str())
520 .context(TableNotFoundSnafu {
521 table_name: &request.table_name,
522 })?;
523 let Some(rows) = &mut request.rows else {
524 continue;
525 };
526 ImpureDefaultFiller::new((*table_info).clone())?.fill_rows(rows);
527 let batch = rows_to_record_batch(rows, table_info)?;
528 prepared.push(((*table_info).clone(), batch));
529 }
530 let metered = InstantAndNormalInsertRequests {
533 normal_requests: RegionInsertRequests {
534 requests: requests
535 .inserts
536 .into_iter()
537 .map(|request| RegionInsertRequest {
538 rows: request.rows,
539 ..Default::default()
540 })
541 .collect(),
542 },
543 instant_requests: RegionInsertRequests::default(),
544 };
545 let table_info = table_infos.values().next();
546 let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str());
547 let schema =
548 table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone());
549 let write_cost = write_meter!(
550 catalog,
551 &schema,
552 metered,
553 ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&metered)?),
554 ctx.channel() as u8
555 )
556 .await
557 .context(WriteRejectedSnafu)?;
558 prepared.retain(|(_, batch)| batch.num_rows() != 0);
559 let results = if prepared.is_empty() {
560 Vec::new()
561 } else {
562 let permit = batcher.acquire().await?;
564 let submissions = prepared.into_iter().map(|(info, batch)| {
565 let mut target_ctx = ctx.fork();
567 target_ctx.set_current_catalog(&info.catalog_name);
568 target_ctx.set_current_schema(&info.schema_name);
569 batcher.submit(info, batch, Arc::new(target_ctx), permit.clone())
570 });
571 future::join_all(submissions).await
573 };
574 let mut affected_rows = 0;
575 let mut meta = OutputMeta::new_with_cost(write_cost as _);
576 for result in results {
577 let output = result?;
578 affected_rows += output.extract_rows_and_cost().0;
579 meta.write_completions.extend(output.meta.write_completions);
580 }
581 Ok(Output::new(OutputData::AffectedRows(affected_rows), meta))
582 }
583
584 pub async fn handle_metric_row_inserts(
586 &self,
587 mut requests: RowInsertRequests,
588 ctx: QueryContextRef,
589 statement_executor: &StatementExecutor,
590 physical_table: String,
591 ) -> Result<Output> {
592 let skip_wal = ctx.skip_wal();
593
594 requests.inserts.retain(|req| {
596 req.rows
597 .as_ref()
598 .map(|r| !r.rows.is_empty())
599 .unwrap_or_default()
600 });
601 validate_column_count_match(&requests)?;
602
603 let physical_table_ref = self
605 .create_physical_table_on_demand(&ctx, physical_table.clone(), statement_executor)
606 .await?;
607
608 let CreateAlterTableResult {
620 instant_table_ids,
621 table_infos,
622 } = self
623 .create_or_alter_tables_on_demand(
624 &mut requests,
625 &ctx,
626 AutoCreateTableType::Logical(physical_table.clone()),
627 statement_executor,
628 true,
629 true,
630 table_time_index_unit(&physical_table_ref),
631 )
632 .await?;
633 let name_to_info = table_infos
634 .values()
635 .map(|info| (info.name.clone(), info.clone()))
636 .collect::<HashMap<_, _>>();
637 let inserts = RowToRegion::new(name_to_info, instant_table_ids, &self.partition_manager)
638 .convert(requests, skip_wal)
639 .await?;
640
641 self.do_request(inserts, &table_infos, &ctx).await
642 }
643
644 fn table_batcher(
645 &self,
646 table_info: &TableInfoRef,
647 ctx: &QueryContextRef,
648 ) -> Option<&Arc<dyn PendingRowsBatcher>> {
649 self.pending_rows_batcher.as_ref().filter(|_| {
650 ctx.batching_enabled()
651 && !table_info.is_ttl_instant_table()
652 && table_info.meta.engine == default_engine()
653 })
654 }
655
656 async fn submit_table_rows(
657 &self,
658 rows: Rows,
659 table_info: TableInfoRef,
660 ctx: QueryContextRef,
661 batcher: &Arc<dyn PendingRowsBatcher>,
662 ) -> Result<Output> {
663 let requests = RowInsertRequests {
664 inserts: vec![RowInsertRequest {
665 table_name: table_info.name.clone(),
666 rows: Some(rows),
667 }],
668 };
669 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info)]);
670 self.submit_pending_rows(requests, table_infos, ctx, batcher)
671 .await
672 }
673
674 pub async fn handle_table_insert(
675 &self,
676 request: TableInsertRequest,
677 ctx: QueryContextRef,
678 ) -> Result<Output> {
679 let catalog = request.catalog_name.as_str();
680 let schema = request.schema_name.as_str();
681 let table_name = request.table_name.as_str();
682 let table = self.get_table(catalog, schema, table_name).await?;
683 let table = table.with_context(|| TableNotFoundSnafu {
684 table_name: common_catalog::format_full_table_name(catalog, schema, table_name),
685 })?;
686 let table_info = table.table_info();
687
688 let converter = TableToRegion::new(&table_info, &self.partition_manager);
689 let skip_wal = request.skip_wal;
690 let rows = converter.prepare(request)?;
691 if let Some(batcher) = self.table_batcher(&table_info, &ctx) {
692 return self.submit_table_rows(rows, table_info, ctx, batcher).await;
693 }
694 let inserts = converter.partition(rows, skip_wal).await?;
695
696 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]);
697
698 self.do_request(inserts, &table_infos, &ctx).await
699 }
700
701 pub async fn handle_statement_insert(
702 &self,
703 insert: &Insert,
704 ctx: &QueryContextRef,
705 ) -> Result<Output> {
706 let converter =
707 StatementToRegion::new(self.catalog_manager.as_ref(), &self.partition_manager, ctx);
708 let (rows, table_info) = converter.prepare(insert, ctx).await?;
709 if let Some(batcher) = self.table_batcher(&table_info, ctx) {
710 return self
711 .submit_table_rows(rows, table_info, ctx.clone(), batcher)
712 .await;
713 }
714 let inserts = converter.partition(rows, table_info.clone(), ctx).await?;
715
716 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]);
717
718 self.do_request(inserts, &table_infos, ctx).await
719 }
720}
721
722pub async fn admit_write(rows: u64, ctx: &QueryContextRef) -> Result<QueryContextRef> {
725 write_meter!(MeterRecord::new(
728 ctx.current_catalog().to_string(),
729 ctx.current_schema(),
730 0,
731 ctx.write_rows_to_admit(ctx.current_catalog(), &ctx.current_schema(), rows),
732 ctx.channel() as u8,
733 ))
734 .await
735 .context(WriteRejectedSnafu)?;
736 Ok(Arc::new(ctx.with_write_admission()))
737}
738
739pub async fn admit_row_insert_batches(
742 batches: &mut [(QueryContextRef, RowInsertRequests)],
743) -> Result<()> {
744 let mut totals = BTreeMap::<_, (QueryContextRef, u64)>::new();
745 for (ctx, requests) in batches.iter() {
746 let catalog = ctx.current_catalog();
747 let schema = ctx.current_schema();
748 if ctx.write_rows_to_admit(catalog, &schema, 1) == 0 {
749 continue;
750 }
751 let (_, total) = totals
752 .entry((catalog.to_string(), schema.clone()))
753 .or_insert_with(|| (ctx.clone(), 0));
754 for rows in requests.inserts.iter().filter_map(|r| r.rows.as_ref()) {
755 *total =
756 total
757 .checked_add(rows.rows.len() as u64)
758 .context(InvalidInsertRequestSnafu {
759 reason: "Insert row count exceeds u64::MAX",
760 })?;
761 }
762 }
763 for (ctx, rows) in totals.values() {
764 admit_write(*rows, ctx).await?;
765 }
766 for (ctx, _) in batches {
767 *ctx = Arc::new(ctx.with_write_admission());
768 }
769 Ok(())
770}
771
772fn count_insert_rows(requests: &InstantAndNormalInsertRequests) -> Result<u64> {
773 requests
774 .normal_requests
775 .requests
776 .iter()
777 .chain(&requests.instant_requests.requests)
778 .filter_map(|request| request.rows.as_ref())
779 .try_fold(0u64, |total, rows| {
780 total
781 .checked_add(rows.rows.len() as u64)
782 .context(InvalidInsertRequestSnafu {
783 reason: "Insert row count exceeds u64::MAX",
784 })
785 })
786}
787
788impl Inserter {
789 async fn do_request(
790 &self,
791 requests: InstantAndNormalInsertRequests,
792 table_infos: &HashMap<TableId, Arc<TableInfo>>,
793 ctx: &QueryContextRef,
794 ) -> Result<Output> {
795 let requests = fill_reqs_with_impure_default(table_infos, requests)?;
797
798 let table_info = table_infos.values().next();
801 let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str());
802 let schema =
803 table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone());
804 let write_cost = write_meter!(
805 catalog,
806 schema.clone(),
807 requests,
808 ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&requests)?),
809 ctx.channel() as u8
810 )
811 .await
812 .context(WriteRejectedSnafu)?;
813 let request_factory = RegionRequestFactory::new(RegionRequestHeader {
814 tracing_context: TracingContext::from_current_span().to_w3c(),
815 dbname: ctx.get_db_string(),
816 ..Default::default()
817 });
818
819 let InstantAndNormalInsertRequests {
820 normal_requests,
821 instant_requests,
822 } = requests;
823
824 let flow_mirror_task = FlowMirrorTask::new(
826 &self.table_flownode_set_cache,
827 normal_requests
828 .requests
829 .iter()
830 .chain(instant_requests.requests.iter()),
831 )
832 .await?;
833 let has_instant_rows = instant_requests.requests.iter().any(|request| {
834 request
835 .rows
836 .as_ref()
837 .is_some_and(|rows| !rows.rows.is_empty())
838 });
839 flow_mirror_task.detach(
840 self.node_manager.clone(),
841 self.mirror_pending_rows.clone(),
842 has_instant_rows,
843 )?;
844
845 #[cfg(feature = "enterprise")]
848 if flow_write::is_child(ctx) {
849 let dispatcher = self.write_dispatcher.clone();
850 let request = request_factory.build_insert(normal_requests);
851 let writes = dispatcher.prepare(vec![request]).await?;
852 let tracked_tasks = writes.into_iter().map(|(peer, mut request)| {
853 let mut query_context = crate::utils::to_meta_query_context(ctx.clone());
854 query_context.extensions.insert(
855 common_grpc::flow_write::CHILD_KEY.to_string(),
856 "1".to_string(),
857 );
858 if let Some(header) = request.header.as_mut() {
859 header.query_context = Some(query_context.into());
860 }
861 let node_manager = self.node_manager.clone();
862 async move { flow_write::dispatch(&node_manager, &peer, request).await }
863 });
864 let results = future::join_all(tracked_tasks).await;
865 let mut affected_rows = 0;
866 let mut first_error = None;
867 for result in results {
868 match result {
869 Ok(response) => affected_rows += response.affected_rows,
870 Err(error) if first_error.is_none() => first_error = Some(error),
871 Err(_) => {}
872 }
873 }
874 if let Some(error) = first_error {
875 return Err(error);
876 }
877 crate::metrics::DIST_INGEST_ROW_COUNT
878 .with_label_values(&[ctx.get_db_string().as_str()])
879 .inc_by(affected_rows as u64);
880 return Ok(Output::new(
881 OutputData::AffectedRows(affected_rows),
882 OutputMeta::new_with_cost(write_cost as _),
883 ));
884 }
885
886 let dispatcher = self.write_dispatcher.clone();
887 let request = request_factory.build_insert(normal_requests);
888 let writes = dispatcher.prepare(vec![request]).await?;
889 let cancellation = CancellationToken::new();
891 let _guard = cancellation.clone().drop_guard();
892 let affected_rows = common_runtime::spawn_global(async move {
893 dispatcher.execute(writes, None, Some(&cancellation)).await
894 })
895 .await
896 .context(JoinTaskSnafu)?
897 .context(RequestInsertsSnafu)?
898 .affected_rows;
899 crate::metrics::DIST_INGEST_ROW_COUNT
900 .with_label_values(&[ctx.get_db_string().as_str()])
901 .inc_by(affected_rows as u64);
902 Ok(Output::new(
903 OutputData::AffectedRows(affected_rows),
904 OutputMeta::new_with_cost(write_cost as _),
905 ))
906 }
907
908 fn auto_create_disabled_reason(&self, ctx: &QueryContextRef) -> Result<Option<&'static str>> {
912 let auto_create_table_hint = ctx
913 .extension(AUTO_CREATE_TABLE_KEY)
914 .map(|v| v.parse::<bool>())
915 .transpose()
916 .map_err(|_| {
917 InvalidInsertRequestSnafu {
918 reason: "`auto_create_table` hint must be a boolean",
919 }
920 .build()
921 })?
922 .unwrap_or(true);
923 Ok(if !self.auto_create_table {
924 Some("auto-create table is disabled by frontend config")
925 } else if !auto_create_table_hint {
926 Some("`auto_create_table` hint is disabled")
927 } else {
928 None
929 })
930 }
931
932 fn is_auto_create_exempt_private_table(schema: &str, table: &str) -> bool {
935 schema == DEFAULT_PRIVATE_SCHEMA_NAME
936 && matches!(
937 table,
938 DEFAULT_EVENTS_TABLE_NAME | SLOW_QUERY_TABLE_NAME | REGION_STATS_HISTORY_TABLE_NAME
939 )
940 }
941
942 pub async fn ensure_bulk_insert_schema(
946 &self,
947 table: TableRef,
948 request_schema: &ArrowSchema,
949 ctx: &QueryContextRef,
950 statement_executor: &StatementExecutor,
951 ) -> Result<TableRef> {
952 let table_info = table.table_info();
953 if self.auto_create_disabled_reason(ctx)?.is_some()
954 && !Self::is_auto_create_exempt_private_table(&table_info.schema_name, &table_info.name)
955 {
956 return Ok(table);
957 }
958
959 let table_schema = table.schema();
960 let schema = request_schema
961 .fields()
962 .iter()
963 .filter(|field| table_schema.column_schema_by_name(field.name()).is_none())
964 .map(|field| {
965 let data_type = field.data_type();
966 ensure!(
969 !data_type.is_nested() && !matches!(data_type, ArrowDataType::Dictionary(..)),
970 crate::error::NotSupportedSnafu {
971 feat: format!(
972 "automatically adding bulk insert column '{}' with type {:?}",
973 field.name(),
974 data_type
975 ),
976 }
977 );
978 let column = datatypes::schema::ColumnSchema::try_from(field.as_ref())
979 .context(crate::error::ConvertSchemaSnafu)?;
980 let column_def =
983 try_as_column_def(&column, false).context(crate::error::ColumnDataTypeSnafu)?;
984 Ok(ColumnSchema {
985 column_name: column_def.name,
986 datatype: column_def.data_type,
987 semantic_type: column_def.semantic_type,
988 datatype_extension: column_def.datatype_extension,
989 options: column_def.options,
990 })
991 })
992 .collect::<Result<Vec<_>>>()?;
993 let mut request = RowInsertRequest {
994 table_name: table_info.name.clone(),
995 rows: Some(Rows {
996 schema,
997 rows: Vec::new(),
998 }),
999 };
1000 let Some(alter_expr) =
1001 self.get_alter_table_expr_on_demand(&mut request, &table, ctx, false, false, true)?
1002 else {
1003 return Ok(table);
1004 };
1005
1006 statement_executor
1007 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1008 .await?;
1009 self.get_table(
1010 &table_info.catalog_name,
1011 &table_info.schema_name,
1012 &table_info.name,
1013 )
1014 .await?
1015 .with_context(|| TableNotFoundSnafu {
1016 table_name: table_info.full_table_name(),
1017 })
1018 }
1019
1020 pub async fn ensure_trace_table_on_demand(
1025 &self,
1026 table_name: &str,
1027 request_schema: Vec<ColumnSchema>,
1028 alter_existing: bool,
1029 ctx: &QueryContextRef,
1030 statement_executor: &StatementExecutor,
1031 ) -> Result<()> {
1032 let mut requests = RowInsertRequests {
1033 inserts: vec![RowInsertRequest {
1034 table_name: table_name.to_string(),
1035 rows: Some(api::v1::Rows {
1036 schema: request_schema,
1037 rows: Vec::new(),
1038 }),
1039 }],
1040 };
1041 self.create_or_alter_tables_on_demand(
1042 &mut requests,
1043 ctx,
1044 AutoCreateTableType::Trace { alter_existing },
1045 statement_executor,
1046 false,
1047 false,
1048 None,
1049 )
1050 .await?;
1051 Ok(())
1052 }
1053
1054 #[allow(clippy::too_many_arguments)]
1073 async fn create_or_alter_tables_on_demand(
1074 &self,
1075 requests: &mut RowInsertRequests,
1076 ctx: &QueryContextRef,
1077 auto_create_table_type: AutoCreateTableType,
1078 statement_executor: &StatementExecutor,
1079 accommodate_existing_schema: bool,
1080 is_single_value: bool,
1081 align_time_index_unit: Option<TimeUnit>,
1082 ) -> Result<CreateAlterTableResult> {
1083 let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND
1084 .with_label_values(&[auto_create_table_type.as_str()])
1085 .start_timer();
1086 let catalog = ctx.current_catalog();
1087 let schema = ctx.current_schema();
1088
1089 let auto_create_disabled_reason = self.auto_create_disabled_reason(ctx)?;
1090 let has_auto_create_exempt_table = auto_create_disabled_reason.is_some()
1092 && requests
1093 .inserts
1094 .iter()
1095 .any(|req| Self::is_auto_create_exempt_private_table(&schema, &req.table_name));
1096 let mut table_infos = HashMap::new();
1097 if let Some(disabled_reason) = auto_create_disabled_reason
1099 && !has_auto_create_exempt_table
1100 {
1101 let mut instant_table_ids = HashSet::new();
1102 for req in &mut requests.inserts {
1103 let table = match self.get_table(catalog, &schema, &req.table_name).await? {
1104 Some(table) => table,
1105 None if is_ddl_reserved_table(&schema, &req.table_name) => {
1108 statement_executor
1109 .create_declared_relationships_table(catalog, ctx.clone())
1110 .await?
1111 }
1112 None => {
1113 return InvalidInsertRequestSnafu {
1114 reason: format!(
1115 "Table `{}` does not exist, and {}",
1116 req.table_name, disabled_reason
1117 ),
1118 }
1119 .fail();
1120 }
1121 };
1122 if align_time_index_unit.is_some()
1126 && let Some(rows) = req.rows.as_mut()
1127 && let Some(target_unit) = table_time_index_unit(&table)
1128 {
1129 convert_rows_time_unit(rows, target_unit)?;
1130 }
1131 let table_info = table.table_info();
1132 if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1133 validate_trace_table_model(&table_info, ctx)?;
1134 }
1135 if table_info.is_ttl_instant_table() {
1136 instant_table_ids.insert(table_info.table_id());
1137 }
1138 table_infos.insert(table_info.table_id(), table.table_info());
1139 }
1140 let ret = CreateAlterTableResult {
1141 instant_table_ids,
1142 table_infos,
1143 };
1144 return Ok(ret);
1145 }
1146
1147 let mut create_tables = vec![];
1148 let mut alter_tables = vec![];
1149 let mut need_refresh_table_infos = HashSet::new();
1150 let mut instant_table_ids = HashSet::new();
1151 let mut per_table_semantics: Option<Option<PerTableSemanticIndex>> = None;
1152
1153 for req in &mut requests.inserts {
1154 let auto_create_allowed = auto_create_disabled_reason.is_none()
1156 || Self::is_auto_create_exempt_private_table(&schema, &req.table_name);
1157 match self.get_table(catalog, &schema, &req.table_name).await? {
1158 Some(table) => {
1159 let table_info = table.table_info();
1160 if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1161 validate_trace_table_model(&table_info, ctx)?;
1162 }
1163 if table_info.is_ttl_instant_table() {
1164 instant_table_ids.insert(table_info.table_id());
1165 }
1166 if align_time_index_unit.is_some()
1170 && let Some(rows) = req.rows.as_mut()
1171 && let Some(target_unit) = table_time_index_unit(&table)
1172 {
1173 convert_rows_time_unit(rows, target_unit)?;
1174 }
1175 if auto_create_allowed
1176 && let Some(alter_expr) = self.get_alter_table_expr_on_demand(
1177 req,
1178 &table,
1179 ctx,
1180 accommodate_existing_schema,
1181 is_single_value,
1182 auto_create_table_type.alter_existing(),
1183 )?
1184 {
1185 alter_tables.push(alter_expr);
1186 need_refresh_table_infos.insert((
1187 catalog.to_string(),
1188 schema.clone(),
1189 req.table_name.clone(),
1190 ));
1191 } else {
1192 table_infos.insert(table_info.table_id(), table.table_info());
1193 }
1194 }
1195 None if is_ddl_reserved_table(&schema, &req.table_name) => {
1199 let table = statement_executor
1200 .create_declared_relationships_table(catalog, ctx.clone())
1201 .await?;
1202 let table_info = table.table_info();
1203 if table_info.is_ttl_instant_table() {
1204 instant_table_ids.insert(table_info.table_id());
1205 }
1206 table_infos.insert(table_info.table_id(), table_info);
1207 }
1208 None if !auto_create_allowed
1209 && let Some(disabled_reason) = auto_create_disabled_reason =>
1210 {
1211 return InvalidInsertRequestSnafu {
1212 reason: format!(
1213 "Table `{}` does not exist, and {}",
1214 req.table_name, disabled_reason,
1215 ),
1216 }
1217 .fail();
1218 }
1219 None => {
1220 if let Some(physical_unit) = align_time_index_unit
1224 && let Some(rows) = req.rows.as_mut()
1225 {
1226 convert_rows_time_unit(rows, physical_unit)?;
1227 }
1228 let semantic_index = per_table_semantics
1229 .get_or_insert_with(|| parse_per_table_semantic_index(ctx))
1230 .as_ref();
1231 let create_expr = self.get_create_table_expr_on_demand(
1232 req,
1233 &auto_create_table_type,
1234 ctx,
1235 semantic_index,
1236 )?;
1237 create_tables.push(create_expr);
1238 }
1239 }
1240 }
1241
1242 match auto_create_table_type {
1243 AutoCreateTableType::Logical(_) => {
1244 if !create_tables.is_empty() {
1245 let tables = self
1247 .create_logical_tables(create_tables, ctx, statement_executor)
1248 .await?;
1249
1250 for table in tables {
1251 let table_info = table.table_info();
1252 if table_info.is_ttl_instant_table() {
1253 instant_table_ids.insert(table_info.table_id());
1254 }
1255 table_infos.insert(table_info.table_id(), table.table_info());
1256 }
1257 }
1258 if !alter_tables.is_empty() {
1259 statement_executor
1261 .alter_logical_tables(alter_tables, ctx.clone(), TriggerReason::AutoAlter)
1262 .await?;
1263 }
1264 }
1265 AutoCreateTableType::Physical
1266 | AutoCreateTableType::Log
1267 | AutoCreateTableType::LastNonNull => {
1268 for create_table in create_tables {
1271 let table = self
1272 .create_physical_table(create_table, None, ctx, statement_executor)
1273 .await?;
1274 let table_info = table.table_info();
1275 if table_info.is_ttl_instant_table() {
1276 instant_table_ids.insert(table_info.table_id());
1277 }
1278 table_infos.insert(table_info.table_id(), table.table_info());
1279 }
1280 for alter_expr in alter_tables.into_iter() {
1281 statement_executor
1282 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1283 .await?;
1284 }
1285 }
1286
1287 AutoCreateTableType::Trace { .. } => {
1288 let trace_table_name = ctx
1289 .extension(TRACE_TABLE_NAME_SESSION_KEY)
1290 .unwrap_or(TRACE_TABLE_NAME);
1291
1292 let trace_table_partitions = if let Some(trace_table_partitions) =
1293 ctx.extension(TRACE_TABLE_PARTITIONS_HINT_KEY)
1294 {
1295 let p = trace_table_partitions.parse::<u32>().map_err(|_| {
1296 InvalidInsertRequestSnafu {
1297 reason: format!(
1298 "Failed to parse trace_table_partitions: {}",
1299 trace_table_partitions
1300 ),
1301 }
1302 .build()
1303 })?;
1304 Some(p)
1305 } else {
1306 None
1307 };
1308
1309 for mut create_table in create_tables {
1312 if create_table.table_name == trace_services_table_name(trace_table_name)
1313 || create_table.table_name == trace_operations_table_name(trace_table_name)
1314 {
1315 create_table
1317 .table_options
1318 .insert(APPEND_MODE_KEY.to_string(), "false".to_string());
1319 create_table.table_options.remove(TTL_KEY);
1321
1322 let table = self
1323 .create_physical_table(create_table, None, ctx, statement_executor)
1324 .await?;
1325 let table_info = table.table_info();
1326 if table_info.is_ttl_instant_table() {
1327 instant_table_ids.insert(table_info.table_id());
1328 }
1329 table_infos.insert(table_info.table_id(), table.table_info());
1330 } else {
1331 let partitions = if matches!(trace_table_partitions, Some(0) | Some(1)) {
1334 None
1336 } else {
1337 let p = partition_rule_for_hexstring(
1338 TRACE_ID_COLUMN,
1339 trace_table_partitions,
1340 )
1341 .context(CreatePartitionRulesSnafu)?;
1342 Some(p)
1343 };
1344
1345 let index_columns =
1350 [TRACE_ID_COLUMN, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN];
1351 for index_column in index_columns {
1352 if let Some(col) = create_table
1353 .column_defs
1354 .iter_mut()
1355 .find(|c| c.name == index_column)
1356 {
1357 col.options =
1358 options_from_skipping(&SkippingIndexOptions::default())
1359 .context(ColumnOptionsSnafu)?;
1360 } else {
1361 warn!(
1362 "Column {} not found when creating index for trace table: {}.",
1363 index_column, create_table.table_name
1364 );
1365 }
1366 }
1367
1368 create_table.table_options.insert(
1370 TABLE_DATA_MODEL.to_string(),
1371 ctx.extension(SEMANTIC_PIPELINE)
1372 .unwrap_or(TABLE_DATA_MODEL_TRACE_V1)
1373 .to_string(),
1374 );
1375
1376 let table = self
1377 .create_physical_table(
1378 create_table,
1379 partitions,
1380 ctx,
1381 statement_executor,
1382 )
1383 .await?;
1384 let table_info = table.table_info();
1385 if table_info.is_ttl_instant_table() {
1386 instant_table_ids.insert(table_info.table_id());
1387 }
1388 table_infos.insert(table_info.table_id(), table.table_info());
1389 }
1390 }
1391 for alter_expr in alter_tables.into_iter() {
1392 statement_executor
1393 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1394 .await?;
1395 }
1396 }
1397 }
1398
1399 for (catalog, schema, table_name) in need_refresh_table_infos {
1401 let table = self
1402 .get_table(&catalog, &schema, &table_name)
1403 .await?
1404 .context(TableNotFoundSnafu {
1405 table_name: common_catalog::format_full_table_name(
1406 &catalog,
1407 &schema,
1408 &table_name,
1409 ),
1410 })?;
1411 let table_info = table.table_info();
1412 table_infos.insert(table_info.table_id(), table.table_info());
1413 }
1414
1415 Ok(CreateAlterTableResult {
1416 instant_table_ids,
1417 table_infos,
1418 })
1419 }
1420
1421 async fn create_physical_table_on_demand(
1422 &self,
1423 ctx: &QueryContextRef,
1424 physical_table: String,
1425 statement_executor: &StatementExecutor,
1426 ) -> Result<TableRef> {
1427 let catalog_name = ctx.current_catalog();
1428 let schema_name = ctx.current_schema();
1429
1430 if let Some(table) = self
1432 .get_table(catalog_name, &schema_name, &physical_table)
1433 .await?
1434 {
1435 return Ok(table);
1436 }
1437
1438 if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? {
1440 return InvalidInsertRequestSnafu {
1441 reason: format!(
1442 "Physical table `{physical_table}` does not exist, and {disabled_reason}"
1443 ),
1444 }
1445 .fail();
1446 }
1447
1448 let table_reference = TableReference::full(catalog_name, &schema_name, &physical_table);
1449 debug!("Ensuring physical metric table `{table_reference}` exists for insert");
1450
1451 let default_schema = vec![
1453 ColumnSchema {
1454 column_name: greptime_timestamp().to_string(),
1455 datatype: ColumnDataType::TimestampMillisecond as _,
1456 semantic_type: SemanticType::Timestamp as _,
1457 datatype_extension: None,
1458 options: None,
1459 },
1460 ColumnSchema {
1461 column_name: greptime_value().to_string(),
1462 datatype: ColumnDataType::Float64 as _,
1463 semantic_type: SemanticType::Field as _,
1464 datatype_extension: None,
1465 options: None,
1466 },
1467 ];
1468 let create_table_expr =
1469 &mut build_create_table_expr(&table_reference, &default_schema, default_engine())?;
1470
1471 create_table_expr.engine = METRIC_ENGINE_NAME.to_string();
1472 create_table_expr
1473 .table_options
1474 .insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), "true".to_string());
1475
1476 let res = statement_executor
1478 .create_table_inner(
1479 create_table_expr,
1480 None,
1481 ctx.clone(),
1482 TriggerReason::AutoCreate,
1483 )
1484 .await;
1485
1486 match res {
1487 Ok(table) => Ok(table),
1488 Err(err) => {
1489 error!(err; "Failed to create table {table_reference}");
1490 Err(err)
1491 }
1492 }
1493 }
1494
1495 async fn get_table(
1496 &self,
1497 catalog: &str,
1498 schema: &str,
1499 table: &str,
1500 ) -> Result<Option<TableRef>> {
1501 self.catalog_manager
1502 .table(catalog, schema, table, None)
1503 .await
1504 .context(CatalogSnafu)
1505 }
1506
1507 fn get_create_table_expr_on_demand(
1508 &self,
1509 req: &RowInsertRequest,
1510 create_type: &AutoCreateTableType,
1511 ctx: &QueryContextRef,
1512 semantic_index: Option<&PerTableSemanticIndex>,
1513 ) -> Result<CreateTableExpr> {
1514 let schema = ctx.current_schema();
1515 let mut table_options = std::collections::HashMap::with_capacity(4);
1516 fill_table_options_for_create(&mut table_options, create_type, ctx);
1517 apply_per_table_semantic_options(
1518 &mut table_options,
1519 semantic_index,
1520 ctx.current_schema().as_str(),
1521 &req.table_name,
1522 );
1523
1524 let engine_name = if let AutoCreateTableType::Logical(_) = create_type {
1525 METRIC_ENGINE_NAME
1527 } else {
1528 default_engine()
1529 };
1530
1531 let table_ref = TableReference::full(ctx.current_catalog(), &schema, &req.table_name);
1532 let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1534 let mut create_table_expr =
1535 build_create_table_expr(&table_ref, request_schema, engine_name)?;
1536
1537 if ctx.extension(SPLUNK_PK_METADATA_ORDER_KEY).is_some() {
1539 reorder_splunk_primary_keys(&mut create_table_expr.primary_keys);
1540 }
1541
1542 debug!("Ensuring table `{table_ref}` exists for insert");
1543 create_table_expr.table_options.extend(table_options);
1544 Ok(create_table_expr)
1545 }
1546
1547 fn get_alter_table_expr_on_demand(
1555 &self,
1556 req: &mut RowInsertRequest,
1557 table: &TableRef,
1558 ctx: &QueryContextRef,
1559 accommodate_existing_schema: bool,
1560 is_single_value: bool,
1561 alter_existing: bool,
1562 ) -> Result<Option<AlterTableExpr>> {
1563 if !alter_existing {
1564 return Ok(None);
1565 }
1566
1567 let catalog_name = ctx.current_catalog();
1568 let schema_name = ctx.current_schema();
1569 let table_name = table.table_info().name.clone();
1570
1571 if is_ddl_reserved_table(&schema_name, &table_name) {
1574 return Ok(None);
1575 }
1576
1577 let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1578 let request_field_count = request_schema
1579 .iter()
1580 .filter(|col| col.semantic_type == SemanticType::Field as i32)
1581 .count();
1582 let column_exprs = ColumnExpr::from_column_schemas(request_schema);
1583 let add_columns = expr_helper::extract_add_columns_expr(&table.schema(), column_exprs)?;
1584 let Some(mut add_columns) = add_columns else {
1585 return Ok(None);
1586 };
1587
1588 if is_single_value {
1589 let request_is_native_histogram = request_is_native_histogram(request_schema);
1590 let table_is_native_histogram = table_is_native_histogram(table);
1591 ensure!(
1592 request_is_native_histogram == table_is_native_histogram,
1593 InvalidInsertRequestSnafu {
1594 reason: format!(
1595 "Table `{table_name}` cannot mix native histogram and float sample fields"
1596 ),
1597 }
1598 );
1599 }
1600
1601 if accommodate_existing_schema {
1603 let table_schema = table.schema();
1604 let ts_col_name = table_schema.timestamp_column().map(|c| c.name.clone());
1606 let mut field_col_name = None;
1608 if is_single_value && request_field_count <= 1 {
1609 let mut multiple_field_cols = false;
1610 table.field_columns().for_each(|col| {
1611 if field_col_name.is_none() {
1612 field_col_name = Some(col.name.clone());
1613 } else {
1614 multiple_field_cols = true;
1615 }
1616 });
1617 if multiple_field_cols {
1618 field_col_name = None;
1619 }
1620 }
1621
1622 if let Some(rows) = req.rows.as_mut() {
1624 for col in &mut rows.schema {
1625 match col.semantic_type {
1626 x if x == SemanticType::Timestamp as i32 => {
1627 if let Some(ref ts_name) = ts_col_name
1628 && col.column_name != *ts_name
1629 {
1630 col.column_name = ts_name.clone();
1631 }
1632 }
1633 x if x == SemanticType::Field as i32 => {
1634 if let Some(ref field_name) = field_col_name
1635 && col.column_name != *field_name
1636 {
1637 col.column_name = field_name.clone();
1638 }
1639 }
1640 _ => {}
1641 }
1642 }
1643 }
1644
1645 add_columns.add_columns.retain(|col| {
1647 let def = col.column_def.as_ref().unwrap();
1648 def.semantic_type == SemanticType::Tag as i32
1649 || (def.semantic_type == SemanticType::Field as i32 && field_col_name.is_none())
1650 });
1651
1652 if add_columns.add_columns.is_empty() {
1653 return Ok(None);
1654 }
1655 }
1656
1657 Ok(Some(AlterTableExpr {
1658 catalog_name: catalog_name.to_string(),
1659 schema_name: schema_name.clone(),
1660 table_name: table_name.clone(),
1661 kind: Some(Kind::AddColumns(add_columns)),
1662 }))
1663 }
1664
1665 async fn create_physical_table(
1667 &self,
1668 mut create_table_expr: CreateTableExpr,
1669 partitions: Option<Partitions>,
1670 ctx: &QueryContextRef,
1671 statement_executor: &StatementExecutor,
1672 ) -> Result<TableRef> {
1673 let res = statement_executor
1674 .create_table_inner(
1675 &mut create_table_expr,
1676 partitions,
1677 ctx.clone(),
1678 TriggerReason::AutoCreate,
1679 )
1680 .await;
1681
1682 let table_ref = TableReference::full(
1683 &create_table_expr.catalog_name,
1684 &create_table_expr.schema_name,
1685 &create_table_expr.table_name,
1686 );
1687
1688 match res {
1689 Ok(table) => {
1690 validate_trace_table_model(&table.table_info(), ctx)?;
1691 Ok(table)
1692 }
1693 Err(err) => {
1694 error!(err; "Failed to create table {}", table_ref);
1695 Err(err)
1696 }
1697 }
1698 }
1699
1700 async fn create_logical_tables(
1701 &self,
1702 create_table_exprs: Vec<CreateTableExpr>,
1703 ctx: &QueryContextRef,
1704 statement_executor: &StatementExecutor,
1705 ) -> Result<Vec<TableRef>> {
1706 let res = statement_executor
1707 .create_logical_tables(&create_table_exprs, ctx.clone(), TriggerReason::AutoCreate)
1708 .await;
1709
1710 match res {
1711 Ok(res) => Ok(res),
1712 Err(err) => {
1713 let failed_tables = create_table_exprs
1714 .into_iter()
1715 .map(|expr| {
1716 format!(
1717 "{}.{}.{}",
1718 expr.catalog_name, expr.schema_name, expr.table_name
1719 )
1720 })
1721 .collect::<Vec<_>>();
1722 error!(
1723 err;
1724 "Failed to create logical tables {:?}",
1725 failed_tables
1726 );
1727 Err(err)
1728 }
1729 }
1730 }
1731
1732 pub fn write_dispatcher(&self) -> &WriteDispatcherRef {
1734 &self.write_dispatcher
1735 }
1736
1737 pub fn node_manager(&self) -> &NodeManagerRef {
1738 &self.node_manager
1739 }
1740
1741 pub fn partition_manager(&self) -> &PartitionRuleManagerRef {
1742 &self.partition_manager
1743 }
1744
1745 pub fn table_flownode_set_cache(&self) -> &TableFlownodeSetCacheRef {
1746 &self.table_flownode_set_cache
1747 }
1748}
1749
1750fn request_is_native_histogram(request_schema: &[ColumnSchema]) -> bool {
1751 let mut fields = request_schema
1752 .iter()
1753 .filter(|col| col.semantic_type == SemanticType::Field as i32);
1754 let Some(col) = fields.next() else {
1755 return false;
1756 };
1757
1758 fields.next().is_none()
1759 && api::helper::is_column_type_value_eq(
1760 col.datatype,
1761 col.datatype_extension.clone(),
1762 native_histogram_value_type(),
1763 )
1764}
1765
1766fn table_time_index_unit(table: &TableRef) -> Option<TimeUnit> {
1770 table
1771 .table_info()
1772 .meta
1773 .schema
1774 .timestamp_column()
1775 .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit()))
1776}
1777
1778fn convert_rows_time_unit(rows: &mut Rows, target_unit: TimeUnit) -> Result<()> {
1779 let Some(ts_index) = rows
1780 .schema
1781 .iter()
1782 .position(|col| col.semantic_type == SemanticType::Timestamp as i32)
1783 else {
1784 return Ok(());
1785 };
1786 let Some(source_unit) = ColumnDataType::try_from(rows.schema[ts_index].datatype)
1787 .ok()
1788 .and_then(api::helper::timestamp_unit)
1789 else {
1790 return Ok(());
1791 };
1792 if source_unit == target_unit {
1793 return Ok(());
1794 }
1795
1796 rows.schema[ts_index].datatype = api::helper::timestamp_datatype(target_unit) as i32;
1797 rows.schema[ts_index].datatype_extension = None;
1799
1800 for row in &mut rows.rows {
1808 debug_assert_eq!(row.values.len(), rows.schema.len());
1809 let value = &mut row.values[ts_index];
1810 let Some(value_data) = value.value_data.take() else {
1811 continue;
1812 };
1813 value.value_data =
1814 convert_timestamp_value_data(value_data, source_unit, target_unit, ts_index)?;
1815 }
1816 Ok(())
1817}
1818
1819fn convert_timestamp_value_data(
1820 value_data: ValueData,
1821 source_unit: TimeUnit,
1822 target_unit: TimeUnit,
1823 column_index: usize,
1824) -> Result<Option<ValueData>> {
1825 let timestamp = match value_data {
1826 ValueData::TimestampSecondValue(v) => Timestamp::new_second(v),
1827 ValueData::TimestampMillisecondValue(v) => Timestamp::new_millisecond(v),
1828 ValueData::TimestampMicrosecondValue(v) => Timestamp::new_microsecond(v),
1829 ValueData::TimestampNanosecondValue(v) => Timestamp::new_nanosecond(v),
1830 other => return Ok(Some(other)),
1832 };
1833 let converted = timestamp
1834 .convert_to(target_unit)
1835 .with_context(|| InvalidInsertRequestSnafu {
1836 reason: format!(
1837 "timestamp column {column_index} value {} in unit {source_unit:?} overflows when converting to unit {target_unit:?}",
1838 timestamp.value()
1839 ),
1840 })?;
1841 Ok(api::helper::to_grpc_value(datatypes::value::Value::Timestamp(converted)).value_data)
1842}
1843
1844fn table_is_native_histogram(table: &TableRef) -> bool {
1845 let mut fields = table.field_columns();
1846 let Some(col) = fields.next() else {
1847 return false;
1848 };
1849
1850 fields.next().is_none() && is_native_histogram_value_type(&col.data_type)
1851}
1852
1853fn validate_column_count_match(requests: &RowInsertRequests) -> Result<()> {
1854 for request in &requests.inserts {
1855 let rows = request.rows.as_ref().unwrap();
1856 let column_count = rows.schema.len();
1857 rows.rows.iter().try_for_each(|r| {
1858 ensure!(
1859 r.values.len() == column_count,
1860 InvalidInsertRequestSnafu {
1861 reason: format!(
1862 "column count mismatch, columns: {}, values: {}",
1863 column_count,
1864 r.values.len()
1865 )
1866 }
1867 );
1868 Ok(())
1869 })?;
1870 }
1871 Ok(())
1872}
1873
1874pub fn validate_trace_table_model(table_info: &TableInfo, ctx: &QueryContextRef) -> Result<()> {
1877 let Some(expected @ (TABLE_DATA_MODEL_TRACE_V1 | TABLE_DATA_MODEL_TRACE_V2)) =
1878 ctx.extension(SEMANTIC_PIPELINE)
1879 else {
1880 return Ok(());
1881 };
1882 if let Some(actual) = table_info.meta.options.data_model() {
1883 ensure!(
1884 actual == expected,
1885 InvalidInsertRequestSnafu {
1886 reason: format!(
1887 "Trace table `{}` uses {actual}, but the request uses {expected}",
1888 table_info.name,
1889 ),
1890 }
1891 );
1892 }
1893 Ok(())
1894}
1895
1896pub fn fill_table_options_for_create(
1898 table_options: &mut std::collections::HashMap<String, String>,
1899 create_type: &AutoCreateTableType,
1900 ctx: &QueryContextRef,
1901) {
1902 for key in VALID_TABLE_OPTION_KEYS {
1903 if let Some(value) = ctx.extension(key) {
1904 table_options.insert(key.to_string(), value.to_string());
1905 }
1906 }
1907
1908 if !matches!(create_type, AutoCreateTableType::Logical(_)) {
1909 let mut has_twcs_option = false;
1910 for key in COMPACTION_OPTION_KEYS {
1911 if let Some(value) = ctx.extension(key) {
1912 has_twcs_option |= key.starts_with(COMPACTION_TWCS_PREFIX);
1913 table_options.insert(key.to_string(), value.to_string());
1914 }
1915 }
1916 if has_twcs_option {
1917 table_options
1918 .entry(COMPACTION_TYPE.to_string())
1919 .or_insert_with(|| COMPACTION_TYPE_TWCS.to_string());
1920 }
1921 }
1922
1923 for (key, value) in ctx.extensions() {
1925 if is_semantic_option_key(&key) && validate_semantic_option(&key, &value) {
1926 table_options.insert(key, value);
1927 }
1928 }
1929
1930 match create_type {
1931 AutoCreateTableType::Logical(physical_table) => {
1932 table_options.insert(
1933 LOGICAL_TABLE_METADATA_KEY.to_string(),
1934 physical_table.clone(),
1935 );
1936 }
1937 AutoCreateTableType::Physical => {
1938 if let Some(append_mode) = ctx.extension(APPEND_MODE_KEY) {
1939 table_options.insert(APPEND_MODE_KEY.to_string(), append_mode.to_string());
1940 }
1941 if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
1942 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
1943 }
1944 }
1945 AutoCreateTableType::Log => {
1948 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1949 }
1950 AutoCreateTableType::LastNonNull => {
1951 if ctx
1952 .extension(APPEND_MODE_KEY)
1953 .is_some_and(|value| value.eq_ignore_ascii_case("true"))
1954 {
1955 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1956 table_options.insert(MERGE_MODE_KEY.to_string(), "last_row".to_string());
1957 } else if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
1958 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
1959 } else {
1960 table_options.insert(MERGE_MODE_KEY.to_string(), "last_non_null".to_string());
1961 }
1962 }
1963 AutoCreateTableType::Trace { .. } => {
1964 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1965 }
1966 }
1967}
1968
1969pub type PerTableSemanticIndex = BTreeMap<String, BTreeMap<String, BTreeMap<String, String>>>;
1975
1976pub fn parse_per_table_semantic_index(ctx: &QueryContextRef) -> Option<PerTableSemanticIndex> {
1981 let raw = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY)?;
1982 match serde_json::from_str(raw) {
1983 Ok(index) => Some(index),
1984 Err(_) => {
1985 warn!("failed to parse semantic per-table index, skipping per-table options");
1986 None
1987 }
1988 }
1989}
1990
1991pub fn apply_per_table_semantic_options(
1999 table_options: &mut std::collections::HashMap<String, String>,
2000 index: Option<&PerTableSemanticIndex>,
2001 schema: &str,
2002 table_name: &str,
2003) {
2004 let Some(entry) = index
2005 .and_then(|index| index.get(schema))
2006 .and_then(|tables| tables.get(table_name))
2007 else {
2008 return;
2009 };
2010 for (key, value) in entry {
2011 if is_semantic_option_key(key) && validate_semantic_option(key, value) {
2012 table_options.insert(key.clone(), value.clone());
2013 }
2014 }
2015}
2016
2017pub fn build_create_table_expr(
2018 table: &TableReference,
2019 request_schema: &[ColumnSchema],
2020 engine: &str,
2021) -> Result<CreateTableExpr> {
2022 expr_helper::create_table_expr_by_column_schemas(table, request_schema, engine, None)
2023}
2024
2025pub const SPLUNK_PK_METADATA_ORDER_KEY: &str = "splunk_pk_metadata_order";
2029
2030fn reorder_splunk_primary_keys(primary_keys: &mut [String]) {
2033 const LEAD: [&str; 3] = ["host", "source", "sourcetype"];
2034 primary_keys.sort_by_key(|name| {
2037 LEAD.iter()
2038 .position(|&lead| lead == name.as_str())
2039 .unwrap_or(LEAD.len())
2040 });
2041}
2042
2043struct CreateAlterTableResult {
2045 instant_table_ids: HashSet<TableId>,
2047 table_infos: HashMap<TableId, Arc<TableInfo>>,
2049}
2050
2051const MAX_MIRROR_PENDING_ROWS: usize = 1_000_000;
2059
2060const MIRROR_LOG_INTERVAL: Duration = Duration::from_secs(10);
2062const MIRROR_LOG_NEVER_REPORTED: u64 = u64::MAX;
2063
2064struct MirrorLog {
2065 last_log_millis: AtomicU64,
2066 events: AtomicU64,
2067}
2068
2069impl MirrorLog {
2070 const fn new() -> Self {
2071 Self {
2072 last_log_millis: AtomicU64::new(MIRROR_LOG_NEVER_REPORTED),
2073 events: AtomicU64::new(0),
2074 }
2075 }
2076
2077 fn claim_report(&self, now_millis: u64) -> Option<u64> {
2080 self.events.fetch_add(1, Ordering::Relaxed);
2081 let last = self.last_log_millis.load(Ordering::Relaxed);
2082 if last != MIRROR_LOG_NEVER_REPORTED
2083 && now_millis.saturating_sub(last) < MIRROR_LOG_INTERVAL.as_millis() as u64
2084 {
2085 return None;
2086 }
2087 if self
2088 .last_log_millis
2089 .compare_exchange(last, now_millis, Ordering::Relaxed, Ordering::Relaxed)
2090 .is_err()
2091 {
2092 return None;
2093 }
2094 Some(self.events.swap(0, Ordering::Relaxed))
2095 }
2096}
2097
2098static MIRROR_DROP_LOG: MirrorLog = MirrorLog::new();
2099static MIRROR_FAILURE_LOG: MirrorLog = MirrorLog::new();
2100
2101fn mirror_drop_log_millis() -> u64 {
2104 static START: LazyLock<Instant> = LazyLock::new(Instant::now);
2105 START.elapsed().as_millis() as u64
2106}
2107
2108struct FlowMirrorTask {
2109 requests: HashMap<Peer, RegionInsertRequests>,
2110}
2111
2112impl FlowMirrorTask {
2113 async fn new(
2114 cache: &TableFlownodeSetCacheRef,
2115 requests: impl Iterator<Item = &RegionInsertRequest>,
2116 ) -> Result<Self> {
2117 let mut src_table_reqs: HashMap<TableId, Option<(Vec<Peer>, RegionInsertRequests)>> =
2118 HashMap::new();
2119
2120 for req in requests {
2121 let table_id = RegionId::from_u64(req.region_id).table_id();
2122 match src_table_reqs.get_mut(&table_id) {
2123 Some(Some((_peers, reqs))) => reqs.requests.push(req.clone()),
2124 Some(None) => continue,
2126 _ => {
2127 let peers = cache
2129 .get(table_id)
2130 .await
2131 .context(RequestInsertsSnafu)?
2132 .unwrap_or_default()
2133 .values()
2134 .cloned()
2135 .collect::<HashSet<_>>()
2136 .into_iter()
2137 .collect::<Vec<_>>();
2138
2139 if !peers.is_empty() {
2140 let mut reqs = RegionInsertRequests::default();
2141 reqs.requests.push(req.clone());
2142 src_table_reqs.insert(table_id, Some((peers, reqs)));
2143 } else {
2144 src_table_reqs.insert(table_id, None);
2146 }
2147 }
2148 }
2149 }
2150
2151 let mut inserts: HashMap<Peer, RegionInsertRequests> = HashMap::new();
2152
2153 for (_table_id, (peers, reqs)) in src_table_reqs
2154 .into_iter()
2155 .filter_map(|(k, v)| v.map(|v| (k, v)))
2156 {
2157 if peers.len() == 1 {
2158 inserts
2160 .entry(peers[0].clone())
2161 .or_default()
2162 .requests
2163 .extend(reqs.requests);
2164 continue;
2165 } else {
2166 for flownode in peers {
2168 inserts
2169 .entry(flownode.clone())
2170 .or_default()
2171 .requests
2172 .extend(reqs.requests.clone());
2173 }
2174 }
2175 }
2176
2177 Ok(Self { requests: inserts })
2178 }
2179
2180 fn pending_rows(&self) -> u64 {
2189 self.requests.values().map(region_inserts_rows).sum()
2190 }
2191
2192 fn detach(
2193 self,
2194 node_manager: NodeManagerRef,
2195 mirror_pending_rows: Arc<AtomicU64>,
2196 has_instant_rows: bool,
2197 ) -> Result<()> {
2198 let num_rows = self.pending_rows();
2201 if num_rows == 0 {
2202 return Ok(());
2203 }
2204 if has_instant_rows && num_rows > MAX_MIRROR_PENDING_ROWS as u64 {
2205 return InvalidInsertRequestSnafu {
2206 reason: format!(
2207 "flow mirror batch has {num_rows} peer-fan-out rows, exceeding the limit of {}; reduce the batch size",
2208 MAX_MIRROR_PENDING_ROWS
2209 ),
2210 }
2211 .fail();
2212 }
2213
2214 let pending = reserve_mirror_pending_rows(&mirror_pending_rows, num_rows);
2215 if pending > MAX_MIRROR_PENDING_ROWS as u64 {
2216 release_mirror_pending_rows(&mirror_pending_rows, num_rows);
2217 if has_instant_rows {
2218 return Err(meter_core::collect::WriteRejected::new(
2219 "flow mirror pending rows limit exceeded for instant-TTL insert",
2220 ))
2221 .context(WriteRejectedSnafu);
2222 }
2223 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.inc_by(num_rows);
2224 if let Some(events) = MIRROR_DROP_LOG.claim_report(mirror_drop_log_millis()) {
2225 warn!(
2226 "Flow mirror write dropped: {} pending rows exceeds limit {}, {} drop events since last report",
2227 pending, MAX_MIRROR_PENDING_ROWS, events
2228 );
2229 }
2230 return Ok(());
2231 }
2232
2233 for (peer, inserts) in self.requests {
2234 let peer_rows = region_inserts_rows(&inserts);
2238 let node_manager = node_manager.clone();
2239 let mirror_pending_rows = mirror_pending_rows.clone();
2240 common_runtime::spawn_global(async move {
2241 let result = node_manager
2242 .flownode(&peer)
2243 .await
2244 .handle_inserts(inserts)
2245 .await
2246 .context(RequestInsertsSnafu);
2247
2248 match result {
2249 Ok(resp) => {
2250 let affected_rows = resp.affected_rows;
2251 crate::metrics::DIST_MIRROR_ROW_COUNT.inc_by(affected_rows);
2252 }
2253 Err(err) => {
2254 if let Some(events) =
2255 MIRROR_FAILURE_LOG.claim_report(mirror_drop_log_millis())
2256 {
2257 error!(err; flownode_id = peer.id, flownode_addr = %peer.addr, "Failed to insert data into flownode ({} total mirror failures across peers since last report)", events);
2258 }
2259 }
2260 }
2261 release_mirror_pending_rows(&mirror_pending_rows, peer_rows);
2264 });
2265 }
2266
2267 Ok(())
2268 }
2269}
2270
2271fn region_inserts_rows(inserts: &RegionInsertRequests) -> u64 {
2274 inserts
2275 .requests
2276 .iter()
2277 .filter_map(|req| req.rows.as_ref())
2278 .map(|rows| rows.rows.len() as u64)
2279 .sum()
2280}
2281
2282fn reserve_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) -> u64 {
2287 let pending = mirror_pending_rows.fetch_add(rows, Ordering::Relaxed) + rows;
2288 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(rows as i64);
2291 pending
2292}
2293
2294fn release_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) {
2300 let prev = mirror_pending_rows
2301 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |pending| {
2302 Some(pending.saturating_sub(rows))
2303 })
2304 .unwrap_or_else(|prev| prev);
2305 let released = rows.min(prev);
2309 if released > 0 {
2310 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(released as i64);
2311 }
2312}
2313
2314#[cfg(test)]
2315mod tests {
2316 use std::sync::Arc;
2317 use std::sync::atomic::{AtomicU64, Ordering};
2318
2319 use api::helper::ColumnDataTypeWrapper;
2320 use api::v1::flow::FlowResponse;
2321 use api::v1::helper::{field_column_schema, time_index_column_schema};
2322 use api::v1::{Row, RowInsertRequest, Rows, Value};
2323 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
2324 use common_meta::cache::new_table_flownode_set_cache;
2325 use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler;
2326 use common_meta::test_util::{MockDatanodeManager, MockFlownodeHandler, MockFlownodeManager};
2327 use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD;
2328 use common_query::prelude::{greptime_native_histogram, set_default_prefix};
2329 use datatypes::data_type::ConcreteDataType;
2330 use datatypes::schema::ColumnSchema;
2331 use moka::future::Cache;
2332 use session::context::QueryContext;
2333 use table::TableRef;
2334 use table::dist_table::DummyDataSource;
2335 use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType};
2336
2337 use crate::insert::*;
2338 use crate::test_util::{
2339 create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
2340 };
2341
2342 fn make_table_ref_with_schema(
2343 ts_name: &str,
2344 field_name: &str,
2345 field_type: ConcreteDataType,
2346 ) -> TableRef {
2347 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2348 ColumnSchema::new(
2349 ts_name,
2350 ConcreteDataType::timestamp_millisecond_datatype(),
2351 false,
2352 )
2353 .with_time_index(true),
2354 ColumnSchema::new(field_name, field_type, true),
2355 ])
2356 .unwrap()
2357 .build()
2358 .unwrap();
2359 let meta = TableMetaBuilder::empty()
2360 .schema(Arc::new(schema))
2361 .primary_key_indices(vec![])
2362 .value_indices(vec![1])
2363 .engine("mito")
2364 .next_column_id(0)
2365 .options(Default::default())
2366 .created_on(Default::default())
2367 .build()
2368 .unwrap();
2369 let info = Arc::new(
2370 TableInfoBuilder::default()
2371 .table_id(1)
2372 .table_version(0)
2373 .name("test_table")
2374 .schema_name(DEFAULT_SCHEMA_NAME)
2375 .catalog_name(DEFAULT_CATALOG_NAME)
2376 .desc(None)
2377 .table_type(TableType::Base)
2378 .meta(meta)
2379 .build()
2380 .unwrap(),
2381 );
2382 Arc::new(table::Table::new(
2383 info,
2384 table::metadata::FilterPushDownType::Unsupported,
2385 Arc::new(DummyDataSource),
2386 ))
2387 }
2388
2389 fn make_metric_physical_table_ref_with_time_unit(unit: TimeUnit) -> TableRef {
2390 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2391 ColumnSchema::new(
2392 greptime_timestamp(),
2393 ConcreteDataType::timestamp_datatype(unit),
2394 false,
2395 )
2396 .with_time_index(true),
2397 ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
2398 ])
2399 .unwrap()
2400 .build()
2401 .unwrap();
2402 let meta = TableMetaBuilder::empty()
2403 .schema(Arc::new(schema))
2404 .primary_key_indices(vec![])
2405 .value_indices(vec![1])
2406 .engine("metric")
2407 .next_column_id(0)
2408 .options(Default::default())
2409 .created_on(Default::default())
2410 .build()
2411 .unwrap();
2412 let info = Arc::new(
2413 TableInfoBuilder::default()
2414 .table_id(1)
2415 .table_version(0)
2416 .name("greptime_physical_table")
2417 .schema_name(DEFAULT_SCHEMA_NAME)
2418 .catalog_name(DEFAULT_CATALOG_NAME)
2419 .desc(None)
2420 .table_type(TableType::Base)
2421 .meta(meta)
2422 .build()
2423 .unwrap(),
2424 );
2425 Arc::new(table::Table::new(
2426 info,
2427 table::metadata::FilterPushDownType::Unsupported,
2428 Arc::new(DummyDataSource),
2429 ))
2430 }
2431
2432 fn ms_row_insert_request(timestamp_ms: i64) -> RowInsertRequest {
2433 ms_row_insert_request_named("my_metric", timestamp_ms)
2434 }
2435
2436 fn ms_row_insert_request_named(table: &str, timestamp_ms: i64) -> RowInsertRequest {
2437 RowInsertRequest {
2438 table_name: table.to_string(),
2439 rows: Some(Rows {
2440 schema: vec![
2441 time_index_column_schema(
2442 greptime_timestamp(),
2443 ColumnDataType::TimestampMillisecond,
2444 ),
2445 field_column_schema(greptime_value(), ColumnDataType::Float64),
2446 ],
2447 rows: vec![Row {
2448 values: vec![
2449 Value {
2450 value_data: Some(ValueData::TimestampMillisecondValue(timestamp_ms)),
2451 },
2452 Value {
2453 value_data: Some(ValueData::F64Value(1.0)),
2454 },
2455 ],
2456 }],
2457 }),
2458 }
2459 }
2460
2461 fn convert_row_insert_requests_time_unit(
2463 requests: &mut RowInsertRequests,
2464 target_unit: TimeUnit,
2465 ) -> Result<()> {
2466 for request in &mut requests.inserts {
2467 let Some(rows) = request.rows.as_mut() else {
2468 continue;
2469 };
2470 convert_rows_time_unit(rows, target_unit)?;
2471 }
2472 Ok(())
2473 }
2474
2475 #[test]
2476 fn test_convert_row_insert_requests_time_unit_noop_when_matching() {
2477 let mut requests = RowInsertRequests {
2478 inserts: vec![ms_row_insert_request(123)],
2479 };
2480 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Millisecond).unwrap();
2481 let rows = requests.inserts[0].rows.as_ref().unwrap();
2482 assert_eq!(
2483 rows.schema[0].datatype,
2484 ColumnDataType::TimestampMillisecond as i32
2485 );
2486 assert!(matches!(
2487 rows.rows[0].values[0].value_data,
2488 Some(ValueData::TimestampMillisecondValue(123))
2489 ));
2490 }
2491
2492 #[test]
2493 fn test_convert_row_insert_requests_time_unit_widens_losslessly() {
2494 let mut requests = RowInsertRequests {
2495 inserts: vec![ms_row_insert_request(123)],
2496 };
2497 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Microsecond).unwrap();
2498 let rows = requests.inserts[0].rows.as_ref().unwrap();
2499 assert_eq!(
2500 rows.schema[0].datatype,
2501 ColumnDataType::TimestampMicrosecond as i32
2502 );
2503 assert!(matches!(
2504 rows.rows[0].values[0].value_data,
2505 Some(ValueData::TimestampMicrosecondValue(123_000))
2506 ));
2507 }
2508
2509 #[test]
2510 fn test_convert_row_insert_requests_time_unit_truncates_on_narrowing() {
2511 let requests = |value_ns: i64| RowInsertRequests {
2514 inserts: vec![RowInsertRequest {
2515 table_name: "my_metric".to_string(),
2516 rows: Some(Rows {
2517 schema: vec![
2518 time_index_column_schema(
2519 greptime_timestamp(),
2520 ColumnDataType::TimestampNanosecond,
2521 ),
2522 field_column_schema(greptime_value(), ColumnDataType::Float64),
2523 ],
2524 rows: vec![Row {
2525 values: vec![
2526 Value {
2527 value_data: Some(ValueData::TimestampNanosecondValue(value_ns)),
2528 },
2529 Value {
2530 value_data: Some(ValueData::F64Value(1.0)),
2531 },
2532 ],
2533 }],
2534 }),
2535 }],
2536 };
2537
2538 let mut reqs = requests(123_456_789);
2539 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
2540 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2541 assert!(matches!(
2542 rows.rows[0].values[0].value_data,
2543 Some(ValueData::TimestampMicrosecondValue(123_456))
2544 ));
2545
2546 let mut reqs = requests(123_456_789);
2547 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Millisecond).unwrap();
2548 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2549 assert!(matches!(
2550 rows.rows[0].values[0].value_data,
2551 Some(ValueData::TimestampMillisecondValue(123))
2552 ));
2553
2554 let mut reqs = requests(-123_456_789);
2555 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
2556 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2557 assert!(matches!(
2558 rows.rows[0].values[0].value_data,
2559 Some(ValueData::TimestampMicrosecondValue(-123_457))
2560 ));
2561 }
2562
2563 #[test]
2564 fn test_convert_row_insert_requests_time_unit_overflow() {
2565 let mut requests = RowInsertRequests {
2566 inserts: vec![ms_row_insert_request(i64::MAX)],
2567 };
2568 let err =
2569 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Nanosecond).unwrap_err();
2570 assert!(err.to_string().contains("overflows"), "{err}");
2571 }
2572
2573 #[test]
2574 fn test_table_time_index_unit() {
2575 assert_eq!(
2576 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
2577 TimeUnit::Microsecond
2578 )),
2579 Some(TimeUnit::Microsecond)
2580 );
2581 assert_eq!(
2582 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
2583 TimeUnit::Millisecond
2584 )),
2585 Some(TimeUnit::Millisecond)
2586 );
2587 }
2588
2589 #[tokio::test]
2590 async fn test_accommodate_existing_schema_and_reject_kind_changes() {
2591 let ts_name = "my_ts";
2592 let field_name = "my_field";
2593 let table =
2594 make_table_ref_with_schema(ts_name, field_name, ConcreteDataType::float64_datatype());
2595
2596 let mut req = RowInsertRequest {
2598 table_name: "test_table".to_string(),
2599 rows: Some(Rows {
2600 schema: vec![
2601 time_index_column_schema("ts_wrong", ColumnDataType::TimestampMillisecond),
2602 field_column_schema("field_wrong", ColumnDataType::Float64),
2603 ],
2604 rows: vec![api::v1::Row {
2605 values: vec![Value::default(), Value::default()],
2606 }],
2607 }),
2608 };
2609 let ctx = Arc::new(QueryContext::with(
2610 DEFAULT_CATALOG_NAME,
2611 DEFAULT_SCHEMA_NAME,
2612 ));
2613
2614 let kv_backend = prepare_mocked_backend().await;
2615 let inserter = Inserter::new(
2616 catalog::memory::MemoryCatalogManager::new(),
2617 create_partition_rule_manager(kv_backend.clone()).await,
2618 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
2619 Arc::new(new_table_flownode_set_cache(
2620 String::new(),
2621 Cache::new(100),
2622 kv_backend.clone(),
2623 )),
2624 true,
2625 );
2626 assert!(
2628 inserter
2629 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, false, false, false)
2630 .unwrap()
2631 .is_none()
2632 );
2633 let alter_expr = inserter
2634 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true, true)
2635 .unwrap();
2636 assert!(alter_expr.is_none());
2637
2638 let req_schema = req.rows.as_ref().unwrap().schema.clone();
2640 assert_eq!(req_schema[0].column_name, ts_name);
2641 assert_eq!(req_schema[1].column_name, field_name);
2642
2643 let (datatype, datatype_extension) =
2644 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
2645 .unwrap()
2646 .into_parts();
2647 let mut histogram_req = RowInsertRequest {
2648 table_name: "test_table".to_string(),
2649 rows: Some(Rows {
2650 schema: vec![
2651 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2652 api::v1::ColumnSchema {
2653 column_name: greptime_native_histogram().to_string(),
2654 datatype: datatype as i32,
2655 semantic_type: SemanticType::Field as i32,
2656 datatype_extension,
2657 options: None,
2658 },
2659 ],
2660 rows: vec![],
2661 }),
2662 };
2663 let error = inserter
2664 .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, false, true, true)
2665 .unwrap_err();
2666 assert!(
2667 error
2668 .to_string()
2669 .contains("cannot mix native histogram and float sample fields")
2670 );
2671
2672 let histogram_table = make_table_ref_with_schema(
2673 "ts",
2674 greptime_native_histogram(),
2675 native_histogram_value_type().clone(),
2676 );
2677 let mut sample_req = RowInsertRequest {
2678 table_name: "test_table".to_string(),
2679 rows: Some(Rows {
2680 schema: vec![
2681 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2682 field_column_schema(greptime_value(), ColumnDataType::Float64),
2683 ],
2684 rows: vec![],
2685 }),
2686 };
2687 let error = inserter
2688 .get_alter_table_expr_on_demand(
2689 &mut sample_req,
2690 &histogram_table,
2691 &ctx,
2692 false,
2693 true,
2694 true,
2695 )
2696 .unwrap_err();
2697 assert!(
2698 error
2699 .to_string()
2700 .contains("cannot mix native histogram and float sample fields")
2701 );
2702 }
2703
2704 #[test]
2705 fn test_native_histogram_detection_survives_prefix_change() {
2706 set_default_prefix(Some("custom")).unwrap();
2707 let table = make_table_ref_with_schema(
2708 "custom_timestamp",
2709 NATIVE_HISTOGRAM_FIELD,
2710 native_histogram_value_type().clone(),
2711 );
2712 let (datatype, datatype_extension) =
2713 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
2714 .unwrap()
2715 .into_parts();
2716 let request_schema = [api::v1::ColumnSchema {
2717 column_name: greptime_native_histogram().to_string(),
2718 datatype: datatype as i32,
2719 semantic_type: SemanticType::Field as i32,
2720 datatype_extension,
2721 options: None,
2722 }];
2723
2724 assert!(request_is_native_histogram(&request_schema));
2725 assert!(table_is_native_histogram(&table));
2726 }
2727
2728 #[tokio::test]
2730 async fn test_write_meter_admission() {
2731 use std::cell::Cell;
2732 use std::sync::Mutex;
2733 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
2734
2735 use api::region::RegionResponse;
2736 use api::v1::region::region_request::Body;
2737 use arrow::array::{Int32Array, TimestampMillisecondArray};
2738 use arrow::record_batch::RecordBatch;
2739 use bytes::Bytes;
2740 use common_error::ext::{ErrorExt, RetryHint};
2741 use common_error::status_code::StatusCode;
2742 use common_grpc::flight::{FlightEncoder, FlightMessage};
2743 use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher;
2744 use futures::future::BoxFuture;
2745 use meter_core::ItemCalculator;
2746 use meter_core::collect::{Collect, WriteRejected};
2747 use meter_core::data::MeterRecord;
2748 use meter_core::global::global_registry;
2749 use session::context::Channel;
2750
2751 const CATALOG: &str = "write_meter_test";
2752
2753 #[derive(Default)]
2754 struct Meter {
2755 reject: AtomicBool,
2756 attempts: Mutex<Vec<MeterRecord>>,
2757 accepted_value: AtomicU64,
2758 }
2759
2760 impl Collect for Meter {
2761 fn on_write(
2762 &self,
2763 record: MeterRecord,
2764 ) -> BoxFuture<'_, std::result::Result<(), WriteRejected>> {
2765 Box::pin(async move {
2766 if record.catalog != CATALOG {
2767 return Ok(());
2768 }
2769 let value = record.value;
2770 self.attempts.lock().unwrap().push(record);
2771 if self.reject.load(Ordering::Relaxed) {
2772 return Err(WriteRejected::new("database row quota exhausted"));
2773 }
2774 self.accepted_value.fetch_add(value, Ordering::Relaxed);
2775 Ok(())
2776 })
2777 }
2778
2779 fn on_read(&self, _: MeterRecord) {}
2780 }
2781
2782 impl ItemCalculator<InstantAndNormalInsertRequests> for Meter {
2783 fn calc(&self, _: &InstantAndNormalInsertRequests) -> u64 {
2784 17
2785 }
2786 }
2787
2788 let kv_backend = prepare_mocked_backend().await;
2789 let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
2790 let (sender, mut dispatched) = tokio::sync::mpsc::channel(16);
2791 let watcher = DatanodeWatcher::new(sender).with_handler(|_, request| {
2792 let rows = match request.body.unwrap() {
2793 Body::Inserts(requests) => requests
2794 .requests
2795 .iter()
2796 .filter_map(|request| request.rows.as_ref())
2797 .map(|rows| rows.rows.len())
2798 .sum(),
2799 Body::BulkInsert(_) => 2,
2801 body => panic!("unexpected request: {body:?}"),
2802 };
2803 Ok(RegionResponse::new(rows))
2804 });
2805 let flow_cache = Cache::new(100);
2806 let inserter = Inserter::new(
2807 catalog::memory::MemoryCatalogManager::new(),
2808 partition_manager,
2809 Arc::new(MockDatanodeManager::new(watcher)),
2810 Arc::new(new_table_flownode_set_cache(
2811 String::new(),
2812 flow_cache.clone(),
2813 kv_backend,
2814 )),
2815 true,
2816 );
2817 let mut table_info = new_test_table_info(1, "table_1", [1].into_iter());
2818 table_info.catalog_name = CATALOG.to_string();
2819 table_info.schema_name = "target_db".to_string();
2820 let table_info = Arc::new(table_info);
2821 let table_infos = HashMap::from_iter([(1, table_info.clone())]);
2822 let ctx = Arc::new(QueryContext::with_channel(
2823 DEFAULT_CATALOG_NAME,
2824 DEFAULT_SCHEMA_NAME,
2825 Channel::Postgres,
2826 ));
2827 let rows_request = || {
2828 let build = |num_rows| RegionInsertRequests {
2829 requests: vec![RegionInsertRequest {
2830 region_id: RegionId::new(1, 1).as_u64(),
2831 rows: Some(Rows {
2832 schema: vec![],
2833 rows: vec![api::v1::Row { values: vec![] }; num_rows],
2834 }),
2835 ..Default::default()
2836 }],
2837 };
2838 InstantAndNormalInsertRequests {
2839 normal_requests: build(3),
2840 instant_requests: build(2),
2841 }
2842 };
2843
2844 let output = inserter
2846 .do_request(rows_request(), &table_infos, &ctx)
2847 .await
2848 .unwrap();
2849 assert_eq!(output.meta.cost, 0);
2850 assert!(matches!(output.data, OutputData::AffectedRows(3)));
2851 dispatched.try_recv().unwrap();
2852
2853 let meter = Arc::new(Meter::default());
2854 global_registry().set_collector(meter.clone());
2855 global_registry().register_calculator(meter.clone());
2856 let enabled = Cell::new(false);
2858 write_meter!({
2859 enabled.set(true);
2860 MeterRecord::new("probe".into(), "probe".into(), 0, 0, 0)
2861 })
2862 .await
2863 .unwrap();
2864 let enabled = enabled.get();
2865
2866 let output = inserter
2867 .do_request(rows_request(), &table_infos, &ctx)
2868 .await
2869 .unwrap();
2870 assert_eq!(output.meta.cost, if enabled { 17 } else { 0 });
2871 assert!(matches!(output.data, OutputData::AffectedRows(3)));
2872 dispatched.try_recv().unwrap();
2873 assert!(flow_cache.contains_key(&1));
2874 flow_cache.invalidate_all();
2875 meter.reject.store(true, Ordering::Relaxed);
2876 let result = inserter
2877 .do_request(rows_request(), &table_infos, &ctx)
2878 .await;
2879 if enabled {
2880 let error = result.unwrap_err();
2881 assert_eq!(error.status_code(), StatusCode::RateLimited);
2882 assert_eq!(error.retry_hint(), RetryHint::Retryable);
2883 assert!(error.to_string().contains("database row quota exhausted"));
2884 assert!(dispatched.try_recv().is_err());
2885 assert!(
2886 !flow_cache.contains_key(&1),
2887 "rejected write reached flow mirroring"
2888 );
2889 assert_eq!(meter.accepted_value.load(Ordering::Relaxed), 17);
2890 let attempts = meter.attempts.lock().unwrap();
2891 assert_eq!(attempts.len(), 2);
2892 for record in attempts.iter() {
2893 assert_eq!(record.catalog, CATALOG);
2894 assert_eq!(record.schema, "target_db");
2895 assert_eq!(
2896 (record.rows, record.value, record.source),
2897 (5, 17, Channel::Postgres as u8)
2898 );
2899 }
2900 } else {
2901 assert_eq!(result.unwrap().meta.cost, 0);
2902 dispatched.try_recv().unwrap();
2903 assert!(meter.attempts.lock().unwrap().is_empty());
2904 }
2905 meter.attempts.lock().unwrap().clear();
2906
2907 if enabled {
2909 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
2910 let rows = Rows {
2911 schema: vec![
2912 api::v1::helper::tag_column_schema("a", ColumnDataType::Int32),
2913 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2914 field_column_schema("b", ColumnDataType::Int32),
2915 ],
2916 rows: vec![api::v1::Row {
2917 values: vec![
2918 api::v1::value::ValueData::I32Value(60).into(),
2919 Value {
2920 value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue(
2921 0,
2922 )),
2923 },
2924 api::v1::value::ValueData::I32Value(0).into(),
2925 ],
2926 }],
2927 };
2928 let error = inserter
2929 .submit_table_rows(rows, table_info.clone(), ctx.clone(), &batcher)
2930 .await
2931 .unwrap_err();
2932 assert_eq!(error.status_code(), StatusCode::RateLimited);
2933 let mut attempts = meter.attempts.lock().unwrap();
2934 assert_eq!(attempts.len(), 1);
2935 assert_eq!(attempts[0].catalog, CATALOG);
2936 assert_eq!(attempts[0].schema, "target_db");
2937 assert_eq!(attempts[0].rows, 1);
2938 attempts.clear();
2939 }
2940
2941 let table = Arc::new(table::Table::new(
2942 table_info.clone(),
2943 table::metadata::FilterPushDownType::Unsupported,
2944 Arc::new(DummyDataSource),
2945 ));
2946 let batch = RecordBatch::try_new(
2947 table_info.meta.schema.arrow_schema().clone(),
2948 vec![
2949 Arc::new(Int32Array::from(vec![60, 70])),
2950 Arc::new(TimestampMillisecondArray::from(vec![0, 1])),
2951 Arc::new(Int32Array::from(vec![0, 0])),
2952 ],
2953 )
2954 .unwrap();
2955 let bulk_insert = |batch: RecordBatch| {
2956 let flight_data = FlightEncoder::default()
2957 .encode(FlightMessage::RecordBatch(batch.clone()))
2958 .into_iter()
2959 .next()
2960 .unwrap();
2961 inserter.handle_bulk_insert(
2962 table.clone(),
2963 flight_data,
2964 batch,
2965 Bytes::new(),
2966 false,
2967 Channel::Grpc,
2968 )
2969 };
2970 assert_eq!(bulk_insert(batch.slice(0, 0)).await.unwrap(), 0);
2972 assert!(meter.attempts.lock().unwrap().is_empty());
2973 assert!(dispatched.try_recv().is_err());
2974
2975 for reject in [false, true, false] {
2977 meter.reject.store(reject, Ordering::Relaxed);
2978 let result = bulk_insert(batch.clone()).await;
2979 if enabled && reject {
2980 let error = result.unwrap_err();
2981 assert_eq!(error.status_code(), StatusCode::RateLimited);
2982 assert_eq!(error.retry_hint(), RetryHint::Retryable);
2983 assert!(dispatched.try_recv().is_err());
2984 } else {
2985 assert_eq!(result.unwrap(), 2);
2986 dispatched.try_recv().unwrap();
2987 }
2988 }
2989 {
2990 let attempts = meter.attempts.lock().unwrap();
2991 assert_eq!(attempts.len(), if enabled { 3 } else { 0 });
2992 for record in attempts.iter() {
2993 assert_eq!(record.catalog, CATALOG);
2994 assert_eq!(record.schema, "target_db");
2995 assert_eq!(
2996 (record.rows, record.value, record.source),
2997 (2, 0, Channel::Grpc as u8)
2998 );
2999 }
3000 }
3001 assert_eq!(
3002 meter.accepted_value.load(Ordering::Relaxed),
3003 if enabled { 17 } else { 0 }
3004 );
3005
3006 meter.attempts.lock().unwrap().clear();
3009 let original = Arc::new(QueryContext::with_channel(
3010 CATALOG,
3011 "a",
3012 Channel::Prometheus,
3013 ));
3014 let mut batches = ["a", "b", "a"].map(|schema| {
3015 let ctx = if schema == "a" {
3016 original.clone()
3017 } else {
3018 Arc::new(QueryContext::with_channel(
3019 CATALOG,
3020 schema,
3021 Channel::Prometheus,
3022 ))
3023 };
3024 (
3025 ctx,
3026 RowInsertRequests {
3027 inserts: vec![RowInsertRequest {
3028 table_name: "data".into(),
3029 rows: Some(Rows {
3030 schema: vec![],
3031 rows: vec![api::v1::Row::default(); 2],
3032 }),
3033 }],
3034 },
3035 )
3036 });
3037 admit_row_insert_batches(&mut batches).await.unwrap();
3038 admit_row_insert_batches(&mut batches).await.unwrap();
3039 assert_eq!(original.write_rows_to_admit(CATALOG, "a", 4), 4);
3040 for (ctx, _) in &batches {
3041 assert_eq!(
3042 ctx.write_rows_to_admit(CATALOG, &ctx.current_schema(), 2),
3043 0
3044 );
3045 assert_eq!(ctx.channel(), Channel::Prometheus);
3046 }
3047 {
3048 let attempts = meter.attempts.lock().unwrap();
3049 let totals = attempts
3050 .iter()
3051 .map(|r| (r.schema.as_str(), r.rows, r.value))
3052 .collect::<Vec<_>>();
3053 assert_eq!(
3054 totals,
3055 if enabled {
3056 vec![("a", 4, 0), ("b", 2, 0)]
3057 } else {
3058 vec![]
3059 }
3060 );
3061 }
3062 meter.attempts.lock().unwrap().clear();
3063 meter.reject.store(false, Ordering::Relaxed);
3064 let ctx = Arc::new(QueryContext::with_channel(
3065 CATALOG,
3066 "logical",
3067 Channel::Otlp,
3068 ));
3069 let admitted = admit_write(2, &ctx).await.unwrap();
3070 let mut requests = RowInsertRequests {
3071 inserts: vec![RowInsertRequest {
3072 table_name: "metric".to_string(),
3073 rows: Some(Rows {
3074 schema: vec![],
3075 rows: vec![api::v1::Row::default(); 2],
3076 }),
3077 }],
3078 };
3079 let original = requests.clone();
3080 let cost = Inserter::meter_row_inserts(&mut requests, &admitted)
3081 .await
3082 .unwrap();
3083 assert_eq!(cost, if enabled { 17 } else { 0 });
3084 assert_eq!(requests, original);
3085 let records = meter
3086 .attempts
3087 .lock()
3088 .unwrap()
3089 .iter()
3090 .map(|record| (record.rows, record.value))
3091 .collect::<Vec<_>>();
3092 assert_eq!(
3093 records,
3094 if enabled {
3095 vec![(2, 0), (0, 17)]
3096 } else {
3097 vec![]
3098 }
3099 );
3100 meter.reject.store(true, Ordering::Relaxed);
3101 let result = Inserter::meter_row_inserts(&mut requests, &admitted).await;
3102 if enabled {
3103 assert_eq!(result.unwrap_err().status_code(), StatusCode::RateLimited);
3104 } else {
3105 assert_eq!(result.unwrap(), 0);
3106 }
3107 assert_eq!(requests, original);
3108 }
3109
3110 #[test]
3111 fn test_trace_create_options_copy_compaction_hints() {
3112 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3113 ctx.set_extension("compaction.twcs.time_window", "1h");
3114 ctx.set_extension("compaction.twcs.max_output_file_size", "256MB");
3115 ctx.set_extension("compaction.twcs.remote_compaction", "true");
3116 ctx.set_extension("compaction.unknown", "ignored");
3117 let ctx = Arc::new(ctx);
3118 let mut table_options = Default::default();
3119
3120 fill_table_options_for_create(
3121 &mut table_options,
3122 &AutoCreateTableType::Trace {
3123 alter_existing: true,
3124 },
3125 &ctx,
3126 );
3127
3128 assert_eq!(
3129 Some("twcs"),
3130 table_options.get(COMPACTION_TYPE).map(String::as_str)
3131 );
3132 assert_eq!(
3133 Some("1h"),
3134 table_options
3135 .get("compaction.twcs.time_window")
3136 .map(String::as_str)
3137 );
3138 assert_eq!(
3139 Some("256MB"),
3140 table_options
3141 .get("compaction.twcs.max_output_file_size")
3142 .map(String::as_str)
3143 );
3144 assert_eq!(
3145 Some("true"),
3146 table_options
3147 .get("compaction.twcs.remote_compaction")
3148 .map(String::as_str)
3149 );
3150 assert!(!table_options.contains_key("compaction.unknown"));
3151 }
3152
3153 #[test]
3154 fn test_skip_wal_does_not_change_table_options() {
3155 check_skip_wal_does_not_change_table_options(false);
3156 check_skip_wal_does_not_change_table_options(true);
3157 }
3158
3159 fn check_skip_wal_does_not_change_table_options(skip_wal: bool) {
3160 let ctx = Arc::new(QueryContext::with(
3161 DEFAULT_CATALOG_NAME,
3162 DEFAULT_SCHEMA_NAME,
3163 ));
3164 ctx.set_skip_wal(skip_wal);
3165 let mut options = Default::default();
3166 fill_table_options_for_create(&mut options, &AutoCreateTableType::Physical, &ctx);
3167 assert!(!options.contains_key(session::hints::INSERT_SKIP_WAL_HINT));
3168 assert!(!options.contains_key("skip_wal"));
3169 }
3170
3171 #[test]
3172 fn test_last_non_null_create_options_preserve_default_without_append_mode() {
3173 let ctx = Arc::new(QueryContext::with(
3174 DEFAULT_CATALOG_NAME,
3175 DEFAULT_SCHEMA_NAME,
3176 ));
3177 let mut table_options = Default::default();
3178
3179 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3180
3181 assert_eq!(
3182 Some("last_non_null"),
3183 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3184 );
3185 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3186 }
3187
3188 #[test]
3189 fn test_fill_table_options_copies_semantic_extensions() {
3190 use table::requests::{
3191 SEMANTIC_METRIC_TYPE, SEMANTIC_PER_TABLE_INDEX_KEY, SEMANTIC_SIGNAL_TYPE,
3192 SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, SOURCE_OPENTELEMETRY,
3193 };
3194
3195 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3196 ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
3197 ctx.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
3198 ctx.set_extension(SEMANTIC_SOURCE_VERSION, "2.0");
3199 ctx.set_extension(SEMANTIC_METRIC_TYPE, "bogus");
3200 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, "{}");
3202 let ctx = Arc::new(ctx);
3203 let mut table_options = Default::default();
3204
3205 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::Physical, &ctx);
3206
3207 assert_eq!(
3208 Some(SIGNAL_TYPE_METRIC),
3209 table_options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str)
3210 );
3211 assert_eq!(
3212 Some(SOURCE_OPENTELEMETRY),
3213 table_options.get(SEMANTIC_SOURCE).map(String::as_str)
3214 );
3215 assert_eq!(
3216 Some("2.0"),
3217 table_options
3218 .get(SEMANTIC_SOURCE_VERSION)
3219 .map(String::as_str)
3220 );
3221 assert!(!table_options.contains_key(SEMANTIC_METRIC_TYPE));
3222 assert!(!table_options.contains_key(SEMANTIC_PER_TABLE_INDEX_KEY));
3223 }
3224
3225 #[test]
3226 fn test_apply_per_table_semantic_options() {
3227 use table::requests::{
3228 SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, SEMANTIC_PER_TABLE_INDEX_KEY,
3229 };
3230
3231 let index = format!(
3232 r#"{{
3233 "{DEFAULT_SCHEMA_NAME}": {{
3234 "http_requests_total": {{
3235 "greptime.semantic.metric.type": "counter",
3236 "greptime.semantic.metric.unit": "By",
3237 "greptime.semantic.metric.type_BOGUS": "x"
3238 }},
3239 "other_table": {{
3240 "greptime.semantic.metric.type": "gauge"
3241 }}
3242 }},
3243 "other_schema": {{
3244 "http_requests_total": {{
3245 "greptime.semantic.metric.type": "gauge"
3246 }}
3247 }}
3248 }}"#
3249 );
3250 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3251 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index);
3252 let ctx = Arc::new(ctx);
3253
3254 let index = parse_per_table_semantic_index(&ctx);
3255 assert!(index.is_some());
3256 let index = index.as_ref();
3257
3258 let mut table_options = std::collections::HashMap::new();
3259 apply_per_table_semantic_options(
3260 &mut table_options,
3261 index,
3262 DEFAULT_SCHEMA_NAME,
3263 "http_requests_total",
3264 );
3265 assert_eq!(
3267 table_options.get(SEMANTIC_METRIC_TYPE).map(String::as_str),
3268 Some("counter")
3269 );
3270 assert_eq!(
3271 table_options.get(SEMANTIC_METRIC_UNIT).map(String::as_str),
3272 Some("By")
3273 );
3274 assert!(!table_options.contains_key("greptime.semantic.metric.type_BOGUS"));
3277 assert_eq!(table_options.len(), 2);
3278
3279 let mut empty = std::collections::HashMap::new();
3280 apply_per_table_semantic_options(&mut empty, index, DEFAULT_SCHEMA_NAME, "not_in_index");
3281 assert!(empty.is_empty());
3282
3283 let mut opts = std::collections::HashMap::new();
3286 apply_per_table_semantic_options(
3287 &mut opts,
3288 index,
3289 "schema_without_entry",
3290 "http_requests_total",
3291 );
3292 assert!(opts.is_empty());
3293
3294 let bare = Arc::new(QueryContext::with(
3296 DEFAULT_CATALOG_NAME,
3297 DEFAULT_SCHEMA_NAME,
3298 ));
3299 assert!(parse_per_table_semantic_index(&bare).is_none());
3300 let mut opts = std::collections::HashMap::new();
3301 apply_per_table_semantic_options(
3302 &mut opts,
3303 None,
3304 DEFAULT_SCHEMA_NAME,
3305 "http_requests_total",
3306 );
3307 assert!(opts.is_empty());
3308 }
3309
3310 #[test]
3311 fn test_last_non_null_create_options_preserve_default_with_append_mode_false() {
3312 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3313 ctx.set_extension(APPEND_MODE_KEY, "false");
3314 let ctx = Arc::new(ctx);
3315 let mut table_options = Default::default();
3316
3317 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3318
3319 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3320 assert_eq!(
3321 Some("last_non_null"),
3322 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3323 );
3324 }
3325
3326 #[test]
3327 fn test_last_non_null_create_options_use_configured_merge_mode() {
3328 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3329 ctx.set_extension(MERGE_MODE_KEY, "last_row");
3330 let ctx = Arc::new(ctx);
3331 let mut table_options = Default::default();
3332
3333 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3334
3335 assert_eq!(
3336 Some("last_row"),
3337 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3338 );
3339 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3340 }
3341
3342 #[test]
3343 fn test_last_non_null_create_options_use_last_row_with_append_mode_true() {
3344 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3345 ctx.set_extension(APPEND_MODE_KEY, "true");
3346 let ctx = Arc::new(ctx);
3347 let mut table_options = Default::default();
3348
3349 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3350
3351 assert_eq!(
3352 Some("true"),
3353 table_options.get(APPEND_MODE_KEY).map(String::as_str)
3354 );
3355 assert_eq!(
3356 Some("last_row"),
3357 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3358 );
3359 }
3360
3361 struct UnexpectedBatcher;
3362
3363 #[async_trait::async_trait]
3364 impl PendingRowsBatcher for UnexpectedBatcher {
3365 async fn acquire(&self) -> Result<Arc<tokio::sync::OwnedSemaphorePermit>> {
3366 panic!("empty writes must not acquire batch admission")
3367 }
3368
3369 async fn submit(
3370 &self,
3371 _table_info: TableInfoRef,
3372 _batch: arrow::record_batch::RecordBatch,
3373 _ctx: QueryContextRef,
3374 _permit: Arc<tokio::sync::OwnedSemaphorePermit>,
3375 ) -> Result<Output> {
3376 panic!("empty writes must not submit a batch")
3377 }
3378 }
3379
3380 async fn batcher_test_inserter() -> Inserter {
3381 let kv_backend = prepare_mocked_backend().await;
3382 Inserter::new(
3383 catalog::memory::MemoryCatalogManager::new(),
3384 create_partition_rule_manager(kv_backend.clone()).await,
3385 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3386 Arc::new(new_table_flownode_set_cache(
3387 String::new(),
3388 Cache::new(100),
3389 kv_backend,
3390 )),
3391 true,
3392 )
3393 }
3394
3395 #[tokio::test]
3396 async fn test_logical_batcher_eligibility() {
3397 use catalog::RegisterTableRequest;
3398 use catalog::memory::MemoryCatalogManager;
3399 use common_meta::instruction::{CacheIdent, CreateFlow};
3400 use common_meta::kv_backend::KvBackendRef;
3401 use common_meta::kv_backend::memory::MemoryKvBackend;
3402 use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder};
3403 let requests = RowInsertRequests {
3404 inserts: vec![RowInsertRequest {
3405 table_name: "test_table".to_string(),
3406 rows: None,
3407 }],
3408 };
3409 let original =
3410 make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3411 .table_info();
3412 for case in [
3413 "eligible",
3414 "physical",
3415 "ordinary",
3416 "instant",
3417 "disabled",
3418 "hint",
3419 "flow",
3420 "required_tag",
3421 "default",
3422 ] {
3423 let mut info = (*original).clone();
3424 info.meta.engine = METRIC_ENGINE_NAME.to_string();
3425 info.meta.options.extra_options.insert(
3426 LOGICAL_TABLE_METADATA_KEY.to_string(),
3427 "physical".to_string(),
3428 );
3429 let mut ctx = QueryContext::arc().fork();
3430 match case {
3431 "physical" => {
3432 info.meta
3433 .options
3434 .extra_options
3435 .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "other".to_string());
3436 }
3437 "ordinary" => info.meta.engine = "mito".to_string(),
3438 "instant" => info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant),
3439 "hint" => ctx.set_extension(AUTO_CREATE_TABLE_KEY, "false"),
3440 "required_tag" | "default" => {
3441 let mut columns = info.meta.schema.column_schemas().to_vec();
3442 let mut tag = ColumnSchema::new(
3443 "tag",
3444 ConcreteDataType::string_datatype(),
3445 case != "required_tag",
3446 );
3447 if case == "default" {
3448 tag = tag
3449 .with_default_constraint(Some(ColumnDefaultConstraint::null_value()))
3450 .unwrap();
3451 }
3452 columns.push(tag);
3453 info.meta.schema = Arc::new(
3454 SchemaBuilder::try_from_columns(columns)
3455 .unwrap()
3456 .build()
3457 .unwrap(),
3458 );
3459 info.meta.primary_key_indices = vec![2];
3460 }
3461 _ => {}
3462 }
3463 let catalog = MemoryCatalogManager::with_default_setup();
3464 let table = Arc::new(table::Table::new(
3465 Arc::new(info),
3466 table::metadata::FilterPushDownType::Unsupported,
3467 Arc::new(DummyDataSource),
3468 ));
3469 catalog
3470 .register_table_sync(RegisterTableRequest {
3471 catalog: DEFAULT_CATALOG_NAME.to_string(),
3472 schema: DEFAULT_SCHEMA_NAME.to_string(),
3473 table_name: "test_table".to_string(),
3474 table_id: 1,
3475 table,
3476 })
3477 .unwrap();
3478 let mut inserter = batcher_test_inserter().await;
3479 inserter.catalog_manager = catalog;
3480 inserter.auto_create_table = case != "disabled";
3481 let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::default());
3482 inserter.table_flownode_set_cache = Arc::new(new_table_flownode_set_cache(
3483 String::new(),
3484 Cache::new(10),
3485 kv_backend,
3486 ));
3487 if case == "flow" {
3488 inserter
3489 .table_flownode_set_cache
3490 .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
3491 flow_id: 1,
3492 source_table_ids: vec![1],
3493 partition_to_peer_mapping: vec![(0, Peer::empty(1))],
3494 })])
3495 .await
3496 .unwrap();
3497 }
3498 assert_eq!(
3499 inserter
3500 .can_batch_metric_rows(&requests, &Arc::new(ctx), "physical")
3501 .await
3502 .unwrap(),
3503 case == "eligible",
3504 "{case}"
3505 );
3506 }
3507 }
3508
3509 #[tokio::test]
3510 async fn test_batcher_meter_preserves_request() {
3511 let mut requests = RowInsertRequests {
3512 inserts: vec![RowInsertRequest {
3513 table_name: "sample".to_string(),
3514 rows: Some(Rows {
3515 schema: vec![],
3516 rows: vec![api::v1::Row {
3517 values: vec![Value {
3518 value_data: Some(api::v1::value::ValueData::F64Value(1.5)),
3519 }],
3520 }],
3521 }),
3522 }],
3523 };
3524 let expected = requests.clone();
3525 Inserter::meter_row_inserts(&mut requests, &QueryContext::arc())
3526 .await
3527 .unwrap();
3528 assert_eq!(requests, expected);
3529 }
3530
3531 #[tokio::test]
3532 async fn test_instant_table_bypasses_batcher() {
3533 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3534 let inserter = batcher_test_inserter()
3535 .await
3536 .with_pending_rows_batcher(Some(batcher));
3537 let mut ctx = session::context::QueryContextBuilder::default().build();
3538 ctx.set_batching_enabled(true);
3539 let ctx = Arc::new(ctx);
3540 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3541 .table_info();
3542 assert!(inserter.table_batcher(&table, &ctx).is_some());
3543 let mut instant = (*table).clone();
3544 instant.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
3545 assert!(inserter.table_batcher(&Arc::new(instant), &ctx).is_none());
3546 }
3547
3548 #[tokio::test]
3549 async fn test_empty_prepared_rows_skip_batcher() {
3550 let inserter = batcher_test_inserter().await;
3551 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3552 .table_info();
3553 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3554 let ctx = QueryContext::arc();
3555 let output = inserter
3556 .submit_table_rows(
3557 Rows {
3558 schema: vec![],
3559 rows: vec![],
3560 },
3561 table.clone(),
3562 ctx.clone(),
3563 &batcher,
3564 )
3565 .await
3566 .unwrap();
3567 assert!(matches!(output.data, OutputData::AffectedRows(0)));
3568 let output = inserter
3569 .submit_pending_rows(
3570 RowInsertRequests {
3571 inserts: vec![RowInsertRequest {
3572 table_name: table.name.clone(),
3573 rows: None,
3574 }],
3575 },
3576 HashMap::from_iter([(table.table_id(), table)]),
3577 ctx,
3578 &batcher,
3579 )
3580 .await
3581 .unwrap();
3582 assert!(matches!(output.data, OutputData::AffectedRows(0)));
3583 }
3584
3585 #[derive(Clone)]
3589 struct GatedFlownodeHandler {
3590 gate: Arc<tokio::sync::Semaphore>,
3591 affected_rows: u64,
3592 }
3593
3594 #[async_trait::async_trait]
3595 impl MockFlownodeHandler for GatedFlownodeHandler {
3596 async fn handle_inserts(
3597 &self,
3598 _peer: &Peer,
3599 _requests: api::v1::region::InsertRequests,
3600 ) -> common_meta::error::Result<FlowResponse> {
3601 let _permit = self.gate.acquire().await.unwrap();
3602 Ok(FlowResponse {
3603 affected_rows: self.affected_rows,
3604 ..Default::default()
3605 })
3606 }
3607 }
3608
3609 fn flownode_peer() -> Peer {
3610 Peer {
3611 id: 1,
3612 addr: "127.0.0.1:4001".to_string(),
3613 }
3614 }
3615
3616 #[derive(Clone)]
3617 struct RecordingNodeManager {
3618 datanode_dispatch: tokio::sync::mpsc::UnboundedSender<usize>,
3619 flownode_dispatch: tokio::sync::mpsc::UnboundedSender<Peer>,
3620 gates: Arc<std::sync::Mutex<HashMap<u64, Arc<tokio::sync::Semaphore>>>>,
3621 fail_flownode: bool,
3622 }
3623
3624 #[async_trait::async_trait]
3625 impl common_meta::node_manager::DatanodeManager for RecordingNodeManager {
3626 async fn datanode(&self, peer: &Peer) -> common_meta::node_manager::DatanodeRef {
3627 Arc::new(RecordingNode {
3628 peer: peer.clone(),
3629 manager: self.clone(),
3630 })
3631 }
3632 }
3633
3634 #[async_trait::async_trait]
3635 impl common_meta::node_manager::FlownodeManager for RecordingNodeManager {
3636 async fn flownode(&self, peer: &Peer) -> common_meta::node_manager::FlownodeRef {
3637 Arc::new(RecordingNode {
3638 peer: peer.clone(),
3639 manager: self.clone(),
3640 })
3641 }
3642 }
3643
3644 struct RecordingNode {
3645 peer: Peer,
3646 manager: RecordingNodeManager,
3647 }
3648
3649 #[async_trait::async_trait]
3650 impl common_meta::node_manager::Datanode for RecordingNode {
3651 async fn handle_request(
3652 &self,
3653 _request: Arc<api::v1::region::RegionRequest>,
3654 ) -> common_meta::error::Result<api::region::RegionResponse> {
3655 let rows = match _request.body.as_ref() {
3656 Some(api::v1::region::region_request::Body::Inserts(requests)) => requests
3657 .requests
3658 .iter()
3659 .filter_map(|request| request.rows.as_ref())
3660 .map(|rows| rows.rows.len())
3661 .sum(),
3662 _ => 0,
3663 };
3664 let _ = self.manager.datanode_dispatch.send(rows);
3665 Ok(api::region::RegionResponse::new(rows))
3666 }
3667
3668 async fn handle_query(
3669 &self,
3670 _request: common_query::request::QueryRequest,
3671 ) -> common_meta::error::Result<common_recordbatch::SendableRecordBatchStream> {
3672 unreachable!()
3673 }
3674 }
3675
3676 #[async_trait::async_trait]
3677 impl common_meta::node_manager::Flownode for RecordingNode {
3678 async fn handle(
3679 &self,
3680 _request: api::v1::flow::FlowRequest,
3681 ) -> common_meta::error::Result<FlowResponse> {
3682 unreachable!()
3683 }
3684
3685 async fn handle_inserts(
3686 &self,
3687 _request: api::v1::region::InsertRequests,
3688 ) -> common_meta::error::Result<FlowResponse> {
3689 let _ = self.manager.flownode_dispatch.send(self.peer.clone());
3690 if self.manager.fail_flownode {
3691 return Err(common_meta::error::UnexpectedSnafu {
3692 err_msg: "test flownode failure".to_string(),
3693 }
3694 .build());
3695 }
3696 let gate = self
3697 .manager
3698 .gates
3699 .lock()
3700 .unwrap()
3701 .get(&self.peer.id)
3702 .cloned();
3703 if let Some(gate) = gate {
3704 let _permit = gate.acquire().await.unwrap();
3705 }
3706 Ok(FlowResponse::default())
3707 }
3708
3709 async fn handle_mark_window_dirty(
3710 &self,
3711 _request: api::v1::flow::DirtyWindowRequests,
3712 ) -> common_meta::error::Result<FlowResponse> {
3713 unreachable!()
3714 }
3715 }
3716
3717 fn mirror_requests(peer: &Peer, num_rows: usize) -> HashMap<Peer, RegionInsertRequests> {
3718 HashMap::from_iter([(
3719 peer.clone(),
3720 RegionInsertRequests {
3721 requests: vec![RegionInsertRequest {
3722 region_id: RegionId::new(1, 1).as_u64(),
3723 rows: Some(Rows {
3724 schema: vec![],
3725 rows: vec![api::v1::Row { values: vec![] }; num_rows],
3726 }),
3727 ..Default::default()
3728 }],
3729 },
3730 )])
3731 }
3732
3733 static MIRROR_METRIC_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
3738
3739 async fn lock_mirror_metrics() -> tokio::sync::MutexGuard<'static, ()> {
3740 MIRROR_METRIC_TEST_LOCK.lock().await
3741 }
3742
3743 #[tokio::test]
3747 async fn flow_mirror_saturation_preserves_instant_and_normal_contracts() {
3748 use common_error::ext::{ErrorExt, RetryHint};
3749 use common_error::status_code::StatusCode;
3750
3751 let _guard = lock_mirror_metrics().await;
3752 let kv_backend = prepare_mocked_backend().await;
3753 let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
3754 let (datanode_tx, mut datanode_rx) = tokio::sync::mpsc::unbounded_channel();
3755 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
3756 let node_manager = Arc::new(RecordingNodeManager {
3757 datanode_dispatch: datanode_tx,
3758 flownode_dispatch: flownode_tx,
3759 gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
3760 fail_flownode: false,
3761 });
3762 let flow_cache = Cache::new(10);
3763 let flow_cache_backend = prepare_mocked_backend().await;
3764 let inserter = Inserter::new(
3765 catalog::memory::MemoryCatalogManager::new(),
3766 partition_manager,
3767 node_manager,
3768 Arc::new(new_table_flownode_set_cache(
3769 String::new(),
3770 flow_cache.clone(),
3771 flow_cache_backend,
3772 )),
3773 true,
3774 );
3775 let peer = flownode_peer();
3776 inserter
3777 .table_flownode_set_cache
3778 .invalidate(&[common_meta::instruction::CacheIdent::CreateFlow(
3779 common_meta::instruction::CreateFlow {
3780 flow_id: 1,
3781 source_table_ids: vec![1, 2],
3782 partition_to_peer_mapping: vec![(0, peer)],
3783 },
3784 )])
3785 .await
3786 .unwrap();
3787 let mut normal_info = new_test_table_info(1, "normal_table", [1].into_iter());
3788 normal_info.catalog_name = DEFAULT_CATALOG_NAME.to_string();
3789 normal_info.schema_name = DEFAULT_SCHEMA_NAME.to_string();
3790 let mut instant_info = new_test_table_info(2, "instant_table", [1].into_iter());
3791 instant_info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
3792 let table_infos =
3793 HashMap::from_iter([(1, Arc::new(normal_info)), (2, Arc::new(instant_info))]);
3794 let ctx = Arc::new(QueryContext::with(
3795 DEFAULT_CATALOG_NAME,
3796 DEFAULT_SCHEMA_NAME,
3797 ));
3798 let insert = |table_id, num_rows| RegionInsertRequests {
3799 requests: vec![RegionInsertRequest {
3800 region_id: RegionId::new(table_id, 1).as_u64(),
3801 rows: Some(Rows {
3802 schema: vec![],
3803 rows: vec![api::v1::Row::default(); num_rows],
3804 }),
3805 ..Default::default()
3806 }],
3807 };
3808 let make_request = |normal, instant| InstantAndNormalInsertRequests {
3809 normal_requests: if normal == 0 {
3810 RegionInsertRequests::default()
3811 } else {
3812 insert(1, normal)
3813 },
3814 instant_requests: if instant == 0 {
3815 RegionInsertRequests::default()
3816 } else {
3817 insert(2, instant)
3818 },
3819 };
3820
3821 let pending = inserter.mirror_pending_rows.clone();
3822 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
3823 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
3824 pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
3825 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
3826 for request in [make_request(0, 1), make_request(1, 1)] {
3827 let error = inserter
3828 .do_request(request, &table_infos, &ctx)
3829 .await
3830 .unwrap_err();
3831 assert_eq!(error.status_code(), StatusCode::RateLimited);
3832 assert_eq!(error.retry_hint(), RetryHint::Retryable);
3833 assert!(datanode_rx.try_recv().is_err());
3834 assert!(flownode_rx.try_recv().is_err());
3835 assert_eq!(
3836 pending.load(Ordering::Relaxed),
3837 MAX_MIRROR_PENDING_ROWS as u64
3838 );
3839 }
3840 assert_eq!(
3841 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
3842 dropped_before
3843 );
3844
3845 for pending_before in [0, MAX_MIRROR_PENDING_ROWS as u64] {
3846 if pending_before == 0 {
3847 pending.store(0, Ordering::Relaxed);
3848 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
3849 }
3850 let gauge_before_rejection = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
3851 for (normal, instant) in [
3852 (0, MAX_MIRROR_PENDING_ROWS + 1),
3853 (MAX_MIRROR_PENDING_ROWS / 2, MAX_MIRROR_PENDING_ROWS / 2 + 1),
3854 ] {
3855 let error = inserter
3856 .do_request(make_request(normal, instant), &table_infos, &ctx)
3857 .await
3858 .unwrap_err();
3859 assert_eq!(error.status_code(), StatusCode::InvalidArguments);
3860 assert_eq!(error.retry_hint(), RetryHint::NonRetryable);
3861 let message = error.to_string();
3862 assert!(message.contains("1000001"), "{message}");
3863 assert!(message.contains("1000000"), "{message}");
3864 assert!(message.contains("reduce the batch size"), "{message}");
3865 assert!(datanode_rx.try_recv().is_err());
3866 assert!(flownode_rx.try_recv().is_err());
3867 assert_eq!(pending.load(Ordering::Relaxed), pending_before);
3868 assert_eq!(
3869 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
3870 gauge_before_rejection
3871 );
3872 assert_eq!(
3873 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
3874 dropped_before
3875 );
3876 }
3877 if pending_before == 0 {
3878 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
3879 pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
3880 }
3881 }
3882
3883 let result = inserter
3884 .do_request(make_request(1, 0), &table_infos, &ctx)
3885 .await
3886 .unwrap();
3887 assert!(matches!(result.data, OutputData::AffectedRows(1)));
3888 assert_eq!(1, datanode_rx.try_recv().unwrap());
3889 assert!(flownode_rx.try_recv().is_err());
3890 assert_eq!(
3891 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
3892 dropped_before + 1
3893 );
3894 pending.store(0, Ordering::Relaxed);
3895 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
3896 assert_eq!(
3897 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
3898 gauge_before
3899 );
3900 }
3901
3902 #[tokio::test]
3903 async fn flow_mirror_new_counts_regions_and_peer_clones() {
3904 use common_meta::instruction::{CacheIdent, CreateFlow};
3905
3906 let peer1 = flownode_peer();
3907 let peer2 = Peer {
3908 id: 2,
3909 addr: "127.0.0.1:4002".to_string(),
3910 };
3911 let cache_backend = prepare_mocked_backend().await;
3912 let cache = Arc::new(new_table_flownode_set_cache(
3913 String::new(),
3914 Cache::new(10),
3915 cache_backend,
3916 ));
3917 cache
3918 .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
3919 flow_id: 1,
3920 source_table_ids: vec![1],
3921 partition_to_peer_mapping: vec![(0, peer1), (1, peer2)],
3922 })])
3923 .await
3924 .unwrap();
3925 let request = |region_id| RegionInsertRequest {
3926 region_id,
3927 rows: Some(Rows {
3928 schema: vec![],
3929 rows: vec![api::v1::Row { values: vec![] }; 3],
3930 }),
3931 ..Default::default()
3932 };
3933 let requests = [
3934 request(RegionId::new(1, 0).as_u64()),
3935 request(RegionId::new(1, 1).as_u64()),
3936 ];
3937 let task = FlowMirrorTask::new(&cache, requests.iter()).await.unwrap();
3938 assert_eq!(12, task.pending_rows());
3939 assert_eq!(2, task.requests.len());
3940 }
3941
3942 #[tokio::test]
3943 async fn flow_mirror_dropped_when_pending_exceeds_limit() {
3944 let _guard = lock_mirror_metrics().await;
3945 let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64));
3946 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
3947 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
3948 let num_rows = 100;
3949
3950 let task = FlowMirrorTask {
3951 requests: mirror_requests(&flownode_peer(), num_rows),
3952 };
3953 task.detach(
3955 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3956 pending.clone(),
3957 false,
3958 )
3959 .unwrap();
3960
3961 assert_eq!(
3962 num_rows as u64,
3963 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get() - dropped_before
3964 );
3965 assert_eq!(
3967 MAX_MIRROR_PENDING_ROWS as u64,
3968 pending.load(Ordering::Relaxed)
3969 );
3970 assert_eq!(
3973 gauge_before,
3974 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
3975 "a dropped batch must not change the pending gauge"
3976 );
3977
3978 let task = FlowMirrorTask {
3979 requests: mirror_requests(&flownode_peer(), num_rows),
3980 };
3981 let error = task
3982 .detach(
3983 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3984 pending.clone(),
3985 true,
3986 )
3987 .err()
3988 .unwrap();
3989 use common_error::ext::{ErrorExt, RetryHint};
3990 use common_error::status_code::StatusCode;
3991 assert_eq!(error.status_code(), StatusCode::RateLimited);
3992 assert_eq!(error.retry_hint(), RetryHint::Retryable);
3993 assert_eq!(
3994 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
3995 dropped_before + num_rows as u64,
3996 "rejected instant rows are not successful best-effort drops"
3997 );
3998 assert_eq!(
3999 MAX_MIRROR_PENDING_ROWS as u64,
4000 pending.load(Ordering::Relaxed)
4001 );
4002 assert_eq!(
4003 gauge_before,
4004 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get()
4005 );
4006 }
4007
4008 #[tokio::test]
4012 async fn flow_mirror_reserves_and_releases_per_peer_shares() {
4013 let _guard = lock_mirror_metrics().await;
4014 let per_peer_rows = 10;
4015 let total_rows = (2 * per_peer_rows) as u64;
4016 let start = MAX_MIRROR_PENDING_ROWS as u64 - total_rows;
4018 let pending = Arc::new(AtomicU64::new(start));
4019 let gate = Arc::new(tokio::sync::Semaphore::new(0));
4020 let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4021 gate: gate.clone(),
4022 affected_rows: per_peer_rows as u64,
4023 }));
4024
4025 let mut requests = mirror_requests(&flownode_peer(), per_peer_rows);
4026 requests.extend(mirror_requests(
4027 &Peer {
4028 id: 2,
4029 addr: "127.0.0.1:4002".to_string(),
4030 },
4031 per_peer_rows,
4032 ));
4033
4034 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4035 FlowMirrorTask { requests }
4036 .detach(node_manager, pending.clone(), false)
4037 .unwrap();
4038
4039 assert_eq!(
4042 MAX_MIRROR_PENDING_ROWS as u64,
4043 pending.load(Ordering::Relaxed)
4044 );
4045
4046 gate.add_permits(2);
4052 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4053 while pending.load(Ordering::Relaxed) != start
4054 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4055 {
4056 assert!(
4057 std::time::Instant::now() < deadline,
4058 "the mirror tasks must release exactly their own shares"
4059 );
4060 std::thread::sleep(std::time::Duration::from_millis(1));
4061 }
4062 }
4063
4064 #[tokio::test]
4065 async fn flow_mirror_fast_peer_releases_only_its_share_before_drop() {
4066 let _guard = lock_mirror_metrics().await;
4067 let peer1 = flownode_peer();
4068 let peer2 = Peer {
4069 id: 2,
4070 addr: "127.0.0.1:4002".to_string(),
4071 };
4072 let gate1 = Arc::new(tokio::sync::Semaphore::new(0));
4073 let gate2 = Arc::new(tokio::sync::Semaphore::new(0));
4074 let gates = Arc::new(std::sync::Mutex::new(HashMap::from_iter([
4075 (peer1.id, gate1.clone()),
4076 (peer2.id, gate2.clone()),
4077 ])));
4078 let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4079 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4080 let node_manager = Arc::new(RecordingNodeManager {
4081 datanode_dispatch: datanode_tx,
4082 flownode_dispatch: flownode_tx,
4083 gates,
4084 fail_flownode: false,
4085 });
4086 let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64 - 20));
4087 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4088 let mut requests = mirror_requests(&peer1, 10);
4089 requests.extend(mirror_requests(&peer2, 10));
4090 let task = FlowMirrorTask { requests };
4091 task.detach(node_manager.clone(), pending.clone(), false)
4092 .unwrap();
4093 let first = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4094 .await
4095 .unwrap()
4096 .unwrap();
4097 let second = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4098 .await
4099 .unwrap()
4100 .unwrap();
4101 assert_ne!(first.id, second.id);
4102 let fast_gate = if first.id == peer1.id {
4103 gate1.clone()
4104 } else {
4105 gate2.clone()
4106 };
4107 let slow_gate = if first.id == peer1.id {
4108 gate2.clone()
4109 } else {
4110 gate1.clone()
4111 };
4112 fast_gate.add_permits(1);
4113 let deadline = std::time::Instant::now() + Duration::from_secs(10);
4114 while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 10 {
4115 assert!(
4116 std::time::Instant::now() < deadline,
4117 "fast peer did not release"
4118 );
4119 std::thread::sleep(Duration::from_millis(1));
4120 }
4121 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4122 FlowMirrorTask {
4123 requests: mirror_requests(&peer1, 11),
4124 }
4125 .detach(node_manager, pending.clone(), false)
4126 .unwrap();
4127 assert_eq!(
4128 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4129 dropped_before + 11
4130 );
4131 assert!(flownode_rx.try_recv().is_err());
4132 slow_gate.add_permits(1);
4133 while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 20
4134 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4135 {
4136 assert!(
4137 std::time::Instant::now() < deadline,
4138 "slow peer did not release"
4139 );
4140 std::thread::sleep(Duration::from_millis(1));
4141 }
4142 }
4143
4144 #[tokio::test]
4145 async fn flow_mirror_failure_releases_budget_and_gauge() {
4146 let _guard = lock_mirror_metrics().await;
4147 let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4148 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4149 let node_manager = Arc::new(RecordingNodeManager {
4150 datanode_dispatch: datanode_tx,
4151 flownode_dispatch: flownode_tx,
4152 gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
4153 fail_flownode: true,
4154 });
4155 let pending = Arc::new(AtomicU64::new(0));
4156 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4157 FlowMirrorTask {
4158 requests: mirror_requests(&flownode_peer(), 5),
4159 }
4160 .detach(node_manager, pending.clone(), false)
4161 .unwrap();
4162 assert_eq!(
4163 tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4164 .await
4165 .unwrap()
4166 .unwrap()
4167 .id,
4168 flownode_peer().id
4169 );
4170 let deadline = std::time::Instant::now() + Duration::from_secs(10);
4171 while pending.load(Ordering::Relaxed) != 0
4172 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4173 {
4174 assert!(
4175 std::time::Instant::now() < deadline,
4176 "failed task leaked reservation"
4177 );
4178 std::thread::sleep(Duration::from_millis(1));
4179 }
4180 }
4181
4182 #[tokio::test]
4183 async fn flow_mirror_pending_counter_consistent() {
4184 let _guard = lock_mirror_metrics().await;
4185 let pending = Arc::new(AtomicU64::new(0));
4186 let gate = Arc::new(tokio::sync::Semaphore::new(0));
4187 let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4188 gate: gate.clone(),
4189 affected_rows: 7,
4190 }));
4191
4192 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4194 FlowMirrorTask {
4195 requests: HashMap::new(),
4196 }
4197 .detach(node_manager.clone(), pending.clone(), false)
4198 .unwrap();
4199 assert_eq!(0, pending.load(Ordering::Relaxed));
4200 assert_eq!(
4201 gauge_before,
4202 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4203 "an empty batch must not touch the pending gauge"
4204 );
4205
4206 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4207 let num_rows = 7;
4208
4209 FlowMirrorTask {
4210 requests: mirror_requests(&flownode_peer(), num_rows),
4211 }
4212 .detach(node_manager, pending.clone(), false)
4213 .unwrap();
4214
4215 assert_eq!(num_rows as u64, pending.load(Ordering::Relaxed));
4218 assert_eq!(
4219 num_rows as i64,
4220 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() - gauge_before,
4221 "the pending gauge must track the pending rows"
4222 );
4223 assert_eq!(
4224 dropped_before,
4225 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4226 "a batch within the limit must not be counted as dropped"
4227 );
4228
4229 gate.add_permits(1);
4232 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4233 while pending.load(Ordering::Relaxed) != 0
4234 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4235 {
4236 assert!(
4237 std::time::Instant::now() < deadline,
4238 "the mirror task did not release its pending reservation"
4239 );
4240 std::thread::sleep(std::time::Duration::from_millis(1));
4241 }
4242 }
4243
4244 #[test]
4245 fn mirror_logs_are_rate_limited() {
4246 let log = Arc::new(MirrorLog::new());
4247 assert_eq!(Some(1), log.claim_report(1_000));
4248 assert_eq!(None, log.claim_report(1_500));
4249 assert_eq!(None, log.claim_report(10_999));
4250 assert_eq!(Some(3), log.claim_report(11_000));
4251 assert_eq!(None, log.claim_report(11_001));
4252 assert_eq!(Some(2), log.claim_report(21_000));
4253
4254 let concurrent = Arc::new(MirrorLog::new());
4255 assert_eq!(Some(1), concurrent.claim_report(50_000));
4256 let barrier = Arc::new(std::sync::Barrier::new(8));
4257 let threads = (0..8)
4258 .map(|_| {
4259 let log = concurrent.clone();
4260 let barrier = barrier.clone();
4261 std::thread::spawn(move || {
4262 barrier.wait();
4263 log.claim_report(60_000)
4264 })
4265 })
4266 .collect::<Vec<_>>();
4267 let winners = threads
4268 .into_iter()
4269 .filter_map(|thread| thread.join().unwrap())
4270 .collect::<Vec<_>>();
4271 assert_eq!(winners.len(), 1);
4272 let residual = concurrent.events.swap(0, Ordering::Relaxed);
4273 assert_eq!(winners[0] + residual, 8);
4274 }
4275}