From 66926c4b674dc548dc564458792270d488912088 Mon Sep 17 00:00:00 2001
From: dtb <2011xuesong@gmail.com>
Date: Wed, 10 Apr 2024 23:32:52 +0800
Subject: [PATCH] add unit test and use old function

---
 .../bdp/waggledance/server/FederatedHMSHandler.java | 10 +---------
 .../waggledance/server/FederatedHMSHandlerTest.java | 13 +++++++++++++
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/waggle-dance-core/src/main/java/com/hotels/bdp/waggledance/server/FederatedHMSHandler.java b/waggle-dance-core/src/main/java/com/hotels/bdp/waggledance/server/FederatedHMSHandler.java
index ec79ee39d..d3439069d 100644
--- a/waggle-dance-core/src/main/java/com/hotels/bdp/waggledance/server/FederatedHMSHandler.java
+++ b/waggle-dance-core/src/main/java/com/hotels/bdp/waggledance/server/FederatedHMSHandler.java
@@ -1341,16 +1341,8 @@ public GrantRevokePrivilegeResponse grant_revoke_privileges(GrantRevokePrivilege
   public GrantRevokePrivilegeResponse refresh_privileges(HiveObjectRef hiveObjectRef, String authorizer,
                                                          GrantRevokePrivilegeRequest grantRevokePrivilegeRequest) throws MetaException, TException {
     DatabaseMapping databaseMapping = checkWritePermissions(hiveObjectRef.getDbName());
-    PrivilegeBag privilegeBag = grantRevokePrivilegeRequest.getPrivileges();
-    if (privilegeBag != null && privilegeBag.getPrivileges() != null) {
-      for (HiveObjectPrivilege hiveObjectPrivilege : privilegeBag.getPrivileges()) {
-        if (hiveObjectPrivilege.getHiveObject() != null) {
-          databaseMapping.transformInboundHiveObjectRef(hiveObjectPrivilege.getHiveObject());
-        }
-      }
-    }
     return databaseMapping.getClient().refresh_privileges(databaseMapping.transformInboundHiveObjectRef(hiveObjectRef),
-            authorizer, grantRevokePrivilegeRequest);
+            authorizer, databaseMapping.transformInboundGrantRevokePrivilegesRequest(grantRevokePrivilegeRequest));
   }
 
   private DatabaseMapping checkWritePermissionsForPrivileges(PrivilegeBag privileges) throws NoSuchObjectException {
diff --git a/waggle-dance-core/src/test/java/com/hotels/bdp/waggledance/server/FederatedHMSHandlerTest.java b/waggle-dance-core/src/test/java/com/hotels/bdp/waggledance/server/FederatedHMSHandlerTest.java
index c7b2029d0..bfafce6a5 100644
--- a/waggle-dance-core/src/test/java/com/hotels/bdp/waggledance/server/FederatedHMSHandlerTest.java
+++ b/waggle-dance-core/src/test/java/com/hotels/bdp/waggledance/server/FederatedHMSHandlerTest.java
@@ -119,6 +119,7 @@
 import org.apache.hadoop.hive.metastore.api.HeartbeatTxnRangeResponse;
 import org.apache.hadoop.hive.metastore.api.HiveObjectPrivilege;
 import org.apache.hadoop.hive.metastore.api.HiveObjectRef;
+import org.apache.hadoop.hive.metastore.api.HiveObjectType;
 import org.apache.hadoop.hive.metastore.api.ISchema;
 import org.apache.hadoop.hive.metastore.api.ISchemaName;
 import org.apache.hadoop.hive.metastore.api.LockComponent;
@@ -1990,9 +1991,21 @@ public void refresh_privileges() throws TException {
     GrantRevokePrivilegeRequest grantRevokePrivilegeRequest = new GrantRevokePrivilegeRequest();
     HiveObjectRef hiveObjectRef = new HiveObjectRef();
     hiveObjectRef.setDbName(DB_P);
+    PrivilegeBag privileges = new PrivilegeBag();
+    List<HiveObjectPrivilege> privilegesList = new ArrayList<>();
+    HiveObjectPrivilege hiveObjectPrivilege = new HiveObjectPrivilege();
+    HiveObjectRef hor = new HiveObjectRef();
+    hor.setDbName(DB_P);
+    hor.setObjectType(HiveObjectType.DATABASE);
+    hor.setObjectName(DB_P);
+    hiveObjectPrivilege.setHiveObject(hor);
+    privilegesList.add(hiveObjectPrivilege);
+    privileges.setPrivileges(privilegesList);
+    grantRevokePrivilegeRequest.setPrivileges(privileges);
     GrantRevokePrivilegeResponse grantRevokePrivilegeResponse = new GrantRevokePrivilegeResponse();
 
     when(primaryMapping.transformInboundHiveObjectRef(hiveObjectRef)).thenReturn(hiveObjectRef);
+    when(primaryMapping.transformInboundGrantRevokePrivilegesRequest(grantRevokePrivilegeRequest)).thenReturn(grantRevokePrivilegeRequest);
     when(primaryClient.refresh_privileges(hiveObjectRef, "dummy", grantRevokePrivilegeRequest)).thenReturn(grantRevokePrivilegeResponse);
     GrantRevokePrivilegeResponse result = handler.refresh_privileges(hiveObjectRef, "dummy", grantRevokePrivilegeRequest);
     assertThat(result, is(grantRevokePrivilegeResponse));