From b6bbd1d14360695621ed77e5af14446dfee372e6 Mon Sep 17 00:00:00 2001
From: Mark Grey <mgthesecond@spotify.com>
Date: Tue, 20 Feb 2024 08:13:07 -0500
Subject: [PATCH] feat: Make thrift transport configurable (#194)

* feat: make transport configurable (#188)

* implement default for HmsThriftTransport
---
 crates/catalog/hms/src/catalog.rs | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/crates/catalog/hms/src/catalog.rs b/crates/catalog/hms/src/catalog.rs
index f3eab62ea..a00aaf337 100644
--- a/crates/catalog/hms/src/catalog.rs
+++ b/crates/catalog/hms/src/catalog.rs
@@ -29,10 +29,22 @@ use std::fmt::{Debug, Formatter};
 use std::net::ToSocketAddrs;
 use typed_builder::TypedBuilder;
 
+/// Which variant of the thrift transport to communicate with HMS
+/// See: <https://github.com/apache/thrift/blob/master/doc/specs/thrift-rpc.md#framed-vs-unframed-transport>
+#[derive(Debug, Default)]
+pub enum HmsThriftTransport {
+    /// Use the framed transport
+    Framed,
+    /// Use the buffered transport (default)
+    #[default]
+    Buffered,
+}
+
 /// Hive metastore Catalog configuration.
 #[derive(Debug, TypedBuilder)]
 pub struct HmsCatalogConfig {
     address: String,
+    thrift_transport: HmsThriftTransport,
 }
 
 struct HmsClient(ThriftHiveMetastoreClient);
@@ -67,12 +79,16 @@ impl HmsCatalog {
                 )
             })?;
 
-        let client = ThriftHiveMetastoreClientBuilder::new("hms")
-            .address(address)
-            // Framed thrift rpc is not enabled by default in HMS, we use
-            // buffered instead.
-            .make_codec(volo_thrift::codec::default::DefaultMakeCodec::buffered())
-            .build();
+        let builder = ThriftHiveMetastoreClientBuilder::new("hms").address(address);
+
+        let client = match &config.thrift_transport {
+            HmsThriftTransport::Framed => builder
+                .make_codec(volo_thrift::codec::default::DefaultMakeCodec::framed())
+                .build(),
+            HmsThriftTransport::Buffered => builder
+                .make_codec(volo_thrift::codec::default::DefaultMakeCodec::buffered())
+                .build(),
+        };
 
         Ok(Self {
             config,