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