Skip to main content

operator/
insert.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::BTreeMap;
16use std::sync::Arc;
17
18use ahash::{HashMap, HashMapExt, HashSet, HashSetExt};
19use api::v1::alter_table_expr::Kind;
20use api::v1::column_def::options_from_skipping;
21use api::v1::region::{
22    InsertRequest as RegionInsertRequest, InsertRequests as RegionInsertRequests,
23    RegionRequestHeader,
24};
25use api::v1::value::ValueData;
26use api::v1::{
27    AlterTableExpr, ColumnDataType, ColumnSchema, CreateTableExpr, InsertRequests,
28    RowInsertRequest, RowInsertRequests, Rows, SemanticType,
29};
30use catalog::CatalogManagerRef;
31use client::{OutputData, OutputMeta};
32use common_catalog::consts::{
33    DEFAULT_PRIVATE_SCHEMA_NAME, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN,
34    TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table,
35    trace_operations_table_name, trace_services_table_name,
36};
37use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME;
38use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME;
39use common_grpc_expr::util::ColumnExpr;
40use common_meta::cache::TableFlownodeSetCacheRef;
41use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME;
42use common_meta::node_manager::{AffectedRows, NodeManagerRef};
43use common_meta::peer::Peer;
44use common_meta::rpc::ddl::TriggerReason;
45use common_query::Output;
46use common_query::native_histogram::{is_native_histogram_value_type, native_histogram_value_type};
47use common_query::prelude::{greptime_timestamp, greptime_value};
48use common_telemetry::tracing_context::TracingContext;
49use common_telemetry::{debug, error, warn};
50use common_time::Timestamp;
51use common_time::timestamp::TimeUnit;
52use datatypes::schema::SkippingIndexOptions;
53use futures_util::future;
54use meter_core::data::MeterRecord;
55use meter_macros::write_meter;
56use partition::manager::PartitionRuleManagerRef;
57use session::context::QueryContextRef;
58use snafu::ResultExt;
59use snafu::prelude::*;
60use sql::partition::partition_rule_for_hexstring;
61use sql::statements::create::Partitions;
62use sql::statements::insert::Insert;
63use store_api::metric_engine_consts::{
64    LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
65};
66use store_api::mito_engine_options::{
67    APPEND_MODE_KEY, COMPACTION_TYPE, COMPACTION_TYPE_TWCS, MERGE_MODE_KEY, TTL_KEY,
68    TWCS_TIME_WINDOW,
69};
70use store_api::storage::{RegionId, TableId};
71use table::TableRef;
72use table::metadata::{TableInfo, TableInfoRef};
73use table::requests::{
74    AUTO_CREATE_TABLE_KEY, InsertRequest as TableInsertRequest, SEMANTIC_PER_TABLE_INDEX_KEY,
75    SEMANTIC_PIPELINE, TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1, TABLE_DATA_MODEL_TRACE_V2,
76    TRACE_TABLE_PARTITIONS_HINT_KEY, VALID_TABLE_OPTION_KEYS, is_semantic_option_key,
77    validate_semantic_option,
78};
79use table::table_reference::TableReference;
80
81use crate::batcher::PendingRowsBatcher;
82use crate::error::{
83    CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, FindRegionLeaderSnafu,
84    InvalidInsertRequestSnafu, JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu,
85    WriteRejectedSnafu,
86};
87use crate::expr_helper;
88use crate::region_req_factory::RegionRequestFactory;
89use crate::req_convert::common::preprocess_row_insert_requests;
90use crate::req_convert::insert::{
91    ColumnToRow, ImpureDefaultFiller, RowToRegion, StatementToRegion, TableToRegion,
92    fill_reqs_with_impure_default, rows_to_record_batch,
93};
94use crate::statement::StatementExecutor;
95
96pub struct Inserter {
97    catalog_manager: CatalogManagerRef,
98    pub(crate) partition_manager: PartitionRuleManagerRef,
99    pub(crate) node_manager: NodeManagerRef,
100    pub(crate) table_flownode_set_cache: TableFlownodeSetCacheRef,
101    /// Server-side upper bound for auto table creation on write.
102    /// When `false`, missing tables are never auto-created regardless of the
103    /// per-request `auto_create_table` hint. When `true`, the hint still applies.
104    auto_create_table: bool,
105    pending_rows_batcher: Option<Arc<dyn PendingRowsBatcher>>,
106}
107
108pub type InserterRef = Arc<Inserter>;
109
110/// Hint for the table type to create automatically.
111#[derive(Clone)]
112pub enum AutoCreateTableType {
113    /// A logical table with the physical table name.
114    Logical(String),
115    /// A physical table.
116    Physical,
117    /// A log table which is append-only.
118    Log,
119    /// A table that merges rows by `last_non_null` strategy.
120    LastNonNull,
121    /// Create table that build index and default partition rules on trace_id
122    Trace { alter_existing: bool },
123}
124
125impl AutoCreateTableType {
126    pub fn as_str(&self) -> &'static str {
127        match self {
128            AutoCreateTableType::Logical(_) => "logical",
129            AutoCreateTableType::Physical => "physical",
130            AutoCreateTableType::Log => "log",
131            AutoCreateTableType::LastNonNull => "last_non_null",
132            AutoCreateTableType::Trace { .. } => "trace",
133        }
134    }
135
136    fn alter_existing(&self) -> bool {
137        !matches!(
138            self,
139            Self::Trace {
140                alter_existing: false
141            }
142        )
143    }
144}
145
146/// Split insert requests into normal and instant requests.
147///
148/// Where instant requests are requests with ttl=instant,
149/// and normal requests are requests with ttl set to other values.
150///
151/// This is used to split requests for different processing.
152#[derive(Clone)]
153pub struct InstantAndNormalInsertRequests {
154    /// Requests with normal ttl.
155    pub normal_requests: RegionInsertRequests,
156    /// Requests with ttl=instant.
157    /// Will be discarded immediately at frontend, wouldn't even insert into memtable, and only sent to flow node if needed.
158    pub instant_requests: RegionInsertRequests,
159}
160
161impl Inserter {
162    /// Checks the assumptions of the logical bulk path without changing tables.
163    /// Unsupported requests retain ordinary insertion, including schema policy,
164    /// defaults, instant TTL and row-based Flow delivery.
165    pub async fn can_batch_metric_rows(
166        &self,
167        requests: &RowInsertRequests,
168        ctx: &QueryContextRef,
169        physical_table: &str,
170    ) -> Result<bool> {
171        if self.auto_create_disabled_reason(ctx)?.is_some() || ctx.extension(TTL_KEY).is_some() {
172            return Ok(false);
173        }
174        for request in &requests.inserts {
175            // The logical bulk encoder only supports scalar metric schemas;
176            // any time index unit is accepted — requests are converted to the
177            // destination table's unit during batch alignment.
178            // Check new tables too, before catalog lookup or schema changes.
179            if request.rows.as_ref().is_some_and(|rows| {
180                rows.schema.iter().any(|column| {
181                    column.datatype_extension.is_some()
182                        || !matches!(
183                            ColumnDataType::try_from(column.datatype),
184                            Ok(ColumnDataType::TimestampSecond
185                                | ColumnDataType::TimestampMillisecond
186                                | ColumnDataType::TimestampMicrosecond
187                                | ColumnDataType::TimestampNanosecond
188                                | ColumnDataType::Float64
189                                | ColumnDataType::String)
190                        )
191                })
192            }) {
193                return Ok(false);
194            }
195
196            let Some(table) = self
197                .get_table(
198                    ctx.current_catalog(),
199                    &ctx.current_schema(),
200                    &request.table_name,
201                )
202                .await?
203            else {
204                continue;
205            };
206            let info = table.table_info();
207            if info.meta.engine != METRIC_ENGINE_NAME
208                || info.is_ttl_instant_table()
209                || info
210                    .meta
211                    .options
212                    .extra_options
213                    .get(LOGICAL_TABLE_METADATA_KEY)
214                    .map(String::as_str)
215                    != Some(physical_table)
216                || info
217                    .meta
218                    .schema
219                    .column_schemas()
220                    .iter()
221                    .any(|column| column.default_constraint().is_some())
222            {
223                return Ok(false);
224            }
225            // Physical metric tags are nullable even when their logical schema
226            // is not. Arrow alignment is stricter than ordinary metric insertion.
227            if info
228                .meta
229                .primary_key_indices
230                .iter()
231                .any(|&index| !info.meta.schema.column_schemas()[index].is_nullable())
232            {
233                return Ok(false);
234            }
235            // The current Flow cache does not distinguish streaming and batch
236            // flows. Keep all Flow sources on the row-based delivery path.
237            match self.table_flownode_set_cache.get(info.table_id()).await {
238                Ok(None) => {}
239                Ok(Some(flows)) if flows.is_empty() => {}
240                _ => return Ok(false),
241            }
242        }
243        Ok(true)
244    }
245
246    /// Meters an original logical-table request before bulk routing, without
247    /// cloning its rows or changing the request boundary used for accounting.
248    pub async fn meter_row_inserts(
249        requests: &mut RowInsertRequests,
250        ctx: &QueryContextRef,
251    ) -> Result<u64> {
252        let metered = InstantAndNormalInsertRequests {
253            normal_requests: RegionInsertRequests {
254                requests: requests
255                    .inserts
256                    .iter_mut()
257                    .map(|request| RegionInsertRequest {
258                        rows: request.rows.take(),
259                        ..Default::default()
260                    })
261                    .collect(),
262            },
263            instant_requests: RegionInsertRequests::default(),
264        };
265        let cost = write_meter!(
266            ctx.current_catalog(),
267            ctx.current_schema(),
268            metered,
269            ctx.write_rows_to_admit(
270                ctx.current_catalog(),
271                &ctx.current_schema(),
272                count_insert_rows(&metered)?
273            ),
274            ctx.channel() as u8
275        )
276        .await
277        .context(WriteRejectedSnafu);
278        for (request, region) in requests
279            .inserts
280            .iter_mut()
281            .zip(metered.normal_requests.requests)
282        {
283            request.rows = region.rows;
284        }
285        cost
286    }
287
288    pub fn new(
289        catalog_manager: CatalogManagerRef,
290        partition_manager: PartitionRuleManagerRef,
291        node_manager: NodeManagerRef,
292        table_flownode_set_cache: TableFlownodeSetCacheRef,
293        auto_create_table: bool,
294    ) -> Self {
295        Self {
296            catalog_manager,
297            partition_manager,
298            node_manager,
299            table_flownode_set_cache,
300            auto_create_table,
301            pending_rows_batcher: None,
302        }
303    }
304
305    /// Installs the shared batcher; callers explicitly select its ingestion entry point.
306    pub fn with_pending_rows_batcher(
307        mut self,
308        batcher: Option<Arc<dyn PendingRowsBatcher>>,
309    ) -> Self {
310        self.pending_rows_batcher = batcher;
311        self
312    }
313
314    pub async fn handle_column_inserts(
315        &self,
316        requests: InsertRequests,
317        ctx: QueryContextRef,
318        statement_executor: &StatementExecutor,
319    ) -> Result<Output> {
320        let row_inserts = ColumnToRow::convert(requests)?;
321        self.handle_row_inserts(row_inserts, ctx, statement_executor, false, false)
322            .await
323    }
324
325    /// Handles row inserts request and creates a physical table on demand.
326    pub async fn handle_row_inserts(
327        &self,
328        mut requests: RowInsertRequests,
329        ctx: QueryContextRef,
330        statement_executor: &StatementExecutor,
331        accommodate_existing_schema: bool,
332        is_single_value: bool,
333    ) -> Result<Output> {
334        preprocess_row_insert_requests(&mut requests.inserts)?;
335        self.handle_row_inserts_with_create_type(
336            requests,
337            ctx,
338            statement_executor,
339            AutoCreateTableType::Physical,
340            accommodate_existing_schema,
341            is_single_value,
342        )
343        .await
344    }
345
346    /// Handles row inserts request and creates a log table on demand.
347    pub async fn handle_log_inserts(
348        &self,
349        requests: RowInsertRequests,
350        ctx: QueryContextRef,
351        statement_executor: &StatementExecutor,
352    ) -> Result<Output> {
353        self.handle_row_inserts_with_create_type(
354            requests,
355            ctx,
356            statement_executor,
357            AutoCreateTableType::Log,
358            false,
359            false,
360        )
361        .await
362    }
363
364    pub async fn handle_trace_inserts(
365        &self,
366        requests: RowInsertRequests,
367        ctx: QueryContextRef,
368        statement_executor: &StatementExecutor,
369    ) -> Result<Output> {
370        self.handle_row_inserts_with_create_type(
371            requests,
372            ctx,
373            statement_executor,
374            AutoCreateTableType::Trace {
375                alter_existing: true,
376            },
377            false,
378            false,
379        )
380        .await
381    }
382
383    /// Handles row inserts request and creates a table with `last_non_null` merge mode on demand.
384    pub async fn handle_last_non_null_inserts(
385        &self,
386        requests: RowInsertRequests,
387        ctx: QueryContextRef,
388        statement_executor: &StatementExecutor,
389        accommodate_existing_schema: bool,
390        is_single_value: bool,
391    ) -> Result<Output> {
392        self.handle_row_inserts_with_create_type(
393            requests,
394            ctx,
395            statement_executor,
396            AutoCreateTableType::LastNonNull,
397            accommodate_existing_schema,
398            is_single_value,
399        )
400        .await
401    }
402
403    /// Handles row inserts request with specified [AutoCreateTableType].
404    async fn handle_row_inserts_with_create_type(
405        &self,
406        mut requests: RowInsertRequests,
407        ctx: QueryContextRef,
408        statement_executor: &StatementExecutor,
409        create_type: AutoCreateTableType,
410        accommodate_existing_schema: bool,
411        is_single_value: bool,
412    ) -> Result<Output> {
413        let skip_wal = ctx.skip_wal();
414
415        let batcher = self
416            .pending_rows_batcher
417            .as_ref()
418            .filter(|_| ctx.batching_enabled());
419
420        // remove empty requests
421        requests.inserts.retain(|req| {
422            req.rows
423                .as_ref()
424                .map(|r| !r.rows.is_empty())
425                .unwrap_or_default()
426        });
427        validate_column_count_match(&requests)?;
428
429        let CreateAlterTableResult {
430            instant_table_ids,
431            table_infos,
432        } = self
433            .create_or_alter_tables_on_demand(
434                &mut requests,
435                &ctx,
436                create_type,
437                statement_executor,
438                accommodate_existing_schema,
439                is_single_value,
440                None,
441            )
442            .await?;
443
444        // Instant tables have no persisted data for dirty-window Flow to read.
445        // Metric tables keep their existing dedicated ingestion path.
446        if let Some(batcher) = batcher
447            && instant_table_ids.is_empty()
448            && table_infos
449                .values()
450                .all(|info| info.meta.engine == default_engine())
451        {
452            return self
453                .submit_pending_rows(requests, table_infos, ctx, batcher)
454                .await;
455        }
456
457        let name_to_info = table_infos
458            .values()
459            .map(|info| (info.name.clone(), info.clone()))
460            .collect::<HashMap<_, _>>();
461        let inserts = RowToRegion::new(
462            name_to_info,
463            instant_table_ids,
464            self.partition_manager.as_ref(),
465        )
466        .convert(requests, skip_wal)
467        .await?;
468
469        self.do_request(inserts, &table_infos, &ctx).await
470    }
471
472    async fn submit_pending_rows(
473        &self,
474        mut requests: RowInsertRequests,
475        table_infos: HashMap<TableId, Arc<TableInfo>>,
476        ctx: QueryContextRef,
477        batcher: &Arc<dyn PendingRowsBatcher>,
478    ) -> Result<Output> {
479        // All entry points, including single-table and SQL writes, skip empty input
480        // before evaluating defaults or converting prepared rows.
481        requests.inserts.retain(|request| {
482            request
483                .rows
484                .as_ref()
485                .is_some_and(|rows| !rows.rows.is_empty())
486        });
487        let by_name = table_infos
488            .values()
489            .map(|info| (info.name.as_str(), info))
490            .collect::<HashMap<_, _>>();
491        let mut prepared = Vec::with_capacity(requests.inserts.len());
492        for request in &mut requests.inserts {
493            let table_info =
494                by_name
495                    .get(request.table_name.as_str())
496                    .context(TableNotFoundSnafu {
497                        table_name: &request.table_name,
498                    })?;
499            let Some(rows) = &mut request.rows else {
500                continue;
501            };
502            ImpureDefaultFiller::new((*table_info).clone())?.fill_rows(rows);
503            let batch = rows_to_record_batch(rows, table_info)?;
504            prepared.push(((*table_info).clone(), batch));
505        }
506        // Preserve the existing meter input and original request boundary. These
507        // envelopes are only for accounting; routing happens after batching.
508        let metered = InstantAndNormalInsertRequests {
509            normal_requests: RegionInsertRequests {
510                requests: requests
511                    .inserts
512                    .into_iter()
513                    .map(|request| RegionInsertRequest {
514                        rows: request.rows,
515                        ..Default::default()
516                    })
517                    .collect(),
518            },
519            instant_requests: RegionInsertRequests::default(),
520        };
521        let table_info = table_infos.values().next();
522        let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str());
523        let schema =
524            table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone());
525        let write_cost = write_meter!(
526            catalog,
527            &schema,
528            metered,
529            ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&metered)?),
530            ctx.channel() as u8
531        )
532        .await
533        .context(WriteRejectedSnafu)?;
534        prepared.retain(|(_, batch)| batch.num_rows() != 0);
535        let results = if prepared.is_empty() {
536            Vec::new()
537        } else {
538            // One original request shares admission across all table submissions.
539            let permit = batcher.acquire().await?;
540            let submissions = prepared.into_iter().map(|(info, batch)| {
541                // Route to the same target database used for admission above.
542                let mut target_ctx = ctx.fork();
543                target_ctx.set_current_catalog(&info.catalog_name);
544                target_ctx.set_current_schema(&info.schema_name);
545                batcher.submit(info, batch, Arc::new(target_ctx), permit.clone())
546            });
547            // Observe every table completion even when another table fails.
548            future::join_all(submissions).await
549        };
550        let affected_rows = results.into_iter().sum::<Result<usize>>()?;
551        Ok(Output::new(
552            OutputData::AffectedRows(affected_rows),
553            OutputMeta::new_with_cost(write_cost as _),
554        ))
555    }
556
557    /// Handles row inserts request with metric engine.
558    pub async fn handle_metric_row_inserts(
559        &self,
560        mut requests: RowInsertRequests,
561        ctx: QueryContextRef,
562        statement_executor: &StatementExecutor,
563        physical_table: String,
564    ) -> Result<Output> {
565        let skip_wal = ctx.skip_wal();
566
567        // remove empty requests
568        requests.inserts.retain(|req| {
569            req.rows
570                .as_ref()
571                .map(|r| !r.rows.is_empty())
572                .unwrap_or_default()
573        });
574        validate_column_count_match(&requests)?;
575
576        // check and create physical table
577        let physical_table_ref = self
578            .create_physical_table_on_demand(&ctx, physical_table.clone(), statement_executor)
579            .await?;
580
581        // check and create logical tables; `create_or_alter_tables_on_demand`
582        // aligns each request's time index unit with the unit of the table it
583        // targets, inside its existing table lookups: existing tables keep
584        // their own unit (which matches the physical table they are bound
585        // to), and new tables use the selected physical table's unit (from
586        // `physical_table_ref`). Ingestion endpoints encode timestamps in a
587        // fixed unit (prometheus remote write always uses millisecond; OTLP
588        // keeps nanosecond precision on the metric engine path), and the
589        // metric engine requires each logical table's requests to match its
590        // time index unit. Narrowing conversions truncate the sub-unit part
591        // (floor), following `Timestamp::convert_to`.
592        let CreateAlterTableResult {
593            instant_table_ids,
594            table_infos,
595        } = self
596            .create_or_alter_tables_on_demand(
597                &mut requests,
598                &ctx,
599                AutoCreateTableType::Logical(physical_table.clone()),
600                statement_executor,
601                true,
602                true,
603                table_time_index_unit(&physical_table_ref),
604            )
605            .await?;
606        let name_to_info = table_infos
607            .values()
608            .map(|info| (info.name.clone(), info.clone()))
609            .collect::<HashMap<_, _>>();
610        let inserts = RowToRegion::new(name_to_info, instant_table_ids, &self.partition_manager)
611            .convert(requests, skip_wal)
612            .await?;
613
614        self.do_request(inserts, &table_infos, &ctx).await
615    }
616
617    fn table_batcher(
618        &self,
619        table_info: &TableInfoRef,
620        ctx: &QueryContextRef,
621    ) -> Option<&Arc<dyn PendingRowsBatcher>> {
622        self.pending_rows_batcher.as_ref().filter(|_| {
623            ctx.batching_enabled()
624                && !table_info.is_ttl_instant_table()
625                && table_info.meta.engine == default_engine()
626        })
627    }
628
629    async fn submit_table_rows(
630        &self,
631        rows: Rows,
632        table_info: TableInfoRef,
633        ctx: QueryContextRef,
634        batcher: &Arc<dyn PendingRowsBatcher>,
635    ) -> Result<Output> {
636        let requests = RowInsertRequests {
637            inserts: vec![RowInsertRequest {
638                table_name: table_info.name.clone(),
639                rows: Some(rows),
640            }],
641        };
642        let table_infos = HashMap::from_iter([(table_info.table_id(), table_info)]);
643        self.submit_pending_rows(requests, table_infos, ctx, batcher)
644            .await
645    }
646
647    pub async fn handle_table_insert(
648        &self,
649        request: TableInsertRequest,
650        ctx: QueryContextRef,
651    ) -> Result<Output> {
652        let catalog = request.catalog_name.as_str();
653        let schema = request.schema_name.as_str();
654        let table_name = request.table_name.as_str();
655        let table = self.get_table(catalog, schema, table_name).await?;
656        let table = table.with_context(|| TableNotFoundSnafu {
657            table_name: common_catalog::format_full_table_name(catalog, schema, table_name),
658        })?;
659        let table_info = table.table_info();
660
661        let converter = TableToRegion::new(&table_info, &self.partition_manager);
662        let skip_wal = request.skip_wal;
663        let rows = converter.prepare(request)?;
664        if let Some(batcher) = self.table_batcher(&table_info, &ctx) {
665            return self.submit_table_rows(rows, table_info, ctx, batcher).await;
666        }
667        let inserts = converter.partition(rows, skip_wal).await?;
668
669        let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]);
670
671        self.do_request(inserts, &table_infos, &ctx).await
672    }
673
674    pub async fn handle_statement_insert(
675        &self,
676        insert: &Insert,
677        ctx: &QueryContextRef,
678    ) -> Result<Output> {
679        let converter =
680            StatementToRegion::new(self.catalog_manager.as_ref(), &self.partition_manager, ctx);
681        let (rows, table_info) = converter.prepare(insert, ctx).await?;
682        if let Some(batcher) = self.table_batcher(&table_info, ctx) {
683            return self
684                .submit_table_rows(rows, table_info, ctx.clone(), batcher)
685                .await;
686        }
687        let inserts = converter.partition(rows, table_info.clone(), ctx).await?;
688
689        let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]);
690
691        self.do_request(inserts, &table_infos, ctx).await
692    }
693}
694
695/// Admits a finite request before it is split into internal writes.
696/// The returned context preserves accounting while preventing a second row debit.
697pub async fn admit_write(rows: u64, ctx: &QueryContextRef) -> Result<QueryContextRef> {
698    // The zero value is WCU: this record only admits rows. Actual inserts retain
699    // their existing WCU accounting, so charging here would count it twice.
700    write_meter!(MeterRecord::new(
701        ctx.current_catalog().to_string(),
702        ctx.current_schema(),
703        0,
704        ctx.write_rows_to_admit(ctx.current_catalog(), &ctx.current_schema(), rows),
705        ctx.channel() as u8,
706    ))
707    .await
708    .context(WriteRejectedSnafu)?;
709    Ok(Arc::new(ctx.with_write_admission()))
710}
711
712/// Admits all database totals before dispatching any batch of a finite request.
713/// Each batch keeps its own protocol options and target database.
714pub async fn admit_row_insert_batches(
715    batches: &mut [(QueryContextRef, RowInsertRequests)],
716) -> Result<()> {
717    let mut totals = BTreeMap::<_, (QueryContextRef, u64)>::new();
718    for (ctx, requests) in batches.iter() {
719        let catalog = ctx.current_catalog();
720        let schema = ctx.current_schema();
721        if ctx.write_rows_to_admit(catalog, &schema, 1) == 0 {
722            continue;
723        }
724        let (_, total) = totals
725            .entry((catalog.to_string(), schema.clone()))
726            .or_insert_with(|| (ctx.clone(), 0));
727        for rows in requests.inserts.iter().filter_map(|r| r.rows.as_ref()) {
728            *total =
729                total
730                    .checked_add(rows.rows.len() as u64)
731                    .context(InvalidInsertRequestSnafu {
732                        reason: "Insert row count exceeds u64::MAX",
733                    })?;
734        }
735    }
736    for (ctx, rows) in totals.values() {
737        admit_write(*rows, ctx).await?;
738    }
739    for (ctx, _) in batches {
740        *ctx = Arc::new(ctx.with_write_admission());
741    }
742    Ok(())
743}
744
745fn count_insert_rows(requests: &InstantAndNormalInsertRequests) -> Result<u64> {
746    requests
747        .normal_requests
748        .requests
749        .iter()
750        .chain(&requests.instant_requests.requests)
751        .filter_map(|request| request.rows.as_ref())
752        .try_fold(0u64, |total, rows| {
753            total
754                .checked_add(rows.rows.len() as u64)
755                .context(InvalidInsertRequestSnafu {
756                    reason: "Insert row count exceeds u64::MAX",
757                })
758        })
759}
760
761impl Inserter {
762    async fn do_request(
763        &self,
764        requests: InstantAndNormalInsertRequests,
765        table_infos: &HashMap<TableId, Arc<TableInfo>>,
766        ctx: &QueryContextRef,
767    ) -> Result<Output> {
768        // Fill impure default values in the request
769        let requests = fill_reqs_with_impure_default(table_infos, requests)?;
770
771        // All tables in a batch resolve to the same database. Qualified SQL
772        // inserts may target a different database than the session's current one.
773        let table_info = table_infos.values().next();
774        let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str());
775        let schema =
776            table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone());
777        let write_cost = write_meter!(
778            catalog,
779            schema.clone(),
780            requests,
781            ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&requests)?),
782            ctx.channel() as u8
783        )
784        .await
785        .context(WriteRejectedSnafu)?;
786        let request_factory = RegionRequestFactory::new(RegionRequestHeader {
787            tracing_context: TracingContext::from_current_span().to_w3c(),
788            dbname: ctx.get_db_string(),
789            ..Default::default()
790        });
791
792        let InstantAndNormalInsertRequests {
793            normal_requests,
794            instant_requests,
795        } = requests;
796
797        // Mirror requests for source table to flownode asynchronously
798        let flow_mirror_task = FlowMirrorTask::new(
799            &self.table_flownode_set_cache,
800            normal_requests
801                .requests
802                .iter()
803                .chain(instant_requests.requests.iter()),
804        )
805        .await?;
806        flow_mirror_task.detach(self.node_manager.clone())?;
807
808        // Write requests to datanode and wait for response
809        let write_tasks = self
810            .group_requests_by_peer(normal_requests)
811            .await?
812            .into_iter()
813            .map(|(peer, inserts)| {
814                let node_manager = self.node_manager.clone();
815                let request = request_factory.build_insert(inserts);
816                common_runtime::spawn_global(async move {
817                    node_manager
818                        .datanode(&peer)
819                        .await
820                        .handle(request)
821                        .await
822                        .context(RequestInsertsSnafu)
823                })
824            });
825        let results = future::try_join_all(write_tasks)
826            .await
827            .context(JoinTaskSnafu)?;
828        let affected_rows = results
829            .into_iter()
830            .map(|resp| resp.map(|r| r.affected_rows))
831            .sum::<Result<AffectedRows>>()?;
832        crate::metrics::DIST_INGEST_ROW_COUNT
833            .with_label_values(&[ctx.get_db_string().as_str()])
834            .inc_by(affected_rows as u64);
835        Ok(Output::new(
836            OutputData::AffectedRows(affected_rows),
837            OutputMeta::new_with_cost(write_cost as _),
838        ))
839    }
840
841    async fn group_requests_by_peer(
842        &self,
843        requests: RegionInsertRequests,
844    ) -> Result<HashMap<Peer, RegionInsertRequests>> {
845        // group by region ids first to reduce repeatedly call `find_region_leader`
846        // TODO(discord9): determine if a addition clone is worth it
847        let mut requests_per_region: HashMap<RegionId, RegionInsertRequests> = HashMap::new();
848        for req in requests.requests {
849            let region_id = RegionId::from_u64(req.region_id);
850            requests_per_region
851                .entry(region_id)
852                .or_default()
853                .requests
854                .push(req);
855        }
856
857        let mut inserts: HashMap<Peer, RegionInsertRequests> = HashMap::new();
858
859        for (region_id, reqs) in requests_per_region {
860            let peer = self
861                .partition_manager
862                .find_region_leader(region_id)
863                .await
864                .context(FindRegionLeaderSnafu)?;
865            inserts
866                .entry(peer)
867                .or_default()
868                .requests
869                .extend(reqs.requests);
870        }
871
872        Ok(inserts)
873    }
874
875    /// Returns `Some(reason)` if the config or request hint disables automatic
876    /// table creation. Exempt private system tables are handled by
877    /// [`Self::is_auto_create_exempt_private_table`].
878    fn auto_create_disabled_reason(&self, ctx: &QueryContextRef) -> Result<Option<&'static str>> {
879        let auto_create_table_hint = ctx
880            .extension(AUTO_CREATE_TABLE_KEY)
881            .map(|v| v.parse::<bool>())
882            .transpose()
883            .map_err(|_| {
884                InvalidInsertRequestSnafu {
885                    reason: "`auto_create_table` hint must be a boolean",
886                }
887                .build()
888            })?
889            .unwrap_or(true);
890        Ok(if !self.auto_create_table {
891            Some("auto-create table is disabled by frontend config")
892        } else if !auto_create_table_hint {
893            Some("`auto_create_table` hint is disabled")
894        } else {
895            None
896        })
897    }
898
899    /// Returns whether a private system table may infer and reconcile its schema
900    /// even when automatic table creation is disabled.
901    fn is_auto_create_exempt_private_table(schema: &str, table: &str) -> bool {
902        schema == DEFAULT_PRIVATE_SCHEMA_NAME
903            && matches!(
904                table,
905                DEFAULT_EVENTS_TABLE_NAME | SLOW_QUERY_TABLE_NAME | REGION_STATS_HISTORY_TABLE_NAME
906            )
907    }
908
909    /// Ensures a trace table has the request-global schema without requiring a
910    /// padded data row to drive on-demand creation or alteration. When
911    /// `alter_existing` is false, a table created after planning is left for the
912    /// caller to re-plan.
913    pub async fn ensure_trace_table_on_demand(
914        &self,
915        table_name: &str,
916        request_schema: Vec<ColumnSchema>,
917        alter_existing: bool,
918        ctx: &QueryContextRef,
919        statement_executor: &StatementExecutor,
920    ) -> Result<()> {
921        let mut requests = RowInsertRequests {
922            inserts: vec![RowInsertRequest {
923                table_name: table_name.to_string(),
924                rows: Some(api::v1::Rows {
925                    schema: request_schema,
926                    rows: Vec::new(),
927                }),
928            }],
929        };
930        self.create_or_alter_tables_on_demand(
931            &mut requests,
932            ctx,
933            AutoCreateTableType::Trace { alter_existing },
934            statement_executor,
935            false,
936            false,
937            None,
938        )
939        .await?;
940        Ok(())
941    }
942
943    /// Creates or alter tables on demand:
944    /// - if table does not exist, create table by inferred CreateExpr
945    /// - if table exist, check if schema matches. If any new column found, alter table by inferred `AlterExpr`
946    ///
947    /// Returns a mapping from table name to table id, where table name is the table name involved in the requests.
948    /// This mapping is used in the conversion of RowToRegion.
949    ///
950    /// `accommodate_existing_schema` is used to determine if the existing schema should override the new schema.
951    /// It only works for TIME_INDEX and single VALUE columns. This is for the case where the user creates a table with
952    /// custom schema, and then inserts data with endpoints that have default schema setting, like prometheus
953    /// remote write. This will modify the `RowInsertRequests` in place.
954    /// `is_single_value` indicates whether the default schema only contains single value column so we can accommodate it.
955    ///
956    /// `align_time_index_unit` is the selected physical metric table's time
957    /// index unit; passing `Some` (metric engine path only) rewrites each
958    /// request's time index column, inside this function's existing table
959    /// lookups (no extra catalog access): existing destination tables are
960    /// converted to their own unit, and new tables to the given unit.
961    #[allow(clippy::too_many_arguments)]
962    async fn create_or_alter_tables_on_demand(
963        &self,
964        requests: &mut RowInsertRequests,
965        ctx: &QueryContextRef,
966        auto_create_table_type: AutoCreateTableType,
967        statement_executor: &StatementExecutor,
968        accommodate_existing_schema: bool,
969        is_single_value: bool,
970        align_time_index_unit: Option<TimeUnit>,
971    ) -> Result<CreateAlterTableResult> {
972        let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND
973            .with_label_values(&[auto_create_table_type.as_str()])
974            .start_timer();
975        let catalog = ctx.current_catalog();
976        let schema = ctx.current_schema();
977
978        let auto_create_disabled_reason = self.auto_create_disabled_reason(ctx)?;
979        // Enabled batches permit every table, so only disabled batches need a whitelist scan.
980        let has_auto_create_exempt_table = auto_create_disabled_reason.is_some()
981            && requests
982                .inserts
983                .iter()
984                .any(|req| Self::is_auto_create_exempt_private_table(&schema, &req.table_name));
985        let mut table_infos = HashMap::new();
986        // Without exempt tables, verify existing tables and reject missing ones without inferring schemas.
987        if let Some(disabled_reason) = auto_create_disabled_reason
988            && !has_auto_create_exempt_table
989        {
990            let mut instant_table_ids = HashSet::new();
991            for req in &mut requests.inserts {
992                let table = match self.get_table(catalog, &schema, &req.table_name).await? {
993                    Some(table) => table,
994                    // System-defined table: created canonically by the system,
995                    // so the auto-create config/hint does not apply.
996                    None if is_ddl_reserved_table(&schema, &req.table_name) => {
997                        statement_executor
998                            .create_declared_relationships_table(catalog, ctx.clone())
999                            .await?
1000                    }
1001                    None => {
1002                        return InvalidInsertRequestSnafu {
1003                            reason: format!(
1004                                "Table `{}` does not exist, and {}",
1005                                req.table_name, disabled_reason
1006                            ),
1007                        }
1008                        .fail();
1009                    }
1010                };
1011                // Metric path: an existing destination table keeps its own
1012                // time index unit (it may be bound to another physical
1013                // table than the one selected by this request).
1014                if align_time_index_unit.is_some()
1015                    && let Some(rows) = req.rows.as_mut()
1016                    && let Some(target_unit) = table_time_index_unit(&table)
1017                {
1018                    convert_rows_time_unit(rows, target_unit)?;
1019                }
1020                let table_info = table.table_info();
1021                if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1022                    validate_trace_table_model(&table_info, ctx)?;
1023                }
1024                if table_info.is_ttl_instant_table() {
1025                    instant_table_ids.insert(table_info.table_id());
1026                }
1027                table_infos.insert(table_info.table_id(), table.table_info());
1028            }
1029            let ret = CreateAlterTableResult {
1030                instant_table_ids,
1031                table_infos,
1032            };
1033            return Ok(ret);
1034        }
1035
1036        let mut create_tables = vec![];
1037        let mut alter_tables = vec![];
1038        let mut need_refresh_table_infos = HashSet::new();
1039        let mut instant_table_ids = HashSet::new();
1040        let mut per_table_semantics: Option<Option<PerTableSemanticIndex>> = None;
1041
1042        for req in &mut requests.inserts {
1043            // Mixed batches need a per-table decision so an exempt table cannot authorize others.
1044            let auto_create_allowed = auto_create_disabled_reason.is_none()
1045                || Self::is_auto_create_exempt_private_table(&schema, &req.table_name);
1046            match self.get_table(catalog, &schema, &req.table_name).await? {
1047                Some(table) => {
1048                    let table_info = table.table_info();
1049                    if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1050                        validate_trace_table_model(&table_info, ctx)?;
1051                    }
1052                    if table_info.is_ttl_instant_table() {
1053                        instant_table_ids.insert(table_info.table_id());
1054                    }
1055                    // Metric path: an existing destination table keeps its
1056                    // own time index unit (it may be bound to another
1057                    // physical table than the one selected by this request).
1058                    if align_time_index_unit.is_some()
1059                        && let Some(rows) = req.rows.as_mut()
1060                        && let Some(target_unit) = table_time_index_unit(&table)
1061                    {
1062                        convert_rows_time_unit(rows, target_unit)?;
1063                    }
1064                    if auto_create_allowed
1065                        && let Some(alter_expr) = self.get_alter_table_expr_on_demand(
1066                            req,
1067                            &table,
1068                            ctx,
1069                            accommodate_existing_schema,
1070                            is_single_value,
1071                            auto_create_table_type.alter_existing(),
1072                        )?
1073                    {
1074                        alter_tables.push(alter_expr);
1075                        need_refresh_table_infos.insert((
1076                            catalog.to_string(),
1077                            schema.clone(),
1078                            req.table_name.clone(),
1079                        ));
1080                    } else {
1081                        table_infos.insert(table_info.table_id(), table.table_info());
1082                    }
1083                }
1084                // A DDL-reserved table's definition never derives from the
1085                // write request; the system creates it canonically, below the
1086                // user-DDL guard that rejects the generic create path.
1087                None if is_ddl_reserved_table(&schema, &req.table_name) => {
1088                    let table = statement_executor
1089                        .create_declared_relationships_table(catalog, ctx.clone())
1090                        .await?;
1091                    let table_info = table.table_info();
1092                    if table_info.is_ttl_instant_table() {
1093                        instant_table_ids.insert(table_info.table_id());
1094                    }
1095                    table_infos.insert(table_info.table_id(), table_info);
1096                }
1097                None if !auto_create_allowed
1098                    && let Some(disabled_reason) = auto_create_disabled_reason =>
1099                {
1100                    return InvalidInsertRequestSnafu {
1101                        reason: format!(
1102                            "Table `{}` does not exist, and {}",
1103                            req.table_name, disabled_reason,
1104                        ),
1105                    }
1106                    .fail();
1107                }
1108                None => {
1109                    // Metric path: a new table uses the selected physical
1110                    // table's unit; convert before the create expression is
1111                    // derived from the request schema.
1112                    if let Some(physical_unit) = align_time_index_unit
1113                        && let Some(rows) = req.rows.as_mut()
1114                    {
1115                        convert_rows_time_unit(rows, physical_unit)?;
1116                    }
1117                    let semantic_index = per_table_semantics
1118                        .get_or_insert_with(|| parse_per_table_semantic_index(ctx))
1119                        .as_ref();
1120                    let create_expr = self.get_create_table_expr_on_demand(
1121                        req,
1122                        &auto_create_table_type,
1123                        ctx,
1124                        semantic_index,
1125                    )?;
1126                    create_tables.push(create_expr);
1127                }
1128            }
1129        }
1130
1131        match auto_create_table_type {
1132            AutoCreateTableType::Logical(_) => {
1133                if !create_tables.is_empty() {
1134                    // Creates logical tables in batch.
1135                    let tables = self
1136                        .create_logical_tables(create_tables, ctx, statement_executor)
1137                        .await?;
1138
1139                    for table in tables {
1140                        let table_info = table.table_info();
1141                        if table_info.is_ttl_instant_table() {
1142                            instant_table_ids.insert(table_info.table_id());
1143                        }
1144                        table_infos.insert(table_info.table_id(), table.table_info());
1145                    }
1146                }
1147                if !alter_tables.is_empty() {
1148                    // Alter logical tables in batch.
1149                    statement_executor
1150                        .alter_logical_tables(alter_tables, ctx.clone(), TriggerReason::AutoAlter)
1151                        .await?;
1152                }
1153            }
1154            AutoCreateTableType::Physical
1155            | AutoCreateTableType::Log
1156            | AutoCreateTableType::LastNonNull => {
1157                // note that auto create table shouldn't be ttl instant table
1158                // for it's a very unexpected behavior and should be set by user explicitly
1159                for create_table in create_tables {
1160                    let table = self
1161                        .create_physical_table(create_table, None, ctx, statement_executor)
1162                        .await?;
1163                    let table_info = table.table_info();
1164                    if table_info.is_ttl_instant_table() {
1165                        instant_table_ids.insert(table_info.table_id());
1166                    }
1167                    table_infos.insert(table_info.table_id(), table.table_info());
1168                }
1169                for alter_expr in alter_tables.into_iter() {
1170                    statement_executor
1171                        .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1172                        .await?;
1173                }
1174            }
1175
1176            AutoCreateTableType::Trace { .. } => {
1177                let trace_table_name = ctx
1178                    .extension(TRACE_TABLE_NAME_SESSION_KEY)
1179                    .unwrap_or(TRACE_TABLE_NAME);
1180
1181                let trace_table_partitions = if let Some(trace_table_partitions) =
1182                    ctx.extension(TRACE_TABLE_PARTITIONS_HINT_KEY)
1183                {
1184                    let p = trace_table_partitions.parse::<u32>().map_err(|_| {
1185                        InvalidInsertRequestSnafu {
1186                            reason: format!(
1187                                "Failed to parse trace_table_partitions: {}",
1188                                trace_table_partitions
1189                            ),
1190                        }
1191                        .build()
1192                    })?;
1193                    Some(p)
1194                } else {
1195                    None
1196                };
1197
1198                // note that auto create table shouldn't be ttl instant table
1199                // for it's a very unexpected behavior and should be set by user explicitly
1200                for mut create_table in create_tables {
1201                    if create_table.table_name == trace_services_table_name(trace_table_name)
1202                        || create_table.table_name == trace_operations_table_name(trace_table_name)
1203                    {
1204                        // Disable append mode for auxiliary tables (services/operations) since they require upsert behavior.
1205                        create_table
1206                            .table_options
1207                            .insert(APPEND_MODE_KEY.to_string(), "false".to_string());
1208                        // Remove `ttl` key from table options if it exists
1209                        create_table.table_options.remove(TTL_KEY);
1210
1211                        let table = self
1212                            .create_physical_table(create_table, None, ctx, statement_executor)
1213                            .await?;
1214                        let table_info = table.table_info();
1215                        if table_info.is_ttl_instant_table() {
1216                            instant_table_ids.insert(table_info.table_id());
1217                        }
1218                        table_infos.insert(table_info.table_id(), table.table_info());
1219                    } else {
1220                        // prebuilt partition rules for uuid data: see the function
1221                        // for more information
1222                        let partitions = if matches!(trace_table_partitions, Some(0) | Some(1)) {
1223                            // disable partitions
1224                            None
1225                        } else {
1226                            let p = partition_rule_for_hexstring(
1227                                TRACE_ID_COLUMN,
1228                                trace_table_partitions,
1229                            )
1230                            .context(CreatePartitionRulesSnafu)?;
1231                            Some(p)
1232                        };
1233
1234                        // add skip index to
1235                        // - trace_id: when searching by trace id
1236                        // - parent_span_id: when searching root span
1237                        // - span_name: when searching certain types of span
1238                        let index_columns =
1239                            [TRACE_ID_COLUMN, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN];
1240                        for index_column in index_columns {
1241                            if let Some(col) = create_table
1242                                .column_defs
1243                                .iter_mut()
1244                                .find(|c| c.name == index_column)
1245                            {
1246                                col.options =
1247                                    options_from_skipping(&SkippingIndexOptions::default())
1248                                        .context(ColumnOptionsSnafu)?;
1249                            } else {
1250                                warn!(
1251                                    "Column {} not found when creating index for trace table: {}.",
1252                                    index_column, create_table.table_name
1253                                );
1254                            }
1255                        }
1256
1257                        // use table_options to mark table model version
1258                        create_table.table_options.insert(
1259                            TABLE_DATA_MODEL.to_string(),
1260                            ctx.extension(SEMANTIC_PIPELINE)
1261                                .unwrap_or(TABLE_DATA_MODEL_TRACE_V1)
1262                                .to_string(),
1263                        );
1264
1265                        let table = self
1266                            .create_physical_table(
1267                                create_table,
1268                                partitions,
1269                                ctx,
1270                                statement_executor,
1271                            )
1272                            .await?;
1273                        let table_info = table.table_info();
1274                        if table_info.is_ttl_instant_table() {
1275                            instant_table_ids.insert(table_info.table_id());
1276                        }
1277                        table_infos.insert(table_info.table_id(), table.table_info());
1278                    }
1279                }
1280                for alter_expr in alter_tables.into_iter() {
1281                    statement_executor
1282                        .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1283                        .await?;
1284                }
1285            }
1286        }
1287
1288        // refresh table infos for altered tables
1289        for (catalog, schema, table_name) in need_refresh_table_infos {
1290            let table = self
1291                .get_table(&catalog, &schema, &table_name)
1292                .await?
1293                .context(TableNotFoundSnafu {
1294                    table_name: common_catalog::format_full_table_name(
1295                        &catalog,
1296                        &schema,
1297                        &table_name,
1298                    ),
1299                })?;
1300            let table_info = table.table_info();
1301            table_infos.insert(table_info.table_id(), table.table_info());
1302        }
1303
1304        Ok(CreateAlterTableResult {
1305            instant_table_ids,
1306            table_infos,
1307        })
1308    }
1309
1310    async fn create_physical_table_on_demand(
1311        &self,
1312        ctx: &QueryContextRef,
1313        physical_table: String,
1314        statement_executor: &StatementExecutor,
1315    ) -> Result<TableRef> {
1316        let catalog_name = ctx.current_catalog();
1317        let schema_name = ctx.current_schema();
1318
1319        // check if exist
1320        if let Some(table) = self
1321            .get_table(catalog_name, &schema_name, &physical_table)
1322            .await?
1323        {
1324            return Ok(table);
1325        }
1326
1327        // Gate here too, otherwise a disabled switch would still leak the physical table.
1328        if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? {
1329            return InvalidInsertRequestSnafu {
1330                reason: format!(
1331                    "Physical table `{physical_table}` does not exist, and {disabled_reason}"
1332                ),
1333            }
1334            .fail();
1335        }
1336
1337        let table_reference = TableReference::full(catalog_name, &schema_name, &physical_table);
1338        debug!("Ensuring physical metric table `{table_reference}` exists for insert");
1339
1340        // schema with timestamp and field column
1341        let default_schema = vec![
1342            ColumnSchema {
1343                column_name: greptime_timestamp().to_string(),
1344                datatype: ColumnDataType::TimestampMillisecond as _,
1345                semantic_type: SemanticType::Timestamp as _,
1346                datatype_extension: None,
1347                options: None,
1348            },
1349            ColumnSchema {
1350                column_name: greptime_value().to_string(),
1351                datatype: ColumnDataType::Float64 as _,
1352                semantic_type: SemanticType::Field as _,
1353                datatype_extension: None,
1354                options: None,
1355            },
1356        ];
1357        let create_table_expr =
1358            &mut build_create_table_expr(&table_reference, &default_schema, default_engine())?;
1359
1360        create_table_expr.engine = METRIC_ENGINE_NAME.to_string();
1361        create_table_expr
1362            .table_options
1363            .insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), "true".to_string());
1364
1365        // create physical table
1366        let res = statement_executor
1367            .create_table_inner(
1368                create_table_expr,
1369                None,
1370                ctx.clone(),
1371                TriggerReason::AutoCreate,
1372            )
1373            .await;
1374
1375        match res {
1376            Ok(table) => Ok(table),
1377            Err(err) => {
1378                error!(err; "Failed to create table {table_reference}");
1379                Err(err)
1380            }
1381        }
1382    }
1383
1384    async fn get_table(
1385        &self,
1386        catalog: &str,
1387        schema: &str,
1388        table: &str,
1389    ) -> Result<Option<TableRef>> {
1390        self.catalog_manager
1391            .table(catalog, schema, table, None)
1392            .await
1393            .context(CatalogSnafu)
1394    }
1395
1396    fn get_create_table_expr_on_demand(
1397        &self,
1398        req: &RowInsertRequest,
1399        create_type: &AutoCreateTableType,
1400        ctx: &QueryContextRef,
1401        semantic_index: Option<&PerTableSemanticIndex>,
1402    ) -> Result<CreateTableExpr> {
1403        let schema = ctx.current_schema();
1404        let mut table_options = std::collections::HashMap::with_capacity(4);
1405        fill_table_options_for_create(&mut table_options, create_type, ctx);
1406        apply_per_table_semantic_options(
1407            &mut table_options,
1408            semantic_index,
1409            ctx.current_schema().as_str(),
1410            &req.table_name,
1411        );
1412
1413        let engine_name = if let AutoCreateTableType::Logical(_) = create_type {
1414            // engine should be metric engine when creating logical tables.
1415            METRIC_ENGINE_NAME
1416        } else {
1417            default_engine()
1418        };
1419
1420        let table_ref = TableReference::full(ctx.current_catalog(), &schema, &req.table_name);
1421        // SAFETY: `req.rows` is guaranteed to be `Some` by `handle_row_inserts_with_create_type()`.
1422        let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1423        let mut create_table_expr =
1424            build_create_table_expr(&table_ref, request_schema, engine_name)?;
1425
1426        // extension set by the Splunk HEC handler for identity path
1427        if ctx.extension(SPLUNK_PK_METADATA_ORDER_KEY).is_some() {
1428            reorder_splunk_primary_keys(&mut create_table_expr.primary_keys);
1429        }
1430
1431        debug!("Ensuring table `{table_ref}` exists for insert");
1432        create_table_expr.table_options.extend(table_options);
1433        Ok(create_table_expr)
1434    }
1435
1436    /// Returns an alter table expression if it finds new columns in the request.
1437    /// When `accommodate_existing_schema` is false, it always adds columns if not exist.
1438    /// When `accommodate_existing_schema` is true, it may modify the input `req` to
1439    /// accommodate it with existing schema. See [`create_or_alter_tables_on_demand`](Self::create_or_alter_tables_on_demand)
1440    /// for more details.
1441    /// When `is_single_value` is true, it also rejects native-histogram/float kind changes.
1442    /// When both options are true, it considers fields when modifying the input `req`.
1443    fn get_alter_table_expr_on_demand(
1444        &self,
1445        req: &mut RowInsertRequest,
1446        table: &TableRef,
1447        ctx: &QueryContextRef,
1448        accommodate_existing_schema: bool,
1449        is_single_value: bool,
1450        alter_existing: bool,
1451    ) -> Result<Option<AlterTableExpr>> {
1452        if !alter_existing {
1453            return Ok(None);
1454        }
1455
1456        let catalog_name = ctx.current_catalog();
1457        let schema_name = ctx.current_schema();
1458        let table_name = table.table_info().name.clone();
1459
1460        // Never auto-alter a system-defined table to fit a write; a request
1461        // with unknown columns fails instead.
1462        if is_ddl_reserved_table(&schema_name, &table_name) {
1463            return Ok(None);
1464        }
1465
1466        let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1467        let request_field_count = request_schema
1468            .iter()
1469            .filter(|col| col.semantic_type == SemanticType::Field as i32)
1470            .count();
1471        let column_exprs = ColumnExpr::from_column_schemas(request_schema);
1472        let add_columns = expr_helper::extract_add_columns_expr(&table.schema(), column_exprs)?;
1473        let Some(mut add_columns) = add_columns else {
1474            return Ok(None);
1475        };
1476
1477        if is_single_value {
1478            let request_is_native_histogram = request_is_native_histogram(request_schema);
1479            let table_is_native_histogram = table_is_native_histogram(table);
1480            ensure!(
1481                request_is_native_histogram == table_is_native_histogram,
1482                InvalidInsertRequestSnafu {
1483                    reason: format!(
1484                        "Table `{table_name}` cannot mix native histogram and float sample fields"
1485                    ),
1486                }
1487            );
1488        }
1489
1490        // If accommodate_existing_schema is true, update request schema for Timestamp/Field columns
1491        if accommodate_existing_schema {
1492            let table_schema = table.schema();
1493            // Find timestamp column name
1494            let ts_col_name = table_schema.timestamp_column().map(|c| c.name.clone());
1495            // Find field column name if there is only one and `is_single_value` is true.
1496            let mut field_col_name = None;
1497            if is_single_value && request_field_count <= 1 {
1498                let mut multiple_field_cols = false;
1499                table.field_columns().for_each(|col| {
1500                    if field_col_name.is_none() {
1501                        field_col_name = Some(col.name.clone());
1502                    } else {
1503                        multiple_field_cols = true;
1504                    }
1505                });
1506                if multiple_field_cols {
1507                    field_col_name = None;
1508                }
1509            }
1510
1511            // Update column name in request schema for Timestamp/Field columns
1512            if let Some(rows) = req.rows.as_mut() {
1513                for col in &mut rows.schema {
1514                    match col.semantic_type {
1515                        x if x == SemanticType::Timestamp as i32 => {
1516                            if let Some(ref ts_name) = ts_col_name
1517                                && col.column_name != *ts_name
1518                            {
1519                                col.column_name = ts_name.clone();
1520                            }
1521                        }
1522                        x if x == SemanticType::Field as i32 => {
1523                            if let Some(ref field_name) = field_col_name
1524                                && col.column_name != *field_name
1525                            {
1526                                col.column_name = field_name.clone();
1527                            }
1528                        }
1529                        _ => {}
1530                    }
1531                }
1532            }
1533
1534            // Only keep columns that are tags or non-single field.
1535            add_columns.add_columns.retain(|col| {
1536                let def = col.column_def.as_ref().unwrap();
1537                def.semantic_type == SemanticType::Tag as i32
1538                    || (def.semantic_type == SemanticType::Field as i32 && field_col_name.is_none())
1539            });
1540
1541            if add_columns.add_columns.is_empty() {
1542                return Ok(None);
1543            }
1544        }
1545
1546        Ok(Some(AlterTableExpr {
1547            catalog_name: catalog_name.to_string(),
1548            schema_name: schema_name.clone(),
1549            table_name: table_name.clone(),
1550            kind: Some(Kind::AddColumns(add_columns)),
1551        }))
1552    }
1553
1554    /// Creates a table with options.
1555    async fn create_physical_table(
1556        &self,
1557        mut create_table_expr: CreateTableExpr,
1558        partitions: Option<Partitions>,
1559        ctx: &QueryContextRef,
1560        statement_executor: &StatementExecutor,
1561    ) -> Result<TableRef> {
1562        let res = statement_executor
1563            .create_table_inner(
1564                &mut create_table_expr,
1565                partitions,
1566                ctx.clone(),
1567                TriggerReason::AutoCreate,
1568            )
1569            .await;
1570
1571        let table_ref = TableReference::full(
1572            &create_table_expr.catalog_name,
1573            &create_table_expr.schema_name,
1574            &create_table_expr.table_name,
1575        );
1576
1577        match res {
1578            Ok(table) => {
1579                validate_trace_table_model(&table.table_info(), ctx)?;
1580                Ok(table)
1581            }
1582            Err(err) => {
1583                error!(err; "Failed to create table {}", table_ref);
1584                Err(err)
1585            }
1586        }
1587    }
1588
1589    async fn create_logical_tables(
1590        &self,
1591        create_table_exprs: Vec<CreateTableExpr>,
1592        ctx: &QueryContextRef,
1593        statement_executor: &StatementExecutor,
1594    ) -> Result<Vec<TableRef>> {
1595        let res = statement_executor
1596            .create_logical_tables(&create_table_exprs, ctx.clone(), TriggerReason::AutoCreate)
1597            .await;
1598
1599        match res {
1600            Ok(res) => Ok(res),
1601            Err(err) => {
1602                let failed_tables = create_table_exprs
1603                    .into_iter()
1604                    .map(|expr| {
1605                        format!(
1606                            "{}.{}.{}",
1607                            expr.catalog_name, expr.schema_name, expr.table_name
1608                        )
1609                    })
1610                    .collect::<Vec<_>>();
1611                error!(
1612                    err;
1613                    "Failed to create logical tables {:?}",
1614                    failed_tables
1615                );
1616                Err(err)
1617            }
1618        }
1619    }
1620
1621    pub fn node_manager(&self) -> &NodeManagerRef {
1622        &self.node_manager
1623    }
1624
1625    pub fn partition_manager(&self) -> &PartitionRuleManagerRef {
1626        &self.partition_manager
1627    }
1628
1629    pub fn table_flownode_set_cache(&self) -> &TableFlownodeSetCacheRef {
1630        &self.table_flownode_set_cache
1631    }
1632}
1633
1634fn request_is_native_histogram(request_schema: &[ColumnSchema]) -> bool {
1635    let mut fields = request_schema
1636        .iter()
1637        .filter(|col| col.semantic_type == SemanticType::Field as i32);
1638    let Some(col) = fields.next() else {
1639        return false;
1640    };
1641
1642    fields.next().is_none()
1643        && api::helper::is_column_type_value_eq(
1644            col.datatype,
1645            col.datatype_extension.clone(),
1646            native_histogram_value_type(),
1647        )
1648}
1649
1650/// Returns the table's time index unit, if any. A metric table without a
1651/// timestamp column is left alone by the unit alignment: the metric engine
1652/// rejects it anyway.
1653fn table_time_index_unit(table: &TableRef) -> Option<TimeUnit> {
1654    table
1655        .table_info()
1656        .meta
1657        .schema
1658        .timestamp_column()
1659        .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit()))
1660}
1661
1662fn convert_rows_time_unit(rows: &mut Rows, target_unit: TimeUnit) -> Result<()> {
1663    let Some(ts_index) = rows
1664        .schema
1665        .iter()
1666        .position(|col| col.semantic_type == SemanticType::Timestamp as i32)
1667    else {
1668        return Ok(());
1669    };
1670    let Some(source_unit) = ColumnDataType::try_from(rows.schema[ts_index].datatype)
1671        .ok()
1672        .and_then(api::helper::timestamp_unit)
1673    else {
1674        return Ok(());
1675    };
1676    if source_unit == target_unit {
1677        return Ok(());
1678    }
1679
1680    rows.schema[ts_index].datatype = api::helper::timestamp_datatype(target_unit) as i32;
1681    // Timestamp columns never carry a datatype extension.
1682    rows.schema[ts_index].datatype_extension = None;
1683
1684    // Note: the schema is rewritten before the rows are converted, so an
1685    // overflow error mid-batch leaves this request half-converted. That is
1686    // harmless: the error aborts the whole insert request.
1687    //
1688    // `validate_column_count_match` guarantees every row carries exactly one
1689    // value per schema column, so the time index position is directly in
1690    // bounds; no per-value search is needed.
1691    for row in &mut rows.rows {
1692        debug_assert_eq!(row.values.len(), rows.schema.len());
1693        let value = &mut row.values[ts_index];
1694        let Some(value_data) = value.value_data.take() else {
1695            continue;
1696        };
1697        value.value_data =
1698            convert_timestamp_value_data(value_data, source_unit, target_unit, ts_index)?;
1699    }
1700    Ok(())
1701}
1702
1703fn convert_timestamp_value_data(
1704    value_data: ValueData,
1705    source_unit: TimeUnit,
1706    target_unit: TimeUnit,
1707    column_index: usize,
1708) -> Result<Option<ValueData>> {
1709    let timestamp = match value_data {
1710        ValueData::TimestampSecondValue(v) => Timestamp::new_second(v),
1711        ValueData::TimestampMillisecondValue(v) => Timestamp::new_millisecond(v),
1712        ValueData::TimestampMicrosecondValue(v) => Timestamp::new_microsecond(v),
1713        ValueData::TimestampNanosecondValue(v) => Timestamp::new_nanosecond(v),
1714        // Null or non-timestamp value; nothing to convert.
1715        other => return Ok(Some(other)),
1716    };
1717    let converted = timestamp
1718        .convert_to(target_unit)
1719        .with_context(|| InvalidInsertRequestSnafu {
1720            reason: format!(
1721                "timestamp column {column_index} value {} in unit {source_unit:?} overflows when converting to unit {target_unit:?}",
1722                timestamp.value()
1723            ),
1724        })?;
1725    Ok(api::helper::to_grpc_value(datatypes::value::Value::Timestamp(converted)).value_data)
1726}
1727
1728fn table_is_native_histogram(table: &TableRef) -> bool {
1729    let mut fields = table.field_columns();
1730    let Some(col) = fields.next() else {
1731        return false;
1732    };
1733
1734    fields.next().is_none() && is_native_histogram_value_type(&col.data_type)
1735}
1736
1737fn validate_column_count_match(requests: &RowInsertRequests) -> Result<()> {
1738    for request in &requests.inserts {
1739        let rows = request.rows.as_ref().unwrap();
1740        let column_count = rows.schema.len();
1741        rows.rows.iter().try_for_each(|r| {
1742            ensure!(
1743                r.values.len() == column_count,
1744                InvalidInsertRequestSnafu {
1745                    reason: format!(
1746                        "column count mismatch, columns: {}, values: {}",
1747                        column_count,
1748                        r.values.len()
1749                    )
1750                }
1751            );
1752            Ok(())
1753        })?;
1754    }
1755    Ok(())
1756}
1757
1758/// Rejects writes from a different built-in trace model before schema mutation.
1759/// Unstamped, explicitly created tables remain subject to normal schema validation.
1760pub fn validate_trace_table_model(table_info: &TableInfo, ctx: &QueryContextRef) -> Result<()> {
1761    let Some(expected @ (TABLE_DATA_MODEL_TRACE_V1 | TABLE_DATA_MODEL_TRACE_V2)) =
1762        ctx.extension(SEMANTIC_PIPELINE)
1763    else {
1764        return Ok(());
1765    };
1766    if let Some(actual) = table_info.meta.options.data_model() {
1767        ensure!(
1768            actual == expected,
1769            InvalidInsertRequestSnafu {
1770                reason: format!(
1771                    "Trace table `{}` uses {actual}, but the request uses {expected}",
1772                    table_info.name,
1773                ),
1774            }
1775        );
1776    }
1777    Ok(())
1778}
1779
1780/// Fill table options for a new table by create type.
1781pub fn fill_table_options_for_create(
1782    table_options: &mut std::collections::HashMap<String, String>,
1783    create_type: &AutoCreateTableType,
1784    ctx: &QueryContextRef,
1785) {
1786    for key in VALID_TABLE_OPTION_KEYS {
1787        if let Some(value) = ctx.extension(key) {
1788            table_options.insert(key.to_string(), value.to_string());
1789        }
1790    }
1791
1792    // Semantic keys use their own vocabulary instead of the fixed option list.
1793    for (key, value) in ctx.extensions() {
1794        if is_semantic_option_key(&key) && validate_semantic_option(&key, &value) {
1795            table_options.insert(key, value);
1796        }
1797    }
1798
1799    match create_type {
1800        AutoCreateTableType::Logical(physical_table) => {
1801            table_options.insert(
1802                LOGICAL_TABLE_METADATA_KEY.to_string(),
1803                physical_table.clone(),
1804            );
1805        }
1806        AutoCreateTableType::Physical => {
1807            if let Some(append_mode) = ctx.extension(APPEND_MODE_KEY) {
1808                table_options.insert(APPEND_MODE_KEY.to_string(), append_mode.to_string());
1809            }
1810            if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
1811                table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
1812            }
1813            if let Some(time_window) = ctx.extension(TWCS_TIME_WINDOW) {
1814                table_options.insert(TWCS_TIME_WINDOW.to_string(), time_window.to_string());
1815                // We need to set the compaction type explicitly.
1816                table_options.insert(
1817                    COMPACTION_TYPE.to_string(),
1818                    COMPACTION_TYPE_TWCS.to_string(),
1819                );
1820            }
1821        }
1822        // Set append_mode to true for log table.
1823        // because log tables should keep rows with the same ts and tags.
1824        AutoCreateTableType::Log => {
1825            table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1826        }
1827        AutoCreateTableType::LastNonNull => {
1828            if ctx
1829                .extension(APPEND_MODE_KEY)
1830                .is_some_and(|value| value.eq_ignore_ascii_case("true"))
1831            {
1832                table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1833                table_options.insert(MERGE_MODE_KEY.to_string(), "last_row".to_string());
1834            } else if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
1835                table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
1836            } else {
1837                table_options.insert(MERGE_MODE_KEY.to_string(), "last_non_null".to_string());
1838            }
1839        }
1840        AutoCreateTableType::Trace { .. } => {
1841            table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1842        }
1843    }
1844}
1845
1846/// The parsed per-table semantic index: `{schema -> {table -> {key -> value}}}`,
1847/// produced by the OTLP metrics encode path (where one metric can fan out into
1848/// several tables with distinct keys) and the Prometheus remote write v2 path
1849/// (where per-series metadata declares type/unit, and a series may override its
1850/// target schema).
1851pub type PerTableSemanticIndex = BTreeMap<String, BTreeMap<String, BTreeMap<String, String>>>;
1852
1853/// Parses the per-table semantic index off the context extension. Call once per
1854/// create-planning round: a first write creating N tables would otherwise
1855/// re-parse the whole index N times. `None` when the request carries no index
1856/// (logs, traces, Prom RW v1) or it fails to parse.
1857pub fn parse_per_table_semantic_index(ctx: &QueryContextRef) -> Option<PerTableSemanticIndex> {
1858    let raw = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY)?;
1859    match serde_json::from_str(raw) {
1860        Ok(index) => Some(index),
1861        Err(_) => {
1862            warn!("failed to parse semantic per-table index, skipping per-table options");
1863            None
1864        }
1865    }
1866}
1867
1868/// Folds the semantic keys of the table being created into `table_options`.
1869///
1870/// Common keys shared by every table in a request travel as plain semantic
1871/// extensions and are handled by [`fill_table_options_for_create`]; this
1872/// carries only the per-table tail and is applied after it, so a per-table
1873/// value (e.g. `declared` quality) wins. Keys are re-checked against the
1874/// vocabulary defensively.
1875pub fn apply_per_table_semantic_options(
1876    table_options: &mut std::collections::HashMap<String, String>,
1877    index: Option<&PerTableSemanticIndex>,
1878    schema: &str,
1879    table_name: &str,
1880) {
1881    let Some(entry) = index
1882        .and_then(|index| index.get(schema))
1883        .and_then(|tables| tables.get(table_name))
1884    else {
1885        return;
1886    };
1887    for (key, value) in entry {
1888        if is_semantic_option_key(key) && validate_semantic_option(key, value) {
1889            table_options.insert(key.clone(), value.clone());
1890        }
1891    }
1892}
1893
1894pub fn build_create_table_expr(
1895    table: &TableReference,
1896    request_schema: &[ColumnSchema],
1897    engine: &str,
1898) -> Result<CreateTableExpr> {
1899    expr_helper::create_table_expr_by_column_schemas(table, request_schema, engine, None)
1900}
1901
1902/// `QueryContext` extension key the Splunk HEC handler sets (to `"true"`) on its identity
1903/// path to request metadata-first primary-key ordering at table creation. It is absent for
1904/// user-supplied pipelines, so their primary-key order is left untouched.
1905pub const SPLUNK_PK_METADATA_ORDER_KEY: &str = "splunk_pk_metadata_order";
1906
1907/// Moves Splunk's metadata tags (`host`, `source`, `sourcetype`) to the front of the
1908/// primary key, keeping the relative order of the remaining tags.
1909fn reorder_splunk_primary_keys(primary_keys: &mut [String]) {
1910    const LEAD: [&str; 3] = ["host", "source", "sourcetype"];
1911    // Stable sort: `LEAD` columns move to the front in `host`/`source`/`sourcetype` order;
1912    // every other column keeps its existing relative position.
1913    primary_keys.sort_by_key(|name| {
1914        LEAD.iter()
1915            .position(|&lead| lead == name.as_str())
1916            .unwrap_or(LEAD.len())
1917    });
1918}
1919
1920/// Result of `create_or_alter_tables_on_demand`.
1921struct CreateAlterTableResult {
1922    /// table ids of ttl=instant tables.
1923    instant_table_ids: HashSet<TableId>,
1924    /// Table Info of the created tables.
1925    table_infos: HashMap<TableId, Arc<TableInfo>>,
1926}
1927
1928struct FlowMirrorTask {
1929    requests: HashMap<Peer, RegionInsertRequests>,
1930    num_rows: usize,
1931}
1932
1933impl FlowMirrorTask {
1934    async fn new(
1935        cache: &TableFlownodeSetCacheRef,
1936        requests: impl Iterator<Item = &RegionInsertRequest>,
1937    ) -> Result<Self> {
1938        let mut src_table_reqs: HashMap<TableId, Option<(Vec<Peer>, RegionInsertRequests)>> =
1939            HashMap::new();
1940        let mut num_rows = 0;
1941
1942        for req in requests {
1943            let table_id = RegionId::from_u64(req.region_id).table_id();
1944            match src_table_reqs.get_mut(&table_id) {
1945                Some(Some((_peers, reqs))) => reqs.requests.push(req.clone()),
1946                // already know this is not source table
1947                Some(None) => continue,
1948                _ => {
1949                    // dedup peers
1950                    let peers = cache
1951                        .get(table_id)
1952                        .await
1953                        .context(RequestInsertsSnafu)?
1954                        .unwrap_or_default()
1955                        .values()
1956                        .cloned()
1957                        .collect::<HashSet<_>>()
1958                        .into_iter()
1959                        .collect::<Vec<_>>();
1960
1961                    if !peers.is_empty() {
1962                        let mut reqs = RegionInsertRequests::default();
1963                        reqs.requests.push(req.clone());
1964                        num_rows += reqs
1965                            .requests
1966                            .iter()
1967                            .map(|r| r.rows.as_ref().unwrap().rows.len())
1968                            .sum::<usize>();
1969                        src_table_reqs.insert(table_id, Some((peers, reqs)));
1970                    } else {
1971                        // insert a empty entry to avoid repeat query
1972                        src_table_reqs.insert(table_id, None);
1973                    }
1974                }
1975            }
1976        }
1977
1978        let mut inserts: HashMap<Peer, RegionInsertRequests> = HashMap::new();
1979
1980        for (_table_id, (peers, reqs)) in src_table_reqs
1981            .into_iter()
1982            .filter_map(|(k, v)| v.map(|v| (k, v)))
1983        {
1984            if peers.len() == 1 {
1985                // fast path, zero copy
1986                inserts
1987                    .entry(peers[0].clone())
1988                    .or_default()
1989                    .requests
1990                    .extend(reqs.requests);
1991                continue;
1992            } else {
1993                // TODO(discord9): need to split requests to multiple flownodes
1994                for flownode in peers {
1995                    inserts
1996                        .entry(flownode.clone())
1997                        .or_default()
1998                        .requests
1999                        .extend(reqs.requests.clone());
2000                }
2001            }
2002        }
2003
2004        Ok(Self {
2005            requests: inserts,
2006            num_rows,
2007        })
2008    }
2009
2010    fn detach(self, node_manager: NodeManagerRef) -> Result<()> {
2011        crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(self.num_rows as i64);
2012        for (peer, inserts) in self.requests {
2013            let node_manager = node_manager.clone();
2014            common_runtime::spawn_global(async move {
2015                let result = node_manager
2016                    .flownode(&peer)
2017                    .await
2018                    .handle_inserts(inserts)
2019                    .await
2020                    .context(RequestInsertsSnafu);
2021
2022                match result {
2023                    Ok(resp) => {
2024                        let affected_rows = resp.affected_rows;
2025                        crate::metrics::DIST_MIRROR_ROW_COUNT.inc_by(affected_rows);
2026                        crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(affected_rows as _);
2027                    }
2028                    Err(err) => {
2029                        error!(err; "Failed to insert data into flownode {}", peer);
2030                    }
2031                }
2032            });
2033        }
2034
2035        Ok(())
2036    }
2037}
2038
2039#[cfg(test)]
2040mod tests {
2041    use std::sync::Arc;
2042
2043    use api::helper::ColumnDataTypeWrapper;
2044    use api::v1::helper::{field_column_schema, time_index_column_schema};
2045    use api::v1::{Row, RowInsertRequest, Rows, Value};
2046    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
2047    use common_meta::cache::new_table_flownode_set_cache;
2048    use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler;
2049    use common_meta::test_util::MockDatanodeManager;
2050    use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD;
2051    use common_query::prelude::{greptime_native_histogram, set_default_prefix};
2052    use datatypes::data_type::ConcreteDataType;
2053    use datatypes::schema::ColumnSchema;
2054    use moka::future::Cache;
2055    use session::context::QueryContext;
2056    use table::TableRef;
2057    use table::dist_table::DummyDataSource;
2058    use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType};
2059
2060    use crate::insert::*;
2061    use crate::test_util::{
2062        create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
2063    };
2064
2065    fn make_table_ref_with_schema(
2066        ts_name: &str,
2067        field_name: &str,
2068        field_type: ConcreteDataType,
2069    ) -> TableRef {
2070        let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2071            ColumnSchema::new(
2072                ts_name,
2073                ConcreteDataType::timestamp_millisecond_datatype(),
2074                false,
2075            )
2076            .with_time_index(true),
2077            ColumnSchema::new(field_name, field_type, true),
2078        ])
2079        .unwrap()
2080        .build()
2081        .unwrap();
2082        let meta = TableMetaBuilder::empty()
2083            .schema(Arc::new(schema))
2084            .primary_key_indices(vec![])
2085            .value_indices(vec![1])
2086            .engine("mito")
2087            .next_column_id(0)
2088            .options(Default::default())
2089            .created_on(Default::default())
2090            .build()
2091            .unwrap();
2092        let info = Arc::new(
2093            TableInfoBuilder::default()
2094                .table_id(1)
2095                .table_version(0)
2096                .name("test_table")
2097                .schema_name(DEFAULT_SCHEMA_NAME)
2098                .catalog_name(DEFAULT_CATALOG_NAME)
2099                .desc(None)
2100                .table_type(TableType::Base)
2101                .meta(meta)
2102                .build()
2103                .unwrap(),
2104        );
2105        Arc::new(table::Table::new(
2106            info,
2107            table::metadata::FilterPushDownType::Unsupported,
2108            Arc::new(DummyDataSource),
2109        ))
2110    }
2111
2112    fn make_metric_physical_table_ref_with_time_unit(unit: TimeUnit) -> TableRef {
2113        let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2114            ColumnSchema::new(
2115                greptime_timestamp(),
2116                ConcreteDataType::timestamp_datatype(unit),
2117                false,
2118            )
2119            .with_time_index(true),
2120            ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
2121        ])
2122        .unwrap()
2123        .build()
2124        .unwrap();
2125        let meta = TableMetaBuilder::empty()
2126            .schema(Arc::new(schema))
2127            .primary_key_indices(vec![])
2128            .value_indices(vec![1])
2129            .engine("metric")
2130            .next_column_id(0)
2131            .options(Default::default())
2132            .created_on(Default::default())
2133            .build()
2134            .unwrap();
2135        let info = Arc::new(
2136            TableInfoBuilder::default()
2137                .table_id(1)
2138                .table_version(0)
2139                .name("greptime_physical_table")
2140                .schema_name(DEFAULT_SCHEMA_NAME)
2141                .catalog_name(DEFAULT_CATALOG_NAME)
2142                .desc(None)
2143                .table_type(TableType::Base)
2144                .meta(meta)
2145                .build()
2146                .unwrap(),
2147        );
2148        Arc::new(table::Table::new(
2149            info,
2150            table::metadata::FilterPushDownType::Unsupported,
2151            Arc::new(DummyDataSource),
2152        ))
2153    }
2154
2155    fn ms_row_insert_request(timestamp_ms: i64) -> RowInsertRequest {
2156        ms_row_insert_request_named("my_metric", timestamp_ms)
2157    }
2158
2159    fn ms_row_insert_request_named(table: &str, timestamp_ms: i64) -> RowInsertRequest {
2160        RowInsertRequest {
2161            table_name: table.to_string(),
2162            rows: Some(Rows {
2163                schema: vec![
2164                    time_index_column_schema(
2165                        greptime_timestamp(),
2166                        ColumnDataType::TimestampMillisecond,
2167                    ),
2168                    field_column_schema(greptime_value(), ColumnDataType::Float64),
2169                ],
2170                rows: vec![Row {
2171                    values: vec![
2172                        Value {
2173                            value_data: Some(ValueData::TimestampMillisecondValue(timestamp_ms)),
2174                        },
2175                        Value {
2176                            value_data: Some(ValueData::F64Value(1.0)),
2177                        },
2178                    ],
2179                }],
2180            }),
2181        }
2182    }
2183
2184    /// Converts the time index column of each request to `target_unit`.
2185    fn convert_row_insert_requests_time_unit(
2186        requests: &mut RowInsertRequests,
2187        target_unit: TimeUnit,
2188    ) -> Result<()> {
2189        for request in &mut requests.inserts {
2190            let Some(rows) = request.rows.as_mut() else {
2191                continue;
2192            };
2193            convert_rows_time_unit(rows, target_unit)?;
2194        }
2195        Ok(())
2196    }
2197
2198    #[test]
2199    fn test_convert_row_insert_requests_time_unit_noop_when_matching() {
2200        let mut requests = RowInsertRequests {
2201            inserts: vec![ms_row_insert_request(123)],
2202        };
2203        convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Millisecond).unwrap();
2204        let rows = requests.inserts[0].rows.as_ref().unwrap();
2205        assert_eq!(
2206            rows.schema[0].datatype,
2207            ColumnDataType::TimestampMillisecond as i32
2208        );
2209        assert!(matches!(
2210            rows.rows[0].values[0].value_data,
2211            Some(ValueData::TimestampMillisecondValue(123))
2212        ));
2213    }
2214
2215    #[test]
2216    fn test_convert_row_insert_requests_time_unit_widens_losslessly() {
2217        let mut requests = RowInsertRequests {
2218            inserts: vec![ms_row_insert_request(123)],
2219        };
2220        convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Microsecond).unwrap();
2221        let rows = requests.inserts[0].rows.as_ref().unwrap();
2222        assert_eq!(
2223            rows.schema[0].datatype,
2224            ColumnDataType::TimestampMicrosecond as i32
2225        );
2226        assert!(matches!(
2227            rows.rows[0].values[0].value_data,
2228            Some(ValueData::TimestampMicrosecondValue(123_000))
2229        ));
2230    }
2231
2232    #[test]
2233    fn test_convert_row_insert_requests_time_unit_truncates_on_narrowing() {
2234        // 123_456_789 ns floors to 123_456 us and 123 ms; negative values
2235        // floor towards negative infinity, matching `Timestamp::convert_to`.
2236        let requests = |value_ns: i64| RowInsertRequests {
2237            inserts: vec![RowInsertRequest {
2238                table_name: "my_metric".to_string(),
2239                rows: Some(Rows {
2240                    schema: vec![
2241                        time_index_column_schema(
2242                            greptime_timestamp(),
2243                            ColumnDataType::TimestampNanosecond,
2244                        ),
2245                        field_column_schema(greptime_value(), ColumnDataType::Float64),
2246                    ],
2247                    rows: vec![Row {
2248                        values: vec![
2249                            Value {
2250                                value_data: Some(ValueData::TimestampNanosecondValue(value_ns)),
2251                            },
2252                            Value {
2253                                value_data: Some(ValueData::F64Value(1.0)),
2254                            },
2255                        ],
2256                    }],
2257                }),
2258            }],
2259        };
2260
2261        let mut reqs = requests(123_456_789);
2262        convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
2263        let rows = reqs.inserts[0].rows.as_ref().unwrap();
2264        assert!(matches!(
2265            rows.rows[0].values[0].value_data,
2266            Some(ValueData::TimestampMicrosecondValue(123_456))
2267        ));
2268
2269        let mut reqs = requests(123_456_789);
2270        convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Millisecond).unwrap();
2271        let rows = reqs.inserts[0].rows.as_ref().unwrap();
2272        assert!(matches!(
2273            rows.rows[0].values[0].value_data,
2274            Some(ValueData::TimestampMillisecondValue(123))
2275        ));
2276
2277        let mut reqs = requests(-123_456_789);
2278        convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
2279        let rows = reqs.inserts[0].rows.as_ref().unwrap();
2280        assert!(matches!(
2281            rows.rows[0].values[0].value_data,
2282            Some(ValueData::TimestampMicrosecondValue(-123_457))
2283        ));
2284    }
2285
2286    #[test]
2287    fn test_convert_row_insert_requests_time_unit_overflow() {
2288        let mut requests = RowInsertRequests {
2289            inserts: vec![ms_row_insert_request(i64::MAX)],
2290        };
2291        let err =
2292            convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Nanosecond).unwrap_err();
2293        assert!(err.to_string().contains("overflows"), "{err}");
2294    }
2295
2296    #[test]
2297    fn test_table_time_index_unit() {
2298        assert_eq!(
2299            table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
2300                TimeUnit::Microsecond
2301            )),
2302            Some(TimeUnit::Microsecond)
2303        );
2304        assert_eq!(
2305            table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
2306                TimeUnit::Millisecond
2307            )),
2308            Some(TimeUnit::Millisecond)
2309        );
2310    }
2311
2312    #[tokio::test]
2313    async fn test_accommodate_existing_schema_and_reject_kind_changes() {
2314        let ts_name = "my_ts";
2315        let field_name = "my_field";
2316        let table =
2317            make_table_ref_with_schema(ts_name, field_name, ConcreteDataType::float64_datatype());
2318
2319        // The request uses different names for timestamp and field columns
2320        let mut req = RowInsertRequest {
2321            table_name: "test_table".to_string(),
2322            rows: Some(Rows {
2323                schema: vec![
2324                    time_index_column_schema("ts_wrong", ColumnDataType::TimestampMillisecond),
2325                    field_column_schema("field_wrong", ColumnDataType::Float64),
2326                ],
2327                rows: vec![api::v1::Row {
2328                    values: vec![Value::default(), Value::default()],
2329                }],
2330            }),
2331        };
2332        let ctx = Arc::new(QueryContext::with(
2333            DEFAULT_CATALOG_NAME,
2334            DEFAULT_SCHEMA_NAME,
2335        ));
2336
2337        let kv_backend = prepare_mocked_backend().await;
2338        let inserter = Inserter::new(
2339            catalog::memory::MemoryCatalogManager::new(),
2340            create_partition_rule_manager(kv_backend.clone()).await,
2341            Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
2342            Arc::new(new_table_flownode_set_cache(
2343                String::new(),
2344                Cache::new(100),
2345                kv_backend.clone(),
2346            )),
2347            true,
2348        );
2349        // Do not apply an absent-table plan to a table that appeared concurrently.
2350        assert!(
2351            inserter
2352                .get_alter_table_expr_on_demand(&mut req, &table, &ctx, false, false, false)
2353                .unwrap()
2354                .is_none()
2355        );
2356        let alter_expr = inserter
2357            .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true, true)
2358            .unwrap();
2359        assert!(alter_expr.is_none());
2360
2361        // The request's schema should have updated names for timestamp and field columns
2362        let req_schema = req.rows.as_ref().unwrap().schema.clone();
2363        assert_eq!(req_schema[0].column_name, ts_name);
2364        assert_eq!(req_schema[1].column_name, field_name);
2365
2366        let (datatype, datatype_extension) =
2367            ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
2368                .unwrap()
2369                .into_parts();
2370        let mut histogram_req = RowInsertRequest {
2371            table_name: "test_table".to_string(),
2372            rows: Some(Rows {
2373                schema: vec![
2374                    time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2375                    api::v1::ColumnSchema {
2376                        column_name: greptime_native_histogram().to_string(),
2377                        datatype: datatype as i32,
2378                        semantic_type: SemanticType::Field as i32,
2379                        datatype_extension,
2380                        options: None,
2381                    },
2382                ],
2383                rows: vec![],
2384            }),
2385        };
2386        let error = inserter
2387            .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, false, true, true)
2388            .unwrap_err();
2389        assert!(
2390            error
2391                .to_string()
2392                .contains("cannot mix native histogram and float sample fields")
2393        );
2394
2395        let histogram_table = make_table_ref_with_schema(
2396            "ts",
2397            greptime_native_histogram(),
2398            native_histogram_value_type().clone(),
2399        );
2400        let mut sample_req = RowInsertRequest {
2401            table_name: "test_table".to_string(),
2402            rows: Some(Rows {
2403                schema: vec![
2404                    time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2405                    field_column_schema(greptime_value(), ColumnDataType::Float64),
2406                ],
2407                rows: vec![],
2408            }),
2409        };
2410        let error = inserter
2411            .get_alter_table_expr_on_demand(
2412                &mut sample_req,
2413                &histogram_table,
2414                &ctx,
2415                false,
2416                true,
2417                true,
2418            )
2419            .unwrap_err();
2420        assert!(
2421            error
2422                .to_string()
2423                .contains("cannot mix native histogram and float sample fields")
2424        );
2425    }
2426
2427    #[test]
2428    fn test_native_histogram_detection_survives_prefix_change() {
2429        set_default_prefix(Some("custom")).unwrap();
2430        let table = make_table_ref_with_schema(
2431            "custom_timestamp",
2432            NATIVE_HISTOGRAM_FIELD,
2433            native_histogram_value_type().clone(),
2434        );
2435        let (datatype, datatype_extension) =
2436            ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
2437                .unwrap()
2438                .into_parts();
2439        let request_schema = [api::v1::ColumnSchema {
2440            column_name: greptime_native_histogram().to_string(),
2441            datatype: datatype as i32,
2442            semantic_type: SemanticType::Field as i32,
2443            datatype_extension,
2444            options: None,
2445        }];
2446
2447        assert!(request_is_native_histogram(&request_schema));
2448        assert!(table_is_native_histogram(&table));
2449    }
2450
2451    // Keep global meter registration in one test, isolated by nextest's per-test process.
2452    #[tokio::test]
2453    async fn test_write_meter_admission() {
2454        use std::cell::Cell;
2455        use std::sync::Mutex;
2456        use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
2457
2458        use api::region::RegionResponse;
2459        use api::v1::region::region_request::Body;
2460        use arrow::array::{Int32Array, TimestampMillisecondArray};
2461        use arrow::record_batch::RecordBatch;
2462        use bytes::Bytes;
2463        use common_error::ext::{ErrorExt, RetryHint};
2464        use common_error::status_code::StatusCode;
2465        use common_grpc::flight::{FlightEncoder, FlightMessage};
2466        use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher;
2467        use futures::future::BoxFuture;
2468        use meter_core::ItemCalculator;
2469        use meter_core::collect::{Collect, WriteRejected};
2470        use meter_core::data::MeterRecord;
2471        use meter_core::global::global_registry;
2472        use session::context::Channel;
2473
2474        const CATALOG: &str = "write_meter_test";
2475
2476        #[derive(Default)]
2477        struct Meter {
2478            reject: AtomicBool,
2479            attempts: Mutex<Vec<MeterRecord>>,
2480            accepted_value: AtomicU64,
2481        }
2482
2483        impl Collect for Meter {
2484            fn on_write(
2485                &self,
2486                record: MeterRecord,
2487            ) -> BoxFuture<'_, std::result::Result<(), WriteRejected>> {
2488                Box::pin(async move {
2489                    if record.catalog != CATALOG {
2490                        return Ok(());
2491                    }
2492                    let value = record.value;
2493                    self.attempts.lock().unwrap().push(record);
2494                    if self.reject.load(Ordering::Relaxed) {
2495                        return Err(WriteRejected::new("database row quota exhausted"));
2496                    }
2497                    self.accepted_value.fetch_add(value, Ordering::Relaxed);
2498                    Ok(())
2499                })
2500            }
2501
2502            fn on_read(&self, _: MeterRecord) {}
2503        }
2504
2505        impl ItemCalculator<InstantAndNormalInsertRequests> for Meter {
2506            fn calc(&self, _: &InstantAndNormalInsertRequests) -> u64 {
2507                17
2508            }
2509        }
2510
2511        let kv_backend = prepare_mocked_backend().await;
2512        let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
2513        let (sender, mut dispatched) = tokio::sync::mpsc::channel(16);
2514        let watcher = DatanodeWatcher::new(sender).with_handler(|_, request| {
2515            let rows = match request.body.unwrap() {
2516                Body::Inserts(requests) => requests
2517                    .requests
2518                    .iter()
2519                    .filter_map(|request| request.rows.as_ref())
2520                    .map(|rows| rows.rows.len())
2521                    .sum(),
2522                // The bulk batches below each contain two rows.
2523                Body::BulkInsert(_) => 2,
2524                body => panic!("unexpected request: {body:?}"),
2525            };
2526            Ok(RegionResponse::new(rows))
2527        });
2528        let flow_cache = Cache::new(100);
2529        let inserter = Inserter::new(
2530            catalog::memory::MemoryCatalogManager::new(),
2531            partition_manager,
2532            Arc::new(MockDatanodeManager::new(watcher)),
2533            Arc::new(new_table_flownode_set_cache(
2534                String::new(),
2535                flow_cache.clone(),
2536                kv_backend,
2537            )),
2538            true,
2539        );
2540        let mut table_info = new_test_table_info(1, "table_1", [1].into_iter());
2541        table_info.catalog_name = CATALOG.to_string();
2542        table_info.schema_name = "target_db".to_string();
2543        let table_info = Arc::new(table_info);
2544        let table_infos = HashMap::from_iter([(1, table_info.clone())]);
2545        let ctx = Arc::new(QueryContext::with_channel(
2546            DEFAULT_CATALOG_NAME,
2547            DEFAULT_SCHEMA_NAME,
2548            Channel::Postgres,
2549        ));
2550        let rows_request = || {
2551            let build = |num_rows| RegionInsertRequests {
2552                requests: vec![RegionInsertRequest {
2553                    region_id: RegionId::new(1, 1).as_u64(),
2554                    rows: Some(Rows {
2555                        schema: vec![],
2556                        rows: vec![api::v1::Row { values: vec![] }; num_rows],
2557                    }),
2558                    ..Default::default()
2559                }],
2560            };
2561            InstantAndNormalInsertRequests {
2562                normal_requests: build(3),
2563                instant_requests: build(2),
2564            }
2565        };
2566
2567        // No collector or calculator: ordinary OSS insertion still succeeds.
2568        let output = inserter
2569            .do_request(rows_request(), &table_infos, &ctx)
2570            .await
2571            .unwrap();
2572        assert_eq!(output.meta.cost, 0);
2573        assert!(matches!(output.data, OutputData::AffectedRows(3)));
2574        dispatched.try_recv().unwrap();
2575
2576        let meter = Arc::new(Meter::default());
2577        global_registry().set_collector(meter.clone());
2578        global_registry().register_calculator(meter.clone());
2579        // The dependency controls noop mode; exercise both builds with this test.
2580        let enabled = Cell::new(false);
2581        write_meter!({
2582            enabled.set(true);
2583            MeterRecord::new("probe".into(), "probe".into(), 0, 0, 0)
2584        })
2585        .await
2586        .unwrap();
2587        let enabled = enabled.get();
2588
2589        let output = inserter
2590            .do_request(rows_request(), &table_infos, &ctx)
2591            .await
2592            .unwrap();
2593        assert_eq!(output.meta.cost, if enabled { 17 } else { 0 });
2594        assert!(matches!(output.data, OutputData::AffectedRows(3)));
2595        dispatched.try_recv().unwrap();
2596        assert!(flow_cache.contains_key(&1));
2597        flow_cache.invalidate_all();
2598        meter.reject.store(true, Ordering::Relaxed);
2599        let result = inserter
2600            .do_request(rows_request(), &table_infos, &ctx)
2601            .await;
2602        if enabled {
2603            let error = result.unwrap_err();
2604            assert_eq!(error.status_code(), StatusCode::RateLimited);
2605            assert_eq!(error.retry_hint(), RetryHint::Retryable);
2606            assert!(error.to_string().contains("database row quota exhausted"));
2607            assert!(dispatched.try_recv().is_err());
2608            assert!(
2609                !flow_cache.contains_key(&1),
2610                "rejected write reached flow mirroring"
2611            );
2612            assert_eq!(meter.accepted_value.load(Ordering::Relaxed), 17);
2613            let attempts = meter.attempts.lock().unwrap();
2614            assert_eq!(attempts.len(), 2);
2615            for record in attempts.iter() {
2616                assert_eq!(record.catalog, CATALOG);
2617                assert_eq!(record.schema, "target_db");
2618                assert_eq!(
2619                    (record.rows, record.value, record.source),
2620                    (5, 17, Channel::Postgres as u8)
2621                );
2622            }
2623        } else {
2624            assert_eq!(result.unwrap().meta.cost, 0);
2625            dispatched.try_recv().unwrap();
2626            assert!(meter.attempts.lock().unwrap().is_empty());
2627        }
2628        meter.attempts.lock().unwrap().clear();
2629
2630        // Rejection must also precede admission to the table batcher's queue.
2631        if enabled {
2632            let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
2633            let rows = Rows {
2634                schema: vec![
2635                    api::v1::helper::tag_column_schema("a", ColumnDataType::Int32),
2636                    time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2637                    field_column_schema("b", ColumnDataType::Int32),
2638                ],
2639                rows: vec![api::v1::Row {
2640                    values: vec![
2641                        api::v1::value::ValueData::I32Value(60).into(),
2642                        Value {
2643                            value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue(
2644                                0,
2645                            )),
2646                        },
2647                        api::v1::value::ValueData::I32Value(0).into(),
2648                    ],
2649                }],
2650            };
2651            let error = inserter
2652                .submit_table_rows(rows, table_info.clone(), ctx.clone(), &batcher)
2653                .await
2654                .unwrap_err();
2655            assert_eq!(error.status_code(), StatusCode::RateLimited);
2656            let mut attempts = meter.attempts.lock().unwrap();
2657            assert_eq!(attempts.len(), 1);
2658            assert_eq!(attempts[0].catalog, CATALOG);
2659            assert_eq!(attempts[0].schema, "target_db");
2660            assert_eq!(attempts[0].rows, 1);
2661            attempts.clear();
2662        }
2663
2664        let table = Arc::new(table::Table::new(
2665            table_info.clone(),
2666            table::metadata::FilterPushDownType::Unsupported,
2667            Arc::new(DummyDataSource),
2668        ));
2669        let batch = RecordBatch::try_new(
2670            table_info.meta.schema.arrow_schema().clone(),
2671            vec![
2672                Arc::new(Int32Array::from(vec![60, 70])),
2673                Arc::new(TimestampMillisecondArray::from(vec![0, 1])),
2674                Arc::new(Int32Array::from(vec![0, 0])),
2675            ],
2676        )
2677        .unwrap();
2678        let bulk_insert = |batch: RecordBatch| {
2679            let flight_data = FlightEncoder::default()
2680                .encode(FlightMessage::RecordBatch(batch.clone()))
2681                .into_iter()
2682                .next()
2683                .unwrap();
2684            inserter.handle_bulk_insert(
2685                table.clone(),
2686                flight_data,
2687                batch,
2688                Bytes::new(),
2689                false,
2690                Channel::Grpc,
2691            )
2692        };
2693        // Empty batches bypass admission even while the collector rejects.
2694        assert_eq!(bulk_insert(batch.slice(0, 0)).await.unwrap(), 0);
2695        assert!(meter.attempts.lock().unwrap().is_empty());
2696        assert!(dispatched.try_recv().is_err());
2697
2698        // A collector change between batches takes effect on the very next batch.
2699        for reject in [false, true, false] {
2700            meter.reject.store(reject, Ordering::Relaxed);
2701            let result = bulk_insert(batch.clone()).await;
2702            if enabled && reject {
2703                let error = result.unwrap_err();
2704                assert_eq!(error.status_code(), StatusCode::RateLimited);
2705                assert_eq!(error.retry_hint(), RetryHint::Retryable);
2706                assert!(dispatched.try_recv().is_err());
2707            } else {
2708                assert_eq!(result.unwrap(), 2);
2709                dispatched.try_recv().unwrap();
2710            }
2711        }
2712        {
2713            let attempts = meter.attempts.lock().unwrap();
2714            assert_eq!(attempts.len(), if enabled { 3 } else { 0 });
2715            for record in attempts.iter() {
2716                assert_eq!(record.catalog, CATALOG);
2717                assert_eq!(record.schema, "target_db");
2718                assert_eq!(
2719                    (record.rows, record.value, record.source),
2720                    (2, 0, Channel::Grpc as u8)
2721                );
2722            }
2723        }
2724        assert_eq!(
2725            meter.accepted_value.load(Ordering::Relaxed),
2726            if enabled { 17 } else { 0 }
2727        );
2728
2729        // Aggregate repeated database targets and retain admission across nested
2730        // batching without changing the caller's reusable context.
2731        meter.attempts.lock().unwrap().clear();
2732        let original = Arc::new(QueryContext::with_channel(
2733            CATALOG,
2734            "a",
2735            Channel::Prometheus,
2736        ));
2737        let mut batches = ["a", "b", "a"].map(|schema| {
2738            let ctx = if schema == "a" {
2739                original.clone()
2740            } else {
2741                Arc::new(QueryContext::with_channel(
2742                    CATALOG,
2743                    schema,
2744                    Channel::Prometheus,
2745                ))
2746            };
2747            (
2748                ctx,
2749                RowInsertRequests {
2750                    inserts: vec![RowInsertRequest {
2751                        table_name: "data".into(),
2752                        rows: Some(Rows {
2753                            schema: vec![],
2754                            rows: vec![api::v1::Row::default(); 2],
2755                        }),
2756                    }],
2757                },
2758            )
2759        });
2760        admit_row_insert_batches(&mut batches).await.unwrap();
2761        admit_row_insert_batches(&mut batches).await.unwrap();
2762        assert_eq!(original.write_rows_to_admit(CATALOG, "a", 4), 4);
2763        for (ctx, _) in &batches {
2764            assert_eq!(
2765                ctx.write_rows_to_admit(CATALOG, &ctx.current_schema(), 2),
2766                0
2767            );
2768            assert_eq!(ctx.channel(), Channel::Prometheus);
2769        }
2770        {
2771            let attempts = meter.attempts.lock().unwrap();
2772            let totals = attempts
2773                .iter()
2774                .map(|r| (r.schema.as_str(), r.rows, r.value))
2775                .collect::<Vec<_>>();
2776            assert_eq!(
2777                totals,
2778                if enabled {
2779                    vec![("a", 4, 0), ("b", 2, 0)]
2780                } else {
2781                    vec![]
2782                }
2783            );
2784        }
2785        meter.attempts.lock().unwrap().clear();
2786        meter.reject.store(false, Ordering::Relaxed);
2787        let ctx = Arc::new(QueryContext::with_channel(
2788            CATALOG,
2789            "logical",
2790            Channel::Otlp,
2791        ));
2792        let admitted = admit_write(2, &ctx).await.unwrap();
2793        let mut requests = RowInsertRequests {
2794            inserts: vec![RowInsertRequest {
2795                table_name: "metric".to_string(),
2796                rows: Some(Rows {
2797                    schema: vec![],
2798                    rows: vec![api::v1::Row::default(); 2],
2799                }),
2800            }],
2801        };
2802        let original = requests.clone();
2803        let cost = Inserter::meter_row_inserts(&mut requests, &admitted)
2804            .await
2805            .unwrap();
2806        assert_eq!(cost, if enabled { 17 } else { 0 });
2807        assert_eq!(requests, original);
2808        let records = meter
2809            .attempts
2810            .lock()
2811            .unwrap()
2812            .iter()
2813            .map(|record| (record.rows, record.value))
2814            .collect::<Vec<_>>();
2815        assert_eq!(
2816            records,
2817            if enabled {
2818                vec![(2, 0), (0, 17)]
2819            } else {
2820                vec![]
2821            }
2822        );
2823        meter.reject.store(true, Ordering::Relaxed);
2824        let result = Inserter::meter_row_inserts(&mut requests, &admitted).await;
2825        if enabled {
2826            assert_eq!(result.unwrap_err().status_code(), StatusCode::RateLimited);
2827        } else {
2828            assert_eq!(result.unwrap(), 0);
2829        }
2830        assert_eq!(requests, original);
2831    }
2832
2833    #[test]
2834    fn test_skip_wal_does_not_change_table_options() {
2835        check_skip_wal_does_not_change_table_options(false);
2836        check_skip_wal_does_not_change_table_options(true);
2837    }
2838
2839    fn check_skip_wal_does_not_change_table_options(skip_wal: bool) {
2840        let ctx = Arc::new(QueryContext::with(
2841            DEFAULT_CATALOG_NAME,
2842            DEFAULT_SCHEMA_NAME,
2843        ));
2844        ctx.set_skip_wal(skip_wal);
2845        let mut options = Default::default();
2846        fill_table_options_for_create(&mut options, &AutoCreateTableType::Physical, &ctx);
2847        assert!(!options.contains_key(session::hints::INSERT_SKIP_WAL_HINT));
2848        assert!(!options.contains_key("skip_wal"));
2849    }
2850
2851    #[test]
2852    fn test_last_non_null_create_options_preserve_default_without_append_mode() {
2853        let ctx = Arc::new(QueryContext::with(
2854            DEFAULT_CATALOG_NAME,
2855            DEFAULT_SCHEMA_NAME,
2856        ));
2857        let mut table_options = Default::default();
2858
2859        fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
2860
2861        assert_eq!(
2862            Some("last_non_null"),
2863            table_options.get(MERGE_MODE_KEY).map(String::as_str)
2864        );
2865        assert!(!table_options.contains_key(APPEND_MODE_KEY));
2866    }
2867
2868    #[test]
2869    fn test_fill_table_options_copies_semantic_extensions() {
2870        use table::requests::{
2871            SEMANTIC_METRIC_TYPE, SEMANTIC_PER_TABLE_INDEX_KEY, SEMANTIC_SIGNAL_TYPE,
2872            SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, SOURCE_OPENTELEMETRY,
2873        };
2874
2875        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
2876        ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
2877        ctx.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
2878        ctx.set_extension(SEMANTIC_SOURCE_VERSION, "2.0");
2879        ctx.set_extension(SEMANTIC_METRIC_TYPE, "bogus");
2880        // The internal transport key must NOT be copied into table options.
2881        ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, "{}");
2882        let ctx = Arc::new(ctx);
2883        let mut table_options = Default::default();
2884
2885        fill_table_options_for_create(&mut table_options, &AutoCreateTableType::Physical, &ctx);
2886
2887        assert_eq!(
2888            Some(SIGNAL_TYPE_METRIC),
2889            table_options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str)
2890        );
2891        assert_eq!(
2892            Some(SOURCE_OPENTELEMETRY),
2893            table_options.get(SEMANTIC_SOURCE).map(String::as_str)
2894        );
2895        assert_eq!(
2896            Some("2.0"),
2897            table_options
2898                .get(SEMANTIC_SOURCE_VERSION)
2899                .map(String::as_str)
2900        );
2901        assert!(!table_options.contains_key(SEMANTIC_METRIC_TYPE));
2902        assert!(!table_options.contains_key(SEMANTIC_PER_TABLE_INDEX_KEY));
2903    }
2904
2905    #[test]
2906    fn test_apply_per_table_semantic_options() {
2907        use table::requests::{
2908            SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, SEMANTIC_PER_TABLE_INDEX_KEY,
2909        };
2910
2911        let index = format!(
2912            r#"{{
2913            "{DEFAULT_SCHEMA_NAME}": {{
2914                "http_requests_total": {{
2915                    "greptime.semantic.metric.type": "counter",
2916                    "greptime.semantic.metric.unit": "By",
2917                    "greptime.semantic.metric.type_BOGUS": "x"
2918                }},
2919                "other_table": {{
2920                    "greptime.semantic.metric.type": "gauge"
2921                }}
2922            }},
2923            "other_schema": {{
2924                "http_requests_total": {{
2925                    "greptime.semantic.metric.type": "gauge"
2926                }}
2927            }}
2928        }}"#
2929        );
2930        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
2931        ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index);
2932        let ctx = Arc::new(ctx);
2933
2934        let index = parse_per_table_semantic_index(&ctx);
2935        assert!(index.is_some());
2936        let index = index.as_ref();
2937
2938        let mut table_options = std::collections::HashMap::new();
2939        apply_per_table_semantic_options(
2940            &mut table_options,
2941            index,
2942            DEFAULT_SCHEMA_NAME,
2943            "http_requests_total",
2944        );
2945        // The write schema's entry applies — not other_schema's `gauge`.
2946        assert_eq!(
2947            table_options.get(SEMANTIC_METRIC_TYPE).map(String::as_str),
2948            Some("counter")
2949        );
2950        assert_eq!(
2951            table_options.get(SEMANTIC_METRIC_UNIT).map(String::as_str),
2952            Some("By")
2953        );
2954        // The unknown key is rejected by the vocabulary check; other tables' keys
2955        // never appear.
2956        assert!(!table_options.contains_key("greptime.semantic.metric.type_BOGUS"));
2957        assert_eq!(table_options.len(), 2);
2958
2959        let mut empty = std::collections::HashMap::new();
2960        apply_per_table_semantic_options(&mut empty, index, DEFAULT_SCHEMA_NAME, "not_in_index");
2961        assert!(empty.is_empty());
2962
2963        // A schema with no entry is a no-op even when the table name matches
2964        // elsewhere.
2965        let mut opts = std::collections::HashMap::new();
2966        apply_per_table_semantic_options(
2967            &mut opts,
2968            index,
2969            "schema_without_entry",
2970            "http_requests_total",
2971        );
2972        assert!(opts.is_empty());
2973
2974        // No extension at all parses to no index (e.g. logs / Prom RW v1).
2975        let bare = Arc::new(QueryContext::with(
2976            DEFAULT_CATALOG_NAME,
2977            DEFAULT_SCHEMA_NAME,
2978        ));
2979        assert!(parse_per_table_semantic_index(&bare).is_none());
2980        let mut opts = std::collections::HashMap::new();
2981        apply_per_table_semantic_options(
2982            &mut opts,
2983            None,
2984            DEFAULT_SCHEMA_NAME,
2985            "http_requests_total",
2986        );
2987        assert!(opts.is_empty());
2988    }
2989
2990    #[test]
2991    fn test_last_non_null_create_options_preserve_default_with_append_mode_false() {
2992        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
2993        ctx.set_extension(APPEND_MODE_KEY, "false");
2994        let ctx = Arc::new(ctx);
2995        let mut table_options = Default::default();
2996
2997        fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
2998
2999        assert!(!table_options.contains_key(APPEND_MODE_KEY));
3000        assert_eq!(
3001            Some("last_non_null"),
3002            table_options.get(MERGE_MODE_KEY).map(String::as_str)
3003        );
3004    }
3005
3006    #[test]
3007    fn test_last_non_null_create_options_use_configured_merge_mode() {
3008        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3009        ctx.set_extension(MERGE_MODE_KEY, "last_row");
3010        let ctx = Arc::new(ctx);
3011        let mut table_options = Default::default();
3012
3013        fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3014
3015        assert_eq!(
3016            Some("last_row"),
3017            table_options.get(MERGE_MODE_KEY).map(String::as_str)
3018        );
3019        assert!(!table_options.contains_key(APPEND_MODE_KEY));
3020    }
3021
3022    #[test]
3023    fn test_last_non_null_create_options_use_last_row_with_append_mode_true() {
3024        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3025        ctx.set_extension(APPEND_MODE_KEY, "true");
3026        let ctx = Arc::new(ctx);
3027        let mut table_options = Default::default();
3028
3029        fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3030
3031        assert_eq!(
3032            Some("true"),
3033            table_options.get(APPEND_MODE_KEY).map(String::as_str)
3034        );
3035        assert_eq!(
3036            Some("last_row"),
3037            table_options.get(MERGE_MODE_KEY).map(String::as_str)
3038        );
3039    }
3040
3041    struct UnexpectedBatcher;
3042
3043    #[async_trait::async_trait]
3044    impl PendingRowsBatcher for UnexpectedBatcher {
3045        async fn acquire(&self) -> Result<Arc<tokio::sync::OwnedSemaphorePermit>> {
3046            panic!("empty writes must not acquire batch admission")
3047        }
3048
3049        async fn submit(
3050            &self,
3051            _table_info: TableInfoRef,
3052            _batch: arrow::record_batch::RecordBatch,
3053            _ctx: QueryContextRef,
3054            _permit: Arc<tokio::sync::OwnedSemaphorePermit>,
3055        ) -> Result<usize> {
3056            panic!("empty writes must not submit a batch")
3057        }
3058    }
3059
3060    async fn batcher_test_inserter() -> Inserter {
3061        let kv_backend = prepare_mocked_backend().await;
3062        Inserter::new(
3063            catalog::memory::MemoryCatalogManager::new(),
3064            create_partition_rule_manager(kv_backend.clone()).await,
3065            Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3066            Arc::new(new_table_flownode_set_cache(
3067                String::new(),
3068                Cache::new(100),
3069                kv_backend,
3070            )),
3071            true,
3072        )
3073    }
3074
3075    #[tokio::test]
3076    async fn test_logical_batcher_eligibility() {
3077        use catalog::RegisterTableRequest;
3078        use catalog::memory::MemoryCatalogManager;
3079        use common_meta::instruction::{CacheIdent, CreateFlow};
3080        use common_meta::kv_backend::KvBackendRef;
3081        use common_meta::kv_backend::memory::MemoryKvBackend;
3082        use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder};
3083        let requests = RowInsertRequests {
3084            inserts: vec![RowInsertRequest {
3085                table_name: "test_table".to_string(),
3086                rows: None,
3087            }],
3088        };
3089        let original =
3090            make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3091                .table_info();
3092        for case in [
3093            "eligible",
3094            "physical",
3095            "ordinary",
3096            "instant",
3097            "disabled",
3098            "hint",
3099            "flow",
3100            "required_tag",
3101            "default",
3102        ] {
3103            let mut info = (*original).clone();
3104            info.meta.engine = METRIC_ENGINE_NAME.to_string();
3105            info.meta.options.extra_options.insert(
3106                LOGICAL_TABLE_METADATA_KEY.to_string(),
3107                "physical".to_string(),
3108            );
3109            let mut ctx = QueryContext::arc().fork();
3110            match case {
3111                "physical" => {
3112                    info.meta
3113                        .options
3114                        .extra_options
3115                        .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "other".to_string());
3116                }
3117                "ordinary" => info.meta.engine = "mito".to_string(),
3118                "instant" => info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant),
3119                "hint" => ctx.set_extension(AUTO_CREATE_TABLE_KEY, "false"),
3120                "required_tag" | "default" => {
3121                    let mut columns = info.meta.schema.column_schemas().to_vec();
3122                    let mut tag = ColumnSchema::new(
3123                        "tag",
3124                        ConcreteDataType::string_datatype(),
3125                        case != "required_tag",
3126                    );
3127                    if case == "default" {
3128                        tag = tag
3129                            .with_default_constraint(Some(ColumnDefaultConstraint::null_value()))
3130                            .unwrap();
3131                    }
3132                    columns.push(tag);
3133                    info.meta.schema = Arc::new(
3134                        SchemaBuilder::try_from_columns(columns)
3135                            .unwrap()
3136                            .build()
3137                            .unwrap(),
3138                    );
3139                    info.meta.primary_key_indices = vec![2];
3140                }
3141                _ => {}
3142            }
3143            let catalog = MemoryCatalogManager::with_default_setup();
3144            let table = Arc::new(table::Table::new(
3145                Arc::new(info),
3146                table::metadata::FilterPushDownType::Unsupported,
3147                Arc::new(DummyDataSource),
3148            ));
3149            catalog
3150                .register_table_sync(RegisterTableRequest {
3151                    catalog: DEFAULT_CATALOG_NAME.to_string(),
3152                    schema: DEFAULT_SCHEMA_NAME.to_string(),
3153                    table_name: "test_table".to_string(),
3154                    table_id: 1,
3155                    table,
3156                })
3157                .unwrap();
3158            let mut inserter = batcher_test_inserter().await;
3159            inserter.catalog_manager = catalog;
3160            inserter.auto_create_table = case != "disabled";
3161            let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::default());
3162            inserter.table_flownode_set_cache = Arc::new(new_table_flownode_set_cache(
3163                String::new(),
3164                Cache::new(10),
3165                kv_backend,
3166            ));
3167            if case == "flow" {
3168                inserter
3169                    .table_flownode_set_cache
3170                    .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
3171                        flow_id: 1,
3172                        source_table_ids: vec![1],
3173                        partition_to_peer_mapping: vec![(0, Peer::empty(1))],
3174                    })])
3175                    .await
3176                    .unwrap();
3177            }
3178            assert_eq!(
3179                inserter
3180                    .can_batch_metric_rows(&requests, &Arc::new(ctx), "physical")
3181                    .await
3182                    .unwrap(),
3183                case == "eligible",
3184                "{case}"
3185            );
3186        }
3187    }
3188
3189    #[tokio::test]
3190    async fn test_batcher_meter_preserves_request() {
3191        let mut requests = RowInsertRequests {
3192            inserts: vec![RowInsertRequest {
3193                table_name: "sample".to_string(),
3194                rows: Some(Rows {
3195                    schema: vec![],
3196                    rows: vec![api::v1::Row {
3197                        values: vec![Value {
3198                            value_data: Some(api::v1::value::ValueData::F64Value(1.5)),
3199                        }],
3200                    }],
3201                }),
3202            }],
3203        };
3204        let expected = requests.clone();
3205        Inserter::meter_row_inserts(&mut requests, &QueryContext::arc())
3206            .await
3207            .unwrap();
3208        assert_eq!(requests, expected);
3209    }
3210
3211    #[tokio::test]
3212    async fn test_instant_table_bypasses_batcher() {
3213        let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3214        let inserter = batcher_test_inserter()
3215            .await
3216            .with_pending_rows_batcher(Some(batcher));
3217        let mut ctx = session::context::QueryContextBuilder::default().build();
3218        ctx.set_batching_enabled(true);
3219        let ctx = Arc::new(ctx);
3220        let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3221            .table_info();
3222        assert!(inserter.table_batcher(&table, &ctx).is_some());
3223        let mut instant = (*table).clone();
3224        instant.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
3225        assert!(inserter.table_batcher(&Arc::new(instant), &ctx).is_none());
3226    }
3227
3228    #[tokio::test]
3229    async fn test_empty_prepared_rows_skip_batcher() {
3230        let inserter = batcher_test_inserter().await;
3231        let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3232            .table_info();
3233        let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3234        let ctx = QueryContext::arc();
3235        let output = inserter
3236            .submit_table_rows(
3237                Rows {
3238                    schema: vec![],
3239                    rows: vec![],
3240                },
3241                table.clone(),
3242                ctx.clone(),
3243                &batcher,
3244            )
3245            .await
3246            .unwrap();
3247        assert!(matches!(output.data, OutputData::AffectedRows(0)));
3248        let output = inserter
3249            .submit_pending_rows(
3250                RowInsertRequests {
3251                    inserts: vec![RowInsertRequest {
3252                        table_name: table.name.clone(),
3253                        rows: None,
3254                    }],
3255                },
3256                HashMap::from_iter([(table.table_id(), table)]),
3257                ctx,
3258                &batcher,
3259            )
3260            .await
3261            .unwrap();
3262        assert!(matches!(output.data, OutputData::AffectedRows(0)));
3263    }
3264}