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