fix(graph, node): add option to authenticate Flight service requests

Author: isum

graph/src/amp/client/flight_client.rs

@@ -33,6 +33,7 @@ use crate::{
 /// using the Apache Arrow Flight protocol.
 pub struct FlightClient {
     channel: Channel,
+    auth_token: Option<String>,
 }
 
 impl FlightClient {
@@ -56,16 +57,27 @@ impl FlightClient {
 
         Ok(Self {
             channel: endpoint.connect().await.map_err(Error::Connection)?,
+            auth_token: None,
         })
     }
 
+    /// Sets the authentication token for requests to the Amp server.
+    pub fn set_auth_token(&mut self, auth_token: impl Into<String>) {
+        self.auth_token = Some(auth_token.into());
+    }
+
     fn raw_client(&self) -> FlightSqlServiceClient<Channel> {
         let channel = self.channel.cheap_clone();
         let client = FlightServiceClient::new(channel)
             .max_encoding_message_size(256 * 1024 * 1024)
             .max_decoding_message_size(256 * 1024 * 1024);
 
-        FlightSqlServiceClient::new_from_inner(client)
+        let mut client = FlightSqlServiceClient::new_from_inner(client);
+        if let Some(auth_token) = &self.auth_token {
+            client.set_token(auth_token.clone());
+        }
+
+        client
     }
 }
 

graph/src/env/amp.rs

@@ -24,6 +24,11 @@ pub struct AmpEnv {
     ///
     /// Defaults to `600` seconds.
     pub query_retry_max_delay: Duration,
+
+    /// Token used to authenticate Amp Flight gRPC service requests.
+    ///
+    /// Defaults to `None`.
+    pub flight_service_token: Option<String>,
 }
 
 impl AmpEnv {
@@ -60,6 +65,12 @@ impl AmpEnv {
                 .amp_query_retry_max_delay_seconds
                 .map(Duration::from_secs)
                 .unwrap_or(Self::DEFAULT_QUERY_RETRY_MAX_DELAY),
+            flight_service_token: raw_env.amp_flight_service_token.as_ref().and_then(|value| {
+                if value.is_empty() {
+                    return None;
+                }
+                Some(value.to_string())
+            }),
         }
     }
 }

graph/src/env/mod.rs

@@ -561,6 +561,8 @@ struct Inner {
     amp_query_retry_min_delay_seconds: Option<u64>,
     #[envconfig(from = "GRAPH_AMP_QUERY_RETRY_MAX_DELAY_SECONDS")]
     amp_query_retry_max_delay_seconds: Option<u64>,
+    #[envconfig(from = "GRAPH_AMP_FLIGHT_SERVICE_TOKEN")]
+    amp_flight_service_token: Option<String>,
 }
 
 #[derive(Clone, Debug)]

node/src/main.rs

@@ -385,12 +385,15 @@ async fn main_inner() {
                     .parse()
                     .expect("Invalid Amp Flight service address");
 
-                let amp_client = Arc::new(
-                    amp::FlightClient::new(addr)
-                        .await
-                        .expect("Failed to connect to Amp Flight service"),
-                );
+                let mut amp_client = amp::FlightClient::new(addr)
+                    .await
+                    .expect("Failed to connect to Amp Flight service");
+
+                if let Some(auth_token) = &env_vars.amp.flight_service_token {
+                    amp_client.set_auth_token(auth_token);
+                }
 
+                let amp_client = Arc::new(amp_client);
                 let amp_instance_manager = graph_core::amp_subgraph::Manager::new(
                     &logger_factory,
                     metrics_registry.cheap_clone(),

node/src/manager/commands/run.rs

@@ -156,12 +156,15 @@ pub async fn run(
                 .parse()
                 .expect("Invalid Amp Flight service address");
 
-            let amp_client = Arc::new(
-                amp::FlightClient::new(addr)
-                    .await
-                    .expect("Failed to connect to Amp Flight service"),
-            );
+            let mut amp_client = amp::FlightClient::new(addr)
+                .await
+                .expect("Failed to connect to Amp Flight service");
+
+            if let Some(auth_token) = &env_vars.amp.flight_service_token {
+                amp_client.set_auth_token(auth_token);
+            }
 
+            let amp_client = Arc::new(amp_client);
             let amp_instance_manager = graph_core::amp_subgraph::Manager::new(
                 &logger_factory,
                 metrics_registry.cheap_clone(),