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, METRICS_EXEMPLARS_TABLE, PARENT_SPAN_ID_COLUMN,
43    SERVICE_NAME_COLUMN, SPAN_ID_COLUMN, TIMESTAMP_END_COLUMN, TRACE_ID_COLUMN, TRACE_TABLE_NAME,
44    TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table,
45    trace_operations_table_name, trace_services_table_name,
46};
47use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME;
48use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME;
49use common_grpc_expr::util::ColumnExpr;
50use common_meta::cache::TableFlownodeSetCacheRef;
51use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME;
52use common_meta::node_manager::NodeManagerRef;
53use common_meta::peer::Peer;
54use common_meta::rpc::ddl::TriggerReason;
55use common_query::Output;
56use common_query::native_histogram::{is_native_histogram_value_type, native_histogram_value_type};
57use common_query::prelude::{
58    EXEMPLAR_LABELS_COLUMN, EXEMPLAR_METRIC_COLUMN, greptime_timestamp, 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: format!(
2212                "{METRICS_EXEMPLARS_TABLE} requires Mito, append_mode=true and primary_key_encoding=sparse"
2213            )
2214        }
2215    );
2216    let columns = meta.schema.column_schemas();
2217    let lookup = |name: &str| {
2218        columns
2219            .iter()
2220            .enumerate()
2221            .find(|(_, c)| c.name == name)
2222            .with_context(|| InvalidInsertRequestSnafu {
2223                reason: format!("missing exemplar column {name}"),
2224            })
2225    };
2226    let (_, timestamp) = lookup(greptime_timestamp())?;
2227    for (position, name, id, data_type) in [
2228        (
2229            0,
2230            DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
2231            ReservedColumnId::table_id(),
2232            datatypes::prelude::ConcreteDataType::uint32_datatype(),
2233        ),
2234        (
2235            1,
2236            DATA_SCHEMA_TSID_COLUMN_NAME,
2237            ReservedColumnId::tsid(),
2238            datatypes::prelude::ConcreteDataType::uint64_datatype(),
2239        ),
2240    ] {
2241        let (index, column) = lookup(name)?;
2242        ensure!(
2243            meta.primary_key_indices.get(position) == Some(&index)
2244                && meta.column_ids.get(index) == Some(&id)
2245                && column.data_type == data_type
2246                && !column.is_nullable()
2247                && column.default_constraint().is_none(),
2248            InvalidInsertRequestSnafu {
2249                reason: format!("invalid exemplar source identity column {name}"),
2250            }
2251        );
2252    }
2253    ensure!(
2254        timestamp.is_time_index() && timestamp.data_type.is_timestamp(),
2255        InvalidInsertRequestSnafu {
2256            reason: "invalid exemplar time index"
2257        }
2258    );
2259    let (metric_index, metric) = lookup(EXEMPLAR_METRIC_COLUMN)?;
2260    ensure!(
2261        metric.data_type.is_string() && meta.primary_key_indices.contains(&metric_index),
2262        InvalidInsertRequestSnafu {
2263            reason: "exemplar metric name must be a string tag"
2264        }
2265    );
2266    let (value_index, value) = lookup(greptime_value())?;
2267    ensure!(
2268        value.data_type == datatypes::prelude::ConcreteDataType::float64_datatype()
2269            && !meta.primary_key_indices.contains(&value_index),
2270        InvalidInsertRequestSnafu {
2271            reason: "exemplar value must be a DOUBLE field"
2272        }
2273    );
2274    let (labels_index, labels) = lookup(EXEMPLAR_LABELS_COLUMN)?;
2275    ensure!(
2276        labels.data_type.is_json2() && !meta.primary_key_indices.contains(&labels_index),
2277        InvalidInsertRequestSnafu {
2278            reason: "exemplar_labels must be a JSON2 field"
2279        }
2280    );
2281    if let Some(extension) = labels
2282        .extension_type::<datatypes::extension::json::Json2ExtensionType>()
2283        .map_err(|error| {
2284            InvalidInsertRequestSnafu {
2285                reason: error.to_string(),
2286            }
2287            .build()
2288        })?
2289    {
2290        ensure!(
2291            extension
2292                .metadata()
2293                .json_settings()
2294                .type_hints()
2295                .iter()
2296                .all(|hint| hint.path.len() == 1 && hint.data_type.is_string()),
2297            InvalidInsertRequestSnafu {
2298                reason: "exemplar JSON2 hints must preserve string labels"
2299            }
2300        );
2301    }
2302    for (index, column) in columns.iter().enumerate() {
2303        if [
2304            greptime_timestamp(),
2305            greptime_value(),
2306            EXEMPLAR_METRIC_COLUMN,
2307            EXEMPLAR_LABELS_COLUMN,
2308            DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
2309            DATA_SCHEMA_TSID_COLUMN_NAME,
2310        ]
2311        .contains(&column.name.as_str())
2312        {
2313            continue;
2314        }
2315        ensure!(
2316            meta.primary_key_indices.contains(&index)
2317                && column.data_type.is_string()
2318                && column.is_nullable(),
2319            InvalidInsertRequestSnafu {
2320                reason: "inherited exemplar labels must be nullable string tags"
2321            }
2322        );
2323        // An absent inherited label must remain absent from the source series identity.
2324        ensure!(
2325            matches!(
2326                column.default_constraint(),
2327                None | Some(ColumnDefaultConstraint::Value(
2328                    datatypes::value::Value::Null
2329                ))
2330            ),
2331            InvalidInsertRequestSnafu {
2332                reason: "inherited exemplar labels must not have non-NULL defaults"
2333            }
2334        );
2335    }
2336    Ok(())
2337}
2338
2339/// Rejects writes from a different built-in trace model before schema mutation.
2340/// Unstamped, explicitly created tables remain subject to normal schema validation.
2341pub fn validate_trace_table_model(table_info: &TableInfo, ctx: &QueryContextRef) -> Result<()> {
2342    let Some(expected @ (TABLE_DATA_MODEL_TRACE_V1 | TABLE_DATA_MODEL_TRACE_V2)) =
2343        ctx.extension(SEMANTIC_PIPELINE)
2344    else {
2345        return Ok(());
2346    };
2347    if let Some(actual) = table_info.meta.options.data_model() {
2348        ensure!(
2349            actual == expected,
2350            InvalidInsertRequestSnafu {
2351                reason: format!(
2352                    "Trace table `{}` uses {actual}, but the request uses {expected}",
2353                    table_info.name,
2354                ),
2355            }
2356        );
2357    }
2358    Ok(())
2359}
2360
2361/// Fill table options for a new table by create type.
2362pub fn fill_table_options_for_create(
2363    table_options: &mut std::collections::HashMap<String, String>,
2364    create_type: &AutoCreateTableType,
2365    ctx: &QueryContextRef,
2366) {
2367    for key in VALID_TABLE_OPTION_KEYS {
2368        if let Some(value) = ctx.extension(key) {
2369            table_options.insert(key.to_string(), value.to_string());
2370        }
2371    }
2372
2373    if !matches!(create_type, AutoCreateTableType::Logical(_)) {
2374        let mut has_twcs_option = false;
2375        for key in COMPACTION_OPTION_KEYS {
2376            if let Some(value) = ctx.extension(key) {
2377                has_twcs_option |= key.starts_with(COMPACTION_TWCS_PREFIX);
2378                table_options.insert(key.to_string(), value.to_string());
2379            }
2380        }
2381        if has_twcs_option {
2382            table_options
2383                .entry(COMPACTION_TYPE.to_string())
2384                .or_insert_with(|| COMPACTION_TYPE_TWCS.to_string());
2385        }
2386    }
2387
2388    // Semantic keys use their own vocabulary instead of the fixed option list.
2389    for (key, value) in ctx.extensions() {
2390        if is_semantic_option_key(&key) && validate_semantic_option(&key, &value) {
2391            table_options.insert(key, value);
2392        }
2393    }
2394
2395    match create_type {
2396        AutoCreateTableType::Logical(physical_table) => {
2397            table_options.insert(
2398                LOGICAL_TABLE_METADATA_KEY.to_string(),
2399                physical_table.clone(),
2400            );
2401        }
2402        AutoCreateTableType::Exemplar => {
2403            table_options.retain(|key, _| {
2404                !is_semantic_option_key(key)
2405                    && key != SEMANTIC_PER_TABLE_INDEX_KEY
2406                    && key != LOGICAL_TABLE_METADATA_KEY
2407                    && key != PHYSICAL_TABLE_METADATA_KEY
2408                    && key != MERGE_MODE_KEY
2409                    && key != "otlp_metric_compat"
2410            });
2411            table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2412            table_options.insert(
2413                store_api::metric_engine_consts::PRIMARY_KEY_ENCODING.to_string(),
2414                "sparse".to_string(),
2415            );
2416        }
2417        AutoCreateTableType::Physical => {
2418            if let Some(append_mode) = ctx.extension(APPEND_MODE_KEY) {
2419                table_options.insert(APPEND_MODE_KEY.to_string(), append_mode.to_string());
2420            }
2421            if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
2422                table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
2423            }
2424        }
2425        // Set append_mode to true for log table.
2426        // because log tables should keep rows with the same ts and tags.
2427        AutoCreateTableType::Log => {
2428            table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2429        }
2430        AutoCreateTableType::LastNonNull => {
2431            if ctx
2432                .extension(APPEND_MODE_KEY)
2433                .is_some_and(|value| value.eq_ignore_ascii_case("true"))
2434            {
2435                table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2436                table_options.insert(MERGE_MODE_KEY.to_string(), "last_row".to_string());
2437            } else if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
2438                table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
2439            } else {
2440                table_options.insert(MERGE_MODE_KEY.to_string(), "last_non_null".to_string());
2441            }
2442        }
2443        AutoCreateTableType::Trace { .. } => {
2444            table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2445        }
2446    }
2447}
2448
2449/// The parsed per-table semantic index: `{schema -> {table -> {key -> value}}}`,
2450/// produced by the OTLP metrics encode path (where one metric can fan out into
2451/// several tables with distinct keys) and the Prometheus remote write v2 path
2452/// (where per-series metadata declares type/unit, and a series may override its
2453/// target schema).
2454pub type PerTableSemanticIndex = BTreeMap<String, BTreeMap<String, BTreeMap<String, String>>>;
2455
2456/// Parses the per-table semantic index off the context extension. Call once per
2457/// create-planning round: a first write creating N tables would otherwise
2458/// re-parse the whole index N times. `None` when the request carries no index
2459/// (logs, traces, Prom RW v1) or it fails to parse.
2460pub fn parse_per_table_semantic_index(ctx: &QueryContextRef) -> Option<PerTableSemanticIndex> {
2461    let raw = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY)?;
2462    match serde_json::from_str(raw) {
2463        Ok(index) => Some(index),
2464        Err(_) => {
2465            warn!("failed to parse semantic per-table index, skipping per-table options");
2466            None
2467        }
2468    }
2469}
2470
2471/// Folds the semantic keys of the table being created into `table_options`.
2472///
2473/// Common keys shared by every table in a request travel as plain semantic
2474/// extensions and are handled by [`fill_table_options_for_create`]; this
2475/// carries only the per-table tail and is applied after it, so a per-table
2476/// value (e.g. `declared` quality) wins. Keys are re-checked against the
2477/// vocabulary defensively.
2478pub fn apply_per_table_semantic_options(
2479    table_options: &mut std::collections::HashMap<String, String>,
2480    index: Option<&PerTableSemanticIndex>,
2481    schema: &str,
2482    table_name: &str,
2483) {
2484    let Some(entry) = index
2485        .and_then(|index| index.get(schema))
2486        .and_then(|tables| tables.get(table_name))
2487    else {
2488        return;
2489    };
2490    for (key, value) in entry {
2491        if is_semantic_option_key(key) && validate_semantic_option(key, value) {
2492            table_options.insert(key.clone(), value.clone());
2493        }
2494    }
2495}
2496
2497pub fn build_create_table_expr(
2498    table: &TableReference,
2499    request_schema: &[ColumnSchema],
2500    engine: &str,
2501) -> Result<CreateTableExpr> {
2502    expr_helper::create_table_expr_by_column_schemas(table, request_schema, engine, None)
2503}
2504
2505/// `QueryContext` extension key the Splunk HEC handler sets (to `"true"`) on its identity
2506/// path to request metadata-first primary-key ordering at table creation. It is absent for
2507/// user-supplied pipelines, so their primary-key order is left untouched.
2508pub const SPLUNK_PK_METADATA_ORDER_KEY: &str = "splunk_pk_metadata_order";
2509
2510/// Moves Splunk's metadata tags (`host`, `source`, `sourcetype`) to the front of the
2511/// primary key, keeping the relative order of the remaining tags.
2512fn reorder_splunk_primary_keys(primary_keys: &mut [String]) {
2513    const LEAD: [&str; 3] = ["host", "source", "sourcetype"];
2514    // Stable sort: `LEAD` columns move to the front in `host`/`source`/`sourcetype` order;
2515    // every other column keeps its existing relative position.
2516    primary_keys.sort_by_key(|name| {
2517        LEAD.iter()
2518            .position(|&lead| lead == name.as_str())
2519            .unwrap_or(LEAD.len())
2520    });
2521}
2522
2523/// Result of `create_or_alter_tables_on_demand`.
2524struct CreateAlterTableResult {
2525    /// table ids of ttl=instant tables.
2526    instant_table_ids: HashSet<TableId>,
2527    /// Table Info of the created tables.
2528    table_infos: HashMap<TableId, Arc<TableInfo>>,
2529}
2530
2531/// Upper bound of rows buffered by detached flow mirror tasks on this frontend.
2532///
2533/// Rows retained by in-flight mirror RPCs count against this bound. If ordinary
2534/// persisted source-table mirroring exceeds it, that mirror batch is best-effort
2535/// dropped while its datanode write continues. Instant-TTL batches fail retryably
2536/// on temporary saturation; oversized batches fail non-retryably so the client
2537/// can reduce the batch size.
2538const MAX_MIRROR_PENDING_ROWS: usize = 1_000_000;
2539
2540/// How often at most each mirror event is reported.
2541const MIRROR_LOG_INTERVAL: Duration = Duration::from_secs(10);
2542const MIRROR_LOG_NEVER_REPORTED: u64 = u64::MAX;
2543
2544struct MirrorLog {
2545    last_log_millis: AtomicU64,
2546    events: AtomicU64,
2547}
2548
2549impl MirrorLog {
2550    const fn new() -> Self {
2551        Self {
2552            last_log_millis: AtomicU64::new(MIRROR_LOG_NEVER_REPORTED),
2553            events: AtomicU64::new(0),
2554        }
2555    }
2556
2557    /// Returns the count since the previous report, including this event, to
2558    /// the single caller that claims the interval.
2559    fn claim_report(&self, now_millis: u64) -> Option<u64> {
2560        self.events.fetch_add(1, Ordering::Relaxed);
2561        let last = self.last_log_millis.load(Ordering::Relaxed);
2562        if last != MIRROR_LOG_NEVER_REPORTED
2563            && now_millis.saturating_sub(last) < MIRROR_LOG_INTERVAL.as_millis() as u64
2564        {
2565            return None;
2566        }
2567        if self
2568            .last_log_millis
2569            .compare_exchange(last, now_millis, Ordering::Relaxed, Ordering::Relaxed)
2570            .is_err()
2571        {
2572            return None;
2573        }
2574        Some(self.events.swap(0, Ordering::Relaxed))
2575    }
2576}
2577
2578static MIRROR_DROP_LOG: MirrorLog = MirrorLog::new();
2579static MIRROR_FAILURE_LOG: MirrorLog = MirrorLog::new();
2580
2581/// Milliseconds since the first call, so the drop log is rate limited on a
2582/// monotonic clock instead of the wall clock, which can jump.
2583fn mirror_drop_log_millis() -> u64 {
2584    static START: LazyLock<Instant> = LazyLock::new(Instant::now);
2585    START.elapsed().as_millis() as u64
2586}
2587
2588struct FlowMirrorTask {
2589    requests: HashMap<Peer, RegionInsertRequests>,
2590}
2591
2592impl FlowMirrorTask {
2593    async fn new(
2594        cache: &TableFlownodeSetCacheRef,
2595        requests: impl Iterator<Item = &RegionInsertRequest>,
2596    ) -> Result<Self> {
2597        let mut src_table_reqs: HashMap<TableId, Option<(Vec<Peer>, RegionInsertRequests)>> =
2598            HashMap::new();
2599
2600        for req in requests {
2601            let table_id = RegionId::from_u64(req.region_id).table_id();
2602            match src_table_reqs.get_mut(&table_id) {
2603                Some(Some((_peers, reqs))) => reqs.requests.push(req.clone()),
2604                // already know this is not source table
2605                Some(None) => continue,
2606                _ => {
2607                    // dedup peers
2608                    let peers = cache
2609                        .get(table_id)
2610                        .await
2611                        .context(RequestInsertsSnafu)?
2612                        .unwrap_or_default()
2613                        .values()
2614                        .cloned()
2615                        .collect::<HashSet<_>>()
2616                        .into_iter()
2617                        .collect::<Vec<_>>();
2618
2619                    if !peers.is_empty() {
2620                        let mut reqs = RegionInsertRequests::default();
2621                        reqs.requests.push(req.clone());
2622                        src_table_reqs.insert(table_id, Some((peers, reqs)));
2623                    } else {
2624                        // insert a empty entry to avoid repeat query
2625                        src_table_reqs.insert(table_id, None);
2626                    }
2627                }
2628            }
2629        }
2630
2631        let mut inserts: HashMap<Peer, RegionInsertRequests> = HashMap::new();
2632
2633        for (_table_id, (peers, reqs)) in src_table_reqs
2634            .into_iter()
2635            .filter_map(|(k, v)| v.map(|v| (k, v)))
2636        {
2637            if peers.len() == 1 {
2638                // fast path, zero copy
2639                inserts
2640                    .entry(peers[0].clone())
2641                    .or_default()
2642                    .requests
2643                    .extend(reqs.requests);
2644                continue;
2645            } else {
2646                // TODO(discord9): need to split requests to multiple flownodes
2647                for flownode in peers {
2648                    inserts
2649                        .entry(flownode.clone())
2650                        .or_default()
2651                        .requests
2652                        .extend(reqs.requests.clone());
2653                }
2654            }
2655        }
2656
2657        Ok(Self { requests: inserts })
2658    }
2659
2660    /// Rows of row data this task keeps alive once detached: the payload it
2661    /// spawns per peer, not the source requests it was built from.
2662    ///
2663    /// A source table mapped to several flownodes clones its requests per peer
2664    /// and every clone owns its own copy of the rows, so reserving from
2665    /// `self.requests` is what the per-peer shares released by [`Self::detach`]
2666    /// sum to. Counting the source requests instead under-reserves whenever a
2667    /// batch holds several requests for the same source table.
2668    fn pending_rows(&self) -> u64 {
2669        self.requests.values().map(region_inserts_rows).sum()
2670    }
2671
2672    fn detach(
2673        self,
2674        node_manager: NodeManagerRef,
2675        mirror_pending_rows: Arc<AtomicU64>,
2676        has_instant_rows: bool,
2677    ) -> Result<()> {
2678        // Reserve the cloned per-peer payload before spawning. A full budget
2679        // drops normal mirrors or rejects instant writes before datanode dispatch.
2680        let num_rows = self.pending_rows();
2681        if num_rows == 0 {
2682            return Ok(());
2683        }
2684        if has_instant_rows && num_rows > MAX_MIRROR_PENDING_ROWS as u64 {
2685            return InvalidInsertRequestSnafu {
2686                reason: format!(
2687                    "flow mirror batch has {num_rows} peer-fan-out rows, exceeding the limit of {}; reduce the batch size",
2688                    MAX_MIRROR_PENDING_ROWS
2689                ),
2690            }
2691            .fail();
2692        }
2693
2694        let pending = reserve_mirror_pending_rows(&mirror_pending_rows, num_rows);
2695        if pending > MAX_MIRROR_PENDING_ROWS as u64 {
2696            release_mirror_pending_rows(&mirror_pending_rows, num_rows);
2697            if has_instant_rows {
2698                return Err(meter_core::collect::WriteRejected::new(
2699                    "flow mirror pending rows limit exceeded for instant-TTL insert",
2700                ))
2701                .context(WriteRejectedSnafu);
2702            }
2703            crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.inc_by(num_rows);
2704            if let Some(events) = MIRROR_DROP_LOG.claim_report(mirror_drop_log_millis()) {
2705                warn!(
2706                    "Flow mirror write dropped: {} pending rows exceeds limit {}, {} drop events since last report",
2707                    pending, MAX_MIRROR_PENDING_ROWS, events
2708                );
2709            }
2710            return Ok(());
2711        }
2712
2713        for (peer, inserts) in self.requests {
2714            // Each spawned task releases exactly its own share of the
2715            // reservation. The shares are computed over the same per-peer
2716            // payload as `pending_rows`, so they sum to the reservation.
2717            let peer_rows = region_inserts_rows(&inserts);
2718            let node_manager = node_manager.clone();
2719            let mirror_pending_rows = mirror_pending_rows.clone();
2720            common_runtime::spawn_global(async move {
2721                let result = node_manager
2722                    .flownode(&peer)
2723                    .await
2724                    .handle_inserts(inserts)
2725                    .await
2726                    .context(RequestInsertsSnafu);
2727
2728                match result {
2729                    Ok(resp) => {
2730                        let affected_rows = resp.affected_rows;
2731                        crate::metrics::DIST_MIRROR_ROW_COUNT.inc_by(affected_rows);
2732                    }
2733                    Err(err) => {
2734                        if let Some(events) =
2735                            MIRROR_FAILURE_LOG.claim_report(mirror_drop_log_millis())
2736                        {
2737                            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);
2738                        }
2739                    }
2740                }
2741                // Release the reservation on both the success and the failure
2742                // path, otherwise a broken flownode would fill the budget forever.
2743                release_mirror_pending_rows(&mirror_pending_rows, peer_rows);
2744            });
2745        }
2746
2747        Ok(())
2748    }
2749}
2750
2751/// Rows of row data carried by `inserts`. A request without a `rows` payload
2752/// carries nothing to keep buffered and contributes zero.
2753fn region_inserts_rows(inserts: &RegionInsertRequests) -> u64 {
2754    inserts
2755        .requests
2756        .iter()
2757        .filter_map(|req| req.rows.as_ref())
2758        .map(|rows| rows.rows.len() as u64)
2759        .sum()
2760}
2761
2762/// Reserves `rows` of the pending mirror budget and returns the resulting
2763/// pending row count. The caller must [`release_mirror_pending_rows`] exactly
2764/// these rows once they are no longer buffered, whether it spawns the batch or
2765/// drops it.
2766fn reserve_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) -> u64 {
2767    let pending = mirror_pending_rows.fetch_add(rows, Ordering::Relaxed) + rows;
2768    // Keep the gauge in step by applying the same increment, never by setting an
2769    // absolute value: another task may move the budget in between.
2770    crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(rows as i64);
2771    pending
2772}
2773
2774/// Releases `rows` of the pending mirror budget.
2775///
2776/// Every release mirrors the reservation it hands back, so a well-behaved
2777/// release never subtracts more than is pending; the saturating update only
2778/// keeps a duplicated release from wrapping the shared budget around.
2779fn release_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) {
2780    let prev = mirror_pending_rows
2781        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |pending| {
2782            Some(pending.saturating_sub(rows))
2783        })
2784        .unwrap_or_else(|prev| prev);
2785    // Subtract from the gauge with the opposite sign of
2786    // [`reserve_mirror_pending_rows`]; reading the budget here and `set`-ting it
2787    // would publish a stale value whenever another task releases concurrently.
2788    let released = rows.min(prev);
2789    if released > 0 {
2790        crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(released as i64);
2791    }
2792}
2793
2794#[cfg(test)]
2795mod tests {
2796    use std::future::Future;
2797    use std::sync::Arc;
2798    use std::sync::atomic::{AtomicU64, Ordering};
2799    use std::task::Poll;
2800
2801    use api::helper::ColumnDataTypeWrapper;
2802    use api::region::RegionResponse;
2803    use api::v1::flow::FlowResponse;
2804    use api::v1::helper::{field_column_schema, time_index_column_schema};
2805    use api::v1::{Row, RowInsertRequest, Rows, Value};
2806    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
2807    use common_error::ext::{ErrorExt, RetryHint};
2808    use common_error::mock::MockError;
2809    use common_error::status_code::StatusCode;
2810    use common_meta::cache::new_table_flownode_set_cache;
2811    use common_meta::cache_invalidator::DummyCacheInvalidator;
2812    use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler;
2813    use common_meta::node_manager::{
2814        Datanode, DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef,
2815    };
2816    use common_meta::test_util::{MockDatanodeManager, MockFlownodeHandler, MockFlownodeManager};
2817    use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD;
2818    use common_query::prelude::{greptime_native_histogram, set_default_prefix};
2819    use datatypes::data_type::ConcreteDataType;
2820    use datatypes::schema::ColumnSchema;
2821    use moka::future::Cache;
2822    use session::context::QueryContext;
2823    use table::TableRef;
2824    use table::dist_table::DummyDataSource;
2825    use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType};
2826    use tokio::sync::{Semaphore, mpsc};
2827
2828    use crate::insert::*;
2829    use crate::test_util::{
2830        create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
2831    };
2832    use crate::write_dispatcher::{TrackedWrite, WriteRetryOptions};
2833
2834    fn make_table_ref_with_schema(
2835        ts_name: &str,
2836        field_name: &str,
2837        field_type: ConcreteDataType,
2838    ) -> TableRef {
2839        let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2840            ColumnSchema::new(
2841                ts_name,
2842                ConcreteDataType::timestamp_millisecond_datatype(),
2843                false,
2844            )
2845            .with_time_index(true),
2846            ColumnSchema::new(field_name, field_type, true),
2847        ])
2848        .unwrap()
2849        .build()
2850        .unwrap();
2851        let meta = TableMetaBuilder::empty()
2852            .schema(Arc::new(schema))
2853            .primary_key_indices(vec![])
2854            .value_indices(vec![1])
2855            .engine("mito")
2856            .next_column_id(0)
2857            .options(Default::default())
2858            .created_on(Default::default())
2859            .build()
2860            .unwrap();
2861        let info = Arc::new(
2862            TableInfoBuilder::default()
2863                .table_id(1)
2864                .table_version(0)
2865                .name("test_table")
2866                .schema_name(DEFAULT_SCHEMA_NAME)
2867                .catalog_name(DEFAULT_CATALOG_NAME)
2868                .desc(None)
2869                .table_type(TableType::Base)
2870                .meta(meta)
2871                .build()
2872                .unwrap(),
2873        );
2874        Arc::new(table::Table::new(
2875            info,
2876            table::metadata::FilterPushDownType::Unsupported,
2877            Arc::new(DummyDataSource),
2878        ))
2879    }
2880
2881    fn make_metric_physical_table_ref_with_time_unit(unit: TimeUnit) -> TableRef {
2882        let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2883            ColumnSchema::new(
2884                greptime_timestamp(),
2885                ConcreteDataType::timestamp_datatype(unit),
2886                false,
2887            )
2888            .with_time_index(true),
2889            ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
2890        ])
2891        .unwrap()
2892        .build()
2893        .unwrap();
2894        let meta = TableMetaBuilder::empty()
2895            .schema(Arc::new(schema))
2896            .primary_key_indices(vec![])
2897            .value_indices(vec![1])
2898            .engine("metric")
2899            .next_column_id(0)
2900            .options(Default::default())
2901            .created_on(Default::default())
2902            .build()
2903            .unwrap();
2904        let info = Arc::new(
2905            TableInfoBuilder::default()
2906                .table_id(1)
2907                .table_version(0)
2908                .name("greptime_physical_table")
2909                .schema_name(DEFAULT_SCHEMA_NAME)
2910                .catalog_name(DEFAULT_CATALOG_NAME)
2911                .desc(None)
2912                .table_type(TableType::Base)
2913                .meta(meta)
2914                .build()
2915                .unwrap(),
2916        );
2917        Arc::new(table::Table::new(
2918            info,
2919            table::metadata::FilterPushDownType::Unsupported,
2920            Arc::new(DummyDataSource),
2921        ))
2922    }
2923
2924    fn ms_row_insert_request(timestamp_ms: i64) -> RowInsertRequest {
2925        ms_row_insert_request_named("my_metric", timestamp_ms)
2926    }
2927
2928    fn ms_row_insert_request_named(table: &str, timestamp_ms: i64) -> RowInsertRequest {
2929        RowInsertRequest {
2930            table_name: table.to_string(),
2931            rows: Some(Rows {
2932                schema: vec![
2933                    time_index_column_schema(
2934                        greptime_timestamp(),
2935                        ColumnDataType::TimestampMillisecond,
2936                    ),
2937                    field_column_schema(greptime_value(), ColumnDataType::Float64),
2938                ],
2939                rows: vec![Row {
2940                    values: vec![
2941                        Value {
2942                            value_data: Some(ValueData::TimestampMillisecondValue(timestamp_ms)),
2943                        },
2944                        Value {
2945                            value_data: Some(ValueData::F64Value(1.0)),
2946                        },
2947                    ],
2948                }],
2949            }),
2950        }
2951    }
2952
2953    /// Converts the time index column of each request to `target_unit`.
2954    fn convert_row_insert_requests_time_unit(
2955        requests: &mut RowInsertRequests,
2956        target_unit: TimeUnit,
2957    ) -> Result<()> {
2958        for request in &mut requests.inserts {
2959            let Some(rows) = request.rows.as_mut() else {
2960                continue;
2961            };
2962            convert_rows_time_unit(rows, target_unit)?;
2963        }
2964        Ok(())
2965    }
2966
2967    #[test]
2968    fn test_convert_row_insert_requests_time_unit_noop_when_matching() {
2969        let mut requests = RowInsertRequests {
2970            inserts: vec![ms_row_insert_request(123)],
2971        };
2972        convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Millisecond).unwrap();
2973        let rows = requests.inserts[0].rows.as_ref().unwrap();
2974        assert_eq!(
2975            rows.schema[0].datatype,
2976            ColumnDataType::TimestampMillisecond as i32
2977        );
2978        assert!(matches!(
2979            rows.rows[0].values[0].value_data,
2980            Some(ValueData::TimestampMillisecondValue(123))
2981        ));
2982    }
2983
2984    #[test]
2985    fn test_convert_row_insert_requests_time_unit_widens_losslessly() {
2986        let mut requests = RowInsertRequests {
2987            inserts: vec![ms_row_insert_request(123)],
2988        };
2989        convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Microsecond).unwrap();
2990        let rows = requests.inserts[0].rows.as_ref().unwrap();
2991        assert_eq!(
2992            rows.schema[0].datatype,
2993            ColumnDataType::TimestampMicrosecond as i32
2994        );
2995        assert!(matches!(
2996            rows.rows[0].values[0].value_data,
2997            Some(ValueData::TimestampMicrosecondValue(123_000))
2998        ));
2999    }
3000
3001    #[test]
3002    fn test_convert_row_insert_requests_time_unit_truncates_on_narrowing() {
3003        // 123_456_789 ns floors to 123_456 us and 123 ms; negative values
3004        // floor towards negative infinity, matching `Timestamp::convert_to`.
3005        let requests = |value_ns: i64| RowInsertRequests {
3006            inserts: vec![RowInsertRequest {
3007                table_name: "my_metric".to_string(),
3008                rows: Some(Rows {
3009                    schema: vec![
3010                        time_index_column_schema(
3011                            greptime_timestamp(),
3012                            ColumnDataType::TimestampNanosecond,
3013                        ),
3014                        field_column_schema(greptime_value(), ColumnDataType::Float64),
3015                    ],
3016                    rows: vec![Row {
3017                        values: vec![
3018                            Value {
3019                                value_data: Some(ValueData::TimestampNanosecondValue(value_ns)),
3020                            },
3021                            Value {
3022                                value_data: Some(ValueData::F64Value(1.0)),
3023                            },
3024                        ],
3025                    }],
3026                }),
3027            }],
3028        };
3029
3030        let mut reqs = requests(123_456_789);
3031        convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
3032        let rows = reqs.inserts[0].rows.as_ref().unwrap();
3033        assert!(matches!(
3034            rows.rows[0].values[0].value_data,
3035            Some(ValueData::TimestampMicrosecondValue(123_456))
3036        ));
3037
3038        let mut reqs = requests(123_456_789);
3039        convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Millisecond).unwrap();
3040        let rows = reqs.inserts[0].rows.as_ref().unwrap();
3041        assert!(matches!(
3042            rows.rows[0].values[0].value_data,
3043            Some(ValueData::TimestampMillisecondValue(123))
3044        ));
3045
3046        let mut reqs = requests(-123_456_789);
3047        convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
3048        let rows = reqs.inserts[0].rows.as_ref().unwrap();
3049        assert!(matches!(
3050            rows.rows[0].values[0].value_data,
3051            Some(ValueData::TimestampMicrosecondValue(-123_457))
3052        ));
3053    }
3054
3055    #[test]
3056    fn test_convert_row_insert_requests_time_unit_overflow() {
3057        let mut requests = RowInsertRequests {
3058            inserts: vec![ms_row_insert_request(i64::MAX)],
3059        };
3060        let err =
3061            convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Nanosecond).unwrap_err();
3062        assert!(err.to_string().contains("overflows"), "{err}");
3063    }
3064
3065    #[test]
3066    fn test_table_time_index_unit() {
3067        assert_eq!(
3068            table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
3069                TimeUnit::Microsecond
3070            )),
3071            Some(TimeUnit::Microsecond)
3072        );
3073        assert_eq!(
3074            table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
3075                TimeUnit::Millisecond
3076            )),
3077            Some(TimeUnit::Millisecond)
3078        );
3079    }
3080
3081    #[tokio::test]
3082    async fn test_accommodate_existing_schema_and_reject_kind_changes() {
3083        let ts_name = "my_ts";
3084        let field_name = "my_field";
3085        let table =
3086            make_table_ref_with_schema(ts_name, field_name, ConcreteDataType::float64_datatype());
3087
3088        // The request uses different names for timestamp and field columns
3089        let mut req = RowInsertRequest {
3090            table_name: "test_table".to_string(),
3091            rows: Some(Rows {
3092                schema: vec![
3093                    time_index_column_schema("ts_wrong", ColumnDataType::TimestampMillisecond),
3094                    field_column_schema("field_wrong", ColumnDataType::Float64),
3095                ],
3096                rows: vec![api::v1::Row {
3097                    values: vec![Value::default(), Value::default()],
3098                }],
3099            }),
3100        };
3101        let ctx = Arc::new(QueryContext::with(
3102            DEFAULT_CATALOG_NAME,
3103            DEFAULT_SCHEMA_NAME,
3104        ));
3105
3106        let kv_backend = prepare_mocked_backend().await;
3107        let inserter = Inserter::new(
3108            catalog::memory::MemoryCatalogManager::new(),
3109            create_partition_rule_manager(kv_backend.clone()).await,
3110            Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3111            Arc::new(new_table_flownode_set_cache(
3112                String::new(),
3113                Cache::new(100),
3114                kv_backend.clone(),
3115            )),
3116            true,
3117        );
3118        // Do not apply an absent-table plan to a table that appeared concurrently.
3119        assert!(
3120            inserter
3121                .get_alter_table_expr_on_demand(&mut req, &table, &ctx, false, false, false)
3122                .unwrap()
3123                .is_none()
3124        );
3125        let alter_expr = inserter
3126            .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true, true)
3127            .unwrap();
3128        assert!(alter_expr.is_none());
3129
3130        // The request's schema should have updated names for timestamp and field columns
3131        let req_schema = req.rows.as_ref().unwrap().schema.clone();
3132        assert_eq!(req_schema[0].column_name, ts_name);
3133        assert_eq!(req_schema[1].column_name, field_name);
3134
3135        let (datatype, datatype_extension) =
3136            ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
3137                .unwrap()
3138                .into_parts();
3139        let mut histogram_req = RowInsertRequest {
3140            table_name: "test_table".to_string(),
3141            rows: Some(Rows {
3142                schema: vec![
3143                    time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
3144                    api::v1::ColumnSchema {
3145                        column_name: greptime_native_histogram().to_string(),
3146                        datatype: datatype as i32,
3147                        semantic_type: SemanticType::Field as i32,
3148                        datatype_extension,
3149                        options: None,
3150                    },
3151                ],
3152                rows: vec![],
3153            }),
3154        };
3155        let error = inserter
3156            .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, false, true, true)
3157            .unwrap_err();
3158        assert!(
3159            error
3160                .to_string()
3161                .contains("cannot mix native histogram and float sample fields")
3162        );
3163
3164        let histogram_table = make_table_ref_with_schema(
3165            "ts",
3166            greptime_native_histogram(),
3167            native_histogram_value_type().clone(),
3168        );
3169        let mut sample_req = RowInsertRequest {
3170            table_name: "test_table".to_string(),
3171            rows: Some(Rows {
3172                schema: vec![
3173                    time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
3174                    field_column_schema(greptime_value(), ColumnDataType::Float64),
3175                ],
3176                rows: vec![],
3177            }),
3178        };
3179        let error = inserter
3180            .get_alter_table_expr_on_demand(
3181                &mut sample_req,
3182                &histogram_table,
3183                &ctx,
3184                false,
3185                true,
3186                true,
3187            )
3188            .unwrap_err();
3189        assert!(
3190            error
3191                .to_string()
3192                .contains("cannot mix native histogram and float sample fields")
3193        );
3194    }
3195
3196    #[test]
3197    fn test_native_histogram_detection_survives_prefix_change() {
3198        set_default_prefix(Some("custom")).unwrap();
3199        let table = make_table_ref_with_schema(
3200            "custom_timestamp",
3201            NATIVE_HISTOGRAM_FIELD,
3202            native_histogram_value_type().clone(),
3203        );
3204        let (datatype, datatype_extension) =
3205            ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
3206                .unwrap()
3207                .into_parts();
3208        let request_schema = [api::v1::ColumnSchema {
3209            column_name: greptime_native_histogram().to_string(),
3210            datatype: datatype as i32,
3211            semantic_type: SemanticType::Field as i32,
3212            datatype_extension,
3213            options: None,
3214        }];
3215
3216        assert!(request_is_native_histogram(&request_schema));
3217        assert!(table_is_native_histogram(&table));
3218    }
3219
3220    #[derive(Clone)]
3221    struct GatedDatanodeHandler {
3222        started: mpsc::UnboundedSender<u32>,
3223        completed: mpsc::UnboundedSender<u32>,
3224        release: Arc<Semaphore>,
3225    }
3226
3227    #[async_trait::async_trait]
3228    impl Datanode for GatedDatanodeHandler {
3229        async fn handle_request(
3230            &self,
3231            request: Arc<api::v1::region::RegionRequest>,
3232        ) -> common_meta::error::Result<RegionResponse> {
3233            let Some(api::v1::region::region_request::Body::Inserts(requests)) = &request.body
3234            else {
3235                panic!("expected inserts");
3236            };
3237            let region = RegionId::from(requests.requests[0].region_id).region_number();
3238            self.started.send(region).unwrap();
3239            self.release.acquire().await.unwrap().forget();
3240            self.completed.send(region).unwrap();
3241            if region == 2 {
3242                Err(common_meta::error::Error::retry_later(MockError::new(
3243                    StatusCode::RegionBusy,
3244                )))
3245            } else {
3246                Ok(RegionResponse::new(1))
3247            }
3248        }
3249
3250        async fn handle_query(
3251            &self,
3252            _request: common_query::request::QueryRequest,
3253        ) -> common_meta::error::Result<common_recordbatch::SendableRecordBatchStream> {
3254            unreachable!()
3255        }
3256    }
3257
3258    #[async_trait::async_trait]
3259    impl DatanodeManager for GatedDatanodeHandler {
3260        async fn datanode(&self, _peer: &Peer) -> DatanodeRef {
3261            Arc::new(self.clone())
3262        }
3263    }
3264
3265    #[async_trait::async_trait]
3266    impl FlownodeManager for GatedDatanodeHandler {
3267        async fn flownode(&self, _peer: &Peer) -> FlownodeRef {
3268            unreachable!()
3269        }
3270    }
3271
3272    #[tokio::test]
3273    async fn test_insert_cancellation_drains_inflight_writes() {
3274        for confirmed in [false, true] {
3275            for max_retries in [0, 2] {
3276                let mut inserter = batcher_test_inserter().await;
3277                let (started, mut starts) = mpsc::unbounded_channel();
3278                let (completed, mut completions) = mpsc::unbounded_channel();
3279                let release = Arc::new(Semaphore::new(0));
3280                let node_manager = Arc::new(GatedDatanodeHandler {
3281                    started,
3282                    completed,
3283                    release: release.clone(),
3284                });
3285                inserter.node_manager = node_manager.clone();
3286                let dispatcher = Arc::new(
3287                    WriteDispatcher::new(node_manager, inserter.partition_manager.clone())
3288                        .with_retry(
3289                            WriteRetryOptions {
3290                                max_retries,
3291                                initial_backoff: Duration::ZERO,
3292                                max_backoff: Duration::ZERO,
3293                                ..Default::default()
3294                            },
3295                            inserter.catalog_manager.clone(),
3296                            Arc::new(DummyCacheInvalidator),
3297                        ),
3298                );
3299                let inserter = inserter.with_write_dispatcher(dispatcher);
3300                let table_infos = HashMap::from_iter([(
3301                    1,
3302                    Arc::new(new_test_table_info(1, "table_1", [1].into_iter())),
3303                )]);
3304                let counter = Arc::new(session::write_progress::WriteProgress::default());
3305                let mut ctx = QueryContext::arc().fork();
3306                ctx.set_write_progress(counter.clone());
3307                let ctx = Arc::new(ctx);
3308                let ctx = if confirmed {
3309                    Arc::new(ctx.with_confirmed_writes())
3310                } else {
3311                    ctx
3312                };
3313                let requests = InstantAndNormalInsertRequests {
3314                    normal_requests: RegionInsertRequests {
3315                        requests: [1, 2]
3316                            .map(|region| RegionInsertRequest {
3317                                region_id: RegionId::new(1, region).as_u64(),
3318                                rows: Some(Rows {
3319                                    schema: vec![],
3320                                    rows: vec![Row { values: vec![] }],
3321                                }),
3322                                ..Default::default()
3323                            })
3324                            .to_vec(),
3325                    },
3326                    instant_requests: RegionInsertRequests::default(),
3327                };
3328                let caller =
3329                    tokio::spawn(
3330                        async move { inserter.do_request(requests, &table_infos, &ctx).await },
3331                    );
3332                for _ in 0..2 {
3333                    tokio::time::timeout(Duration::from_secs(5), starts.recv())
3334                        .await
3335                        .unwrap()
3336                        .unwrap();
3337                }
3338                caller.abort();
3339                assert!(caller.await.unwrap_err().is_cancelled());
3340                release.add_permits(2);
3341                // Closing the channel proves the detached dispatcher has finished.
3342                let mut completed = tokio::time::timeout(Duration::from_secs(5), async {
3343                    let mut regions = Vec::new();
3344                    while let Some(region) = completions.recv().await {
3345                        regions.push(region);
3346                    }
3347                    regions
3348                })
3349                .await
3350                .unwrap();
3351                completed.sort_unstable();
3352                assert_eq!(
3353                    completed,
3354                    vec![1, 2],
3355                    "confirmed={confirmed}, retries={max_retries}"
3356                );
3357                assert!(
3358                    starts.try_recv().is_err(),
3359                    "cancelled caller started another retry"
3360                );
3361                assert_eq!(counter.written(), 1);
3362            }
3363        }
3364    }
3365
3366    #[tokio::test]
3367    async fn test_tracked_cancellation_stops_backoff() {
3368        let inserter = batcher_test_inserter().await;
3369        let (started, mut starts) = mpsc::unbounded_channel();
3370        let (completed, mut completions) = mpsc::unbounded_channel();
3371        let node_manager = Arc::new(GatedDatanodeHandler {
3372            started,
3373            completed,
3374            release: Arc::new(Semaphore::new(2)),
3375        });
3376        let dispatcher = WriteDispatcher::new(node_manager, inserter.partition_manager.clone())
3377            .with_retry(
3378                WriteRetryOptions {
3379                    max_retries: 2,
3380                    initial_backoff: Duration::from_secs(60),
3381                    max_backoff: Duration::from_secs(60),
3382                    ..Default::default()
3383                },
3384                inserter.catalog_manager.clone(),
3385                Arc::new(DummyCacheInvalidator),
3386            );
3387        let factory = RegionRequestFactory::new(RegionRequestHeader::default());
3388        let requests = [1, 2]
3389            .into_iter()
3390            .map(|region| TrackedWrite {
3391                request: factory.build_insert(RegionInsertRequests {
3392                    requests: vec![RegionInsertRequest {
3393                        region_id: RegionId::new(1, region).as_u64(),
3394                        ..Default::default()
3395                    }],
3396                }),
3397                owners: vec![(region - 1) as usize],
3398            })
3399            .collect();
3400        let (writes, progress) = dispatcher.prepare_tracked(requests, 2).await;
3401        let cancellation = CancellationToken::new();
3402        let execute = dispatcher.execute_tracked(writes, progress, None, Some(&cancellation));
3403        tokio::pin!(execute);
3404        // Wait for both sends to complete before cancelling the backoff.
3405        tokio::time::timeout(
3406            Duration::from_secs(1),
3407            future::poll_fn(|cx| {
3408                assert!(execute.as_mut().poll(cx).is_pending());
3409                if completions.len() == 2 {
3410                    Poll::Ready(())
3411                } else {
3412                    Poll::Pending
3413                }
3414            }),
3415        )
3416        .await
3417        .unwrap();
3418        for region in [1, 2] {
3419            assert_eq!(starts.try_recv().unwrap(), region);
3420            assert_eq!(completions.try_recv().unwrap(), region);
3421        }
3422        cancellation.cancel();
3423        let progress = tokio::time::timeout(Duration::from_secs(1), execute)
3424            .await
3425            .unwrap();
3426        assert_eq!(progress[0].affected_rows, 1);
3427        assert!(progress[0].error.is_none());
3428        assert_eq!(progress[1].affected_rows, 0);
3429        assert_eq!(
3430            progress[1].error.as_ref().unwrap().retry_hint(),
3431            RetryHint::Retryable
3432        );
3433        assert!(starts.try_recv().is_err());
3434    }
3435
3436    // Keep global meter registration in one test, isolated by nextest's per-test process.
3437    #[tokio::test]
3438    async fn test_write_meter_admission() {
3439        use std::cell::Cell;
3440        use std::sync::Mutex;
3441        use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
3442
3443        use api::region::RegionResponse;
3444        use api::v1::region::region_request::Body;
3445        use arrow::array::{Int32Array, TimestampMillisecondArray};
3446        use arrow::record_batch::RecordBatch;
3447        use bytes::Bytes;
3448        use common_error::ext::{ErrorExt, RetryHint};
3449        use common_error::status_code::StatusCode;
3450        use common_grpc::flight::{FlightEncoder, FlightMessage};
3451        use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher;
3452        use futures::future::BoxFuture;
3453        use meter_core::ItemCalculator;
3454        use meter_core::collect::{Collect, WriteRejected};
3455        use meter_core::data::MeterRecord;
3456        use meter_core::global::global_registry;
3457        use session::context::Channel;
3458
3459        const CATALOG: &str = "write_meter_test";
3460
3461        #[derive(Default)]
3462        struct Meter {
3463            reject: AtomicBool,
3464            attempts: Mutex<Vec<MeterRecord>>,
3465            accepted_value: AtomicU64,
3466        }
3467
3468        impl Collect for Meter {
3469            fn on_write(
3470                &self,
3471                record: MeterRecord,
3472            ) -> BoxFuture<'_, std::result::Result<(), WriteRejected>> {
3473                Box::pin(async move {
3474                    if record.catalog != CATALOG {
3475                        return Ok(());
3476                    }
3477                    let value = record.value;
3478                    self.attempts.lock().unwrap().push(record);
3479                    if self.reject.load(Ordering::Relaxed) {
3480                        return Err(WriteRejected::new("database row quota exhausted"));
3481                    }
3482                    self.accepted_value.fetch_add(value, Ordering::Relaxed);
3483                    Ok(())
3484                })
3485            }
3486
3487            fn on_read(&self, _: MeterRecord) {}
3488        }
3489
3490        impl ItemCalculator<InstantAndNormalInsertRequests> for Meter {
3491            fn calc(&self, _: &InstantAndNormalInsertRequests) -> u64 {
3492                17
3493            }
3494        }
3495
3496        let kv_backend = prepare_mocked_backend().await;
3497        let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
3498        let (sender, mut dispatched) = tokio::sync::mpsc::channel(16);
3499        let watcher = DatanodeWatcher::new(sender).with_handler(|_, request| {
3500            let rows = match request.body.unwrap() {
3501                Body::Inserts(requests) => requests
3502                    .requests
3503                    .iter()
3504                    .filter_map(|request| request.rows.as_ref())
3505                    .map(|rows| rows.rows.len())
3506                    .sum(),
3507                // The bulk batches below each contain two rows.
3508                Body::BulkInsert(_) => 2,
3509                body => panic!("unexpected request: {body:?}"),
3510            };
3511            Ok(RegionResponse::new(rows))
3512        });
3513        let flow_cache = Cache::new(100);
3514        let inserter = Inserter::new(
3515            catalog::memory::MemoryCatalogManager::new(),
3516            partition_manager,
3517            Arc::new(MockDatanodeManager::new(watcher)),
3518            Arc::new(new_table_flownode_set_cache(
3519                String::new(),
3520                flow_cache.clone(),
3521                kv_backend,
3522            )),
3523            true,
3524        );
3525        let mut table_info = new_test_table_info(1, "table_1", [1].into_iter());
3526        table_info.catalog_name = CATALOG.to_string();
3527        table_info.schema_name = "target_db".to_string();
3528        let table_info = Arc::new(table_info);
3529        let table_infos = HashMap::from_iter([(1, table_info.clone())]);
3530        let ctx = Arc::new(QueryContext::with_channel(
3531            DEFAULT_CATALOG_NAME,
3532            DEFAULT_SCHEMA_NAME,
3533            Channel::Postgres,
3534        ));
3535        let rows_request = || {
3536            let build = |num_rows| RegionInsertRequests {
3537                requests: vec![RegionInsertRequest {
3538                    region_id: RegionId::new(1, 1).as_u64(),
3539                    rows: Some(Rows {
3540                        schema: vec![],
3541                        rows: vec![api::v1::Row { values: vec![] }; num_rows],
3542                    }),
3543                    ..Default::default()
3544                }],
3545            };
3546            InstantAndNormalInsertRequests {
3547                normal_requests: build(3),
3548                instant_requests: build(2),
3549            }
3550        };
3551
3552        // No collector or calculator: ordinary OSS insertion still succeeds.
3553        let output = inserter
3554            .do_request(rows_request(), &table_infos, &ctx)
3555            .await
3556            .unwrap();
3557        assert_eq!(output.meta.cost, 0);
3558        assert!(matches!(output.data, OutputData::AffectedRows(3)));
3559        dispatched.try_recv().unwrap();
3560
3561        let meter = Arc::new(Meter::default());
3562        global_registry().set_collector(meter.clone());
3563        global_registry().register_calculator(meter.clone());
3564        // The dependency controls noop mode; exercise both builds with this test.
3565        let enabled = Cell::new(false);
3566        write_meter!({
3567            enabled.set(true);
3568            MeterRecord::new("probe".into(), "probe".into(), 0, 0, 0)
3569        })
3570        .await
3571        .unwrap();
3572        let enabled = enabled.get();
3573
3574        let output = inserter
3575            .do_request(rows_request(), &table_infos, &ctx)
3576            .await
3577            .unwrap();
3578        assert_eq!(output.meta.cost, if enabled { 17 } else { 0 });
3579        assert!(matches!(output.data, OutputData::AffectedRows(3)));
3580        dispatched.try_recv().unwrap();
3581        assert!(flow_cache.contains_key(&1));
3582        flow_cache.invalidate_all();
3583        meter.reject.store(true, Ordering::Relaxed);
3584        let result = inserter
3585            .do_request(rows_request(), &table_infos, &ctx)
3586            .await;
3587        if enabled {
3588            let error = result.unwrap_err();
3589            assert_eq!(error.status_code(), StatusCode::RateLimited);
3590            assert_eq!(error.retry_hint(), RetryHint::Retryable);
3591            assert!(error.to_string().contains("database row quota exhausted"));
3592            assert!(dispatched.try_recv().is_err());
3593            assert!(
3594                !flow_cache.contains_key(&1),
3595                "rejected write reached flow mirroring"
3596            );
3597            assert_eq!(meter.accepted_value.load(Ordering::Relaxed), 17);
3598            let attempts = meter.attempts.lock().unwrap();
3599            assert_eq!(attempts.len(), 2);
3600            for record in attempts.iter() {
3601                assert_eq!(record.catalog, CATALOG);
3602                assert_eq!(record.schema, "target_db");
3603                assert_eq!(
3604                    (record.rows, record.value, record.source),
3605                    (5, 17, Channel::Postgres as u8)
3606                );
3607            }
3608        } else {
3609            assert_eq!(result.unwrap().meta.cost, 0);
3610            dispatched.try_recv().unwrap();
3611            assert!(meter.attempts.lock().unwrap().is_empty());
3612        }
3613        meter.attempts.lock().unwrap().clear();
3614
3615        // Rejection must also precede admission to the table batcher's queue.
3616        if enabled {
3617            let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3618            let rows = Rows {
3619                schema: vec![
3620                    api::v1::helper::tag_column_schema("a", ColumnDataType::Int32),
3621                    time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
3622                    field_column_schema("b", ColumnDataType::Int32),
3623                ],
3624                rows: vec![api::v1::Row {
3625                    values: vec![
3626                        api::v1::value::ValueData::I32Value(60).into(),
3627                        Value {
3628                            value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue(
3629                                0,
3630                            )),
3631                        },
3632                        api::v1::value::ValueData::I32Value(0).into(),
3633                    ],
3634                }],
3635            };
3636            let error = inserter
3637                .submit_table_rows(rows, table_info.clone(), ctx.clone(), &batcher)
3638                .await
3639                .unwrap_err();
3640            assert_eq!(error.status_code(), StatusCode::RateLimited);
3641            let mut attempts = meter.attempts.lock().unwrap();
3642            assert_eq!(attempts.len(), 1);
3643            assert_eq!(attempts[0].catalog, CATALOG);
3644            assert_eq!(attempts[0].schema, "target_db");
3645            assert_eq!(attempts[0].rows, 1);
3646            attempts.clear();
3647        }
3648
3649        let table = Arc::new(table::Table::new(
3650            table_info.clone(),
3651            table::metadata::FilterPushDownType::Unsupported,
3652            Arc::new(DummyDataSource),
3653        ));
3654        let batch = RecordBatch::try_new(
3655            table_info.meta.schema.arrow_schema().clone(),
3656            vec![
3657                Arc::new(Int32Array::from(vec![60, 70])),
3658                Arc::new(TimestampMillisecondArray::from(vec![0, 1])),
3659                Arc::new(Int32Array::from(vec![0, 0])),
3660            ],
3661        )
3662        .unwrap();
3663        let bulk_insert = |batch: RecordBatch| {
3664            let flight_data = FlightEncoder::default()
3665                .encode(FlightMessage::RecordBatch(batch.clone()))
3666                .into_iter()
3667                .next()
3668                .unwrap();
3669            inserter.handle_bulk_insert(
3670                table.clone(),
3671                flight_data,
3672                batch,
3673                Bytes::new(),
3674                false,
3675                Channel::Grpc,
3676            )
3677        };
3678        // Empty batches bypass admission even while the collector rejects.
3679        assert_eq!(bulk_insert(batch.slice(0, 0)).await.unwrap(), 0);
3680        assert!(meter.attempts.lock().unwrap().is_empty());
3681        assert!(dispatched.try_recv().is_err());
3682
3683        // A collector change between batches takes effect on the very next batch.
3684        for reject in [false, true, false] {
3685            meter.reject.store(reject, Ordering::Relaxed);
3686            let result = bulk_insert(batch.clone()).await;
3687            if enabled && reject {
3688                let error = result.unwrap_err();
3689                assert_eq!(error.status_code(), StatusCode::RateLimited);
3690                assert_eq!(error.retry_hint(), RetryHint::Retryable);
3691                assert!(dispatched.try_recv().is_err());
3692            } else {
3693                assert_eq!(result.unwrap(), 2);
3694                dispatched.try_recv().unwrap();
3695            }
3696        }
3697        {
3698            let attempts = meter.attempts.lock().unwrap();
3699            assert_eq!(attempts.len(), if enabled { 3 } else { 0 });
3700            for record in attempts.iter() {
3701                assert_eq!(record.catalog, CATALOG);
3702                assert_eq!(record.schema, "target_db");
3703                assert_eq!(
3704                    (record.rows, record.value, record.source),
3705                    (2, 0, Channel::Grpc as u8)
3706                );
3707            }
3708        }
3709        assert_eq!(
3710            meter.accepted_value.load(Ordering::Relaxed),
3711            if enabled { 17 } else { 0 }
3712        );
3713
3714        // Aggregate repeated database targets and retain admission across nested
3715        // batching without changing the caller's reusable context.
3716        meter.attempts.lock().unwrap().clear();
3717        let original = Arc::new(QueryContext::with_channel(
3718            CATALOG,
3719            "a",
3720            Channel::Prometheus,
3721        ));
3722        let mut batches = ["a", "b", "a"].map(|schema| {
3723            let ctx = if schema == "a" {
3724                original.clone()
3725            } else {
3726                Arc::new(QueryContext::with_channel(
3727                    CATALOG,
3728                    schema,
3729                    Channel::Prometheus,
3730                ))
3731            };
3732            (
3733                ctx,
3734                RowInsertRequests {
3735                    inserts: vec![RowInsertRequest {
3736                        table_name: "data".into(),
3737                        rows: Some(Rows {
3738                            schema: vec![],
3739                            rows: vec![api::v1::Row::default(); 2],
3740                        }),
3741                    }],
3742                },
3743            )
3744        });
3745        admit_row_insert_batches(&mut batches).await.unwrap();
3746        admit_row_insert_batches(&mut batches).await.unwrap();
3747        assert_eq!(original.write_rows_to_admit(CATALOG, "a", 4), 4);
3748        for (ctx, _) in &batches {
3749            assert_eq!(
3750                ctx.write_rows_to_admit(CATALOG, &ctx.current_schema(), 2),
3751                0
3752            );
3753            assert_eq!(ctx.channel(), Channel::Prometheus);
3754        }
3755        {
3756            let attempts = meter.attempts.lock().unwrap();
3757            let totals = attempts
3758                .iter()
3759                .map(|r| (r.schema.as_str(), r.rows, r.value))
3760                .collect::<Vec<_>>();
3761            assert_eq!(
3762                totals,
3763                if enabled {
3764                    vec![("a", 4, 0), ("b", 2, 0)]
3765                } else {
3766                    vec![]
3767                }
3768            );
3769        }
3770        meter.attempts.lock().unwrap().clear();
3771        meter.reject.store(false, Ordering::Relaxed);
3772        let ctx = Arc::new(QueryContext::with_channel(
3773            CATALOG,
3774            "logical",
3775            Channel::Otlp,
3776        ));
3777        let admitted = admit_write(2, &ctx).await.unwrap();
3778        let mut requests = RowInsertRequests {
3779            inserts: vec![RowInsertRequest {
3780                table_name: "metric".to_string(),
3781                rows: Some(Rows {
3782                    schema: vec![],
3783                    rows: vec![api::v1::Row::default(); 2],
3784                }),
3785            }],
3786        };
3787        let original = requests.clone();
3788        let cost = Inserter::meter_row_inserts(&mut requests, &admitted)
3789            .await
3790            .unwrap();
3791        assert_eq!(cost, if enabled { 17 } else { 0 });
3792        assert_eq!(requests, original);
3793        let records = meter
3794            .attempts
3795            .lock()
3796            .unwrap()
3797            .iter()
3798            .map(|record| (record.rows, record.value))
3799            .collect::<Vec<_>>();
3800        assert_eq!(
3801            records,
3802            if enabled {
3803                vec![(2, 0), (0, 17)]
3804            } else {
3805                vec![]
3806            }
3807        );
3808        meter.reject.store(true, Ordering::Relaxed);
3809        let result = Inserter::meter_row_inserts(&mut requests, &admitted).await;
3810        if enabled {
3811            assert_eq!(result.unwrap_err().status_code(), StatusCode::RateLimited);
3812        } else {
3813            assert_eq!(result.unwrap(), 0);
3814        }
3815        assert_eq!(requests, original);
3816    }
3817
3818    #[test]
3819    fn test_trace_create_options_copy_compaction_hints() {
3820        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3821        ctx.set_extension("compaction.twcs.time_window", "1h");
3822        ctx.set_extension("compaction.twcs.max_output_file_size", "256MB");
3823        ctx.set_extension("compaction.twcs.remote_compaction", "true");
3824        ctx.set_extension("compaction.unknown", "ignored");
3825        let ctx = Arc::new(ctx);
3826        let mut table_options = Default::default();
3827
3828        fill_table_options_for_create(
3829            &mut table_options,
3830            &AutoCreateTableType::Trace {
3831                alter_existing: true,
3832            },
3833            &ctx,
3834        );
3835
3836        assert_eq!(
3837            Some("twcs"),
3838            table_options.get(COMPACTION_TYPE).map(String::as_str)
3839        );
3840        assert_eq!(
3841            Some("1h"),
3842            table_options
3843                .get("compaction.twcs.time_window")
3844                .map(String::as_str)
3845        );
3846        assert_eq!(
3847            Some("256MB"),
3848            table_options
3849                .get("compaction.twcs.max_output_file_size")
3850                .map(String::as_str)
3851        );
3852        assert_eq!(
3853            Some("true"),
3854            table_options
3855                .get("compaction.twcs.remote_compaction")
3856                .map(String::as_str)
3857        );
3858        assert!(!table_options.contains_key("compaction.unknown"));
3859    }
3860
3861    #[test]
3862    fn test_exemplar_create_options_isolate_metric_context() {
3863        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3864        for (key, value) in [
3865            (APPEND_MODE_KEY, "false"),
3866            (MERGE_MODE_KEY, "last_non_null"),
3867            (PHYSICAL_TABLE_METADATA_KEY, "true"),
3868            (LOGICAL_TABLE_METADATA_KEY, "physical_metrics"),
3869            ("otlp_metric_compat", "prometheus"),
3870            (table::requests::SEMANTIC_SIGNAL_TYPE, "metric"),
3871            (SEMANTIC_PER_TABLE_INDEX_KEY, "{}"),
3872            ("ttl", "7d"),
3873            ("compaction.twcs.time_window", "1h"),
3874        ] {
3875            ctx.set_extension(key, value);
3876        }
3877        let mut options = Default::default();
3878        fill_table_options_for_create(&mut options, &AutoCreateTableType::Exemplar, &Arc::new(ctx));
3879        assert_eq!(
3880            options,
3881            std::collections::HashMap::from([
3882                (APPEND_MODE_KEY.into(), "true".into()),
3883                (
3884                    store_api::metric_engine_consts::PRIMARY_KEY_ENCODING.into(),
3885                    "sparse".into()
3886                ),
3887                ("ttl".into(), "7d".into()),
3888                (COMPACTION_TYPE.into(), "twcs".into()),
3889                ("compaction.twcs.time_window".into(), "1h".into()),
3890            ])
3891        );
3892    }
3893
3894    #[test]
3895    fn test_skip_wal_does_not_change_table_options() {
3896        check_skip_wal_does_not_change_table_options(false);
3897        check_skip_wal_does_not_change_table_options(true);
3898    }
3899
3900    fn check_skip_wal_does_not_change_table_options(skip_wal: bool) {
3901        let ctx = Arc::new(QueryContext::with(
3902            DEFAULT_CATALOG_NAME,
3903            DEFAULT_SCHEMA_NAME,
3904        ));
3905        ctx.set_skip_wal(skip_wal);
3906        let mut options = Default::default();
3907        fill_table_options_for_create(&mut options, &AutoCreateTableType::Physical, &ctx);
3908        assert!(!options.contains_key(session::hints::INSERT_SKIP_WAL_HINT));
3909        assert!(!options.contains_key("skip_wal"));
3910    }
3911
3912    #[test]
3913    fn test_last_non_null_create_options_preserve_default_without_append_mode() {
3914        let ctx = Arc::new(QueryContext::with(
3915            DEFAULT_CATALOG_NAME,
3916            DEFAULT_SCHEMA_NAME,
3917        ));
3918        let mut table_options = Default::default();
3919
3920        fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3921
3922        assert_eq!(
3923            Some("last_non_null"),
3924            table_options.get(MERGE_MODE_KEY).map(String::as_str)
3925        );
3926        assert!(!table_options.contains_key(APPEND_MODE_KEY));
3927    }
3928
3929    #[test]
3930    fn test_fill_table_options_copies_semantic_extensions() {
3931        use table::requests::{
3932            SEMANTIC_METRIC_TYPE, SEMANTIC_PER_TABLE_INDEX_KEY, SEMANTIC_SIGNAL_TYPE,
3933            SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, SOURCE_OPENTELEMETRY,
3934        };
3935
3936        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3937        ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
3938        ctx.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
3939        ctx.set_extension(SEMANTIC_SOURCE_VERSION, "2.0");
3940        ctx.set_extension(SEMANTIC_METRIC_TYPE, "bogus");
3941        // The internal transport key must NOT be copied into table options.
3942        ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, "{}");
3943        let ctx = Arc::new(ctx);
3944        let mut table_options = Default::default();
3945
3946        fill_table_options_for_create(&mut table_options, &AutoCreateTableType::Physical, &ctx);
3947
3948        assert_eq!(
3949            Some(SIGNAL_TYPE_METRIC),
3950            table_options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str)
3951        );
3952        assert_eq!(
3953            Some(SOURCE_OPENTELEMETRY),
3954            table_options.get(SEMANTIC_SOURCE).map(String::as_str)
3955        );
3956        assert_eq!(
3957            Some("2.0"),
3958            table_options
3959                .get(SEMANTIC_SOURCE_VERSION)
3960                .map(String::as_str)
3961        );
3962        assert!(!table_options.contains_key(SEMANTIC_METRIC_TYPE));
3963        assert!(!table_options.contains_key(SEMANTIC_PER_TABLE_INDEX_KEY));
3964    }
3965
3966    #[test]
3967    fn test_apply_per_table_semantic_options() {
3968        use table::requests::{
3969            SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, SEMANTIC_PER_TABLE_INDEX_KEY,
3970        };
3971
3972        let index = format!(
3973            r#"{{
3974            "{DEFAULT_SCHEMA_NAME}": {{
3975                "http_requests_total": {{
3976                    "greptime.semantic.metric.type": "counter",
3977                    "greptime.semantic.metric.unit": "By",
3978                    "greptime.semantic.metric.type_BOGUS": "x"
3979                }},
3980                "other_table": {{
3981                    "greptime.semantic.metric.type": "gauge"
3982                }}
3983            }},
3984            "other_schema": {{
3985                "http_requests_total": {{
3986                    "greptime.semantic.metric.type": "gauge"
3987                }}
3988            }}
3989        }}"#
3990        );
3991        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3992        ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index);
3993        let ctx = Arc::new(ctx);
3994
3995        let index = parse_per_table_semantic_index(&ctx);
3996        assert!(index.is_some());
3997        let index = index.as_ref();
3998
3999        let mut table_options = std::collections::HashMap::new();
4000        apply_per_table_semantic_options(
4001            &mut table_options,
4002            index,
4003            DEFAULT_SCHEMA_NAME,
4004            "http_requests_total",
4005        );
4006        // The write schema's entry applies — not other_schema's `gauge`.
4007        assert_eq!(
4008            table_options.get(SEMANTIC_METRIC_TYPE).map(String::as_str),
4009            Some("counter")
4010        );
4011        assert_eq!(
4012            table_options.get(SEMANTIC_METRIC_UNIT).map(String::as_str),
4013            Some("By")
4014        );
4015        // The unknown key is rejected by the vocabulary check; other tables' keys
4016        // never appear.
4017        assert!(!table_options.contains_key("greptime.semantic.metric.type_BOGUS"));
4018        assert_eq!(table_options.len(), 2);
4019
4020        let mut empty = std::collections::HashMap::new();
4021        apply_per_table_semantic_options(&mut empty, index, DEFAULT_SCHEMA_NAME, "not_in_index");
4022        assert!(empty.is_empty());
4023
4024        // A schema with no entry is a no-op even when the table name matches
4025        // elsewhere.
4026        let mut opts = std::collections::HashMap::new();
4027        apply_per_table_semantic_options(
4028            &mut opts,
4029            index,
4030            "schema_without_entry",
4031            "http_requests_total",
4032        );
4033        assert!(opts.is_empty());
4034
4035        // No extension at all parses to no index (e.g. logs / Prom RW v1).
4036        let bare = Arc::new(QueryContext::with(
4037            DEFAULT_CATALOG_NAME,
4038            DEFAULT_SCHEMA_NAME,
4039        ));
4040        assert!(parse_per_table_semantic_index(&bare).is_none());
4041        let mut opts = std::collections::HashMap::new();
4042        apply_per_table_semantic_options(
4043            &mut opts,
4044            None,
4045            DEFAULT_SCHEMA_NAME,
4046            "http_requests_total",
4047        );
4048        assert!(opts.is_empty());
4049    }
4050
4051    #[test]
4052    fn test_last_non_null_create_options_preserve_default_with_append_mode_false() {
4053        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
4054        ctx.set_extension(APPEND_MODE_KEY, "false");
4055        let ctx = Arc::new(ctx);
4056        let mut table_options = Default::default();
4057
4058        fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
4059
4060        assert!(!table_options.contains_key(APPEND_MODE_KEY));
4061        assert_eq!(
4062            Some("last_non_null"),
4063            table_options.get(MERGE_MODE_KEY).map(String::as_str)
4064        );
4065    }
4066
4067    #[test]
4068    fn test_last_non_null_create_options_use_configured_merge_mode() {
4069        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
4070        ctx.set_extension(MERGE_MODE_KEY, "last_row");
4071        let ctx = Arc::new(ctx);
4072        let mut table_options = Default::default();
4073
4074        fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
4075
4076        assert_eq!(
4077            Some("last_row"),
4078            table_options.get(MERGE_MODE_KEY).map(String::as_str)
4079        );
4080        assert!(!table_options.contains_key(APPEND_MODE_KEY));
4081    }
4082
4083    #[test]
4084    fn test_last_non_null_create_options_use_last_row_with_append_mode_true() {
4085        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
4086        ctx.set_extension(APPEND_MODE_KEY, "true");
4087        let ctx = Arc::new(ctx);
4088        let mut table_options = Default::default();
4089
4090        fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
4091
4092        assert_eq!(
4093            Some("true"),
4094            table_options.get(APPEND_MODE_KEY).map(String::as_str)
4095        );
4096        assert_eq!(
4097            Some("last_row"),
4098            table_options.get(MERGE_MODE_KEY).map(String::as_str)
4099        );
4100    }
4101
4102    struct UnexpectedBatcher;
4103
4104    #[async_trait::async_trait]
4105    impl PendingRowsBatcher for UnexpectedBatcher {
4106        async fn acquire(&self) -> Result<Arc<tokio::sync::OwnedSemaphorePermit>> {
4107            panic!("empty writes must not acquire batch admission")
4108        }
4109
4110        async fn submit(
4111            &self,
4112            _table_info: TableInfoRef,
4113            _batch: arrow::record_batch::RecordBatch,
4114            _ctx: QueryContextRef,
4115            _permit: Arc<tokio::sync::OwnedSemaphorePermit>,
4116        ) -> Result<Output> {
4117            panic!("empty writes must not submit a batch")
4118        }
4119    }
4120
4121    async fn batcher_test_inserter() -> Inserter {
4122        let kv_backend = prepare_mocked_backend().await;
4123        Inserter::new(
4124            catalog::memory::MemoryCatalogManager::new(),
4125            create_partition_rule_manager(kv_backend.clone()).await,
4126            Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
4127            Arc::new(new_table_flownode_set_cache(
4128                String::new(),
4129                Cache::new(100),
4130                kv_backend,
4131            )),
4132            true,
4133        )
4134    }
4135
4136    #[tokio::test]
4137    async fn test_logical_batcher_eligibility() {
4138        use catalog::RegisterTableRequest;
4139        use catalog::memory::MemoryCatalogManager;
4140        use common_meta::instruction::{CacheIdent, CreateFlow};
4141        use common_meta::kv_backend::KvBackendRef;
4142        use common_meta::kv_backend::memory::MemoryKvBackend;
4143        use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder};
4144        let requests = RowInsertRequests {
4145            inserts: vec![RowInsertRequest {
4146                table_name: "test_table".to_string(),
4147                rows: None,
4148            }],
4149        };
4150        let original =
4151            make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
4152                .table_info();
4153        for case in [
4154            "eligible",
4155            "physical",
4156            "ordinary",
4157            "instant",
4158            "disabled",
4159            "hint",
4160            "flow",
4161            "required_tag",
4162            "default",
4163        ] {
4164            let mut info = (*original).clone();
4165            info.meta.engine = METRIC_ENGINE_NAME.to_string();
4166            info.meta.options.extra_options.insert(
4167                LOGICAL_TABLE_METADATA_KEY.to_string(),
4168                "physical".to_string(),
4169            );
4170            let mut ctx = QueryContext::arc().fork();
4171            match case {
4172                "physical" => {
4173                    info.meta
4174                        .options
4175                        .extra_options
4176                        .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "other".to_string());
4177                }
4178                "ordinary" => info.meta.engine = "mito".to_string(),
4179                "instant" => info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant),
4180                "hint" => ctx.set_extension(AUTO_CREATE_TABLE_KEY, "false"),
4181                "required_tag" | "default" => {
4182                    let mut columns = info.meta.schema.column_schemas().to_vec();
4183                    let mut tag = ColumnSchema::new(
4184                        "tag",
4185                        ConcreteDataType::string_datatype(),
4186                        case != "required_tag",
4187                    );
4188                    if case == "default" {
4189                        tag = tag
4190                            .with_default_constraint(Some(ColumnDefaultConstraint::null_value()))
4191                            .unwrap();
4192                    }
4193                    columns.push(tag);
4194                    info.meta.schema = Arc::new(
4195                        SchemaBuilder::try_from_columns(columns)
4196                            .unwrap()
4197                            .build()
4198                            .unwrap(),
4199                    );
4200                    info.meta.primary_key_indices = vec![2];
4201                }
4202                _ => {}
4203            }
4204            let catalog = MemoryCatalogManager::with_default_setup();
4205            let table = Arc::new(table::Table::new(
4206                Arc::new(info),
4207                table::metadata::FilterPushDownType::Unsupported,
4208                Arc::new(DummyDataSource),
4209            ));
4210            catalog
4211                .register_table_sync(RegisterTableRequest {
4212                    catalog: DEFAULT_CATALOG_NAME.to_string(),
4213                    schema: DEFAULT_SCHEMA_NAME.to_string(),
4214                    table_name: "test_table".to_string(),
4215                    table_id: 1,
4216                    table,
4217                })
4218                .unwrap();
4219            let mut inserter = batcher_test_inserter().await;
4220            inserter.catalog_manager = catalog;
4221            inserter.auto_create_table = case != "disabled";
4222            let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::default());
4223            inserter.table_flownode_set_cache = Arc::new(new_table_flownode_set_cache(
4224                String::new(),
4225                Cache::new(10),
4226                kv_backend,
4227            ));
4228            if case == "flow" {
4229                inserter
4230                    .table_flownode_set_cache
4231                    .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
4232                        flow_id: 1,
4233                        source_table_ids: vec![1],
4234                        partition_to_peer_mapping: vec![(0, Peer::empty(1))],
4235                    })])
4236                    .await
4237                    .unwrap();
4238            }
4239            assert_eq!(
4240                inserter
4241                    .can_batch_metric_rows(&requests, &Arc::new(ctx), "physical")
4242                    .await
4243                    .unwrap(),
4244                case == "eligible",
4245                "{case}"
4246            );
4247        }
4248    }
4249
4250    #[tokio::test]
4251    async fn test_batcher_meter_preserves_request() {
4252        let mut requests = RowInsertRequests {
4253            inserts: vec![RowInsertRequest {
4254                table_name: "sample".to_string(),
4255                rows: Some(Rows {
4256                    schema: vec![],
4257                    rows: vec![api::v1::Row {
4258                        values: vec![Value {
4259                            value_data: Some(api::v1::value::ValueData::F64Value(1.5)),
4260                        }],
4261                    }],
4262                }),
4263            }],
4264        };
4265        let expected = requests.clone();
4266        Inserter::meter_row_inserts(&mut requests, &QueryContext::arc())
4267            .await
4268            .unwrap();
4269        assert_eq!(requests, expected);
4270    }
4271
4272    #[tokio::test]
4273    async fn test_instant_table_bypasses_batcher() {
4274        let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
4275        let inserter = batcher_test_inserter()
4276            .await
4277            .with_pending_rows_batcher(Some(batcher));
4278        let mut ctx = session::context::QueryContextBuilder::default().build();
4279        ctx.set_batching_enabled(true);
4280        let ctx = Arc::new(ctx);
4281        let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
4282            .table_info();
4283        assert!(inserter.table_batcher(&table, &ctx).is_some());
4284        let mut instant = (*table).clone();
4285        instant.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
4286        assert!(inserter.table_batcher(&Arc::new(instant), &ctx).is_none());
4287    }
4288
4289    #[tokio::test]
4290    async fn test_empty_prepared_rows_skip_batcher() {
4291        let inserter = batcher_test_inserter().await;
4292        let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
4293            .table_info();
4294        let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
4295        let ctx = QueryContext::arc();
4296        let output = inserter
4297            .submit_table_rows(
4298                Rows {
4299                    schema: vec![],
4300                    rows: vec![],
4301                },
4302                table.clone(),
4303                ctx.clone(),
4304                &batcher,
4305            )
4306            .await
4307            .unwrap();
4308        assert!(matches!(output.data, OutputData::AffectedRows(0)));
4309        let output = inserter
4310            .submit_pending_rows(
4311                RowInsertRequests {
4312                    inserts: vec![RowInsertRequest {
4313                        table_name: table.name.clone(),
4314                        rows: None,
4315                    }],
4316                },
4317                HashMap::from_iter([(table.table_id(), table)]),
4318                ctx,
4319                &batcher,
4320            )
4321            .await
4322            .unwrap();
4323        assert!(matches!(output.data, OutputData::AffectedRows(0)));
4324    }
4325
4326    /// Flownode handler that keeps mirror inserts in flight until the test
4327    /// releases the gate, so the pending mirror budget can be observed
4328    /// deterministically.
4329    #[derive(Clone)]
4330    struct GatedFlownodeHandler {
4331        gate: Arc<tokio::sync::Semaphore>,
4332        affected_rows: u64,
4333    }
4334
4335    #[async_trait::async_trait]
4336    impl MockFlownodeHandler for GatedFlownodeHandler {
4337        async fn handle_inserts(
4338            &self,
4339            _peer: &Peer,
4340            _requests: api::v1::region::InsertRequests,
4341        ) -> common_meta::error::Result<FlowResponse> {
4342            let _permit = self.gate.acquire().await.unwrap();
4343            Ok(FlowResponse {
4344                affected_rows: self.affected_rows,
4345                ..Default::default()
4346            })
4347        }
4348    }
4349
4350    fn flownode_peer() -> Peer {
4351        Peer {
4352            id: 1,
4353            addr: "127.0.0.1:4001".to_string(),
4354        }
4355    }
4356
4357    #[derive(Clone)]
4358    struct RecordingNodeManager {
4359        datanode_dispatch: tokio::sync::mpsc::UnboundedSender<usize>,
4360        flownode_dispatch: tokio::sync::mpsc::UnboundedSender<Peer>,
4361        gates: Arc<std::sync::Mutex<HashMap<u64, Arc<tokio::sync::Semaphore>>>>,
4362        fail_flownode: bool,
4363    }
4364
4365    #[async_trait::async_trait]
4366    impl common_meta::node_manager::DatanodeManager for RecordingNodeManager {
4367        async fn datanode(&self, peer: &Peer) -> common_meta::node_manager::DatanodeRef {
4368            Arc::new(RecordingNode {
4369                peer: peer.clone(),
4370                manager: self.clone(),
4371            })
4372        }
4373    }
4374
4375    #[async_trait::async_trait]
4376    impl common_meta::node_manager::FlownodeManager for RecordingNodeManager {
4377        async fn flownode(&self, peer: &Peer) -> common_meta::node_manager::FlownodeRef {
4378            Arc::new(RecordingNode {
4379                peer: peer.clone(),
4380                manager: self.clone(),
4381            })
4382        }
4383    }
4384
4385    struct RecordingNode {
4386        peer: Peer,
4387        manager: RecordingNodeManager,
4388    }
4389
4390    #[async_trait::async_trait]
4391    impl common_meta::node_manager::Datanode for RecordingNode {
4392        async fn handle_request(
4393            &self,
4394            _request: Arc<api::v1::region::RegionRequest>,
4395        ) -> common_meta::error::Result<api::region::RegionResponse> {
4396            let rows = match _request.body.as_ref() {
4397                Some(api::v1::region::region_request::Body::Inserts(requests)) => requests
4398                    .requests
4399                    .iter()
4400                    .filter_map(|request| request.rows.as_ref())
4401                    .map(|rows| rows.rows.len())
4402                    .sum(),
4403                _ => 0,
4404            };
4405            let _ = self.manager.datanode_dispatch.send(rows);
4406            Ok(api::region::RegionResponse::new(rows))
4407        }
4408
4409        async fn handle_query(
4410            &self,
4411            _request: common_query::request::QueryRequest,
4412        ) -> common_meta::error::Result<common_recordbatch::SendableRecordBatchStream> {
4413            unreachable!()
4414        }
4415    }
4416
4417    #[async_trait::async_trait]
4418    impl common_meta::node_manager::Flownode for RecordingNode {
4419        async fn handle(
4420            &self,
4421            _request: api::v1::flow::FlowRequest,
4422        ) -> common_meta::error::Result<FlowResponse> {
4423            unreachable!()
4424        }
4425
4426        async fn handle_inserts(
4427            &self,
4428            _request: api::v1::region::InsertRequests,
4429        ) -> common_meta::error::Result<FlowResponse> {
4430            let _ = self.manager.flownode_dispatch.send(self.peer.clone());
4431            if self.manager.fail_flownode {
4432                return Err(common_meta::error::UnexpectedSnafu {
4433                    err_msg: "test flownode failure".to_string(),
4434                }
4435                .build());
4436            }
4437            let gate = self
4438                .manager
4439                .gates
4440                .lock()
4441                .unwrap()
4442                .get(&self.peer.id)
4443                .cloned();
4444            if let Some(gate) = gate {
4445                let _permit = gate.acquire().await.unwrap();
4446            }
4447            Ok(FlowResponse::default())
4448        }
4449
4450        async fn handle_mark_window_dirty(
4451            &self,
4452            _request: api::v1::flow::DirtyWindowRequests,
4453        ) -> common_meta::error::Result<FlowResponse> {
4454            unreachable!()
4455        }
4456    }
4457
4458    fn mirror_requests(peer: &Peer, num_rows: usize) -> HashMap<Peer, RegionInsertRequests> {
4459        HashMap::from_iter([(
4460            peer.clone(),
4461            RegionInsertRequests {
4462                requests: vec![RegionInsertRequest {
4463                    region_id: RegionId::new(1, 1).as_u64(),
4464                    rows: Some(Rows {
4465                        schema: vec![],
4466                        rows: vec![api::v1::Row { values: vec![] }; num_rows],
4467                    }),
4468                    ..Default::default()
4469                }],
4470            },
4471        )])
4472    }
4473
4474    /// Serializes the mirror metric tests below: the metrics are process-wide
4475    /// singletons, so these tests would otherwise observe each other's updates
4476    /// when the test binary runs tests in parallel (they are isolated per
4477    /// process only under nextest).
4478    static MIRROR_METRIC_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
4479
4480    async fn lock_mirror_metrics() -> tokio::sync::MutexGuard<'static, ()> {
4481        MIRROR_METRIC_TEST_LOCK.lock().await
4482    }
4483
4484    // The tests below share the process-wide mirror metrics, so they take
4485    // `MIRROR_METRIC_TEST_LOCK` and assert on the difference their own budget
4486    // publishes instead of on absolute values.
4487    #[tokio::test]
4488    async fn flow_mirror_saturation_preserves_instant_and_normal_contracts() {
4489        use common_error::ext::{ErrorExt, RetryHint};
4490        use common_error::status_code::StatusCode;
4491
4492        let _guard = lock_mirror_metrics().await;
4493        let kv_backend = prepare_mocked_backend().await;
4494        let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
4495        let (datanode_tx, mut datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4496        let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4497        let node_manager = Arc::new(RecordingNodeManager {
4498            datanode_dispatch: datanode_tx,
4499            flownode_dispatch: flownode_tx,
4500            gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
4501            fail_flownode: false,
4502        });
4503        let flow_cache = Cache::new(10);
4504        let flow_cache_backend = prepare_mocked_backend().await;
4505        let inserter = Inserter::new(
4506            catalog::memory::MemoryCatalogManager::new(),
4507            partition_manager,
4508            node_manager,
4509            Arc::new(new_table_flownode_set_cache(
4510                String::new(),
4511                flow_cache.clone(),
4512                flow_cache_backend,
4513            )),
4514            true,
4515        );
4516        let peer = flownode_peer();
4517        inserter
4518            .table_flownode_set_cache
4519            .invalidate(&[common_meta::instruction::CacheIdent::CreateFlow(
4520                common_meta::instruction::CreateFlow {
4521                    flow_id: 1,
4522                    source_table_ids: vec![1, 2],
4523                    partition_to_peer_mapping: vec![(0, peer)],
4524                },
4525            )])
4526            .await
4527            .unwrap();
4528        let mut normal_info = new_test_table_info(1, "normal_table", [1].into_iter());
4529        normal_info.catalog_name = DEFAULT_CATALOG_NAME.to_string();
4530        normal_info.schema_name = DEFAULT_SCHEMA_NAME.to_string();
4531        let mut instant_info = new_test_table_info(2, "instant_table", [1].into_iter());
4532        instant_info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
4533        let table_infos =
4534            HashMap::from_iter([(1, Arc::new(normal_info)), (2, Arc::new(instant_info))]);
4535        let ctx = Arc::new(QueryContext::with(
4536            DEFAULT_CATALOG_NAME,
4537            DEFAULT_SCHEMA_NAME,
4538        ));
4539        let insert = |table_id, num_rows| RegionInsertRequests {
4540            requests: vec![RegionInsertRequest {
4541                region_id: RegionId::new(table_id, 1).as_u64(),
4542                rows: Some(Rows {
4543                    schema: vec![],
4544                    rows: vec![api::v1::Row::default(); num_rows],
4545                }),
4546                ..Default::default()
4547            }],
4548        };
4549        let make_request = |normal, instant| InstantAndNormalInsertRequests {
4550            normal_requests: if normal == 0 {
4551                RegionInsertRequests::default()
4552            } else {
4553                insert(1, normal)
4554            },
4555            instant_requests: if instant == 0 {
4556                RegionInsertRequests::default()
4557            } else {
4558                insert(2, instant)
4559            },
4560        };
4561
4562        let pending = inserter.mirror_pending_rows.clone();
4563        let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4564        crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
4565        pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
4566        let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4567        for request in [make_request(0, 1), make_request(1, 1)] {
4568            let error = inserter
4569                .do_request(request, &table_infos, &ctx)
4570                .await
4571                .unwrap_err();
4572            assert_eq!(error.status_code(), StatusCode::RateLimited);
4573            assert_eq!(error.retry_hint(), RetryHint::Retryable);
4574            assert!(datanode_rx.try_recv().is_err());
4575            assert!(flownode_rx.try_recv().is_err());
4576            assert_eq!(
4577                pending.load(Ordering::Relaxed),
4578                MAX_MIRROR_PENDING_ROWS as u64
4579            );
4580        }
4581        assert_eq!(
4582            crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4583            dropped_before
4584        );
4585
4586        for pending_before in [0, MAX_MIRROR_PENDING_ROWS as u64] {
4587            if pending_before == 0 {
4588                pending.store(0, Ordering::Relaxed);
4589                crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
4590            }
4591            let gauge_before_rejection = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4592            for (normal, instant) in [
4593                (0, MAX_MIRROR_PENDING_ROWS + 1),
4594                (MAX_MIRROR_PENDING_ROWS / 2, MAX_MIRROR_PENDING_ROWS / 2 + 1),
4595            ] {
4596                let error = inserter
4597                    .do_request(make_request(normal, instant), &table_infos, &ctx)
4598                    .await
4599                    .unwrap_err();
4600                assert_eq!(error.status_code(), StatusCode::InvalidArguments);
4601                assert_eq!(error.retry_hint(), RetryHint::NonRetryable);
4602                let message = error.to_string();
4603                assert!(message.contains("1000001"), "{message}");
4604                assert!(message.contains("1000000"), "{message}");
4605                assert!(message.contains("reduce the batch size"), "{message}");
4606                assert!(datanode_rx.try_recv().is_err());
4607                assert!(flownode_rx.try_recv().is_err());
4608                assert_eq!(pending.load(Ordering::Relaxed), pending_before);
4609                assert_eq!(
4610                    crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4611                    gauge_before_rejection
4612                );
4613                assert_eq!(
4614                    crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4615                    dropped_before
4616                );
4617            }
4618            if pending_before == 0 {
4619                crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
4620                pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
4621            }
4622        }
4623
4624        let result = inserter
4625            .do_request(make_request(1, 0), &table_infos, &ctx)
4626            .await
4627            .unwrap();
4628        assert!(matches!(result.data, OutputData::AffectedRows(1)));
4629        assert_eq!(1, datanode_rx.try_recv().unwrap());
4630        assert!(flownode_rx.try_recv().is_err());
4631        assert_eq!(
4632            crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4633            dropped_before + 1
4634        );
4635        pending.store(0, Ordering::Relaxed);
4636        crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
4637        assert_eq!(
4638            crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4639            gauge_before
4640        );
4641    }
4642
4643    #[tokio::test]
4644    async fn flow_mirror_new_counts_regions_and_peer_clones() {
4645        use common_meta::instruction::{CacheIdent, CreateFlow};
4646
4647        let peer1 = flownode_peer();
4648        let peer2 = Peer {
4649            id: 2,
4650            addr: "127.0.0.1:4002".to_string(),
4651        };
4652        let cache_backend = prepare_mocked_backend().await;
4653        let cache = Arc::new(new_table_flownode_set_cache(
4654            String::new(),
4655            Cache::new(10),
4656            cache_backend,
4657        ));
4658        cache
4659            .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
4660                flow_id: 1,
4661                source_table_ids: vec![1],
4662                partition_to_peer_mapping: vec![(0, peer1), (1, peer2)],
4663            })])
4664            .await
4665            .unwrap();
4666        let request = |region_id| RegionInsertRequest {
4667            region_id,
4668            rows: Some(Rows {
4669                schema: vec![],
4670                rows: vec![api::v1::Row { values: vec![] }; 3],
4671            }),
4672            ..Default::default()
4673        };
4674        let requests = [
4675            request(RegionId::new(1, 0).as_u64()),
4676            request(RegionId::new(1, 1).as_u64()),
4677        ];
4678        let task = FlowMirrorTask::new(&cache, requests.iter()).await.unwrap();
4679        assert_eq!(12, task.pending_rows());
4680        assert_eq!(2, task.requests.len());
4681    }
4682
4683    #[tokio::test]
4684    async fn flow_mirror_dropped_when_pending_exceeds_limit() {
4685        let _guard = lock_mirror_metrics().await;
4686        let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64));
4687        let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4688        let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4689        let num_rows = 100;
4690
4691        let task = FlowMirrorTask {
4692            requests: mirror_requests(&flownode_peer(), num_rows),
4693        };
4694        // The budget is already full, so nothing may be spawned for this batch.
4695        task.detach(
4696            Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
4697            pending.clone(),
4698            false,
4699        )
4700        .unwrap();
4701
4702        assert_eq!(
4703            num_rows as u64,
4704            crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get() - dropped_before
4705        );
4706        // The dropped batch must not leak its reservation.
4707        assert_eq!(
4708            MAX_MIRROR_PENDING_ROWS as u64,
4709            pending.load(Ordering::Relaxed)
4710        );
4711        // The reservation taken before the limit check is released again, so the
4712        // gauge nets out and keeps matching the pending budget.
4713        assert_eq!(
4714            gauge_before,
4715            crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4716            "a dropped batch must not change the pending gauge"
4717        );
4718
4719        let task = FlowMirrorTask {
4720            requests: mirror_requests(&flownode_peer(), num_rows),
4721        };
4722        let error = task
4723            .detach(
4724                Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
4725                pending.clone(),
4726                true,
4727            )
4728            .err()
4729            .unwrap();
4730        use common_error::ext::{ErrorExt, RetryHint};
4731        use common_error::status_code::StatusCode;
4732        assert_eq!(error.status_code(), StatusCode::RateLimited);
4733        assert_eq!(error.retry_hint(), RetryHint::Retryable);
4734        assert_eq!(
4735            crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4736            dropped_before + num_rows as u64,
4737            "rejected instant rows are not successful best-effort drops"
4738        );
4739        assert_eq!(
4740            MAX_MIRROR_PENDING_ROWS as u64,
4741            pending.load(Ordering::Relaxed)
4742        );
4743        assert_eq!(
4744            gauge_before,
4745            crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get()
4746        );
4747    }
4748
4749    /// A source table mapped to several flownodes clones its requests per peer,
4750    /// so the reservation must cover every clone and each spawned task must hand
4751    /// back exactly its own share; otherwise the shared budget drifts.
4752    #[tokio::test]
4753    async fn flow_mirror_reserves_and_releases_per_peer_shares() {
4754        let _guard = lock_mirror_metrics().await;
4755        let per_peer_rows = 10;
4756        let total_rows = (2 * per_peer_rows) as u64;
4757        // Exactly enough room for the whole batch, cloned rows included.
4758        let start = MAX_MIRROR_PENDING_ROWS as u64 - total_rows;
4759        let pending = Arc::new(AtomicU64::new(start));
4760        let gate = Arc::new(tokio::sync::Semaphore::new(0));
4761        let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4762            gate: gate.clone(),
4763            affected_rows: per_peer_rows as u64,
4764        }));
4765
4766        let mut requests = mirror_requests(&flownode_peer(), per_peer_rows);
4767        requests.extend(mirror_requests(
4768            &Peer {
4769                id: 2,
4770                addr: "127.0.0.1:4002".to_string(),
4771            },
4772            per_peer_rows,
4773        ));
4774
4775        let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4776        FlowMirrorTask { requests }
4777            .detach(node_manager, pending.clone(), false)
4778            .unwrap();
4779
4780        // Both per-peer payloads are counted, so the budget is exactly full
4781        // instead of over-reserved by the duplicated rows.
4782        assert_eq!(
4783            MAX_MIRROR_PENDING_ROWS as u64,
4784            pending.load(Ordering::Relaxed)
4785        );
4786
4787        // Let both tasks finish and check that the budget lands back on its
4788        // starting value: over-releasing a shared reservation would saturate it.
4789        // The gauge is checked too because `release_mirror_pending_rows` moves
4790        // the budget before the gauge; exiting on the budget alone could leak a
4791        // pending `gauge.sub` into the next test holding the metrics lock.
4792        gate.add_permits(2);
4793        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4794        while pending.load(Ordering::Relaxed) != start
4795            || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4796        {
4797            assert!(
4798                std::time::Instant::now() < deadline,
4799                "the mirror tasks must release exactly their own shares"
4800            );
4801            std::thread::sleep(std::time::Duration::from_millis(1));
4802        }
4803    }
4804
4805    #[tokio::test]
4806    async fn flow_mirror_fast_peer_releases_only_its_share_before_drop() {
4807        let _guard = lock_mirror_metrics().await;
4808        let peer1 = flownode_peer();
4809        let peer2 = Peer {
4810            id: 2,
4811            addr: "127.0.0.1:4002".to_string(),
4812        };
4813        let gate1 = Arc::new(tokio::sync::Semaphore::new(0));
4814        let gate2 = Arc::new(tokio::sync::Semaphore::new(0));
4815        let gates = Arc::new(std::sync::Mutex::new(HashMap::from_iter([
4816            (peer1.id, gate1.clone()),
4817            (peer2.id, gate2.clone()),
4818        ])));
4819        let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4820        let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4821        let node_manager = Arc::new(RecordingNodeManager {
4822            datanode_dispatch: datanode_tx,
4823            flownode_dispatch: flownode_tx,
4824            gates,
4825            fail_flownode: false,
4826        });
4827        let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64 - 20));
4828        let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4829        let mut requests = mirror_requests(&peer1, 10);
4830        requests.extend(mirror_requests(&peer2, 10));
4831        let task = FlowMirrorTask { requests };
4832        task.detach(node_manager.clone(), pending.clone(), false)
4833            .unwrap();
4834        let first = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4835            .await
4836            .unwrap()
4837            .unwrap();
4838        let second = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4839            .await
4840            .unwrap()
4841            .unwrap();
4842        assert_ne!(first.id, second.id);
4843        let fast_gate = if first.id == peer1.id {
4844            gate1.clone()
4845        } else {
4846            gate2.clone()
4847        };
4848        let slow_gate = if first.id == peer1.id {
4849            gate2.clone()
4850        } else {
4851            gate1.clone()
4852        };
4853        fast_gate.add_permits(1);
4854        let deadline = std::time::Instant::now() + Duration::from_secs(10);
4855        while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 10 {
4856            assert!(
4857                std::time::Instant::now() < deadline,
4858                "fast peer did not release"
4859            );
4860            std::thread::sleep(Duration::from_millis(1));
4861        }
4862        let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4863        FlowMirrorTask {
4864            requests: mirror_requests(&peer1, 11),
4865        }
4866        .detach(node_manager, pending.clone(), false)
4867        .unwrap();
4868        assert_eq!(
4869            crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4870            dropped_before + 11
4871        );
4872        assert!(flownode_rx.try_recv().is_err());
4873        slow_gate.add_permits(1);
4874        while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 20
4875            || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4876        {
4877            assert!(
4878                std::time::Instant::now() < deadline,
4879                "slow peer did not release"
4880            );
4881            std::thread::sleep(Duration::from_millis(1));
4882        }
4883    }
4884
4885    #[tokio::test]
4886    async fn flow_mirror_failure_releases_budget_and_gauge() {
4887        let _guard = lock_mirror_metrics().await;
4888        let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4889        let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4890        let node_manager = Arc::new(RecordingNodeManager {
4891            datanode_dispatch: datanode_tx,
4892            flownode_dispatch: flownode_tx,
4893            gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
4894            fail_flownode: true,
4895        });
4896        let pending = Arc::new(AtomicU64::new(0));
4897        let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4898        FlowMirrorTask {
4899            requests: mirror_requests(&flownode_peer(), 5),
4900        }
4901        .detach(node_manager, pending.clone(), false)
4902        .unwrap();
4903        assert_eq!(
4904            tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4905                .await
4906                .unwrap()
4907                .unwrap()
4908                .id,
4909            flownode_peer().id
4910        );
4911        let deadline = std::time::Instant::now() + Duration::from_secs(10);
4912        while pending.load(Ordering::Relaxed) != 0
4913            || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4914        {
4915            assert!(
4916                std::time::Instant::now() < deadline,
4917                "failed task leaked reservation"
4918            );
4919            std::thread::sleep(Duration::from_millis(1));
4920        }
4921    }
4922
4923    #[tokio::test]
4924    async fn flow_mirror_pending_counter_consistent() {
4925        let _guard = lock_mirror_metrics().await;
4926        let pending = Arc::new(AtomicU64::new(0));
4927        let gate = Arc::new(tokio::sync::Semaphore::new(0));
4928        let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4929            gate: gate.clone(),
4930            affected_rows: 7,
4931        }));
4932
4933        // A batch without rows short-circuits and leaves the budget untouched.
4934        let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4935        FlowMirrorTask {
4936            requests: HashMap::new(),
4937        }
4938        .detach(node_manager.clone(), pending.clone(), false)
4939        .unwrap();
4940        assert_eq!(0, pending.load(Ordering::Relaxed));
4941        assert_eq!(
4942            gauge_before,
4943            crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4944            "an empty batch must not touch the pending gauge"
4945        );
4946
4947        let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4948        let num_rows = 7;
4949
4950        FlowMirrorTask {
4951            requests: mirror_requests(&flownode_peer(), num_rows),
4952        }
4953        .detach(node_manager, pending.clone(), false)
4954        .unwrap();
4955
4956        // The spawned task is still waiting on the gated flownode, so the
4957        // reservation made by `detach` is the value observable here.
4958        assert_eq!(num_rows as u64, pending.load(Ordering::Relaxed));
4959        assert_eq!(
4960            num_rows as i64,
4961            crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() - gauge_before,
4962            "the pending gauge must track the pending rows"
4963        );
4964        assert_eq!(
4965            dropped_before,
4966            crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4967            "a batch within the limit must not be counted as dropped"
4968        );
4969
4970        // Release the in-flight task and wait for it to hand its reservation
4971        // back, so no task is left holding the shared metrics after this test.
4972        gate.add_permits(1);
4973        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4974        while pending.load(Ordering::Relaxed) != 0
4975            || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4976        {
4977            assert!(
4978                std::time::Instant::now() < deadline,
4979                "the mirror task did not release its pending reservation"
4980            );
4981            std::thread::sleep(std::time::Duration::from_millis(1));
4982        }
4983    }
4984
4985    #[test]
4986    fn mirror_logs_are_rate_limited() {
4987        let log = Arc::new(MirrorLog::new());
4988        assert_eq!(Some(1), log.claim_report(1_000));
4989        assert_eq!(None, log.claim_report(1_500));
4990        assert_eq!(None, log.claim_report(10_999));
4991        assert_eq!(Some(3), log.claim_report(11_000));
4992        assert_eq!(None, log.claim_report(11_001));
4993        assert_eq!(Some(2), log.claim_report(21_000));
4994
4995        let concurrent = Arc::new(MirrorLog::new());
4996        assert_eq!(Some(1), concurrent.claim_report(50_000));
4997        let barrier = Arc::new(std::sync::Barrier::new(8));
4998        let threads = (0..8)
4999            .map(|_| {
5000                let log = concurrent.clone();
5001                let barrier = barrier.clone();
5002                std::thread::spawn(move || {
5003                    barrier.wait();
5004                    log.claim_report(60_000)
5005                })
5006            })
5007            .collect::<Vec<_>>();
5008        let winners = threads
5009            .into_iter()
5010            .filter_map(|thread| thread.join().unwrap())
5011            .collect::<Vec<_>>();
5012        assert_eq!(winners.len(), 1);
5013        let residual = concurrent.events.swap(0, Ordering::Relaxed);
5014        assert_eq!(winners[0] + residual, 8);
5015    }
5016}