refactor

Author: fivetran-arunsuri

persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java

@@ -756,7 +756,8 @@ public PolarisPrincipalSecrets resetPrincipalSecrets(
       @Nonnull String clientId,
       long principalId,
       String customClientId,
-      String customClientSecret) {
+      String customClientSecret,
+      boolean customReset) {
     throw new UnsupportedOperationException(
         "This method is not supported for EclipseLink as metastore");
   }

persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java

@@ -780,7 +780,8 @@ public PolarisPrincipalSecrets resetPrincipalSecrets(
       @Nonnull String clientId,
       long principalId,
       String customClientId,
-      String customClientSecret) {
+      String customClientSecret,
+      boolean customReset) {
     PolarisPrincipalSecrets principalSecrets = loadPrincipalSecrets(callCtx, clientId);
 
     // should be found
@@ -803,7 +804,7 @@ public PolarisPrincipalSecrets resetPrincipalSecrets(
             principalId,
             principalSecrets.getPrincipalId());
 
-    if (customClientId != null && customClientSecret != null) {
+    if (customReset) {
       principalSecrets =
           new PolarisPrincipalSecrets(
               principalSecrets.getPrincipalId(), customClientId, customClientSecret, null);

polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisPrincipalSecrets.java

@@ -36,7 +36,7 @@ public class PolarisPrincipalSecrets {
   private final long principalId;
 
   // the client id for that principal
-  private String principalClientId;
+  private final String principalClientId;
 
   // the main secret hash for that principal
   private String mainSecret;

polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java

@@ -916,14 +916,20 @@ private void revokeGrantRecord(
         PolarisObjectMapperUtil.deserializeProperties(
             principal.getInternalProperties() == null ? "{}" : principal.getInternalProperties());
 
+    boolean customReset = customClientId != null && customClientSecret != null;
     PolarisPrincipalSecrets secrets =
         ((IntegrationPersistence) ms)
             .resetPrincipalSecrets(
-                callCtx, clientId, principalId, customClientId, customClientSecret);
+                callCtx, clientId, principalId, customClientId, customClientSecret, customReset);
 
     PolarisBaseEntity.Builder principalBuilder = new PolarisBaseEntity.Builder(principal);
     principalBuilder.internalProperties(PolarisObjectMapperUtil.serializeProperties(internalProps));
-    principalBuilder.entityVersion(principal.getEntityVersion() + 1);
+    // To avoid incrementing entity version twice
+    if (customReset) {
+      principalBuilder.entityVersion(principal.getEntityVersion());
+    } else {
+      principalBuilder.entityVersion(principal.getEntityVersion() + 1);
+    }
     ms.writeEntity(callCtx, principalBuilder.build(), true, principal);
     return (secrets == null)
         ? new PrincipalSecretsResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null)

polaris-core/src/main/java/org/apache/polaris/core/persistence/IntegrationPersistence.java

@@ -90,14 +90,16 @@ PolarisPrincipalSecrets rotatePrincipalSecrets(
    * @param principalId principal id
    * @param customClientId the principal secret's old main secret hash
    * @param customClientSecret the principal secret's old main secret hash
+   * @param customReset
    */
   @Nullable
   PolarisPrincipalSecrets resetPrincipalSecrets(
       @Nonnull PolarisCallContext callCtx,
       @Nonnull String clientId,
       long principalId,
       String customClientId,
-      String customClientSecret);
+      String customClientSecret,
+      boolean customReset);
 
   /**
    * When dropping a principal, we also need to drop the secrets of that principal

polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java

@@ -942,11 +942,18 @@ private void bootstrapPolarisService(
         PolarisObjectMapperUtil.deserializeProperties(
             principal.getInternalProperties() == null ? "{}" : principal.getInternalProperties());
 
+    boolean customReset = customClientId != null && customClientSecret != null;
+
     PolarisPrincipalSecrets secrets =
         ms.resetPrincipalSecrets(
-            callCtx, clientId, principalId, customClientId, customClientSecret);
+            callCtx, clientId, principalId, customClientId, customClientSecret, customReset);
     principalBuilder.internalProperties(PolarisObjectMapperUtil.serializeProperties(internalProps));
-    principalBuilder.entityVersion(principal.getEntityVersion() + 1);
+    // To avoid incrementing entity version twice
+    if (customReset) {
+      principalBuilder.entityVersion(principal.getEntityVersion());
+    } else {
+      principalBuilder.entityVersion(principal.getEntityVersion() + 1);
+    }
     ms.writeEntityInCurrentTxn(callCtx, principalBuilder.build(), true, principal);
     return secrets;
   }

polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TreeMapTransactionalPersistenceImpl.java

@@ -500,7 +500,8 @@ public int lookupEntityGrantRecordsVersionInCurrentTxn(
       @Nonnull String clientId,
       long principalId,
       String customClientId,
-      String customClientSecret) {
+      String customClientSecret,
+      boolean customReset) {
     PolarisPrincipalSecrets principalSecrets = this.store.getSlicePrincipalSecrets().read(clientId);
 
     // should be found
@@ -524,7 +525,7 @@ public int lookupEntityGrantRecordsVersionInCurrentTxn(
             principalSecrets.getPrincipalId());
 
     // reset the secrets
-    if (customClientId != null && customClientSecret != null) {
+    if (customReset) {
       principalSecrets =
           new PolarisPrincipalSecrets(
               principalSecrets.getPrincipalId(), customClientId, customClientSecret, null);