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