Add DataClient struct, allowing users to query Grafana for data over gRPC

WIP, not even tested, feels like it might just work though?

TODO:
- [ ] Finish off `Frame::from_arrow` so it actually populates the chunks
- [ ] Test it out

Depends on grafana/grafana#55781.

Author: sd2k

crates/grafana-plugin-sdk/Cargo.toml

@@ -10,6 +10,7 @@ description = "SDK for building Grafana backend plugins."
 
 [dependencies]
 arrow2 = { version = "0.14.0", features = ["io_ipc"] }
+bytes = "1.2.1"
 chrono = "0.4.19"
 futures-core = "0.3.17"
 futures-util = "0.3.17"
@@ -40,7 +41,7 @@ tracing-subscriber = { version = "0.3.1", features = [
 
 [dev-dependencies]
 async-stream = "0.3.2"
-bytes = "1.1.0"
+bytes = "1.2.1"
 futures = "0.3.17"
 paste = "1.0.6"
 pretty_assertions = "1.0.0"

crates/grafana-plugin-sdk/examples/main.rs

@@ -10,7 +10,7 @@ use bytes::Bytes;
 use chrono::prelude::*;
 use futures_util::stream::FuturesOrdered;
 use http::Response;
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tokio_stream::StreamExt;
 use tracing::{debug, info};
@@ -21,28 +21,42 @@ use grafana_plugin_sdk::{
     prelude::*,
 };
 
-#[derive(Clone, Debug, Default)]
-struct MyPluginService(Arc<AtomicUsize>);
+#[derive(Clone, Debug)]
+struct MyPluginService {
+    counter: Arc<AtomicUsize>,
+    data_client: backend::DataClient,
+}
 
 impl MyPluginService {
     fn new() -> Self {
-        Self(Arc::new(AtomicUsize::new(0)))
+        Self {
+            counter: Arc::new(AtomicUsize::new(0)),
+            data_client: backend::DataClient::new("localhost:10000").expect("valid URL"),
+        }
     }
 }
 
 // Data service implementation.
 
-#[derive(Debug, Deserialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 #[serde(rename_all = "camelCase")]
 struct Query {
     pub expression: String,
     pub other_user_input: u64,
 }
 
+#[derive(Debug, Error)]
+enum SDKError {
+    #[error("SDK error: {0}")]
+    Data(data::Error),
+    #[error("Query data error: {0}")]
+    QueryData(backend::QueryDataError),
+}
+
 #[derive(Debug, Error)]
 #[error("Error querying backend for query {ref_id}: {source}")]
 struct QueryError {
-    source: data::Error,
+    source: SDKError,
     ref_id: String,
 }
 
@@ -58,40 +72,60 @@ impl backend::DataService for MyPluginService {
     type QueryError = QueryError;
     type Stream = backend::BoxDataResponseStream<Self::QueryError>;
     async fn query_data(&self, request: backend::QueryDataRequest<Self::Query>) -> Self::Stream {
+        let client = self.data_client.clone();
+        let transport_metadata = request.transport_metadata().clone();
         Box::pin(
             request
                 .queries
                 .into_iter()
-                .map(|x: DataQuery<Self::Query>| async move {
+                .map(|x: DataQuery<Self::Query>| {
                     // We can see the user's query in `x.query`:
                     debug!(
                         expression = x.query.expression,
                         other_user_input = x.query.other_user_input,
                         "Got backend query",
                     );
-                    // Here we create a single response Frame for each query.
-                    // Frames can be created from iterators of fields using [`IntoFrame`].
-                    Ok(backend::DataResponse::new(
-                        x.ref_id.clone(),
-                        vec![[
-                            // Fields can be created from iterators of a variety of
-                            // relevant datatypes.
-                            [
-                                Utc.ymd(2021, 1, 1).and_hms(12, 0, 0),
-                                Utc.ymd(2021, 1, 1).and_hms(12, 0, 1),
-                                Utc.ymd(2021, 1, 1).and_hms(12, 0, 2),
+
+                    let ref_id = x.ref_id.clone();
+                    let transport_metadata = transport_metadata.clone();
+                    let mut client = client.clone();
+
+                    async move {
+                        // We can proxy this request to a different datasource if we like.
+                        // Here we do that, but ignore the response.
+                        let proxied = client
+                            .query_data(vec![x.clone()], &transport_metadata)
+                            .await
+                            .map_err(|source| QueryError {
+                                source: SDKError::QueryData(source),
+                                ref_id: ref_id.clone(),
+                            })?;
+                        info!("Got proxied response: {:?}", proxied);
+
+                        // Here we create a single response Frame for each query.
+                        // Frames can be created from iterators of fields using [`IntoFrame`].
+                        Ok(backend::DataResponse::new(
+                            ref_id.clone(),
+                            vec![[
+                                // Fields can be created from iterators of a variety of
+                                // relevant datatypes.
+                                [
+                                    Utc.ymd(2021, 1, 1).and_hms(12, 0, 0),
+                                    Utc.ymd(2021, 1, 1).and_hms(12, 0, 1),
+                                    Utc.ymd(2021, 1, 1).and_hms(12, 0, 2),
+                                ]
+                                .into_field("time"),
+                                [1_u32, 2, 3].into_field("x"),
+                                ["a", "b", "c"].into_field("y"),
                             ]
-                            .into_field("time"),
-                            [1_u32, 2, 3].into_field("x"),
-                            ["a", "b", "c"].into_field("y"),
-                        ]
-                        .into_frame("foo")
-                        .check()
-                        .map_err(|source| QueryError {
-                            ref_id: x.ref_id,
-                            source,
-                        })?],
-                    ))
+                            .into_frame("foo")
+                            .check()
+                            .map_err(|source| QueryError {
+                                source: SDKError::Data(source),
+                                ref_id,
+                            })?],
+                        ))
+                    }
                 })
                 .collect::<FuturesOrdered<_>>(),
         )
@@ -195,7 +229,7 @@ impl backend::ResourceService for MyPluginService {
         &self,
         r: backend::CallResourceRequest,
     ) -> Result<(Self::InitialResponse, Self::Stream), Self::Error> {
-        let count = Arc::clone(&self.0);
+        let count = Arc::clone(&self.counter);
         let response_and_stream = match r.request.uri().path() {
             // Just send back a single response.
             "/echo" => Ok((

crates/grafana-plugin-sdk/src/backend/data.rs

@@ -1,25 +1,33 @@
 //! SDK types and traits relevant to plugins that query data.
 use std::{collections::HashMap, pin::Pin, time::Duration};
 
+use bytes::Bytes;
 use futures_core::Stream;
 use futures_util::StreamExt;
-use serde::de::DeserializeOwned;
+use serde::{de::DeserializeOwned, Serialize};
+use tonic::transport::Endpoint;
 
 use crate::{
     backend::{self, ConvertFromError, TimeRange},
     data, pluginv2,
 };
 
+use super::ConvertToError;
+
+/// Transport metadata sent alongside a query.
+///
+/// This should be passed into [`DataClient::query_data`] when forwarding
+/// a query to the same Grafana instance as the query originated.
+#[derive(Clone, Debug)]
+pub struct TransportMetadata(tonic::metadata::MetadataMap);
+
 /// A request for data made by Grafana.
 ///
 /// Details of the request source can be found in `plugin_context`,
 /// while the actual plugins themselves are in `queries`.
 #[derive(Debug)]
 #[non_exhaustive]
-pub struct QueryDataRequest<Q>
-where
-    Q: DeserializeOwned,
-{
+pub struct QueryDataRequest<Q> {
     /// Details of the plugin instance from which the request originated.
     ///
     /// If the request originates from a datasource instance, this will
@@ -34,38 +42,51 @@ where
     /// the query to the frontend; this should be included in the corresponding
     /// `DataResponse` for each query.
     pub queries: Vec<DataQuery<Q>>,
+
+    grpc_meta: TransportMetadata,
+}
+
+impl<Q> QueryDataRequest<Q> {
+    /// Get the transport metadata in this request.
+    ///
+    /// This is required when using the [`DataClient`] to query
+    /// other Grafana datasources.
+    pub fn transport_metadata(&self) -> &TransportMetadata {
+        &self.grpc_meta
+    }
 }
 
-impl<Q> TryFrom<pluginv2::QueryDataRequest> for QueryDataRequest<Q>
+impl<Q> TryFrom<tonic::Request<pluginv2::QueryDataRequest>> for QueryDataRequest<Q>
 where
     Q: DeserializeOwned,
 {
     type Error = ConvertFromError;
-    fn try_from(other: pluginv2::QueryDataRequest) -> Result<Self, Self::Error> {
+    fn try_from(other: tonic::Request<pluginv2::QueryDataRequest>) -> Result<Self, Self::Error> {
+        // Clone is required until https://github.com/hyperium/tonic/pull/1118 is released.
+        let grpc_meta = other.metadata().clone();
+        let request = other.into_inner();
         Ok(Self {
-            plugin_context: other
+            plugin_context: request
                 .plugin_context
                 .ok_or(ConvertFromError::MissingPluginContext)
                 .and_then(TryInto::try_into)?,
-            headers: other.headers,
-            queries: other
+            headers: request.headers,
+            queries: request
                 .queries
                 .into_iter()
                 .map(DataQuery::try_from)
                 .collect::<Result<Vec<_>, _>>()?,
+            grpc_meta: TransportMetadata(grpc_meta),
         })
     }
 }
 
 /// A query made by Grafana to the plugin as part of a [`QueryDataRequest`].
 ///
 /// The `query` field contains any fields set by the plugin's UI.
-#[derive(Debug)]
 #[non_exhaustive]
-pub struct DataQuery<Q>
-where
-    Q: DeserializeOwned,
-{
+#[derive(Clone, Debug)]
+pub struct DataQuery<Q> {
     /// The unique identifier of the query, set by the frontend call.
     ///
     /// This should be included in the corresponding [`DataResponse`].
@@ -111,6 +132,46 @@ where
     }
 }
 
+impl<Q> TryFrom<DataQuery<Q>> for pluginv2::DataQuery
+where
+    Q: Serialize,
+{
+    type Error = ConvertToError;
+
+    fn try_from(other: DataQuery<Q>) -> Result<Self, Self::Error> {
+        Ok(Self {
+            ref_id: other.ref_id,
+            max_data_points: other.max_data_points,
+            interval_ms: other.interval.as_millis() as i64,
+            time_range: Some(other.time_range.into()),
+            json: serde_json::to_vec(&other.query)
+                .map_err(|err| ConvertToError::InvalidJson { err })?,
+            query_type: other.query_type,
+        })
+    }
+}
+
+#[derive(Debug)]
+enum DataResponseFrames {
+    Deserialized(Vec<data::Frame>),
+    Serialized(Result<Vec<Vec<u8>>, data::Error>),
+}
+
+impl DataResponseFrames {
+    fn into_serialized(self) -> Result<Vec<Vec<u8>>, data::Error> {
+        match self {
+            Self::Deserialized(frames) => to_arrow(
+                frames
+                    .iter()
+                    .map(|x| x.check())
+                    .collect::<Result<Vec<_>, _>>()?,
+                &None,
+            ),
+            Self::Serialized(x) => x,
+        }
+    }
+}
+
 /// The results from a [`DataQuery`].
 #[derive(Debug)]
 pub struct DataResponse {
@@ -121,7 +182,7 @@ pub struct DataResponse {
     ref_id: String,
 
     /// The data returned from the query.
-    frames: Result<Vec<Vec<u8>>, data::Error>,
+    frames: DataResponseFrames,
 }
 
 impl DataResponse {
@@ -130,7 +191,7 @@ impl DataResponse {
     pub fn new(ref_id: String, frames: Vec<data::CheckedFrame<'_>>) -> Self {
         Self {
             ref_id: ref_id.clone(),
-            frames: to_arrow(frames, &Some(ref_id)),
+            frames: DataResponseFrames::Serialized(to_arrow(frames, &Some(ref_id))),
         }
     }
 }
@@ -286,15 +347,15 @@ where
         let responses = DataService::query_data(
             self,
             request
-                .into_inner()
+                // .into_inner()
                 .try_into()
                 .map_err(ConvertFromError::into_tonic_status)?,
         )
         .await
         .map(|resp| match resp {
             Ok(x) => {
                 let ref_id = x.ref_id;
-                x.frames.map_or_else(
+                x.frames.into_serialized().map_or_else(
                     |e| {
                         (
                             ref_id.clone(),
@@ -336,3 +397,89 @@ where
         }))
     }
 }
+
+/// A client for querying data from Grafana.
+///
+/// This can be used by plugins which need to query data from
+/// a different datasource of the same Grafana instance.
+#[derive(Debug, Clone)]
+pub struct DataClient<T = tonic::transport::Channel> {
+    inner: pluginv2::data_client::DataClient<T>,
+}
+
+impl DataClient<tonic::transport::Channel> {
+    /// Create a new DataClient to connect to the given endpoint.
+    ///
+    /// This constructor uses the default [`tonic::transport::Channel`] as the
+    /// transport. Use [`DataClient::with_channel`] to provide your own channel.
+    pub fn new(url: impl Into<Bytes>) -> Result<Self, tonic::transport::Error> {
+        let endpoint = Endpoint::from_shared(url)?;
+        let channel = endpoint.connect_lazy();
+        Ok(Self {
+            inner: pluginv2::data_client::DataClient::new(channel),
+        })
+    }
+}
+
+/// Errors which can occur when querying data.
+#[derive(Clone, Debug, thiserror::Error)]
+#[error("Error querying data")]
+pub struct QueryDataError;
+
+impl<T> DataClient<T> {
+    /// Query for data from a Grafana datasource.
+    pub async fn query_data<Q>(
+        &mut self,
+        queries: Vec<DataQuery<Q>>,
+        upstream_metadata: &TransportMetadata,
+    ) -> Result<HashMap<String, DataResponse>, QueryDataError>
+    where
+        Q: Serialize,
+        T: tonic::client::GrpcService<tonic::body::BoxBody> + Clone + Send + Sync + 'static,
+        T::ResponseBody: tonic::codegen::Body<Data = tonic::codegen::Bytes> + Send + Sync + 'static,
+        <T::ResponseBody as tonic::codegen::Body>::Error: Into<tonic::codegen::StdError> + Send,
+    {
+        let queries: Vec<_> = queries
+            .into_iter()
+            // TODO: add enum variant
+            .map(|q| q.try_into().map_err(|_| QueryDataError))
+            .collect::<Result<_, _>>()?;
+        let query_data_request = pluginv2::QueryDataRequest {
+            plugin_context: None,
+            headers: Default::default(),
+            queries,
+        };
+        let mut request = tonic::Request::new(query_data_request);
+        for kv in upstream_metadata.0.iter() {
+            match kv {
+                tonic::metadata::KeyAndValueRef::Ascii(k, v) => {
+                    request.metadata_mut().insert(k, v.clone());
+                }
+                tonic::metadata::KeyAndValueRef::Binary(k, v) => {
+                    request.metadata_mut().insert_bin(k, v.clone());
+                }
+            }
+        }
+        let responses = self
+            .inner
+            .query_data(request)
+            .await
+            // TODO: add enum variant
+            .map_err(|_| QueryDataError)?
+            .into_inner()
+            .responses
+            .into_iter()
+            .map(|(k, v)| {
+                let frames = v
+                    .frames
+                    .into_iter()
+                    .map(|f| data::Frame::from_arrow(f).map_err(|_| QueryDataError))
+                    .collect::<Result<Vec<_>, _>>()
+                    .map(DataResponseFrames::Deserialized)
+                    .unwrap();
+                Ok((k.clone(), DataResponse { ref_id: k, frames }))
+            })
+            .collect::<Result<_, _>>()?;
+        Ok(responses)
+    }
+}

crates/grafana-plugin-sdk/src/backend/mod.rs

@@ -159,7 +159,8 @@ mod stream;
 mod tracing_fmt;
 
 pub use data::{
-    BoxDataResponseStream, DataQuery, DataQueryError, DataResponse, DataService, QueryDataRequest,
+    BoxDataResponseStream, DataClient, DataQuery, DataQueryError, DataResponse, DataService,
+    QueryDataError, QueryDataRequest, TransportMetadata,
 };
 pub use diagnostics::{
     CheckHealthRequest, CheckHealthResponse, CollectMetricsRequest, CollectMetricsResponse,
@@ -701,6 +702,15 @@ impl From<pluginv2::TimeRange> for TimeRange {
     }
 }
 
+impl From<TimeRange> for pluginv2::TimeRange {
+    fn from(other: TimeRange) -> Self {
+        Self {
+            from_epoch_ms: other.from.timestamp_millis(),
+            to_epoch_ms: other.to.timestamp_millis(),
+        }
+    }
+}
+
 /// A role within Grafana.
 #[derive(Clone, Copy, Debug)]
 #[non_exhaustive]

crates/grafana-plugin-sdk/src/data/error.rs

@@ -3,15 +3,15 @@ use arrow2::datatypes::DataType;
 use itertools::Itertools;
 use thiserror::Error;
 
-use super::frame::to_arrow;
+use super::frame::arrow;
 
 /// Errors that can occur when interacting with the Grafana plugin SDK.
 #[derive(Debug, Error)]
 #[non_exhaustive]
 pub enum Error {
     /// An error has occurred when serializing to Arrow IPC format.
     #[error("Arrow serialization error: {0}")]
-    ArrowSerialization(#[from] to_arrow::Error),
+    ArrowSerialization(#[from] arrow::Error),
 
     /// There is a datatype mismatch in a field.
     ///

crates/grafana-plugin-sdk/src/data/field.rs

@@ -5,7 +5,7 @@ use std::{
 };
 
 use arrow2::{
-    array::Array,
+    array::{new_empty_array, Array},
     datatypes::{DataType, Field as ArrowField, TimeUnit},
 };
 use serde::{Deserialize, Serialize};
@@ -70,6 +70,27 @@ impl Field {
         })
     }
 
+    pub(crate) fn from_arrow(f: ArrowField) -> Result<Self, serde_json::Error> {
+        let data_type: TypeInfoType = f.data_type().try_into().unwrap();
+        Ok(Self {
+            name: f.name,
+            labels: f
+                .metadata
+                .get("labels")
+                .and_then(|c| serde_json::from_str(c).ok())
+                .unwrap_or_default(),
+            config: f
+                .metadata
+                .get("config")
+                .and_then(|c| serde_json::from_str(c).ok()),
+            values: data_type.empty_array(),
+            type_info: TypeInfo {
+                frame: data_type,
+                nullable: Some(f.is_nullable),
+            },
+        })
+    }
+
     /// Return a new field with the given name.
     ///
     /// # Example
@@ -476,6 +497,26 @@ impl TryFrom<&DataType> for TypeInfoType {
     }
 }
 
+impl From<&TypeInfoType> for DataType {
+    fn from(other: &TypeInfoType) -> Self {
+        match other {
+            TypeInfoType::Int8 => Self::Int8,
+            TypeInfoType::Int16 => Self::Int16,
+            TypeInfoType::Int32 => Self::Int32,
+            TypeInfoType::Int64 => Self::Int64,
+            TypeInfoType::UInt8 => Self::UInt8,
+            TypeInfoType::UInt16 => Self::UInt16,
+            TypeInfoType::UInt32 => Self::UInt32,
+            TypeInfoType::UInt64 => Self::UInt64,
+            TypeInfoType::Float32 => Self::Float32,
+            TypeInfoType::Float64 => Self::Float64,
+            TypeInfoType::String => Self::Utf8,
+            TypeInfoType::Bool => Self::Boolean,
+            TypeInfoType::Time => Self::Timestamp(TimeUnit::Nanosecond, None),
+        }
+    }
+}
+
 impl From<TypeInfoType> for DataType {
     fn from(other: TypeInfoType) -> Self {
         match other {
@@ -515,6 +556,11 @@ impl TypeInfoType {
             Self::Time => SimpleType::Time,
         }
     }
+
+    #[must_use]
+    pub(crate) fn empty_array(&self) -> Box<dyn Array> {
+        new_empty_array(self.into())
+    }
 }
 
 /// The 'simple' type of this data.

crates/grafana-plugin-sdk/src/data/frame/to_arrow.rs crates/grafana-plugin-sdk/src/data/frame/arrow.rs

@@ -1,11 +1,20 @@
 //! Conversion of [`Frame`][crate::data::Frame]s to the Arrow IPC format.
 use std::{collections::BTreeMap, sync::Arc};
 
-use arrow2::{chunk::Chunk, datatypes::Schema, io::ipc::write::FileWriter};
+use arrow2::{
+    chunk::Chunk,
+    datatypes::Schema,
+    io::ipc::{
+        read::{read_file_metadata, FileReader},
+        write::FileWriter,
+    },
+};
 use thiserror::Error;
 
 use crate::data::{field::Field, frame::CheckedFrame};
 
+use super::Frame;
+
 /// Errors occurring when serializing a [`Frame`] to the Arrow IPC format.
 #[derive(Debug, Error)]
 #[non_exhaustive]
@@ -83,6 +92,36 @@ impl CheckedFrame<'_> {
     }
 }
 
+impl Frame {
+    pub(crate) fn from_arrow(data: Vec<u8>) -> Result<Self, Error> {
+        let mut cursor = std::io::Cursor::new(data);
+        let metadata = read_file_metadata(&mut cursor).unwrap();
+        let arrow_schema = metadata.schema.clone();
+        let mut frame = Frame {
+            name: arrow_schema.metadata.get("name").unwrap().to_string(),
+            ref_id: arrow_schema.metadata.get("refId").map(Clone::clone),
+            meta: arrow_schema
+                .metadata
+                .get("meta")
+                .and_then(|m| serde_json::from_str(m).ok()),
+            fields: vec![],
+        };
+
+        let fields = arrow_schema
+            .fields
+            .into_iter()
+            .map(Field::from_arrow)
+            .collect::<Result<Vec<_>, _>>()
+            .unwrap();
+        let reader = FileReader::new(cursor, metadata, None, None);
+        // TODO: are these chunks of columns? need to actually read them in...
+        let chunks = reader.collect::<Result<Vec<_>, _>>().unwrap();
+
+        frame.fields = fields;
+        Ok(frame)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use crate::data;

crates/grafana-plugin-sdk/src/data/frame/mod.rs

@@ -11,9 +11,9 @@ use crate::{
     live::Channel,
 };
 
+pub(crate) mod arrow;
 pub(self) mod de;
 pub(self) mod ser;
-pub(crate) mod to_arrow;
 
 use ser::{SerializableField, SerializableFrame, SerializableFrameData, SerializableFrameSchema};