add customization point for ctx

Author: robtandy

src/bin/distributed-datafusion.rs

@@ -33,7 +33,7 @@ async fn main() -> Result<()> {
 
     match args.mode.as_str() {
         "proxy" => {
-            let service = DDProxyService::new(new_friendly_name()?, args.port).await?;
+            let service = DDProxyService::new(new_friendly_name()?, args.port, None).await?;
             service.serve().await?;
         }
         "worker" => {

src/lib.rs

@@ -24,6 +24,7 @@ pub use proto::generated::protobuf;
 
 pub mod analyze;
 pub mod codec;
+pub mod ctx_customizer;
 pub mod explain;
 pub mod flight;
 pub mod friendly;

src/proxy_service.rs

@@ -40,6 +40,7 @@ use tokio::{
 use tonic::{async_trait, transport::Server, Request, Response, Status};
 
 use crate::{
+    ctx_customizer::CtxCustomizer,
     flight::{FlightSqlHandler, FlightSqlServ},
     logging::{debug, info, trace},
     planning::{add_ctx_extentions, get_ctx},
@@ -57,10 +58,16 @@ pub struct DDProxyHandler {
     pub host: Host,
 
     pub planner: QueryPlanner,
+
+    pub ctx_customizer: Option<Arc<dyn CtxCustomizer + Send + Sync>>,
 }
 
 impl DDProxyHandler {
-    pub fn new(name: String, addr: String) -> Self {
+    pub fn new(
+        name: String,
+        addr: String,
+        ctx_customizer: Option<Arc<dyn CtxCustomizer + Send + Sync>>,
+    ) -> Self {
         // call this function to bootstrap the worker discovery mechanism
         get_worker_addresses().expect("Could not get worker addresses upon startup");
 
@@ -71,6 +78,7 @@ impl DDProxyHandler {
         Self {
             host: host.clone(),
             planner: QueryPlanner::new(),
+            ctx_customizer,
         }
     }
 
@@ -118,6 +126,11 @@ impl DDProxyHandler {
         add_ctx_extentions(&mut ctx, &self.host, &query_id, stage_id, addrs, vec![])
             .map_err(|e| Status::internal(format!("Could not add context extensions {e:?}")))?;
 
+        if let Some(ref c) = self.ctx_customizer {
+            c.customize(&mut ctx)
+                .map_err(|e| Status::internal(format!("Could not customize context {e:?}")))?;
+        }
+
         // TODO: revisit this to allow for consuming a partitular partition
         trace!("calling execute plan");
         let partition = 0;
@@ -278,7 +291,11 @@ pub struct DDProxyService {
 }
 
 impl DDProxyService {
-    pub async fn new(name: String, port: usize) -> Result<Self> {
+    pub async fn new(
+        name: String,
+        port: usize,
+        ctx_customizer: Option<Arc<dyn CtxCustomizer + Send + Sync>>,
+    ) -> Result<Self> {
         debug!("Creating DDProxyService!");
 
         let (all_done_tx, all_done_rx) = channel(1);
@@ -290,7 +307,7 @@ impl DDProxyService {
 
         info!("DDProxyService bound to {addr}");
 
-        let handler = Arc::new(DDProxyHandler::new(name, addr.clone()));
+        let handler = Arc::new(DDProxyHandler::new(name, addr.clone(), ctx_customizer));
 
         Ok(Self {
             listener,