Skip to main content

operator/
insert.rs

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