1use std::collections::HashMap;
16
17use api::helper::ColumnDataTypeWrapper;
18use api::v1::{
19 ColumnSchema, PrimaryKeyEncoding as PrimaryKeyEncodingProto, Row, Rows, SemanticType, Value,
20 WriteHint,
21};
22use common_telemetry::{error, info};
23use fxhash::FxHashMap;
24use snafu::{OptionExt, ResultExt, ensure};
25use store_api::codec::PrimaryKeyEncoding;
26use store_api::metadata::ColumnMetadata;
27use store_api::region_request::{
28 AffectedRows, RegionDeleteRequest, RegionPutRequest, RegionRequest,
29};
30use store_api::storage::{RegionId, TableId};
31
32use crate::engine::MetricEngineInner;
33use crate::error::{
34 ColumnNotFoundSnafu, CreateDefaultSnafu, ForbiddenPhysicalWriteSnafu, InvalidRequestSnafu,
35 LogicalRegionNotFoundSnafu, PhysicalRegionNotFoundSnafu, Result, UnexpectedRequestSnafu,
36 UnsupportedRegionRequestSnafu,
37};
38use crate::metrics::{FORBIDDEN_OPERATION_COUNT, MITO_OPERATION_ELAPSED};
39use crate::row_modifier::{RowsIter, TableIdInput};
40use crate::utils::to_data_region_id;
41
42impl MetricEngineInner {
43 pub async fn put_region(
45 &self,
46 region_id: RegionId,
47 request: RegionPutRequest,
48 ) -> Result<AffectedRows> {
49 let is_putting_physical_region =
50 self.state.read().unwrap().exist_physical_region(region_id);
51
52 if is_putting_physical_region {
53 info!(
54 "Metric region received put request {request:?} on physical region {region_id:?}"
55 );
56 FORBIDDEN_OPERATION_COUNT.inc();
57
58 ForbiddenPhysicalWriteSnafu.fail()
59 } else {
60 self.put_logical_region(region_id, request).await
61 }
62 }
63
64 pub async fn put_regions_batch(
74 &self,
75 requests: impl ExactSizeIterator<Item = (RegionId, RegionPutRequest)>,
76 ) -> Result<AffectedRows> {
77 let len = requests.len();
78
79 if len == 0 {
80 return Ok(0);
81 }
82
83 let _timer = MITO_OPERATION_ELAPSED
84 .with_label_values(&["put_batch"])
85 .start_timer();
86
87 if len == 1 {
89 let (region_id, req) = requests.into_iter().next().unwrap();
90 let is_putting_physical_region =
91 self.state.read().unwrap().exist_physical_region(region_id);
92 if is_putting_physical_region {
93 FORBIDDEN_OPERATION_COUNT.inc();
94 return ForbiddenPhysicalWriteSnafu.fail();
95 }
96
97 return self.put_logical_region(region_id, req).await;
98 }
99
100 let mut requests_per_physical: HashMap<RegionId, Vec<(RegionId, RegionPutRequest)>> =
101 HashMap::new();
102 for (region_id, request) in requests {
103 let is_putting_physical_region =
104 self.state.read().unwrap().exist_physical_region(region_id);
105 if is_putting_physical_region {
106 FORBIDDEN_OPERATION_COUNT.inc();
107 return ForbiddenPhysicalWriteSnafu.fail();
108 }
109 let physical_region_id = self.find_physical_region_id(region_id)?;
110 requests_per_physical
111 .entry(physical_region_id)
112 .or_default()
113 .push((region_id, request));
114 }
115
116 let mut total_affected_rows: AffectedRows = 0;
117 for (physical_region_id, requests) in requests_per_physical {
118 let affected_rows = self
119 .put_regions_batch_single_physical(physical_region_id, requests)
120 .await?;
121 total_affected_rows += affected_rows;
122 }
123
124 Ok(total_affected_rows)
125 }
126
127 async fn put_regions_batch_single_physical(
134 &self,
135 physical_region_id: RegionId,
136 mut requests: Vec<(RegionId, RegionPutRequest)>,
137 ) -> Result<AffectedRows> {
138 if requests.is_empty() {
139 return Ok(0);
140 }
141
142 let data_region_id = to_data_region_id(physical_region_id);
143 let primary_key_encoding = self.get_primary_key_encoding(data_region_id)?;
144
145 self.validate_batch_requests(physical_region_id, &mut requests)
147 .await?;
148
149 let (merged_request, total_affected_rows) = match primary_key_encoding {
151 PrimaryKeyEncoding::Sparse => self.merge_sparse_batch(physical_region_id, requests)?,
152 PrimaryKeyEncoding::Dense => self.merge_dense_batch(data_region_id, requests)?,
153 };
154
155 self.data_region
157 .write_data(data_region_id, RegionRequest::Put(merged_request))
158 .await?;
159
160 Ok(total_affected_rows)
161 }
162
163 fn get_primary_key_encoding(&self, data_region_id: RegionId) -> Result<PrimaryKeyEncoding> {
165 let state = self.state.read().unwrap();
166 state
167 .get_primary_key_encoding(data_region_id)
168 .context(PhysicalRegionNotFoundSnafu {
169 region_id: data_region_id,
170 })
171 }
172
173 async fn validate_batch_requests(
175 &self,
176 physical_region_id: RegionId,
177 requests: &mut [(RegionId, RegionPutRequest)],
178 ) -> Result<()> {
179 for (logical_region_id, request) in requests {
180 self.verify_rows(
181 *logical_region_id,
182 physical_region_id,
183 &mut request.rows,
184 true,
185 )
186 .await?;
187 }
188 Ok(())
189 }
190
191 fn merge_sparse_batch(
193 &self,
194 physical_region_id: RegionId,
195 requests: Vec<(RegionId, RegionPutRequest)>,
196 ) -> Result<(RegionPutRequest, AffectedRows)> {
197 let total_rows: usize = requests.iter().map(|(_, req)| req.rows.rows.len()).sum();
198 let mut modified_requests = Vec::with_capacity(requests.len());
199 let mut total_affected_rows: AffectedRows = 0;
200 let mut merged_version: Option<u64> = None;
201
202 for (logical_region_id, mut request) in requests {
203 if let Some(request_version) = request.partition_expr_version {
204 if let Some(merged_version) = merged_version {
205 ensure!(
206 merged_version == request_version,
207 InvalidRequestSnafu {
208 region_id: physical_region_id,
209 reason: "inconsistent partition expr version in batch"
210 }
211 );
212 } else {
213 merged_version = Some(request_version);
214 }
215 }
216 self.modify_rows(
217 physical_region_id,
218 logical_region_id.table_id(),
219 &mut request.rows,
220 PrimaryKeyEncoding::Sparse,
221 )?;
222
223 let row_count = request.rows.rows.len();
224 total_affected_rows += row_count as AffectedRows;
225 modified_requests.push(request.rows);
226 }
227
228 let schema =
229 Self::build_union_schema(modified_requests.iter().map(|rows| rows.schema.as_slice()));
230 let mut merged_rows = Vec::with_capacity(total_rows);
231 for rows in modified_requests {
232 merged_rows.extend(Self::align_rows_to_schema(rows, &schema));
233 }
234
235 let merged_request = RegionPutRequest {
236 rows: Rows {
237 schema,
238 rows: merged_rows,
239 },
240 hint: Some(WriteHint {
241 primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(),
242 }),
243 partition_expr_version: merged_version,
244 };
245
246 Ok((merged_request, total_affected_rows))
247 }
248
249 fn merge_dense_batch(
255 &self,
256 data_region_id: RegionId,
257 requests: Vec<(RegionId, RegionPutRequest)>,
258 ) -> Result<(RegionPutRequest, AffectedRows)> {
259 let merged_schema =
261 Self::build_union_schema(requests.iter().map(|(_, req)| req.rows.schema.as_slice()));
262
263 let (merged_rows, table_ids, merged_version) =
265 Self::align_requests_to_schema(requests, &merged_schema)?;
266
267 let final_rows = {
269 let state = self.state.read().unwrap();
270 let physical_columns = state
271 .physical_region_states()
272 .get(&data_region_id)
273 .with_context(|| PhysicalRegionNotFoundSnafu {
274 region_id: data_region_id,
275 })?
276 .physical_columns();
277
278 let iter = RowsIter::new(
279 Rows {
280 schema: merged_schema,
281 rows: merged_rows,
282 },
283 physical_columns,
284 );
285
286 self.row_modifier.modify_rows(
287 iter,
288 TableIdInput::Batch(&table_ids),
289 PrimaryKeyEncoding::Dense,
290 )?
291 };
292
293 let merged_request = RegionPutRequest {
294 rows: final_rows,
295 hint: None,
296 partition_expr_version: merged_version,
297 };
298
299 Ok((merged_request, table_ids.len() as AffectedRows))
300 }
301
302 fn build_union_schema<'a>(
303 schemas: impl IntoIterator<Item = &'a [ColumnSchema]>,
304 ) -> Vec<ColumnSchema> {
305 let mut schema = Vec::new();
306 for columns in schemas {
307 for col in columns {
308 if !schema
309 .iter()
310 .any(|existing: &ColumnSchema| existing.column_name == col.column_name)
311 {
312 schema.push(col.clone());
313 }
314 }
315 }
316 schema
317 }
318
319 fn align_requests_to_schema(
320 requests: Vec<(RegionId, RegionPutRequest)>,
321 merged_schema: &[ColumnSchema],
322 ) -> Result<(Vec<Row>, Vec<TableId>, Option<u64>)> {
323 let total_rows: usize = requests.iter().map(|(_, req)| req.rows.rows.len()).sum();
325 let mut merged_rows = Vec::with_capacity(total_rows);
326 let mut table_ids = Vec::with_capacity(total_rows);
327 let mut merged_version: Option<u64> = None;
328
329 for (logical_region_id, request) in requests {
330 if let Some(request_version) = request.partition_expr_version {
331 if let Some(merged_version) = merged_version {
332 ensure!(
333 merged_version == request_version,
334 InvalidRequestSnafu {
335 region_id: logical_region_id,
336 reason: "inconsistent partition expr version in batch"
337 }
338 );
339 } else {
340 merged_version = Some(request_version);
341 }
342 }
343 let table_id = logical_region_id.table_id();
344 let row_count = request.rows.rows.len();
345 merged_rows.extend(Self::align_rows_to_schema(request.rows, merged_schema));
346 table_ids.extend(std::iter::repeat_n(table_id, row_count));
347 }
348
349 Ok((merged_rows, table_ids, merged_version))
350 }
351
352 fn align_rows_to_schema(rows: Rows, merged_schema: &[ColumnSchema]) -> Vec<Row> {
353 let Rows { schema, rows } = rows;
354 if schema.len() == merged_schema.len()
355 && schema
356 .iter()
357 .zip(merged_schema)
358 .all(|(left, right)| left.column_name == right.column_name)
359 {
360 return rows;
361 }
362
363 let col_name_to_idx: FxHashMap<&str, usize> = schema
364 .iter()
365 .enumerate()
366 .map(|(idx, col)| (col.column_name.as_str(), idx))
367 .collect();
368 let col_mapping: Vec<Option<usize>> = merged_schema
369 .iter()
370 .map(|merged_col| {
371 col_name_to_idx
372 .get(merged_col.column_name.as_str())
373 .copied()
374 })
375 .collect();
376 let null_value = Value { value_data: None };
377
378 rows.into_iter()
379 .map(|mut row| {
380 let values = col_mapping
381 .iter()
382 .map(|opt_idx| match opt_idx {
383 Some(idx) => std::mem::take(&mut row.values[*idx]),
384 None => null_value.clone(),
385 })
386 .collect();
387 Row { values }
388 })
389 .collect()
390 }
391
392 fn find_physical_region_id(&self, logical_region_id: RegionId) -> Result<RegionId> {
394 let state = self.state.read().unwrap();
395 state
396 .logical_regions()
397 .get(&logical_region_id)
398 .copied()
399 .context(LogicalRegionNotFoundSnafu {
400 region_id: logical_region_id,
401 })
402 }
403
404 pub async fn delete_region(
406 &self,
407 region_id: RegionId,
408 request: RegionDeleteRequest,
409 ) -> Result<AffectedRows> {
410 if self.is_physical_region(region_id) {
411 info!(
412 "Metric region received delete request {request:?} on physical region {region_id:?}"
413 );
414 FORBIDDEN_OPERATION_COUNT.inc();
415
416 UnsupportedRegionRequestSnafu {
417 request: RegionRequest::Delete(request),
418 }
419 .fail()
420 } else {
421 self.delete_logical_region(region_id, request).await
422 }
423 }
424
425 async fn put_logical_region(
426 &self,
427 logical_region_id: RegionId,
428 mut request: RegionPutRequest,
429 ) -> Result<AffectedRows> {
430 let _timer = MITO_OPERATION_ELAPSED
431 .with_label_values(&["put"])
432 .start_timer();
433
434 let (physical_region_id, data_region_id, primary_key_encoding) =
435 self.find_data_region_meta(logical_region_id)?;
436
437 self.verify_rows(
438 logical_region_id,
439 physical_region_id,
440 &mut request.rows,
441 true,
442 )
443 .await?;
444
445 self.modify_rows(
448 physical_region_id,
449 logical_region_id.table_id(),
450 &mut request.rows,
451 primary_key_encoding,
452 )?;
453 if primary_key_encoding == PrimaryKeyEncoding::Sparse {
454 request.hint = Some(WriteHint {
455 primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(),
456 });
457 }
458 self.data_region
459 .write_data(data_region_id, RegionRequest::Put(request))
460 .await
461 }
462
463 async fn delete_logical_region(
464 &self,
465 logical_region_id: RegionId,
466 mut request: RegionDeleteRequest,
467 ) -> Result<AffectedRows> {
468 let _timer = MITO_OPERATION_ELAPSED
469 .with_label_values(&["delete"])
470 .start_timer();
471
472 let (physical_region_id, data_region_id, primary_key_encoding) =
473 self.find_data_region_meta(logical_region_id)?;
474
475 self.verify_rows(
476 logical_region_id,
477 physical_region_id,
478 &mut request.rows,
479 false,
480 )
481 .await?;
482
483 self.modify_rows(
486 physical_region_id,
487 logical_region_id.table_id(),
488 &mut request.rows,
489 primary_key_encoding,
490 )?;
491 if primary_key_encoding == PrimaryKeyEncoding::Sparse {
492 request.hint = Some(WriteHint {
493 primary_key_encoding: PrimaryKeyEncodingProto::Sparse.into(),
494 });
495 }
496 self.data_region
497 .write_data(data_region_id, RegionRequest::Delete(request))
498 .await
499 }
500
501 pub(crate) fn find_data_region_meta(
502 &self,
503 logical_region_id: RegionId,
504 ) -> Result<(RegionId, RegionId, PrimaryKeyEncoding)> {
505 let state = self.state.read().unwrap();
506 let physical_region_id = *state
507 .logical_regions()
508 .get(&logical_region_id)
509 .with_context(|| LogicalRegionNotFoundSnafu {
510 region_id: logical_region_id,
511 })?;
512 let data_region_id = to_data_region_id(physical_region_id);
513 let primary_key_encoding = state.get_primary_key_encoding(data_region_id).context(
514 PhysicalRegionNotFoundSnafu {
515 region_id: data_region_id,
516 },
517 )?;
518 Ok((physical_region_id, data_region_id, primary_key_encoding))
519 }
520
521 async fn verify_rows(
532 &self,
533 logical_region_id: RegionId,
534 physical_region_id: RegionId,
535 rows: &mut Rows,
536 check_fields: bool,
537 ) -> Result<()> {
538 let data_region_id = to_data_region_id(physical_region_id);
540 let (physical_columns, ts_name) = {
541 let state = self.state.read().unwrap();
542 if !state.is_logical_region_exist(logical_region_id) {
543 error!("Trying to write to an nonexistent region {logical_region_id}");
544 return LogicalRegionNotFoundSnafu {
545 region_id: logical_region_id,
546 }
547 .fail();
548 }
549
550 let physical_state = state
551 .physical_region_states()
552 .get(&data_region_id)
553 .context(PhysicalRegionNotFoundSnafu {
554 region_id: data_region_id,
555 })?;
556 (
557 physical_state.physical_columns().clone(),
558 physical_state.time_index_column_name().to_string(),
559 )
560 };
561
562 for col in &rows.schema {
564 let info = physical_columns
565 .get(&col.column_name)
566 .context(ColumnNotFoundSnafu {
567 name: &col.column_name,
568 region_id: logical_region_id,
569 })?;
570
571 ensure!(
572 api::helper::is_column_type_value_eq(
573 col.datatype,
574 col.datatype_extension.clone(),
575 &info.column_schema.data_type
576 ),
577 InvalidRequestSnafu {
578 region_id: logical_region_id,
579 reason: format!(
580 "column {} expect type {:?}, given: {}({})",
581 col.column_name,
582 info.column_schema.data_type,
583 api::v1::ColumnDataType::try_from(col.datatype)
584 .map(|v| v.as_str_name())
585 .unwrap_or("Unknown"),
586 col.datatype,
587 ),
588 }
589 );
590
591 ensure!(
592 api::helper::is_semantic_type_eq(col.semantic_type, info.semantic_type),
593 InvalidRequestSnafu {
594 region_id: logical_region_id,
595 reason: format!(
596 "column {} expect semantic type {:?}, given: {}({})",
597 col.column_name,
598 info.semantic_type,
599 api::v1::SemanticType::try_from(col.semantic_type)
600 .map(|v| v.as_str_name())
601 .unwrap_or("Unknown"),
602 col.semantic_type,
603 ),
604 }
605 );
606 }
607
608 ensure!(
609 rows.schema.iter().any(|col| col.column_name == ts_name),
610 InvalidRequestSnafu {
611 region_id: logical_region_id,
612 reason: format!("missing required time index column {ts_name}"),
613 }
614 );
615
616 let logical_columns = self
617 .load_logical_columns(physical_region_id, logical_region_id)
618 .await?;
619 let logical_fields = logical_columns
620 .iter()
621 .filter(|col| col.semantic_type == SemanticType::Field)
622 .map(|col| (col.column_schema.name.as_str(), col))
623 .collect::<HashMap<_, _>>();
624
625 for col in &rows.schema {
626 if api::helper::is_semantic_type_eq(col.semantic_type, SemanticType::Field) {
627 ensure!(
628 logical_fields.contains_key(col.column_name.as_str()),
629 InvalidRequestSnafu {
630 region_id: logical_region_id,
631 reason: format!(
632 "field column {} does not belong to logical region {logical_region_id}",
633 col.column_name,
634 ),
635 }
636 );
637 }
638 }
639
640 if check_fields {
641 for (field_name, field_meta) in logical_fields {
644 if !rows.schema.iter().any(|col| col.column_name == field_name) {
645 Self::fill_missing_field_column(
646 logical_region_id,
647 field_name,
648 field_meta,
649 rows,
650 )?;
651 }
652 }
653
654 for (field_name, field_meta) in physical_columns
655 .iter()
656 .filter(|(_, col)| col.semantic_type == SemanticType::Field)
657 {
658 if !rows.schema.iter().any(|col| col.column_name == *field_name) {
659 Self::fill_missing_field_column(
660 logical_region_id,
661 field_name,
662 field_meta,
663 rows,
664 )?;
665 }
666 }
667 }
668
669 Ok(())
670 }
671
672 fn fill_missing_field_column(
673 logical_region_id: RegionId,
674 field_name: &str,
675 field_meta: &ColumnMetadata,
676 rows: &mut Rows,
677 ) -> Result<()> {
678 ensure!(
681 !field_meta.column_schema.is_default_impure(),
682 UnexpectedRequestSnafu {
683 reason: format!(
684 "unexpected impure default value with region_id: {logical_region_id}, column: {field_name}, default_value: {:?}",
685 field_meta.column_schema.default_constraint(),
686 ),
687 }
688 );
689
690 let default_value = field_meta
691 .column_schema
692 .create_default()
693 .context(CreateDefaultSnafu {
694 region_id: logical_region_id,
695 column: field_name,
696 })?
697 .with_context(|| InvalidRequestSnafu {
698 region_id: logical_region_id,
699 reason: format!("missing required field column {field_name}"),
700 })?;
701 let default_value = api::helper::to_grpc_value(default_value);
702 let (datatype, datatype_extension) =
703 ColumnDataTypeWrapper::try_from(field_meta.column_schema.data_type.clone())
704 .map_err(|e| {
705 InvalidRequestSnafu {
706 region_id: logical_region_id,
707 reason: format!(
708 "no protobuf type for field column {field_name} ({:?}): {e}",
709 field_meta.column_schema.data_type
710 ),
711 }
712 .build()
713 })?
714 .to_parts();
715
716 rows.schema.push(ColumnSchema {
717 column_name: field_name.to_string(),
718 datatype: datatype as i32,
719 semantic_type: SemanticType::Field as i32,
720 datatype_extension,
721 options: None,
722 });
723
724 for row in &mut rows.rows {
725 row.values.push(default_value.clone());
726 }
727
728 Ok(())
729 }
730
731 fn modify_rows(
735 &self,
736 physical_region_id: RegionId,
737 table_id: TableId,
738 rows: &mut Rows,
739 encoding: PrimaryKeyEncoding,
740 ) -> Result<()> {
741 let input = std::mem::take(rows);
742 let iter = {
743 let state = self.state.read().unwrap();
744 let physical_columns = state
745 .physical_region_states()
746 .get(&physical_region_id)
747 .with_context(|| PhysicalRegionNotFoundSnafu {
748 region_id: physical_region_id,
749 })?
750 .physical_columns();
751 RowsIter::new(input, physical_columns)
752 };
753 let output =
754 self.row_modifier
755 .modify_rows(iter, TableIdInput::Single(table_id), encoding)?;
756 *rows = output;
757 Ok(())
758 }
759}
760
761#[cfg(test)]
762mod tests {
763 use std::collections::HashSet;
764
765 use api::v1::value::ValueData;
766 use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema};
767 use common_error::ext::ErrorExt;
768 use common_error::status_code::StatusCode;
769 use common_function::utils::partition_expr_version;
770 use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value};
771 use common_recordbatch::RecordBatches;
772 use datatypes::prelude::ConcreteDataType;
773 use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema};
774 use datatypes::value::Value as PartitionValue;
775 use partition::expr::col;
776 use store_api::metadata::ColumnMetadata;
777 use store_api::metric_engine_consts::{
778 DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, PRIMARY_KEY_ENCODING,
779 };
780 use store_api::path_utils::table_dir;
781 use store_api::region_engine::RegionEngine;
782 use store_api::region_request::{
783 EnterStagingRequest, RegionRequest, StagingPartitionDirective,
784 };
785 use store_api::storage::ScanRequest;
786 use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
787
788 use super::*;
789 use crate::test_util::{self, TestEnv};
790
791 fn assert_merged_schema(rows: &Rows, expect_sparse: bool) {
792 let column_names: HashSet<String> = rows
793 .schema
794 .iter()
795 .map(|col| col.column_name.clone())
796 .collect();
797
798 if expect_sparse {
799 assert!(
800 column_names.contains(PRIMARY_KEY_COLUMN_NAME),
801 "sparse encoding should include primary key column"
802 );
803 assert!(
804 !column_names.contains(DATA_SCHEMA_TABLE_ID_COLUMN_NAME),
805 "sparse encoding should not include table id column"
806 );
807 assert!(
808 !column_names.contains(DATA_SCHEMA_TSID_COLUMN_NAME),
809 "sparse encoding should not include tsid column"
810 );
811 assert!(
812 !column_names.contains("job"),
813 "sparse encoding should not include tag columns"
814 );
815 assert!(
816 !column_names.contains("instance"),
817 "sparse encoding should not include tag columns"
818 );
819 } else {
820 assert!(
821 !column_names.contains(PRIMARY_KEY_COLUMN_NAME),
822 "dense encoding should not include primary key column"
823 );
824 assert!(
825 column_names.contains(DATA_SCHEMA_TABLE_ID_COLUMN_NAME),
826 "dense encoding should include table id column"
827 );
828 assert!(
829 column_names.contains(DATA_SCHEMA_TSID_COLUMN_NAME),
830 "dense encoding should include tsid column"
831 );
832 assert!(
833 column_names.contains("job"),
834 "dense encoding should keep tag columns"
835 );
836 assert!(
837 column_names.contains("instance"),
838 "dense encoding should keep tag columns"
839 );
840 }
841 }
842
843 fn job_partition_expr_json() -> String {
844 let expr = col("job")
845 .gt_eq(PartitionValue::String("job-0".into()))
846 .and(col("job").lt(PartitionValue::String("job-9".into())));
847 expr.as_json_str().unwrap()
848 }
849
850 async fn create_logical_region_with_tags(
851 env: &TestEnv,
852 physical_region_id: RegionId,
853 logical_region_id: RegionId,
854 tags: &[&str],
855 ) {
856 let region_create_request = test_util::create_logical_region_request(
857 tags,
858 physical_region_id,
859 &table_dir("test", logical_region_id.table_id()),
860 );
861 env.metric()
862 .handle_request(
863 logical_region_id,
864 RegionRequest::Create(region_create_request),
865 )
866 .await
867 .unwrap();
868 }
869
870 fn column_index(rows: &Rows, name: &str) -> usize {
871 rows.schema
872 .iter()
873 .position(|col| col.column_name == name)
874 .unwrap()
875 }
876
877 async fn run_batch_write_with_schema_variants(
878 env: &TestEnv,
879 physical_region_id: RegionId,
880 options: Vec<(String, String)>,
881 expect_sparse: bool,
882 ) {
883 env.create_physical_region(physical_region_id, &TestEnv::default_table_dir(), options)
884 .await;
885
886 let logical_region_1 = env.default_logical_region_id();
887 let logical_region_2 = RegionId::new(1024, 1);
888
889 create_logical_region_with_tags(env, physical_region_id, logical_region_1, &["job"]).await;
890 create_logical_region_with_tags(
891 env,
892 physical_region_id,
893 logical_region_2,
894 &["job", "instance"],
895 )
896 .await;
897
898 let schema_1 = test_util::row_schema_with_tags(&["job"]);
899 let schema_2 = test_util::row_schema_with_tags(&["job", "instance"]);
900
901 let data_region_id = RegionId::new(physical_region_id.table_id(), 2);
902 let primary_key_encoding = env
903 .metric()
904 .inner
905 .get_primary_key_encoding(data_region_id)
906 .unwrap();
907 assert_eq!(
908 primary_key_encoding,
909 if expect_sparse {
910 PrimaryKeyEncoding::Sparse
911 } else {
912 PrimaryKeyEncoding::Dense
913 }
914 );
915
916 let build_requests = || {
917 let rows_1 = test_util::build_rows(1, 3);
918 let rows_2 = test_util::build_rows(2, 2);
919
920 vec![
921 (
922 logical_region_1,
923 RegionPutRequest {
924 rows: Rows {
925 schema: schema_1.clone(),
926 rows: rows_1,
927 },
928 hint: None,
929 partition_expr_version: None,
930 },
931 ),
932 (
933 logical_region_2,
934 RegionPutRequest {
935 rows: Rows {
936 schema: schema_2.clone(),
937 rows: rows_2,
938 },
939 hint: None,
940 partition_expr_version: None,
941 },
942 ),
943 ]
944 };
945
946 let merged_request = if expect_sparse {
947 let (merged_request, _) = env
948 .metric()
949 .inner
950 .merge_sparse_batch(physical_region_id, build_requests())
951 .unwrap();
952 let hint = merged_request
953 .hint
954 .as_ref()
955 .expect("missing sparse write hint");
956 assert_eq!(
957 hint.primary_key_encoding,
958 PrimaryKeyEncodingProto::Sparse as i32
959 );
960 merged_request
961 } else {
962 let (merged_request, _) = env
963 .metric()
964 .inner
965 .merge_dense_batch(data_region_id, build_requests())
966 .unwrap();
967 assert!(merged_request.hint.is_none());
968 merged_request
969 };
970
971 assert_merged_schema(&merged_request.rows, expect_sparse);
972
973 let affected_rows = env
974 .metric()
975 .inner
976 .put_regions_batch(build_requests().into_iter())
977 .await
978 .unwrap();
979 assert_eq!(affected_rows, 5);
980
981 let request = ScanRequest::default();
982 let stream = env
983 .mito()
984 .scan_to_stream(data_region_id, request)
985 .await
986 .unwrap();
987 let batches = RecordBatches::try_collect(stream).await.unwrap();
988
989 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 5);
990 }
991
992 #[test]
993 fn test_sparse_batch_aligns_mixed_field_order() {
994 let primary_key = PbColumnSchema {
995 column_name: PRIMARY_KEY_COLUMN_NAME.to_string(),
996 datatype: ColumnDataType::Binary as i32,
997 semantic_type: SemanticType::Tag as _,
998 datatype_extension: None,
999 options: None,
1000 };
1001 let timestamp = PbColumnSchema {
1002 column_name: greptime_timestamp().to_string(),
1003 datatype: ColumnDataType::TimestampMillisecond as i32,
1004 semantic_type: SemanticType::Timestamp as _,
1005 datatype_extension: None,
1006 options: None,
1007 };
1008 let value = PbColumnSchema {
1009 column_name: greptime_value().to_string(),
1010 datatype: ColumnDataType::Float64 as i32,
1011 semantic_type: SemanticType::Field as _,
1012 datatype_extension: None,
1013 options: None,
1014 };
1015 let histogram = PbColumnSchema {
1016 column_name: greptime_native_histogram().to_string(),
1017 datatype: ColumnDataType::Struct as i32,
1018 semantic_type: SemanticType::Field as _,
1019 datatype_extension: None,
1020 options: None,
1021 };
1022
1023 let sample_rows = Rows {
1024 schema: vec![
1025 primary_key.clone(),
1026 timestamp.clone(),
1027 value.clone(),
1028 histogram.clone(),
1029 ],
1030 rows: vec![Row {
1031 values: vec![
1032 ValueData::BinaryValue(vec![1]).into(),
1033 ValueData::TimestampMillisecondValue(0).into(),
1034 ValueData::F64Value(1.0).into(),
1035 Value { value_data: None },
1036 ],
1037 }],
1038 };
1039 let histogram_rows = Rows {
1040 schema: vec![primary_key, timestamp, histogram, value],
1041 rows: vec![Row {
1042 values: vec![
1043 ValueData::BinaryValue(vec![2]).into(),
1044 ValueData::TimestampMillisecondValue(0).into(),
1045 ValueData::StructValue(api::v1::StructValue { items: vec![] }).into(),
1046 Value { value_data: None },
1047 ],
1048 }],
1049 };
1050
1051 let schema = MetricEngineInner::build_union_schema([
1052 sample_rows.schema.as_slice(),
1053 histogram_rows.schema.as_slice(),
1054 ]);
1055 let merged_rows = MetricEngineInner::align_rows_to_schema(sample_rows, &schema)
1056 .into_iter()
1057 .chain(MetricEngineInner::align_rows_to_schema(
1058 histogram_rows,
1059 &schema,
1060 ))
1061 .collect();
1062 let merged_request = Rows {
1063 schema,
1064 rows: merged_rows,
1065 };
1066
1067 let value_idx = column_index(&merged_request, greptime_value());
1068 let histogram_idx = column_index(&merged_request, greptime_native_histogram());
1069 assert!(matches!(
1070 merged_request.rows[0].values[value_idx].value_data,
1071 Some(ValueData::F64Value(_))
1072 ));
1073 assert!(
1074 merged_request.rows[0].values[histogram_idx]
1075 .value_data
1076 .is_none()
1077 );
1078 assert!(
1079 merged_request.rows[1].values[value_idx]
1080 .value_data
1081 .is_none()
1082 );
1083 assert!(matches!(
1084 merged_request.rows[1].values[histogram_idx].value_data,
1085 Some(ValueData::StructValue(_))
1086 ));
1087 }
1088
1089 #[tokio::test]
1090 async fn test_write_logical_region() {
1091 let env = TestEnv::new().await;
1092 env.init_metric_region().await;
1093
1094 let schema = test_util::row_schema_with_tags(&["job"]);
1096 let rows = test_util::build_rows(1, 5);
1097 let request = RegionRequest::Put(RegionPutRequest {
1098 rows: Rows { schema, rows },
1099 hint: None,
1100 partition_expr_version: None,
1101 });
1102
1103 let logical_region_id = env.default_logical_region_id();
1105 let result = env
1106 .metric()
1107 .handle_request(logical_region_id, request)
1108 .await
1109 .unwrap();
1110 assert_eq!(result.affected_rows, 5);
1111
1112 let physical_region_id = env.default_physical_region_id();
1114 let request = ScanRequest::default();
1115 let stream = env
1116 .metric()
1117 .scan_to_stream(physical_region_id, request)
1118 .await
1119 .unwrap();
1120 let batches = RecordBatches::try_collect(stream).await.unwrap();
1121 let expected = "\
1122+-------------------------+----------------+------------+---------------------+-------+
1123| greptime_timestamp | greptime_value | __table_id | __tsid | job |
1124+-------------------------+----------------+------------+---------------------+-------+
1125| 1970-01-01T00:00:00 | 0.0 | 3 | 2955007454552897459 | tag_0 |
1126| 1970-01-01T00:00:00.001 | 1.0 | 3 | 2955007454552897459 | tag_0 |
1127| 1970-01-01T00:00:00.002 | 2.0 | 3 | 2955007454552897459 | tag_0 |
1128| 1970-01-01T00:00:00.003 | 3.0 | 3 | 2955007454552897459 | tag_0 |
1129| 1970-01-01T00:00:00.004 | 4.0 | 3 | 2955007454552897459 | tag_0 |
1130+-------------------------+----------------+------------+---------------------+-------+";
1131 assert_eq!(expected, batches.pretty_print().unwrap(), "physical region");
1132
1133 let request = ScanRequest::default();
1135 let stream = env
1136 .metric()
1137 .scan_to_stream(logical_region_id, request)
1138 .await
1139 .unwrap();
1140 let batches = RecordBatches::try_collect(stream).await.unwrap();
1141 let expected = "\
1142+-------------------------+----------------+-------+
1143| greptime_timestamp | greptime_value | job |
1144+-------------------------+----------------+-------+
1145| 1970-01-01T00:00:00 | 0.0 | tag_0 |
1146| 1970-01-01T00:00:00.001 | 1.0 | tag_0 |
1147| 1970-01-01T00:00:00.002 | 2.0 | tag_0 |
1148| 1970-01-01T00:00:00.003 | 3.0 | tag_0 |
1149| 1970-01-01T00:00:00.004 | 4.0 | tag_0 |
1150+-------------------------+----------------+-------+";
1151 assert_eq!(expected, batches.pretty_print().unwrap(), "logical region");
1152 }
1153
1154 #[tokio::test]
1155 async fn test_write_logical_region_row_count() {
1156 let env = TestEnv::new().await;
1157 env.init_metric_region().await;
1158 let engine = env.metric();
1159
1160 let logical_region_id = env.default_logical_region_id();
1162 let columns = &["odd", "even", "Ev_En"];
1163 let alter_request = test_util::alter_logical_region_add_tag_columns(123456, columns);
1164 engine
1165 .handle_request(logical_region_id, RegionRequest::Alter(alter_request))
1166 .await
1167 .unwrap();
1168
1169 let schema = test_util::row_schema_with_tags(columns);
1171 let rows = test_util::build_rows(3, 100);
1172 let request = RegionRequest::Put(RegionPutRequest {
1173 rows: Rows { schema, rows },
1174 hint: None,
1175 partition_expr_version: None,
1176 });
1177
1178 let result = engine
1180 .handle_request(logical_region_id, request)
1181 .await
1182 .unwrap();
1183 assert_eq!(100, result.affected_rows);
1184 }
1185
1186 #[tokio::test]
1187 async fn test_write_physical_region() {
1188 let env = TestEnv::new().await;
1189 env.init_metric_region().await;
1190 let engine = env.metric();
1191
1192 let physical_region_id = env.default_physical_region_id();
1193 let schema = test_util::row_schema_with_tags(&["abc"]);
1194 let rows = test_util::build_rows(1, 100);
1195 let request = RegionRequest::Put(RegionPutRequest {
1196 rows: Rows { schema, rows },
1197 hint: None,
1198 partition_expr_version: None,
1199 });
1200
1201 engine
1202 .handle_request(physical_region_id, request)
1203 .await
1204 .unwrap_err();
1205 }
1206
1207 #[tokio::test]
1208 async fn test_write_nonexist_logical_region() {
1209 let env = TestEnv::new().await;
1210 env.init_metric_region().await;
1211 let engine = env.metric();
1212
1213 let logical_region_id = RegionId::new(175, 8345);
1214 let schema = test_util::row_schema_with_tags(&["def"]);
1215 let rows = test_util::build_rows(1, 100);
1216 let request = RegionRequest::Put(RegionPutRequest {
1217 rows: Rows { schema, rows },
1218 hint: None,
1219 partition_expr_version: None,
1220 });
1221
1222 engine
1223 .handle_request(logical_region_id, request)
1224 .await
1225 .unwrap_err();
1226 }
1227
1228 #[tokio::test]
1229 async fn test_batch_write_multiple_logical_regions() {
1230 let env = TestEnv::new().await;
1231 env.init_metric_region().await;
1232 let engine = env.metric();
1233
1234 let physical_region_id = env.default_physical_region_id();
1236 let logical_region_1 = env.default_logical_region_id();
1237 let logical_region_2 = RegionId::new(1024, 1);
1238 let logical_region_3 = RegionId::new(1024, 2);
1239
1240 env.create_logical_region(physical_region_id, logical_region_2)
1241 .await;
1242 env.create_logical_region(physical_region_id, logical_region_3)
1243 .await;
1244
1245 let schema = test_util::row_schema_with_tags(&["job"]);
1247
1248 let rows1 = test_util::build_rows(1, 3);
1253 let mut rows2 = test_util::build_rows(1, 2);
1254 let mut rows3 = test_util::build_rows(1, 5);
1255
1256 use api::v1::value::ValueData;
1258 for (i, row) in rows2.iter_mut().enumerate() {
1259 if let Some(ValueData::TimestampMillisecondValue(ts)) =
1260 row.values.get_mut(0).and_then(|v| v.value_data.as_mut())
1261 {
1262 *ts = (10 + i) as i64;
1263 }
1264 }
1265 for (i, row) in rows3.iter_mut().enumerate() {
1266 if let Some(ValueData::TimestampMillisecondValue(ts)) =
1267 row.values.get_mut(0).and_then(|v| v.value_data.as_mut())
1268 {
1269 *ts = (20 + i) as i64;
1270 }
1271 }
1272
1273 let requests = vec![
1274 (
1275 logical_region_1,
1276 RegionPutRequest {
1277 rows: Rows {
1278 schema: schema.clone(),
1279 rows: rows1,
1280 },
1281 hint: None,
1282 partition_expr_version: None,
1283 },
1284 ),
1285 (
1286 logical_region_2,
1287 RegionPutRequest {
1288 rows: Rows {
1289 schema: schema.clone(),
1290 rows: rows2,
1291 },
1292 hint: None,
1293 partition_expr_version: None,
1294 },
1295 ),
1296 (
1297 logical_region_3,
1298 RegionPutRequest {
1299 rows: Rows {
1300 schema: schema.clone(),
1301 rows: rows3,
1302 },
1303 hint: None,
1304 partition_expr_version: None,
1305 },
1306 ),
1307 ];
1308
1309 let affected_rows = engine
1311 .inner
1312 .put_regions_batch(requests.into_iter())
1313 .await
1314 .unwrap();
1315 assert_eq!(affected_rows, 10);
1316
1317 let request = ScanRequest::default();
1319 let stream = env
1320 .metric()
1321 .scan_to_stream(physical_region_id, request)
1322 .await
1323 .unwrap();
1324 let batches = RecordBatches::try_collect(stream).await.unwrap();
1325
1326 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 10);
1328 }
1329
1330 #[tokio::test]
1331 async fn test_batch_write_with_partial_failure() {
1332 let env = TestEnv::new().await;
1333 env.init_metric_region().await;
1334 let engine = env.metric();
1335
1336 let physical_region_id = env.default_physical_region_id();
1337 let logical_region_1 = env.default_logical_region_id();
1338 let logical_region_2 = RegionId::new(1024, 1);
1339 let nonexistent_region = RegionId::new(9999, 9999);
1340
1341 env.create_logical_region(physical_region_id, logical_region_2)
1342 .await;
1343
1344 let schema = test_util::row_schema_with_tags(&["job"]);
1346 let requests = vec![
1347 (
1348 logical_region_1,
1349 RegionPutRequest {
1350 rows: Rows {
1351 schema: schema.clone(),
1352 rows: test_util::build_rows(1, 3),
1353 },
1354 hint: None,
1355 partition_expr_version: None,
1356 },
1357 ),
1358 (
1359 nonexistent_region,
1360 RegionPutRequest {
1361 rows: Rows {
1362 schema: schema.clone(),
1363 rows: test_util::build_rows(1, 2),
1364 },
1365 hint: None,
1366 partition_expr_version: None,
1367 },
1368 ),
1369 (
1370 logical_region_2,
1371 RegionPutRequest {
1372 rows: Rows {
1373 schema: schema.clone(),
1374 rows: test_util::build_rows(1, 5),
1375 },
1376 hint: None,
1377 partition_expr_version: None,
1378 },
1379 ),
1380 ];
1381
1382 let result = engine.inner.put_regions_batch(requests.into_iter()).await;
1384 assert!(result.is_err());
1385
1386 let request = ScanRequest::default();
1389 let stream = env
1390 .metric()
1391 .scan_to_stream(physical_region_id, request)
1392 .await
1393 .unwrap();
1394 let batches = RecordBatches::try_collect(stream).await.unwrap();
1395
1396 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 0);
1397 }
1398
1399 #[tokio::test]
1400 async fn test_batch_write_single_physical_region_forbidden() {
1401 let env = TestEnv::new().await;
1402 env.init_metric_region().await;
1403 let engine = env.metric();
1404
1405 let physical_region_id = env.default_physical_region_id();
1406 let schema = test_util::row_schema_with_tags(&["job"]);
1407 let requests = vec![(
1408 physical_region_id,
1409 RegionPutRequest {
1410 rows: Rows {
1411 schema,
1412 rows: test_util::build_rows(1, 1),
1413 },
1414 hint: None,
1415 partition_expr_version: None,
1416 },
1417 )];
1418
1419 let err = engine
1420 .inner
1421 .put_regions_batch(requests.into_iter())
1422 .await
1423 .unwrap_err();
1424
1425 assert!(matches!(
1426 err,
1427 crate::error::Error::ForbiddenPhysicalWrite { .. }
1428 ));
1429 }
1430
1431 #[tokio::test]
1432 async fn test_batch_write_physical_region_forbidden() {
1433 let env = TestEnv::new().await;
1434 env.init_metric_region().await;
1435 let engine = env.metric();
1436
1437 let physical_region_id = env.default_physical_region_id();
1438 let logical_region_id = env.default_logical_region_id();
1439 let schema = test_util::row_schema_with_tags(&["job"]);
1440 let requests = vec![
1441 (
1442 logical_region_id,
1443 RegionPutRequest {
1444 rows: Rows {
1445 schema: schema.clone(),
1446 rows: test_util::build_rows(1, 1),
1447 },
1448 hint: None,
1449 partition_expr_version: None,
1450 },
1451 ),
1452 (
1453 physical_region_id,
1454 RegionPutRequest {
1455 rows: Rows {
1456 schema,
1457 rows: test_util::build_rows(1, 1),
1458 },
1459 hint: None,
1460 partition_expr_version: None,
1461 },
1462 ),
1463 ];
1464
1465 let err = engine
1466 .inner
1467 .put_regions_batch(requests.into_iter())
1468 .await
1469 .unwrap_err();
1470
1471 assert!(matches!(
1472 err,
1473 crate::error::Error::ForbiddenPhysicalWrite { .. }
1474 ));
1475 }
1476
1477 #[tokio::test]
1478 async fn test_batch_write_single_request_fast_path() {
1479 let env = TestEnv::new().await;
1480 env.init_metric_region().await;
1481 let engine = env.metric();
1482
1483 let logical_region_id = env.default_logical_region_id();
1484 let schema = test_util::row_schema_with_tags(&["job"]);
1485
1486 let requests = vec![(
1488 logical_region_id,
1489 RegionPutRequest {
1490 rows: Rows {
1491 schema,
1492 rows: test_util::build_rows(1, 5),
1493 },
1494 hint: None,
1495 partition_expr_version: None,
1496 },
1497 )];
1498
1499 let affected_rows = engine
1500 .inner
1501 .put_regions_batch(requests.into_iter())
1502 .await
1503 .unwrap();
1504 assert_eq!(affected_rows, 5);
1505 }
1506
1507 #[tokio::test]
1508 async fn test_batch_write_empty_requests() {
1509 let env = TestEnv::new().await;
1510 env.init_metric_region().await;
1511 let engine = env.metric();
1512
1513 let requests = vec![];
1515 let affected_rows = engine
1516 .inner
1517 .put_regions_batch(requests.into_iter())
1518 .await
1519 .unwrap();
1520
1521 assert_eq!(affected_rows, 0);
1522 }
1523
1524 #[tokio::test]
1525 async fn test_batch_write_sparse_encoding() {
1526 let env = TestEnv::new().await;
1527 let physical_region_id = env.default_physical_region_id();
1528
1529 run_batch_write_with_schema_variants(
1530 &env,
1531 physical_region_id,
1532 vec![(PRIMARY_KEY_ENCODING.to_string(), "sparse".to_string())],
1533 true,
1534 )
1535 .await;
1536 }
1537
1538 #[tokio::test]
1539 async fn test_batch_write_dense_encoding() {
1540 let env = TestEnv::new().await;
1541 let physical_region_id = env.default_physical_region_id();
1542
1543 run_batch_write_with_schema_variants(
1544 &env,
1545 physical_region_id,
1546 vec![(PRIMARY_KEY_ENCODING.to_string(), "dense".to_string())],
1547 false,
1548 )
1549 .await;
1550 }
1551
1552 #[tokio::test]
1553 async fn test_metric_put_rejects_bad_partition_expr_version() {
1554 let env = TestEnv::new().await;
1555 env.init_metric_region().await;
1556
1557 let logical_region_id = env.default_logical_region_id();
1558 let rows = Rows {
1559 schema: test_util::row_schema_with_tags(&["job"]),
1560 rows: test_util::build_rows(1, 3),
1561 };
1562
1563 let err = env
1564 .metric()
1565 .handle_request(
1566 logical_region_id,
1567 RegionRequest::Put(RegionPutRequest {
1568 rows,
1569 hint: None,
1570 partition_expr_version: Some(1),
1571 }),
1572 )
1573 .await
1574 .unwrap_err();
1575
1576 assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1577 }
1578
1579 #[tokio::test]
1580 async fn test_metric_put_respects_staging_partition_expr_version() {
1581 let env = TestEnv::new().await;
1582 env.init_metric_region().await;
1583
1584 let logical_region_id = env.default_logical_region_id();
1585 let physical_region_id = env.default_physical_region_id();
1586 let partition_expr = job_partition_expr_json();
1587 env.metric()
1588 .handle_request(
1589 physical_region_id,
1590 RegionRequest::EnterStaging(EnterStagingRequest {
1591 partition_directive: StagingPartitionDirective::UpdatePartitionExpr(
1592 partition_expr.clone(),
1593 ),
1594 }),
1595 )
1596 .await
1597 .unwrap();
1598
1599 let expected_version = partition_expr_version(Some(&partition_expr));
1600 let rows = Rows {
1601 schema: test_util::row_schema_with_tags(&["job"]),
1602 rows: test_util::build_rows(1, 3),
1603 };
1604
1605 let err = env
1606 .metric()
1607 .handle_request(
1608 logical_region_id,
1609 RegionRequest::Put(RegionPutRequest {
1610 rows: rows.clone(),
1611 hint: None,
1612 partition_expr_version: Some(expected_version.wrapping_add(1)),
1613 }),
1614 )
1615 .await
1616 .unwrap_err();
1617 assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1618
1619 let response = env
1620 .metric()
1621 .handle_request(
1622 logical_region_id,
1623 RegionRequest::Put(RegionPutRequest {
1624 rows: rows.clone(),
1625 hint: None,
1626 partition_expr_version: None,
1627 }),
1628 )
1629 .await
1630 .unwrap();
1631 assert_eq!(response.affected_rows, 3);
1632
1633 let response = env
1634 .metric()
1635 .handle_request(
1636 logical_region_id,
1637 RegionRequest::Put(RegionPutRequest {
1638 rows,
1639 hint: None,
1640 partition_expr_version: Some(expected_version),
1641 }),
1642 )
1643 .await
1644 .unwrap();
1645 assert_eq!(response.affected_rows, 3);
1646 }
1647
1648 #[tokio::test]
1652 async fn test_verify_rows_rejects_wrong_type() {
1653 use api::v1::value::ValueData;
1654 use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema, SemanticType};
1655 use common_query::prelude::{greptime_timestamp, greptime_value};
1656
1657 let env = TestEnv::new().await;
1658 env.init_metric_region().await;
1659
1660 let logical_region_id = env.default_logical_region_id();
1661
1662 let schema = vec![
1665 PbColumnSchema {
1666 column_name: greptime_timestamp().to_string(),
1667 datatype: ColumnDataType::String as i32,
1668 semantic_type: SemanticType::Timestamp as _,
1669 datatype_extension: None,
1670 options: None,
1671 },
1672 PbColumnSchema {
1673 column_name: greptime_value().to_string(),
1674 datatype: ColumnDataType::Float64 as i32,
1675 semantic_type: SemanticType::Field as _,
1676 datatype_extension: None,
1677 options: None,
1678 },
1679 PbColumnSchema {
1680 column_name: "job".to_string(),
1681 datatype: ColumnDataType::String as i32,
1682 semantic_type: SemanticType::Tag as _,
1683 datatype_extension: None,
1684 options: None,
1685 },
1686 ];
1687 let rows = vec![Row {
1688 values: vec![
1689 Value {
1690 value_data: Some(ValueData::StringValue("not-a-timestamp".to_string())),
1691 },
1692 Value {
1693 value_data: Some(ValueData::F64Value(1.0)),
1694 },
1695 Value {
1696 value_data: Some(ValueData::StringValue("tag_0".to_string())),
1697 },
1698 ],
1699 }];
1700
1701 let err = env
1702 .metric()
1703 .handle_request(
1704 logical_region_id,
1705 RegionRequest::Put(RegionPutRequest {
1706 rows: Rows { schema, rows },
1707 hint: None,
1708 partition_expr_version: None,
1709 }),
1710 )
1711 .await
1712 .unwrap_err();
1713 assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1714 }
1715
1716 #[tokio::test]
1720 async fn test_verify_rows_rejects_missing_time_index() {
1721 use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema, SemanticType};
1722 use common_query::prelude::greptime_value;
1723
1724 let env = TestEnv::new().await;
1725 env.init_metric_region().await;
1726
1727 let logical_region_id = env.default_logical_region_id();
1728
1729 let schema = vec![
1731 PbColumnSchema {
1732 column_name: greptime_value().to_string(),
1733 datatype: ColumnDataType::Float64 as i32,
1734 semantic_type: SemanticType::Field as _,
1735 datatype_extension: None,
1736 options: None,
1737 },
1738 PbColumnSchema {
1739 column_name: "job".to_string(),
1740 datatype: ColumnDataType::String as i32,
1741 semantic_type: SemanticType::Tag as _,
1742 datatype_extension: None,
1743 options: None,
1744 },
1745 ];
1746 let rows = vec![Row {
1747 values: vec![
1748 Value {
1749 value_data: Some(api::v1::value::ValueData::F64Value(1.0)),
1750 },
1751 Value {
1752 value_data: Some(api::v1::value::ValueData::StringValue("tag_0".to_string())),
1753 },
1754 ],
1755 }];
1756
1757 let err = env
1758 .metric()
1759 .handle_request(
1760 logical_region_id,
1761 RegionRequest::Put(RegionPutRequest {
1762 rows: Rows { schema, rows },
1763 hint: None,
1764 partition_expr_version: None,
1765 }),
1766 )
1767 .await
1768 .unwrap_err();
1769 assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1770 }
1771
1772 #[tokio::test]
1773 async fn test_verify_rows_rejects_missing_field() {
1774 use api::v1::value::ValueData;
1775 use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema, SemanticType};
1776 use common_query::prelude::greptime_timestamp;
1777
1778 let env = TestEnv::new().await;
1779 env.init_metric_region().await;
1780
1781 let logical_region_id = env.default_logical_region_id();
1782
1783 let schema = vec![
1785 PbColumnSchema {
1786 column_name: greptime_timestamp().to_string(),
1787 datatype: ColumnDataType::TimestampMillisecond as i32,
1788 semantic_type: SemanticType::Timestamp as _,
1789 datatype_extension: None,
1790 options: None,
1791 },
1792 PbColumnSchema {
1793 column_name: "job".to_string(),
1794 datatype: ColumnDataType::String as i32,
1795 semantic_type: SemanticType::Tag as _,
1796 datatype_extension: None,
1797 options: None,
1798 },
1799 ];
1800 let rows = vec![Row {
1801 values: vec![
1802 Value {
1803 value_data: Some(ValueData::TimestampMillisecondValue(0)),
1804 },
1805 Value {
1806 value_data: Some(ValueData::StringValue("tag_0".to_string())),
1807 },
1808 ],
1809 }];
1810
1811 let err = env
1812 .metric()
1813 .handle_request(
1814 logical_region_id,
1815 RegionRequest::Put(RegionPutRequest {
1816 rows: Rows { schema, rows },
1817 hint: None,
1818 partition_expr_version: None,
1819 }),
1820 )
1821 .await
1822 .unwrap_err();
1823 let message = err.to_string();
1824 assert!(
1825 message.contains("missing required field column"),
1826 "expected field-completeness rejection, got: {message}"
1827 );
1828 assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1829 }
1830
1831 #[test]
1832 fn test_fill_missing_field_column_nullable_no_default() {
1833 let field_meta = ColumnMetadata {
1834 column_id: 1,
1835 semantic_type: SemanticType::Field,
1836 column_schema: ColumnSchema::new(
1837 "greptime_value".to_string(),
1838 ConcreteDataType::float64_datatype(),
1839 true, ),
1841 };
1842 let mut rows = Rows {
1843 schema: vec![PbColumnSchema {
1844 column_name: "ts".to_string(),
1845 datatype: ColumnDataType::TimestampMillisecond as i32,
1846 semantic_type: SemanticType::Timestamp as _,
1847 datatype_extension: None,
1848 options: None,
1849 }],
1850 rows: vec![Row {
1851 values: vec![Value {
1852 value_data: Some(ValueData::TimestampMillisecondValue(0)),
1853 }],
1854 }],
1855 };
1856
1857 MetricEngineInner::fill_missing_field_column(
1858 RegionId::new(1, 1),
1859 "greptime_value",
1860 &field_meta,
1861 &mut rows,
1862 )
1863 .unwrap();
1864
1865 assert_eq!(rows.schema.len(), 2);
1866 assert_eq!(rows.schema[1].column_name, "greptime_value");
1867 assert_eq!(rows.rows[0].values.len(), 2);
1868 assert!(
1869 rows.rows[0].values[1].value_data.is_none(),
1870 "missing nullable field should be filled with null"
1871 );
1872 }
1873
1874 #[test]
1875 fn test_fill_missing_field_column_rejects_impure_default() {
1876 let field_meta = ColumnMetadata {
1877 column_id: 1,
1878 semantic_type: SemanticType::Field,
1879 column_schema: ColumnSchema::new(
1880 "greptime_value".to_string(),
1881 ConcreteDataType::timestamp_millisecond_datatype(),
1882 false,
1883 )
1884 .with_default_constraint(Some(ColumnDefaultConstraint::Function("now()".to_string())))
1885 .unwrap(),
1886 };
1887 let mut rows = Rows {
1888 schema: vec![PbColumnSchema {
1889 column_name: "ts".to_string(),
1890 datatype: api::v1::ColumnDataType::TimestampMillisecond as i32,
1891 semantic_type: SemanticType::Timestamp as _,
1892 datatype_extension: None,
1893 options: None,
1894 }],
1895 rows: vec![Row {
1896 values: vec![Value {
1897 value_data: Some(ValueData::TimestampMillisecondValue(0)),
1898 }],
1899 }],
1900 };
1901
1902 let err = MetricEngineInner::fill_missing_field_column(
1903 RegionId::new(1, 1),
1904 "greptime_value",
1905 &field_meta,
1906 &mut rows,
1907 )
1908 .unwrap_err();
1909 assert!(
1910 err.to_string().contains("impure default value"),
1911 "expected impure-default rejection, got: {err}"
1912 );
1913 }
1914}