Add support for federated principal and role with block for manual role assignment (#1353)

* Add support for federated principal and role with block for manual role assignment

* Update spec to distinguish federated and non-federated entities

* Changed builder to allow setting federated status twice

* Revert spec changes - add 'federated' property back to Principal entity

* Fixed builder to remove federated property

* Removed unnecessary openapi config flags

* Fix compilation issue in test

* Remove federated flag from principal entity

* Fixed builder oversight

* Fix compilation failures and rebase on main

Author: collado-mike

api/iceberg-service/build.gradle.kts

@@ -66,7 +66,6 @@ openApiGenerate {
   configOptions.put("useBeanValidation", "false")
   configOptions.put("sourceFolder", "src/main/java")
   configOptions.put("useJakartaEe", "true")
-  openapiNormalizer.put("REFACTOR_ALLOF_WITH_PROPERTIES_ONLY", "true")
   additionalProperties.put("apiNamePrefix", "IcebergRest")
   additionalProperties.put("apiNameSuffix", "")
   additionalProperties.put("metricsPrefix", "polaris")

api/polaris-catalog-service/build.gradle.kts

@@ -96,7 +96,6 @@ openApiGenerate {
   configOptions.put("generateBuilders", "true")
   configOptions.put("generateConstructorWithAllArgs", "true")
   configOptions.put("openApiNullable", "false")
-  openapiNormalizer.put("REFACTOR_ALLOF_WITH_PROPERTIES_ONLY", "true")
   additionalProperties.put("apiNamePrefix", "PolarisCatalog")
   additionalProperties.put("apiNameSuffix", "")
   additionalProperties.put("metricsPrefix", "polaris")

integration-tests/src/main/java/org/apache/polaris/service/it/test/PolarisManagementServiceIntegrationTest.java

@@ -18,6 +18,7 @@
  */
 package org.apache.polaris.service.it.test;
 
+import static javax.ws.rs.core.Response.Status.CREATED;
 import static javax.ws.rs.core.Response.Status.FORBIDDEN;
 import static org.apache.polaris.service.it.env.PolarisClient.polarisClient;
 import static org.apache.polaris.service.it.test.PolarisApplicationIntegrationTest.PRINCIPAL_ROLE_ALL;
@@ -872,6 +873,27 @@ public void testCreatePrincipalAndRotateCredentials() {
     // rotation that makes the old secret fall off retention.
   }
 
+  @Test
+  public void testCreateFederatedPrincipalRoleSucceeds() {
+    // Create a federated Principal Role
+    PrincipalRole federatedPrincipalRole =
+        new PrincipalRole(
+            client.newEntityName("federatedRole"),
+            true,
+            Map.of(),
+            Instant.now().toEpochMilli(),
+            Instant.now().toEpochMilli(),
+            1);
+
+    // Attempt to create the federated Principal using the managementApi
+    try (Response createResponse =
+        managementApi
+            .request("v1/principal-roles")
+            .post(Entity.json(new CreatePrincipalRoleRequest(federatedPrincipalRole)))) {
+      assertThat(createResponse).returns(CREATED.getStatusCode(), Response::getStatus);
+    }
+  }
+
   @Test
   public void testCreateListUpdateAndDeletePrincipal() {
     Principal principal =
@@ -1023,7 +1045,7 @@ public void testGetPrincipalWithInvalidName() {
   public void testCreateListUpdateAndDeletePrincipalRole() {
     PrincipalRole principalRole =
         new PrincipalRole(
-            client.newEntityName("myprincipalrole"), Map.of("custom-tag", "foo"), 0L, 0L, 1);
+            client.newEntityName("myprincipalrole"), false, Map.of("custom-tag", "foo"), 0L, 0L, 1);
     managementApi.createPrincipalRole(principalRole);
 
     // Second attempt to create the same entity should fail with CONFLICT.
@@ -1115,7 +1137,7 @@ public void testCreateListUpdateAndDeletePrincipalRole() {
   public void testCreatePrincipalRoleInvalidName() {
     String goodName = RandomStringUtils.random(MAX_IDENTIFIER_LENGTH, true, true);
     PrincipalRole principalRole =
-        new PrincipalRole(goodName, Map.of("custom-tag", "good_principal_role"), 0L, 0L, 1);
+        new PrincipalRole(goodName, false, Map.of("custom-tag", "good_principal_role"), 0L, 0L, 1);
     managementApi.createPrincipalRole(principalRole);
 
     String longInvalidName = RandomStringUtils.random(MAX_IDENTIFIER_LENGTH + 1, true, true);
@@ -1131,7 +1153,12 @@ public void testCreatePrincipalRoleInvalidName() {
     for (String invalidPrincipalRoleName : invalidPrincipalRoleNames) {
       principalRole =
           new PrincipalRole(
-              invalidPrincipalRoleName, Map.of("custom-tag", "bad_principal_role"), 0L, 0L, 1);
+              invalidPrincipalRoleName,
+              false,
+              Map.of("custom-tag", "bad_principal_role"),
+              0L,
+              0L,
+              1);
 
       try (Response response =
           managementApi
@@ -2154,7 +2181,12 @@ public void testCreateAndUpdatePrincipalRoleWithReservedProperties() {
 
     PrincipalRole badPrincipalRole =
         new PrincipalRole(
-            client.newEntityName("myprincipalrole"), Map.of("polaris.reserved", "foo"), 0L, 0L, 1);
+            client.newEntityName("myprincipalrole"),
+            false,
+            Map.of("polaris.reserved", "foo"),
+            0L,
+            0L,
+            1);
     try (Response response =
         managementApi
             .request("v1/principal-roles")
@@ -2165,7 +2197,12 @@ public void testCreateAndUpdatePrincipalRoleWithReservedProperties() {
 
     PrincipalRole goodPrincipalRole =
         new PrincipalRole(
-            client.newEntityName("myprincipalrole"), Map.of("not.reserved", "foo"), 0L, 0L, 1);
+            client.newEntityName("myprincipalrole"),
+            false,
+            Map.of("not.reserved", "foo"),
+            0L,
+            0L,
+            1);
     try (Response response =
         managementApi
             .request("v1/principal-roles")

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

@@ -19,6 +19,7 @@
 package org.apache.polaris.core.entity;
 
 import org.apache.polaris.core.admin.model.PrincipalRole;
+import org.apache.polaris.core.entity.table.federated.FederatedEntities;
 
 /**
  * Wrapper for translating between the REST PrincipalRole object and the base PolarisEntity type.
@@ -38,19 +39,19 @@ public static PrincipalRoleEntity of(PolarisBaseEntity sourceEntity) {
   public static PrincipalRoleEntity fromPrincipalRole(PrincipalRole principalRole) {
     return new Builder()
         .setName(principalRole.getName())
+        .setFederated(principalRole.getFederated())
         .setProperties(principalRole.getProperties())
         .build();
   }
 
   public PrincipalRole asPrincipalRole() {
-    PrincipalRole principalRole =
-        new PrincipalRole(
-            getName(),
-            getPropertiesAsMap(),
-            getCreateTimestamp(),
-            getLastUpdateTimestamp(),
-            getEntityVersion());
-    return principalRole;
+    return new PrincipalRole(
+        getName(),
+        FederatedEntities.isFederated(this),
+        getPropertiesAsMap(),
+        getCreateTimestamp(),
+        getLastUpdateTimestamp(),
+        getEntityVersion());
   }
 
   public static class Builder extends PolarisEntity.BaseBuilder<PrincipalRoleEntity, Builder> {
@@ -65,6 +66,15 @@ public Builder(PrincipalRoleEntity original) {
       super(original);
     }
 
+    public Builder setFederated(Boolean isFederated) {
+      if (isFederated != null && isFederated) {
+        internalProperties.put(FederatedEntities.FEDERATED_ENTITY, "true");
+      } else {
+        internalProperties.remove(FederatedEntities.FEDERATED_ENTITY);
+      }
+      return this;
+    }
+
     @Override
     public PrincipalRoleEntity build() {
       return new PrincipalRoleEntity(buildBase());

polaris-core/src/main/java/org/apache/polaris/core/entity/table/federated/FederatedEntities.java

@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.polaris.core.entity.table.federated;
+
+import java.util.Optional;
+import org.apache.polaris.core.entity.PolarisBaseEntity;
+
+public final class FederatedEntities {
+
+  public static final String FEDERATED_ENTITY = "federated";
+
+  public static boolean isFederated(PolarisBaseEntity entity) {
+    return Optional.ofNullable(entity.getInternalPropertiesAsMap())
+        .map(map -> Boolean.parseBoolean(map.get(FEDERATED_ENTITY)))
+        .orElse(false);
+  }
+
+  private FederatedEntities() {}
+}

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/ManagementServiceTest.java

@@ -22,9 +22,14 @@
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import jakarta.ws.rs.core.Response;
+import jakarta.ws.rs.core.SecurityContext;
+import java.security.Principal;
 import java.time.Clock;
+import java.time.Instant;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import org.apache.iceberg.exceptions.ValidationException;
 import org.apache.polaris.core.PolarisCallContext;
 import org.apache.polaris.core.admin.model.AwsStorageConfigInfo;
 import org.apache.polaris.core.admin.model.Catalog;
@@ -34,8 +39,20 @@
 import org.apache.polaris.core.admin.model.PolarisCatalog;
 import org.apache.polaris.core.admin.model.StorageConfigInfo;
 import org.apache.polaris.core.admin.model.UpdateCatalogRequest;
+import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal;
+import org.apache.polaris.core.auth.PolarisAuthorizerImpl;
 import org.apache.polaris.core.context.CallContext;
+import org.apache.polaris.core.context.RealmContext;
+import org.apache.polaris.core.entity.PrincipalEntity;
+import org.apache.polaris.core.entity.PrincipalRoleEntity;
+import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
+import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
+import org.apache.polaris.core.persistence.dao.entity.EntityResult;
+import org.apache.polaris.core.secrets.UnsafeInMemorySecretsManager;
 import org.apache.polaris.service.TestServices;
+import org.apache.polaris.service.admin.PolarisAdminService;
+import org.apache.polaris.service.config.DefaultConfigurationStore;
+import org.apache.polaris.service.config.ReservedProperties;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
@@ -158,4 +175,105 @@ public void testUpdateCatalogWithDisallowedStorageConfig() {
         .isInstanceOf(IllegalArgumentException.class)
         .hasMessage("Unsupported storage type: FILE");
   }
+
+  private PolarisMetaStoreManager setupMetaStoreManager() {
+    MetaStoreManagerFactory metaStoreManagerFactory = services.metaStoreManagerFactory();
+    RealmContext realmContext = services.realmContext();
+    return metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext);
+  }
+
+  private PolarisCallContext setupCallContext(PolarisMetaStoreManager metaStoreManager) {
+    MetaStoreManagerFactory metaStoreManagerFactory = services.metaStoreManagerFactory();
+    RealmContext realmContext = services.realmContext();
+    return new PolarisCallContext(
+        metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(),
+        services.polarisDiagnostics());
+  }
+
+  private PolarisAdminService setupPolarisAdminService(
+      PolarisMetaStoreManager metaStoreManager, PolarisCallContext callContext) {
+    RealmContext realmContext = services.realmContext();
+    return new PolarisAdminService(
+        CallContext.of(realmContext, callContext),
+        services.entityManagerFactory().getOrCreateEntityManager(realmContext),
+        metaStoreManager,
+        new UnsafeInMemorySecretsManager(),
+        new SecurityContext() {
+          @Override
+          public Principal getUserPrincipal() {
+            return new AuthenticatedPolarisPrincipal(
+                new PrincipalEntity.Builder().setName("root").build(), Set.of("service_admin"));
+          }
+
+          @Override
+          public boolean isUserInRole(String role) {
+            return true;
+          }
+
+          @Override
+          public boolean isSecure() {
+            return false;
+          }
+
+          @Override
+          public String getAuthenticationScheme() {
+            return "";
+          }
+        },
+        new PolarisAuthorizerImpl(new DefaultConfigurationStore(Map.of())),
+        new ReservedProperties() {
+          @Override
+          public List<String> prefixes() {
+            return List.of();
+          }
+
+          @Override
+          public Set<String> allowlist() {
+            return Set.of();
+          }
+        });
+  }
+
+  private PrincipalEntity createPrincipal(
+      PolarisMetaStoreManager metaStoreManager, PolarisCallContext callContext, String name) {
+    return new PrincipalEntity.Builder()
+        .setName(name)
+        .setCreateTimestamp(Instant.now().toEpochMilli())
+        .setId(metaStoreManager.generateNewEntityId(callContext).getId())
+        .build();
+  }
+
+  private PrincipalRoleEntity createRole(
+      PolarisMetaStoreManager metaStoreManager,
+      PolarisCallContext callContext,
+      String name,
+      boolean isFederated) {
+    return new PrincipalRoleEntity.Builder()
+        .setId(metaStoreManager.generateNewEntityId(callContext).getId())
+        .setName(name)
+        .setFederated(isFederated)
+        .setProperties(Map.of())
+        .setCreateTimestamp(Instant.now().toEpochMilli())
+        .setLastUpdateTimestamp(Instant.now().toEpochMilli())
+        .build();
+  }
+
+  @Test
+  public void testCannotAssignFederatedEntities() {
+    PolarisMetaStoreManager metaStoreManager = setupMetaStoreManager();
+    PolarisCallContext callContext = setupCallContext(metaStoreManager);
+    PolarisAdminService polarisAdminService =
+        setupPolarisAdminService(metaStoreManager, callContext);
+
+    PrincipalEntity principal = createPrincipal(metaStoreManager, callContext, "principal_id");
+    metaStoreManager.createPrincipal(callContext, principal);
+
+    PrincipalRoleEntity role = createRole(metaStoreManager, callContext, "federated_role_id", true);
+    EntityResult result = metaStoreManager.createEntityIfNotExists(callContext, null, role);
+    assertThat(result.isSuccess()).isTrue();
+
+    assertThatThrownBy(
+            () -> polarisAdminService.assignPrincipalRole(principal.getName(), role.getName()))
+        .isInstanceOf(ValidationException.class);
+  }
 }