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