feat: allow customized time index unit for metric engine table

tests-integration/src/otlp.rs

@@ -19,6 +19,7 @@ mod test {
     use client::{DEFAULT_CATALOG_NAME, OutputData};
     use common_recordbatch::RecordBatches;
     use datatypes::arrow::array::AsArray;
+    use datatypes::arrow::datatypes::TimestampMicrosecondType;
     use frontend::instance::Instance;
     use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
     use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest;
@@ -37,6 +38,7 @@ mod test {
     use servers::query_handler::OpenTelemetryProtocolHandler;
     use servers::query_handler::sql::SqlQueryHandler;
     use session::context::QueryContext;
+    use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx};
 
     use crate::standalone::GreptimeDbStandaloneBuilder;
     use crate::tests;
@@ -888,6 +890,99 @@ WITH(
         );
     }
 
+    #[tokio::test(flavor = "multi_thread")]
+    pub async fn test_otlp_metrics_into_microsecond_physical_table_on_standalone() {
+        let standalone = GreptimeDbStandaloneBuilder::new("test_otlp_us_physical")
+            .build()
+            .await;
+
+        test_otlp_metrics_into_microsecond_physical_table(standalone.fe_instance()).await;
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    pub async fn test_otlp_metrics_into_microsecond_physical_table_on_distributed() {
+        let instance = tests::create_distributed_instance("test_otlp_us_physical_dist").await;
+
+        test_otlp_metrics_into_microsecond_physical_table(&instance.frontend()).await;
+    }
+
+    /// Regression test for <https://github.com/GreptimeTeam/greptimedb/issues/9231>:
+    /// a physical metric table pre-created with a TIMESTAMP(6) (microsecond)
+    /// time index accepts OTLP ingestion; nanosecond samples are converted to
+    /// the physical table's unit, keeping microsecond precision.
+    async fn test_otlp_metrics_into_microsecond_physical_table(instance: &Arc<Instance>) {
+        let db = "otlp_us_physical";
+        let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, db);
+        // Route the request to the metric engine, like the OTLP HTTP handler
+        // does with the default `prom_store.with_metric_engine = true`.
+        ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx {
+            with_metric_engine: true,
+            ..Default::default()
+        }));
+        let ctx = Arc::new(ctx);
+        assert!(
+            SqlQueryHandler::do_query(
+                instance.as_ref(),
+                &format!("CREATE DATABASE IF NOT EXISTS {db}"),
+                ctx.clone(),
+            )
+            .await
+            .first()
+            .unwrap()
+            .is_ok()
+        );
+
+        let mut output = instance
+            .do_query(
+                "CREATE TABLE greptime_physical_table (\
+                 greptime_timestamp TIMESTAMP(6) NOT NULL, \
+                 greptime_value DOUBLE NULL, \
+                 TIME INDEX (greptime_timestamp)) \
+                 ENGINE = metric WITH ('physical_metric_table' = 'true')",
+                ctx.clone(),
+            )
+            .await;
+        assert!(output.remove(0).is_ok());
+
+        let request = ExportMetricsServiceRequest {
+            resource_metrics: vec![ResourceMetrics {
+                scope_metrics: vec![ScopeMetrics {
+                    metrics: vec![Metric {
+                        name: "my_gauge".to_string(),
+                        data: Some(metric::Data::Gauge(Gauge {
+                            data_points: vec![NumberDataPoint {
+                                attributes: vec![keyvalue("host", "h1")],
+                                time_unix_nano: 1_704_067_200_123_456_789,
+                                value: Some(Value::AsDouble(1.0)),
+                                ..Default::default()
+                            }],
+                        })),
+                        ..Default::default()
+                    }],
+                    ..Default::default()
+                }],
+                ..Default::default()
+            }],
+        };
+        instance.metrics(request, ctx.clone()).await.unwrap();
+
+        let mut output = instance
+            .do_query("SELECT greptime_timestamp FROM my_gauge", ctx.clone())
+            .await;
+        let OutputData::Stream(stream) = output.remove(0).unwrap().data else {
+            unreachable!()
+        };
+        let batches = RecordBatches::try_collect(stream).await.unwrap().take();
+        assert_eq!(batches[0].num_rows(), 1);
+        // The logical table's time index keeps the physical table's
+        // microsecond unit and precision: 1704067200123456789ns ->
+        // 1704067200123456us.
+        let timestamps = batches[0]
+            .column(0)
+            .as_primitive::<TimestampMicrosecondType>();
+        assert_eq!(timestamps.value(0), 1_704_067_200_123_456);
+    }
+
     fn build_request() -> ExportMetricsServiceRequest {
         let data_points = vec![
             NumberDataPoint {

tests-integration/src/prom_store.rs

@@ -20,7 +20,10 @@ mod tests {
     use api::prom_store::remote::{
         Label, LabelMatcher, Query, ReadRequest, ReadResponse, Sample, WriteRequest,
     };
+    use client::OutputData;
     use common_catalog::consts::DEFAULT_CATALOG_NAME;
+    use datatypes::arrow::array::AsArray;
+    use datatypes::arrow::datatypes::TimestampMicrosecondType;
     use frontend::instance::Instance;
     use prost::Message;
     use servers::http::prom_store::PHYSICAL_TABLE_PARAM;
@@ -83,6 +86,121 @@ mod tests {
         .await;
     }
 
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_standalone_prom_store_remote_rw_microsecond_physical_table() {
+        common_telemetry::init_default_ut_logging();
+        let standalone =
+            GreptimeDbStandaloneBuilder::new("test_prom_store_remote_rw_us_physical_table")
+                .build()
+                .await;
+        let instance = standalone.fe_instance();
+
+        test_prom_store_remote_rw_microsecond_physical_table(instance).await;
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_distributed_prom_store_remote_rw_microsecond_physical_table() {
+        common_telemetry::init_default_ut_logging();
+        let distributed =
+            tests::create_distributed_instance("test_prom_store_remote_rw_us_physical_table").await;
+        test_prom_store_remote_rw_microsecond_physical_table(&distributed.frontend()).await;
+    }
+
+    /// Regression test for <https://github.com/GreptimeTeam/greptimedb/issues/9231>:
+    /// prometheus remote write/read against a physical metric table
+    /// pre-created with a TIMESTAMP(6) (microsecond) time index. Millisecond
+    /// samples are widened to the physical table's unit on write and narrowed
+    /// back to milliseconds on remote read.
+    async fn test_prom_store_remote_rw_microsecond_physical_table(instance: &Arc<Instance>) {
+        let db = "prometheus_us";
+        let ctx = Arc::new(QueryContext::with(DEFAULT_CATALOG_NAME, db));
+        assert!(
+            SqlQueryHandler::do_query(
+                instance.as_ref(),
+                &format!("CREATE DATABASE IF NOT EXISTS {db}"),
+                ctx.clone(),
+            )
+            .await
+            .first()
+            .unwrap()
+            .is_ok()
+        );
+
+        let mut output = instance
+            .do_query(
+                "CREATE TABLE greptime_physical_table (\
+                 greptime_timestamp TIMESTAMP(6) NOT NULL, \
+                 greptime_value DOUBLE NULL, \
+                 TIME INDEX (greptime_timestamp)) \
+                 ENGINE = metric WITH ('physical_metric_table' = 'true')",
+                ctx.clone(),
+            )
+            .await;
+        assert!(output.remove(0).is_ok());
+
+        let write_request = WriteRequest {
+            timeseries: vec![prom_store::mock_timeseries()[0].clone()],
+            ..Default::default()
+        };
+        let (row_inserts, _) = to_grpc_row_insert_requests(&write_request).unwrap();
+        instance
+            .write(row_inserts, ctx.clone(), true)
+            .await
+            .unwrap();
+
+        let read_request = ReadRequest {
+            queries: vec![Query {
+                start_timestamp_ms: 1000,
+                end_timestamp_ms: 2000,
+                matchers: vec![LabelMatcher {
+                    name: prom_store::METRIC_NAME_LABEL.to_string(),
+                    value: "metric1".to_string(),
+                    r#type: 0,
+                }],
+                ..Default::default()
+            }],
+            ..Default::default()
+        };
+        let resp = instance.read(read_request, ctx.clone()).await.unwrap();
+        let body = prom_store::snappy_decompress(&resp.body).unwrap();
+        let read_response = ReadResponse::decode(&body[..]).unwrap();
+        assert_eq!(1, read_response.results.len());
+        assert_eq!(1, read_response.results[0].timeseries.len());
+        let timeseries = &read_response.results[0].timeseries[0];
+        assert_eq!(
+            timeseries.samples,
+            vec![
+                Sample {
+                    value: 1.0,
+                    timestamp: 1000,
+                },
+                Sample {
+                    value: 2.0,
+                    timestamp: 2000,
+                },
+            ]
+        );
+
+        // The stored timestamps keep the physical table's microsecond unit.
+        let mut output = instance
+            .do_query("SELECT greptime_timestamp FROM metric1", ctx.clone())
+            .await;
+        let OutputData::Stream(stream) = output.remove(0).unwrap().data else {
+            unreachable!()
+        };
+        let batches = common_recordbatch::RecordBatches::try_collect(stream)
+            .await
+            .unwrap()
+            .take();
+        assert_eq!(batches[0].num_rows(), 2);
+        let timestamps = batches[0]
+            .column(0)
+            .as_primitive::<TimestampMicrosecondType>();
+        // Millisecond samples are widened to the microsecond time index.
+        assert_eq!(timestamps.value(0), 1_000_000);
+        assert_eq!(timestamps.value(1), 2_000_000);
+    }
+
     async fn test_prom_store_remote_rw(instance: &Arc<Instance>, physical_table: Option<String>) {
         let write_request = WriteRequest {
             timeseries: prom_store::mock_timeseries(),