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