Skip to main content

operator/
insert.rs

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