1use std::collections::HashMap;
16use std::sync::Arc;
17
18use api::prom_store::remote::read_request::ResponseType;
19use api::prom_store::remote::{Query, QueryResult, ReadRequest, ReadResponse};
20use api::v1::alter_table_expr::Kind;
21use api::v1::{
22 AddColumn, AddColumns, AlterTableExpr, ColumnDataType, ColumnDef, CreateTableExpr,
23 RowInsertRequests, SemanticType,
24};
25use async_trait::async_trait;
26use auth::{
27 PROM_STORE_READ, PROM_STORE_WRITE, PermissionChecker, PermissionCheckerRef, PermissionReq,
28 PermissionTableTarget, PermissionTableTargets,
29};
30use client::OutputData;
31use common_catalog::{format_full_table_name, parse_optional_catalog_and_schema_from_db_string};
32use common_error::ext::BoxedError;
33use common_meta::rpc::ddl::TriggerReason;
34use common_query::Output;
35use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_value};
36use common_recordbatch::RecordBatches;
37use common_telemetry::{debug, tracing};
38use datatypes::data_type::ConcreteDataType;
39use datatypes::schema::Schema;
40use operator::insert::{
41 AutoCreateTableType, InserterRef, PerTableSemanticIndex, apply_per_table_semantic_options,
42 build_create_table_expr, fill_table_options_for_create, parse_per_table_semantic_index,
43};
44use operator::statement::StatementExecutor;
45use prost::Message;
46use query::query_engine::options::{QueryOptions, validate_catalog_and_schema};
47use servers::batcher::logical_table::PendingRowsSchemaAlterer;
48use servers::error::{self, AuthSnafu, Result as ServerResult};
49use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF, collect_plan_metrics};
50use servers::http::prom_store::PHYSICAL_TABLE_PARAM;
51use servers::interceptor::{PromStoreProtocolInterceptor, PromStoreProtocolInterceptorRef};
52use servers::prom_store;
53use servers::query_handler::{
54 PromStoreProtocolHandler, PromStoreProtocolHandlerRef, PromStoreResponse,
55};
56use session::context::QueryContextRef;
57use snafu::{OptionExt, ResultExt, ensure};
58use store_api::metric_engine_consts::{METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY};
59use store_api::mito_engine_options::SST_FORMAT_KEY;
60use table::TableRef;
61use table::table_reference::TableReference;
62use tracing::instrument;
63
64use crate::error::{
65 AmbiguousValueColumnSnafu, CatalogSnafu, ColumnNotFoundSnafu, ExecLogicalPlanSnafu,
66 InvalidPromRemoteReadSchemaSnafu, PromStoreRemoteQueryPlanSnafu, ReadTableSnafu, Result,
67};
68use crate::instance::Instance;
69
70const SAMPLES_RESPONSE_TYPE: i32 = ResponseType::Samples as i32;
71
72struct RemoteQueryOutput {
73 table_name: String,
74 timestamp_column_name: String,
75 value_column_name: String,
76 output: Output,
77}
78
79fn auto_create_table_type_for_prom_remote_write(
80 ctx: &QueryContextRef,
81 with_metric_engine: bool,
82) -> AutoCreateTableType {
83 if with_metric_engine {
84 let physical_table = ctx
85 .extension(PHYSICAL_TABLE_PARAM)
86 .unwrap_or(GREPTIME_PHYSICAL_TABLE)
87 .to_string();
88 AutoCreateTableType::Logical(physical_table)
89 } else {
90 AutoCreateTableType::Physical
91 }
92}
93
94fn required_physical_table_for_create_type(create_type: &AutoCreateTableType) -> Option<&str> {
95 match create_type {
96 AutoCreateTableType::Logical(physical_table) => Some(physical_table.as_str()),
97 _ => None,
98 }
99}
100
101fn fill_metric_physical_table_options(table_options: &mut HashMap<String, String>) {
102 table_options.insert(SST_FORMAT_KEY.to_string(), "flat".to_string());
104 table_options.insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), "true".to_string());
105}
106
107#[inline]
108fn is_supported(response_type: i32) -> bool {
109 response_type == SAMPLES_RESPONSE_TYPE
111}
112
113fn negotiate_response_type(accepted_response_types: &[i32]) -> ServerResult<ResponseType> {
119 if accepted_response_types.is_empty() {
120 return Ok(ResponseType::Samples);
121 }
122
123 let response_type = accepted_response_types
124 .iter()
125 .find(|t| is_supported(**t))
126 .with_context(|| error::NotSupportedSnafu {
127 feat: format!(
128 "server does not support any of the requested response types: {accepted_response_types:?}",
129 ),
130 })?;
131
132 Ok(ResponseType::try_from(*response_type).unwrap())
134}
135
136fn resolve_remote_query_target(
137 ctx: &QueryContextRef,
138 query: &Query,
139 table_name: &str,
140) -> PermissionTableTarget {
141 match prom_store::extract_schema_from_query(query) {
142 Some(database) => {
143 let (catalog, schema) = parse_optional_catalog_and_schema_from_db_string(&database);
144 PermissionTableTarget::new(
145 catalog.unwrap_or_else(|| ctx.current_catalog().to_string()),
146 schema,
147 table_name,
148 )
149 }
150 None => PermissionTableTarget::new(ctx.current_catalog(), ctx.current_schema(), table_name),
151 }
152}
153
154#[instrument(skip_all, fields(table_name))]
155async fn to_query_result(
156 table_name: &str,
157 timestamp_column_name: &str,
158 value_column_name: &str,
159 output: Output,
160) -> ServerResult<QueryResult> {
161 let OutputData::Stream(stream) = output.data else {
162 unreachable!()
163 };
164 let recordbatches = RecordBatches::try_collect(stream)
165 .await
166 .context(error::CollectRecordbatchSnafu)?;
167 Ok(QueryResult {
168 timeseries: prom_store::recordbatches_to_timeseries(
169 table_name,
170 timestamp_column_name,
171 value_column_name,
172 recordbatches,
173 )?,
174 })
175}
176
177fn resolve_value_column_name(table_name: &str, table: &TableRef) -> Result<String> {
178 let columns = table
179 .field_columns()
180 .map(|column| column.name)
181 .collect::<Vec<_>>();
182
183 match columns.as_slice() {
184 [] => ColumnNotFoundSnafu {
185 msg: format!("value field in table '{table_name}'"),
186 }
187 .fail(),
188
189 [only] => Ok(only.clone()),
190
191 columns if columns.iter().any(|name| name == greptime_value()) => {
192 Ok(greptime_value().to_string())
193 }
194
195 columns => AmbiguousValueColumnSnafu {
196 table_name: table_name.to_string(),
197 field_columns: columns.to_vec(),
198 }
199 .fail(),
200 }
201}
202
203fn ensure_string_label_columns(
204 schema: &Schema,
205 timestamp_column_name: &str,
206 value_column_name: &str,
207) -> Result<()> {
208 for column in schema.column_schemas() {
209 if column.name == timestamp_column_name || column.name == value_column_name {
210 continue;
211 }
212 ensure!(
214 matches!(column.data_type, ConcreteDataType::String(_)),
215 InvalidPromRemoteReadSchemaSnafu {
216 detail: format!(
217 "Prometheus remote read only supports string labels: column '{}' has type {:?}",
218 column.name, column.data_type
219 ),
220 }
221 );
222 }
223 Ok(())
224}
225
226impl Instance {
227 #[tracing::instrument(skip_all)]
228 async fn handle_remote_query(
229 &self,
230 ctx: &QueryContextRef,
231 catalog_name: &str,
232 schema_name: &str,
233 table_name: &str,
234 query: &Query,
235 ) -> Result<Option<RemoteQueryOutput>> {
236 let Some(table) = self
237 .catalog_manager
238 .table(catalog_name, schema_name, table_name, Some(ctx))
239 .await
240 .context(CatalogSnafu)?
241 else {
242 debug!(
243 catalog = catalog_name,
244 schema = schema_name,
245 table = table_name,
246 "Table not found during remote read, returning empty result"
247 );
248 return Ok(None);
249 };
250
251 let timestamp_column_name = table
252 .schema()
253 .timestamp_column()
254 .with_context(|| ColumnNotFoundSnafu {
255 msg: format!("time index in table '{table_name}'"),
256 })?
257 .name
258 .clone();
259
260 let value_column_name = resolve_value_column_name(table_name, &table)?;
261
262 ensure_string_label_columns(&table.schema(), ×tamp_column_name, &value_column_name)?;
263
264 let dataframe = self
265 .query_engine
266 .read_table(table)
267 .with_context(|_| ReadTableSnafu {
268 table_name: format_full_table_name(catalog_name, schema_name, table_name),
269 })?;
270
271 let logical_plan = prom_store::query_to_plan(dataframe, query, ×tamp_column_name)
272 .context(PromStoreRemoteQueryPlanSnafu)?;
273
274 debug!(
275 "Prometheus remote read, table: {}, logical plan: {}",
276 table_name,
277 logical_plan.display_indent(),
278 );
279
280 let output = self
281 .query_engine
282 .execute(logical_plan, ctx.clone())
283 .await
284 .context(ExecLogicalPlanSnafu)?;
285
286 Ok(Some(RemoteQueryOutput {
287 table_name: table_name.to_string(),
288 timestamp_column_name,
289 value_column_name,
290 output,
291 }))
292 }
293
294 #[tracing::instrument(skip_all)]
295 async fn handle_remote_queries(
296 &self,
297 ctx: QueryContextRef,
298 queries: &[Query],
299 query_targets: &[PermissionTableTarget],
300 ) -> ServerResult<Vec<Option<RemoteQueryOutput>>> {
301 let mut results = Vec::with_capacity(queries.len());
302
303 for (query, target) in queries.iter().zip(query_targets) {
304 let result = self
305 .handle_remote_query(&ctx, &target.catalog, &target.schema, &target.table, query)
306 .await
307 .map_err(BoxedError::new)
308 .context(error::ExecuteQuerySnafu)?;
309
310 results.push(result);
311 }
312 Ok(results)
313 }
314}
315
316#[async_trait]
317impl PendingRowsSchemaAlterer for Instance {
318 async fn create_tables_if_missing_batch(
319 &self,
320 catalog: &str,
321 schema: &str,
322 tables: &[(&str, &[api::v1::ColumnSchema])],
323 with_metric_engine: bool,
324 ctx: QueryContextRef,
325 ) -> ServerResult<()> {
326 if tables.is_empty() {
327 return Ok(());
328 }
329
330 let create_type = auto_create_table_type_for_prom_remote_write(&ctx, with_metric_engine);
331 if let Some(physical_table) = required_physical_table_for_create_type(&create_type) {
332 self.create_metric_physical_table_if_missing(
333 catalog,
334 schema,
335 physical_table,
336 ctx.clone(),
337 )
338 .await?;
339 }
340
341 let engine = if matches!(create_type, AutoCreateTableType::Logical(_)) {
342 METRIC_ENGINE_NAME
343 } else {
344 common_catalog::consts::default_engine()
345 };
346
347 let mut create_exprs: Vec<CreateTableExpr> = Vec::with_capacity(tables.len());
350 let mut per_table_semantics: Option<Option<PerTableSemanticIndex>> = None;
351 for &(table_name, request_schema) in tables {
352 let existing = self
353 .catalog_manager()
354 .table(catalog, schema, table_name, Some(ctx.as_ref()))
355 .await
356 .map_err(BoxedError::new)
357 .context(error::ExecuteGrpcQuerySnafu)?;
358 if existing.is_some() {
359 continue;
360 }
361
362 let table_ref = TableReference::full(catalog, schema, table_name);
363 let mut create_table_expr = build_create_table_expr(&table_ref, request_schema, engine)
364 .map_err(BoxedError::new)
365 .context(error::ExecuteGrpcQuerySnafu)?;
366
367 let mut table_options = std::collections::HashMap::with_capacity(4);
368 fill_table_options_for_create(&mut table_options, &create_type, &ctx);
369 let semantic_index = per_table_semantics
372 .get_or_insert_with(|| parse_per_table_semantic_index(&ctx))
373 .as_ref();
374 apply_per_table_semantic_options(
375 &mut table_options,
376 semantic_index,
377 schema,
378 table_name,
379 );
380 create_table_expr.table_options.extend(table_options);
381 create_exprs.push(create_table_expr);
382 }
383
384 if create_exprs.is_empty() {
385 return Ok(());
386 }
387
388 match create_type {
389 AutoCreateTableType::Logical(_) => {
390 self.statement_executor
392 .create_logical_tables(&create_exprs, ctx, TriggerReason::AutoCreate)
393 .await
394 .map_err(BoxedError::new)
395 .context(error::ExecuteGrpcQuerySnafu)?;
396 }
397 AutoCreateTableType::Physical => {
398 for mut expr in create_exprs {
400 expr.table_options
401 .insert(SST_FORMAT_KEY.to_string(), "flat".to_string());
402 self.statement_executor
403 .create_table_inner(&mut expr, None, ctx.clone(), TriggerReason::AutoCreate)
404 .await
405 .map_err(BoxedError::new)
406 .context(error::ExecuteGrpcQuerySnafu)?;
407 }
408 }
409 create_type => {
410 return error::InvalidPromRemoteRequestSnafu {
411 msg: format!(
412 "prom remote write only supports logical or physical auto-create: {}",
413 create_type.as_str()
414 ),
415 }
416 .fail();
417 }
418 }
419
420 Ok(())
421 }
422
423 async fn add_missing_prom_tag_columns_batch(
424 &self,
425 catalog: &str,
426 schema: &str,
427 tables: &[(&str, &[String])],
428 ctx: QueryContextRef,
429 ) -> ServerResult<()> {
430 if tables.is_empty() {
431 return Ok(());
432 }
433
434 let alter_exprs: Vec<AlterTableExpr> = tables
435 .iter()
436 .filter(|(_, columns)| !columns.is_empty())
437 .map(|&(table_name, columns)| {
438 let add_columns = AddColumns {
439 add_columns: columns
440 .iter()
441 .map(|column_name| AddColumn {
442 column_def: Some(ColumnDef {
443 name: column_name.clone(),
444 data_type: ColumnDataType::String as i32,
445 is_nullable: true,
446 semantic_type: SemanticType::Tag as i32,
447 comment: String::new(),
448 ..Default::default()
449 }),
450 location: None,
451 add_if_not_exists: true,
452 })
453 .collect(),
454 };
455
456 AlterTableExpr {
457 catalog_name: catalog.to_string(),
458 schema_name: schema.to_string(),
459 table_name: table_name.to_string(),
460 kind: Some(Kind::AddColumns(add_columns)),
461 }
462 })
463 .collect();
464
465 if alter_exprs.is_empty() {
466 return Ok(());
467 }
468
469 self.statement_executor
470 .alter_logical_tables(alter_exprs, ctx, TriggerReason::AutoAlter)
471 .await
472 .map_err(BoxedError::new)
473 .context(error::ExecuteGrpcQuerySnafu)?;
474
475 Ok(())
476 }
477}
478
479impl Instance {
480 async fn prepare_prom_store_write(
481 &self,
482 request: RowInsertRequests,
483 ctx: QueryContextRef,
484 ) -> ServerResult<(RowInsertRequests, QueryContextRef)> {
485 PromStoreProtocolHandler::pre_write(self, &request, ctx.clone()).await?;
486 Ok((request, Arc::new(ctx.fork())))
487 }
488
489 async fn execute_prom_store_write(
490 &self,
491 request: RowInsertRequests,
492 ctx: QueryContextRef,
493 with_metric_engine: bool,
494 ) -> ServerResult<Output> {
495 let output = if with_metric_engine {
496 let physical_table = ctx
497 .extension(PHYSICAL_TABLE_PARAM)
498 .unwrap_or(GREPTIME_PHYSICAL_TABLE)
499 .to_string();
500 self.handle_metric_row_inserts(request, ctx.clone(), physical_table.clone())
501 .await
502 .map_err(BoxedError::new)
503 .context(error::ExecuteGrpcQuerySnafu)?
504 } else {
505 self.handle_row_inserts(request, ctx.clone(), true, true)
506 .await
507 .map_err(BoxedError::new)
508 .context(error::ExecuteGrpcQuerySnafu)?
509 };
510
511 Ok(output)
512 }
513}
514
515#[async_trait]
516impl PromStoreProtocolHandler for Instance {
517 async fn pre_write(
518 &self,
519 request: &RowInsertRequests,
520 ctx: QueryContextRef,
521 ) -> ServerResult<()> {
522 self.plugins
523 .get::<PermissionCheckerRef>()
524 .as_ref()
525 .check_permission(ctx.current_user(), PermissionReq::Action(PROM_STORE_WRITE))
526 .context(AuthSnafu)?;
527 let interceptor_ref = self
528 .plugins
529 .get::<PromStoreProtocolInterceptorRef<servers::error::Error>>();
530 interceptor_ref.pre_write(request, ctx.clone())?;
531 self.check_row_insert_permission(request, &ctx, PermissionReq::Action(PROM_STORE_WRITE))
532 .context(AuthSnafu)?;
533 Ok(())
534 }
535
536 async fn write_prepared(
537 &self,
538 request: RowInsertRequests,
539 ctx: QueryContextRef,
540 with_metric_engine: bool,
541 ) -> ServerResult<Output> {
542 self.execute_prom_store_write(request, ctx, with_metric_engine)
543 .await
544 }
545
546 async fn write(
547 &self,
548 request: RowInsertRequests,
549 ctx: QueryContextRef,
550 with_metric_engine: bool,
551 ) -> ServerResult<Output> {
552 let (request, ctx) = self.prepare_prom_store_write(request, ctx).await?;
553 self.write_prepared(request, ctx, with_metric_engine).await
554 }
555
556 async fn write_all(
557 &self,
558 requests: Vec<(QueryContextRef, RowInsertRequests)>,
559 with_metric_engine: bool,
560 ) -> ServerResult<Vec<ServerResult<Output>>> {
561 let mut prepared = Vec::with_capacity(requests.len());
562 for (ctx, request) in requests {
563 let (request, ctx) = self.prepare_prom_store_write(request, ctx).await?;
564 prepared.push((ctx, request));
565 }
566 operator::insert::admit_row_insert_batches(&mut prepared)
567 .await
568 .map_err(BoxedError::new)
569 .context(error::ExecuteGrpcQuerySnafu)?;
570
571 let mut outputs = Vec::with_capacity(prepared.len());
572 for (ctx, request) in prepared {
573 let output = self.write_prepared(request, ctx, with_metric_engine).await;
574 let failed = output.is_err();
575 outputs.push(output);
576 if failed {
577 break;
578 }
579 }
580 Ok(outputs)
581 }
582
583 #[instrument(skip_all, fields(table_name))]
584 async fn read(
585 &self,
586 request: ReadRequest,
587 ctx: QueryContextRef,
588 ) -> ServerResult<PromStoreResponse> {
589 self.plugins
590 .get::<PermissionCheckerRef>()
591 .as_ref()
592 .check_permission(ctx.current_user(), PermissionReq::Action(PROM_STORE_READ))
593 .context(AuthSnafu)?;
594
595 let interceptor_ref = self
596 .plugins
597 .get::<PromStoreProtocolInterceptorRef<servers::error::Error>>();
598 interceptor_ref.pre_read(&request, ctx.clone())?;
599 let ctx = Arc::new(ctx.fork());
600
601 let table_names = request
602 .queries
603 .iter()
604 .map(prom_store::table_name)
605 .collect::<ServerResult<Vec<_>>>()?;
606 let query_targets = request
607 .queries
608 .iter()
609 .zip(&table_names)
610 .map(|(query, table_name)| resolve_remote_query_target(&ctx, query, table_name))
611 .collect::<Vec<_>>();
612 let disallow_cross_catalog_query = self
613 .plugins
614 .get::<QueryOptions>()
615 .map(|opts| opts.disallow_cross_catalog_query)
616 .unwrap_or_default();
617 if disallow_cross_catalog_query {
618 for target in &query_targets {
619 validate_catalog_and_schema(&target.catalog, &target.schema, &ctx)
620 .map_err(BoxedError::new)
621 .context(error::ExecuteQuerySnafu)?;
622 }
623 }
624 let targets = self
625 .resolve_query_permission_targets(
626 PermissionTableTargets::resolved(query_targets.clone()),
627 &ctx,
628 )
629 .await?;
630 self.check_table_permission(&ctx, PermissionReq::Action(PROM_STORE_READ), targets)
631 .context(AuthSnafu)?;
632
633 let response_type = negotiate_response_type(&request.accepted_response_types)?;
634
635 for query in &request.queries {
636 prom_store::validate_matchers(&query.matchers)?;
637 }
638
639 let results = self
641 .handle_remote_queries(ctx, &request.queries, &query_targets)
642 .await?;
643
644 match response_type {
645 ResponseType::Samples => {
646 let mut query_results = Vec::with_capacity(results.len());
647 let mut map = HashMap::new();
648 for result in results {
649 let Some(result) = result else {
650 query_results.push(QueryResult { timeseries: vec![] });
651 continue;
652 };
653 let RemoteQueryOutput {
654 table_name,
655 timestamp_column_name,
656 value_column_name,
657 output,
658 } = result;
659 let plan = output.meta.plan.clone();
660 query_results.push(
661 to_query_result(
662 &table_name,
663 ×tamp_column_name,
664 &value_column_name,
665 output,
666 )
667 .await?,
668 );
669 if let Some(ref plan) = plan {
670 collect_plan_metrics(plan, &mut [&mut map]);
671 }
672 }
673
674 let response = ReadResponse {
675 results: query_results,
676 };
677
678 let resp_metrics = map
679 .into_iter()
680 .map(|(k, v)| (k, v.into()))
681 .collect::<HashMap<_, _>>();
682
683 Ok(PromStoreResponse {
685 content_type: CONTENT_TYPE_PROTOBUF.clone(),
686 content_encoding: CONTENT_ENCODING_SNAPPY.clone(),
687 resp_metrics,
688 body: prom_store::snappy_compress(&response.encode_to_vec())?,
689 })
690 }
691 ResponseType::StreamedXorChunks => error::NotSupportedSnafu {
692 feat: "streamed remote read",
693 }
694 .fail(),
695 }
696 }
697}
698
699impl Instance {
700 async fn create_metric_physical_table_if_missing(
701 &self,
702 catalog: &str,
703 schema: &str,
704 physical_table: &str,
705 ctx: QueryContextRef,
706 ) -> ServerResult<()> {
707 let table = self
708 .catalog_manager()
709 .table(catalog, schema, physical_table, Some(ctx.as_ref()))
710 .await
711 .map_err(BoxedError::new)
712 .context(error::ExecuteGrpcQuerySnafu)?;
713 if table.is_some() {
714 return Ok(());
715 }
716
717 let table_ref = TableReference::full(catalog, schema, physical_table);
718 let default_schema = vec![
719 api::v1::ColumnSchema {
720 column_name: common_query::prelude::greptime_timestamp().to_string(),
721 datatype: api::v1::ColumnDataType::TimestampMillisecond as i32,
722 semantic_type: api::v1::SemanticType::Timestamp as i32,
723 datatype_extension: None,
724 options: None,
725 },
726 api::v1::ColumnSchema {
727 column_name: common_query::prelude::greptime_value().to_string(),
728 datatype: api::v1::ColumnDataType::Float64 as i32,
729 semantic_type: api::v1::SemanticType::Field as i32,
730 datatype_extension: None,
731 options: None,
732 },
733 ];
734 let mut create_table_expr = build_create_table_expr(
735 &table_ref,
736 &default_schema,
737 common_catalog::consts::default_engine(),
738 )
739 .map_err(BoxedError::new)
740 .context(error::ExecuteGrpcQuerySnafu)?;
741 create_table_expr.engine = METRIC_ENGINE_NAME.to_string();
742 fill_metric_physical_table_options(&mut create_table_expr.table_options);
743
744 self.statement_executor
745 .create_table_inner(&mut create_table_expr, None, ctx, TriggerReason::AutoCreate)
746 .await
747 .map_err(BoxedError::new)
748 .context(error::ExecuteGrpcQuerySnafu)?;
749
750 Ok(())
751 }
752}
753
754pub struct ExportMetricHandler {
759 inserter: InserterRef,
760 statement_executor: Arc<StatementExecutor>,
761}
762
763impl ExportMetricHandler {
764 pub fn new_handler(
765 inserter: InserterRef,
766 statement_executor: Arc<StatementExecutor>,
767 ) -> PromStoreProtocolHandlerRef {
768 Arc::new(Self {
769 inserter,
770 statement_executor,
771 })
772 }
773}
774
775#[async_trait]
776impl PromStoreProtocolHandler for ExportMetricHandler {
777 async fn pre_write(
778 &self,
779 _request: &RowInsertRequests,
780 _ctx: QueryContextRef,
781 ) -> ServerResult<()> {
782 Ok(())
783 }
784
785 async fn write_prepared(
786 &self,
787 request: RowInsertRequests,
788 ctx: QueryContextRef,
789 _: bool,
790 ) -> ServerResult<Output> {
791 self.inserter
792 .handle_metric_row_inserts(
793 request,
794 ctx,
795 &self.statement_executor,
796 GREPTIME_PHYSICAL_TABLE.to_string(),
797 )
798 .await
799 .map_err(BoxedError::new)
800 .context(error::ExecuteGrpcQuerySnafu)
801 }
802
803 async fn write(
804 &self,
805 request: RowInsertRequests,
806 ctx: QueryContextRef,
807 with_metric_engine: bool,
808 ) -> ServerResult<Output> {
809 self.write_prepared(request, ctx, with_metric_engine).await
810 }
811
812 async fn write_all(
813 &self,
814 mut requests: Vec<(QueryContextRef, RowInsertRequests)>,
815 with_metric_engine: bool,
816 ) -> ServerResult<Vec<ServerResult<Output>>> {
817 operator::insert::admit_row_insert_batches(&mut requests)
818 .await
819 .map_err(BoxedError::new)
820 .context(error::ExecuteGrpcQuerySnafu)?;
821 let mut outputs = Vec::with_capacity(requests.len());
822 for (ctx, request) in requests {
823 let output = self.write_prepared(request, ctx, with_metric_engine).await;
824 let failed = output.is_err();
825 outputs.push(output);
826 if failed {
827 break;
828 }
829 }
830 Ok(outputs)
831 }
832
833 async fn read(
834 &self,
835 _request: ReadRequest,
836 _ctx: QueryContextRef,
837 ) -> ServerResult<PromStoreResponse> {
838 unreachable!();
839 }
840}
841
842#[cfg(test)]
843mod tests {
844 use std::sync::Arc;
845
846 use api::prom_store::remote::LabelMatcher;
847 use common_error::ext::ErrorExt;
848 use common_error::status_code::StatusCode;
849 use session::context::QueryContext;
850
851 use super::*;
852
853 #[test]
854 fn test_remote_read_label_types() {
855 use datatypes::schema::ColumnSchema;
856
857 for (label_type, supported) in [
858 (ConcreteDataType::string_datatype(), true),
859 (ConcreteDataType::large_string_datatype(), true),
860 (ConcreteDataType::utf8_view_datatype(), true),
861 (ConcreteDataType::binary_datatype(), false),
862 (ConcreteDataType::binary_view_datatype(), false),
863 (ConcreteDataType::json_datatype(), false),
864 (ConcreteDataType::vector_datatype(3), false),
865 (ConcreteDataType::int64_datatype(), false),
866 (ConcreteDataType::boolean_datatype(), false),
867 ] {
868 let schema = Schema::new(vec![
869 ColumnSchema::new(
870 "ts",
871 ConcreteDataType::timestamp_millisecond_datatype(),
872 false,
873 ),
874 ColumnSchema::new("value", label_type.clone(), true),
875 ColumnSchema::new("label", ConcreteDataType::string_datatype(), true),
876 ]);
877 ensure_string_label_columns(&schema, "ts", "value").unwrap();
879
880 let result = ensure_string_label_columns(&schema, "ts", "label");
881 assert_eq!(supported, result.is_ok(), "{label_type:?}");
882 if let Err(err) = result {
883 assert!(matches!(
884 err,
885 crate::error::Error::InvalidPromRemoteReadSchema { .. }
886 ));
887 assert_eq!(err.status_code(), StatusCode::InvalidArguments);
888 assert!(
889 err.to_string()
890 .starts_with("Invalid Prometheus remote read schema: ")
891 );
892 assert!(
893 err.to_string()
894 .contains("only supports string labels: column 'value'")
895 );
896 }
897 }
898 }
899
900 #[test]
901 fn test_auto_create_table_type_for_prom_remote_write_metric_engine() {
902 let mut query_ctx = QueryContext::with(
903 common_catalog::consts::DEFAULT_CATALOG_NAME,
904 common_catalog::consts::DEFAULT_SCHEMA_NAME,
905 );
906 query_ctx.set_extension(PHYSICAL_TABLE_PARAM, "metric_physical".to_string());
907 let ctx = Arc::new(query_ctx);
908
909 let create_type = auto_create_table_type_for_prom_remote_write(&ctx, true);
910 match create_type {
911 AutoCreateTableType::Logical(physical) => assert_eq!(physical, "metric_physical"),
912 _ => panic!("expected logical table create type"),
913 }
914 }
915
916 #[test]
917 fn test_auto_create_table_type_for_prom_remote_write_without_metric_engine() {
918 let ctx = Arc::new(QueryContext::with(
919 common_catalog::consts::DEFAULT_CATALOG_NAME,
920 common_catalog::consts::DEFAULT_SCHEMA_NAME,
921 ));
922
923 let create_type = auto_create_table_type_for_prom_remote_write(&ctx, false);
924 match create_type {
925 AutoCreateTableType::Physical => {}
926 _ => panic!("expected physical table create type"),
927 }
928 }
929
930 #[test]
931 fn test_required_physical_table_for_create_type() {
932 let logical = AutoCreateTableType::Logical("phy_table".to_string());
933 assert_eq!(
934 Some("phy_table"),
935 required_physical_table_for_create_type(&logical)
936 );
937
938 let physical = AutoCreateTableType::Physical;
939 assert_eq!(None, required_physical_table_for_create_type(&physical));
940 }
941
942 #[test]
943 fn test_metric_physical_table_options_forces_flat_sst_format() {
944 let mut table_options = HashMap::new();
945
946 fill_metric_physical_table_options(&mut table_options);
947
948 assert_eq!(
949 Some("flat"),
950 table_options.get(SST_FORMAT_KEY).map(String::as_str)
951 );
952 assert_eq!(
953 Some("true"),
954 table_options
955 .get(PHYSICAL_TABLE_METADATA_KEY)
956 .map(String::as_str)
957 );
958 }
959
960 fn query_with_database(database: &str) -> Query {
961 Query {
962 matchers: vec![LabelMatcher {
963 name: servers::prom_store::DATABASE_LABEL.to_string(),
964 value: database.to_string(),
965 r#type: api::prom_store::remote::label_matcher::Type::Eq as i32,
966 }],
967 ..Default::default()
968 }
969 }
970
971 #[test]
972 fn test_resolve_remote_query_target() {
973 let ctx = Arc::new(QueryContext::with("request_catalog", "request_schema"));
974
975 assert_eq!(
976 PermissionTableTarget::new("request_catalog", "request_schema", "fallback_metric"),
977 resolve_remote_query_target(&ctx, &Query::default(), "fallback_metric")
978 );
979 assert_eq!(
980 PermissionTableTarget::new("request_catalog", "selected_schema", "schema_metric"),
981 resolve_remote_query_target(
982 &ctx,
983 &query_with_database("selected_schema"),
984 "schema_metric"
985 )
986 );
987 assert_eq!(
988 PermissionTableTarget::new("selected_catalog", "selected_schema", "catalog_metric"),
989 resolve_remote_query_target(
990 &ctx,
991 &query_with_database("selected_catalog-selected_schema"),
992 "catalog_metric"
993 )
994 );
995 }
996
997 #[test]
998 fn test_resolve_remote_query_target_preserves_context() {
999 let ctx = Arc::new(QueryContext::with("request_catalog", "request_schema"));
1000
1001 resolve_remote_query_target(
1002 &ctx,
1003 &query_with_database("selected_catalog-selected_schema"),
1004 "metric",
1005 );
1006
1007 assert_eq!("request_catalog", ctx.current_catalog());
1008 assert_eq!("request_schema", ctx.current_schema());
1009 }
1010}