Jefffrey commented on code in PR #11180:
URL: https://github.com/apache/arrow-rs/pull/11180#discussion_r4130964152


##########
arrow-flight/src/bin/flight_sql_client.rs:
##########
@@ -398,12 +430,24 @@ fn setup_logging(args: LoggingArgs) -> Result<()> {
     Ok(())
 }
 
-async fn setup_client(args: ClientArgs) -> 
Result<FlightSqlServiceClient<Channel>> {
+async fn setup_client(args: &ClientArgs) -> 
Result<FlightSqlServiceClient<Channel>> {
     let port = args.port.unwrap_or(if args.tls { 443 } else { 80 });
 
     let protocol = if args.tls { "https" } else { "http" };
 
-    let mut endpoint = Endpoint::new(format!("{}://{}:{}", protocol, 
args.host, port))
+    setup_client_for_uri(args, &format!("{}://{}:{}", protocol, args.host, 
port)).await
+}
+
+/// Connect a client to `uri`, applying the headers, token, handshake and 
compression settings from
+/// `args`. TLS is used when `uri` has an `https` scheme, so that an endpoint 
location may differ
+/// from the main connection.
+async fn setup_client_for_uri(
+    args: &ClientArgs,
+    uri: &str,
+) -> Result<FlightSqlServiceClient<Channel>> {
+    let tls = uri.starts_with("https://";);

Review Comment:
   if a user specifies `tls` in the cli args, but a server responds with a 
location that is just `http`, it might seen surprising to just accept and use 
this. perhaps we should try to look only for `https` in cases where user 
specifies `tls`, and fail explicitly if one cannot be found 🤔 



##########
arrow-flight/src/bin/flight_sql_client.rs:
##########
@@ -329,18 +337,42 @@ async fn main() -> Result<()> {
 
 async fn execute_flight(
     client: &mut FlightSqlServiceClient<Channel>,
+    client_args: &ClientArgs,
     info: FlightInfo,
 ) -> Result<Vec<RecordBatch>> {
     let schema = Arc::new(Schema::try_from(info.clone()).context("valid 
schema")?);
     let mut batches = Vec::with_capacity(info.endpoint.len() + 1);
     batches.push(RecordBatch::new_empty(schema));
     info!("decoded schema");
 
+    let mut location_clients = HashMap::new();
+
     for endpoint in info.endpoint {
         let Some(ticket) = &endpoint.ticket else {
             bail!("did not get ticket");
         };
 
+        // `None` means no location was given, or only the reserved 
reuse-connection form, so
+        // the ticket is redeemed on the server that returned the `FlightInfo`.
+        let location = endpoint
+            .location
+            .iter()
+            .map(|location| location.uri.as_str())
+            .find(|uri| !uri.is_empty() && *uri != REUSE_CONNECTION_URI);
+
+        let client = match location {
+            None => &mut *client,
+            Some(uri) => match location_clients.entry(uri.to_owned()) {
+                Entry::Occupied(entry) => entry.into_mut(),
+                Entry::Vacant(entry) => {
+                    let client = setup_client_for_uri(client_args, uri)

Review Comment:
   does the spec technically require us to do authentication handshake for 
these endpoint servers? i tried taking a look but it seemed a bit vague



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to