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