From 83dda05e8b752ea2f360bd9aa5bd0fe9e06d3f2b Mon Sep 17 00:00:00 2001
From: "Paul J. Davis" <paul.joseph.davis@gmail.com>
Date: Wed, 4 Dec 2024 08:55:58 -0600
Subject: [PATCH] fix(flightsql): remove Any encoding of DoPutUpdateResult
 (#6825)

---
 arrow-flight/src/sql/client.rs | 12 ++++++------
 arrow-flight/src/sql/server.rs |  8 ++++----
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/arrow-flight/src/sql/client.rs b/arrow-flight/src/sql/client.rs
index e45e505b2b61..a6e228737b3f 100644
--- a/arrow-flight/src/sql/client.rs
+++ b/arrow-flight/src/sql/client.rs
@@ -228,8 +228,8 @@ impl FlightSqlServiceClient<Channel> {
             .await
             .map_err(status_to_arrow_error)?
             .unwrap();
-        let any = Any::decode(&*result.app_metadata).map_err(decode_error_to_arrow_error)?;
-        let result: DoPutUpdateResult = any.unpack()?.unwrap();
+        let result: DoPutUpdateResult =
+            Message::decode(&*result.app_metadata).map_err(decode_error_to_arrow_error)?;
         Ok(result.record_count)
     }
 
@@ -274,8 +274,8 @@ impl FlightSqlServiceClient<Channel> {
             .await
             .map_err(status_to_arrow_error)?
             .unwrap();
-        let any = Any::decode(&*result.app_metadata).map_err(decode_error_to_arrow_error)?;
-        let result: DoPutUpdateResult = any.unpack()?.unwrap();
+        let result: DoPutUpdateResult =
+            Message::decode(&*result.app_metadata).map_err(decode_error_to_arrow_error)?;
         Ok(result.record_count)
     }
 
@@ -593,8 +593,8 @@ impl PreparedStatement<Channel> {
             .await
             .map_err(status_to_arrow_error)?
             .unwrap();
-        let any = Any::decode(&*result.app_metadata).map_err(decode_error_to_arrow_error)?;
-        let result: DoPutUpdateResult = any.unpack()?.unwrap();
+        let result: DoPutUpdateResult =
+            Message::decode(&*result.app_metadata).map_err(decode_error_to_arrow_error)?;
         Ok(result.record_count)
     }
 
diff --git a/arrow-flight/src/sql/server.rs b/arrow-flight/src/sql/server.rs
index 6b9befa63600..8ab8a16dbb50 100644
--- a/arrow-flight/src/sql/server.rs
+++ b/arrow-flight/src/sql/server.rs
@@ -719,7 +719,7 @@ where
                 let record_count = self.do_put_statement_update(command, request).await?;
                 let result = DoPutUpdateResult { record_count };
                 let output = futures::stream::iter(vec![Ok(PutResult {
-                    app_metadata: result.as_any().encode_to_vec().into(),
+                    app_metadata: result.encode_to_vec().into(),
                 })]);
                 Ok(Response::new(Box::pin(output)))
             }
@@ -727,7 +727,7 @@ where
                 let record_count = self.do_put_statement_ingest(command, request).await?;
                 let result = DoPutUpdateResult { record_count };
                 let output = futures::stream::iter(vec![Ok(PutResult {
-                    app_metadata: result.as_any().encode_to_vec().into(),
+                    app_metadata: result.encode_to_vec().into(),
                 })]);
                 Ok(Response::new(Box::pin(output)))
             }
@@ -744,7 +744,7 @@ where
                 let record_count = self.do_put_substrait_plan(command, request).await?;
                 let result = DoPutUpdateResult { record_count };
                 let output = futures::stream::iter(vec![Ok(PutResult {
-                    app_metadata: result.as_any().encode_to_vec().into(),
+                    app_metadata: result.encode_to_vec().into(),
                 })]);
                 Ok(Response::new(Box::pin(output)))
             }
@@ -754,7 +754,7 @@ where
                     .await?;
                 let result = DoPutUpdateResult { record_count };
                 let output = futures::stream::iter(vec![Ok(PutResult {
-                    app_metadata: result.as_any().encode_to_vec().into(),
+                    app_metadata: result.encode_to_vec().into(),
                 })]);
                 Ok(Response::new(Box::pin(output)))
             }