WIP - compatibility with GAT in backend::DataService

Means we don't need to clone so many things or connect
to the database as many times.

Author: sd2k

backend/src/data.rs

@@ -1,14 +1,12 @@
-use std::{collections::HashMap, sync::Arc};
+use std::sync::Arc;
 
 use futures_util::stream::FuturesOrdered;
 
-use grafana_plugin_sdk::backend;
-use tokio::sync::RwLock;
-use tokio_postgres::Client;
+use grafana_plugin_sdk::backend::{self, DataSourceInstanceSettings};
 
 use crate::{
-    path::{self, PathDisplay, QueryId},
-    queries::{Query, SelectStatement, TailTarget},
+    path::{PathDisplay, QueryId},
+    queries::{Query, TailTarget},
     rows_to_frame, Error, MaterializePlugin,
 };
 
@@ -31,32 +29,37 @@ impl backend::DataQueryError for QueryError {
 // Unfortunately this has to take all of its arguments by value until we have
 // GATs, since the `DataService::Stream` associated type can't contain references.
 // Ideally we'd just borrow the query/uid etc but it's really not a big deal.
-async fn query_data_single(
-    client: &Client,
-    uid: &str,
-    query: backend::DataQuery<Query>,
-    queries: Arc<RwLock<HashMap<path::QueryId, SelectStatement>>>,
-) -> Result<backend::DataResponse, Error> {
-    let q = query.query;
-    let target = q.as_tail()?;
-    let rows = target.select_all(client).await?;
-    let mut frame = rows_to_frame(&rows);
+impl MaterializePlugin {
+    async fn query_data_single(
+        &self,
+        datasource_instance_settings: Arc<DataSourceInstanceSettings>,
+        query: backend::DataQuery<Query>,
+    ) -> Result<backend::DataResponse, Error> {
+        let q = query.query;
+        let client = self.get_client(&datasource_instance_settings).await?;
+        let target = q.as_tail()?;
+        let rows = target.select_all(&client).await?;
+        let mut frame = rows_to_frame(&rows);
 
-    if let TailTarget::Select { statement } = target {
-        let query_id = QueryId::from_statement(statement);
-        queries.write().await.insert(query_id, statement.clone());
-    }
+        if let TailTarget::Select { statement } = target {
+            let query_id = QueryId::from_statement(statement);
+            self.sql_queries
+                .write()
+                .await
+                .insert(query_id, statement.clone());
+        }
 
-    let path = q.to_path();
-    // Set the channel of the frame, indicating to Grafana that it should switch to
-    // streaming.
-    let channel = format!("ds/{}/{}", uid, path)
-        .parse()
-        .map_err(Error::CreatingChannel)?;
-    frame.set_channel(channel);
-    let frame = frame.check()?;
+        let path = q.to_path();
+        // Set the channel of the frame, indicating to Grafana that it should switch to
+        // streaming.
+        let channel = format!("ds/{}/{}", datasource_instance_settings.uid, path)
+            .parse()
+            .map_err(Error::CreatingChannel)?;
+        frame.set_channel(channel);
+        let frame = frame.check()?;
 
-    Ok(backend::DataResponse::new(query.ref_id, vec![frame]))
+        Ok(backend::DataResponse::new(query.ref_id, vec![frame]))
+    }
 }
 
 #[backend::async_trait]
@@ -69,36 +72,27 @@ impl backend::DataService for MaterializePlugin {
         &self,
         request: backend::QueryDataRequest<Self::Query>,
     ) -> Self::Stream<'_> {
-        let datasource_settings = request
-            .plugin_context
-            .datasource_instance_settings
-            .clone()
-            .ok_or(Error::MissingDatasource)
-            .unwrap();
-        // let clients: Vec<_> = request
-        //     .queries
-        //     .iter()
-        //     .map(|_| self.get_client(&datasource_settings))
-        //     .collect::<FuturesUnordered<_>>()
-        //     .collect()
-        //     .await;
-        let client = Arc::new(self.get_client(&datasource_settings).await.unwrap());
-        let queries = self.sql_queries.clone();
+        let datasource_settings = Arc::new(
+            request
+                .plugin_context
+                .datasource_instance_settings
+                .clone()
+                .ok_or(Error::MissingDatasource)
+                .unwrap(),
+        );
         Box::pin(
             request
                 .queries
                 .into_iter()
                 .map(|x| {
-                    let client = client.clone();
-                    let queries = queries.clone();
+                    let datasource_settings = datasource_settings.clone();
                     let ref_id = x.ref_id.clone();
-                    let uid = datasource_settings.uid.clone();
-                    async move {
+                    async {
                         // let client = client.map_err(|source| QueryError {
                         //     ref_id: ref_id.clone(),
                         //     source,
-                        // })?;
-                        query_data_single(&client, &uid, x, queries)
+                        self.query_data_single(datasource_settings, x)
+                            // })?;
                             .await
                             .map_err(|source| QueryError { ref_id, source })
                     }