Store multiple hooks

Author: gabotechs

src/flight_service/do_get.rs

@@ -94,8 +94,10 @@ impl ArrowFlightEndpoint {
         let stage_data = once
             .get_or_try_init(|| async {
                 let proto_node = PhysicalPlanNode::try_decode(doget.plan_proto.as_ref())?;
-                let plan = proto_node.try_into_physical_plan(&ctx, &self.runtime, &codec)?;
-                let plan = (self.hooks.on_plan)(plan);
+                let mut plan = proto_node.try_into_physical_plan(&ctx, &self.runtime, &codec)?;
+                for hook in self.hooks.on_plan.iter() {
+                    plan = hook(plan)
+                }
 
                 // Initialize partition count to the number of partitions in the stage
                 let total_partitions = plan.properties().partitioning.partition_count();
@@ -232,7 +234,7 @@ mod tests {
         let plans_received = Arc::new(AtomicUsize::default());
         {
             let plans_received = Arc::clone(&plans_received);
-            endpoint.on_plan(move |plan| {
+            endpoint.add_on_plan_hook(move |plan| {
                 plans_received.fetch_add(1, Ordering::SeqCst);
                 plan
             });

src/flight_service/service.rs

@@ -17,16 +17,10 @@ use tokio::sync::OnceCell;
 use tonic::{Request, Response, Status, Streaming};
 
 #[allow(clippy::type_complexity)]
+#[derive(Default)]
 pub(super) struct ArrowFlightEndpointHooks {
-    pub(super) on_plan: Arc<dyn Fn(Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> + Sync + Send>,
-}
-
-impl Default for ArrowFlightEndpointHooks {
-    fn default() -> Self {
-        Self {
-            on_plan: Arc::new(|plan| plan),
-        }
-    }
+    pub(super) on_plan:
+        Vec<Arc<dyn Fn(Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> + Sync + Send>>,
 }
 
 pub struct ArrowFlightEndpoint {
@@ -54,11 +48,11 @@ impl ArrowFlightEndpoint {
     /// The callback takes the plan and returns another plan that must be either the same,
     /// or equivalent in terms of execution. Mutating the plan by adding nodes or removing them
     /// will make the query blow up in unexpected ways.
-    pub fn on_plan(
+    pub fn add_on_plan_hook(
         &mut self,
-        cbk: impl Fn(Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> + Sync + Send + 'static,
+        hook: impl Fn(Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> + Sync + Send + 'static,
     ) {
-        self.hooks.on_plan = Arc::new(cbk);
+        self.hooks.on_plan.push(Arc::new(hook));
     }
 }