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