1use std::collections::BTreeMap;
16use std::sync::Arc;
17
18use ahash::{HashMap, HashMapExt, HashSet, HashSetExt};
19use api::v1::alter_table_expr::Kind;
20use api::v1::column_def::options_from_skipping;
21use api::v1::region::{
22 InsertRequest as RegionInsertRequest, InsertRequests as RegionInsertRequests,
23 RegionRequestHeader,
24};
25use api::v1::value::ValueData;
26use api::v1::{
27 AlterTableExpr, ColumnDataType, ColumnSchema, CreateTableExpr, InsertRequests,
28 RowInsertRequest, RowInsertRequests, Rows, SemanticType,
29};
30use catalog::CatalogManagerRef;
31use client::{OutputData, OutputMeta};
32use common_catalog::consts::{
33 DEFAULT_PRIVATE_SCHEMA_NAME, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN,
34 TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table,
35 trace_operations_table_name, trace_services_table_name,
36};
37use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME;
38use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME;
39use common_grpc_expr::util::ColumnExpr;
40use common_meta::cache::TableFlownodeSetCacheRef;
41use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME;
42use common_meta::node_manager::{AffectedRows, NodeManagerRef};
43use common_meta::peer::Peer;
44use common_meta::rpc::ddl::TriggerReason;
45use common_query::Output;
46use common_query::native_histogram::{is_native_histogram_value_type, native_histogram_value_type};
47use common_query::prelude::{greptime_timestamp, greptime_value};
48use common_telemetry::tracing_context::TracingContext;
49use common_telemetry::{debug, error, warn};
50use common_time::Timestamp;
51use common_time::timestamp::TimeUnit;
52use datatypes::schema::SkippingIndexOptions;
53use futures_util::future;
54use meter_core::data::MeterRecord;
55use meter_macros::write_meter;
56use partition::manager::PartitionRuleManagerRef;
57use session::context::QueryContextRef;
58use snafu::ResultExt;
59use snafu::prelude::*;
60use sql::partition::partition_rule_for_hexstring;
61use sql::statements::create::Partitions;
62use sql::statements::insert::Insert;
63use store_api::metric_engine_consts::{
64 LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
65};
66use store_api::mito_engine_options::{
67 APPEND_MODE_KEY, COMPACTION_TYPE, COMPACTION_TYPE_TWCS, MERGE_MODE_KEY, TTL_KEY,
68 TWCS_TIME_WINDOW,
69};
70use store_api::storage::{RegionId, TableId};
71use table::TableRef;
72use table::metadata::{TableInfo, TableInfoRef};
73use table::requests::{
74 AUTO_CREATE_TABLE_KEY, InsertRequest as TableInsertRequest, SEMANTIC_PER_TABLE_INDEX_KEY,
75 SEMANTIC_PIPELINE, TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1, TABLE_DATA_MODEL_TRACE_V2,
76 TRACE_TABLE_PARTITIONS_HINT_KEY, VALID_TABLE_OPTION_KEYS, is_semantic_option_key,
77 validate_semantic_option,
78};
79use table::table_reference::TableReference;
80
81use crate::batcher::PendingRowsBatcher;
82use crate::error::{
83 CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, FindRegionLeaderSnafu,
84 InvalidInsertRequestSnafu, JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu,
85 WriteRejectedSnafu,
86};
87use crate::expr_helper;
88use crate::region_req_factory::RegionRequestFactory;
89use crate::req_convert::common::preprocess_row_insert_requests;
90use crate::req_convert::insert::{
91 ColumnToRow, ImpureDefaultFiller, RowToRegion, StatementToRegion, TableToRegion,
92 fill_reqs_with_impure_default, rows_to_record_batch,
93};
94use crate::statement::StatementExecutor;
95
96pub struct Inserter {
97 catalog_manager: CatalogManagerRef,
98 pub(crate) partition_manager: PartitionRuleManagerRef,
99 pub(crate) node_manager: NodeManagerRef,
100 pub(crate) table_flownode_set_cache: TableFlownodeSetCacheRef,
101 auto_create_table: bool,
105 pending_rows_batcher: Option<Arc<dyn PendingRowsBatcher>>,
106}
107
108pub type InserterRef = Arc<Inserter>;
109
110#[derive(Clone)]
112pub enum AutoCreateTableType {
113 Logical(String),
115 Physical,
117 Log,
119 LastNonNull,
121 Trace { alter_existing: bool },
123}
124
125impl AutoCreateTableType {
126 pub fn as_str(&self) -> &'static str {
127 match self {
128 AutoCreateTableType::Logical(_) => "logical",
129 AutoCreateTableType::Physical => "physical",
130 AutoCreateTableType::Log => "log",
131 AutoCreateTableType::LastNonNull => "last_non_null",
132 AutoCreateTableType::Trace { .. } => "trace",
133 }
134 }
135
136 fn alter_existing(&self) -> bool {
137 !matches!(
138 self,
139 Self::Trace {
140 alter_existing: false
141 }
142 )
143 }
144}
145
146#[derive(Clone)]
153pub struct InstantAndNormalInsertRequests {
154 pub normal_requests: RegionInsertRequests,
156 pub instant_requests: RegionInsertRequests,
159}
160
161impl Inserter {
162 pub async fn can_batch_metric_rows(
166 &self,
167 requests: &RowInsertRequests,
168 ctx: &QueryContextRef,
169 physical_table: &str,
170 ) -> Result<bool> {
171 if self.auto_create_disabled_reason(ctx)?.is_some() || ctx.extension(TTL_KEY).is_some() {
172 return Ok(false);
173 }
174 for request in &requests.inserts {
175 if request.rows.as_ref().is_some_and(|rows| {
180 rows.schema.iter().any(|column| {
181 column.datatype_extension.is_some()
182 || !matches!(
183 ColumnDataType::try_from(column.datatype),
184 Ok(ColumnDataType::TimestampSecond
185 | ColumnDataType::TimestampMillisecond
186 | ColumnDataType::TimestampMicrosecond
187 | ColumnDataType::TimestampNanosecond
188 | ColumnDataType::Float64
189 | ColumnDataType::String)
190 )
191 })
192 }) {
193 return Ok(false);
194 }
195
196 let Some(table) = self
197 .get_table(
198 ctx.current_catalog(),
199 &ctx.current_schema(),
200 &request.table_name,
201 )
202 .await?
203 else {
204 continue;
205 };
206 let info = table.table_info();
207 if info.meta.engine != METRIC_ENGINE_NAME
208 || info.is_ttl_instant_table()
209 || info
210 .meta
211 .options
212 .extra_options
213 .get(LOGICAL_TABLE_METADATA_KEY)
214 .map(String::as_str)
215 != Some(physical_table)
216 || info
217 .meta
218 .schema
219 .column_schemas()
220 .iter()
221 .any(|column| column.default_constraint().is_some())
222 {
223 return Ok(false);
224 }
225 if info
228 .meta
229 .primary_key_indices
230 .iter()
231 .any(|&index| !info.meta.schema.column_schemas()[index].is_nullable())
232 {
233 return Ok(false);
234 }
235 match self.table_flownode_set_cache.get(info.table_id()).await {
238 Ok(None) => {}
239 Ok(Some(flows)) if flows.is_empty() => {}
240 _ => return Ok(false),
241 }
242 }
243 Ok(true)
244 }
245
246 pub async fn meter_row_inserts(
249 requests: &mut RowInsertRequests,
250 ctx: &QueryContextRef,
251 ) -> Result<u64> {
252 let metered = InstantAndNormalInsertRequests {
253 normal_requests: RegionInsertRequests {
254 requests: requests
255 .inserts
256 .iter_mut()
257 .map(|request| RegionInsertRequest {
258 rows: request.rows.take(),
259 ..Default::default()
260 })
261 .collect(),
262 },
263 instant_requests: RegionInsertRequests::default(),
264 };
265 let cost = write_meter!(
266 ctx.current_catalog(),
267 ctx.current_schema(),
268 metered,
269 ctx.write_rows_to_admit(
270 ctx.current_catalog(),
271 &ctx.current_schema(),
272 count_insert_rows(&metered)?
273 ),
274 ctx.channel() as u8
275 )
276 .await
277 .context(WriteRejectedSnafu);
278 for (request, region) in requests
279 .inserts
280 .iter_mut()
281 .zip(metered.normal_requests.requests)
282 {
283 request.rows = region.rows;
284 }
285 cost
286 }
287
288 pub fn new(
289 catalog_manager: CatalogManagerRef,
290 partition_manager: PartitionRuleManagerRef,
291 node_manager: NodeManagerRef,
292 table_flownode_set_cache: TableFlownodeSetCacheRef,
293 auto_create_table: bool,
294 ) -> Self {
295 Self {
296 catalog_manager,
297 partition_manager,
298 node_manager,
299 table_flownode_set_cache,
300 auto_create_table,
301 pending_rows_batcher: None,
302 }
303 }
304
305 pub fn with_pending_rows_batcher(
307 mut self,
308 batcher: Option<Arc<dyn PendingRowsBatcher>>,
309 ) -> Self {
310 self.pending_rows_batcher = batcher;
311 self
312 }
313
314 pub async fn handle_column_inserts(
315 &self,
316 requests: InsertRequests,
317 ctx: QueryContextRef,
318 statement_executor: &StatementExecutor,
319 ) -> Result<Output> {
320 let row_inserts = ColumnToRow::convert(requests)?;
321 self.handle_row_inserts(row_inserts, ctx, statement_executor, false, false)
322 .await
323 }
324
325 pub async fn handle_row_inserts(
327 &self,
328 mut requests: RowInsertRequests,
329 ctx: QueryContextRef,
330 statement_executor: &StatementExecutor,
331 accommodate_existing_schema: bool,
332 is_single_value: bool,
333 ) -> Result<Output> {
334 preprocess_row_insert_requests(&mut requests.inserts)?;
335 self.handle_row_inserts_with_create_type(
336 requests,
337 ctx,
338 statement_executor,
339 AutoCreateTableType::Physical,
340 accommodate_existing_schema,
341 is_single_value,
342 )
343 .await
344 }
345
346 pub async fn handle_log_inserts(
348 &self,
349 requests: RowInsertRequests,
350 ctx: QueryContextRef,
351 statement_executor: &StatementExecutor,
352 ) -> Result<Output> {
353 self.handle_row_inserts_with_create_type(
354 requests,
355 ctx,
356 statement_executor,
357 AutoCreateTableType::Log,
358 false,
359 false,
360 )
361 .await
362 }
363
364 pub async fn handle_trace_inserts(
365 &self,
366 requests: RowInsertRequests,
367 ctx: QueryContextRef,
368 statement_executor: &StatementExecutor,
369 ) -> Result<Output> {
370 self.handle_row_inserts_with_create_type(
371 requests,
372 ctx,
373 statement_executor,
374 AutoCreateTableType::Trace {
375 alter_existing: true,
376 },
377 false,
378 false,
379 )
380 .await
381 }
382
383 pub async fn handle_last_non_null_inserts(
385 &self,
386 requests: RowInsertRequests,
387 ctx: QueryContextRef,
388 statement_executor: &StatementExecutor,
389 accommodate_existing_schema: bool,
390 is_single_value: bool,
391 ) -> Result<Output> {
392 self.handle_row_inserts_with_create_type(
393 requests,
394 ctx,
395 statement_executor,
396 AutoCreateTableType::LastNonNull,
397 accommodate_existing_schema,
398 is_single_value,
399 )
400 .await
401 }
402
403 async fn handle_row_inserts_with_create_type(
405 &self,
406 mut requests: RowInsertRequests,
407 ctx: QueryContextRef,
408 statement_executor: &StatementExecutor,
409 create_type: AutoCreateTableType,
410 accommodate_existing_schema: bool,
411 is_single_value: bool,
412 ) -> Result<Output> {
413 let skip_wal = ctx.skip_wal();
414
415 let batcher = self
416 .pending_rows_batcher
417 .as_ref()
418 .filter(|_| ctx.batching_enabled());
419
420 requests.inserts.retain(|req| {
422 req.rows
423 .as_ref()
424 .map(|r| !r.rows.is_empty())
425 .unwrap_or_default()
426 });
427 validate_column_count_match(&requests)?;
428
429 let CreateAlterTableResult {
430 instant_table_ids,
431 table_infos,
432 } = self
433 .create_or_alter_tables_on_demand(
434 &mut requests,
435 &ctx,
436 create_type,
437 statement_executor,
438 accommodate_existing_schema,
439 is_single_value,
440 None,
441 )
442 .await?;
443
444 if let Some(batcher) = batcher
447 && instant_table_ids.is_empty()
448 && table_infos
449 .values()
450 .all(|info| info.meta.engine == default_engine())
451 {
452 return self
453 .submit_pending_rows(requests, table_infos, ctx, batcher)
454 .await;
455 }
456
457 let name_to_info = table_infos
458 .values()
459 .map(|info| (info.name.clone(), info.clone()))
460 .collect::<HashMap<_, _>>();
461 let inserts = RowToRegion::new(
462 name_to_info,
463 instant_table_ids,
464 self.partition_manager.as_ref(),
465 )
466 .convert(requests, skip_wal)
467 .await?;
468
469 self.do_request(inserts, &table_infos, &ctx).await
470 }
471
472 async fn submit_pending_rows(
473 &self,
474 mut requests: RowInsertRequests,
475 table_infos: HashMap<TableId, Arc<TableInfo>>,
476 ctx: QueryContextRef,
477 batcher: &Arc<dyn PendingRowsBatcher>,
478 ) -> Result<Output> {
479 requests.inserts.retain(|request| {
482 request
483 .rows
484 .as_ref()
485 .is_some_and(|rows| !rows.rows.is_empty())
486 });
487 let by_name = table_infos
488 .values()
489 .map(|info| (info.name.as_str(), info))
490 .collect::<HashMap<_, _>>();
491 let mut prepared = Vec::with_capacity(requests.inserts.len());
492 for request in &mut requests.inserts {
493 let table_info =
494 by_name
495 .get(request.table_name.as_str())
496 .context(TableNotFoundSnafu {
497 table_name: &request.table_name,
498 })?;
499 let Some(rows) = &mut request.rows else {
500 continue;
501 };
502 ImpureDefaultFiller::new((*table_info).clone())?.fill_rows(rows);
503 let batch = rows_to_record_batch(rows, table_info)?;
504 prepared.push(((*table_info).clone(), batch));
505 }
506 let metered = InstantAndNormalInsertRequests {
509 normal_requests: RegionInsertRequests {
510 requests: requests
511 .inserts
512 .into_iter()
513 .map(|request| RegionInsertRequest {
514 rows: request.rows,
515 ..Default::default()
516 })
517 .collect(),
518 },
519 instant_requests: RegionInsertRequests::default(),
520 };
521 let table_info = table_infos.values().next();
522 let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str());
523 let schema =
524 table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone());
525 let write_cost = write_meter!(
526 catalog,
527 &schema,
528 metered,
529 ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&metered)?),
530 ctx.channel() as u8
531 )
532 .await
533 .context(WriteRejectedSnafu)?;
534 prepared.retain(|(_, batch)| batch.num_rows() != 0);
535 let results = if prepared.is_empty() {
536 Vec::new()
537 } else {
538 let permit = batcher.acquire().await?;
540 let submissions = prepared.into_iter().map(|(info, batch)| {
541 let mut target_ctx = ctx.fork();
543 target_ctx.set_current_catalog(&info.catalog_name);
544 target_ctx.set_current_schema(&info.schema_name);
545 batcher.submit(info, batch, Arc::new(target_ctx), permit.clone())
546 });
547 future::join_all(submissions).await
549 };
550 let affected_rows = results.into_iter().sum::<Result<usize>>()?;
551 Ok(Output::new(
552 OutputData::AffectedRows(affected_rows),
553 OutputMeta::new_with_cost(write_cost as _),
554 ))
555 }
556
557 pub async fn handle_metric_row_inserts(
559 &self,
560 mut requests: RowInsertRequests,
561 ctx: QueryContextRef,
562 statement_executor: &StatementExecutor,
563 physical_table: String,
564 ) -> Result<Output> {
565 let skip_wal = ctx.skip_wal();
566
567 requests.inserts.retain(|req| {
569 req.rows
570 .as_ref()
571 .map(|r| !r.rows.is_empty())
572 .unwrap_or_default()
573 });
574 validate_column_count_match(&requests)?;
575
576 let physical_table_ref = self
578 .create_physical_table_on_demand(&ctx, physical_table.clone(), statement_executor)
579 .await?;
580
581 let CreateAlterTableResult {
593 instant_table_ids,
594 table_infos,
595 } = self
596 .create_or_alter_tables_on_demand(
597 &mut requests,
598 &ctx,
599 AutoCreateTableType::Logical(physical_table.clone()),
600 statement_executor,
601 true,
602 true,
603 table_time_index_unit(&physical_table_ref),
604 )
605 .await?;
606 let name_to_info = table_infos
607 .values()
608 .map(|info| (info.name.clone(), info.clone()))
609 .collect::<HashMap<_, _>>();
610 let inserts = RowToRegion::new(name_to_info, instant_table_ids, &self.partition_manager)
611 .convert(requests, skip_wal)
612 .await?;
613
614 self.do_request(inserts, &table_infos, &ctx).await
615 }
616
617 fn table_batcher(
618 &self,
619 table_info: &TableInfoRef,
620 ctx: &QueryContextRef,
621 ) -> Option<&Arc<dyn PendingRowsBatcher>> {
622 self.pending_rows_batcher.as_ref().filter(|_| {
623 ctx.batching_enabled()
624 && !table_info.is_ttl_instant_table()
625 && table_info.meta.engine == default_engine()
626 })
627 }
628
629 async fn submit_table_rows(
630 &self,
631 rows: Rows,
632 table_info: TableInfoRef,
633 ctx: QueryContextRef,
634 batcher: &Arc<dyn PendingRowsBatcher>,
635 ) -> Result<Output> {
636 let requests = RowInsertRequests {
637 inserts: vec![RowInsertRequest {
638 table_name: table_info.name.clone(),
639 rows: Some(rows),
640 }],
641 };
642 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info)]);
643 self.submit_pending_rows(requests, table_infos, ctx, batcher)
644 .await
645 }
646
647 pub async fn handle_table_insert(
648 &self,
649 request: TableInsertRequest,
650 ctx: QueryContextRef,
651 ) -> Result<Output> {
652 let catalog = request.catalog_name.as_str();
653 let schema = request.schema_name.as_str();
654 let table_name = request.table_name.as_str();
655 let table = self.get_table(catalog, schema, table_name).await?;
656 let table = table.with_context(|| TableNotFoundSnafu {
657 table_name: common_catalog::format_full_table_name(catalog, schema, table_name),
658 })?;
659 let table_info = table.table_info();
660
661 let converter = TableToRegion::new(&table_info, &self.partition_manager);
662 let skip_wal = request.skip_wal;
663 let rows = converter.prepare(request)?;
664 if let Some(batcher) = self.table_batcher(&table_info, &ctx) {
665 return self.submit_table_rows(rows, table_info, ctx, batcher).await;
666 }
667 let inserts = converter.partition(rows, skip_wal).await?;
668
669 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]);
670
671 self.do_request(inserts, &table_infos, &ctx).await
672 }
673
674 pub async fn handle_statement_insert(
675 &self,
676 insert: &Insert,
677 ctx: &QueryContextRef,
678 ) -> Result<Output> {
679 let converter =
680 StatementToRegion::new(self.catalog_manager.as_ref(), &self.partition_manager, ctx);
681 let (rows, table_info) = converter.prepare(insert, ctx).await?;
682 if let Some(batcher) = self.table_batcher(&table_info, ctx) {
683 return self
684 .submit_table_rows(rows, table_info, ctx.clone(), batcher)
685 .await;
686 }
687 let inserts = converter.partition(rows, table_info.clone(), ctx).await?;
688
689 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]);
690
691 self.do_request(inserts, &table_infos, ctx).await
692 }
693}
694
695pub async fn admit_write(rows: u64, ctx: &QueryContextRef) -> Result<QueryContextRef> {
698 write_meter!(MeterRecord::new(
701 ctx.current_catalog().to_string(),
702 ctx.current_schema(),
703 0,
704 ctx.write_rows_to_admit(ctx.current_catalog(), &ctx.current_schema(), rows),
705 ctx.channel() as u8,
706 ))
707 .await
708 .context(WriteRejectedSnafu)?;
709 Ok(Arc::new(ctx.with_write_admission()))
710}
711
712pub async fn admit_row_insert_batches(
715 batches: &mut [(QueryContextRef, RowInsertRequests)],
716) -> Result<()> {
717 let mut totals = BTreeMap::<_, (QueryContextRef, u64)>::new();
718 for (ctx, requests) in batches.iter() {
719 let catalog = ctx.current_catalog();
720 let schema = ctx.current_schema();
721 if ctx.write_rows_to_admit(catalog, &schema, 1) == 0 {
722 continue;
723 }
724 let (_, total) = totals
725 .entry((catalog.to_string(), schema.clone()))
726 .or_insert_with(|| (ctx.clone(), 0));
727 for rows in requests.inserts.iter().filter_map(|r| r.rows.as_ref()) {
728 *total =
729 total
730 .checked_add(rows.rows.len() as u64)
731 .context(InvalidInsertRequestSnafu {
732 reason: "Insert row count exceeds u64::MAX",
733 })?;
734 }
735 }
736 for (ctx, rows) in totals.values() {
737 admit_write(*rows, ctx).await?;
738 }
739 for (ctx, _) in batches {
740 *ctx = Arc::new(ctx.with_write_admission());
741 }
742 Ok(())
743}
744
745fn count_insert_rows(requests: &InstantAndNormalInsertRequests) -> Result<u64> {
746 requests
747 .normal_requests
748 .requests
749 .iter()
750 .chain(&requests.instant_requests.requests)
751 .filter_map(|request| request.rows.as_ref())
752 .try_fold(0u64, |total, rows| {
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
761impl Inserter {
762 async fn do_request(
763 &self,
764 requests: InstantAndNormalInsertRequests,
765 table_infos: &HashMap<TableId, Arc<TableInfo>>,
766 ctx: &QueryContextRef,
767 ) -> Result<Output> {
768 let requests = fill_reqs_with_impure_default(table_infos, requests)?;
770
771 let table_info = table_infos.values().next();
774 let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str());
775 let schema =
776 table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone());
777 let write_cost = write_meter!(
778 catalog,
779 schema.clone(),
780 requests,
781 ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&requests)?),
782 ctx.channel() as u8
783 )
784 .await
785 .context(WriteRejectedSnafu)?;
786 let request_factory = RegionRequestFactory::new(RegionRequestHeader {
787 tracing_context: TracingContext::from_current_span().to_w3c(),
788 dbname: ctx.get_db_string(),
789 ..Default::default()
790 });
791
792 let InstantAndNormalInsertRequests {
793 normal_requests,
794 instant_requests,
795 } = requests;
796
797 let flow_mirror_task = FlowMirrorTask::new(
799 &self.table_flownode_set_cache,
800 normal_requests
801 .requests
802 .iter()
803 .chain(instant_requests.requests.iter()),
804 )
805 .await?;
806 flow_mirror_task.detach(self.node_manager.clone())?;
807
808 let write_tasks = self
810 .group_requests_by_peer(normal_requests)
811 .await?
812 .into_iter()
813 .map(|(peer, inserts)| {
814 let node_manager = self.node_manager.clone();
815 let request = request_factory.build_insert(inserts);
816 common_runtime::spawn_global(async move {
817 node_manager
818 .datanode(&peer)
819 .await
820 .handle(request)
821 .await
822 .context(RequestInsertsSnafu)
823 })
824 });
825 let results = future::try_join_all(write_tasks)
826 .await
827 .context(JoinTaskSnafu)?;
828 let affected_rows = results
829 .into_iter()
830 .map(|resp| resp.map(|r| r.affected_rows))
831 .sum::<Result<AffectedRows>>()?;
832 crate::metrics::DIST_INGEST_ROW_COUNT
833 .with_label_values(&[ctx.get_db_string().as_str()])
834 .inc_by(affected_rows as u64);
835 Ok(Output::new(
836 OutputData::AffectedRows(affected_rows),
837 OutputMeta::new_with_cost(write_cost as _),
838 ))
839 }
840
841 async fn group_requests_by_peer(
842 &self,
843 requests: RegionInsertRequests,
844 ) -> Result<HashMap<Peer, RegionInsertRequests>> {
845 let mut requests_per_region: HashMap<RegionId, RegionInsertRequests> = HashMap::new();
848 for req in requests.requests {
849 let region_id = RegionId::from_u64(req.region_id);
850 requests_per_region
851 .entry(region_id)
852 .or_default()
853 .requests
854 .push(req);
855 }
856
857 let mut inserts: HashMap<Peer, RegionInsertRequests> = HashMap::new();
858
859 for (region_id, reqs) in requests_per_region {
860 let peer = self
861 .partition_manager
862 .find_region_leader(region_id)
863 .await
864 .context(FindRegionLeaderSnafu)?;
865 inserts
866 .entry(peer)
867 .or_default()
868 .requests
869 .extend(reqs.requests);
870 }
871
872 Ok(inserts)
873 }
874
875 fn auto_create_disabled_reason(&self, ctx: &QueryContextRef) -> Result<Option<&'static str>> {
879 let auto_create_table_hint = ctx
880 .extension(AUTO_CREATE_TABLE_KEY)
881 .map(|v| v.parse::<bool>())
882 .transpose()
883 .map_err(|_| {
884 InvalidInsertRequestSnafu {
885 reason: "`auto_create_table` hint must be a boolean",
886 }
887 .build()
888 })?
889 .unwrap_or(true);
890 Ok(if !self.auto_create_table {
891 Some("auto-create table is disabled by frontend config")
892 } else if !auto_create_table_hint {
893 Some("`auto_create_table` hint is disabled")
894 } else {
895 None
896 })
897 }
898
899 fn is_auto_create_exempt_private_table(schema: &str, table: &str) -> bool {
902 schema == DEFAULT_PRIVATE_SCHEMA_NAME
903 && matches!(
904 table,
905 DEFAULT_EVENTS_TABLE_NAME | SLOW_QUERY_TABLE_NAME | REGION_STATS_HISTORY_TABLE_NAME
906 )
907 }
908
909 pub async fn ensure_trace_table_on_demand(
914 &self,
915 table_name: &str,
916 request_schema: Vec<ColumnSchema>,
917 alter_existing: bool,
918 ctx: &QueryContextRef,
919 statement_executor: &StatementExecutor,
920 ) -> Result<()> {
921 let mut requests = RowInsertRequests {
922 inserts: vec![RowInsertRequest {
923 table_name: table_name.to_string(),
924 rows: Some(api::v1::Rows {
925 schema: request_schema,
926 rows: Vec::new(),
927 }),
928 }],
929 };
930 self.create_or_alter_tables_on_demand(
931 &mut requests,
932 ctx,
933 AutoCreateTableType::Trace { alter_existing },
934 statement_executor,
935 false,
936 false,
937 None,
938 )
939 .await?;
940 Ok(())
941 }
942
943 #[allow(clippy::too_many_arguments)]
962 async fn create_or_alter_tables_on_demand(
963 &self,
964 requests: &mut RowInsertRequests,
965 ctx: &QueryContextRef,
966 auto_create_table_type: AutoCreateTableType,
967 statement_executor: &StatementExecutor,
968 accommodate_existing_schema: bool,
969 is_single_value: bool,
970 align_time_index_unit: Option<TimeUnit>,
971 ) -> Result<CreateAlterTableResult> {
972 let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND
973 .with_label_values(&[auto_create_table_type.as_str()])
974 .start_timer();
975 let catalog = ctx.current_catalog();
976 let schema = ctx.current_schema();
977
978 let auto_create_disabled_reason = self.auto_create_disabled_reason(ctx)?;
979 let has_auto_create_exempt_table = auto_create_disabled_reason.is_some()
981 && requests
982 .inserts
983 .iter()
984 .any(|req| Self::is_auto_create_exempt_private_table(&schema, &req.table_name));
985 let mut table_infos = HashMap::new();
986 if let Some(disabled_reason) = auto_create_disabled_reason
988 && !has_auto_create_exempt_table
989 {
990 let mut instant_table_ids = HashSet::new();
991 for req in &mut requests.inserts {
992 let table = match self.get_table(catalog, &schema, &req.table_name).await? {
993 Some(table) => table,
994 None if is_ddl_reserved_table(&schema, &req.table_name) => {
997 statement_executor
998 .create_declared_relationships_table(catalog, ctx.clone())
999 .await?
1000 }
1001 None => {
1002 return InvalidInsertRequestSnafu {
1003 reason: format!(
1004 "Table `{}` does not exist, and {}",
1005 req.table_name, disabled_reason
1006 ),
1007 }
1008 .fail();
1009 }
1010 };
1011 if align_time_index_unit.is_some()
1015 && let Some(rows) = req.rows.as_mut()
1016 && let Some(target_unit) = table_time_index_unit(&table)
1017 {
1018 convert_rows_time_unit(rows, target_unit)?;
1019 }
1020 let table_info = table.table_info();
1021 if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1022 validate_trace_table_model(&table_info, ctx)?;
1023 }
1024 if table_info.is_ttl_instant_table() {
1025 instant_table_ids.insert(table_info.table_id());
1026 }
1027 table_infos.insert(table_info.table_id(), table.table_info());
1028 }
1029 let ret = CreateAlterTableResult {
1030 instant_table_ids,
1031 table_infos,
1032 };
1033 return Ok(ret);
1034 }
1035
1036 let mut create_tables = vec![];
1037 let mut alter_tables = vec![];
1038 let mut need_refresh_table_infos = HashSet::new();
1039 let mut instant_table_ids = HashSet::new();
1040 let mut per_table_semantics: Option<Option<PerTableSemanticIndex>> = None;
1041
1042 for req in &mut requests.inserts {
1043 let auto_create_allowed = auto_create_disabled_reason.is_none()
1045 || Self::is_auto_create_exempt_private_table(&schema, &req.table_name);
1046 match self.get_table(catalog, &schema, &req.table_name).await? {
1047 Some(table) => {
1048 let table_info = table.table_info();
1049 if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1050 validate_trace_table_model(&table_info, ctx)?;
1051 }
1052 if table_info.is_ttl_instant_table() {
1053 instant_table_ids.insert(table_info.table_id());
1054 }
1055 if align_time_index_unit.is_some()
1059 && let Some(rows) = req.rows.as_mut()
1060 && let Some(target_unit) = table_time_index_unit(&table)
1061 {
1062 convert_rows_time_unit(rows, target_unit)?;
1063 }
1064 if auto_create_allowed
1065 && let Some(alter_expr) = self.get_alter_table_expr_on_demand(
1066 req,
1067 &table,
1068 ctx,
1069 accommodate_existing_schema,
1070 is_single_value,
1071 auto_create_table_type.alter_existing(),
1072 )?
1073 {
1074 alter_tables.push(alter_expr);
1075 need_refresh_table_infos.insert((
1076 catalog.to_string(),
1077 schema.clone(),
1078 req.table_name.clone(),
1079 ));
1080 } else {
1081 table_infos.insert(table_info.table_id(), table.table_info());
1082 }
1083 }
1084 None if is_ddl_reserved_table(&schema, &req.table_name) => {
1088 let table = statement_executor
1089 .create_declared_relationships_table(catalog, ctx.clone())
1090 .await?;
1091 let table_info = table.table_info();
1092 if table_info.is_ttl_instant_table() {
1093 instant_table_ids.insert(table_info.table_id());
1094 }
1095 table_infos.insert(table_info.table_id(), table_info);
1096 }
1097 None if !auto_create_allowed
1098 && let Some(disabled_reason) = auto_create_disabled_reason =>
1099 {
1100 return InvalidInsertRequestSnafu {
1101 reason: format!(
1102 "Table `{}` does not exist, and {}",
1103 req.table_name, disabled_reason,
1104 ),
1105 }
1106 .fail();
1107 }
1108 None => {
1109 if let Some(physical_unit) = align_time_index_unit
1113 && let Some(rows) = req.rows.as_mut()
1114 {
1115 convert_rows_time_unit(rows, physical_unit)?;
1116 }
1117 let semantic_index = per_table_semantics
1118 .get_or_insert_with(|| parse_per_table_semantic_index(ctx))
1119 .as_ref();
1120 let create_expr = self.get_create_table_expr_on_demand(
1121 req,
1122 &auto_create_table_type,
1123 ctx,
1124 semantic_index,
1125 )?;
1126 create_tables.push(create_expr);
1127 }
1128 }
1129 }
1130
1131 match auto_create_table_type {
1132 AutoCreateTableType::Logical(_) => {
1133 if !create_tables.is_empty() {
1134 let tables = self
1136 .create_logical_tables(create_tables, ctx, statement_executor)
1137 .await?;
1138
1139 for table in tables {
1140 let table_info = table.table_info();
1141 if table_info.is_ttl_instant_table() {
1142 instant_table_ids.insert(table_info.table_id());
1143 }
1144 table_infos.insert(table_info.table_id(), table.table_info());
1145 }
1146 }
1147 if !alter_tables.is_empty() {
1148 statement_executor
1150 .alter_logical_tables(alter_tables, ctx.clone(), TriggerReason::AutoAlter)
1151 .await?;
1152 }
1153 }
1154 AutoCreateTableType::Physical
1155 | AutoCreateTableType::Log
1156 | AutoCreateTableType::LastNonNull => {
1157 for create_table in create_tables {
1160 let table = self
1161 .create_physical_table(create_table, None, ctx, statement_executor)
1162 .await?;
1163 let table_info = table.table_info();
1164 if table_info.is_ttl_instant_table() {
1165 instant_table_ids.insert(table_info.table_id());
1166 }
1167 table_infos.insert(table_info.table_id(), table.table_info());
1168 }
1169 for alter_expr in alter_tables.into_iter() {
1170 statement_executor
1171 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1172 .await?;
1173 }
1174 }
1175
1176 AutoCreateTableType::Trace { .. } => {
1177 let trace_table_name = ctx
1178 .extension(TRACE_TABLE_NAME_SESSION_KEY)
1179 .unwrap_or(TRACE_TABLE_NAME);
1180
1181 let trace_table_partitions = if let Some(trace_table_partitions) =
1182 ctx.extension(TRACE_TABLE_PARTITIONS_HINT_KEY)
1183 {
1184 let p = trace_table_partitions.parse::<u32>().map_err(|_| {
1185 InvalidInsertRequestSnafu {
1186 reason: format!(
1187 "Failed to parse trace_table_partitions: {}",
1188 trace_table_partitions
1189 ),
1190 }
1191 .build()
1192 })?;
1193 Some(p)
1194 } else {
1195 None
1196 };
1197
1198 for mut create_table in create_tables {
1201 if create_table.table_name == trace_services_table_name(trace_table_name)
1202 || create_table.table_name == trace_operations_table_name(trace_table_name)
1203 {
1204 create_table
1206 .table_options
1207 .insert(APPEND_MODE_KEY.to_string(), "false".to_string());
1208 create_table.table_options.remove(TTL_KEY);
1210
1211 let table = self
1212 .create_physical_table(create_table, None, ctx, statement_executor)
1213 .await?;
1214 let table_info = table.table_info();
1215 if table_info.is_ttl_instant_table() {
1216 instant_table_ids.insert(table_info.table_id());
1217 }
1218 table_infos.insert(table_info.table_id(), table.table_info());
1219 } else {
1220 let partitions = if matches!(trace_table_partitions, Some(0) | Some(1)) {
1223 None
1225 } else {
1226 let p = partition_rule_for_hexstring(
1227 TRACE_ID_COLUMN,
1228 trace_table_partitions,
1229 )
1230 .context(CreatePartitionRulesSnafu)?;
1231 Some(p)
1232 };
1233
1234 let index_columns =
1239 [TRACE_ID_COLUMN, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN];
1240 for index_column in index_columns {
1241 if let Some(col) = create_table
1242 .column_defs
1243 .iter_mut()
1244 .find(|c| c.name == index_column)
1245 {
1246 col.options =
1247 options_from_skipping(&SkippingIndexOptions::default())
1248 .context(ColumnOptionsSnafu)?;
1249 } else {
1250 warn!(
1251 "Column {} not found when creating index for trace table: {}.",
1252 index_column, create_table.table_name
1253 );
1254 }
1255 }
1256
1257 create_table.table_options.insert(
1259 TABLE_DATA_MODEL.to_string(),
1260 ctx.extension(SEMANTIC_PIPELINE)
1261 .unwrap_or(TABLE_DATA_MODEL_TRACE_V1)
1262 .to_string(),
1263 );
1264
1265 let table = self
1266 .create_physical_table(
1267 create_table,
1268 partitions,
1269 ctx,
1270 statement_executor,
1271 )
1272 .await?;
1273 let table_info = table.table_info();
1274 if table_info.is_ttl_instant_table() {
1275 instant_table_ids.insert(table_info.table_id());
1276 }
1277 table_infos.insert(table_info.table_id(), table.table_info());
1278 }
1279 }
1280 for alter_expr in alter_tables.into_iter() {
1281 statement_executor
1282 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1283 .await?;
1284 }
1285 }
1286 }
1287
1288 for (catalog, schema, table_name) in need_refresh_table_infos {
1290 let table = self
1291 .get_table(&catalog, &schema, &table_name)
1292 .await?
1293 .context(TableNotFoundSnafu {
1294 table_name: common_catalog::format_full_table_name(
1295 &catalog,
1296 &schema,
1297 &table_name,
1298 ),
1299 })?;
1300 let table_info = table.table_info();
1301 table_infos.insert(table_info.table_id(), table.table_info());
1302 }
1303
1304 Ok(CreateAlterTableResult {
1305 instant_table_ids,
1306 table_infos,
1307 })
1308 }
1309
1310 async fn create_physical_table_on_demand(
1311 &self,
1312 ctx: &QueryContextRef,
1313 physical_table: String,
1314 statement_executor: &StatementExecutor,
1315 ) -> Result<TableRef> {
1316 let catalog_name = ctx.current_catalog();
1317 let schema_name = ctx.current_schema();
1318
1319 if let Some(table) = self
1321 .get_table(catalog_name, &schema_name, &physical_table)
1322 .await?
1323 {
1324 return Ok(table);
1325 }
1326
1327 if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? {
1329 return InvalidInsertRequestSnafu {
1330 reason: format!(
1331 "Physical table `{physical_table}` does not exist, and {disabled_reason}"
1332 ),
1333 }
1334 .fail();
1335 }
1336
1337 let table_reference = TableReference::full(catalog_name, &schema_name, &physical_table);
1338 debug!("Ensuring physical metric table `{table_reference}` exists for insert");
1339
1340 let default_schema = vec![
1342 ColumnSchema {
1343 column_name: greptime_timestamp().to_string(),
1344 datatype: ColumnDataType::TimestampMillisecond as _,
1345 semantic_type: SemanticType::Timestamp as _,
1346 datatype_extension: None,
1347 options: None,
1348 },
1349 ColumnSchema {
1350 column_name: greptime_value().to_string(),
1351 datatype: ColumnDataType::Float64 as _,
1352 semantic_type: SemanticType::Field as _,
1353 datatype_extension: None,
1354 options: None,
1355 },
1356 ];
1357 let create_table_expr =
1358 &mut build_create_table_expr(&table_reference, &default_schema, default_engine())?;
1359
1360 create_table_expr.engine = METRIC_ENGINE_NAME.to_string();
1361 create_table_expr
1362 .table_options
1363 .insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), "true".to_string());
1364
1365 let res = statement_executor
1367 .create_table_inner(
1368 create_table_expr,
1369 None,
1370 ctx.clone(),
1371 TriggerReason::AutoCreate,
1372 )
1373 .await;
1374
1375 match res {
1376 Ok(table) => Ok(table),
1377 Err(err) => {
1378 error!(err; "Failed to create table {table_reference}");
1379 Err(err)
1380 }
1381 }
1382 }
1383
1384 async fn get_table(
1385 &self,
1386 catalog: &str,
1387 schema: &str,
1388 table: &str,
1389 ) -> Result<Option<TableRef>> {
1390 self.catalog_manager
1391 .table(catalog, schema, table, None)
1392 .await
1393 .context(CatalogSnafu)
1394 }
1395
1396 fn get_create_table_expr_on_demand(
1397 &self,
1398 req: &RowInsertRequest,
1399 create_type: &AutoCreateTableType,
1400 ctx: &QueryContextRef,
1401 semantic_index: Option<&PerTableSemanticIndex>,
1402 ) -> Result<CreateTableExpr> {
1403 let schema = ctx.current_schema();
1404 let mut table_options = std::collections::HashMap::with_capacity(4);
1405 fill_table_options_for_create(&mut table_options, create_type, ctx);
1406 apply_per_table_semantic_options(
1407 &mut table_options,
1408 semantic_index,
1409 ctx.current_schema().as_str(),
1410 &req.table_name,
1411 );
1412
1413 let engine_name = if let AutoCreateTableType::Logical(_) = create_type {
1414 METRIC_ENGINE_NAME
1416 } else {
1417 default_engine()
1418 };
1419
1420 let table_ref = TableReference::full(ctx.current_catalog(), &schema, &req.table_name);
1421 let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1423 let mut create_table_expr =
1424 build_create_table_expr(&table_ref, request_schema, engine_name)?;
1425
1426 if ctx.extension(SPLUNK_PK_METADATA_ORDER_KEY).is_some() {
1428 reorder_splunk_primary_keys(&mut create_table_expr.primary_keys);
1429 }
1430
1431 debug!("Ensuring table `{table_ref}` exists for insert");
1432 create_table_expr.table_options.extend(table_options);
1433 Ok(create_table_expr)
1434 }
1435
1436 fn get_alter_table_expr_on_demand(
1444 &self,
1445 req: &mut RowInsertRequest,
1446 table: &TableRef,
1447 ctx: &QueryContextRef,
1448 accommodate_existing_schema: bool,
1449 is_single_value: bool,
1450 alter_existing: bool,
1451 ) -> Result<Option<AlterTableExpr>> {
1452 if !alter_existing {
1453 return Ok(None);
1454 }
1455
1456 let catalog_name = ctx.current_catalog();
1457 let schema_name = ctx.current_schema();
1458 let table_name = table.table_info().name.clone();
1459
1460 if is_ddl_reserved_table(&schema_name, &table_name) {
1463 return Ok(None);
1464 }
1465
1466 let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1467 let request_field_count = request_schema
1468 .iter()
1469 .filter(|col| col.semantic_type == SemanticType::Field as i32)
1470 .count();
1471 let column_exprs = ColumnExpr::from_column_schemas(request_schema);
1472 let add_columns = expr_helper::extract_add_columns_expr(&table.schema(), column_exprs)?;
1473 let Some(mut add_columns) = add_columns else {
1474 return Ok(None);
1475 };
1476
1477 if is_single_value {
1478 let request_is_native_histogram = request_is_native_histogram(request_schema);
1479 let table_is_native_histogram = table_is_native_histogram(table);
1480 ensure!(
1481 request_is_native_histogram == table_is_native_histogram,
1482 InvalidInsertRequestSnafu {
1483 reason: format!(
1484 "Table `{table_name}` cannot mix native histogram and float sample fields"
1485 ),
1486 }
1487 );
1488 }
1489
1490 if accommodate_existing_schema {
1492 let table_schema = table.schema();
1493 let ts_col_name = table_schema.timestamp_column().map(|c| c.name.clone());
1495 let mut field_col_name = None;
1497 if is_single_value && request_field_count <= 1 {
1498 let mut multiple_field_cols = false;
1499 table.field_columns().for_each(|col| {
1500 if field_col_name.is_none() {
1501 field_col_name = Some(col.name.clone());
1502 } else {
1503 multiple_field_cols = true;
1504 }
1505 });
1506 if multiple_field_cols {
1507 field_col_name = None;
1508 }
1509 }
1510
1511 if let Some(rows) = req.rows.as_mut() {
1513 for col in &mut rows.schema {
1514 match col.semantic_type {
1515 x if x == SemanticType::Timestamp as i32 => {
1516 if let Some(ref ts_name) = ts_col_name
1517 && col.column_name != *ts_name
1518 {
1519 col.column_name = ts_name.clone();
1520 }
1521 }
1522 x if x == SemanticType::Field as i32 => {
1523 if let Some(ref field_name) = field_col_name
1524 && col.column_name != *field_name
1525 {
1526 col.column_name = field_name.clone();
1527 }
1528 }
1529 _ => {}
1530 }
1531 }
1532 }
1533
1534 add_columns.add_columns.retain(|col| {
1536 let def = col.column_def.as_ref().unwrap();
1537 def.semantic_type == SemanticType::Tag as i32
1538 || (def.semantic_type == SemanticType::Field as i32 && field_col_name.is_none())
1539 });
1540
1541 if add_columns.add_columns.is_empty() {
1542 return Ok(None);
1543 }
1544 }
1545
1546 Ok(Some(AlterTableExpr {
1547 catalog_name: catalog_name.to_string(),
1548 schema_name: schema_name.clone(),
1549 table_name: table_name.clone(),
1550 kind: Some(Kind::AddColumns(add_columns)),
1551 }))
1552 }
1553
1554 async fn create_physical_table(
1556 &self,
1557 mut create_table_expr: CreateTableExpr,
1558 partitions: Option<Partitions>,
1559 ctx: &QueryContextRef,
1560 statement_executor: &StatementExecutor,
1561 ) -> Result<TableRef> {
1562 let res = statement_executor
1563 .create_table_inner(
1564 &mut create_table_expr,
1565 partitions,
1566 ctx.clone(),
1567 TriggerReason::AutoCreate,
1568 )
1569 .await;
1570
1571 let table_ref = TableReference::full(
1572 &create_table_expr.catalog_name,
1573 &create_table_expr.schema_name,
1574 &create_table_expr.table_name,
1575 );
1576
1577 match res {
1578 Ok(table) => {
1579 validate_trace_table_model(&table.table_info(), ctx)?;
1580 Ok(table)
1581 }
1582 Err(err) => {
1583 error!(err; "Failed to create table {}", table_ref);
1584 Err(err)
1585 }
1586 }
1587 }
1588
1589 async fn create_logical_tables(
1590 &self,
1591 create_table_exprs: Vec<CreateTableExpr>,
1592 ctx: &QueryContextRef,
1593 statement_executor: &StatementExecutor,
1594 ) -> Result<Vec<TableRef>> {
1595 let res = statement_executor
1596 .create_logical_tables(&create_table_exprs, ctx.clone(), TriggerReason::AutoCreate)
1597 .await;
1598
1599 match res {
1600 Ok(res) => Ok(res),
1601 Err(err) => {
1602 let failed_tables = create_table_exprs
1603 .into_iter()
1604 .map(|expr| {
1605 format!(
1606 "{}.{}.{}",
1607 expr.catalog_name, expr.schema_name, expr.table_name
1608 )
1609 })
1610 .collect::<Vec<_>>();
1611 error!(
1612 err;
1613 "Failed to create logical tables {:?}",
1614 failed_tables
1615 );
1616 Err(err)
1617 }
1618 }
1619 }
1620
1621 pub fn node_manager(&self) -> &NodeManagerRef {
1622 &self.node_manager
1623 }
1624
1625 pub fn partition_manager(&self) -> &PartitionRuleManagerRef {
1626 &self.partition_manager
1627 }
1628
1629 pub fn table_flownode_set_cache(&self) -> &TableFlownodeSetCacheRef {
1630 &self.table_flownode_set_cache
1631 }
1632}
1633
1634fn request_is_native_histogram(request_schema: &[ColumnSchema]) -> bool {
1635 let mut fields = request_schema
1636 .iter()
1637 .filter(|col| col.semantic_type == SemanticType::Field as i32);
1638 let Some(col) = fields.next() else {
1639 return false;
1640 };
1641
1642 fields.next().is_none()
1643 && api::helper::is_column_type_value_eq(
1644 col.datatype,
1645 col.datatype_extension.clone(),
1646 native_histogram_value_type(),
1647 )
1648}
1649
1650fn table_time_index_unit(table: &TableRef) -> Option<TimeUnit> {
1654 table
1655 .table_info()
1656 .meta
1657 .schema
1658 .timestamp_column()
1659 .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit()))
1660}
1661
1662fn convert_rows_time_unit(rows: &mut Rows, target_unit: TimeUnit) -> Result<()> {
1663 let Some(ts_index) = rows
1664 .schema
1665 .iter()
1666 .position(|col| col.semantic_type == SemanticType::Timestamp as i32)
1667 else {
1668 return Ok(());
1669 };
1670 let Some(source_unit) = ColumnDataType::try_from(rows.schema[ts_index].datatype)
1671 .ok()
1672 .and_then(api::helper::timestamp_unit)
1673 else {
1674 return Ok(());
1675 };
1676 if source_unit == target_unit {
1677 return Ok(());
1678 }
1679
1680 rows.schema[ts_index].datatype = api::helper::timestamp_datatype(target_unit) as i32;
1681 rows.schema[ts_index].datatype_extension = None;
1683
1684 for row in &mut rows.rows {
1692 debug_assert_eq!(row.values.len(), rows.schema.len());
1693 let value = &mut row.values[ts_index];
1694 let Some(value_data) = value.value_data.take() else {
1695 continue;
1696 };
1697 value.value_data =
1698 convert_timestamp_value_data(value_data, source_unit, target_unit, ts_index)?;
1699 }
1700 Ok(())
1701}
1702
1703fn convert_timestamp_value_data(
1704 value_data: ValueData,
1705 source_unit: TimeUnit,
1706 target_unit: TimeUnit,
1707 column_index: usize,
1708) -> Result<Option<ValueData>> {
1709 let timestamp = match value_data {
1710 ValueData::TimestampSecondValue(v) => Timestamp::new_second(v),
1711 ValueData::TimestampMillisecondValue(v) => Timestamp::new_millisecond(v),
1712 ValueData::TimestampMicrosecondValue(v) => Timestamp::new_microsecond(v),
1713 ValueData::TimestampNanosecondValue(v) => Timestamp::new_nanosecond(v),
1714 other => return Ok(Some(other)),
1716 };
1717 let converted = timestamp
1718 .convert_to(target_unit)
1719 .with_context(|| InvalidInsertRequestSnafu {
1720 reason: format!(
1721 "timestamp column {column_index} value {} in unit {source_unit:?} overflows when converting to unit {target_unit:?}",
1722 timestamp.value()
1723 ),
1724 })?;
1725 Ok(api::helper::to_grpc_value(datatypes::value::Value::Timestamp(converted)).value_data)
1726}
1727
1728fn table_is_native_histogram(table: &TableRef) -> bool {
1729 let mut fields = table.field_columns();
1730 let Some(col) = fields.next() else {
1731 return false;
1732 };
1733
1734 fields.next().is_none() && is_native_histogram_value_type(&col.data_type)
1735}
1736
1737fn validate_column_count_match(requests: &RowInsertRequests) -> Result<()> {
1738 for request in &requests.inserts {
1739 let rows = request.rows.as_ref().unwrap();
1740 let column_count = rows.schema.len();
1741 rows.rows.iter().try_for_each(|r| {
1742 ensure!(
1743 r.values.len() == column_count,
1744 InvalidInsertRequestSnafu {
1745 reason: format!(
1746 "column count mismatch, columns: {}, values: {}",
1747 column_count,
1748 r.values.len()
1749 )
1750 }
1751 );
1752 Ok(())
1753 })?;
1754 }
1755 Ok(())
1756}
1757
1758pub fn validate_trace_table_model(table_info: &TableInfo, ctx: &QueryContextRef) -> Result<()> {
1761 let Some(expected @ (TABLE_DATA_MODEL_TRACE_V1 | TABLE_DATA_MODEL_TRACE_V2)) =
1762 ctx.extension(SEMANTIC_PIPELINE)
1763 else {
1764 return Ok(());
1765 };
1766 if let Some(actual) = table_info.meta.options.data_model() {
1767 ensure!(
1768 actual == expected,
1769 InvalidInsertRequestSnafu {
1770 reason: format!(
1771 "Trace table `{}` uses {actual}, but the request uses {expected}",
1772 table_info.name,
1773 ),
1774 }
1775 );
1776 }
1777 Ok(())
1778}
1779
1780pub fn fill_table_options_for_create(
1782 table_options: &mut std::collections::HashMap<String, String>,
1783 create_type: &AutoCreateTableType,
1784 ctx: &QueryContextRef,
1785) {
1786 for key in VALID_TABLE_OPTION_KEYS {
1787 if let Some(value) = ctx.extension(key) {
1788 table_options.insert(key.to_string(), value.to_string());
1789 }
1790 }
1791
1792 for (key, value) in ctx.extensions() {
1794 if is_semantic_option_key(&key) && validate_semantic_option(&key, &value) {
1795 table_options.insert(key, value);
1796 }
1797 }
1798
1799 match create_type {
1800 AutoCreateTableType::Logical(physical_table) => {
1801 table_options.insert(
1802 LOGICAL_TABLE_METADATA_KEY.to_string(),
1803 physical_table.clone(),
1804 );
1805 }
1806 AutoCreateTableType::Physical => {
1807 if let Some(append_mode) = ctx.extension(APPEND_MODE_KEY) {
1808 table_options.insert(APPEND_MODE_KEY.to_string(), append_mode.to_string());
1809 }
1810 if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
1811 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
1812 }
1813 if let Some(time_window) = ctx.extension(TWCS_TIME_WINDOW) {
1814 table_options.insert(TWCS_TIME_WINDOW.to_string(), time_window.to_string());
1815 table_options.insert(
1817 COMPACTION_TYPE.to_string(),
1818 COMPACTION_TYPE_TWCS.to_string(),
1819 );
1820 }
1821 }
1822 AutoCreateTableType::Log => {
1825 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1826 }
1827 AutoCreateTableType::LastNonNull => {
1828 if ctx
1829 .extension(APPEND_MODE_KEY)
1830 .is_some_and(|value| value.eq_ignore_ascii_case("true"))
1831 {
1832 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1833 table_options.insert(MERGE_MODE_KEY.to_string(), "last_row".to_string());
1834 } else if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
1835 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
1836 } else {
1837 table_options.insert(MERGE_MODE_KEY.to_string(), "last_non_null".to_string());
1838 }
1839 }
1840 AutoCreateTableType::Trace { .. } => {
1841 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1842 }
1843 }
1844}
1845
1846pub type PerTableSemanticIndex = BTreeMap<String, BTreeMap<String, BTreeMap<String, String>>>;
1852
1853pub fn parse_per_table_semantic_index(ctx: &QueryContextRef) -> Option<PerTableSemanticIndex> {
1858 let raw = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY)?;
1859 match serde_json::from_str(raw) {
1860 Ok(index) => Some(index),
1861 Err(_) => {
1862 warn!("failed to parse semantic per-table index, skipping per-table options");
1863 None
1864 }
1865 }
1866}
1867
1868pub fn apply_per_table_semantic_options(
1876 table_options: &mut std::collections::HashMap<String, String>,
1877 index: Option<&PerTableSemanticIndex>,
1878 schema: &str,
1879 table_name: &str,
1880) {
1881 let Some(entry) = index
1882 .and_then(|index| index.get(schema))
1883 .and_then(|tables| tables.get(table_name))
1884 else {
1885 return;
1886 };
1887 for (key, value) in entry {
1888 if is_semantic_option_key(key) && validate_semantic_option(key, value) {
1889 table_options.insert(key.clone(), value.clone());
1890 }
1891 }
1892}
1893
1894pub fn build_create_table_expr(
1895 table: &TableReference,
1896 request_schema: &[ColumnSchema],
1897 engine: &str,
1898) -> Result<CreateTableExpr> {
1899 expr_helper::create_table_expr_by_column_schemas(table, request_schema, engine, None)
1900}
1901
1902pub const SPLUNK_PK_METADATA_ORDER_KEY: &str = "splunk_pk_metadata_order";
1906
1907fn reorder_splunk_primary_keys(primary_keys: &mut [String]) {
1910 const LEAD: [&str; 3] = ["host", "source", "sourcetype"];
1911 primary_keys.sort_by_key(|name| {
1914 LEAD.iter()
1915 .position(|&lead| lead == name.as_str())
1916 .unwrap_or(LEAD.len())
1917 });
1918}
1919
1920struct CreateAlterTableResult {
1922 instant_table_ids: HashSet<TableId>,
1924 table_infos: HashMap<TableId, Arc<TableInfo>>,
1926}
1927
1928struct FlowMirrorTask {
1929 requests: HashMap<Peer, RegionInsertRequests>,
1930 num_rows: usize,
1931}
1932
1933impl FlowMirrorTask {
1934 async fn new(
1935 cache: &TableFlownodeSetCacheRef,
1936 requests: impl Iterator<Item = &RegionInsertRequest>,
1937 ) -> Result<Self> {
1938 let mut src_table_reqs: HashMap<TableId, Option<(Vec<Peer>, RegionInsertRequests)>> =
1939 HashMap::new();
1940 let mut num_rows = 0;
1941
1942 for req in requests {
1943 let table_id = RegionId::from_u64(req.region_id).table_id();
1944 match src_table_reqs.get_mut(&table_id) {
1945 Some(Some((_peers, reqs))) => reqs.requests.push(req.clone()),
1946 Some(None) => continue,
1948 _ => {
1949 let peers = cache
1951 .get(table_id)
1952 .await
1953 .context(RequestInsertsSnafu)?
1954 .unwrap_or_default()
1955 .values()
1956 .cloned()
1957 .collect::<HashSet<_>>()
1958 .into_iter()
1959 .collect::<Vec<_>>();
1960
1961 if !peers.is_empty() {
1962 let mut reqs = RegionInsertRequests::default();
1963 reqs.requests.push(req.clone());
1964 num_rows += reqs
1965 .requests
1966 .iter()
1967 .map(|r| r.rows.as_ref().unwrap().rows.len())
1968 .sum::<usize>();
1969 src_table_reqs.insert(table_id, Some((peers, reqs)));
1970 } else {
1971 src_table_reqs.insert(table_id, None);
1973 }
1974 }
1975 }
1976 }
1977
1978 let mut inserts: HashMap<Peer, RegionInsertRequests> = HashMap::new();
1979
1980 for (_table_id, (peers, reqs)) in src_table_reqs
1981 .into_iter()
1982 .filter_map(|(k, v)| v.map(|v| (k, v)))
1983 {
1984 if peers.len() == 1 {
1985 inserts
1987 .entry(peers[0].clone())
1988 .or_default()
1989 .requests
1990 .extend(reqs.requests);
1991 continue;
1992 } else {
1993 for flownode in peers {
1995 inserts
1996 .entry(flownode.clone())
1997 .or_default()
1998 .requests
1999 .extend(reqs.requests.clone());
2000 }
2001 }
2002 }
2003
2004 Ok(Self {
2005 requests: inserts,
2006 num_rows,
2007 })
2008 }
2009
2010 fn detach(self, node_manager: NodeManagerRef) -> Result<()> {
2011 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(self.num_rows as i64);
2012 for (peer, inserts) in self.requests {
2013 let node_manager = node_manager.clone();
2014 common_runtime::spawn_global(async move {
2015 let result = node_manager
2016 .flownode(&peer)
2017 .await
2018 .handle_inserts(inserts)
2019 .await
2020 .context(RequestInsertsSnafu);
2021
2022 match result {
2023 Ok(resp) => {
2024 let affected_rows = resp.affected_rows;
2025 crate::metrics::DIST_MIRROR_ROW_COUNT.inc_by(affected_rows);
2026 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(affected_rows as _);
2027 }
2028 Err(err) => {
2029 error!(err; "Failed to insert data into flownode {}", peer);
2030 }
2031 }
2032 });
2033 }
2034
2035 Ok(())
2036 }
2037}
2038
2039#[cfg(test)]
2040mod tests {
2041 use std::sync::Arc;
2042
2043 use api::helper::ColumnDataTypeWrapper;
2044 use api::v1::helper::{field_column_schema, time_index_column_schema};
2045 use api::v1::{Row, RowInsertRequest, Rows, Value};
2046 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
2047 use common_meta::cache::new_table_flownode_set_cache;
2048 use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler;
2049 use common_meta::test_util::MockDatanodeManager;
2050 use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD;
2051 use common_query::prelude::{greptime_native_histogram, set_default_prefix};
2052 use datatypes::data_type::ConcreteDataType;
2053 use datatypes::schema::ColumnSchema;
2054 use moka::future::Cache;
2055 use session::context::QueryContext;
2056 use table::TableRef;
2057 use table::dist_table::DummyDataSource;
2058 use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType};
2059
2060 use crate::insert::*;
2061 use crate::test_util::{
2062 create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
2063 };
2064
2065 fn make_table_ref_with_schema(
2066 ts_name: &str,
2067 field_name: &str,
2068 field_type: ConcreteDataType,
2069 ) -> TableRef {
2070 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2071 ColumnSchema::new(
2072 ts_name,
2073 ConcreteDataType::timestamp_millisecond_datatype(),
2074 false,
2075 )
2076 .with_time_index(true),
2077 ColumnSchema::new(field_name, field_type, true),
2078 ])
2079 .unwrap()
2080 .build()
2081 .unwrap();
2082 let meta = TableMetaBuilder::empty()
2083 .schema(Arc::new(schema))
2084 .primary_key_indices(vec![])
2085 .value_indices(vec![1])
2086 .engine("mito")
2087 .next_column_id(0)
2088 .options(Default::default())
2089 .created_on(Default::default())
2090 .build()
2091 .unwrap();
2092 let info = Arc::new(
2093 TableInfoBuilder::default()
2094 .table_id(1)
2095 .table_version(0)
2096 .name("test_table")
2097 .schema_name(DEFAULT_SCHEMA_NAME)
2098 .catalog_name(DEFAULT_CATALOG_NAME)
2099 .desc(None)
2100 .table_type(TableType::Base)
2101 .meta(meta)
2102 .build()
2103 .unwrap(),
2104 );
2105 Arc::new(table::Table::new(
2106 info,
2107 table::metadata::FilterPushDownType::Unsupported,
2108 Arc::new(DummyDataSource),
2109 ))
2110 }
2111
2112 fn make_metric_physical_table_ref_with_time_unit(unit: TimeUnit) -> TableRef {
2113 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2114 ColumnSchema::new(
2115 greptime_timestamp(),
2116 ConcreteDataType::timestamp_datatype(unit),
2117 false,
2118 )
2119 .with_time_index(true),
2120 ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
2121 ])
2122 .unwrap()
2123 .build()
2124 .unwrap();
2125 let meta = TableMetaBuilder::empty()
2126 .schema(Arc::new(schema))
2127 .primary_key_indices(vec![])
2128 .value_indices(vec![1])
2129 .engine("metric")
2130 .next_column_id(0)
2131 .options(Default::default())
2132 .created_on(Default::default())
2133 .build()
2134 .unwrap();
2135 let info = Arc::new(
2136 TableInfoBuilder::default()
2137 .table_id(1)
2138 .table_version(0)
2139 .name("greptime_physical_table")
2140 .schema_name(DEFAULT_SCHEMA_NAME)
2141 .catalog_name(DEFAULT_CATALOG_NAME)
2142 .desc(None)
2143 .table_type(TableType::Base)
2144 .meta(meta)
2145 .build()
2146 .unwrap(),
2147 );
2148 Arc::new(table::Table::new(
2149 info,
2150 table::metadata::FilterPushDownType::Unsupported,
2151 Arc::new(DummyDataSource),
2152 ))
2153 }
2154
2155 fn ms_row_insert_request(timestamp_ms: i64) -> RowInsertRequest {
2156 ms_row_insert_request_named("my_metric", timestamp_ms)
2157 }
2158
2159 fn ms_row_insert_request_named(table: &str, timestamp_ms: i64) -> RowInsertRequest {
2160 RowInsertRequest {
2161 table_name: table.to_string(),
2162 rows: Some(Rows {
2163 schema: vec![
2164 time_index_column_schema(
2165 greptime_timestamp(),
2166 ColumnDataType::TimestampMillisecond,
2167 ),
2168 field_column_schema(greptime_value(), ColumnDataType::Float64),
2169 ],
2170 rows: vec![Row {
2171 values: vec![
2172 Value {
2173 value_data: Some(ValueData::TimestampMillisecondValue(timestamp_ms)),
2174 },
2175 Value {
2176 value_data: Some(ValueData::F64Value(1.0)),
2177 },
2178 ],
2179 }],
2180 }),
2181 }
2182 }
2183
2184 fn convert_row_insert_requests_time_unit(
2186 requests: &mut RowInsertRequests,
2187 target_unit: TimeUnit,
2188 ) -> Result<()> {
2189 for request in &mut requests.inserts {
2190 let Some(rows) = request.rows.as_mut() else {
2191 continue;
2192 };
2193 convert_rows_time_unit(rows, target_unit)?;
2194 }
2195 Ok(())
2196 }
2197
2198 #[test]
2199 fn test_convert_row_insert_requests_time_unit_noop_when_matching() {
2200 let mut requests = RowInsertRequests {
2201 inserts: vec![ms_row_insert_request(123)],
2202 };
2203 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Millisecond).unwrap();
2204 let rows = requests.inserts[0].rows.as_ref().unwrap();
2205 assert_eq!(
2206 rows.schema[0].datatype,
2207 ColumnDataType::TimestampMillisecond as i32
2208 );
2209 assert!(matches!(
2210 rows.rows[0].values[0].value_data,
2211 Some(ValueData::TimestampMillisecondValue(123))
2212 ));
2213 }
2214
2215 #[test]
2216 fn test_convert_row_insert_requests_time_unit_widens_losslessly() {
2217 let mut requests = RowInsertRequests {
2218 inserts: vec![ms_row_insert_request(123)],
2219 };
2220 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Microsecond).unwrap();
2221 let rows = requests.inserts[0].rows.as_ref().unwrap();
2222 assert_eq!(
2223 rows.schema[0].datatype,
2224 ColumnDataType::TimestampMicrosecond as i32
2225 );
2226 assert!(matches!(
2227 rows.rows[0].values[0].value_data,
2228 Some(ValueData::TimestampMicrosecondValue(123_000))
2229 ));
2230 }
2231
2232 #[test]
2233 fn test_convert_row_insert_requests_time_unit_truncates_on_narrowing() {
2234 let requests = |value_ns: i64| RowInsertRequests {
2237 inserts: vec![RowInsertRequest {
2238 table_name: "my_metric".to_string(),
2239 rows: Some(Rows {
2240 schema: vec![
2241 time_index_column_schema(
2242 greptime_timestamp(),
2243 ColumnDataType::TimestampNanosecond,
2244 ),
2245 field_column_schema(greptime_value(), ColumnDataType::Float64),
2246 ],
2247 rows: vec![Row {
2248 values: vec![
2249 Value {
2250 value_data: Some(ValueData::TimestampNanosecondValue(value_ns)),
2251 },
2252 Value {
2253 value_data: Some(ValueData::F64Value(1.0)),
2254 },
2255 ],
2256 }],
2257 }),
2258 }],
2259 };
2260
2261 let mut reqs = requests(123_456_789);
2262 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
2263 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2264 assert!(matches!(
2265 rows.rows[0].values[0].value_data,
2266 Some(ValueData::TimestampMicrosecondValue(123_456))
2267 ));
2268
2269 let mut reqs = requests(123_456_789);
2270 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Millisecond).unwrap();
2271 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2272 assert!(matches!(
2273 rows.rows[0].values[0].value_data,
2274 Some(ValueData::TimestampMillisecondValue(123))
2275 ));
2276
2277 let mut reqs = requests(-123_456_789);
2278 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
2279 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2280 assert!(matches!(
2281 rows.rows[0].values[0].value_data,
2282 Some(ValueData::TimestampMicrosecondValue(-123_457))
2283 ));
2284 }
2285
2286 #[test]
2287 fn test_convert_row_insert_requests_time_unit_overflow() {
2288 let mut requests = RowInsertRequests {
2289 inserts: vec![ms_row_insert_request(i64::MAX)],
2290 };
2291 let err =
2292 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Nanosecond).unwrap_err();
2293 assert!(err.to_string().contains("overflows"), "{err}");
2294 }
2295
2296 #[test]
2297 fn test_table_time_index_unit() {
2298 assert_eq!(
2299 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
2300 TimeUnit::Microsecond
2301 )),
2302 Some(TimeUnit::Microsecond)
2303 );
2304 assert_eq!(
2305 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
2306 TimeUnit::Millisecond
2307 )),
2308 Some(TimeUnit::Millisecond)
2309 );
2310 }
2311
2312 #[tokio::test]
2313 async fn test_accommodate_existing_schema_and_reject_kind_changes() {
2314 let ts_name = "my_ts";
2315 let field_name = "my_field";
2316 let table =
2317 make_table_ref_with_schema(ts_name, field_name, ConcreteDataType::float64_datatype());
2318
2319 let mut req = RowInsertRequest {
2321 table_name: "test_table".to_string(),
2322 rows: Some(Rows {
2323 schema: vec![
2324 time_index_column_schema("ts_wrong", ColumnDataType::TimestampMillisecond),
2325 field_column_schema("field_wrong", ColumnDataType::Float64),
2326 ],
2327 rows: vec![api::v1::Row {
2328 values: vec![Value::default(), Value::default()],
2329 }],
2330 }),
2331 };
2332 let ctx = Arc::new(QueryContext::with(
2333 DEFAULT_CATALOG_NAME,
2334 DEFAULT_SCHEMA_NAME,
2335 ));
2336
2337 let kv_backend = prepare_mocked_backend().await;
2338 let inserter = Inserter::new(
2339 catalog::memory::MemoryCatalogManager::new(),
2340 create_partition_rule_manager(kv_backend.clone()).await,
2341 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
2342 Arc::new(new_table_flownode_set_cache(
2343 String::new(),
2344 Cache::new(100),
2345 kv_backend.clone(),
2346 )),
2347 true,
2348 );
2349 assert!(
2351 inserter
2352 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, false, false, false)
2353 .unwrap()
2354 .is_none()
2355 );
2356 let alter_expr = inserter
2357 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true, true)
2358 .unwrap();
2359 assert!(alter_expr.is_none());
2360
2361 let req_schema = req.rows.as_ref().unwrap().schema.clone();
2363 assert_eq!(req_schema[0].column_name, ts_name);
2364 assert_eq!(req_schema[1].column_name, field_name);
2365
2366 let (datatype, datatype_extension) =
2367 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
2368 .unwrap()
2369 .into_parts();
2370 let mut histogram_req = RowInsertRequest {
2371 table_name: "test_table".to_string(),
2372 rows: Some(Rows {
2373 schema: vec![
2374 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2375 api::v1::ColumnSchema {
2376 column_name: greptime_native_histogram().to_string(),
2377 datatype: datatype as i32,
2378 semantic_type: SemanticType::Field as i32,
2379 datatype_extension,
2380 options: None,
2381 },
2382 ],
2383 rows: vec![],
2384 }),
2385 };
2386 let error = inserter
2387 .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, false, true, true)
2388 .unwrap_err();
2389 assert!(
2390 error
2391 .to_string()
2392 .contains("cannot mix native histogram and float sample fields")
2393 );
2394
2395 let histogram_table = make_table_ref_with_schema(
2396 "ts",
2397 greptime_native_histogram(),
2398 native_histogram_value_type().clone(),
2399 );
2400 let mut sample_req = RowInsertRequest {
2401 table_name: "test_table".to_string(),
2402 rows: Some(Rows {
2403 schema: vec![
2404 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2405 field_column_schema(greptime_value(), ColumnDataType::Float64),
2406 ],
2407 rows: vec![],
2408 }),
2409 };
2410 let error = inserter
2411 .get_alter_table_expr_on_demand(
2412 &mut sample_req,
2413 &histogram_table,
2414 &ctx,
2415 false,
2416 true,
2417 true,
2418 )
2419 .unwrap_err();
2420 assert!(
2421 error
2422 .to_string()
2423 .contains("cannot mix native histogram and float sample fields")
2424 );
2425 }
2426
2427 #[test]
2428 fn test_native_histogram_detection_survives_prefix_change() {
2429 set_default_prefix(Some("custom")).unwrap();
2430 let table = make_table_ref_with_schema(
2431 "custom_timestamp",
2432 NATIVE_HISTOGRAM_FIELD,
2433 native_histogram_value_type().clone(),
2434 );
2435 let (datatype, datatype_extension) =
2436 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
2437 .unwrap()
2438 .into_parts();
2439 let request_schema = [api::v1::ColumnSchema {
2440 column_name: greptime_native_histogram().to_string(),
2441 datatype: datatype as i32,
2442 semantic_type: SemanticType::Field as i32,
2443 datatype_extension,
2444 options: None,
2445 }];
2446
2447 assert!(request_is_native_histogram(&request_schema));
2448 assert!(table_is_native_histogram(&table));
2449 }
2450
2451 #[tokio::test]
2453 async fn test_write_meter_admission() {
2454 use std::cell::Cell;
2455 use std::sync::Mutex;
2456 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
2457
2458 use api::region::RegionResponse;
2459 use api::v1::region::region_request::Body;
2460 use arrow::array::{Int32Array, TimestampMillisecondArray};
2461 use arrow::record_batch::RecordBatch;
2462 use bytes::Bytes;
2463 use common_error::ext::{ErrorExt, RetryHint};
2464 use common_error::status_code::StatusCode;
2465 use common_grpc::flight::{FlightEncoder, FlightMessage};
2466 use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher;
2467 use futures::future::BoxFuture;
2468 use meter_core::ItemCalculator;
2469 use meter_core::collect::{Collect, WriteRejected};
2470 use meter_core::data::MeterRecord;
2471 use meter_core::global::global_registry;
2472 use session::context::Channel;
2473
2474 const CATALOG: &str = "write_meter_test";
2475
2476 #[derive(Default)]
2477 struct Meter {
2478 reject: AtomicBool,
2479 attempts: Mutex<Vec<MeterRecord>>,
2480 accepted_value: AtomicU64,
2481 }
2482
2483 impl Collect for Meter {
2484 fn on_write(
2485 &self,
2486 record: MeterRecord,
2487 ) -> BoxFuture<'_, std::result::Result<(), WriteRejected>> {
2488 Box::pin(async move {
2489 if record.catalog != CATALOG {
2490 return Ok(());
2491 }
2492 let value = record.value;
2493 self.attempts.lock().unwrap().push(record);
2494 if self.reject.load(Ordering::Relaxed) {
2495 return Err(WriteRejected::new("database row quota exhausted"));
2496 }
2497 self.accepted_value.fetch_add(value, Ordering::Relaxed);
2498 Ok(())
2499 })
2500 }
2501
2502 fn on_read(&self, _: MeterRecord) {}
2503 }
2504
2505 impl ItemCalculator<InstantAndNormalInsertRequests> for Meter {
2506 fn calc(&self, _: &InstantAndNormalInsertRequests) -> u64 {
2507 17
2508 }
2509 }
2510
2511 let kv_backend = prepare_mocked_backend().await;
2512 let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
2513 let (sender, mut dispatched) = tokio::sync::mpsc::channel(16);
2514 let watcher = DatanodeWatcher::new(sender).with_handler(|_, request| {
2515 let rows = match request.body.unwrap() {
2516 Body::Inserts(requests) => requests
2517 .requests
2518 .iter()
2519 .filter_map(|request| request.rows.as_ref())
2520 .map(|rows| rows.rows.len())
2521 .sum(),
2522 Body::BulkInsert(_) => 2,
2524 body => panic!("unexpected request: {body:?}"),
2525 };
2526 Ok(RegionResponse::new(rows))
2527 });
2528 let flow_cache = Cache::new(100);
2529 let inserter = Inserter::new(
2530 catalog::memory::MemoryCatalogManager::new(),
2531 partition_manager,
2532 Arc::new(MockDatanodeManager::new(watcher)),
2533 Arc::new(new_table_flownode_set_cache(
2534 String::new(),
2535 flow_cache.clone(),
2536 kv_backend,
2537 )),
2538 true,
2539 );
2540 let mut table_info = new_test_table_info(1, "table_1", [1].into_iter());
2541 table_info.catalog_name = CATALOG.to_string();
2542 table_info.schema_name = "target_db".to_string();
2543 let table_info = Arc::new(table_info);
2544 let table_infos = HashMap::from_iter([(1, table_info.clone())]);
2545 let ctx = Arc::new(QueryContext::with_channel(
2546 DEFAULT_CATALOG_NAME,
2547 DEFAULT_SCHEMA_NAME,
2548 Channel::Postgres,
2549 ));
2550 let rows_request = || {
2551 let build = |num_rows| RegionInsertRequests {
2552 requests: vec![RegionInsertRequest {
2553 region_id: RegionId::new(1, 1).as_u64(),
2554 rows: Some(Rows {
2555 schema: vec![],
2556 rows: vec![api::v1::Row { values: vec![] }; num_rows],
2557 }),
2558 ..Default::default()
2559 }],
2560 };
2561 InstantAndNormalInsertRequests {
2562 normal_requests: build(3),
2563 instant_requests: build(2),
2564 }
2565 };
2566
2567 let output = inserter
2569 .do_request(rows_request(), &table_infos, &ctx)
2570 .await
2571 .unwrap();
2572 assert_eq!(output.meta.cost, 0);
2573 assert!(matches!(output.data, OutputData::AffectedRows(3)));
2574 dispatched.try_recv().unwrap();
2575
2576 let meter = Arc::new(Meter::default());
2577 global_registry().set_collector(meter.clone());
2578 global_registry().register_calculator(meter.clone());
2579 let enabled = Cell::new(false);
2581 write_meter!({
2582 enabled.set(true);
2583 MeterRecord::new("probe".into(), "probe".into(), 0, 0, 0)
2584 })
2585 .await
2586 .unwrap();
2587 let enabled = enabled.get();
2588
2589 let output = inserter
2590 .do_request(rows_request(), &table_infos, &ctx)
2591 .await
2592 .unwrap();
2593 assert_eq!(output.meta.cost, if enabled { 17 } else { 0 });
2594 assert!(matches!(output.data, OutputData::AffectedRows(3)));
2595 dispatched.try_recv().unwrap();
2596 assert!(flow_cache.contains_key(&1));
2597 flow_cache.invalidate_all();
2598 meter.reject.store(true, Ordering::Relaxed);
2599 let result = inserter
2600 .do_request(rows_request(), &table_infos, &ctx)
2601 .await;
2602 if enabled {
2603 let error = result.unwrap_err();
2604 assert_eq!(error.status_code(), StatusCode::RateLimited);
2605 assert_eq!(error.retry_hint(), RetryHint::Retryable);
2606 assert!(error.to_string().contains("database row quota exhausted"));
2607 assert!(dispatched.try_recv().is_err());
2608 assert!(
2609 !flow_cache.contains_key(&1),
2610 "rejected write reached flow mirroring"
2611 );
2612 assert_eq!(meter.accepted_value.load(Ordering::Relaxed), 17);
2613 let attempts = meter.attempts.lock().unwrap();
2614 assert_eq!(attempts.len(), 2);
2615 for record in attempts.iter() {
2616 assert_eq!(record.catalog, CATALOG);
2617 assert_eq!(record.schema, "target_db");
2618 assert_eq!(
2619 (record.rows, record.value, record.source),
2620 (5, 17, Channel::Postgres as u8)
2621 );
2622 }
2623 } else {
2624 assert_eq!(result.unwrap().meta.cost, 0);
2625 dispatched.try_recv().unwrap();
2626 assert!(meter.attempts.lock().unwrap().is_empty());
2627 }
2628 meter.attempts.lock().unwrap().clear();
2629
2630 if enabled {
2632 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
2633 let rows = Rows {
2634 schema: vec![
2635 api::v1::helper::tag_column_schema("a", ColumnDataType::Int32),
2636 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2637 field_column_schema("b", ColumnDataType::Int32),
2638 ],
2639 rows: vec![api::v1::Row {
2640 values: vec![
2641 api::v1::value::ValueData::I32Value(60).into(),
2642 Value {
2643 value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue(
2644 0,
2645 )),
2646 },
2647 api::v1::value::ValueData::I32Value(0).into(),
2648 ],
2649 }],
2650 };
2651 let error = inserter
2652 .submit_table_rows(rows, table_info.clone(), ctx.clone(), &batcher)
2653 .await
2654 .unwrap_err();
2655 assert_eq!(error.status_code(), StatusCode::RateLimited);
2656 let mut attempts = meter.attempts.lock().unwrap();
2657 assert_eq!(attempts.len(), 1);
2658 assert_eq!(attempts[0].catalog, CATALOG);
2659 assert_eq!(attempts[0].schema, "target_db");
2660 assert_eq!(attempts[0].rows, 1);
2661 attempts.clear();
2662 }
2663
2664 let table = Arc::new(table::Table::new(
2665 table_info.clone(),
2666 table::metadata::FilterPushDownType::Unsupported,
2667 Arc::new(DummyDataSource),
2668 ));
2669 let batch = RecordBatch::try_new(
2670 table_info.meta.schema.arrow_schema().clone(),
2671 vec![
2672 Arc::new(Int32Array::from(vec![60, 70])),
2673 Arc::new(TimestampMillisecondArray::from(vec![0, 1])),
2674 Arc::new(Int32Array::from(vec![0, 0])),
2675 ],
2676 )
2677 .unwrap();
2678 let bulk_insert = |batch: RecordBatch| {
2679 let flight_data = FlightEncoder::default()
2680 .encode(FlightMessage::RecordBatch(batch.clone()))
2681 .into_iter()
2682 .next()
2683 .unwrap();
2684 inserter.handle_bulk_insert(
2685 table.clone(),
2686 flight_data,
2687 batch,
2688 Bytes::new(),
2689 false,
2690 Channel::Grpc,
2691 )
2692 };
2693 assert_eq!(bulk_insert(batch.slice(0, 0)).await.unwrap(), 0);
2695 assert!(meter.attempts.lock().unwrap().is_empty());
2696 assert!(dispatched.try_recv().is_err());
2697
2698 for reject in [false, true, false] {
2700 meter.reject.store(reject, Ordering::Relaxed);
2701 let result = bulk_insert(batch.clone()).await;
2702 if enabled && reject {
2703 let error = result.unwrap_err();
2704 assert_eq!(error.status_code(), StatusCode::RateLimited);
2705 assert_eq!(error.retry_hint(), RetryHint::Retryable);
2706 assert!(dispatched.try_recv().is_err());
2707 } else {
2708 assert_eq!(result.unwrap(), 2);
2709 dispatched.try_recv().unwrap();
2710 }
2711 }
2712 {
2713 let attempts = meter.attempts.lock().unwrap();
2714 assert_eq!(attempts.len(), if enabled { 3 } else { 0 });
2715 for record in attempts.iter() {
2716 assert_eq!(record.catalog, CATALOG);
2717 assert_eq!(record.schema, "target_db");
2718 assert_eq!(
2719 (record.rows, record.value, record.source),
2720 (2, 0, Channel::Grpc as u8)
2721 );
2722 }
2723 }
2724 assert_eq!(
2725 meter.accepted_value.load(Ordering::Relaxed),
2726 if enabled { 17 } else { 0 }
2727 );
2728
2729 meter.attempts.lock().unwrap().clear();
2732 let original = Arc::new(QueryContext::with_channel(
2733 CATALOG,
2734 "a",
2735 Channel::Prometheus,
2736 ));
2737 let mut batches = ["a", "b", "a"].map(|schema| {
2738 let ctx = if schema == "a" {
2739 original.clone()
2740 } else {
2741 Arc::new(QueryContext::with_channel(
2742 CATALOG,
2743 schema,
2744 Channel::Prometheus,
2745 ))
2746 };
2747 (
2748 ctx,
2749 RowInsertRequests {
2750 inserts: vec![RowInsertRequest {
2751 table_name: "data".into(),
2752 rows: Some(Rows {
2753 schema: vec![],
2754 rows: vec![api::v1::Row::default(); 2],
2755 }),
2756 }],
2757 },
2758 )
2759 });
2760 admit_row_insert_batches(&mut batches).await.unwrap();
2761 admit_row_insert_batches(&mut batches).await.unwrap();
2762 assert_eq!(original.write_rows_to_admit(CATALOG, "a", 4), 4);
2763 for (ctx, _) in &batches {
2764 assert_eq!(
2765 ctx.write_rows_to_admit(CATALOG, &ctx.current_schema(), 2),
2766 0
2767 );
2768 assert_eq!(ctx.channel(), Channel::Prometheus);
2769 }
2770 {
2771 let attempts = meter.attempts.lock().unwrap();
2772 let totals = attempts
2773 .iter()
2774 .map(|r| (r.schema.as_str(), r.rows, r.value))
2775 .collect::<Vec<_>>();
2776 assert_eq!(
2777 totals,
2778 if enabled {
2779 vec![("a", 4, 0), ("b", 2, 0)]
2780 } else {
2781 vec![]
2782 }
2783 );
2784 }
2785 meter.attempts.lock().unwrap().clear();
2786 meter.reject.store(false, Ordering::Relaxed);
2787 let ctx = Arc::new(QueryContext::with_channel(
2788 CATALOG,
2789 "logical",
2790 Channel::Otlp,
2791 ));
2792 let admitted = admit_write(2, &ctx).await.unwrap();
2793 let mut requests = RowInsertRequests {
2794 inserts: vec![RowInsertRequest {
2795 table_name: "metric".to_string(),
2796 rows: Some(Rows {
2797 schema: vec![],
2798 rows: vec![api::v1::Row::default(); 2],
2799 }),
2800 }],
2801 };
2802 let original = requests.clone();
2803 let cost = Inserter::meter_row_inserts(&mut requests, &admitted)
2804 .await
2805 .unwrap();
2806 assert_eq!(cost, if enabled { 17 } else { 0 });
2807 assert_eq!(requests, original);
2808 let records = meter
2809 .attempts
2810 .lock()
2811 .unwrap()
2812 .iter()
2813 .map(|record| (record.rows, record.value))
2814 .collect::<Vec<_>>();
2815 assert_eq!(
2816 records,
2817 if enabled {
2818 vec![(2, 0), (0, 17)]
2819 } else {
2820 vec![]
2821 }
2822 );
2823 meter.reject.store(true, Ordering::Relaxed);
2824 let result = Inserter::meter_row_inserts(&mut requests, &admitted).await;
2825 if enabled {
2826 assert_eq!(result.unwrap_err().status_code(), StatusCode::RateLimited);
2827 } else {
2828 assert_eq!(result.unwrap(), 0);
2829 }
2830 assert_eq!(requests, original);
2831 }
2832
2833 #[test]
2834 fn test_skip_wal_does_not_change_table_options() {
2835 check_skip_wal_does_not_change_table_options(false);
2836 check_skip_wal_does_not_change_table_options(true);
2837 }
2838
2839 fn check_skip_wal_does_not_change_table_options(skip_wal: bool) {
2840 let ctx = Arc::new(QueryContext::with(
2841 DEFAULT_CATALOG_NAME,
2842 DEFAULT_SCHEMA_NAME,
2843 ));
2844 ctx.set_skip_wal(skip_wal);
2845 let mut options = Default::default();
2846 fill_table_options_for_create(&mut options, &AutoCreateTableType::Physical, &ctx);
2847 assert!(!options.contains_key(session::hints::INSERT_SKIP_WAL_HINT));
2848 assert!(!options.contains_key("skip_wal"));
2849 }
2850
2851 #[test]
2852 fn test_last_non_null_create_options_preserve_default_without_append_mode() {
2853 let ctx = Arc::new(QueryContext::with(
2854 DEFAULT_CATALOG_NAME,
2855 DEFAULT_SCHEMA_NAME,
2856 ));
2857 let mut table_options = Default::default();
2858
2859 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
2860
2861 assert_eq!(
2862 Some("last_non_null"),
2863 table_options.get(MERGE_MODE_KEY).map(String::as_str)
2864 );
2865 assert!(!table_options.contains_key(APPEND_MODE_KEY));
2866 }
2867
2868 #[test]
2869 fn test_fill_table_options_copies_semantic_extensions() {
2870 use table::requests::{
2871 SEMANTIC_METRIC_TYPE, SEMANTIC_PER_TABLE_INDEX_KEY, SEMANTIC_SIGNAL_TYPE,
2872 SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, SOURCE_OPENTELEMETRY,
2873 };
2874
2875 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
2876 ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
2877 ctx.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
2878 ctx.set_extension(SEMANTIC_SOURCE_VERSION, "2.0");
2879 ctx.set_extension(SEMANTIC_METRIC_TYPE, "bogus");
2880 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, "{}");
2882 let ctx = Arc::new(ctx);
2883 let mut table_options = Default::default();
2884
2885 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::Physical, &ctx);
2886
2887 assert_eq!(
2888 Some(SIGNAL_TYPE_METRIC),
2889 table_options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str)
2890 );
2891 assert_eq!(
2892 Some(SOURCE_OPENTELEMETRY),
2893 table_options.get(SEMANTIC_SOURCE).map(String::as_str)
2894 );
2895 assert_eq!(
2896 Some("2.0"),
2897 table_options
2898 .get(SEMANTIC_SOURCE_VERSION)
2899 .map(String::as_str)
2900 );
2901 assert!(!table_options.contains_key(SEMANTIC_METRIC_TYPE));
2902 assert!(!table_options.contains_key(SEMANTIC_PER_TABLE_INDEX_KEY));
2903 }
2904
2905 #[test]
2906 fn test_apply_per_table_semantic_options() {
2907 use table::requests::{
2908 SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, SEMANTIC_PER_TABLE_INDEX_KEY,
2909 };
2910
2911 let index = format!(
2912 r#"{{
2913 "{DEFAULT_SCHEMA_NAME}": {{
2914 "http_requests_total": {{
2915 "greptime.semantic.metric.type": "counter",
2916 "greptime.semantic.metric.unit": "By",
2917 "greptime.semantic.metric.type_BOGUS": "x"
2918 }},
2919 "other_table": {{
2920 "greptime.semantic.metric.type": "gauge"
2921 }}
2922 }},
2923 "other_schema": {{
2924 "http_requests_total": {{
2925 "greptime.semantic.metric.type": "gauge"
2926 }}
2927 }}
2928 }}"#
2929 );
2930 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
2931 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index);
2932 let ctx = Arc::new(ctx);
2933
2934 let index = parse_per_table_semantic_index(&ctx);
2935 assert!(index.is_some());
2936 let index = index.as_ref();
2937
2938 let mut table_options = std::collections::HashMap::new();
2939 apply_per_table_semantic_options(
2940 &mut table_options,
2941 index,
2942 DEFAULT_SCHEMA_NAME,
2943 "http_requests_total",
2944 );
2945 assert_eq!(
2947 table_options.get(SEMANTIC_METRIC_TYPE).map(String::as_str),
2948 Some("counter")
2949 );
2950 assert_eq!(
2951 table_options.get(SEMANTIC_METRIC_UNIT).map(String::as_str),
2952 Some("By")
2953 );
2954 assert!(!table_options.contains_key("greptime.semantic.metric.type_BOGUS"));
2957 assert_eq!(table_options.len(), 2);
2958
2959 let mut empty = std::collections::HashMap::new();
2960 apply_per_table_semantic_options(&mut empty, index, DEFAULT_SCHEMA_NAME, "not_in_index");
2961 assert!(empty.is_empty());
2962
2963 let mut opts = std::collections::HashMap::new();
2966 apply_per_table_semantic_options(
2967 &mut opts,
2968 index,
2969 "schema_without_entry",
2970 "http_requests_total",
2971 );
2972 assert!(opts.is_empty());
2973
2974 let bare = Arc::new(QueryContext::with(
2976 DEFAULT_CATALOG_NAME,
2977 DEFAULT_SCHEMA_NAME,
2978 ));
2979 assert!(parse_per_table_semantic_index(&bare).is_none());
2980 let mut opts = std::collections::HashMap::new();
2981 apply_per_table_semantic_options(
2982 &mut opts,
2983 None,
2984 DEFAULT_SCHEMA_NAME,
2985 "http_requests_total",
2986 );
2987 assert!(opts.is_empty());
2988 }
2989
2990 #[test]
2991 fn test_last_non_null_create_options_preserve_default_with_append_mode_false() {
2992 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
2993 ctx.set_extension(APPEND_MODE_KEY, "false");
2994 let ctx = Arc::new(ctx);
2995 let mut table_options = Default::default();
2996
2997 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
2998
2999 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3000 assert_eq!(
3001 Some("last_non_null"),
3002 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3003 );
3004 }
3005
3006 #[test]
3007 fn test_last_non_null_create_options_use_configured_merge_mode() {
3008 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3009 ctx.set_extension(MERGE_MODE_KEY, "last_row");
3010 let ctx = Arc::new(ctx);
3011 let mut table_options = Default::default();
3012
3013 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3014
3015 assert_eq!(
3016 Some("last_row"),
3017 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3018 );
3019 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3020 }
3021
3022 #[test]
3023 fn test_last_non_null_create_options_use_last_row_with_append_mode_true() {
3024 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3025 ctx.set_extension(APPEND_MODE_KEY, "true");
3026 let ctx = Arc::new(ctx);
3027 let mut table_options = Default::default();
3028
3029 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3030
3031 assert_eq!(
3032 Some("true"),
3033 table_options.get(APPEND_MODE_KEY).map(String::as_str)
3034 );
3035 assert_eq!(
3036 Some("last_row"),
3037 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3038 );
3039 }
3040
3041 struct UnexpectedBatcher;
3042
3043 #[async_trait::async_trait]
3044 impl PendingRowsBatcher for UnexpectedBatcher {
3045 async fn acquire(&self) -> Result<Arc<tokio::sync::OwnedSemaphorePermit>> {
3046 panic!("empty writes must not acquire batch admission")
3047 }
3048
3049 async fn submit(
3050 &self,
3051 _table_info: TableInfoRef,
3052 _batch: arrow::record_batch::RecordBatch,
3053 _ctx: QueryContextRef,
3054 _permit: Arc<tokio::sync::OwnedSemaphorePermit>,
3055 ) -> Result<usize> {
3056 panic!("empty writes must not submit a batch")
3057 }
3058 }
3059
3060 async fn batcher_test_inserter() -> Inserter {
3061 let kv_backend = prepare_mocked_backend().await;
3062 Inserter::new(
3063 catalog::memory::MemoryCatalogManager::new(),
3064 create_partition_rule_manager(kv_backend.clone()).await,
3065 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3066 Arc::new(new_table_flownode_set_cache(
3067 String::new(),
3068 Cache::new(100),
3069 kv_backend,
3070 )),
3071 true,
3072 )
3073 }
3074
3075 #[tokio::test]
3076 async fn test_logical_batcher_eligibility() {
3077 use catalog::RegisterTableRequest;
3078 use catalog::memory::MemoryCatalogManager;
3079 use common_meta::instruction::{CacheIdent, CreateFlow};
3080 use common_meta::kv_backend::KvBackendRef;
3081 use common_meta::kv_backend::memory::MemoryKvBackend;
3082 use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder};
3083 let requests = RowInsertRequests {
3084 inserts: vec![RowInsertRequest {
3085 table_name: "test_table".to_string(),
3086 rows: None,
3087 }],
3088 };
3089 let original =
3090 make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3091 .table_info();
3092 for case in [
3093 "eligible",
3094 "physical",
3095 "ordinary",
3096 "instant",
3097 "disabled",
3098 "hint",
3099 "flow",
3100 "required_tag",
3101 "default",
3102 ] {
3103 let mut info = (*original).clone();
3104 info.meta.engine = METRIC_ENGINE_NAME.to_string();
3105 info.meta.options.extra_options.insert(
3106 LOGICAL_TABLE_METADATA_KEY.to_string(),
3107 "physical".to_string(),
3108 );
3109 let mut ctx = QueryContext::arc().fork();
3110 match case {
3111 "physical" => {
3112 info.meta
3113 .options
3114 .extra_options
3115 .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "other".to_string());
3116 }
3117 "ordinary" => info.meta.engine = "mito".to_string(),
3118 "instant" => info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant),
3119 "hint" => ctx.set_extension(AUTO_CREATE_TABLE_KEY, "false"),
3120 "required_tag" | "default" => {
3121 let mut columns = info.meta.schema.column_schemas().to_vec();
3122 let mut tag = ColumnSchema::new(
3123 "tag",
3124 ConcreteDataType::string_datatype(),
3125 case != "required_tag",
3126 );
3127 if case == "default" {
3128 tag = tag
3129 .with_default_constraint(Some(ColumnDefaultConstraint::null_value()))
3130 .unwrap();
3131 }
3132 columns.push(tag);
3133 info.meta.schema = Arc::new(
3134 SchemaBuilder::try_from_columns(columns)
3135 .unwrap()
3136 .build()
3137 .unwrap(),
3138 );
3139 info.meta.primary_key_indices = vec![2];
3140 }
3141 _ => {}
3142 }
3143 let catalog = MemoryCatalogManager::with_default_setup();
3144 let table = Arc::new(table::Table::new(
3145 Arc::new(info),
3146 table::metadata::FilterPushDownType::Unsupported,
3147 Arc::new(DummyDataSource),
3148 ));
3149 catalog
3150 .register_table_sync(RegisterTableRequest {
3151 catalog: DEFAULT_CATALOG_NAME.to_string(),
3152 schema: DEFAULT_SCHEMA_NAME.to_string(),
3153 table_name: "test_table".to_string(),
3154 table_id: 1,
3155 table,
3156 })
3157 .unwrap();
3158 let mut inserter = batcher_test_inserter().await;
3159 inserter.catalog_manager = catalog;
3160 inserter.auto_create_table = case != "disabled";
3161 let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::default());
3162 inserter.table_flownode_set_cache = Arc::new(new_table_flownode_set_cache(
3163 String::new(),
3164 Cache::new(10),
3165 kv_backend,
3166 ));
3167 if case == "flow" {
3168 inserter
3169 .table_flownode_set_cache
3170 .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
3171 flow_id: 1,
3172 source_table_ids: vec![1],
3173 partition_to_peer_mapping: vec![(0, Peer::empty(1))],
3174 })])
3175 .await
3176 .unwrap();
3177 }
3178 assert_eq!(
3179 inserter
3180 .can_batch_metric_rows(&requests, &Arc::new(ctx), "physical")
3181 .await
3182 .unwrap(),
3183 case == "eligible",
3184 "{case}"
3185 );
3186 }
3187 }
3188
3189 #[tokio::test]
3190 async fn test_batcher_meter_preserves_request() {
3191 let mut requests = RowInsertRequests {
3192 inserts: vec![RowInsertRequest {
3193 table_name: "sample".to_string(),
3194 rows: Some(Rows {
3195 schema: vec![],
3196 rows: vec![api::v1::Row {
3197 values: vec![Value {
3198 value_data: Some(api::v1::value::ValueData::F64Value(1.5)),
3199 }],
3200 }],
3201 }),
3202 }],
3203 };
3204 let expected = requests.clone();
3205 Inserter::meter_row_inserts(&mut requests, &QueryContext::arc())
3206 .await
3207 .unwrap();
3208 assert_eq!(requests, expected);
3209 }
3210
3211 #[tokio::test]
3212 async fn test_instant_table_bypasses_batcher() {
3213 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3214 let inserter = batcher_test_inserter()
3215 .await
3216 .with_pending_rows_batcher(Some(batcher));
3217 let mut ctx = session::context::QueryContextBuilder::default().build();
3218 ctx.set_batching_enabled(true);
3219 let ctx = Arc::new(ctx);
3220 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3221 .table_info();
3222 assert!(inserter.table_batcher(&table, &ctx).is_some());
3223 let mut instant = (*table).clone();
3224 instant.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
3225 assert!(inserter.table_batcher(&Arc::new(instant), &ctx).is_none());
3226 }
3227
3228 #[tokio::test]
3229 async fn test_empty_prepared_rows_skip_batcher() {
3230 let inserter = batcher_test_inserter().await;
3231 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3232 .table_info();
3233 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3234 let ctx = QueryContext::arc();
3235 let output = inserter
3236 .submit_table_rows(
3237 Rows {
3238 schema: vec![],
3239 rows: vec![],
3240 },
3241 table.clone(),
3242 ctx.clone(),
3243 &batcher,
3244 )
3245 .await
3246 .unwrap();
3247 assert!(matches!(output.data, OutputData::AffectedRows(0)));
3248 let output = inserter
3249 .submit_pending_rows(
3250 RowInsertRequests {
3251 inserts: vec![RowInsertRequest {
3252 table_name: table.name.clone(),
3253 rows: None,
3254 }],
3255 },
3256 HashMap::from_iter([(table.table_id(), table)]),
3257 ctx,
3258 &batcher,
3259 )
3260 .await
3261 .unwrap();
3262 assert!(matches!(output.data, OutputData::AffectedRows(0)));
3263 }
3264}