Add PolarisDiagnostics field to TransactionWorkspaceMetaStoreManager (#2359)

the ultimate goal is removing the `PolarisCallContext` parameter from every
`PolarisMetaStoreManager` interface method, so we make steps towards
reducing its usage first.

Author: XN137

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

@@ -27,6 +27,7 @@
 import java.util.Optional;
 import java.util.Set;
 import org.apache.polaris.core.PolarisCallContext;
+import org.apache.polaris.core.PolarisDiagnostics;
 import org.apache.polaris.core.entity.LocationBasedEntity;
 import org.apache.polaris.core.entity.PolarisBaseEntity;
 import org.apache.polaris.core.entity.PolarisEntity;
@@ -72,6 +73,7 @@
  * be reused between requests.
  */
 public class TransactionWorkspaceMetaStoreManager implements PolarisMetaStoreManager {
+  private final PolarisDiagnostics diagnostics;
   private final PolarisMetaStoreManager delegate;
 
   // TODO: If we want to support the semantic of opening a transaction in which multiple
@@ -83,7 +85,9 @@ public class TransactionWorkspaceMetaStoreManager implements PolarisMetaStoreMan
   // pendingUpdates that ultimately need to be applied in order within the real MetaStoreManager.
   private final List<EntityWithPath> pendingUpdates = new ArrayList<>();
 
-  public TransactionWorkspaceMetaStoreManager(PolarisMetaStoreManager delegate) {
+  public TransactionWorkspaceMetaStoreManager(
+      PolarisDiagnostics diagnostics, PolarisMetaStoreManager delegate) {
+    this.diagnostics = diagnostics;
     this.delegate = delegate;
   }
 
@@ -93,15 +97,13 @@ public List<EntityWithPath> getPendingUpdates() {
 
   @Override
   public BaseResult bootstrapPolarisService(@Nonnull PolarisCallContext callCtx) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "bootstrapPolarisService");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "bootstrapPolarisService");
     return null;
   }
 
   @Override
   public BaseResult purge(@Nonnull PolarisCallContext callCtx) {
-    callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "purge");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "purge");
     return null;
   }
 
@@ -112,7 +114,7 @@ public EntityResult readEntityByName(
       @Nonnull PolarisEntityType entityType,
       @Nonnull PolarisEntitySubType entitySubType,
       @Nonnull String name) {
-    callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "readEntityByName");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "readEntityByName");
     return null;
   }
 
@@ -123,31 +125,27 @@ public ListEntitiesResult listEntities(
       @Nonnull PolarisEntityType entityType,
       @Nonnull PolarisEntitySubType entitySubType,
       @Nonnull PageToken pageToken) {
-    callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "listEntities");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "listEntities");
     return null;
   }
 
   @Override
   public GenerateEntityIdResult generateNewEntityId(@Nonnull PolarisCallContext callCtx) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "generateNewEntityId");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "generateNewEntityId");
     return null;
   }
 
   @Override
   public CreatePrincipalResult createPrincipal(
       @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity principal) {
-    callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "createPrincipal");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "createPrincipal");
     return null;
   }
 
   @Override
   public PrincipalSecretsResult loadPrincipalSecrets(
       @Nonnull PolarisCallContext callCtx, @Nonnull String clientId) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "loadPrincipalSecrets");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadPrincipalSecrets");
     return null;
   }
 
@@ -158,9 +156,7 @@ public PrincipalSecretsResult rotatePrincipalSecrets(
       long principalId,
       boolean reset,
       @Nonnull String oldSecretHash) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "rotatePrincipalSecrets");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "rotatePrincipalSecrets");
     return null;
   }
 
@@ -169,7 +165,7 @@ public CreateCatalogResult createCatalog(
       @Nonnull PolarisCallContext callCtx,
       @Nonnull PolarisBaseEntity catalog,
       @Nonnull List<PolarisEntityCore> principalRoles) {
-    callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "createCatalog");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "createCatalog");
     return null;
   }
 
@@ -178,9 +174,7 @@ public EntityResult createEntityIfNotExists(
       @Nonnull PolarisCallContext callCtx,
       @Nullable List<PolarisEntityCore> catalogPath,
       @Nonnull PolarisBaseEntity entity) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "createEntityIfNotExists");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "createEntityIfNotExists");
     return null;
   }
 
@@ -189,9 +183,7 @@ public EntitiesResult createEntitiesIfNotExist(
       @Nonnull PolarisCallContext callCtx,
       @Nullable List<PolarisEntityCore> catalogPath,
       @Nonnull List<? extends PolarisBaseEntity> entities) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "createEntitiesIfNotExist");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "createEntitiesIfNotExist");
     return null;
   }
 
@@ -207,9 +199,8 @@ public EntityResult updateEntityPropertiesIfNotChanged(
   @Override
   public EntitiesResult updateEntitiesPropertiesIfNotChanged(
       @Nonnull PolarisCallContext callCtx, @Nonnull List<EntityWithPath> entities) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "updateEntitiesPropertiesIfNotChanged");
+    diagnostics.fail(
+        "illegal_method_in_transaction_workspace", "updateEntitiesPropertiesIfNotChanged");
     return null;
   }
 
@@ -220,7 +211,7 @@ public EntityResult renameEntity(
       @Nonnull PolarisBaseEntity entityToRename,
       @Nullable List<PolarisEntityCore> newCatalogPath,
       @Nonnull PolarisEntity renamedEntity) {
-    callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "renameEntity");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "renameEntity");
     return null;
   }
 
@@ -231,7 +222,7 @@ public DropEntityResult dropEntityIfExists(
       @Nonnull PolarisBaseEntity entityToDrop,
       @Nullable Map<String, String> cleanupProperties,
       boolean cleanup) {
-    callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "dropEntityIfExists");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "dropEntityIfExists");
     return null;
   }
 
@@ -241,9 +232,7 @@ public PrivilegeResult grantUsageOnRoleToGrantee(
       @Nullable PolarisEntityCore catalog,
       @Nonnull PolarisEntityCore role,
       @Nonnull PolarisEntityCore grantee) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "grantUsageOnRoleToGrantee");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "grantUsageOnRoleToGrantee");
     return null;
   }
 
@@ -253,9 +242,7 @@ public PrivilegeResult revokeUsageOnRoleFromGrantee(
       @Nullable PolarisEntityCore catalog,
       @Nonnull PolarisEntityCore role,
       @Nonnull PolarisEntityCore grantee) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "revokeUsageOnRoleFromGrantee");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "revokeUsageOnRoleFromGrantee");
     return null;
   }
 
@@ -266,9 +253,7 @@ public PrivilegeResult grantPrivilegeOnSecurableToRole(
       @Nullable List<PolarisEntityCore> catalogPath,
       @Nonnull PolarisEntityCore securable,
       @Nonnull PolarisPrivilege privilege) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "grantPrivilegeOnSecurableToRole");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "grantPrivilegeOnSecurableToRole");
     return null;
   }
 
@@ -279,36 +264,29 @@ public PrivilegeResult revokePrivilegeOnSecurableFromRole(
       @Nullable List<PolarisEntityCore> catalogPath,
       @Nonnull PolarisEntityCore securable,
       @Nonnull PolarisPrivilege privilege) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "revokePrivilegeOnSecurableFromRole");
+    diagnostics.fail(
+        "illegal_method_in_transaction_workspace", "revokePrivilegeOnSecurableFromRole");
     return null;
   }
 
   @Override
   public @Nonnull LoadGrantsResult loadGrantsOnSecurable(
       @Nonnull PolarisCallContext callCtx, PolarisEntityCore securable) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "loadGrantsOnSecurable");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadGrantsOnSecurable");
     return null;
   }
 
   @Override
   public @Nonnull LoadGrantsResult loadGrantsToGrantee(
       @Nonnull PolarisCallContext callCtx, PolarisEntityCore grantee) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "loadGrantsToGrantee");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadGrantsToGrantee");
     return null;
   }
 
   @Override
   public ChangeTrackingResult loadEntitiesChangeTracking(
       @Nonnull PolarisCallContext callCtx, @Nonnull List<PolarisEntityId> entityIds) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "loadEntitiesChangeTracking");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadEntitiesChangeTracking");
     return null;
   }
 
@@ -318,14 +296,14 @@ public EntityResult loadEntity(
       long entityCatalogId,
       long entityId,
       PolarisEntityType entityType) {
-    callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "loadEntity");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadEntity");
     return null;
   }
 
   @Override
   public EntitiesResult loadTasks(
       @Nonnull PolarisCallContext callCtx, String executorId, PageToken pageToken) {
-    callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "loadTasks");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadTasks");
     return null;
   }
 
@@ -354,9 +332,7 @@ public ResolvedEntityResult loadResolvedEntityById(
       long entityCatalogId,
       long entityId,
       PolarisEntityType entityType) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "loadResolvedEntityById");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadResolvedEntityById");
     return null;
   }
 
@@ -367,9 +343,7 @@ public ResolvedEntityResult loadResolvedEntityByName(
       long parentId,
       @Nonnull PolarisEntityType entityType,
       @Nonnull String entityName) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "loadResolvedEntityByName");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadResolvedEntityByName");
     return null;
   }
 
@@ -381,9 +355,7 @@ public ResolvedEntityResult refreshResolvedEntity(
       @Nonnull PolarisEntityType entityType,
       long entityCatalogId,
       long entityId) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "refreshResolvedEntity");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "refreshResolvedEntity");
     return null;
   }
 
@@ -392,9 +364,7 @@ public ResolvedEntityResult refreshResolvedEntity(
   public <T extends PolarisEntity & LocationBasedEntity>
       Optional<Optional<String>> hasOverlappingSiblings(
           @Nonnull PolarisCallContext callContext, T entity) {
-    callContext
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "hasOverlappingSiblings");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "hasOverlappingSiblings");
     return Optional.empty();
   }
 
@@ -406,9 +376,7 @@ Optional<Optional<String>> hasOverlappingSiblings(
       @Nonnull List<PolarisEntityCore> policyCatalogPath,
       @Nonnull PolicyEntity policy,
       Map<String, String> parameters) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "attachPolicyToEntity");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "attachPolicyToEntity");
     return null;
   }
 
@@ -419,18 +387,14 @@ Optional<Optional<String>> hasOverlappingSiblings(
       @Nonnull PolarisEntityCore target,
       @Nonnull List<PolarisEntityCore> policyCatalogPath,
       @Nonnull PolicyEntity policy) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "detachPolicyFromEntity");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "detachPolicyFromEntity");
     return null;
   }
 
   @Override
   public @Nonnull LoadPolicyMappingsResult loadPoliciesOnEntity(
       @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore target) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "loadPoliciesOnEntity");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadPoliciesOnEntity");
     return null;
   }
 
@@ -439,9 +403,7 @@ Optional<Optional<String>> hasOverlappingSiblings(
       @Nonnull PolarisCallContext callCtx,
       @Nonnull PolarisEntityCore target,
       @Nonnull PolicyType policyType) {
-    callCtx
-        .getDiagServices()
-        .fail("illegal_method_in_transaction_workspace", "loadPoliciesOnEntityByType");
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadPoliciesOnEntityByType");
     return null;
   }
 }

runtime/service/src/main/java/org/apache/polaris/service/catalog/common/CatalogHandler.java

@@ -59,6 +59,7 @@ public abstract class CatalogHandler {
   protected final String catalogName;
   protected final PolarisAuthorizer authorizer;
 
+  protected final PolarisDiagnostics diagnostics;
   protected final CallContext callContext;
   protected final PolarisPrincipal polarisPrincipal;
   protected final SecurityContext securityContext;
@@ -70,12 +71,12 @@ public CatalogHandler(
       String catalogName,
       PolarisAuthorizer authorizer) {
     this.callContext = callContext;
+    this.diagnostics = callContext.getPolarisCallContext().getDiagServices();
     this.resolutionManifestFactory = resolutionManifestFactory;
     this.catalogName = catalogName;
-    PolarisDiagnostics diagServices = callContext.getPolarisCallContext().getDiagServices();
-    diagServices.checkNotNull(securityContext, "null_security_context");
-    diagServices.checkNotNull(securityContext.getUserPrincipal(), "null_user_principal");
-    diagServices.check(
+    diagnostics.checkNotNull(securityContext, "null_security_context");
+    diagnostics.checkNotNull(securityContext.getUserPrincipal(), "null_user_principal");
+    diagnostics.check(
         securityContext.getUserPrincipal() instanceof PolarisPrincipal,
         "invalid_principal_type",
         "Principal must be a PolarisPrincipal");

runtime/service/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogHandler.java

@@ -685,10 +685,7 @@ public Optional<LoadTableResponse> loadTableWithAccessDelegationIfStale(
     }
 
     PolarisResolvedPathWrapper catalogPath = resolutionManifest.getResolvedReferenceCatalogEntity();
-    callContext
-        .getPolarisCallContext()
-        .getDiagServices()
-        .checkNotNull(catalogPath, "No catalog available for loadTable request");
+    diagnostics.checkNotNull(catalogPath, "No catalog available for loadTable request");
     CatalogEntity catalogEntity = CatalogEntity.of(catalogPath.getRawLeafEntity());
     LOGGER.info("Catalog type: {}", catalogEntity.getCatalogType());
     LOGGER.info(
@@ -921,7 +918,7 @@ public void commitTransaction(CommitTransactionRequest commitTransactionRequest)
     // only go into an in-memory collection that we can commit as a single atomic unit after all
     // validations.
     TransactionWorkspaceMetaStoreManager transactionMetaStoreManager =
-        new TransactionWorkspaceMetaStoreManager(metaStoreManager);
+        new TransactionWorkspaceMetaStoreManager(diagnostics, metaStoreManager);
     ((IcebergCatalog) baseCatalog).setMetaStoreManager(transactionMetaStoreManager);
 
     commitTransactionRequest.tableChanges().stream()