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