feat: gRPC compression support for flight CLI (#8240)

# Which issue does this PR close?
\-

# Rationale for this change
Some services support gRPC compression. Expose this to the CLI client
for:

- testing
- more efficient data transfer over slow internet connections

# What changes are included in this PR?
CLI argument wiring.

# Are these changes tested?
No automated tests. I think we can assume that the libraries we use do
what they promise to do. But I also verified that this works by
inspecting the traffic using Wireshark.

# Are there any user-facing changes?
They now have more options.

Author: crepererum

arrow-flight/Cargo.toml

@@ -70,7 +70,7 @@ tls-ring = ["tonic/tls-ring"]
 tls-webpki-roots = ["tonic/tls-webpki-roots"]
 
 # Enable CLI tools
-cli = ["arrow-array/chrono-tz", "arrow-cast/prettyprint", "tonic/tls-webpki-roots", "dep:anyhow", "dep:clap", "dep:tracing-log", "dep:tracing-subscriber", "dep:tokio"]
+cli = ["arrow-array/chrono-tz", "arrow-cast/prettyprint", "tonic/tls-webpki-roots", "tonic/gzip", "tonic/deflate", "tonic/zstd", "dep:anyhow", "dep:clap", "dep:tracing-log", "dep:tracing-subscriber", "dep:tokio"]
 
 [dev-dependencies]
 arrow-cast = { workspace = true, features = ["prettyprint"] }

arrow-flight/src/bin/flight_sql_client.rs

@@ -21,11 +21,12 @@ use anyhow::{bail, Context, Result};
 use arrow_array::{ArrayRef, Datum, RecordBatch, StringArray};
 use arrow_cast::{cast_with_options, pretty::pretty_format_batches, CastOptions};
 use arrow_flight::{
+    flight_service_client::FlightServiceClient,
     sql::{client::FlightSqlServiceClient, CommandGetDbSchemas, CommandGetTables},
     FlightInfo,
 };
 use arrow_schema::Schema;
-use clap::{Parser, Subcommand};
+use clap::{Parser, Subcommand, ValueEnum};
 use core::str;
 use futures::TryStreamExt;
 use tonic::{
@@ -53,6 +54,24 @@ pub struct LoggingArgs {
     log_verbose_count: u8,
 }
 
+/// gRPC/HTTP compression algorithms.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)]
+pub enum CompressionEncoding {
+    Gzip,
+    Deflate,
+    Zstd,
+}
+
+impl From<CompressionEncoding> for tonic::codec::CompressionEncoding {
+    fn from(encoding: CompressionEncoding) -> Self {
+        match encoding {
+            CompressionEncoding::Gzip => Self::Gzip,
+            CompressionEncoding::Deflate => Self::Deflate,
+            CompressionEncoding::Zstd => Self::Zstd,
+        }
+    }
+}
+
 #[derive(Debug, Parser)]
 struct ClientArgs {
     /// Additional headers.
@@ -96,6 +115,34 @@ struct ClientArgs {
     /// Defaults to `443` if `tls` is set, otherwise defaults to `80`.
     #[clap(long)]
     port: Option<u16>,
+
+    /// Compression accepted by the client for responses sent by the server.
+    ///
+    /// The client will send this information to the server as part of the request. The server is free to pick an
+    /// algorithm from that list or use no compression (called "identity" encoding).
+    ///
+    /// You may define multiple algorithms by using a comma-separated list.
+    #[clap(long, value_delimiter = ',')]
+    accept_compression: Vec<CompressionEncoding>,
+
+    /// Compression of requests sent by the client to the server.
+    ///
+    /// Since the client needs to decide on the compression before sending the request, there is no client<->server
+    /// negotiation. If the server does NOT support the chosen compression, it will respond with an error a la:
+    ///
+    /// ```
+    /// Ipc error: Status {
+    ///     code: Unimplemented,
+    ///     message: "Content is compressed with `zstd` which isn't supported",
+    ///     metadata: MetadataMap { headers: {"grpc-accept-encoding": "identity", ...} },
+    ///     ...
+    /// }
+    /// ```
+    ///
+    /// Based on the algorithms listed in the `grpc-accept-encoding` header, you may make a more educated guess for
+    /// your next request. Note that `identity` is a synonym for "no compression".
+    #[clap(long)]
+    send_compression: Option<CompressionEncoding>,
 }
 
 #[derive(Debug, Parser)]
@@ -365,7 +412,14 @@ async fn setup_client(args: ClientArgs) -> Result<FlightSqlServiceClient<Channel
 
     let channel = endpoint.connect().await.context("connect to endpoint")?;
 
-    let mut client = FlightSqlServiceClient::new(channel);
+    let mut client = FlightServiceClient::new(channel);
+    for encoding in args.accept_compression {
+        client = client.accept_compressed(encoding.into());
+    }
+    if let Some(encoding) = args.send_compression {
+        client = client.send_compressed(encoding.into());
+    }
+    let mut client = FlightSqlServiceClient::new_from_inner(client);
     info!("connected");
 
     for (k, v) in args.headers {