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 affected_rows = common_runtime::spawn_global(async move {
865 dispatcher.execute(writes, None, Some(&cancellation)).await
866 })
867 .await
868 .context(JoinTaskSnafu)?
869 .context(RequestInsertsSnafu)?
870 .affected_rows;
871 crate::metrics::DIST_INGEST_ROW_COUNT
872 .with_label_values(&[ctx.get_db_string().as_str()])
873 .inc_by(affected_rows as u64);
874 Ok(Output::new(
875 OutputData::AffectedRows(affected_rows),
876 OutputMeta::new_with_cost(write_cost as _),
877 ))
878 }
879
880 fn auto_create_disabled_reason(&self, ctx: &QueryContextRef) -> Result<Option<&'static str>> {
884 let auto_create_table_hint = ctx
885 .extension(AUTO_CREATE_TABLE_KEY)
886 .map(|v| v.parse::<bool>())
887 .transpose()
888 .map_err(|_| {
889 InvalidInsertRequestSnafu {
890 reason: "`auto_create_table` hint must be a boolean",
891 }
892 .build()
893 })?
894 .unwrap_or(true);
895 Ok(if !self.auto_create_table {
896 Some("auto-create table is disabled by frontend config")
897 } else if !auto_create_table_hint {
898 Some("`auto_create_table` hint is disabled")
899 } else {
900 None
901 })
902 }
903
904 fn is_auto_create_exempt_private_table(schema: &str, table: &str) -> bool {
907 schema == DEFAULT_PRIVATE_SCHEMA_NAME
908 && matches!(
909 table,
910 DEFAULT_EVENTS_TABLE_NAME | SLOW_QUERY_TABLE_NAME | REGION_STATS_HISTORY_TABLE_NAME
911 )
912 }
913
914 pub async fn ensure_bulk_insert_schema(
918 &self,
919 table: TableRef,
920 request_schema: &ArrowSchema,
921 ctx: &QueryContextRef,
922 statement_executor: &StatementExecutor,
923 ) -> Result<TableRef> {
924 let table_info = table.table_info();
925 if self.auto_create_disabled_reason(ctx)?.is_some()
926 && !Self::is_auto_create_exempt_private_table(&table_info.schema_name, &table_info.name)
927 {
928 return Ok(table);
929 }
930
931 let table_schema = table.schema();
932 let schema = request_schema
933 .fields()
934 .iter()
935 .filter(|field| table_schema.column_schema_by_name(field.name()).is_none())
936 .map(|field| {
937 let data_type = field.data_type();
938 ensure!(
941 !data_type.is_nested() && !matches!(data_type, ArrowDataType::Dictionary(..)),
942 crate::error::NotSupportedSnafu {
943 feat: format!(
944 "automatically adding bulk insert column '{}' with type {:?}",
945 field.name(),
946 data_type
947 ),
948 }
949 );
950 let column = datatypes::schema::ColumnSchema::try_from(field.as_ref())
951 .context(crate::error::ConvertSchemaSnafu)?;
952 let column_def =
955 try_as_column_def(&column, false).context(crate::error::ColumnDataTypeSnafu)?;
956 Ok(ColumnSchema {
957 column_name: column_def.name,
958 datatype: column_def.data_type,
959 semantic_type: column_def.semantic_type,
960 datatype_extension: column_def.datatype_extension,
961 options: column_def.options,
962 })
963 })
964 .collect::<Result<Vec<_>>>()?;
965 let mut request = RowInsertRequest {
966 table_name: table_info.name.clone(),
967 rows: Some(Rows {
968 schema,
969 rows: Vec::new(),
970 }),
971 };
972 let Some(alter_expr) =
973 self.get_alter_table_expr_on_demand(&mut request, &table, ctx, false, false, true)?
974 else {
975 return Ok(table);
976 };
977
978 statement_executor
979 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
980 .await?;
981 self.get_table(
982 &table_info.catalog_name,
983 &table_info.schema_name,
984 &table_info.name,
985 )
986 .await?
987 .with_context(|| TableNotFoundSnafu {
988 table_name: table_info.full_table_name(),
989 })
990 }
991
992 pub async fn ensure_trace_table_on_demand(
997 &self,
998 table_name: &str,
999 request_schema: Vec<ColumnSchema>,
1000 alter_existing: bool,
1001 ctx: &QueryContextRef,
1002 statement_executor: &StatementExecutor,
1003 ) -> Result<()> {
1004 let mut requests = RowInsertRequests {
1005 inserts: vec![RowInsertRequest {
1006 table_name: table_name.to_string(),
1007 rows: Some(api::v1::Rows {
1008 schema: request_schema,
1009 rows: Vec::new(),
1010 }),
1011 }],
1012 };
1013 self.create_or_alter_tables_on_demand(
1014 &mut requests,
1015 ctx,
1016 AutoCreateTableType::Trace { alter_existing },
1017 statement_executor,
1018 false,
1019 false,
1020 None,
1021 )
1022 .await?;
1023 Ok(())
1024 }
1025
1026 #[allow(clippy::too_many_arguments)]
1044 async fn create_or_alter_tables_on_demand(
1045 &self,
1046 requests: &mut RowInsertRequests,
1047 ctx: &QueryContextRef,
1048 auto_create_table_type: AutoCreateTableType,
1049 statement_executor: &StatementExecutor,
1050 accommodate_existing_schema: bool,
1051 is_single_value: bool,
1052 align_time_index_unit: Option<TimeUnit>,
1053 ) -> Result<CreateAlterTableResult> {
1054 let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND
1055 .with_label_values(&[auto_create_table_type.as_str()])
1056 .start_timer();
1057 let catalog = ctx.current_catalog();
1058 let schema = ctx.current_schema();
1059
1060 let auto_create_disabled_reason = self.auto_create_disabled_reason(ctx)?;
1061 let has_auto_create_exempt_table = auto_create_disabled_reason.is_some()
1063 && requests
1064 .inserts
1065 .iter()
1066 .any(|req| Self::is_auto_create_exempt_private_table(&schema, &req.table_name));
1067 let mut table_infos = HashMap::new();
1068 if let Some(disabled_reason) = auto_create_disabled_reason
1070 && !has_auto_create_exempt_table
1071 {
1072 let mut instant_table_ids = HashSet::new();
1073 for req in &mut requests.inserts {
1074 let table = match self.get_table(catalog, &schema, &req.table_name).await? {
1075 Some(table) => table,
1076 None if is_ddl_reserved_table(&schema, &req.table_name) => {
1079 statement_executor
1080 .create_declared_relationships_table(catalog, ctx.clone())
1081 .await?
1082 }
1083 None => {
1084 return InvalidInsertRequestSnafu {
1085 reason: format!(
1086 "Table `{}` does not exist, and {}",
1087 req.table_name, disabled_reason
1088 ),
1089 }
1090 .fail();
1091 }
1092 };
1093 if matches!(auto_create_table_type, AutoCreateTableType::Logical(_))
1097 && let Some(rows) = req.rows.as_mut()
1098 && let Some(target_unit) = table_time_index_unit(&table)
1099 {
1100 convert_rows_time_unit(rows, target_unit)?;
1101 }
1102 let table_info = table.table_info();
1103 if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1104 validate_trace_table_model(&table_info, ctx)?;
1105 }
1106 if table_info.is_ttl_instant_table() {
1107 instant_table_ids.insert(table_info.table_id());
1108 }
1109 table_infos.insert(table_info.table_id(), table.table_info());
1110 }
1111 let ret = CreateAlterTableResult {
1112 instant_table_ids,
1113 table_infos,
1114 };
1115 return Ok(ret);
1116 }
1117
1118 let mut create_tables = vec![];
1119 let mut alter_tables = vec![];
1120 let mut need_refresh_table_infos = HashSet::new();
1121 let mut instant_table_ids = HashSet::new();
1122 let mut per_table_semantics: Option<Option<PerTableSemanticIndex>> = None;
1123
1124 for req in &mut requests.inserts {
1125 let auto_create_allowed = auto_create_disabled_reason.is_none()
1127 || Self::is_auto_create_exempt_private_table(&schema, &req.table_name);
1128 match self.get_table(catalog, &schema, &req.table_name).await? {
1129 Some(table) => {
1130 let table_info = table.table_info();
1131 if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1132 validate_trace_table_model(&table_info, ctx)?;
1133 }
1134 if table_info.is_ttl_instant_table() {
1135 instant_table_ids.insert(table_info.table_id());
1136 }
1137 if matches!(auto_create_table_type, AutoCreateTableType::Logical(_))
1141 && let Some(rows) = req.rows.as_mut()
1142 && let Some(target_unit) = table_time_index_unit(&table)
1143 {
1144 convert_rows_time_unit(rows, target_unit)?;
1145 }
1146 if auto_create_allowed
1147 && let Some(alter_expr) = self.get_alter_table_expr_on_demand(
1148 req,
1149 &table,
1150 ctx,
1151 accommodate_existing_schema,
1152 is_single_value,
1153 auto_create_table_type.alter_existing(),
1154 )?
1155 {
1156 alter_tables.push(alter_expr);
1157 need_refresh_table_infos.insert((
1158 catalog.to_string(),
1159 schema.clone(),
1160 req.table_name.clone(),
1161 ));
1162 } else {
1163 table_infos.insert(table_info.table_id(), table.table_info());
1164 }
1165 }
1166 None if is_ddl_reserved_table(&schema, &req.table_name) => {
1170 let table = statement_executor
1171 .create_declared_relationships_table(catalog, ctx.clone())
1172 .await?;
1173 let table_info = table.table_info();
1174 if table_info.is_ttl_instant_table() {
1175 instant_table_ids.insert(table_info.table_id());
1176 }
1177 table_infos.insert(table_info.table_id(), table_info);
1178 }
1179 None if !auto_create_allowed
1180 && let Some(disabled_reason) = auto_create_disabled_reason =>
1181 {
1182 return InvalidInsertRequestSnafu {
1183 reason: format!(
1184 "Table `{}` does not exist, and {}",
1185 req.table_name, disabled_reason,
1186 ),
1187 }
1188 .fail();
1189 }
1190 None => {
1191 if let Some(physical_unit) = align_time_index_unit
1195 && let Some(rows) = req.rows.as_mut()
1196 {
1197 convert_rows_time_unit(rows, physical_unit)?;
1198 }
1199 let semantic_index = per_table_semantics
1200 .get_or_insert_with(|| parse_per_table_semantic_index(ctx))
1201 .as_ref();
1202 let create_expr = self.get_create_table_expr_on_demand(
1203 req,
1204 &auto_create_table_type,
1205 ctx,
1206 semantic_index,
1207 )?;
1208 create_tables.push(create_expr);
1209 }
1210 }
1211 }
1212
1213 match auto_create_table_type {
1214 AutoCreateTableType::Logical(_) => {
1215 if !create_tables.is_empty() {
1216 let tables = self
1218 .create_logical_tables(create_tables, ctx, statement_executor)
1219 .await?;
1220
1221 for table in tables {
1222 let table_info = table.table_info();
1223 if table_info.is_ttl_instant_table() {
1224 instant_table_ids.insert(table_info.table_id());
1225 }
1226 table_infos.insert(table_info.table_id(), table.table_info());
1227 }
1228 }
1229 if !alter_tables.is_empty() {
1230 statement_executor
1232 .alter_logical_tables(alter_tables, ctx.clone(), TriggerReason::AutoAlter)
1233 .await?;
1234 }
1235 }
1236 AutoCreateTableType::Physical
1237 | AutoCreateTableType::Log
1238 | AutoCreateTableType::LastNonNull => {
1239 for create_table in create_tables {
1242 let table = self
1243 .create_physical_table(create_table, None, ctx, statement_executor)
1244 .await?;
1245 let table_info = table.table_info();
1246 if table_info.is_ttl_instant_table() {
1247 instant_table_ids.insert(table_info.table_id());
1248 }
1249 table_infos.insert(table_info.table_id(), table.table_info());
1250 }
1251 for alter_expr in alter_tables.into_iter() {
1252 statement_executor
1253 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1254 .await?;
1255 }
1256 }
1257
1258 AutoCreateTableType::Trace { .. } => {
1259 let trace_table_name = ctx
1260 .extension(TRACE_TABLE_NAME_SESSION_KEY)
1261 .unwrap_or(TRACE_TABLE_NAME);
1262
1263 let trace_table_partitions = if let Some(trace_table_partitions) =
1264 ctx.extension(TRACE_TABLE_PARTITIONS_HINT_KEY)
1265 {
1266 let p = trace_table_partitions.parse::<u32>().map_err(|_| {
1267 InvalidInsertRequestSnafu {
1268 reason: format!(
1269 "Failed to parse trace_table_partitions: {}",
1270 trace_table_partitions
1271 ),
1272 }
1273 .build()
1274 })?;
1275 Some(p)
1276 } else {
1277 None
1278 };
1279
1280 for mut create_table in create_tables {
1283 if create_table.table_name == trace_services_table_name(trace_table_name)
1284 || create_table.table_name == trace_operations_table_name(trace_table_name)
1285 {
1286 create_table
1288 .table_options
1289 .insert(APPEND_MODE_KEY.to_string(), "false".to_string());
1290 create_table.table_options.remove(TTL_KEY);
1292
1293 let table = self
1294 .create_physical_table(create_table, None, ctx, statement_executor)
1295 .await?;
1296 let table_info = table.table_info();
1297 if table_info.is_ttl_instant_table() {
1298 instant_table_ids.insert(table_info.table_id());
1299 }
1300 table_infos.insert(table_info.table_id(), table.table_info());
1301 } else {
1302 let partitions = if matches!(trace_table_partitions, Some(0) | Some(1)) {
1305 None
1307 } else {
1308 let p = partition_rule_for_hexstring(
1309 TRACE_ID_COLUMN,
1310 trace_table_partitions,
1311 )
1312 .context(CreatePartitionRulesSnafu)?;
1313 Some(p)
1314 };
1315
1316 let index_columns =
1321 [TRACE_ID_COLUMN, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN];
1322 for index_column in index_columns {
1323 if let Some(col) = create_table
1324 .column_defs
1325 .iter_mut()
1326 .find(|c| c.name == index_column)
1327 {
1328 col.options =
1329 options_from_skipping(&SkippingIndexOptions::default())
1330 .context(ColumnOptionsSnafu)?;
1331 } else {
1332 warn!(
1333 "Column {} not found when creating index for trace table: {}.",
1334 index_column, create_table.table_name
1335 );
1336 }
1337 }
1338
1339 let id_compression = CompressionOptions {
1343 encoding: Some(ColumnEncoding::Plain),
1344 page_size: Some(ReadableSize::kb(64)),
1345 ..Default::default()
1346 };
1347 let compressed_columns = [
1348 (SPAN_ID_COLUMN, id_compression),
1349 (PARENT_SPAN_ID_COLUMN, id_compression),
1350 (
1351 TIMESTAMP_END_COLUMN,
1352 CompressionOptions {
1353 encoding: Some(ColumnEncoding::DeltaBinaryPacked),
1354 ..Default::default()
1355 },
1356 ),
1357 ];
1358 for (name, compression) in compressed_columns {
1359 if let Some(col) =
1360 create_table.column_defs.iter_mut().find(|c| c.name == name)
1361 {
1362 let compression = options_from_compression(&compression)
1363 .context(ColumnOptionsSnafu)?;
1364 col.options
1365 .get_or_insert_with(Default::default)
1366 .options
1367 .extend(compression.options);
1368 }
1369 }
1370
1371 create_table.table_options.insert(
1373 TABLE_DATA_MODEL.to_string(),
1374 ctx.extension(SEMANTIC_PIPELINE)
1375 .unwrap_or(TABLE_DATA_MODEL_TRACE_V1)
1376 .to_string(),
1377 );
1378
1379 let table = self
1380 .create_physical_table(
1381 create_table,
1382 partitions,
1383 ctx,
1384 statement_executor,
1385 )
1386 .await?;
1387 let table_info = table.table_info();
1388 if table_info.is_ttl_instant_table() {
1389 instant_table_ids.insert(table_info.table_id());
1390 }
1391 table_infos.insert(table_info.table_id(), table.table_info());
1392 }
1393 }
1394 for alter_expr in alter_tables.into_iter() {
1395 statement_executor
1396 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1397 .await?;
1398 }
1399 }
1400 }
1401
1402 for (catalog, schema, table_name) in need_refresh_table_infos {
1404 let table = self
1405 .get_table(&catalog, &schema, &table_name)
1406 .await?
1407 .context(TableNotFoundSnafu {
1408 table_name: common_catalog::format_full_table_name(
1409 &catalog,
1410 &schema,
1411 &table_name,
1412 ),
1413 })?;
1414 let table_info = table.table_info();
1415 table_infos.insert(table_info.table_id(), table.table_info());
1416 }
1417
1418 Ok(CreateAlterTableResult {
1419 instant_table_ids,
1420 table_infos,
1421 })
1422 }
1423
1424 async fn create_physical_table_on_demand(
1425 &self,
1426 ctx: &QueryContextRef,
1427 physical_table: String,
1428 statement_executor: &StatementExecutor,
1429 ) -> Result<TableRef> {
1430 let catalog_name = ctx.current_catalog();
1431 let schema_name = ctx.current_schema();
1432
1433 if let Some(table) = self
1435 .get_table(catalog_name, &schema_name, &physical_table)
1436 .await?
1437 {
1438 return Ok(table);
1439 }
1440
1441 if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? {
1443 return InvalidInsertRequestSnafu {
1444 reason: format!(
1445 "Physical table `{physical_table}` does not exist, and {disabled_reason}"
1446 ),
1447 }
1448 .fail();
1449 }
1450
1451 let table_reference = TableReference::full(catalog_name, &schema_name, &physical_table);
1452 debug!("Ensuring physical metric table `{table_reference}` exists for insert");
1453
1454 let default_schema = vec![
1456 ColumnSchema {
1457 column_name: greptime_timestamp().to_string(),
1458 datatype: ColumnDataType::TimestampMillisecond as _,
1459 semantic_type: SemanticType::Timestamp as _,
1460 datatype_extension: None,
1461 options: None,
1462 },
1463 ColumnSchema {
1464 column_name: greptime_value().to_string(),
1465 datatype: ColumnDataType::Float64 as _,
1466 semantic_type: SemanticType::Field as _,
1467 datatype_extension: None,
1468 options: None,
1469 },
1470 ];
1471 let create_table_expr =
1472 &mut build_create_table_expr(&table_reference, &default_schema, default_engine())?;
1473
1474 create_table_expr.engine = METRIC_ENGINE_NAME.to_string();
1475 create_table_expr
1476 .table_options
1477 .insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), "true".to_string());
1478
1479 let res = statement_executor
1481 .create_table_inner(
1482 create_table_expr,
1483 None,
1484 ctx.clone(),
1485 TriggerReason::AutoCreate,
1486 )
1487 .await;
1488
1489 match res {
1490 Ok(table) => Ok(table),
1491 Err(err) => {
1492 error!(err; "Failed to create table {table_reference}");
1493 Err(err)
1494 }
1495 }
1496 }
1497
1498 async fn get_table(
1499 &self,
1500 catalog: &str,
1501 schema: &str,
1502 table: &str,
1503 ) -> Result<Option<TableRef>> {
1504 self.catalog_manager
1505 .table(catalog, schema, table, None)
1506 .await
1507 .context(CatalogSnafu)
1508 }
1509
1510 fn get_create_table_expr_on_demand(
1511 &self,
1512 req: &RowInsertRequest,
1513 create_type: &AutoCreateTableType,
1514 ctx: &QueryContextRef,
1515 semantic_index: Option<&PerTableSemanticIndex>,
1516 ) -> Result<CreateTableExpr> {
1517 let schema = ctx.current_schema();
1518 let mut table_options = std::collections::HashMap::with_capacity(4);
1519 fill_table_options_for_create(&mut table_options, create_type, ctx);
1520 apply_per_table_semantic_options(
1521 &mut table_options,
1522 semantic_index,
1523 ctx.current_schema().as_str(),
1524 &req.table_name,
1525 );
1526
1527 let engine_name = if let AutoCreateTableType::Logical(_) = create_type {
1528 METRIC_ENGINE_NAME
1530 } else {
1531 default_engine()
1532 };
1533
1534 let table_ref = TableReference::full(ctx.current_catalog(), &schema, &req.table_name);
1535 let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1537 let mut create_table_expr =
1538 build_create_table_expr(&table_ref, request_schema, engine_name)?;
1539
1540 if ctx.extension(SPLUNK_PK_METADATA_ORDER_KEY).is_some() {
1542 reorder_splunk_primary_keys(&mut create_table_expr.primary_keys);
1543 }
1544
1545 debug!("Ensuring table `{table_ref}` exists for insert");
1546 create_table_expr.table_options.extend(table_options);
1547 Ok(create_table_expr)
1548 }
1549
1550 fn get_alter_table_expr_on_demand(
1559 &self,
1560 req: &mut RowInsertRequest,
1561 table: &TableRef,
1562 ctx: &QueryContextRef,
1563 accommodate_existing_schema: bool,
1564 is_single_value: bool,
1565 alter_existing: bool,
1566 ) -> Result<Option<AlterTableExpr>> {
1567 if !alter_existing {
1568 return Ok(None);
1569 }
1570
1571 let catalog_name = ctx.current_catalog();
1572 let schema_name = ctx.current_schema();
1573 let table_name = table.table_info().name.clone();
1574
1575 if is_ddl_reserved_table(&schema_name, &table_name) {
1578 return Ok(None);
1579 }
1580
1581 let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1582 let request_field_count = request_schema
1583 .iter()
1584 .filter(|col| col.semantic_type == SemanticType::Field as i32)
1585 .count();
1586 if is_single_value {
1587 let request_is_native_histogram = request_is_native_histogram(request_schema);
1588 let table_is_native_histogram = table_is_native_histogram(table);
1589 ensure!(
1590 request_is_native_histogram == table_is_native_histogram,
1591 InvalidInsertRequestSnafu {
1592 reason: format!(
1593 "Table `{table_name}` cannot mix native histogram and float sample fields"
1594 ),
1595 }
1596 );
1597 }
1598
1599 let column_exprs = ColumnExpr::from_column_schemas(request_schema);
1600 let add_columns = expr_helper::extract_add_columns_expr(&table.schema(), column_exprs)?;
1601 let Some(mut add_columns) = add_columns else {
1602 return Ok(None);
1603 };
1604
1605 if accommodate_existing_schema {
1607 let table_schema = table.schema();
1608 let ts_col_name = table_schema.timestamp_column().map(|c| c.name.clone());
1610 let mut field_col_name = None;
1612 if is_single_value && request_field_count <= 1 {
1613 let mut multiple_field_cols = false;
1614 table.field_columns().for_each(|col| {
1615 if field_col_name.is_none() {
1616 field_col_name = Some(col.name.clone());
1617 } else {
1618 multiple_field_cols = true;
1619 }
1620 });
1621 if multiple_field_cols {
1622 field_col_name = None;
1623 }
1624 }
1625
1626 if let Some(rows) = req.rows.as_mut() {
1628 for col in &mut rows.schema {
1629 match col.semantic_type {
1630 x if x == SemanticType::Timestamp as i32 => {
1631 if let Some(ref ts_name) = ts_col_name
1632 && col.column_name != *ts_name
1633 {
1634 col.column_name = ts_name.clone();
1635 }
1636 }
1637 x if x == SemanticType::Field as i32 => {
1638 if let Some(ref field_name) = field_col_name
1639 && col.column_name != *field_name
1640 {
1641 col.column_name = field_name.clone();
1642 }
1643 }
1644 _ => {}
1645 }
1646 }
1647 }
1648
1649 add_columns.add_columns.retain(|col| {
1651 let def = col.column_def.as_ref().unwrap();
1652 def.semantic_type == SemanticType::Tag as i32
1653 || (def.semantic_type == SemanticType::Field as i32 && field_col_name.is_none())
1654 });
1655
1656 if add_columns.add_columns.is_empty() {
1657 return Ok(None);
1658 }
1659 }
1660
1661 Ok(Some(AlterTableExpr {
1662 catalog_name: catalog_name.to_string(),
1663 schema_name: schema_name.clone(),
1664 table_name: table_name.clone(),
1665 kind: Some(Kind::AddColumns(add_columns)),
1666 }))
1667 }
1668
1669 async fn create_physical_table(
1671 &self,
1672 mut create_table_expr: CreateTableExpr,
1673 partitions: Option<Partitions>,
1674 ctx: &QueryContextRef,
1675 statement_executor: &StatementExecutor,
1676 ) -> Result<TableRef> {
1677 let res = statement_executor
1678 .create_table_inner(
1679 &mut create_table_expr,
1680 partitions,
1681 ctx.clone(),
1682 TriggerReason::AutoCreate,
1683 )
1684 .await;
1685
1686 let table_ref = TableReference::full(
1687 &create_table_expr.catalog_name,
1688 &create_table_expr.schema_name,
1689 &create_table_expr.table_name,
1690 );
1691
1692 match res {
1693 Ok(table) => {
1694 validate_trace_table_model(&table.table_info(), ctx)?;
1695 Ok(table)
1696 }
1697 Err(err) => {
1698 error!(err; "Failed to create table {}", table_ref);
1699 Err(err)
1700 }
1701 }
1702 }
1703
1704 async fn create_logical_tables(
1705 &self,
1706 create_table_exprs: Vec<CreateTableExpr>,
1707 ctx: &QueryContextRef,
1708 statement_executor: &StatementExecutor,
1709 ) -> Result<Vec<TableRef>> {
1710 let res = statement_executor
1711 .create_logical_tables(&create_table_exprs, ctx.clone(), TriggerReason::AutoCreate)
1712 .await;
1713
1714 match res {
1715 Ok(res) => Ok(res),
1716 Err(err) => {
1717 let failed_tables = create_table_exprs
1718 .into_iter()
1719 .map(|expr| {
1720 format!(
1721 "{}.{}.{}",
1722 expr.catalog_name, expr.schema_name, expr.table_name
1723 )
1724 })
1725 .collect::<Vec<_>>();
1726 error!(
1727 err;
1728 "Failed to create logical tables {:?}",
1729 failed_tables
1730 );
1731 Err(err)
1732 }
1733 }
1734 }
1735
1736 pub fn write_dispatcher(&self) -> &WriteDispatcherRef {
1738 &self.write_dispatcher
1739 }
1740
1741 pub fn node_manager(&self) -> &NodeManagerRef {
1742 &self.node_manager
1743 }
1744
1745 pub fn partition_manager(&self) -> &PartitionRuleManagerRef {
1746 &self.partition_manager
1747 }
1748
1749 pub fn table_flownode_set_cache(&self) -> &TableFlownodeSetCacheRef {
1750 &self.table_flownode_set_cache
1751 }
1752}
1753
1754fn request_is_native_histogram(request_schema: &[ColumnSchema]) -> bool {
1755 let mut fields = request_schema
1756 .iter()
1757 .filter(|col| col.semantic_type == SemanticType::Field as i32);
1758 let Some(col) = fields.next() else {
1759 return false;
1760 };
1761
1762 fields.next().is_none()
1763 && api::helper::is_column_type_value_eq(
1764 col.datatype,
1765 col.datatype_extension.clone(),
1766 native_histogram_value_type(),
1767 )
1768}
1769
1770fn table_time_index_unit(table: &TableRef) -> Option<TimeUnit> {
1774 table
1775 .table_info()
1776 .meta
1777 .schema
1778 .timestamp_column()
1779 .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit()))
1780}
1781
1782fn convert_rows_time_unit(rows: &mut Rows, target_unit: TimeUnit) -> Result<()> {
1783 let Some(ts_index) = rows
1784 .schema
1785 .iter()
1786 .position(|col| col.semantic_type == SemanticType::Timestamp as i32)
1787 else {
1788 return Ok(());
1789 };
1790 let Some(source_unit) = ColumnDataType::try_from(rows.schema[ts_index].datatype)
1791 .ok()
1792 .and_then(api::helper::timestamp_unit)
1793 else {
1794 return Ok(());
1795 };
1796 if source_unit == target_unit {
1797 return Ok(());
1798 }
1799
1800 rows.schema[ts_index].datatype = api::helper::timestamp_datatype(target_unit) as i32;
1801 rows.schema[ts_index].datatype_extension = None;
1803
1804 for row in &mut rows.rows {
1812 debug_assert_eq!(row.values.len(), rows.schema.len());
1813 let value = &mut row.values[ts_index];
1814 let Some(value_data) = value.value_data.take() else {
1815 continue;
1816 };
1817 value.value_data =
1818 convert_timestamp_value_data(value_data, source_unit, target_unit, ts_index)?;
1819 }
1820 Ok(())
1821}
1822
1823fn convert_timestamp_value_data(
1824 value_data: ValueData,
1825 source_unit: TimeUnit,
1826 target_unit: TimeUnit,
1827 column_index: usize,
1828) -> Result<Option<ValueData>> {
1829 let timestamp = match value_data {
1830 ValueData::TimestampSecondValue(v) => Timestamp::new_second(v),
1831 ValueData::TimestampMillisecondValue(v) => Timestamp::new_millisecond(v),
1832 ValueData::TimestampMicrosecondValue(v) => Timestamp::new_microsecond(v),
1833 ValueData::TimestampNanosecondValue(v) => Timestamp::new_nanosecond(v),
1834 other => return Ok(Some(other)),
1836 };
1837 let converted = timestamp
1838 .convert_to(target_unit)
1839 .with_context(|| InvalidInsertRequestSnafu {
1840 reason: format!(
1841 "timestamp column {column_index} value {} in unit {source_unit:?} overflows when converting to unit {target_unit:?}",
1842 timestamp.value()
1843 ),
1844 })?;
1845 Ok(api::helper::to_grpc_value(datatypes::value::Value::Timestamp(converted)).value_data)
1846}
1847
1848fn table_is_native_histogram(table: &TableRef) -> bool {
1849 let mut fields = table.field_columns();
1850 let Some(col) = fields.next() else {
1851 return false;
1852 };
1853
1854 fields.next().is_none() && is_native_histogram_value_type(&col.data_type)
1855}
1856
1857fn validate_column_count_match(requests: &RowInsertRequests) -> Result<()> {
1858 for request in &requests.inserts {
1859 let rows = request.rows.as_ref().unwrap();
1860 let column_count = rows.schema.len();
1861 rows.rows.iter().try_for_each(|r| {
1862 ensure!(
1863 r.values.len() == column_count,
1864 InvalidInsertRequestSnafu {
1865 reason: format!(
1866 "column count mismatch, columns: {}, values: {}",
1867 column_count,
1868 r.values.len()
1869 )
1870 }
1871 );
1872 Ok(())
1873 })?;
1874 }
1875 Ok(())
1876}
1877
1878pub fn validate_trace_table_model(table_info: &TableInfo, ctx: &QueryContextRef) -> Result<()> {
1881 let Some(expected @ (TABLE_DATA_MODEL_TRACE_V1 | TABLE_DATA_MODEL_TRACE_V2)) =
1882 ctx.extension(SEMANTIC_PIPELINE)
1883 else {
1884 return Ok(());
1885 };
1886 if let Some(actual) = table_info.meta.options.data_model() {
1887 ensure!(
1888 actual == expected,
1889 InvalidInsertRequestSnafu {
1890 reason: format!(
1891 "Trace table `{}` uses {actual}, but the request uses {expected}",
1892 table_info.name,
1893 ),
1894 }
1895 );
1896 }
1897 Ok(())
1898}
1899
1900pub fn fill_table_options_for_create(
1902 table_options: &mut std::collections::HashMap<String, String>,
1903 create_type: &AutoCreateTableType,
1904 ctx: &QueryContextRef,
1905) {
1906 for key in VALID_TABLE_OPTION_KEYS {
1907 if let Some(value) = ctx.extension(key) {
1908 table_options.insert(key.to_string(), value.to_string());
1909 }
1910 }
1911
1912 if !matches!(create_type, AutoCreateTableType::Logical(_)) {
1913 let mut has_twcs_option = false;
1914 for key in COMPACTION_OPTION_KEYS {
1915 if let Some(value) = ctx.extension(key) {
1916 has_twcs_option |= key.starts_with(COMPACTION_TWCS_PREFIX);
1917 table_options.insert(key.to_string(), value.to_string());
1918 }
1919 }
1920 if has_twcs_option {
1921 table_options
1922 .entry(COMPACTION_TYPE.to_string())
1923 .or_insert_with(|| COMPACTION_TYPE_TWCS.to_string());
1924 }
1925 }
1926
1927 for (key, value) in ctx.extensions() {
1929 if is_semantic_option_key(&key) && validate_semantic_option(&key, &value) {
1930 table_options.insert(key, value);
1931 }
1932 }
1933
1934 match create_type {
1935 AutoCreateTableType::Logical(physical_table) => {
1936 table_options.insert(
1937 LOGICAL_TABLE_METADATA_KEY.to_string(),
1938 physical_table.clone(),
1939 );
1940 }
1941 AutoCreateTableType::Physical => {
1942 if let Some(append_mode) = ctx.extension(APPEND_MODE_KEY) {
1943 table_options.insert(APPEND_MODE_KEY.to_string(), append_mode.to_string());
1944 }
1945 if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
1946 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
1947 }
1948 }
1949 AutoCreateTableType::Log => {
1952 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1953 }
1954 AutoCreateTableType::LastNonNull => {
1955 if ctx
1956 .extension(APPEND_MODE_KEY)
1957 .is_some_and(|value| value.eq_ignore_ascii_case("true"))
1958 {
1959 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1960 table_options.insert(MERGE_MODE_KEY.to_string(), "last_row".to_string());
1961 } else if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
1962 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
1963 } else {
1964 table_options.insert(MERGE_MODE_KEY.to_string(), "last_non_null".to_string());
1965 }
1966 }
1967 AutoCreateTableType::Trace { .. } => {
1968 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1969 }
1970 }
1971}
1972
1973pub type PerTableSemanticIndex = BTreeMap<String, BTreeMap<String, BTreeMap<String, String>>>;
1979
1980pub fn parse_per_table_semantic_index(ctx: &QueryContextRef) -> Option<PerTableSemanticIndex> {
1985 let raw = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY)?;
1986 match serde_json::from_str(raw) {
1987 Ok(index) => Some(index),
1988 Err(_) => {
1989 warn!("failed to parse semantic per-table index, skipping per-table options");
1990 None
1991 }
1992 }
1993}
1994
1995pub fn apply_per_table_semantic_options(
2003 table_options: &mut std::collections::HashMap<String, String>,
2004 index: Option<&PerTableSemanticIndex>,
2005 schema: &str,
2006 table_name: &str,
2007) {
2008 let Some(entry) = index
2009 .and_then(|index| index.get(schema))
2010 .and_then(|tables| tables.get(table_name))
2011 else {
2012 return;
2013 };
2014 for (key, value) in entry {
2015 if is_semantic_option_key(key) && validate_semantic_option(key, value) {
2016 table_options.insert(key.clone(), value.clone());
2017 }
2018 }
2019}
2020
2021pub fn build_create_table_expr(
2022 table: &TableReference,
2023 request_schema: &[ColumnSchema],
2024 engine: &str,
2025) -> Result<CreateTableExpr> {
2026 expr_helper::create_table_expr_by_column_schemas(table, request_schema, engine, None)
2027}
2028
2029pub const SPLUNK_PK_METADATA_ORDER_KEY: &str = "splunk_pk_metadata_order";
2033
2034fn reorder_splunk_primary_keys(primary_keys: &mut [String]) {
2037 const LEAD: [&str; 3] = ["host", "source", "sourcetype"];
2038 primary_keys.sort_by_key(|name| {
2041 LEAD.iter()
2042 .position(|&lead| lead == name.as_str())
2043 .unwrap_or(LEAD.len())
2044 });
2045}
2046
2047struct CreateAlterTableResult {
2049 instant_table_ids: HashSet<TableId>,
2051 table_infos: HashMap<TableId, Arc<TableInfo>>,
2053}
2054
2055const MAX_MIRROR_PENDING_ROWS: usize = 1_000_000;
2063
2064const MIRROR_LOG_INTERVAL: Duration = Duration::from_secs(10);
2066const MIRROR_LOG_NEVER_REPORTED: u64 = u64::MAX;
2067
2068struct MirrorLog {
2069 last_log_millis: AtomicU64,
2070 events: AtomicU64,
2071}
2072
2073impl MirrorLog {
2074 const fn new() -> Self {
2075 Self {
2076 last_log_millis: AtomicU64::new(MIRROR_LOG_NEVER_REPORTED),
2077 events: AtomicU64::new(0),
2078 }
2079 }
2080
2081 fn claim_report(&self, now_millis: u64) -> Option<u64> {
2084 self.events.fetch_add(1, Ordering::Relaxed);
2085 let last = self.last_log_millis.load(Ordering::Relaxed);
2086 if last != MIRROR_LOG_NEVER_REPORTED
2087 && now_millis.saturating_sub(last) < MIRROR_LOG_INTERVAL.as_millis() as u64
2088 {
2089 return None;
2090 }
2091 if self
2092 .last_log_millis
2093 .compare_exchange(last, now_millis, Ordering::Relaxed, Ordering::Relaxed)
2094 .is_err()
2095 {
2096 return None;
2097 }
2098 Some(self.events.swap(0, Ordering::Relaxed))
2099 }
2100}
2101
2102static MIRROR_DROP_LOG: MirrorLog = MirrorLog::new();
2103static MIRROR_FAILURE_LOG: MirrorLog = MirrorLog::new();
2104
2105fn mirror_drop_log_millis() -> u64 {
2108 static START: LazyLock<Instant> = LazyLock::new(Instant::now);
2109 START.elapsed().as_millis() as u64
2110}
2111
2112struct FlowMirrorTask {
2113 requests: HashMap<Peer, RegionInsertRequests>,
2114}
2115
2116impl FlowMirrorTask {
2117 async fn new(
2118 cache: &TableFlownodeSetCacheRef,
2119 requests: impl Iterator<Item = &RegionInsertRequest>,
2120 ) -> Result<Self> {
2121 let mut src_table_reqs: HashMap<TableId, Option<(Vec<Peer>, RegionInsertRequests)>> =
2122 HashMap::new();
2123
2124 for req in requests {
2125 let table_id = RegionId::from_u64(req.region_id).table_id();
2126 match src_table_reqs.get_mut(&table_id) {
2127 Some(Some((_peers, reqs))) => reqs.requests.push(req.clone()),
2128 Some(None) => continue,
2130 _ => {
2131 let peers = cache
2133 .get(table_id)
2134 .await
2135 .context(RequestInsertsSnafu)?
2136 .unwrap_or_default()
2137 .values()
2138 .cloned()
2139 .collect::<HashSet<_>>()
2140 .into_iter()
2141 .collect::<Vec<_>>();
2142
2143 if !peers.is_empty() {
2144 let mut reqs = RegionInsertRequests::default();
2145 reqs.requests.push(req.clone());
2146 src_table_reqs.insert(table_id, Some((peers, reqs)));
2147 } else {
2148 src_table_reqs.insert(table_id, None);
2150 }
2151 }
2152 }
2153 }
2154
2155 let mut inserts: HashMap<Peer, RegionInsertRequests> = HashMap::new();
2156
2157 for (_table_id, (peers, reqs)) in src_table_reqs
2158 .into_iter()
2159 .filter_map(|(k, v)| v.map(|v| (k, v)))
2160 {
2161 if peers.len() == 1 {
2162 inserts
2164 .entry(peers[0].clone())
2165 .or_default()
2166 .requests
2167 .extend(reqs.requests);
2168 continue;
2169 } else {
2170 for flownode in peers {
2172 inserts
2173 .entry(flownode.clone())
2174 .or_default()
2175 .requests
2176 .extend(reqs.requests.clone());
2177 }
2178 }
2179 }
2180
2181 Ok(Self { requests: inserts })
2182 }
2183
2184 fn pending_rows(&self) -> u64 {
2193 self.requests.values().map(region_inserts_rows).sum()
2194 }
2195
2196 fn detach(
2197 self,
2198 node_manager: NodeManagerRef,
2199 mirror_pending_rows: Arc<AtomicU64>,
2200 has_instant_rows: bool,
2201 ) -> Result<()> {
2202 let num_rows = self.pending_rows();
2205 if num_rows == 0 {
2206 return Ok(());
2207 }
2208 if has_instant_rows && num_rows > MAX_MIRROR_PENDING_ROWS as u64 {
2209 return InvalidInsertRequestSnafu {
2210 reason: format!(
2211 "flow mirror batch has {num_rows} peer-fan-out rows, exceeding the limit of {}; reduce the batch size",
2212 MAX_MIRROR_PENDING_ROWS
2213 ),
2214 }
2215 .fail();
2216 }
2217
2218 let pending = reserve_mirror_pending_rows(&mirror_pending_rows, num_rows);
2219 if pending > MAX_MIRROR_PENDING_ROWS as u64 {
2220 release_mirror_pending_rows(&mirror_pending_rows, num_rows);
2221 if has_instant_rows {
2222 return Err(meter_core::collect::WriteRejected::new(
2223 "flow mirror pending rows limit exceeded for instant-TTL insert",
2224 ))
2225 .context(WriteRejectedSnafu);
2226 }
2227 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.inc_by(num_rows);
2228 if let Some(events) = MIRROR_DROP_LOG.claim_report(mirror_drop_log_millis()) {
2229 warn!(
2230 "Flow mirror write dropped: {} pending rows exceeds limit {}, {} drop events since last report",
2231 pending, MAX_MIRROR_PENDING_ROWS, events
2232 );
2233 }
2234 return Ok(());
2235 }
2236
2237 for (peer, inserts) in self.requests {
2238 let peer_rows = region_inserts_rows(&inserts);
2242 let node_manager = node_manager.clone();
2243 let mirror_pending_rows = mirror_pending_rows.clone();
2244 common_runtime::spawn_global(async move {
2245 let result = node_manager
2246 .flownode(&peer)
2247 .await
2248 .handle_inserts(inserts)
2249 .await
2250 .context(RequestInsertsSnafu);
2251
2252 match result {
2253 Ok(resp) => {
2254 let affected_rows = resp.affected_rows;
2255 crate::metrics::DIST_MIRROR_ROW_COUNT.inc_by(affected_rows);
2256 }
2257 Err(err) => {
2258 if let Some(events) =
2259 MIRROR_FAILURE_LOG.claim_report(mirror_drop_log_millis())
2260 {
2261 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);
2262 }
2263 }
2264 }
2265 release_mirror_pending_rows(&mirror_pending_rows, peer_rows);
2268 });
2269 }
2270
2271 Ok(())
2272 }
2273}
2274
2275fn region_inserts_rows(inserts: &RegionInsertRequests) -> u64 {
2278 inserts
2279 .requests
2280 .iter()
2281 .filter_map(|req| req.rows.as_ref())
2282 .map(|rows| rows.rows.len() as u64)
2283 .sum()
2284}
2285
2286fn reserve_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) -> u64 {
2291 let pending = mirror_pending_rows.fetch_add(rows, Ordering::Relaxed) + rows;
2292 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(rows as i64);
2295 pending
2296}
2297
2298fn release_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) {
2304 let prev = mirror_pending_rows
2305 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |pending| {
2306 Some(pending.saturating_sub(rows))
2307 })
2308 .unwrap_or_else(|prev| prev);
2309 let released = rows.min(prev);
2313 if released > 0 {
2314 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(released as i64);
2315 }
2316}
2317
2318#[cfg(test)]
2319mod tests {
2320 use std::sync::Arc;
2321 use std::sync::atomic::{AtomicU64, Ordering};
2322
2323 use api::helper::ColumnDataTypeWrapper;
2324 use api::v1::flow::FlowResponse;
2325 use api::v1::helper::{field_column_schema, time_index_column_schema};
2326 use api::v1::{Row, RowInsertRequest, Rows, Value};
2327 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
2328 use common_meta::cache::new_table_flownode_set_cache;
2329 use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler;
2330 use common_meta::test_util::{MockDatanodeManager, MockFlownodeHandler, MockFlownodeManager};
2331 use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD;
2332 use common_query::prelude::{greptime_native_histogram, set_default_prefix};
2333 use datatypes::data_type::ConcreteDataType;
2334 use datatypes::schema::ColumnSchema;
2335 use moka::future::Cache;
2336 use session::context::QueryContext;
2337 use table::TableRef;
2338 use table::dist_table::DummyDataSource;
2339 use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType};
2340
2341 use crate::insert::*;
2342 use crate::test_util::{
2343 create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
2344 };
2345
2346 fn make_table_ref_with_schema(
2347 ts_name: &str,
2348 field_name: &str,
2349 field_type: ConcreteDataType,
2350 ) -> TableRef {
2351 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2352 ColumnSchema::new(
2353 ts_name,
2354 ConcreteDataType::timestamp_millisecond_datatype(),
2355 false,
2356 )
2357 .with_time_index(true),
2358 ColumnSchema::new(field_name, field_type, true),
2359 ])
2360 .unwrap()
2361 .build()
2362 .unwrap();
2363 let meta = TableMetaBuilder::empty()
2364 .schema(Arc::new(schema))
2365 .primary_key_indices(vec![])
2366 .value_indices(vec![1])
2367 .engine("mito")
2368 .next_column_id(0)
2369 .options(Default::default())
2370 .created_on(Default::default())
2371 .build()
2372 .unwrap();
2373 let info = Arc::new(
2374 TableInfoBuilder::default()
2375 .table_id(1)
2376 .table_version(0)
2377 .name("test_table")
2378 .schema_name(DEFAULT_SCHEMA_NAME)
2379 .catalog_name(DEFAULT_CATALOG_NAME)
2380 .desc(None)
2381 .table_type(TableType::Base)
2382 .meta(meta)
2383 .build()
2384 .unwrap(),
2385 );
2386 Arc::new(table::Table::new(
2387 info,
2388 table::metadata::FilterPushDownType::Unsupported,
2389 Arc::new(DummyDataSource),
2390 ))
2391 }
2392
2393 fn make_metric_physical_table_ref_with_time_unit(unit: TimeUnit) -> TableRef {
2394 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2395 ColumnSchema::new(
2396 greptime_timestamp(),
2397 ConcreteDataType::timestamp_datatype(unit),
2398 false,
2399 )
2400 .with_time_index(true),
2401 ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
2402 ])
2403 .unwrap()
2404 .build()
2405 .unwrap();
2406 let meta = TableMetaBuilder::empty()
2407 .schema(Arc::new(schema))
2408 .primary_key_indices(vec![])
2409 .value_indices(vec![1])
2410 .engine("metric")
2411 .next_column_id(0)
2412 .options(Default::default())
2413 .created_on(Default::default())
2414 .build()
2415 .unwrap();
2416 let info = Arc::new(
2417 TableInfoBuilder::default()
2418 .table_id(1)
2419 .table_version(0)
2420 .name("greptime_physical_table")
2421 .schema_name(DEFAULT_SCHEMA_NAME)
2422 .catalog_name(DEFAULT_CATALOG_NAME)
2423 .desc(None)
2424 .table_type(TableType::Base)
2425 .meta(meta)
2426 .build()
2427 .unwrap(),
2428 );
2429 Arc::new(table::Table::new(
2430 info,
2431 table::metadata::FilterPushDownType::Unsupported,
2432 Arc::new(DummyDataSource),
2433 ))
2434 }
2435
2436 fn ms_row_insert_request(timestamp_ms: i64) -> RowInsertRequest {
2437 ms_row_insert_request_named("my_metric", timestamp_ms)
2438 }
2439
2440 fn ms_row_insert_request_named(table: &str, timestamp_ms: i64) -> RowInsertRequest {
2441 RowInsertRequest {
2442 table_name: table.to_string(),
2443 rows: Some(Rows {
2444 schema: vec![
2445 time_index_column_schema(
2446 greptime_timestamp(),
2447 ColumnDataType::TimestampMillisecond,
2448 ),
2449 field_column_schema(greptime_value(), ColumnDataType::Float64),
2450 ],
2451 rows: vec![Row {
2452 values: vec![
2453 Value {
2454 value_data: Some(ValueData::TimestampMillisecondValue(timestamp_ms)),
2455 },
2456 Value {
2457 value_data: Some(ValueData::F64Value(1.0)),
2458 },
2459 ],
2460 }],
2461 }),
2462 }
2463 }
2464
2465 fn convert_row_insert_requests_time_unit(
2467 requests: &mut RowInsertRequests,
2468 target_unit: TimeUnit,
2469 ) -> Result<()> {
2470 for request in &mut requests.inserts {
2471 let Some(rows) = request.rows.as_mut() else {
2472 continue;
2473 };
2474 convert_rows_time_unit(rows, target_unit)?;
2475 }
2476 Ok(())
2477 }
2478
2479 #[test]
2480 fn test_convert_row_insert_requests_time_unit_noop_when_matching() {
2481 let mut requests = RowInsertRequests {
2482 inserts: vec![ms_row_insert_request(123)],
2483 };
2484 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Millisecond).unwrap();
2485 let rows = requests.inserts[0].rows.as_ref().unwrap();
2486 assert_eq!(
2487 rows.schema[0].datatype,
2488 ColumnDataType::TimestampMillisecond as i32
2489 );
2490 assert!(matches!(
2491 rows.rows[0].values[0].value_data,
2492 Some(ValueData::TimestampMillisecondValue(123))
2493 ));
2494 }
2495
2496 #[test]
2497 fn test_convert_row_insert_requests_time_unit_widens_losslessly() {
2498 let mut requests = RowInsertRequests {
2499 inserts: vec![ms_row_insert_request(123)],
2500 };
2501 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Microsecond).unwrap();
2502 let rows = requests.inserts[0].rows.as_ref().unwrap();
2503 assert_eq!(
2504 rows.schema[0].datatype,
2505 ColumnDataType::TimestampMicrosecond as i32
2506 );
2507 assert!(matches!(
2508 rows.rows[0].values[0].value_data,
2509 Some(ValueData::TimestampMicrosecondValue(123_000))
2510 ));
2511 }
2512
2513 #[test]
2514 fn test_convert_row_insert_requests_time_unit_truncates_on_narrowing() {
2515 let requests = |value_ns: i64| RowInsertRequests {
2518 inserts: vec![RowInsertRequest {
2519 table_name: "my_metric".to_string(),
2520 rows: Some(Rows {
2521 schema: vec![
2522 time_index_column_schema(
2523 greptime_timestamp(),
2524 ColumnDataType::TimestampNanosecond,
2525 ),
2526 field_column_schema(greptime_value(), ColumnDataType::Float64),
2527 ],
2528 rows: vec![Row {
2529 values: vec![
2530 Value {
2531 value_data: Some(ValueData::TimestampNanosecondValue(value_ns)),
2532 },
2533 Value {
2534 value_data: Some(ValueData::F64Value(1.0)),
2535 },
2536 ],
2537 }],
2538 }),
2539 }],
2540 };
2541
2542 let mut reqs = requests(123_456_789);
2543 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
2544 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2545 assert!(matches!(
2546 rows.rows[0].values[0].value_data,
2547 Some(ValueData::TimestampMicrosecondValue(123_456))
2548 ));
2549
2550 let mut reqs = requests(123_456_789);
2551 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Millisecond).unwrap();
2552 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2553 assert!(matches!(
2554 rows.rows[0].values[0].value_data,
2555 Some(ValueData::TimestampMillisecondValue(123))
2556 ));
2557
2558 let mut reqs = requests(-123_456_789);
2559 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
2560 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2561 assert!(matches!(
2562 rows.rows[0].values[0].value_data,
2563 Some(ValueData::TimestampMicrosecondValue(-123_457))
2564 ));
2565 }
2566
2567 #[test]
2568 fn test_convert_row_insert_requests_time_unit_overflow() {
2569 let mut requests = RowInsertRequests {
2570 inserts: vec![ms_row_insert_request(i64::MAX)],
2571 };
2572 let err =
2573 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Nanosecond).unwrap_err();
2574 assert!(err.to_string().contains("overflows"), "{err}");
2575 }
2576
2577 #[test]
2578 fn test_table_time_index_unit() {
2579 assert_eq!(
2580 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
2581 TimeUnit::Microsecond
2582 )),
2583 Some(TimeUnit::Microsecond)
2584 );
2585 assert_eq!(
2586 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
2587 TimeUnit::Millisecond
2588 )),
2589 Some(TimeUnit::Millisecond)
2590 );
2591 }
2592
2593 #[tokio::test]
2594 async fn test_accommodate_existing_schema_and_reject_kind_changes() {
2595 let ts_name = "my_ts";
2596 let field_name = "my_field";
2597 let table =
2598 make_table_ref_with_schema(ts_name, field_name, ConcreteDataType::float64_datatype());
2599
2600 let mut req = RowInsertRequest {
2602 table_name: "test_table".to_string(),
2603 rows: Some(Rows {
2604 schema: vec![
2605 time_index_column_schema("ts_wrong", ColumnDataType::TimestampMillisecond),
2606 field_column_schema("field_wrong", ColumnDataType::Float64),
2607 ],
2608 rows: vec![api::v1::Row {
2609 values: vec![Value::default(), Value::default()],
2610 }],
2611 }),
2612 };
2613 let ctx = Arc::new(QueryContext::with(
2614 DEFAULT_CATALOG_NAME,
2615 DEFAULT_SCHEMA_NAME,
2616 ));
2617
2618 let kv_backend = prepare_mocked_backend().await;
2619 let inserter = Inserter::new(
2620 catalog::memory::MemoryCatalogManager::new(),
2621 create_partition_rule_manager(kv_backend.clone()).await,
2622 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
2623 Arc::new(new_table_flownode_set_cache(
2624 String::new(),
2625 Cache::new(100),
2626 kv_backend.clone(),
2627 )),
2628 true,
2629 );
2630 assert!(
2632 inserter
2633 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, false, false, false)
2634 .unwrap()
2635 .is_none()
2636 );
2637 let alter_expr = inserter
2638 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true, true)
2639 .unwrap();
2640 assert!(alter_expr.is_none());
2641
2642 let req_schema = req.rows.as_ref().unwrap().schema.clone();
2644 assert_eq!(req_schema[0].column_name, ts_name);
2645 assert_eq!(req_schema[1].column_name, field_name);
2646
2647 let (datatype, datatype_extension) =
2648 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
2649 .unwrap()
2650 .into_parts();
2651 let mut histogram_req = RowInsertRequest {
2652 table_name: "test_table".to_string(),
2653 rows: Some(Rows {
2654 schema: vec![
2655 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2656 api::v1::ColumnSchema {
2657 column_name: greptime_native_histogram().to_string(),
2658 datatype: datatype as i32,
2659 semantic_type: SemanticType::Field as i32,
2660 datatype_extension,
2661 options: None,
2662 },
2663 ],
2664 rows: vec![],
2665 }),
2666 };
2667 let error = inserter
2668 .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, false, true, true)
2669 .unwrap_err();
2670 assert!(
2671 error
2672 .to_string()
2673 .contains("cannot mix native histogram and float sample fields")
2674 );
2675
2676 let histogram_table = make_table_ref_with_schema(
2677 "ts",
2678 greptime_native_histogram(),
2679 native_histogram_value_type().clone(),
2680 );
2681 let mut sample_req = RowInsertRequest {
2682 table_name: "test_table".to_string(),
2683 rows: Some(Rows {
2684 schema: vec![
2685 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2686 field_column_schema(greptime_value(), ColumnDataType::Float64),
2687 ],
2688 rows: vec![],
2689 }),
2690 };
2691 let error = inserter
2692 .get_alter_table_expr_on_demand(
2693 &mut sample_req,
2694 &histogram_table,
2695 &ctx,
2696 false,
2697 true,
2698 true,
2699 )
2700 .unwrap_err();
2701 assert!(
2702 error
2703 .to_string()
2704 .contains("cannot mix native histogram and float sample fields")
2705 );
2706 }
2707
2708 #[test]
2709 fn test_native_histogram_detection_survives_prefix_change() {
2710 set_default_prefix(Some("custom")).unwrap();
2711 let table = make_table_ref_with_schema(
2712 "custom_timestamp",
2713 NATIVE_HISTOGRAM_FIELD,
2714 native_histogram_value_type().clone(),
2715 );
2716 let (datatype, datatype_extension) =
2717 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
2718 .unwrap()
2719 .into_parts();
2720 let request_schema = [api::v1::ColumnSchema {
2721 column_name: greptime_native_histogram().to_string(),
2722 datatype: datatype as i32,
2723 semantic_type: SemanticType::Field as i32,
2724 datatype_extension,
2725 options: None,
2726 }];
2727
2728 assert!(request_is_native_histogram(&request_schema));
2729 assert!(table_is_native_histogram(&table));
2730 }
2731
2732 #[tokio::test]
2734 async fn test_write_meter_admission() {
2735 use std::cell::Cell;
2736 use std::sync::Mutex;
2737 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
2738
2739 use api::region::RegionResponse;
2740 use api::v1::region::region_request::Body;
2741 use arrow::array::{Int32Array, TimestampMillisecondArray};
2742 use arrow::record_batch::RecordBatch;
2743 use bytes::Bytes;
2744 use common_error::ext::{ErrorExt, RetryHint};
2745 use common_error::status_code::StatusCode;
2746 use common_grpc::flight::{FlightEncoder, FlightMessage};
2747 use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher;
2748 use futures::future::BoxFuture;
2749 use meter_core::ItemCalculator;
2750 use meter_core::collect::{Collect, WriteRejected};
2751 use meter_core::data::MeterRecord;
2752 use meter_core::global::global_registry;
2753 use session::context::Channel;
2754
2755 const CATALOG: &str = "write_meter_test";
2756
2757 #[derive(Default)]
2758 struct Meter {
2759 reject: AtomicBool,
2760 attempts: Mutex<Vec<MeterRecord>>,
2761 accepted_value: AtomicU64,
2762 }
2763
2764 impl Collect for Meter {
2765 fn on_write(
2766 &self,
2767 record: MeterRecord,
2768 ) -> BoxFuture<'_, std::result::Result<(), WriteRejected>> {
2769 Box::pin(async move {
2770 if record.catalog != CATALOG {
2771 return Ok(());
2772 }
2773 let value = record.value;
2774 self.attempts.lock().unwrap().push(record);
2775 if self.reject.load(Ordering::Relaxed) {
2776 return Err(WriteRejected::new("database row quota exhausted"));
2777 }
2778 self.accepted_value.fetch_add(value, Ordering::Relaxed);
2779 Ok(())
2780 })
2781 }
2782
2783 fn on_read(&self, _: MeterRecord) {}
2784 }
2785
2786 impl ItemCalculator<InstantAndNormalInsertRequests> for Meter {
2787 fn calc(&self, _: &InstantAndNormalInsertRequests) -> u64 {
2788 17
2789 }
2790 }
2791
2792 let kv_backend = prepare_mocked_backend().await;
2793 let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
2794 let (sender, mut dispatched) = tokio::sync::mpsc::channel(16);
2795 let watcher = DatanodeWatcher::new(sender).with_handler(|_, request| {
2796 let rows = match request.body.unwrap() {
2797 Body::Inserts(requests) => requests
2798 .requests
2799 .iter()
2800 .filter_map(|request| request.rows.as_ref())
2801 .map(|rows| rows.rows.len())
2802 .sum(),
2803 Body::BulkInsert(_) => 2,
2805 body => panic!("unexpected request: {body:?}"),
2806 };
2807 Ok(RegionResponse::new(rows))
2808 });
2809 let flow_cache = Cache::new(100);
2810 let inserter = Inserter::new(
2811 catalog::memory::MemoryCatalogManager::new(),
2812 partition_manager,
2813 Arc::new(MockDatanodeManager::new(watcher)),
2814 Arc::new(new_table_flownode_set_cache(
2815 String::new(),
2816 flow_cache.clone(),
2817 kv_backend,
2818 )),
2819 true,
2820 );
2821 let mut table_info = new_test_table_info(1, "table_1", [1].into_iter());
2822 table_info.catalog_name = CATALOG.to_string();
2823 table_info.schema_name = "target_db".to_string();
2824 let table_info = Arc::new(table_info);
2825 let table_infos = HashMap::from_iter([(1, table_info.clone())]);
2826 let ctx = Arc::new(QueryContext::with_channel(
2827 DEFAULT_CATALOG_NAME,
2828 DEFAULT_SCHEMA_NAME,
2829 Channel::Postgres,
2830 ));
2831 let rows_request = || {
2832 let build = |num_rows| RegionInsertRequests {
2833 requests: vec![RegionInsertRequest {
2834 region_id: RegionId::new(1, 1).as_u64(),
2835 rows: Some(Rows {
2836 schema: vec![],
2837 rows: vec![api::v1::Row { values: vec![] }; num_rows],
2838 }),
2839 ..Default::default()
2840 }],
2841 };
2842 InstantAndNormalInsertRequests {
2843 normal_requests: build(3),
2844 instant_requests: build(2),
2845 }
2846 };
2847
2848 let output = inserter
2850 .do_request(rows_request(), &table_infos, &ctx)
2851 .await
2852 .unwrap();
2853 assert_eq!(output.meta.cost, 0);
2854 assert!(matches!(output.data, OutputData::AffectedRows(3)));
2855 dispatched.try_recv().unwrap();
2856
2857 let meter = Arc::new(Meter::default());
2858 global_registry().set_collector(meter.clone());
2859 global_registry().register_calculator(meter.clone());
2860 let enabled = Cell::new(false);
2862 write_meter!({
2863 enabled.set(true);
2864 MeterRecord::new("probe".into(), "probe".into(), 0, 0, 0)
2865 })
2866 .await
2867 .unwrap();
2868 let enabled = enabled.get();
2869
2870 let output = inserter
2871 .do_request(rows_request(), &table_infos, &ctx)
2872 .await
2873 .unwrap();
2874 assert_eq!(output.meta.cost, if enabled { 17 } else { 0 });
2875 assert!(matches!(output.data, OutputData::AffectedRows(3)));
2876 dispatched.try_recv().unwrap();
2877 assert!(flow_cache.contains_key(&1));
2878 flow_cache.invalidate_all();
2879 meter.reject.store(true, Ordering::Relaxed);
2880 let result = inserter
2881 .do_request(rows_request(), &table_infos, &ctx)
2882 .await;
2883 if enabled {
2884 let error = result.unwrap_err();
2885 assert_eq!(error.status_code(), StatusCode::RateLimited);
2886 assert_eq!(error.retry_hint(), RetryHint::Retryable);
2887 assert!(error.to_string().contains("database row quota exhausted"));
2888 assert!(dispatched.try_recv().is_err());
2889 assert!(
2890 !flow_cache.contains_key(&1),
2891 "rejected write reached flow mirroring"
2892 );
2893 assert_eq!(meter.accepted_value.load(Ordering::Relaxed), 17);
2894 let attempts = meter.attempts.lock().unwrap();
2895 assert_eq!(attempts.len(), 2);
2896 for record in attempts.iter() {
2897 assert_eq!(record.catalog, CATALOG);
2898 assert_eq!(record.schema, "target_db");
2899 assert_eq!(
2900 (record.rows, record.value, record.source),
2901 (5, 17, Channel::Postgres as u8)
2902 );
2903 }
2904 } else {
2905 assert_eq!(result.unwrap().meta.cost, 0);
2906 dispatched.try_recv().unwrap();
2907 assert!(meter.attempts.lock().unwrap().is_empty());
2908 }
2909 meter.attempts.lock().unwrap().clear();
2910
2911 if enabled {
2913 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
2914 let rows = Rows {
2915 schema: vec![
2916 api::v1::helper::tag_column_schema("a", ColumnDataType::Int32),
2917 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2918 field_column_schema("b", ColumnDataType::Int32),
2919 ],
2920 rows: vec![api::v1::Row {
2921 values: vec![
2922 api::v1::value::ValueData::I32Value(60).into(),
2923 Value {
2924 value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue(
2925 0,
2926 )),
2927 },
2928 api::v1::value::ValueData::I32Value(0).into(),
2929 ],
2930 }],
2931 };
2932 let error = inserter
2933 .submit_table_rows(rows, table_info.clone(), ctx.clone(), &batcher)
2934 .await
2935 .unwrap_err();
2936 assert_eq!(error.status_code(), StatusCode::RateLimited);
2937 let mut attempts = meter.attempts.lock().unwrap();
2938 assert_eq!(attempts.len(), 1);
2939 assert_eq!(attempts[0].catalog, CATALOG);
2940 assert_eq!(attempts[0].schema, "target_db");
2941 assert_eq!(attempts[0].rows, 1);
2942 attempts.clear();
2943 }
2944
2945 let table = Arc::new(table::Table::new(
2946 table_info.clone(),
2947 table::metadata::FilterPushDownType::Unsupported,
2948 Arc::new(DummyDataSource),
2949 ));
2950 let batch = RecordBatch::try_new(
2951 table_info.meta.schema.arrow_schema().clone(),
2952 vec![
2953 Arc::new(Int32Array::from(vec![60, 70])),
2954 Arc::new(TimestampMillisecondArray::from(vec![0, 1])),
2955 Arc::new(Int32Array::from(vec![0, 0])),
2956 ],
2957 )
2958 .unwrap();
2959 let bulk_insert = |batch: RecordBatch| {
2960 let flight_data = FlightEncoder::default()
2961 .encode(FlightMessage::RecordBatch(batch.clone()))
2962 .into_iter()
2963 .next()
2964 .unwrap();
2965 inserter.handle_bulk_insert(
2966 table.clone(),
2967 flight_data,
2968 batch,
2969 Bytes::new(),
2970 false,
2971 Channel::Grpc,
2972 )
2973 };
2974 assert_eq!(bulk_insert(batch.slice(0, 0)).await.unwrap(), 0);
2976 assert!(meter.attempts.lock().unwrap().is_empty());
2977 assert!(dispatched.try_recv().is_err());
2978
2979 for reject in [false, true, false] {
2981 meter.reject.store(reject, Ordering::Relaxed);
2982 let result = bulk_insert(batch.clone()).await;
2983 if enabled && reject {
2984 let error = result.unwrap_err();
2985 assert_eq!(error.status_code(), StatusCode::RateLimited);
2986 assert_eq!(error.retry_hint(), RetryHint::Retryable);
2987 assert!(dispatched.try_recv().is_err());
2988 } else {
2989 assert_eq!(result.unwrap(), 2);
2990 dispatched.try_recv().unwrap();
2991 }
2992 }
2993 {
2994 let attempts = meter.attempts.lock().unwrap();
2995 assert_eq!(attempts.len(), if enabled { 3 } else { 0 });
2996 for record in attempts.iter() {
2997 assert_eq!(record.catalog, CATALOG);
2998 assert_eq!(record.schema, "target_db");
2999 assert_eq!(
3000 (record.rows, record.value, record.source),
3001 (2, 0, Channel::Grpc as u8)
3002 );
3003 }
3004 }
3005 assert_eq!(
3006 meter.accepted_value.load(Ordering::Relaxed),
3007 if enabled { 17 } else { 0 }
3008 );
3009
3010 meter.attempts.lock().unwrap().clear();
3013 let original = Arc::new(QueryContext::with_channel(
3014 CATALOG,
3015 "a",
3016 Channel::Prometheus,
3017 ));
3018 let mut batches = ["a", "b", "a"].map(|schema| {
3019 let ctx = if schema == "a" {
3020 original.clone()
3021 } else {
3022 Arc::new(QueryContext::with_channel(
3023 CATALOG,
3024 schema,
3025 Channel::Prometheus,
3026 ))
3027 };
3028 (
3029 ctx,
3030 RowInsertRequests {
3031 inserts: vec![RowInsertRequest {
3032 table_name: "data".into(),
3033 rows: Some(Rows {
3034 schema: vec![],
3035 rows: vec![api::v1::Row::default(); 2],
3036 }),
3037 }],
3038 },
3039 )
3040 });
3041 admit_row_insert_batches(&mut batches).await.unwrap();
3042 admit_row_insert_batches(&mut batches).await.unwrap();
3043 assert_eq!(original.write_rows_to_admit(CATALOG, "a", 4), 4);
3044 for (ctx, _) in &batches {
3045 assert_eq!(
3046 ctx.write_rows_to_admit(CATALOG, &ctx.current_schema(), 2),
3047 0
3048 );
3049 assert_eq!(ctx.channel(), Channel::Prometheus);
3050 }
3051 {
3052 let attempts = meter.attempts.lock().unwrap();
3053 let totals = attempts
3054 .iter()
3055 .map(|r| (r.schema.as_str(), r.rows, r.value))
3056 .collect::<Vec<_>>();
3057 assert_eq!(
3058 totals,
3059 if enabled {
3060 vec![("a", 4, 0), ("b", 2, 0)]
3061 } else {
3062 vec![]
3063 }
3064 );
3065 }
3066 meter.attempts.lock().unwrap().clear();
3067 meter.reject.store(false, Ordering::Relaxed);
3068 let ctx = Arc::new(QueryContext::with_channel(
3069 CATALOG,
3070 "logical",
3071 Channel::Otlp,
3072 ));
3073 let admitted = admit_write(2, &ctx).await.unwrap();
3074 let mut requests = RowInsertRequests {
3075 inserts: vec![RowInsertRequest {
3076 table_name: "metric".to_string(),
3077 rows: Some(Rows {
3078 schema: vec![],
3079 rows: vec![api::v1::Row::default(); 2],
3080 }),
3081 }],
3082 };
3083 let original = requests.clone();
3084 let cost = Inserter::meter_row_inserts(&mut requests, &admitted)
3085 .await
3086 .unwrap();
3087 assert_eq!(cost, if enabled { 17 } else { 0 });
3088 assert_eq!(requests, original);
3089 let records = meter
3090 .attempts
3091 .lock()
3092 .unwrap()
3093 .iter()
3094 .map(|record| (record.rows, record.value))
3095 .collect::<Vec<_>>();
3096 assert_eq!(
3097 records,
3098 if enabled {
3099 vec![(2, 0), (0, 17)]
3100 } else {
3101 vec![]
3102 }
3103 );
3104 meter.reject.store(true, Ordering::Relaxed);
3105 let result = Inserter::meter_row_inserts(&mut requests, &admitted).await;
3106 if enabled {
3107 assert_eq!(result.unwrap_err().status_code(), StatusCode::RateLimited);
3108 } else {
3109 assert_eq!(result.unwrap(), 0);
3110 }
3111 assert_eq!(requests, original);
3112 }
3113
3114 #[test]
3115 fn test_trace_create_options_copy_compaction_hints() {
3116 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3117 ctx.set_extension("compaction.twcs.time_window", "1h");
3118 ctx.set_extension("compaction.twcs.max_output_file_size", "256MB");
3119 ctx.set_extension("compaction.twcs.remote_compaction", "true");
3120 ctx.set_extension("compaction.unknown", "ignored");
3121 let ctx = Arc::new(ctx);
3122 let mut table_options = Default::default();
3123
3124 fill_table_options_for_create(
3125 &mut table_options,
3126 &AutoCreateTableType::Trace {
3127 alter_existing: true,
3128 },
3129 &ctx,
3130 );
3131
3132 assert_eq!(
3133 Some("twcs"),
3134 table_options.get(COMPACTION_TYPE).map(String::as_str)
3135 );
3136 assert_eq!(
3137 Some("1h"),
3138 table_options
3139 .get("compaction.twcs.time_window")
3140 .map(String::as_str)
3141 );
3142 assert_eq!(
3143 Some("256MB"),
3144 table_options
3145 .get("compaction.twcs.max_output_file_size")
3146 .map(String::as_str)
3147 );
3148 assert_eq!(
3149 Some("true"),
3150 table_options
3151 .get("compaction.twcs.remote_compaction")
3152 .map(String::as_str)
3153 );
3154 assert!(!table_options.contains_key("compaction.unknown"));
3155 }
3156
3157 #[test]
3158 fn test_skip_wal_does_not_change_table_options() {
3159 check_skip_wal_does_not_change_table_options(false);
3160 check_skip_wal_does_not_change_table_options(true);
3161 }
3162
3163 fn check_skip_wal_does_not_change_table_options(skip_wal: bool) {
3164 let ctx = Arc::new(QueryContext::with(
3165 DEFAULT_CATALOG_NAME,
3166 DEFAULT_SCHEMA_NAME,
3167 ));
3168 ctx.set_skip_wal(skip_wal);
3169 let mut options = Default::default();
3170 fill_table_options_for_create(&mut options, &AutoCreateTableType::Physical, &ctx);
3171 assert!(!options.contains_key(session::hints::INSERT_SKIP_WAL_HINT));
3172 assert!(!options.contains_key("skip_wal"));
3173 }
3174
3175 #[test]
3176 fn test_last_non_null_create_options_preserve_default_without_append_mode() {
3177 let ctx = Arc::new(QueryContext::with(
3178 DEFAULT_CATALOG_NAME,
3179 DEFAULT_SCHEMA_NAME,
3180 ));
3181 let mut table_options = Default::default();
3182
3183 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3184
3185 assert_eq!(
3186 Some("last_non_null"),
3187 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3188 );
3189 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3190 }
3191
3192 #[test]
3193 fn test_fill_table_options_copies_semantic_extensions() {
3194 use table::requests::{
3195 SEMANTIC_METRIC_TYPE, SEMANTIC_PER_TABLE_INDEX_KEY, SEMANTIC_SIGNAL_TYPE,
3196 SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, SOURCE_OPENTELEMETRY,
3197 };
3198
3199 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3200 ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
3201 ctx.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
3202 ctx.set_extension(SEMANTIC_SOURCE_VERSION, "2.0");
3203 ctx.set_extension(SEMANTIC_METRIC_TYPE, "bogus");
3204 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, "{}");
3206 let ctx = Arc::new(ctx);
3207 let mut table_options = Default::default();
3208
3209 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::Physical, &ctx);
3210
3211 assert_eq!(
3212 Some(SIGNAL_TYPE_METRIC),
3213 table_options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str)
3214 );
3215 assert_eq!(
3216 Some(SOURCE_OPENTELEMETRY),
3217 table_options.get(SEMANTIC_SOURCE).map(String::as_str)
3218 );
3219 assert_eq!(
3220 Some("2.0"),
3221 table_options
3222 .get(SEMANTIC_SOURCE_VERSION)
3223 .map(String::as_str)
3224 );
3225 assert!(!table_options.contains_key(SEMANTIC_METRIC_TYPE));
3226 assert!(!table_options.contains_key(SEMANTIC_PER_TABLE_INDEX_KEY));
3227 }
3228
3229 #[test]
3230 fn test_apply_per_table_semantic_options() {
3231 use table::requests::{
3232 SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, SEMANTIC_PER_TABLE_INDEX_KEY,
3233 };
3234
3235 let index = format!(
3236 r#"{{
3237 "{DEFAULT_SCHEMA_NAME}": {{
3238 "http_requests_total": {{
3239 "greptime.semantic.metric.type": "counter",
3240 "greptime.semantic.metric.unit": "By",
3241 "greptime.semantic.metric.type_BOGUS": "x"
3242 }},
3243 "other_table": {{
3244 "greptime.semantic.metric.type": "gauge"
3245 }}
3246 }},
3247 "other_schema": {{
3248 "http_requests_total": {{
3249 "greptime.semantic.metric.type": "gauge"
3250 }}
3251 }}
3252 }}"#
3253 );
3254 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3255 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index);
3256 let ctx = Arc::new(ctx);
3257
3258 let index = parse_per_table_semantic_index(&ctx);
3259 assert!(index.is_some());
3260 let index = index.as_ref();
3261
3262 let mut table_options = std::collections::HashMap::new();
3263 apply_per_table_semantic_options(
3264 &mut table_options,
3265 index,
3266 DEFAULT_SCHEMA_NAME,
3267 "http_requests_total",
3268 );
3269 assert_eq!(
3271 table_options.get(SEMANTIC_METRIC_TYPE).map(String::as_str),
3272 Some("counter")
3273 );
3274 assert_eq!(
3275 table_options.get(SEMANTIC_METRIC_UNIT).map(String::as_str),
3276 Some("By")
3277 );
3278 assert!(!table_options.contains_key("greptime.semantic.metric.type_BOGUS"));
3281 assert_eq!(table_options.len(), 2);
3282
3283 let mut empty = std::collections::HashMap::new();
3284 apply_per_table_semantic_options(&mut empty, index, DEFAULT_SCHEMA_NAME, "not_in_index");
3285 assert!(empty.is_empty());
3286
3287 let mut opts = std::collections::HashMap::new();
3290 apply_per_table_semantic_options(
3291 &mut opts,
3292 index,
3293 "schema_without_entry",
3294 "http_requests_total",
3295 );
3296 assert!(opts.is_empty());
3297
3298 let bare = Arc::new(QueryContext::with(
3300 DEFAULT_CATALOG_NAME,
3301 DEFAULT_SCHEMA_NAME,
3302 ));
3303 assert!(parse_per_table_semantic_index(&bare).is_none());
3304 let mut opts = std::collections::HashMap::new();
3305 apply_per_table_semantic_options(
3306 &mut opts,
3307 None,
3308 DEFAULT_SCHEMA_NAME,
3309 "http_requests_total",
3310 );
3311 assert!(opts.is_empty());
3312 }
3313
3314 #[test]
3315 fn test_last_non_null_create_options_preserve_default_with_append_mode_false() {
3316 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3317 ctx.set_extension(APPEND_MODE_KEY, "false");
3318 let ctx = Arc::new(ctx);
3319 let mut table_options = Default::default();
3320
3321 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3322
3323 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3324 assert_eq!(
3325 Some("last_non_null"),
3326 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3327 );
3328 }
3329
3330 #[test]
3331 fn test_last_non_null_create_options_use_configured_merge_mode() {
3332 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3333 ctx.set_extension(MERGE_MODE_KEY, "last_row");
3334 let ctx = Arc::new(ctx);
3335 let mut table_options = Default::default();
3336
3337 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3338
3339 assert_eq!(
3340 Some("last_row"),
3341 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3342 );
3343 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3344 }
3345
3346 #[test]
3347 fn test_last_non_null_create_options_use_last_row_with_append_mode_true() {
3348 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3349 ctx.set_extension(APPEND_MODE_KEY, "true");
3350 let ctx = Arc::new(ctx);
3351 let mut table_options = Default::default();
3352
3353 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3354
3355 assert_eq!(
3356 Some("true"),
3357 table_options.get(APPEND_MODE_KEY).map(String::as_str)
3358 );
3359 assert_eq!(
3360 Some("last_row"),
3361 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3362 );
3363 }
3364
3365 struct UnexpectedBatcher;
3366
3367 #[async_trait::async_trait]
3368 impl PendingRowsBatcher for UnexpectedBatcher {
3369 async fn acquire(&self) -> Result<Arc<tokio::sync::OwnedSemaphorePermit>> {
3370 panic!("empty writes must not acquire batch admission")
3371 }
3372
3373 async fn submit(
3374 &self,
3375 _table_info: TableInfoRef,
3376 _batch: arrow::record_batch::RecordBatch,
3377 _ctx: QueryContextRef,
3378 _permit: Arc<tokio::sync::OwnedSemaphorePermit>,
3379 ) -> Result<Output> {
3380 panic!("empty writes must not submit a batch")
3381 }
3382 }
3383
3384 async fn batcher_test_inserter() -> Inserter {
3385 let kv_backend = prepare_mocked_backend().await;
3386 Inserter::new(
3387 catalog::memory::MemoryCatalogManager::new(),
3388 create_partition_rule_manager(kv_backend.clone()).await,
3389 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3390 Arc::new(new_table_flownode_set_cache(
3391 String::new(),
3392 Cache::new(100),
3393 kv_backend,
3394 )),
3395 true,
3396 )
3397 }
3398
3399 #[tokio::test]
3400 async fn test_logical_batcher_eligibility() {
3401 use catalog::RegisterTableRequest;
3402 use catalog::memory::MemoryCatalogManager;
3403 use common_meta::instruction::{CacheIdent, CreateFlow};
3404 use common_meta::kv_backend::KvBackendRef;
3405 use common_meta::kv_backend::memory::MemoryKvBackend;
3406 use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder};
3407 let requests = RowInsertRequests {
3408 inserts: vec![RowInsertRequest {
3409 table_name: "test_table".to_string(),
3410 rows: None,
3411 }],
3412 };
3413 let original =
3414 make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3415 .table_info();
3416 for case in [
3417 "eligible",
3418 "physical",
3419 "ordinary",
3420 "instant",
3421 "disabled",
3422 "hint",
3423 "flow",
3424 "required_tag",
3425 "default",
3426 ] {
3427 let mut info = (*original).clone();
3428 info.meta.engine = METRIC_ENGINE_NAME.to_string();
3429 info.meta.options.extra_options.insert(
3430 LOGICAL_TABLE_METADATA_KEY.to_string(),
3431 "physical".to_string(),
3432 );
3433 let mut ctx = QueryContext::arc().fork();
3434 match case {
3435 "physical" => {
3436 info.meta
3437 .options
3438 .extra_options
3439 .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "other".to_string());
3440 }
3441 "ordinary" => info.meta.engine = "mito".to_string(),
3442 "instant" => info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant),
3443 "hint" => ctx.set_extension(AUTO_CREATE_TABLE_KEY, "false"),
3444 "required_tag" | "default" => {
3445 let mut columns = info.meta.schema.column_schemas().to_vec();
3446 let mut tag = ColumnSchema::new(
3447 "tag",
3448 ConcreteDataType::string_datatype(),
3449 case != "required_tag",
3450 );
3451 if case == "default" {
3452 tag = tag
3453 .with_default_constraint(Some(ColumnDefaultConstraint::null_value()))
3454 .unwrap();
3455 }
3456 columns.push(tag);
3457 info.meta.schema = Arc::new(
3458 SchemaBuilder::try_from_columns(columns)
3459 .unwrap()
3460 .build()
3461 .unwrap(),
3462 );
3463 info.meta.primary_key_indices = vec![2];
3464 }
3465 _ => {}
3466 }
3467 let catalog = MemoryCatalogManager::with_default_setup();
3468 let table = Arc::new(table::Table::new(
3469 Arc::new(info),
3470 table::metadata::FilterPushDownType::Unsupported,
3471 Arc::new(DummyDataSource),
3472 ));
3473 catalog
3474 .register_table_sync(RegisterTableRequest {
3475 catalog: DEFAULT_CATALOG_NAME.to_string(),
3476 schema: DEFAULT_SCHEMA_NAME.to_string(),
3477 table_name: "test_table".to_string(),
3478 table_id: 1,
3479 table,
3480 })
3481 .unwrap();
3482 let mut inserter = batcher_test_inserter().await;
3483 inserter.catalog_manager = catalog;
3484 inserter.auto_create_table = case != "disabled";
3485 let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::default());
3486 inserter.table_flownode_set_cache = Arc::new(new_table_flownode_set_cache(
3487 String::new(),
3488 Cache::new(10),
3489 kv_backend,
3490 ));
3491 if case == "flow" {
3492 inserter
3493 .table_flownode_set_cache
3494 .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
3495 flow_id: 1,
3496 source_table_ids: vec![1],
3497 partition_to_peer_mapping: vec![(0, Peer::empty(1))],
3498 })])
3499 .await
3500 .unwrap();
3501 }
3502 assert_eq!(
3503 inserter
3504 .can_batch_metric_rows(&requests, &Arc::new(ctx), "physical")
3505 .await
3506 .unwrap(),
3507 case == "eligible",
3508 "{case}"
3509 );
3510 }
3511 }
3512
3513 #[tokio::test]
3514 async fn test_batcher_meter_preserves_request() {
3515 let mut requests = RowInsertRequests {
3516 inserts: vec![RowInsertRequest {
3517 table_name: "sample".to_string(),
3518 rows: Some(Rows {
3519 schema: vec![],
3520 rows: vec![api::v1::Row {
3521 values: vec![Value {
3522 value_data: Some(api::v1::value::ValueData::F64Value(1.5)),
3523 }],
3524 }],
3525 }),
3526 }],
3527 };
3528 let expected = requests.clone();
3529 Inserter::meter_row_inserts(&mut requests, &QueryContext::arc())
3530 .await
3531 .unwrap();
3532 assert_eq!(requests, expected);
3533 }
3534
3535 #[tokio::test]
3536 async fn test_instant_table_bypasses_batcher() {
3537 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3538 let inserter = batcher_test_inserter()
3539 .await
3540 .with_pending_rows_batcher(Some(batcher));
3541 let mut ctx = session::context::QueryContextBuilder::default().build();
3542 ctx.set_batching_enabled(true);
3543 let ctx = Arc::new(ctx);
3544 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3545 .table_info();
3546 assert!(inserter.table_batcher(&table, &ctx).is_some());
3547 let mut instant = (*table).clone();
3548 instant.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
3549 assert!(inserter.table_batcher(&Arc::new(instant), &ctx).is_none());
3550 }
3551
3552 #[tokio::test]
3553 async fn test_empty_prepared_rows_skip_batcher() {
3554 let inserter = batcher_test_inserter().await;
3555 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3556 .table_info();
3557 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3558 let ctx = QueryContext::arc();
3559 let output = inserter
3560 .submit_table_rows(
3561 Rows {
3562 schema: vec![],
3563 rows: vec![],
3564 },
3565 table.clone(),
3566 ctx.clone(),
3567 &batcher,
3568 )
3569 .await
3570 .unwrap();
3571 assert!(matches!(output.data, OutputData::AffectedRows(0)));
3572 let output = inserter
3573 .submit_pending_rows(
3574 RowInsertRequests {
3575 inserts: vec![RowInsertRequest {
3576 table_name: table.name.clone(),
3577 rows: None,
3578 }],
3579 },
3580 HashMap::from_iter([(table.table_id(), table)]),
3581 ctx,
3582 &batcher,
3583 )
3584 .await
3585 .unwrap();
3586 assert!(matches!(output.data, OutputData::AffectedRows(0)));
3587 }
3588
3589 #[derive(Clone)]
3593 struct GatedFlownodeHandler {
3594 gate: Arc<tokio::sync::Semaphore>,
3595 affected_rows: u64,
3596 }
3597
3598 #[async_trait::async_trait]
3599 impl MockFlownodeHandler for GatedFlownodeHandler {
3600 async fn handle_inserts(
3601 &self,
3602 _peer: &Peer,
3603 _requests: api::v1::region::InsertRequests,
3604 ) -> common_meta::error::Result<FlowResponse> {
3605 let _permit = self.gate.acquire().await.unwrap();
3606 Ok(FlowResponse {
3607 affected_rows: self.affected_rows,
3608 ..Default::default()
3609 })
3610 }
3611 }
3612
3613 fn flownode_peer() -> Peer {
3614 Peer {
3615 id: 1,
3616 addr: "127.0.0.1:4001".to_string(),
3617 }
3618 }
3619
3620 #[derive(Clone)]
3621 struct RecordingNodeManager {
3622 datanode_dispatch: tokio::sync::mpsc::UnboundedSender<usize>,
3623 flownode_dispatch: tokio::sync::mpsc::UnboundedSender<Peer>,
3624 gates: Arc<std::sync::Mutex<HashMap<u64, Arc<tokio::sync::Semaphore>>>>,
3625 fail_flownode: bool,
3626 }
3627
3628 #[async_trait::async_trait]
3629 impl common_meta::node_manager::DatanodeManager for RecordingNodeManager {
3630 async fn datanode(&self, peer: &Peer) -> common_meta::node_manager::DatanodeRef {
3631 Arc::new(RecordingNode {
3632 peer: peer.clone(),
3633 manager: self.clone(),
3634 })
3635 }
3636 }
3637
3638 #[async_trait::async_trait]
3639 impl common_meta::node_manager::FlownodeManager for RecordingNodeManager {
3640 async fn flownode(&self, peer: &Peer) -> common_meta::node_manager::FlownodeRef {
3641 Arc::new(RecordingNode {
3642 peer: peer.clone(),
3643 manager: self.clone(),
3644 })
3645 }
3646 }
3647
3648 struct RecordingNode {
3649 peer: Peer,
3650 manager: RecordingNodeManager,
3651 }
3652
3653 #[async_trait::async_trait]
3654 impl common_meta::node_manager::Datanode for RecordingNode {
3655 async fn handle_request(
3656 &self,
3657 _request: Arc<api::v1::region::RegionRequest>,
3658 ) -> common_meta::error::Result<api::region::RegionResponse> {
3659 let rows = match _request.body.as_ref() {
3660 Some(api::v1::region::region_request::Body::Inserts(requests)) => requests
3661 .requests
3662 .iter()
3663 .filter_map(|request| request.rows.as_ref())
3664 .map(|rows| rows.rows.len())
3665 .sum(),
3666 _ => 0,
3667 };
3668 let _ = self.manager.datanode_dispatch.send(rows);
3669 Ok(api::region::RegionResponse::new(rows))
3670 }
3671
3672 async fn handle_query(
3673 &self,
3674 _request: common_query::request::QueryRequest,
3675 ) -> common_meta::error::Result<common_recordbatch::SendableRecordBatchStream> {
3676 unreachable!()
3677 }
3678 }
3679
3680 #[async_trait::async_trait]
3681 impl common_meta::node_manager::Flownode for RecordingNode {
3682 async fn handle(
3683 &self,
3684 _request: api::v1::flow::FlowRequest,
3685 ) -> common_meta::error::Result<FlowResponse> {
3686 unreachable!()
3687 }
3688
3689 async fn handle_inserts(
3690 &self,
3691 _request: api::v1::region::InsertRequests,
3692 ) -> common_meta::error::Result<FlowResponse> {
3693 let _ = self.manager.flownode_dispatch.send(self.peer.clone());
3694 if self.manager.fail_flownode {
3695 return Err(common_meta::error::UnexpectedSnafu {
3696 err_msg: "test flownode failure".to_string(),
3697 }
3698 .build());
3699 }
3700 let gate = self
3701 .manager
3702 .gates
3703 .lock()
3704 .unwrap()
3705 .get(&self.peer.id)
3706 .cloned();
3707 if let Some(gate) = gate {
3708 let _permit = gate.acquire().await.unwrap();
3709 }
3710 Ok(FlowResponse::default())
3711 }
3712
3713 async fn handle_mark_window_dirty(
3714 &self,
3715 _request: api::v1::flow::DirtyWindowRequests,
3716 ) -> common_meta::error::Result<FlowResponse> {
3717 unreachable!()
3718 }
3719 }
3720
3721 fn mirror_requests(peer: &Peer, num_rows: usize) -> HashMap<Peer, RegionInsertRequests> {
3722 HashMap::from_iter([(
3723 peer.clone(),
3724 RegionInsertRequests {
3725 requests: vec![RegionInsertRequest {
3726 region_id: RegionId::new(1, 1).as_u64(),
3727 rows: Some(Rows {
3728 schema: vec![],
3729 rows: vec![api::v1::Row { values: vec![] }; num_rows],
3730 }),
3731 ..Default::default()
3732 }],
3733 },
3734 )])
3735 }
3736
3737 static MIRROR_METRIC_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
3742
3743 async fn lock_mirror_metrics() -> tokio::sync::MutexGuard<'static, ()> {
3744 MIRROR_METRIC_TEST_LOCK.lock().await
3745 }
3746
3747 #[tokio::test]
3751 async fn flow_mirror_saturation_preserves_instant_and_normal_contracts() {
3752 use common_error::ext::{ErrorExt, RetryHint};
3753 use common_error::status_code::StatusCode;
3754
3755 let _guard = lock_mirror_metrics().await;
3756 let kv_backend = prepare_mocked_backend().await;
3757 let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
3758 let (datanode_tx, mut datanode_rx) = tokio::sync::mpsc::unbounded_channel();
3759 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
3760 let node_manager = Arc::new(RecordingNodeManager {
3761 datanode_dispatch: datanode_tx,
3762 flownode_dispatch: flownode_tx,
3763 gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
3764 fail_flownode: false,
3765 });
3766 let flow_cache = Cache::new(10);
3767 let flow_cache_backend = prepare_mocked_backend().await;
3768 let inserter = Inserter::new(
3769 catalog::memory::MemoryCatalogManager::new(),
3770 partition_manager,
3771 node_manager,
3772 Arc::new(new_table_flownode_set_cache(
3773 String::new(),
3774 flow_cache.clone(),
3775 flow_cache_backend,
3776 )),
3777 true,
3778 );
3779 let peer = flownode_peer();
3780 inserter
3781 .table_flownode_set_cache
3782 .invalidate(&[common_meta::instruction::CacheIdent::CreateFlow(
3783 common_meta::instruction::CreateFlow {
3784 flow_id: 1,
3785 source_table_ids: vec![1, 2],
3786 partition_to_peer_mapping: vec![(0, peer)],
3787 },
3788 )])
3789 .await
3790 .unwrap();
3791 let mut normal_info = new_test_table_info(1, "normal_table", [1].into_iter());
3792 normal_info.catalog_name = DEFAULT_CATALOG_NAME.to_string();
3793 normal_info.schema_name = DEFAULT_SCHEMA_NAME.to_string();
3794 let mut instant_info = new_test_table_info(2, "instant_table", [1].into_iter());
3795 instant_info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
3796 let table_infos =
3797 HashMap::from_iter([(1, Arc::new(normal_info)), (2, Arc::new(instant_info))]);
3798 let ctx = Arc::new(QueryContext::with(
3799 DEFAULT_CATALOG_NAME,
3800 DEFAULT_SCHEMA_NAME,
3801 ));
3802 let insert = |table_id, num_rows| RegionInsertRequests {
3803 requests: vec![RegionInsertRequest {
3804 region_id: RegionId::new(table_id, 1).as_u64(),
3805 rows: Some(Rows {
3806 schema: vec![],
3807 rows: vec![api::v1::Row::default(); num_rows],
3808 }),
3809 ..Default::default()
3810 }],
3811 };
3812 let make_request = |normal, instant| InstantAndNormalInsertRequests {
3813 normal_requests: if normal == 0 {
3814 RegionInsertRequests::default()
3815 } else {
3816 insert(1, normal)
3817 },
3818 instant_requests: if instant == 0 {
3819 RegionInsertRequests::default()
3820 } else {
3821 insert(2, instant)
3822 },
3823 };
3824
3825 let pending = inserter.mirror_pending_rows.clone();
3826 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
3827 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
3828 pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
3829 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
3830 for request in [make_request(0, 1), make_request(1, 1)] {
3831 let error = inserter
3832 .do_request(request, &table_infos, &ctx)
3833 .await
3834 .unwrap_err();
3835 assert_eq!(error.status_code(), StatusCode::RateLimited);
3836 assert_eq!(error.retry_hint(), RetryHint::Retryable);
3837 assert!(datanode_rx.try_recv().is_err());
3838 assert!(flownode_rx.try_recv().is_err());
3839 assert_eq!(
3840 pending.load(Ordering::Relaxed),
3841 MAX_MIRROR_PENDING_ROWS as u64
3842 );
3843 }
3844 assert_eq!(
3845 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
3846 dropped_before
3847 );
3848
3849 for pending_before in [0, MAX_MIRROR_PENDING_ROWS as u64] {
3850 if pending_before == 0 {
3851 pending.store(0, Ordering::Relaxed);
3852 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
3853 }
3854 let gauge_before_rejection = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
3855 for (normal, instant) in [
3856 (0, MAX_MIRROR_PENDING_ROWS + 1),
3857 (MAX_MIRROR_PENDING_ROWS / 2, MAX_MIRROR_PENDING_ROWS / 2 + 1),
3858 ] {
3859 let error = inserter
3860 .do_request(make_request(normal, instant), &table_infos, &ctx)
3861 .await
3862 .unwrap_err();
3863 assert_eq!(error.status_code(), StatusCode::InvalidArguments);
3864 assert_eq!(error.retry_hint(), RetryHint::NonRetryable);
3865 let message = error.to_string();
3866 assert!(message.contains("1000001"), "{message}");
3867 assert!(message.contains("1000000"), "{message}");
3868 assert!(message.contains("reduce the batch size"), "{message}");
3869 assert!(datanode_rx.try_recv().is_err());
3870 assert!(flownode_rx.try_recv().is_err());
3871 assert_eq!(pending.load(Ordering::Relaxed), pending_before);
3872 assert_eq!(
3873 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
3874 gauge_before_rejection
3875 );
3876 assert_eq!(
3877 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
3878 dropped_before
3879 );
3880 }
3881 if pending_before == 0 {
3882 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
3883 pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
3884 }
3885 }
3886
3887 let result = inserter
3888 .do_request(make_request(1, 0), &table_infos, &ctx)
3889 .await
3890 .unwrap();
3891 assert!(matches!(result.data, OutputData::AffectedRows(1)));
3892 assert_eq!(1, datanode_rx.try_recv().unwrap());
3893 assert!(flownode_rx.try_recv().is_err());
3894 assert_eq!(
3895 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
3896 dropped_before + 1
3897 );
3898 pending.store(0, Ordering::Relaxed);
3899 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
3900 assert_eq!(
3901 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
3902 gauge_before
3903 );
3904 }
3905
3906 #[tokio::test]
3907 async fn flow_mirror_new_counts_regions_and_peer_clones() {
3908 use common_meta::instruction::{CacheIdent, CreateFlow};
3909
3910 let peer1 = flownode_peer();
3911 let peer2 = Peer {
3912 id: 2,
3913 addr: "127.0.0.1:4002".to_string(),
3914 };
3915 let cache_backend = prepare_mocked_backend().await;
3916 let cache = Arc::new(new_table_flownode_set_cache(
3917 String::new(),
3918 Cache::new(10),
3919 cache_backend,
3920 ));
3921 cache
3922 .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
3923 flow_id: 1,
3924 source_table_ids: vec![1],
3925 partition_to_peer_mapping: vec![(0, peer1), (1, peer2)],
3926 })])
3927 .await
3928 .unwrap();
3929 let request = |region_id| RegionInsertRequest {
3930 region_id,
3931 rows: Some(Rows {
3932 schema: vec![],
3933 rows: vec![api::v1::Row { values: vec![] }; 3],
3934 }),
3935 ..Default::default()
3936 };
3937 let requests = [
3938 request(RegionId::new(1, 0).as_u64()),
3939 request(RegionId::new(1, 1).as_u64()),
3940 ];
3941 let task = FlowMirrorTask::new(&cache, requests.iter()).await.unwrap();
3942 assert_eq!(12, task.pending_rows());
3943 assert_eq!(2, task.requests.len());
3944 }
3945
3946 #[tokio::test]
3947 async fn flow_mirror_dropped_when_pending_exceeds_limit() {
3948 let _guard = lock_mirror_metrics().await;
3949 let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64));
3950 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
3951 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
3952 let num_rows = 100;
3953
3954 let task = FlowMirrorTask {
3955 requests: mirror_requests(&flownode_peer(), num_rows),
3956 };
3957 task.detach(
3959 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3960 pending.clone(),
3961 false,
3962 )
3963 .unwrap();
3964
3965 assert_eq!(
3966 num_rows as u64,
3967 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get() - dropped_before
3968 );
3969 assert_eq!(
3971 MAX_MIRROR_PENDING_ROWS as u64,
3972 pending.load(Ordering::Relaxed)
3973 );
3974 assert_eq!(
3977 gauge_before,
3978 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
3979 "a dropped batch must not change the pending gauge"
3980 );
3981
3982 let task = FlowMirrorTask {
3983 requests: mirror_requests(&flownode_peer(), num_rows),
3984 };
3985 let error = task
3986 .detach(
3987 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3988 pending.clone(),
3989 true,
3990 )
3991 .err()
3992 .unwrap();
3993 use common_error::ext::{ErrorExt, RetryHint};
3994 use common_error::status_code::StatusCode;
3995 assert_eq!(error.status_code(), StatusCode::RateLimited);
3996 assert_eq!(error.retry_hint(), RetryHint::Retryable);
3997 assert_eq!(
3998 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
3999 dropped_before + num_rows as u64,
4000 "rejected instant rows are not successful best-effort drops"
4001 );
4002 assert_eq!(
4003 MAX_MIRROR_PENDING_ROWS as u64,
4004 pending.load(Ordering::Relaxed)
4005 );
4006 assert_eq!(
4007 gauge_before,
4008 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get()
4009 );
4010 }
4011
4012 #[tokio::test]
4016 async fn flow_mirror_reserves_and_releases_per_peer_shares() {
4017 let _guard = lock_mirror_metrics().await;
4018 let per_peer_rows = 10;
4019 let total_rows = (2 * per_peer_rows) as u64;
4020 let start = MAX_MIRROR_PENDING_ROWS as u64 - total_rows;
4022 let pending = Arc::new(AtomicU64::new(start));
4023 let gate = Arc::new(tokio::sync::Semaphore::new(0));
4024 let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4025 gate: gate.clone(),
4026 affected_rows: per_peer_rows as u64,
4027 }));
4028
4029 let mut requests = mirror_requests(&flownode_peer(), per_peer_rows);
4030 requests.extend(mirror_requests(
4031 &Peer {
4032 id: 2,
4033 addr: "127.0.0.1:4002".to_string(),
4034 },
4035 per_peer_rows,
4036 ));
4037
4038 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4039 FlowMirrorTask { requests }
4040 .detach(node_manager, pending.clone(), false)
4041 .unwrap();
4042
4043 assert_eq!(
4046 MAX_MIRROR_PENDING_ROWS as u64,
4047 pending.load(Ordering::Relaxed)
4048 );
4049
4050 gate.add_permits(2);
4056 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4057 while pending.load(Ordering::Relaxed) != start
4058 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4059 {
4060 assert!(
4061 std::time::Instant::now() < deadline,
4062 "the mirror tasks must release exactly their own shares"
4063 );
4064 std::thread::sleep(std::time::Duration::from_millis(1));
4065 }
4066 }
4067
4068 #[tokio::test]
4069 async fn flow_mirror_fast_peer_releases_only_its_share_before_drop() {
4070 let _guard = lock_mirror_metrics().await;
4071 let peer1 = flownode_peer();
4072 let peer2 = Peer {
4073 id: 2,
4074 addr: "127.0.0.1:4002".to_string(),
4075 };
4076 let gate1 = Arc::new(tokio::sync::Semaphore::new(0));
4077 let gate2 = Arc::new(tokio::sync::Semaphore::new(0));
4078 let gates = Arc::new(std::sync::Mutex::new(HashMap::from_iter([
4079 (peer1.id, gate1.clone()),
4080 (peer2.id, gate2.clone()),
4081 ])));
4082 let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4083 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4084 let node_manager = Arc::new(RecordingNodeManager {
4085 datanode_dispatch: datanode_tx,
4086 flownode_dispatch: flownode_tx,
4087 gates,
4088 fail_flownode: false,
4089 });
4090 let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64 - 20));
4091 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4092 let mut requests = mirror_requests(&peer1, 10);
4093 requests.extend(mirror_requests(&peer2, 10));
4094 let task = FlowMirrorTask { requests };
4095 task.detach(node_manager.clone(), pending.clone(), false)
4096 .unwrap();
4097 let first = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4098 .await
4099 .unwrap()
4100 .unwrap();
4101 let second = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4102 .await
4103 .unwrap()
4104 .unwrap();
4105 assert_ne!(first.id, second.id);
4106 let fast_gate = if first.id == peer1.id {
4107 gate1.clone()
4108 } else {
4109 gate2.clone()
4110 };
4111 let slow_gate = if first.id == peer1.id {
4112 gate2.clone()
4113 } else {
4114 gate1.clone()
4115 };
4116 fast_gate.add_permits(1);
4117 let deadline = std::time::Instant::now() + Duration::from_secs(10);
4118 while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 10 {
4119 assert!(
4120 std::time::Instant::now() < deadline,
4121 "fast peer did not release"
4122 );
4123 std::thread::sleep(Duration::from_millis(1));
4124 }
4125 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4126 FlowMirrorTask {
4127 requests: mirror_requests(&peer1, 11),
4128 }
4129 .detach(node_manager, pending.clone(), false)
4130 .unwrap();
4131 assert_eq!(
4132 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4133 dropped_before + 11
4134 );
4135 assert!(flownode_rx.try_recv().is_err());
4136 slow_gate.add_permits(1);
4137 while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 20
4138 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4139 {
4140 assert!(
4141 std::time::Instant::now() < deadline,
4142 "slow peer did not release"
4143 );
4144 std::thread::sleep(Duration::from_millis(1));
4145 }
4146 }
4147
4148 #[tokio::test]
4149 async fn flow_mirror_failure_releases_budget_and_gauge() {
4150 let _guard = lock_mirror_metrics().await;
4151 let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4152 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4153 let node_manager = Arc::new(RecordingNodeManager {
4154 datanode_dispatch: datanode_tx,
4155 flownode_dispatch: flownode_tx,
4156 gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
4157 fail_flownode: true,
4158 });
4159 let pending = Arc::new(AtomicU64::new(0));
4160 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4161 FlowMirrorTask {
4162 requests: mirror_requests(&flownode_peer(), 5),
4163 }
4164 .detach(node_manager, pending.clone(), false)
4165 .unwrap();
4166 assert_eq!(
4167 tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4168 .await
4169 .unwrap()
4170 .unwrap()
4171 .id,
4172 flownode_peer().id
4173 );
4174 let deadline = std::time::Instant::now() + Duration::from_secs(10);
4175 while pending.load(Ordering::Relaxed) != 0
4176 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4177 {
4178 assert!(
4179 std::time::Instant::now() < deadline,
4180 "failed task leaked reservation"
4181 );
4182 std::thread::sleep(Duration::from_millis(1));
4183 }
4184 }
4185
4186 #[tokio::test]
4187 async fn flow_mirror_pending_counter_consistent() {
4188 let _guard = lock_mirror_metrics().await;
4189 let pending = Arc::new(AtomicU64::new(0));
4190 let gate = Arc::new(tokio::sync::Semaphore::new(0));
4191 let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4192 gate: gate.clone(),
4193 affected_rows: 7,
4194 }));
4195
4196 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4198 FlowMirrorTask {
4199 requests: HashMap::new(),
4200 }
4201 .detach(node_manager.clone(), pending.clone(), false)
4202 .unwrap();
4203 assert_eq!(0, pending.load(Ordering::Relaxed));
4204 assert_eq!(
4205 gauge_before,
4206 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4207 "an empty batch must not touch the pending gauge"
4208 );
4209
4210 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4211 let num_rows = 7;
4212
4213 FlowMirrorTask {
4214 requests: mirror_requests(&flownode_peer(), num_rows),
4215 }
4216 .detach(node_manager, pending.clone(), false)
4217 .unwrap();
4218
4219 assert_eq!(num_rows as u64, pending.load(Ordering::Relaxed));
4222 assert_eq!(
4223 num_rows as i64,
4224 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() - gauge_before,
4225 "the pending gauge must track the pending rows"
4226 );
4227 assert_eq!(
4228 dropped_before,
4229 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4230 "a batch within the limit must not be counted as dropped"
4231 );
4232
4233 gate.add_permits(1);
4236 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4237 while pending.load(Ordering::Relaxed) != 0
4238 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4239 {
4240 assert!(
4241 std::time::Instant::now() < deadline,
4242 "the mirror task did not release its pending reservation"
4243 );
4244 std::thread::sleep(std::time::Duration::from_millis(1));
4245 }
4246 }
4247
4248 #[test]
4249 fn mirror_logs_are_rate_limited() {
4250 let log = Arc::new(MirrorLog::new());
4251 assert_eq!(Some(1), log.claim_report(1_000));
4252 assert_eq!(None, log.claim_report(1_500));
4253 assert_eq!(None, log.claim_report(10_999));
4254 assert_eq!(Some(3), log.claim_report(11_000));
4255 assert_eq!(None, log.claim_report(11_001));
4256 assert_eq!(Some(2), log.claim_report(21_000));
4257
4258 let concurrent = Arc::new(MirrorLog::new());
4259 assert_eq!(Some(1), concurrent.claim_report(50_000));
4260 let barrier = Arc::new(std::sync::Barrier::new(8));
4261 let threads = (0..8)
4262 .map(|_| {
4263 let log = concurrent.clone();
4264 let barrier = barrier.clone();
4265 std::thread::spawn(move || {
4266 barrier.wait();
4267 log.claim_report(60_000)
4268 })
4269 })
4270 .collect::<Vec<_>>();
4271 let winners = threads
4272 .into_iter()
4273 .filter_map(|thread| thread.join().unwrap())
4274 .collect::<Vec<_>>();
4275 assert_eq!(winners.len(), 1);
4276 let residual = concurrent.events.swap(0, Ordering::Relaxed);
4277 assert_eq!(winners[0] + residual, 8);
4278 }
4279}