Skip to content

Commit 80f0456

Browse files
XN137XN137
authored
Use PrincipalEntity in PolarisMetaStoreManager.createPrincipal (#2518)
* Use PrincipalEntity in PolarisMetaStoreManager.createPrincipal if we pass and return a more specific `PrincipalEntity` we can simplify the surrounding code and implementation
1 parent 7bfe2f6 commit 80f0456

File tree

11 files changed

+77
-163
lines changed

11 files changed

+77
-163
lines changed

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

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@
3636
import org.apache.polaris.core.config.PolarisConfigurationStore;
3737
import org.apache.polaris.core.config.RealmConfig;
3838
import org.apache.polaris.core.context.RealmContext;
39-
import org.apache.polaris.core.entity.PolarisEntityConstants;
4039
import org.apache.polaris.core.entity.PrincipalEntity;
4140
import org.apache.polaris.core.persistence.AtomicOperationMetaStoreManager;
4241
import org.apache.polaris.core.persistence.BasePersistence;
@@ -257,11 +256,7 @@ private PrincipalSecretsResult bootstrapServiceAndCreatePolarisPrincipalForRealm
257256

258257
PrincipalEntity rootPrincipal =
259258
metaStoreManager.findRootPrincipal(polarisContext).orElseThrow();
260-
return metaStoreManager.loadPrincipalSecrets(
261-
polarisContext,
262-
rootPrincipal
263-
.getInternalPropertiesAsMap()
264-
.get(PolarisEntityConstants.getClientIdPropertyName()));
259+
return metaStoreManager.loadPrincipalSecrets(polarisContext, rootPrincipal.getClientId());
265260
}
266261

267262
/**

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

Lines changed: 27 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@
4949
import org.apache.polaris.core.entity.PolarisPrincipalSecrets;
5050
import org.apache.polaris.core.entity.PolarisPrivilege;
5151
import org.apache.polaris.core.entity.PolarisTaskConstants;
52+
import org.apache.polaris.core.entity.PrincipalEntity;
53+
import org.apache.polaris.core.entity.PrincipalRoleEntity;
5254
import org.apache.polaris.core.persistence.dao.entity.BaseResult;
5355
import org.apache.polaris.core.persistence.dao.entity.ChangeTrackingResult;
5456
import org.apache.polaris.core.persistence.dao.entity.CreateCatalogResult;
@@ -253,12 +255,8 @@ private void dropEntity(
253255

254256
// if it is a principal, we also need to drop the secrets
255257
if (entity.getType() == PolarisEntityType.PRINCIPAL) {
256-
// get internal properties
257-
Map<String, String> properties = entity.getInternalPropertiesAsMap();
258-
259-
// get client_id
260-
String clientId = properties.get(PolarisEntityConstants.getClientIdPropertyName());
261-
258+
PrincipalEntity principalEntity = PrincipalEntity.of(entity);
259+
String clientId = principalEntity.getClientId();
262260
// delete it from the secret slice
263261
((IntegrationPersistence) ms).deletePrincipalSecrets(callCtx, clientId, entity.getId());
264262
}
@@ -577,28 +575,22 @@ private void revokeGrantRecord(
577575
// role. The principal role will be granted to that root principal and the root catalog admin
578576
// of the root catalog will be granted to that principal role.
579577
long rootPrincipalId = ms.generateNewId(callCtx);
580-
PolarisBaseEntity rootPrincipal =
581-
new PolarisBaseEntity(
582-
PolarisEntityConstants.getNullId(),
583-
rootPrincipalId,
584-
PolarisEntityType.PRINCIPAL,
585-
PolarisEntitySubType.NULL_SUBTYPE,
586-
PolarisEntityConstants.getRootEntityId(),
587-
PolarisEntityConstants.getRootPrincipalName());
588-
589-
// create this principal
578+
PrincipalEntity rootPrincipal =
579+
new PrincipalEntity.Builder()
580+
.setId(rootPrincipalId)
581+
.setName(PolarisEntityConstants.getRootPrincipalName())
582+
.setCreateTimestamp(System.currentTimeMillis())
583+
.build();
590584
this.createPrincipal(callCtx, rootPrincipal);
591585

592586
// now create the account admin principal role
593587
long serviceAdminPrincipalRoleId = ms.generateNewId(callCtx);
594-
PolarisBaseEntity serviceAdminPrincipalRole =
595-
new PolarisBaseEntity(
596-
PolarisEntityConstants.getNullId(),
597-
serviceAdminPrincipalRoleId,
598-
PolarisEntityType.PRINCIPAL_ROLE,
599-
PolarisEntitySubType.NULL_SUBTYPE,
600-
PolarisEntityConstants.getRootEntityId(),
601-
PolarisEntityConstants.getNameOfPrincipalServiceAdminRole());
588+
PrincipalRoleEntity serviceAdminPrincipalRole =
589+
new PrincipalRoleEntity.Builder()
590+
.setId(serviceAdminPrincipalRoleId)
591+
.setName(PolarisEntityConstants.getNameOfPrincipalServiceAdminRole())
592+
.setCreateTimestamp(System.currentTimeMillis())
593+
.build();
602594
this.persistNewEntity(callCtx, ms, serviceAdminPrincipalRole);
603595

604596
// we also need to grant usage on the account-admin principal to the principal
@@ -748,51 +740,28 @@ private void revokeGrantRecord(
748740
/** {@inheritDoc} */
749741
@Override
750742
public @Nonnull CreatePrincipalResult createPrincipal(
751-
@Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity principal) {
743+
@Nonnull PolarisCallContext callCtx, @Nonnull PrincipalEntity principal) {
752744
// get metastore we should be using
753745
BasePersistence ms = callCtx.getMetaStore();
754746

755747
// validate input
756748
getDiagnostics().checkNotNull(principal, "unexpected_null_principal");
757749

758750
// check if that catalog has already been created
759-
PolarisBaseEntity refreshPrincipal =
760-
ms.lookupEntity(
761-
callCtx, principal.getCatalogId(), principal.getId(), principal.getTypeCode());
751+
PrincipalEntity refreshPrincipal =
752+
PrincipalEntity.of(
753+
ms.lookupEntity(
754+
callCtx, principal.getCatalogId(), principal.getId(), principal.getTypeCode()));
762755

763756
// if found, probably a retry, simply return the previously created principal
764757
// This can be done safely as a separate atomic operation before trying to create the principal
765758
// because same-id idempotent-retry collisions of this sort are necessarily sequential, so
766759
// there is no concurrency conflict for something else creating a principal of this same id.
767760
if (refreshPrincipal != null) {
768-
// if found, ensure it is indeed a principal
769-
getDiagnostics()
770-
.check(
771-
principal.getTypeCode() == PolarisEntityType.PRINCIPAL.getCode(),
772-
"not_a_principal",
773-
"principal={}",
774-
principal);
775-
776-
// get internal properties
777-
Map<String, String> properties = refreshPrincipal.getInternalPropertiesAsMap();
778-
779-
// get client_id
780-
String clientId = properties.get(PolarisEntityConstants.getClientIdPropertyName());
781-
782-
// should not be null
783-
getDiagnostics()
784-
.checkNotNull(
785-
clientId,
786-
"null_client_id",
787-
"properties={}",
788-
refreshPrincipal.getInternalProperties());
789-
// ensure non null and non empty
761+
String clientId = refreshPrincipal.getClientId();
762+
getDiagnostics().checkNotNull(clientId, "null_client_id", "principal={}", refreshPrincipal);
790763
getDiagnostics()
791-
.check(
792-
!clientId.isEmpty(),
793-
"empty_client_id",
794-
"properties={}",
795-
refreshPrincipal.getInternalProperties());
764+
.check(!clientId.isEmpty(), "empty_client_id", "principal={}", refreshPrincipal);
796765

797766
// get the main and secondary secrets for that client
798767
PolarisPrincipalSecrets principalSecrets =
@@ -816,15 +785,10 @@ private void revokeGrantRecord(
816785
((IntegrationPersistence) ms)
817786
.generateNewPrincipalSecrets(callCtx, principal.getName(), principal.getId());
818787

819-
// generate properties
820-
Map<String, String> internalProperties = principal.getInternalPropertiesAsMap();
821-
internalProperties.put(
822-
PolarisEntityConstants.getClientIdPropertyName(), principalSecrets.getPrincipalClientId());
823-
824788
// remember client id
825-
PolarisBaseEntity updatedPrincipal =
826-
new PolarisBaseEntity.Builder(principal)
827-
.internalPropertiesAsMap(internalProperties)
789+
PrincipalEntity updatedPrincipal =
790+
new PrincipalEntity.Builder(principal)
791+
.setClientId(principalSecrets.getPrincipalClientId())
828792
.build();
829793
// now create and persist new catalog entity
830794
EntityResult lowLevelResult = this.persistNewEntity(callCtx, ms, updatedPrincipal);

‎polaris-core/src/main/java/org/apache/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929
import org.apache.polaris.core.PolarisDiagnostics;
3030
import org.apache.polaris.core.config.RealmConfig;
3131
import org.apache.polaris.core.context.RealmContext;
32-
import org.apache.polaris.core.entity.PolarisEntityConstants;
3332
import org.apache.polaris.core.entity.PrincipalEntity;
3433
import org.apache.polaris.core.persistence.bootstrap.RootCredentialsSet;
3534
import org.apache.polaris.core.persistence.cache.EntityCache;
@@ -202,11 +201,7 @@ private PrincipalSecretsResult bootstrapServiceAndCreatePolarisPrincipalForRealm
202201

203202
PrincipalEntity rootPrincipal =
204203
metaStoreManager.findRootPrincipal(polarisContext).orElseThrow();
205-
return metaStoreManager.loadPrincipalSecrets(
206-
polarisContext,
207-
rootPrincipal
208-
.getInternalPropertiesAsMap()
209-
.get(PolarisEntityConstants.getClientIdPropertyName()));
204+
return metaStoreManager.loadPrincipalSecrets(polarisContext, rootPrincipal.getClientId());
210205
}
211206

212207
/**

‎polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ Page<PolarisBaseEntity> loadEntities(
194194
*/
195195
@Nonnull
196196
CreatePrincipalResult createPrincipal(
197-
@Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity principal);
197+
@Nonnull PolarisCallContext callCtx, @Nonnull PrincipalEntity principal);
198198

199199
/**
200200
* Create a new catalog. This not only creates the new catalog entity but also the initial admin

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import org.apache.polaris.core.entity.PolarisEntityType;
3838
import org.apache.polaris.core.entity.PolarisEvent;
3939
import org.apache.polaris.core.entity.PolarisPrivilege;
40+
import org.apache.polaris.core.entity.PrincipalEntity;
4041
import org.apache.polaris.core.persistence.dao.entity.BaseResult;
4142
import org.apache.polaris.core.persistence.dao.entity.ChangeTrackingResult;
4243
import org.apache.polaris.core.persistence.dao.entity.CreateCatalogResult;
@@ -150,7 +151,7 @@ public GenerateEntityIdResult generateNewEntityId(@Nonnull PolarisCallContext ca
150151

151152
@Override
152153
public CreatePrincipalResult createPrincipal(
153-
@Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity principal) {
154+
@Nonnull PolarisCallContext callCtx, @Nonnull PrincipalEntity principal) {
154155
diagnostics.fail("illegal_method_in_transaction_workspace", "createPrincipal");
155156
return null;
156157
}

‎polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/CreatePrincipalResult.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import jakarta.annotation.Nullable;
2525
import org.apache.polaris.core.entity.PolarisBaseEntity;
2626
import org.apache.polaris.core.entity.PolarisPrincipalSecrets;
27+
import org.apache.polaris.core.entity.PrincipalEntity;
2728

2829
/** the return the result of a create-principal method */
2930
public class CreatePrincipalResult extends BaseResult {
@@ -69,8 +70,8 @@ private CreatePrincipalResult(
6970
this.principalSecrets = principalSecrets;
7071
}
7172

72-
public PolarisBaseEntity getPrincipal() {
73-
return principal;
73+
public PrincipalEntity getPrincipal() {
74+
return PrincipalEntity.of(principal);
7475
}
7576

7677
public PolarisPrincipalSecrets getPrincipalSecrets() {

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

Lines changed: 29 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@
4848
import org.apache.polaris.core.entity.PolarisPrincipalSecrets;
4949
import org.apache.polaris.core.entity.PolarisPrivilege;
5050
import org.apache.polaris.core.entity.PolarisTaskConstants;
51+
import org.apache.polaris.core.entity.PrincipalEntity;
52+
import org.apache.polaris.core.entity.PrincipalRoleEntity;
5153
import org.apache.polaris.core.persistence.BaseMetaStoreManager;
5254
import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
5355
import org.apache.polaris.core.persistence.PolarisObjectMapperUtil;
@@ -238,12 +240,8 @@ private void dropEntity(
238240

239241
// if it is a principal, we also need to drop the secrets
240242
if (entity.getType() == PolarisEntityType.PRINCIPAL) {
241-
// get internal properties
242-
Map<String, String> properties = entity.getInternalPropertiesAsMap();
243-
244-
// get client_id
245-
String clientId = properties.get(PolarisEntityConstants.getClientIdPropertyName());
246-
243+
PrincipalEntity principalEntity = PrincipalEntity.of(entity);
244+
String clientId = principalEntity.getClientId();
247245
// delete it from the secret slice
248246
ms.deletePrincipalSecretsInCurrentTxn(callCtx, clientId, entity.getId());
249247
}
@@ -552,28 +550,22 @@ private void bootstrapPolarisService(
552550
// role. The principal role will be granted to that root principal and the root catalog admin
553551
// of the root catalog will be granted to that principal role.
554552
long rootPrincipalId = ms.generateNewIdInCurrentTxn(callCtx);
555-
PolarisBaseEntity rootPrincipal =
556-
new PolarisBaseEntity(
557-
PolarisEntityConstants.getNullId(),
558-
rootPrincipalId,
559-
PolarisEntityType.PRINCIPAL,
560-
PolarisEntitySubType.NULL_SUBTYPE,
561-
PolarisEntityConstants.getRootEntityId(),
562-
PolarisEntityConstants.getRootPrincipalName());
563-
564-
// create this principal
553+
PrincipalEntity rootPrincipal =
554+
new PrincipalEntity.Builder()
555+
.setId(rootPrincipalId)
556+
.setName(PolarisEntityConstants.getRootPrincipalName())
557+
.setCreateTimestamp(System.currentTimeMillis())
558+
.build();
565559
this.createPrincipal(callCtx, ms, rootPrincipal);
566560

567561
// now create the account admin principal role
568562
long serviceAdminPrincipalRoleId = ms.generateNewIdInCurrentTxn(callCtx);
569-
PolarisBaseEntity serviceAdminPrincipalRole =
570-
new PolarisBaseEntity(
571-
PolarisEntityConstants.getNullId(),
572-
serviceAdminPrincipalRoleId,
573-
PolarisEntityType.PRINCIPAL_ROLE,
574-
PolarisEntitySubType.NULL_SUBTYPE,
575-
PolarisEntityConstants.getRootEntityId(),
576-
PolarisEntityConstants.getNameOfPrincipalServiceAdminRole());
563+
PrincipalRoleEntity serviceAdminPrincipalRole =
564+
new PrincipalRoleEntity.Builder()
565+
.setId(serviceAdminPrincipalRoleId)
566+
.setName(PolarisEntityConstants.getNameOfPrincipalServiceAdminRole())
567+
.setCreateTimestamp(System.currentTimeMillis())
568+
.build();
577569
this.persistNewEntity(callCtx, ms, serviceAdminPrincipalRole);
578570

579571
// we also need to grant usage on the account-admin principal to the principal
@@ -777,49 +769,26 @@ private void bootstrapPolarisService(
777769
() -> loadEntities(callCtx, ms, catalogPath, entityType, entitySubType, pageToken));
778770
}
779771

780-
/** {@link #createPrincipal(PolarisCallContext, PolarisBaseEntity)} */
772+
/** {@link #createPrincipal(PolarisCallContext, PrincipalEntity)} */
781773
private @Nonnull CreatePrincipalResult createPrincipal(
782774
@Nonnull PolarisCallContext callCtx,
783775
@Nonnull TransactionalPersistence ms,
784-
@Nonnull PolarisBaseEntity principal) {
776+
@Nonnull PrincipalEntity principal) {
785777
// validate input
786778
getDiagnostics().checkNotNull(principal, "unexpected_null_principal");
787779

788780
// check if that catalog has already been created
789-
PolarisBaseEntity refreshPrincipal =
790-
ms.lookupEntityInCurrentTxn(
791-
callCtx, principal.getCatalogId(), principal.getId(), principal.getTypeCode());
781+
PrincipalEntity refreshPrincipal =
782+
PrincipalEntity.of(
783+
ms.lookupEntityInCurrentTxn(
784+
callCtx, principal.getCatalogId(), principal.getId(), principal.getTypeCode()));
792785

793786
// if found, probably a retry, simply return the previously created principal
794787
if (refreshPrincipal != null) {
795-
// if found, ensure it is indeed a principal
796-
getDiagnostics()
797-
.check(
798-
principal.getTypeCode() == PolarisEntityType.PRINCIPAL.getCode(),
799-
"not_a_principal",
800-
"principal={}",
801-
principal);
802-
803-
// get internal properties
804-
Map<String, String> properties = refreshPrincipal.getInternalPropertiesAsMap();
805-
806-
// get client_id
807-
String clientId = properties.get(PolarisEntityConstants.getClientIdPropertyName());
808-
809-
// should not be null
810-
getDiagnostics()
811-
.checkNotNull(
812-
clientId,
813-
"null_client_id",
814-
"properties={}",
815-
refreshPrincipal.getInternalProperties());
816-
// ensure non null and non empty
788+
String clientId = refreshPrincipal.getClientId();
789+
getDiagnostics().checkNotNull(clientId, "null_client_id", "principal={}", refreshPrincipal);
817790
getDiagnostics()
818-
.check(
819-
!clientId.isEmpty(),
820-
"empty_client_id",
821-
"properties={}",
822-
refreshPrincipal.getInternalProperties());
791+
.check(!clientId.isEmpty(), "empty_client_id", "principal={}", refreshPrincipal);
823792

824793
// get the main and secondary secrets for that client
825794
PolarisPrincipalSecrets principalSecrets =
@@ -854,15 +823,10 @@ private void bootstrapPolarisService(
854823
PolarisPrincipalSecrets principalSecrets =
855824
ms.generateNewPrincipalSecretsInCurrentTxn(callCtx, principal.getName(), principal.getId());
856825

857-
// generate properties
858-
Map<String, String> internalProperties = principal.getInternalPropertiesAsMap();
859-
internalProperties.put(
860-
PolarisEntityConstants.getClientIdPropertyName(), principalSecrets.getPrincipalClientId());
861-
862826
// remember client id
863-
PolarisBaseEntity updatedPrincipal =
864-
new PolarisBaseEntity.Builder(principal)
865-
.internalPropertiesAsMap(internalProperties)
827+
PrincipalEntity updatedPrincipal =
828+
new PrincipalEntity.Builder(principal)
829+
.setClientId(principalSecrets.getPrincipalClientId())
866830
.build();
867831

868832
// now create and persist new catalog entity
@@ -875,7 +839,7 @@ private void bootstrapPolarisService(
875839
/** {@inheritDoc} */
876840
@Override
877841
public @Nonnull CreatePrincipalResult createPrincipal(
878-
@Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity principal) {
842+
@Nonnull PolarisCallContext callCtx, @Nonnull PrincipalEntity principal) {
879843
// get metastore we should be using
880844
TransactionalPersistence ms = ((TransactionalPersistence) callCtx.getMetaStore());
881845

0 commit comments

Comments
 (0)