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