Skip to main content

operator/
insert.rs

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