refactor: associated QueryExec return type

Allow the return type of the QueryExec trait's query_exec() method to be
parametrised by the implementer.

This allows the trait to be reused across different data sources that
return differing concrete types.

ingester2/src/query/exec.rs

@@ -17,13 +17,15 @@ impl QueryRunner {
 
 #[async_trait]
 impl QueryExec for QueryRunner {
+    type Response = QueryResponse;
+
     async fn query_exec(
         &self,
         namespace_id: NamespaceId,
         table_id: TableId,
         columns: Vec<String>,
         span: Option<Span>,
-    ) -> Result<QueryResponse, QueryError> {
+    ) -> Result<Self::Response, QueryError> {
         let mut _span_recorder = SpanRecorder::new(span);
 
         info!(

ingester2/src/query/instrumentation.rs

@@ -4,7 +4,7 @@ use iox_time::{SystemProvider, TimeProvider};
 use metric::{DurationHistogram, Metric};
 use trace::span::Span;
 
-use super::{response::QueryResponse, QueryExec};
+use super::QueryExec;
 use crate::query::QueryError;
 
 /// An instrumentation decorator over a [`QueryExec`] implementation.
@@ -49,6 +49,8 @@ where
     T: QueryExec,
     P: TimeProvider,
 {
+    type Response = T::Response;
+
     #[inline(always)]
     async fn query_exec(
         &self,
@@ -56,7 +58,7 @@ where
         table_id: TableId,
         columns: Vec<String>,
         span: Option<Span>,
-    ) -> Result<QueryResponse, QueryError> {
+    ) -> Result<Self::Response, QueryError> {
         let t = self.time_provider.now();
 
         let res = self
@@ -83,7 +85,10 @@ mod tests {
     use metric::Attributes;
 
     use super::*;
-    use crate::query::{mock_query_exec::MockQueryExec, response::PartitionStream};
+    use crate::query::{
+        mock_query_exec::MockQueryExec,
+        response::{PartitionStream, QueryResponse},
+    };
 
     macro_rules! test_metric {
         (

ingester2/src/query/mock_query_exec.rs

@@ -19,13 +19,15 @@ impl MockQueryExec {
 
 #[async_trait]
 impl QueryExec for MockQueryExec {
+    type Response = QueryResponse;
+
     async fn query_exec(
         &self,
         _namespace_id: NamespaceId,
         _table_id: TableId,
         _columns: Vec<String>,
         _span: Option<Span>,
-    ) -> Result<QueryResponse, QueryError> {
+    ) -> Result<Self::Response, QueryError> {
         self.response
             .lock()
             .take()

ingester2/src/query/tracing.rs

@@ -4,7 +4,7 @@ use async_trait::async_trait;
 use data_types::{NamespaceId, TableId};
 use trace::span::{Span, SpanRecorder};
 
-use super::{response::QueryResponse, QueryExec};
+use super::QueryExec;
 use crate::query::QueryError;
 
 /// An tracing decorator over a [`QueryExec`] implementation.
@@ -33,6 +33,8 @@ impl<T> QueryExec for QueryExecTracing<T>
 where
     T: QueryExec,
 {
+    type Response = T::Response;
+
     #[inline(always)]
     async fn query_exec(
         &self,
@@ -40,7 +42,7 @@ where
         table_id: TableId,
         columns: Vec<String>,
         span: Option<Span>,
-    ) -> Result<QueryResponse, QueryError> {
+    ) -> Result<Self::Response, QueryError> {
         let span = span.map(|s| s.child(self.name.clone()));
         let mut recorder = SpanRecorder::new(span.clone());
 

ingester2/src/query/trait.rs

@@ -5,8 +5,6 @@ use data_types::{NamespaceId, TableId};
 use thiserror::Error;
 use trace::span::Span;
 
-use super::response::QueryResponse;
-
 #[derive(Debug, Error)]
 #[allow(missing_copy_implementations)]
 pub(crate) enum QueryError {
@@ -19,13 +17,15 @@ pub(crate) enum QueryError {
 
 #[async_trait]
 pub(crate) trait QueryExec: Send + Sync + Debug {
+    type Response: Send + Debug;
+
     async fn query_exec(
         &self,
         namespace_id: NamespaceId,
         table_id: TableId,
         columns: Vec<String>,
         span: Option<Span>,
-    ) -> Result<QueryResponse, QueryError>;
+    ) -> Result<Self::Response, QueryError>;
 }
 
 #[async_trait]
@@ -33,13 +33,15 @@ impl<T> QueryExec for Arc<T>
 where
     T: QueryExec,
 {
+    type Response = T::Response;
+
     async fn query_exec(
         &self,
         namespace_id: NamespaceId,
         table_id: TableId,
         columns: Vec<String>,
         span: Option<Span>,
-    ) -> Result<QueryResponse, QueryError> {
+    ) -> Result<Self::Response, QueryError> {
         self.deref()
             .query_exec(namespace_id, table_id, columns, span)
             .await

ingester2/src/server/grpc.rs

@@ -13,7 +13,11 @@ use generated_types::influxdata::iox::{
 use iox_catalog::interface::Catalog;
 use service_grpc_catalog::CatalogService;
 
-use crate::{dml_sink::DmlSink, init::IngesterRpcInterface, query::QueryExec};
+use crate::{
+    dml_sink::DmlSink,
+    init::IngesterRpcInterface,
+    query::{response::QueryResponse, QueryExec},
+};
 
 use self::rpc_write::RpcWrite;
 
@@ -31,7 +35,7 @@ pub(crate) struct GrpcDelegate<D, Q> {
 impl<D, Q> GrpcDelegate<D, Q>
 where
     D: DmlSink + 'static,
-    Q: QueryExec + 'static,
+    Q: QueryExec<Response = QueryResponse> + 'static,
 {
     /// Initialise a new [`GrpcDelegate`].
     pub(crate) fn new(dml_sink: Arc<D>, query_exec: Arc<Q>) -> Self {
@@ -47,7 +51,7 @@ where
 impl<D, Q> IngesterRpcInterface for GrpcDelegate<D, Q>
 where
     D: DmlSink + 'static,
-    Q: QueryExec + 'static,
+    Q: QueryExec<Response = QueryResponse> + 'static,
 {
     type CatalogHandler = CatalogService;
     type WriteHandler = RpcWrite<Arc<D>>;

ingester2/src/server/grpc/query.rs

@@ -133,7 +133,7 @@ type TonicStream<T> = Pin<Box<dyn Stream<Item = Result<T, tonic::Status>> + Send
 #[tonic::async_trait]
 impl<Q> Flight for FlightService<Q>
 where
-    Q: QueryExec + 'static,
+    Q: QueryExec<Response = QueryResponse> + 'static,
 {
     type HandshakeStream = TonicStream<HandshakeResponse>;
     type ListFlightsStream = TonicStream<FlightInfo>;