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