Add type-check to PolarisEntity subclass ctors (#2302)

XN137 · web-flow · commit ac31963c9eef · 2025-09-02T16:58:01.000-04:00
currently one can freely "cast" any `PolarisEntity` to a more specific type via their constructors. this can lead to subtle bugs like we fixed in a29f800 by adding type checks we discover a few more places where we need to be more careful about how we construct new or handle existing entities. note that we can add a check for `PolarisEntitySubType` in a followup, but it requires more fixes currently.
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/CatalogEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/CatalogEntity.java
@@ -20,6 +20,7 @@
 
 import static org.apache.polaris.core.admin.model.StorageConfigInfo.StorageTypeEnum.AZURE;
 
+import com.google.common.base.Preconditions;
 import jakarta.annotation.Nonnull;
 import jakarta.annotation.Nullable;
 import java.util.Collection;
@@ -70,9 +71,11 @@ public class CatalogEntity extends PolarisEntity implements LocationBasedEntity
 
   public CatalogEntity(PolarisBaseEntity sourceEntity) {
     super(sourceEntity);
+    Preconditions.checkState(
+        getType() == PolarisEntityType.CATALOG, "Invalid entity type: %s", getType());
   }
 
-  public static CatalogEntity of(PolarisBaseEntity sourceEntity) {
+  public static @Nullable CatalogEntity of(@Nullable PolarisBaseEntity sourceEntity) {
     if (sourceEntity != null) {
       return new CatalogEntity(sourceEntity);
     }
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/CatalogRoleEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/CatalogRoleEntity.java
@@ -18,15 +18,19 @@
  */
 package org.apache.polaris.core.entity;
 
+import com.google.common.base.Preconditions;
+import jakarta.annotation.Nullable;
 import org.apache.polaris.core.admin.model.CatalogRole;
 
 /** Wrapper for translating between the REST CatalogRole object and the base PolarisEntity type. */
 public class CatalogRoleEntity extends PolarisEntity {
   public CatalogRoleEntity(PolarisBaseEntity sourceEntity) {
     super(sourceEntity);
+    Preconditions.checkState(
+        getType() == PolarisEntityType.CATALOG_ROLE, "Invalid entity type: %s", getType());
   }
 
-  public static CatalogRoleEntity of(PolarisBaseEntity sourceEntity) {
+  public static @Nullable CatalogRoleEntity of(@Nullable PolarisBaseEntity sourceEntity) {
     if (sourceEntity != null) {
       return new CatalogRoleEntity(sourceEntity);
     }
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/NamespaceEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/NamespaceEntity.java
@@ -19,6 +19,8 @@
 package org.apache.polaris.core.entity;
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.google.common.base.Preconditions;
+import jakarta.annotation.Nullable;
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.rest.RESTUtil;
 import org.apache.polaris.core.catalog.PolarisCatalogHelpers;
@@ -33,9 +35,11 @@ public class NamespaceEntity extends PolarisEntity implements LocationBasedEntit
 
   public NamespaceEntity(PolarisBaseEntity sourceEntity) {
     super(sourceEntity);
+    Preconditions.checkState(
+        getType() == PolarisEntityType.NAMESPACE, "Invalid entity type: %s", getType());
   }
 
-  public static NamespaceEntity of(PolarisBaseEntity sourceEntity) {
+  public static @Nullable NamespaceEntity of(@Nullable PolarisBaseEntity sourceEntity) {
     if (sourceEntity != null) {
       return new NamespaceEntity(sourceEntity);
     }
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisEntity.java
@@ -22,6 +22,7 @@
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import jakarta.annotation.Nonnull;
+import jakarta.annotation.Nullable;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -151,14 +152,14 @@ public PolarisEntity(
         entityVersion);
   }
 
-  public static PolarisEntity of(PolarisBaseEntity sourceEntity) {
+  public static @Nullable PolarisEntity of(@Nullable PolarisBaseEntity sourceEntity) {
     if (sourceEntity != null) {
       return new PolarisEntity(sourceEntity);
     }
     return null;
   }
 
-  public static PolarisEntity of(EntityResult result) {
+  public static @Nullable PolarisEntity of(EntityResult result) {
     if (result.isSuccess()) {
       return new PolarisEntity(result.getEntity());
     }
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/PrincipalEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/PrincipalEntity.java
@@ -18,15 +18,19 @@
  */
 package org.apache.polaris.core.entity;
 
+import com.google.common.base.Preconditions;
+import jakarta.annotation.Nullable;
 import org.apache.polaris.core.admin.model.Principal;
 
 /** Wrapper for translating between the REST Principal object and the base PolarisEntity type. */
 public class PrincipalEntity extends PolarisEntity {
   public PrincipalEntity(PolarisBaseEntity sourceEntity) {
     super(sourceEntity);
+    Preconditions.checkState(
+        getType() == PolarisEntityType.PRINCIPAL, "Invalid entity type: %s", getType());
   }
 
-  public static PrincipalEntity of(PolarisBaseEntity sourceEntity) {
+  public static @Nullable PrincipalEntity of(@Nullable PolarisBaseEntity sourceEntity) {
     if (sourceEntity != null) {
       return new PrincipalEntity(sourceEntity);
     }
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/PrincipalRoleEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/PrincipalRoleEntity.java
@@ -18,6 +18,8 @@
  */
 package org.apache.polaris.core.entity;
 
+import com.google.common.base.Preconditions;
+import jakarta.annotation.Nullable;
 import org.apache.polaris.core.admin.model.PrincipalRole;
 import org.apache.polaris.core.entity.table.federated.FederatedEntities;
 
@@ -27,9 +29,11 @@
 public class PrincipalRoleEntity extends PolarisEntity {
   public PrincipalRoleEntity(PolarisBaseEntity sourceEntity) {
     super(sourceEntity);
+    Preconditions.checkState(
+        getType() == PolarisEntityType.PRINCIPAL_ROLE, "Invalid entity type: %s", getType());
   }
 
-  public static PrincipalRoleEntity of(PolarisBaseEntity sourceEntity) {
+  public static @Nullable PrincipalRoleEntity of(@Nullable PolarisBaseEntity sourceEntity) {
     if (sourceEntity != null) {
       return new PrincipalRoleEntity(sourceEntity);
     }
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/TaskEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/TaskEntity.java
@@ -18,6 +18,8 @@
  */
 package org.apache.polaris.core.entity;
 
+import com.google.common.base.Preconditions;
+import jakarta.annotation.Nullable;
 import org.apache.polaris.core.persistence.PolarisObjectMapperUtil;
 
 /**
@@ -27,14 +29,15 @@
 public class TaskEntity extends PolarisEntity {
   public TaskEntity(PolarisBaseEntity sourceEntity) {
     super(sourceEntity);
+    Preconditions.checkState(
+        getType() == PolarisEntityType.TASK, "Invalid entity type: %s", getType());
   }
 
-  public static TaskEntity of(PolarisBaseEntity polarisEntity) {
-    if (polarisEntity != null) {
-      return new TaskEntity(polarisEntity);
-    } else {
-      return null;
+  public static @Nullable TaskEntity of(@Nullable PolarisBaseEntity sourceEntity) {
+    if (sourceEntity != null) {
+      return new TaskEntity(sourceEntity);
     }
+    return null;
   }
 
   public <T> T readData(Class<T> klass) {
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/table/GenericTableEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/table/GenericTableEntity.java
@@ -19,6 +19,7 @@
 package org.apache.polaris.core.entity.table;
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
+import jakarta.annotation.Nullable;
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.apache.iceberg.rest.RESTUtil;
@@ -42,7 +43,7 @@ public GenericTableEntity(PolarisBaseEntity sourceEntity) {
     super(sourceEntity);
   }
 
-  public static GenericTableEntity of(PolarisBaseEntity sourceEntity) {
+  public static @Nullable GenericTableEntity of(@Nullable PolarisBaseEntity sourceEntity) {
     if (sourceEntity != null) {
       return new GenericTableEntity(sourceEntity);
     }
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/table/IcebergTableLikeEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/table/IcebergTableLikeEntity.java
@@ -19,6 +19,7 @@
 package org.apache.polaris.core.entity.table;
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
+import jakarta.annotation.Nullable;
 import java.util.Optional;
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.catalog.TableIdentifier;
@@ -48,7 +49,7 @@ public IcebergTableLikeEntity(PolarisBaseEntity sourceEntity) {
     super(sourceEntity);
   }
 
-  public static IcebergTableLikeEntity of(PolarisBaseEntity sourceEntity) {
+  public static @Nullable IcebergTableLikeEntity of(@Nullable PolarisBaseEntity sourceEntity) {
     if (sourceEntity != null) {
       return new IcebergTableLikeEntity(sourceEntity);
     }
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/table/TableLikeEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/table/TableLikeEntity.java
@@ -19,6 +19,7 @@
 package org.apache.polaris.core.entity.table;
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.google.common.base.Preconditions;
 import jakarta.annotation.Nonnull;
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.catalog.TableIdentifier;
@@ -37,6 +38,8 @@ public abstract class TableLikeEntity extends PolarisEntity implements LocationB
 
   public TableLikeEntity(@Nonnull PolarisBaseEntity sourceEntity) {
     super(sourceEntity);
+    Preconditions.checkState(
+        getType() == PolarisEntityType.TABLE_LIKE, "Invalid entity type: %s", getType());
   }
 
   @JsonIgnore
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/policy/PolicyEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/policy/PolicyEntity.java
@@ -20,6 +20,7 @@
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.google.common.base.Preconditions;
+import jakarta.annotation.Nullable;
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.rest.RESTUtil;
 import org.apache.polaris.core.entity.NamespaceEntity;
@@ -36,13 +37,14 @@ public class PolicyEntity extends PolarisEntity {
 
   PolicyEntity(PolarisBaseEntity sourceEntity) {
     super(sourceEntity);
+    Preconditions.checkState(
+        getType() == PolarisEntityType.POLICY, "Invalid entity type: %s", getType());
   }
 
-  public static PolicyEntity of(PolarisBaseEntity sourceEntity) {
+  public static @Nullable PolicyEntity of(@Nullable PolarisBaseEntity sourceEntity) {
     if (sourceEntity != null) {
       return new PolicyEntity(sourceEntity);
     }
-
     return null;
   }
 
diff --git a/runtime/service/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalog.java b/runtime/service/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalog.java
@@ -1049,6 +1049,7 @@ private void validateNoLocationOverlap(
       IcebergTableLikeEntity virtualEntity =
           IcebergTableLikeEntity.of(
               new PolarisEntity.Builder()
+                  .setType(PolarisEntityType.TABLE_LIKE)
                   .setParentId(resolvedNamespace.getLast().getId())
                   .setProperties(Map.of(PolarisEntityConstants.ENTITY_BASE_LOCATION, location))
                   .build());
diff --git a/runtime/service/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogHandler.java b/runtime/service/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogHandler.java
@@ -24,6 +24,7 @@
 import com.google.common.collect.Maps;
 import io.smallrye.common.annotation.Identifier;
 import jakarta.annotation.Nonnull;
+import jakarta.annotation.Nullable;
 import jakarta.enterprise.inject.Instance;
 import jakarta.ws.rs.core.SecurityContext;
 import java.io.Closeable;
@@ -82,7 +83,9 @@
 import org.apache.polaris.core.connection.ConnectionType;
 import org.apache.polaris.core.context.CallContext;
 import org.apache.polaris.core.entity.CatalogEntity;
+import org.apache.polaris.core.entity.PolarisEntity;
 import org.apache.polaris.core.entity.PolarisEntitySubType;
+import org.apache.polaris.core.entity.PolarisEntityType;
 import org.apache.polaris.core.entity.table.IcebergTableLikeEntity;
 import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
 import org.apache.polaris.core.persistence.PolarisResolvedPathWrapper;
@@ -582,12 +585,15 @@ public boolean sendNotification(TableIdentifier identifier, NotificationRequest
    * Fetch the metastore table entity for the given table identifier
    *
    * @param tableIdentifier The identifier of the table
-   * @return the Polaris table entity for the table
+   * @return the Polaris table entity for the table or null for external catalogs
    */
-  private IcebergTableLikeEntity getTableEntity(TableIdentifier tableIdentifier) {
+  private @Nullable IcebergTableLikeEntity getTableEntity(TableIdentifier tableIdentifier) {
     PolarisResolvedPathWrapper target = resolutionManifest.getResolvedPath(tableIdentifier);
-
-    return IcebergTableLikeEntity.of(target.getRawLeafEntity());
+    PolarisEntity rawLeafEntity = target.getRawLeafEntity();
+    if (rawLeafEntity.getType() == PolarisEntityType.TABLE_LIKE) {
+      return IcebergTableLikeEntity.of(rawLeafEntity);
+    }
+    return null; // could be an external catalog
   }
 
   public LoadTableResponse loadTable(TableIdentifier tableIdentifier, String snapshots) {

Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,7 @@`
`22`	`22`	`import com.fasterxml.jackson.annotation.JsonIgnore;`
`23`	`23`	`import com.fasterxml.jackson.annotation.JsonProperty;`
`24`	`24`	`import jakarta.annotation.Nonnull;`
	`25`	`+import jakarta.annotation.Nullable;`
`25`	`26`	`import java.util.HashMap;`
`26`	`27`	`import java.util.List;`
`27`	`28`	`import java.util.Map;`
`@@ -151,14 +152,14 @@ public PolarisEntity(`
`151`	`152`	`entityVersion);`
`152`	`153`	`}`
`153`	`154`
`154`		`- public static PolarisEntity of(PolarisBaseEntity sourceEntity) {`
	`155`	`+ public static @Nullable PolarisEntity of(@Nullable PolarisBaseEntity sourceEntity) {`
`155`	`156`	`if (sourceEntity != null) {`
`156`	`157`	`return new PolarisEntity(sourceEntity);`
`157`	`158`	`}`
`158`	`159`	`return null;`
`159`	`160`	`}`
`160`	`161`
`161`		`- public static PolarisEntity of(EntityResult result) {`
	`162`	`+ public static @Nullable PolarisEntity of(EntityResult result) {`
`162`	`163`	`if (result.isSuccess()) {`
`163`	`164`	`return new PolarisEntity(result.getEntity());`
`164`	`165`	`}`
Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,7 @@`
`19`	`19`	`package org.apache.polaris.core.entity.table;`
`20`	`20`
`21`	`21`	`import com.fasterxml.jackson.annotation.JsonIgnore;`
	`22`	`+import jakarta.annotation.Nullable;`
`22`	`23`	`import org.apache.iceberg.catalog.Namespace;`
`23`	`24`	`import org.apache.iceberg.catalog.TableIdentifier;`
`24`	`25`	`import org.apache.iceberg.rest.RESTUtil;`
`@@ -42,7 +43,7 @@ public GenericTableEntity(PolarisBaseEntity sourceEntity) {`
`42`	`43`	`super(sourceEntity);`
`43`	`44`	`}`
`44`	`45`
`45`		`- public static GenericTableEntity of(PolarisBaseEntity sourceEntity) {`
	`46`	`+ public static @Nullable GenericTableEntity of(@Nullable PolarisBaseEntity sourceEntity) {`
`46`	`47`	`if (sourceEntity != null) {`
`47`	`48`	`return new GenericTableEntity(sourceEntity);`
`48`	`49`	`}`