From 46d23ffd5649097909feb6ef6cc48b4a12804eff Mon Sep 17 00:00:00 2001
From: fwang12 <fwang12@ebay.com>
Date: Fri, 24 Mar 2023 15:37:51 +0800
Subject: [PATCH] [KYUUBI #4594] Support PUT method in REST client

### _Why are the changes needed?_

We have api that need `PUT` method.
<img width="802" alt="image" src="https://user-images.githubusercontent.com/6757692/227441606-100bb7ff-2a51-400d-a1c9-dcd4217da305.png">

### _How was this patch tested?_
- [ ] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [ ] Add screenshots for manual tests if appropriate

- [ ] [Run test](https://kyuubi.readthedocs.io/en/master/develop_tools/testing.html#running-tests) locally before make a pull request

Closes #4594 from turboFei/support_put.

Closes #4594

26fab9e26 [fwang12] add put in rest client

Authored-by: fwang12 <fwang12@ebay.com>
Signed-off-by: fwang12 <fwang12@ebay.com>
---
 .../org/apache/kyuubi/client/IRestClient.java     |  4 ++++
 .../java/org/apache/kyuubi/client/RestClient.java | 15 +++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/IRestClient.java b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/IRestClient.java
index 66897df54..0eaffebd2 100644
--- a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/IRestClient.java
+++ b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/IRestClient.java
@@ -32,6 +32,10 @@ public interface IRestClient extends AutoCloseable {
 
   String post(String path, String body, String authHeader);
 
+  <T> T put(String path, String body, Class<T> type, String authHeader);
+
+  String put(String path, String body, String authHeader);
+
   <T> T delete(String path, Map<String, Object> params, Class<T> type, String authHeader);
 
   String delete(String path, Map<String, Object> params, String authHeader);
diff --git a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/RestClient.java b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/RestClient.java
index 20e57b963..6447d5477 100644
--- a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/RestClient.java
+++ b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/RestClient.java
@@ -126,6 +126,21 @@ public class RestClient implements IRestClient {
     return JsonUtils.fromJson(responseJson, type);
   }
 
+  @Override
+  public <T> T put(String path, String body, Class<T> type, String authHeader) {
+    String responseJson = put(path, body, authHeader);
+    return JsonUtils.fromJson(responseJson, type);
+  }
+
+  @Override
+  public String put(String path, String body, String authHeader) {
+    RequestBuilder putRequestBuilder = RequestBuilder.put();
+    if (body != null) {
+      putRequestBuilder.setEntity(new StringEntity(body, StandardCharsets.UTF_8));
+    }
+    return doRequest(buildURI(path), authHeader, putRequestBuilder);
+  }
+
   @Override
   public <T> T delete(String path, Map<String, Object> params, Class<T> type, String authHeader) {
     String responseJson = delete(path, params, authHeader);