amitvijapur commented on code in PR #11180:
URL: https://github.com/apache/arrow-rs/pull/11180#discussion_r4137860241
##########
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:
Updated in 18be6d2. With `--tls`, the CLI now selects the first HTTPS
location, including one listed after plaintext locations. If there is no HTTPS
location, it can reuse the original TLS connection when the endpoint advertises
the reuse form; otherwise it returns an explicit error. Added regression tests
for these cases and preserved selection without `--tls`. The Arrow Flight crate
tests and Clippy pass.
##########
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:
I found no requirement to handshake on each endpoint server. The [Flight
authentication
section](https://arrow.apache.org/docs/format/Flight.html#authentication)
explicitly allows token-based calls without a handshake, and
`FlightSqlServiceClient::token()` is documented for sharing a bearer token with
other `DoGet` clients. In 18be6d2, new endpoint clients copy the original
client's token instead of replaying username/password. This covers both an
explicit `--token` and a token issued by the initial handshake. The integration
test uses a data server that requires that bearer token and rejects handshakes,
and runs both authentication modes. Token acceptance across the advertised
servers is still application-dependent; this does not add cookie-based session
handling. This follow-up was implemented with Codex assistance and separately
reviewed.
--
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]
