diff --git a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java index e19b0ef721..c0aadf3fc5 100644 --- a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java +++ b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java @@ -52,11 +52,14 @@ import org.apache.polaris.core.persistence.BaseMetaStoreManager; import org.apache.polaris.core.persistence.PrincipalSecretsGenerator; import org.apache.polaris.core.persistence.RetryOnConcurrencyException; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; import org.apache.polaris.core.persistence.transactional.AbstractTransactionalPersistence; import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; import org.apache.polaris.core.storage.PolarisStorageIntegration; import org.apache.polaris.core.storage.PolarisStorageIntegrationProvider; +import org.apache.polaris.jpa.models.EntityIdPageToken; import org.apache.polaris.jpa.models.ModelEntity; import org.apache.polaris.jpa.models.ModelEntityActive; import org.apache.polaris.jpa.models.ModelEntityChangeTracking; @@ -419,29 +422,30 @@ public List lookupEntityActiveBatchInCurrentTxn( /** {@inheritDoc} */ @Override - public @Nonnull List listEntitiesInCurrentTxn( + public @Nonnull PolarisPage listEntitiesInCurrentTxn( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, - @Nonnull PolarisEntityType entityType) { + @Nonnull PolarisEntityType entityType, + @Nonnull PageToken pageToken) { return this.listEntitiesInCurrentTxn( - callCtx, catalogId, parentId, entityType, Predicates.alwaysTrue()); + callCtx, catalogId, parentId, entityType, Predicates.alwaysTrue(), pageToken); } @Override - public @Nonnull List listEntitiesInCurrentTxn( + public @Nonnull PolarisPage listEntitiesInCurrentTxn( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, - @Nonnull Predicate entityFilter) { + @Nonnull Predicate entityFilter, + @Nonnull PageToken pageToken) { // full range scan under the parent for that type return this.listEntitiesInCurrentTxn( callCtx, catalogId, parentId, entityType, - Integer.MAX_VALUE, entityFilter, entity -> new EntityNameLookupRecord( @@ -450,27 +454,57 @@ public List lookupEntityActiveBatchInCurrentTxn( entity.getParentId(), entity.getName(), entity.getTypeCode(), - entity.getSubTypeCode())); + entity.getSubTypeCode()), + pageToken); } @Override - public @Nonnull List listEntitiesInCurrentTxn( + public @Nonnull PolarisPage listEntitiesInCurrentTxn( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, - int limit, @Nonnull Predicate entityFilter, - @Nonnull Function transformer) { - // full range scan under the parent for that type - return this.store - .lookupFullEntitiesActive(localSession.get(), catalogId, parentId, entityType) - .stream() - .map(ModelEntity::toEntity) - .filter(entityFilter) - .limit(limit) - .map(transformer) - .collect(Collectors.toList()); + @Nonnull Function transformer, + @Nonnull PageToken pageToken) { + List data; + if (entityFilter.equals(Predicates.alwaysTrue())) { + // In this case, we can push the filter down into the query + data = + this.store + .lookupFullEntitiesActive( + localSession.get(), catalogId, parentId, entityType, pageToken) + .stream() + .map(ModelEntity::toEntity) + .filter(entityFilter) + .map(transformer) + .collect(Collectors.toList()); + } else { + // In this case, we cannot push the filter down into the query. We must therefore remove + // the page size limit from the PageToken and filter on the client side. + // TODO Implement a generic predicate that can be pushed down into different metastores + PageToken unlimitedPageSizeToken = pageToken.withPageSize(Integer.MAX_VALUE); + List rawData = + this.store.lookupFullEntitiesActive( + localSession.get(), catalogId, parentId, entityType, unlimitedPageSizeToken); + if (pageToken.pageSize < Integer.MAX_VALUE && rawData.size() > pageToken.pageSize) { + LOGGER.info( + "A page token could not be respected due to a predicate. " + + "{} records were read but the client was asked to return {}.", + rawData.size(), + pageToken.pageSize); + } + + data = + rawData.stream() + .map(ModelEntity::toEntity) + .filter(entityFilter) + .limit(pageToken.pageSize) + .map(transformer) + .collect(Collectors.toList()); + } + + return pageToken.buildNextPage(data); } /** {@inheritDoc} */ @@ -762,4 +796,9 @@ public void rollback() { session.getTransaction().rollback(); } } + + @Override + public @Nonnull PageToken.PageTokenBuilder pageTokenBuilder() { + return EntityIdPageToken.builder(); + } } diff --git a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkStore.java b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkStore.java index 9ebb4e8164..ca3693b2ac 100644 --- a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkStore.java +++ b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkStore.java @@ -35,7 +35,10 @@ import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PolarisGrantRecord; import org.apache.polaris.core.entity.PolarisPrincipalSecrets; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; +import org.apache.polaris.jpa.models.EntityIdPageToken; import org.apache.polaris.jpa.models.ModelEntity; import org.apache.polaris.jpa.models.ModelEntityActive; import org.apache.polaris.jpa.models.ModelEntityChangeTracking; @@ -282,21 +285,40 @@ long countActiveChildEntities( } List lookupFullEntitiesActive( - EntityManager session, long catalogId, long parentId, @Nonnull PolarisEntityType entityType) { + EntityManager session, + long catalogId, + long parentId, + @Nonnull PolarisEntityType entityType, + @Nonnull PageToken pageToken) { diagnosticServices.check(session != null, "session_is_null"); + diagnosticServices.check( + (pageToken instanceof EntityIdPageToken || pageToken instanceof ReadEverythingPageToken), + "unexpected_page_token"); checkInitialized(); // Currently check against ENTITIES not joining with ENTITIES_ACTIVE String hql = - "SELECT m from ModelEntity m where m.catalogId=:catalogId and m.parentId=:parentId and m.typeCode=:typeCode"; + "SELECT m from ModelEntity m " + + "where m.catalogId=:catalogId and m.parentId=:parentId and m.typeCode=:typeCode and m.id > :tokenId"; + + if (pageToken instanceof EntityIdPageToken) { + hql += " order by m.id asc"; + } TypedQuery query = session .createQuery(hql, ModelEntity.class) .setParameter("catalogId", catalogId) .setParameter("parentId", parentId) - .setParameter("typeCode", entityType.getCode()); - + .setParameter("typeCode", entityType.getCode()) + .setParameter("tokenId", -1L); + + if (pageToken instanceof EntityIdPageToken) { + query = + query + .setParameter("tokenId", ((EntityIdPageToken) pageToken).id) + .setMaxResults(pageToken.pageSize); + } return query.getResultList(); } diff --git a/extension/persistence/jpa-model/src/main/java/org/apache/polaris/jpa/models/EntityIdPageToken.java b/extension/persistence/jpa-model/src/main/java/org/apache/polaris/jpa/models/EntityIdPageToken.java new file mode 100644 index 0000000000..7035420fd6 --- /dev/null +++ b/extension/persistence/jpa-model/src/main/java/org/apache/polaris/jpa/models/EntityIdPageToken.java @@ -0,0 +1,114 @@ +/* + * 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.jpa.models; + +import java.util.List; +import org.apache.polaris.core.entity.PolarisBaseEntity; +import org.apache.polaris.core.persistence.pagination.PageToken; + +/** + * A {@link PageToken} implementation that tracks the greatest ID from either {@link + * PolarisBaseEntity} or {@link ModelEntity} objects supplied in updates. Entities are meant to be + * filtered during listing such that only entities with and ID greater than the ID of the token are + * returned. + */ +public class EntityIdPageToken extends PageToken { + public long id; + + private EntityIdPageToken(long id, int pageSize) { + this.id = id; + this.pageSize = pageSize; + validate(); + } + + /** The minimum ID that could be attached to an entity */ + private static final long MINIMUM_ID = 0; + + /** The entity ID to use to start with. */ + private static final long BASE_ID = MINIMUM_ID - 1; + + @Override + protected List getComponents() { + return List.of(String.valueOf(id), String.valueOf(pageSize)); + } + + /** Get a new `EntityIdPageTokenBuilder` instance */ + public static PageTokenBuilder builder() { + return new EntityIdPageToken.EntityIdPageTokenBuilder(); + } + + @Override + protected PageTokenBuilder getBuilder() { + return EntityIdPageToken.builder(); + } + + /** A {@link PageTokenBuilder} implementation for {@link EntityIdPageToken} */ + public static class EntityIdPageTokenBuilder extends PageTokenBuilder { + + @Override + public String tokenPrefix() { + return "polaris-entity-id"; + } + + @Override + public int expectedComponents() { + // id, pageSize + return 2; + } + + @Override + protected EntityIdPageToken fromStringComponents(List components) { + return new EntityIdPageToken( + Integer.parseInt(components.get(0)), Integer.parseInt(components.get(1))); + } + + @Override + protected EntityIdPageToken fromLimitImpl(int limit) { + return new EntityIdPageToken(BASE_ID, limit); + } + } + + @Override + public PageToken updated(List newData) { + if (newData == null || newData.size() < this.pageSize) { + return DONE; + } else { + var head = newData.getFirst(); + if (head instanceof ModelEntity) { + return new EntityIdPageToken(((ModelEntity) newData.getLast()).getId(), this.pageSize); + } else if (head instanceof PolarisBaseEntity) { + // Assumed to be sorted with the greatest entity ID last + return new EntityIdPageToken( + ((PolarisBaseEntity) newData.getLast()).getId(), this.pageSize); + } else { + throw new IllegalArgumentException( + "Cannot build a page token from: " + newData.getFirst().getClass().getSimpleName()); + } + } + } + + @Override + public PageToken withPageSize(Integer pageSize) { + if (pageSize == null) { + return new EntityIdPageToken(BASE_ID, this.pageSize); + } else { + return new EntityIdPageToken(this.id, pageSize); + } + } +} diff --git a/extension/persistence/relational-jdbc/build.gradle.kts b/extension/persistence/relational-jdbc/build.gradle.kts index 82f67c8a5a..59f0454f3d 100644 --- a/extension/persistence/relational-jdbc/build.gradle.kts +++ b/extension/persistence/relational-jdbc/build.gradle.kts @@ -24,6 +24,7 @@ plugins { dependencies { implementation(project(":polaris-core")) + implementation(project(":polaris-jpa-model")) implementation(libs.slf4j.api) implementation(libs.guava) diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/DatasourceOperations.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/DatasourceOperations.java index 6fec8e67fe..e29bf02cb8 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/DatasourceOperations.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/DatasourceOperations.java @@ -34,6 +34,7 @@ import java.util.function.Function; import java.util.function.Predicate; import javax.sql.DataSource; +import org.apache.polaris.core.persistence.pagination.PageToken; import org.apache.polaris.extension.persistence.relational.jdbc.models.Converter; public class DatasourceOperations { @@ -95,8 +96,8 @@ public void executeScript(String scriptFilePath) throws SQLException { * @param query : Query to executed * @param entityClass : Class of the entity being selected * @param transformer : Transformation of entity class to Result class - * @param entityFilter : Filter to applied on the Result class - * @param limit : Limit to to enforced. + * @param entityFilter : Client-side filter to applied on the Result class + * @param pageToken : Page token to be enforced. * @return List of Result class objects * @param : Entity class * @param : Result class @@ -107,13 +108,13 @@ public List executeSelect( @Nonnull Class entityClass, @Nonnull Function transformer, Predicate entityFilter, - int limit) + PageToken pageToken) throws SQLException { try (Connection connection = borrowConnection(); Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(query)) { - List resultList = new ArrayList<>(); - while (resultSet.next() && resultList.size() < limit) { + List resultList = new ArrayList<>(PageToken.DEFAULT_PAGE_SIZE); + while (resultSet.next() && resultList.size() < pageToken.pageSize) { Converter object = (Converter) entityClass.getDeclaredConstructor().newInstance(); // Create a new instance diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index 5ffce813f9..7fcd1dfc49 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java @@ -21,6 +21,7 @@ import static org.apache.polaris.extension.persistence.relational.jdbc.QueryGenerator.*; import com.google.common.base.Preconditions; +import com.google.common.base.Predicates; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; import java.sql.SQLException; @@ -34,6 +35,7 @@ import java.util.function.Predicate; import java.util.stream.Collectors; import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisChangeTrackingVersions; @@ -49,6 +51,9 @@ import org.apache.polaris.core.persistence.PolicyMappingAlreadyExistsException; import org.apache.polaris.core.persistence.PrincipalSecretsGenerator; import org.apache.polaris.core.persistence.RetryOnConcurrencyException; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; import org.apache.polaris.core.policy.PolicyType; import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; @@ -58,6 +63,7 @@ import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelGrantRecord; import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPolicyMappingRecord; import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPrincipalAuthenticationData; +import org.apache.polaris.jpa.models.EntityIdPageToken; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -69,16 +75,19 @@ public class JdbcBasePersistenceImpl implements BasePersistence, IntegrationPers private final PrincipalSecretsGenerator secretsGenerator; private final PolarisStorageIntegrationProvider storageIntegrationProvider; private final String realmId; + private final PolarisDiagnostics polarisDiagnostics; public JdbcBasePersistenceImpl( DatasourceOperations databaseOperations, PrincipalSecretsGenerator secretsGenerator, PolarisStorageIntegrationProvider storageIntegrationProvider, - String realmId) { + String realmId, + PolarisDiagnostics polarisDiagnostics) { this.datasourceOperations = databaseOperations; this.secretsGenerator = secretsGenerator; this.storageIntegrationProvider = storageIntegrationProvider; this.realmId = realmId; + this.polarisDiagnostics = polarisDiagnostics; } @Override @@ -303,7 +312,7 @@ private PolarisBaseEntity getPolarisBaseEntity(String query) { try { List results = datasourceOperations.executeSelect( - query, ModelEntity.class, ModelEntity::toEntity, null, Integer.MAX_VALUE); + query, ModelEntity.class, ModelEntity::toEntity, null, ReadEverythingPageToken.get()); if (results.isEmpty()) { return null; } else if (results.size() > 1) { @@ -328,7 +337,7 @@ public List lookupEntities( String query = generateSelectQueryWithEntityIds(realmId, entityIds); try { return datasourceOperations.executeSelect( - query, ModelEntity.class, ModelEntity::toEntity, null, Integer.MAX_VALUE); + query, ModelEntity.class, ModelEntity::toEntity, null, ReadEverythingPageToken.get()); } catch (SQLException e) { throw new RuntimeException( String.format("Failed to retrieve polaris entities due to %s", e.getMessage()), e); @@ -359,49 +368,55 @@ public List lookupEntityVersions( @Nonnull @Override - public List listEntities( + public PolarisPage listEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, - @Nonnull PolarisEntityType entityType) { + @Nonnull PolarisEntityType entityType, + @Nonnull PageToken pageToken) { return listEntities( callCtx, catalogId, parentId, entityType, - Integer.MAX_VALUE, entity -> true, - EntityNameLookupRecord::new); + EntityNameLookupRecord::new, + pageToken); } @Nonnull @Override - public List listEntities( + public PolarisPage listEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, - @Nonnull Predicate entityFilter) { + @Nonnull Predicate entityFilter, + @Nonnull PageToken pageToken) { return listEntities( callCtx, catalogId, parentId, entityType, - Integer.MAX_VALUE, entityFilter, - EntityNameLookupRecord::new); + EntityNameLookupRecord::new, + pageToken); } @Nonnull @Override - public List listEntities( + public PolarisPage listEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, PolarisEntityType entityType, - int limit, @Nonnull Predicate entityFilter, - @Nonnull Function transformer) { + @Nonnull Function transformer, + PageToken pageToken) { + polarisDiagnostics.check( + (pageToken instanceof EntityIdPageToken || pageToken instanceof ReadEverythingPageToken), + "unexpected_page_token"); + Map params = Map.of( "catalog_id", @@ -416,17 +431,53 @@ public List listEntities( // Limit can't be pushed down, due to client side filtering // absence of transaction. String query = QueryGenerator.generateSelectQuery(ModelEntity.class, params); + if (!(pageToken instanceof ReadEverythingPageToken)) { + query = QueryGenerator.updateQueryWithPageToken(query, pageToken); + } + final List data; try { - List results = - datasourceOperations.executeSelect( - query, ModelEntity.class, ModelEntity::toEntity, entityFilter, limit); - return results == null - ? Collections.emptyList() - : results.stream().filter(entityFilter).map(transformer).collect(Collectors.toList()); + if (entityFilter.equals(Predicates.alwaysTrue())) { + // In this case, we can push the filter down into the query + data = + datasourceOperations + .executeSelect( + query, ModelEntity.class, ModelEntity::toEntity, entityFilter, pageToken) + .stream() + .map(transformer) + .toList(); + } else { + // In this case, we cannot push the filter down into the query. We must therefore remove + // the page size limit from the PageToken and filter on the client side. + // TODO Implement a generic predicate that can be pushed down into different metastores + PageToken unlimitedPageSizeToken = pageToken.withPageSize(Integer.MAX_VALUE); + List rawData = + datasourceOperations.executeSelect( + query, + ModelEntity.class, + ModelEntity::toEntity, + entityFilter, + unlimitedPageSizeToken); + if (pageToken.pageSize < Integer.MAX_VALUE && rawData.size() > pageToken.pageSize) { + LOGGER.info( + "A page token could not be respected due to a predicate. " + + "{} records were read but the client was asked to return {}.", + rawData.size(), + pageToken.pageSize); + } + + data = + rawData.stream() + .filter(entityFilter) + .limit(pageToken.pageSize) + .map(transformer) + .collect(Collectors.toList()); + } } catch (SQLException e) { throw new RuntimeException( String.format("Failed to retrieve polaris entities due to %s", e.getMessage()), e); } + + return pageToken.buildNextPage(data); } @Override @@ -470,7 +521,7 @@ public PolarisGrantRecord lookupGrantRecord( ModelGrantRecord.class, ModelGrantRecord::toGrantRecord, null, - Integer.MAX_VALUE); + ReadEverythingPageToken.get()); if (results.size() > 1) { throw new IllegalStateException( String.format( @@ -505,7 +556,7 @@ public List loadAllGrantRecordsOnSecurable( ModelGrantRecord.class, ModelGrantRecord::toGrantRecord, null, - Integer.MAX_VALUE); + ReadEverythingPageToken.get()); return results == null ? Collections.emptyList() : results; } catch (SQLException e) { throw new RuntimeException( @@ -531,7 +582,7 @@ public List loadAllGrantRecordsOnGrantee( ModelGrantRecord.class, ModelGrantRecord::toGrantRecord, null, - Integer.MAX_VALUE); + ReadEverythingPageToken.get()); return results == null ? Collections.emptyList() : results; } catch (SQLException e) { throw new RuntimeException( @@ -559,7 +610,7 @@ public boolean hasChildren( try { List results = datasourceOperations.executeSelect( - query, ModelEntity.class, Function.identity(), null, Integer.MAX_VALUE); + query, ModelEntity.class, Function.identity(), null, ReadEverythingPageToken.get()); return results != null && !results.isEmpty(); } catch (SQLException e) { throw new RuntimeException( @@ -582,7 +633,7 @@ public PolarisPrincipalSecrets loadPrincipalSecrets( ModelPrincipalAuthenticationData.class, ModelPrincipalAuthenticationData::toPrincipalAuthenticationData, null, - Integer.MAX_VALUE); + ReadEverythingPageToken.get()); return results == null || results.isEmpty() ? null : results.getFirst(); } catch (SQLException e) { LOGGER.error( @@ -888,7 +939,7 @@ private List fetchPolicyMappingRecords(String query) ModelPolicyMappingRecord.class, ModelPolicyMappingRecord::toPolicyMappingRecord, null, - Integer.MAX_VALUE); + ReadEverythingPageToken.get()); return results == null ? Collections.emptyList() : results; } catch (SQLException e) { throw new RuntimeException( @@ -923,4 +974,10 @@ PolarisStorageIntegration loadPolarisStorageIntegration( BaseMetaStoreManager.extractStorageConfiguration(callContext, entity); return storageIntegrationProvider.getStorageIntegrationForConfig(storageConfig); } + + @Nonnull + @Override + public PageToken.PageTokenBuilder pageTokenBuilder() { + return EntityIdPageToken.builder(); + } } diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcMetaStoreManagerFactory.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcMetaStoreManagerFactory.java index f19194f33d..638ab5d581 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcMetaStoreManagerFactory.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcMetaStoreManagerFactory.java @@ -101,7 +101,8 @@ private void initializeForRealm( databaseOperations, secretsGenerator(realmContext, rootCredentialsSet), storageIntegrationProvider, - realmContext.getRealmIdentifier())); + realmContext.getRealmIdentifier(), + diagServices)); PolarisMetaStoreManager metaStoreManager = createNewMetaStoreManager(); metaStoreManagerMap.put(realmContext.getRealmIdentifier(), metaStoreManager); diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java index 7d0b5ec928..896079780c 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java @@ -26,13 +26,19 @@ import java.util.Map; import org.apache.polaris.core.entity.PolarisEntityCore; import org.apache.polaris.core.entity.PolarisEntityId; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; import org.apache.polaris.extension.persistence.relational.jdbc.models.Converter; import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelEntity; import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelGrantRecord; import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPolicyMappingRecord; import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPrincipalAuthenticationData; +import org.apache.polaris.jpa.models.EntityIdPageToken; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class QueryGenerator { + private static final Logger LOGGER = LoggerFactory.getLogger(QueryGenerator.class); public static String generateSelectQuery( @Nonnull Class entityClass, @Nonnull Map whereClause) { @@ -218,4 +224,17 @@ public static String getTableName(@Nonnull Class entityClass) { return tableName; } + + public static String updateQueryWithPageToken(String existingQuery, PageToken pageToken) { + if (pageToken instanceof ReadEverythingPageToken) { + return existingQuery; + } else if (pageToken instanceof EntityIdPageToken) { + long previousPageEntityId = ((EntityIdPageToken) pageToken).id; + return String.format("%s AND id > %d ORDER BY id ASC", existingQuery, previousPageEntityId); + } else { + // The caller of this method is supposed to already have validated the PageToken! + LOGGER.error("Unsupported page token: {}", pageToken); + return existingQuery; + } + } } diff --git a/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java b/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java index 92d31d8343..57df0febb9 100644 --- a/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java +++ b/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java @@ -58,7 +58,12 @@ protected PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() { } JdbcBasePersistenceImpl basePersistence = - new JdbcBasePersistenceImpl(datasourceOperations, RANDOM_SECRETS, Mockito.mock(), "REALM"); + new JdbcBasePersistenceImpl( + datasourceOperations, + RANDOM_SECRETS, + Mockito.mock(), + "REALM", + new PolarisDefaultDiagServiceImpl()); return new PolarisTestMetaStoreManager( new AtomicOperationMetaStoreManager(), new PolarisCallContext( diff --git a/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/DatasourceOperationsTest.java b/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/DatasourceOperationsTest.java index 48d2aa39bf..5544033330 100644 --- a/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/DatasourceOperationsTest.java +++ b/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/DatasourceOperationsTest.java @@ -26,6 +26,7 @@ import java.sql.Statement; import java.util.function.Function; import javax.sql.DataSource; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; import org.apache.polaris.extension.persistence.relational.jdbc.DatasourceOperations; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -78,7 +79,7 @@ void testExecuteSelect_exception() throws Exception { SQLException.class, () -> datasourceOperations.executeSelect( - query, Object.class, Function.identity(), null, Integer.MAX_VALUE)); + query, Object.class, Function.identity(), null, ReadEverythingPageToken.get())); } @Test diff --git a/polaris-core/src/main/java/org/apache/polaris/core/config/FeatureConfiguration.java b/polaris-core/src/main/java/org/apache/polaris/core/config/FeatureConfiguration.java index ee663dbb19..3ae1914129 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/config/FeatureConfiguration.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/config/FeatureConfiguration.java @@ -202,6 +202,16 @@ public static void enforceFeatureEnabledOrThrow( .defaultValue(2) .buildFeatureConfiguration(); + public static final PolarisConfiguration LIST_PAGINATION_ENABLED = + PolarisConfiguration.builder() + .key("LIST_PAGINATION_ENABLED") + .catalogConfig("list-pagination.enabled") + .description( + "If set to true, pagination for APIs like listTables is enabled. The APIs that" + + " currently support pagination are listTables, listViews, and listNamespaces.") + .defaultValue(false) + .buildFeatureConfiguration(); + public static final FeatureConfiguration ENABLE_GENERIC_TABLES = PolarisConfiguration.builder() .key("ENABLE_GENERIC_TABLES") diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java index d6cfd84292..3e7fffb632 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java @@ -62,6 +62,8 @@ import org.apache.polaris.core.persistence.dao.entity.PrivilegeResult; import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult; import org.apache.polaris.core.persistence.dao.entity.ScopedCredentialsResult; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; import org.apache.polaris.core.policy.PolicyEntity; import org.apache.polaris.core.policy.PolicyType; @@ -665,7 +667,8 @@ private void revokeGrantRecord( @Nonnull PolarisCallContext callCtx, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, - @Nonnull PolarisEntitySubType entitySubType) { + @Nonnull PolarisEntitySubType entitySubType, + @Nonnull PageToken pageToken) { // get meta store we should be using BasePersistence ms = callCtx.getMetaStore(); @@ -677,15 +680,16 @@ private void revokeGrantRecord( catalogPath == null || catalogPath.size() == 0 ? 0l : catalogPath.get(catalogPath.size() - 1).getId(); - List toreturnList = - ms.listEntities(callCtx, catalogId, parentId, entityType); + PolarisPage resultPage = + ms.listEntities(callCtx, catalogId, parentId, entityType, pageToken); // prune the returned list with only entities matching the entity subtype if (entitySubType != PolarisEntitySubType.ANY_SUBTYPE) { - toreturnList = - toreturnList.stream() - .filter(rec -> rec.getSubTypeCode() == entitySubType.getCode()) - .collect(Collectors.toList()); + resultPage = + pageToken.buildNextPage( + resultPage.data.stream() + .filter(rec -> rec.getSubTypeCode() == entitySubType.getCode()) + .collect(Collectors.toList())); } // TODO: Use post-validation to enforce consistent view against catalogPath. In the @@ -695,7 +699,7 @@ private void revokeGrantRecord( // in-flight request (the cache-based resolution follows a different path entirely). // done - return new ListEntitiesResult(toreturnList); + return ListEntitiesResult.fromPolarisPage(resultPage); } /** {@inheritDoc} */ @@ -1154,13 +1158,14 @@ private void revokeGrantRecord( // get the list of catalog roles, at most 2 List catalogRoles = ms.listEntities( - callCtx, - catalogId, - catalogId, - PolarisEntityType.CATALOG_ROLE, - 2, - entity -> true, - Function.identity()); + callCtx, + catalogId, + catalogId, + PolarisEntityType.CATALOG_ROLE, + entity -> true, + Function.identity(), + ms.pageTokenBuilder().fromLimit(2)) + .data; // if we have 2, we cannot drop the catalog. If only one left, better be the admin role if (catalogRoles.size() > 1) { @@ -1454,17 +1459,16 @@ private void revokeGrantRecord( @Override public @Nonnull EntitiesResult loadTasks( - @Nonnull PolarisCallContext callCtx, String executorId, int limit) { + @Nonnull PolarisCallContext callCtx, String executorId, PageToken pageToken) { BasePersistence ms = callCtx.getMetaStore(); // find all available tasks - List availableTasks = + PolarisPage availableTasks = ms.listEntities( callCtx, PolarisEntityConstants.getRootEntityId(), PolarisEntityConstants.getRootEntityId(), PolarisEntityType.TASK, - limit, entity -> { PolarisObjectMapperUtil.TaskExecutionState taskState = PolarisObjectMapperUtil.parseTaskState(entity); @@ -1479,11 +1483,12 @@ private void revokeGrantRecord( || taskState.executor == null || callCtx.getClock().millis() - taskState.lastAttemptStartTime > taskAgeTimeout; }, - Function.identity()); + Function.identity(), + pageToken); List loadedTasks = new ArrayList<>(); final AtomicInteger failedLeaseCount = new AtomicInteger(0); - availableTasks.forEach( + availableTasks.data.forEach( task -> { PolarisBaseEntity updatedTask = new PolarisBaseEntity(task); Map properties = @@ -1520,7 +1525,7 @@ private void revokeGrantRecord( throw new RetryOnConcurrencyException( "Failed to lease any of %s tasks due to concurrent leases", failedLeaseCount.get()); } - return new EntitiesResult(loadedTasks); + return EntitiesResult.fromPolarisPage(PolarisPage.fromData(loadedTasks)); } /** {@inheritDoc} */ diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java index 75b18fb45c..d0f79fd146 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java @@ -31,6 +31,8 @@ import org.apache.polaris.core.entity.PolarisEntityId; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PolarisGrantRecord; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; import org.apache.polaris.core.policy.PolicyMappingPersistence; /** @@ -270,14 +272,16 @@ List lookupEntityVersions( * @param catalogId catalog id for that entity, NULL_ID if the entity is top-level * @param parentId id of the parent, can be the special 0 value representing the root entity * @param entityType type of entities to list + * @param pageToken the token to start listing after * @return the list of entities for the specified list operation */ @Nonnull - List listEntities( + PolarisPage listEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, - @Nonnull PolarisEntityType entityType); + @Nonnull PolarisEntityType entityType, + @Nonnull PageToken pageToken); /** * List entities where some predicate returns true @@ -288,15 +292,17 @@ List listEntities( * @param entityType type of entities to list * @param entityFilter the filter to be applied to each entity. Only entities where the predicate * returns true are returned in the list + * @param pageToken the token to start listing after * @return the list of entities for which the predicate returns true */ @Nonnull - List listEntities( + PolarisPage listEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, - @Nonnull Predicate entityFilter); + @Nonnull Predicate entityFilter, + @Nonnull PageToken pageToken); /** * List entities where some predicate returns true and transform the entities with a function @@ -313,14 +319,14 @@ List listEntities( * @return the list of entities for which the predicate returns true */ @Nonnull - List listEntities( + PolarisPage listEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, - int limit, @Nonnull Predicate entityFilter, - @Nonnull Function transformer); + @Nonnull Function transformer, + PageToken pageToken); /** * Lookup the current entityGrantRecordsVersion for the specified entity. That version is changed @@ -408,4 +414,7 @@ boolean hasChildren( default BasePersistence detach() { return this; } + + @Nonnull + PageToken.PageTokenBuilder pageTokenBuilder(); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java index da2ab521e1..2a20ad5c1e 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java @@ -42,6 +42,7 @@ import org.apache.polaris.core.persistence.dao.entity.GenerateEntityIdResult; import org.apache.polaris.core.persistence.dao.entity.ListEntitiesResult; import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult; +import org.apache.polaris.core.persistence.pagination.PageToken; import org.apache.polaris.core.policy.PolarisPolicyMappingManager; import org.apache.polaris.core.storage.PolarisCredentialVendor; @@ -120,7 +121,8 @@ ListEntitiesResult listEntities( @Nonnull PolarisCallContext callCtx, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, - @Nonnull PolarisEntitySubType entitySubType); + @Nonnull PolarisEntitySubType entitySubType, + @Nonnull PageToken pageToken); /** * Generate a new unique id that can be used by the Polaris client when it needs to create a new @@ -300,11 +302,12 @@ EntityResult loadEntity( * * @param callCtx call context * @param executorId executor id - * @param limit limit + * @param pageToken page token to start after * @return list of tasks to be completed */ @Nonnull - EntitiesResult loadTasks(@Nonnull PolarisCallContext callCtx, String executorId, int limit); + EntitiesResult loadTasks( + @Nonnull PolarisCallContext callCtx, String executorId, PageToken pageToken); /** * Load change tracking information for a set of entities in one single shot and return for each diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java index 9b5a6b6dbe..b7ba47e83e 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java @@ -50,6 +50,7 @@ import org.apache.polaris.core.persistence.dao.entity.PrivilegeResult; import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult; import org.apache.polaris.core.persistence.dao.entity.ScopedCredentialsResult; +import org.apache.polaris.core.persistence.pagination.PageToken; import org.apache.polaris.core.policy.PolicyEntity; import org.apache.polaris.core.policy.PolicyType; @@ -118,7 +119,8 @@ public ListEntitiesResult listEntities( @Nonnull PolarisCallContext callCtx, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, - @Nonnull PolarisEntitySubType entitySubType) { + @Nonnull PolarisEntitySubType entitySubType, + @Nonnull PageToken pageToken) { callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "listEntities"); return null; } @@ -320,7 +322,7 @@ public EntityResult loadEntity( @Override public EntitiesResult loadTasks( - @Nonnull PolarisCallContext callCtx, String executorId, int limit) { + @Nonnull PolarisCallContext callCtx, String executorId, PageToken pageToken) { callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "loadTasks"); return null; } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/EntitiesResult.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/EntitiesResult.java index 70d9edcf58..3e311b9102 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/EntitiesResult.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/EntitiesResult.java @@ -23,13 +23,21 @@ import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; import java.util.List; +import java.util.Optional; import org.apache.polaris.core.entity.PolarisBaseEntity; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; /** a set of returned entities result */ public class EntitiesResult extends BaseResult { // null if not success. Else the list of entities being returned private final List entities; + private final Optional pageTokenOpt; + + public static EntitiesResult fromPolarisPage(PolarisPage polarisPage) { + return new EntitiesResult(polarisPage.data, Optional.ofNullable(polarisPage.pageToken)); + } /** * Constructor for an error @@ -40,6 +48,11 @@ public class EntitiesResult extends BaseResult { public EntitiesResult(@Nonnull ReturnStatus errorStatus, @Nullable String extraInformation) { super(errorStatus, extraInformation); this.entities = null; + this.pageTokenOpt = Optional.empty(); + } + + public EntitiesResult(@Nonnull List entities) { + this(entities, Optional.empty()); } /** @@ -47,21 +60,29 @@ public EntitiesResult(@Nonnull ReturnStatus errorStatus, @Nullable String extraI * * @param entities list of entities being returned, implies success */ - public EntitiesResult(@Nonnull List entities) { + public EntitiesResult( + @Nonnull List entities, @Nonnull Optional pageTokenOpt) { super(ReturnStatus.SUCCESS); this.entities = entities; + this.pageTokenOpt = pageTokenOpt; } @JsonCreator private EntitiesResult( @JsonProperty("returnStatus") @Nonnull ReturnStatus returnStatus, @JsonProperty("extraInformation") String extraInformation, - @JsonProperty("entities") List entities) { + @JsonProperty("entities") List entities, + @JsonProperty("pageToken") Optional pageTokenOpt) { super(returnStatus, extraInformation); this.entities = entities; + this.pageTokenOpt = pageTokenOpt; } public List getEntities() { return entities; } + + public Optional getPageToken() { + return pageTokenOpt; + } } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ListEntitiesResult.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ListEntitiesResult.java index bc51f4dab5..9d66d5ae51 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ListEntitiesResult.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ListEntitiesResult.java @@ -23,13 +23,23 @@ import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; import java.util.List; +import java.util.Optional; import org.apache.polaris.core.entity.EntityNameLookupRecord; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; /** the return the result for a list entities call */ public class ListEntitiesResult extends BaseResult { // null if not success. Else the list of entities being returned private final List entities; + private final Optional pageTokenOpt; + + /** Create a {@link ListEntitiesResult} from a {@link PolarisPage} */ + public static ListEntitiesResult fromPolarisPage( + PolarisPage polarisPage) { + return new ListEntitiesResult(polarisPage.data, Optional.ofNullable(polarisPage.pageToken)); + } /** * Constructor for an error @@ -37,9 +47,13 @@ public class ListEntitiesResult extends BaseResult { * @param errorCode error code, cannot be SUCCESS * @param extraInformation extra information */ - public ListEntitiesResult(@Nonnull ReturnStatus errorCode, @Nullable String extraInformation) { + public ListEntitiesResult( + @Nonnull ReturnStatus errorCode, + @Nullable String extraInformation, + @Nonnull Optional pageTokenOpt) { super(errorCode, extraInformation); this.entities = null; + this.pageTokenOpt = pageTokenOpt; } /** @@ -47,21 +61,29 @@ public ListEntitiesResult(@Nonnull ReturnStatus errorCode, @Nullable String extr * * @param entities list of entities being returned, implies success */ - public ListEntitiesResult(@Nonnull List entities) { + public ListEntitiesResult( + @Nonnull List entities, @Nonnull Optional pageTokenOpt) { super(ReturnStatus.SUCCESS); this.entities = entities; + this.pageTokenOpt = pageTokenOpt; } @JsonCreator private ListEntitiesResult( @JsonProperty("returnStatus") @Nonnull ReturnStatus returnStatus, @JsonProperty("extraInformation") String extraInformation, - @JsonProperty("entities") List entities) { + @JsonProperty("entities") List entities, + @JsonProperty("pageToken") Optional pageTokenOpt) { super(returnStatus, extraInformation); this.entities = entities; + this.pageTokenOpt = pageTokenOpt; } public List getEntities() { return entities; } + + public Optional getPageToken() { + return pageTokenOpt; + } } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/OffsetPageToken.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/OffsetPageToken.java new file mode 100644 index 0000000000..ea0e8be1c7 --- /dev/null +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/OffsetPageToken.java @@ -0,0 +1,113 @@ +/* + * 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.persistence.pagination; + +import java.util.List; + +/** + * A simple {@link PageToken} implementation that tracks the number of records returned. Entities + * are meant to be filtered during listing such that when a token with offset N is supplied, the + * first N records are omitted from the results. + */ +public class OffsetPageToken extends PageToken { + + /** + * The offset of the token. If this is `5` for example, the first 5 entities returned by a list + * operation that uses this token will be skipped. + */ + public final int offset; + + /** The offset to use to start with. */ + private static final int BASE_OFFSET = 0; + + private OffsetPageToken(int offset, int pageSize) { + this.offset = offset; + this.pageSize = pageSize; + validate(); + } + + @Override + protected void validate() { + if (offset < 0) { + throw new IllegalArgumentException("Offset must be greater than zero"); + } + super.validate(); + } + + /** Get a new `EntityIdPageTokenBuilder` instance */ + public static PageTokenBuilder builder() { + return new OffsetPageTokenBuilder(); + } + + @Override + protected PageTokenBuilder getBuilder() { + return OffsetPageToken.builder(); + } + + @Override + protected List getComponents() { + return List.of(String.valueOf(this.offset), String.valueOf(this.pageSize)); + } + + /** A {@link PageTokenBuilder} implementation for {@link OffsetPageToken} */ + public static class OffsetPageTokenBuilder extends PageTokenBuilder { + + private OffsetPageTokenBuilder() {} + + @Override + public String tokenPrefix() { + return "polaris-offset"; + } + + @Override + public int expectedComponents() { + // offset + limit + return 2; + } + + @Override + protected OffsetPageToken fromStringComponents(List components) { + return new OffsetPageToken( + Integer.parseInt(components.get(0)), Integer.parseInt(components.get(1))); + } + + @Override + protected OffsetPageToken fromLimitImpl(int limit) { + return new OffsetPageToken(BASE_OFFSET, limit); + } + } + + @Override + public PageToken updated(List newData) { + if (newData == null || newData.size() < this.pageSize) { + return PageToken.DONE; + } else { + return new OffsetPageToken(this.offset + newData.size(), pageSize); + } + } + + @Override + public OffsetPageToken withPageSize(Integer pageSize) { + if (pageSize == null) { + return new OffsetPageToken(BASE_OFFSET, this.pageSize); + } else { + return new OffsetPageToken(this.offset, pageSize); + } + } +} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/PageToken.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/PageToken.java new file mode 100644 index 0000000000..a9ba699877 --- /dev/null +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/PageToken.java @@ -0,0 +1,179 @@ +/* + * 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.persistence.pagination; + +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Base64; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +/** + * Represents a page token that can be used by operations like `listTables`. Clients that specify a + * `pageSize` (or a `pageToken`) may receive a `next-page-token` in the response, the content of + * which is a serialized PageToken. + * + *

By providing that in the next query's `pageToken`, the client can resume listing where they + * left off. If the client provides a `pageToken` or `pageSize` but `next-page-token` is null in the + * response, that means there is no more data to read. + */ +public abstract class PageToken { + + public int pageSize; + + public static final PageToken DONE = null; + public static final int DEFAULT_PAGE_SIZE = 1000; + + protected void validate() { + if (pageSize <= 0) { + throw new IllegalArgumentException("Page size must be greater than zero"); + } + } + + /** + * Get a new PageTokenBuilder from a PageToken. The PageTokenBuilder type should match the + * PageToken type. Implementations may also provide a static `builder` method to obtain the same + * PageTokenBuilder. + */ + protected abstract PageTokenBuilder getBuilder(); + + /** Allows `PageToken` implementations to implement methods like `fromLimit` */ + public abstract static class PageTokenBuilder { + + /** + * A prefix that tokens are expected to start with, ideally unique across `PageTokenBuilder` + * implementations. + */ + public abstract String tokenPrefix(); + + /** + * The number of expected components in a token. This should match the number of components + * returned by getComponents and shouldn't account for the prefix or the checksum. + */ + public abstract int expectedComponents(); + + /** Deserialize a string into a {@link PageToken} */ + public final PageToken fromString(String tokenString) { + if (tokenString == null) { + throw new IllegalArgumentException("Cannot build page token from null string"); + } else if (tokenString.isEmpty()) { + if (this instanceof ReadEverythingPageToken.ReadEverythingPageTokenBuilder) { + return ReadEverythingPageToken.get(); + } else { + return fromLimit(DEFAULT_PAGE_SIZE); + } + } else { + try { + String decoded = + new String(Base64.getDecoder().decode(tokenString), StandardCharsets.UTF_8); + String[] parts = decoded.split(":"); + + // +2 to account for the prefix and checksum. + if (parts.length != expectedComponents() + 2 || !parts[0].equals(tokenPrefix())) { + throw new IllegalArgumentException("Invalid token format in token: " + tokenString); + } + + // Cut off prefix and checksum + T result = fromStringComponents(Arrays.asList(parts).subList(1, parts.length - 1)); + result.validate(); + return result; + } catch (Exception e) { + throw new IllegalArgumentException("Failed to decode page token: " + tokenString, e); + } + } + } + + /** Construct a {@link PageToken} from a plain limit */ + public final PageToken fromLimit(Integer limit) { + if (limit == null) { + return ReadEverythingPageToken.get(); + } else { + return fromLimitImpl(limit); + } + } + + /** Construct a {@link PageToken} from a plain limit */ + protected abstract T fromLimitImpl(int limit); + + /** + * {@link PageTokenBuilder} implementations should implement this to build a {@link PageToken} + * from components in a string token. These components should be the same ones returned by + * {@link #getComponents()} and won't include the token prefix or the checksum. + */ + protected abstract T fromStringComponents(List components); + } + + /** Convert this into components that the serialized token string will be built from. */ + protected abstract List getComponents(); + + /** + * Builds a new page token to reflect new data that's been read. If the amount of data read is + * less than the pageSize, this will return {@link PageToken#DONE}(null) + */ + protected abstract PageToken updated(List newData); + + /** + * Builds a {@link PolarisPage} from a {@link List}. The {@link PageToken} attached to the + * new {@link PolarisPage} is the same as the result of calling {@link #updated(List)} on this + * {@link PageToken}. + */ + public final PolarisPage buildNextPage(List data) { + return new PolarisPage(updated(data), data); + } + + /** + * Return a new {@link PageToken} with an updated page size. If the pageSize provided is null, the + * existing page size will be preserved. + */ + public abstract PageToken withPageSize(Integer pageSize); + + /** Serialize a {@link PageToken} into a string */ + @Override + public final String toString() { + List components = getComponents(); + String prefix = getBuilder().tokenPrefix(); + String componentString = String.join(":", components); + String checksum = String.valueOf(componentString.hashCode()); + List allElements = + Stream.of(prefix, componentString, checksum) + .filter(s -> !s.isEmpty()) + .collect(Collectors.toList()); + String rawString = String.join(":", allElements); + return Base64.getEncoder().encodeToString(rawString.getBytes(StandardCharsets.UTF_8)); + } + + @Override + public final boolean equals(Object o) { + if (o instanceof PageToken) { + return this.toString().equals(o.toString()); + } else { + return false; + } + } + + @Override + public final int hashCode() { + if (toString() == null) { + return 0; + } else { + return toString().hashCode(); + } + } +} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/PolarisPage.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/PolarisPage.java new file mode 100644 index 0000000000..085f13ff18 --- /dev/null +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/PolarisPage.java @@ -0,0 +1,42 @@ +/* + * 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.persistence.pagination; + +import java.util.List; + +/** + * A wrapper for a {@link List} of data and a {@link PageToken} that can be used to continue the + * listing operation that generated that data. + */ +public class PolarisPage { + public final PageToken pageToken; + public final List data; + + public PolarisPage(PageToken pageToken, List data) { + this.pageToken = pageToken; + this.data = data; + } + + /** + * Used to wrap a {@link List} of data into a {@link PolarisPage} when there is no more data + */ + public static PolarisPage fromData(List data) { + return new PolarisPage<>(PageToken.DONE, data); + } +} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/ReadEverythingPageToken.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/ReadEverythingPageToken.java new file mode 100644 index 0000000000..8a9edd3db9 --- /dev/null +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/pagination/ReadEverythingPageToken.java @@ -0,0 +1,95 @@ +/* + * 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.persistence.pagination; + +import java.util.List; + +/** + * A {@link PageToken} implementation for readers who want to read everything. The behavior when + * using this token should be the same as when reading without a token. + */ +public class ReadEverythingPageToken extends PageToken { + + private ReadEverythingPageToken() { + this.pageSize = Integer.MAX_VALUE; + validate(); + } + + /** Get a {@link ReadEverythingPageToken} */ + public static ReadEverythingPageToken get() { + return new ReadEverythingPageToken(); + } + + public static PageTokenBuilder builder() { + return new ReadEverythingPageTokenBuilder(); + } + + @Override + protected PageTokenBuilder getBuilder() { + return builder(); + } + + /** A {@link PageTokenBuilder} implementation for {@link ReadEverythingPageToken} */ + public static class ReadEverythingPageTokenBuilder + extends PageTokenBuilder { + + private ReadEverythingPageTokenBuilder() {} + + @Override + public String tokenPrefix() { + return "polaris-read-everything"; + } + + @Override + public int expectedComponents() { + return 0; + } + + @Override + protected ReadEverythingPageToken fromStringComponents(List components) { + return ReadEverythingPageToken.get(); + } + + @Override + protected ReadEverythingPageToken fromLimitImpl(int limit) { + throw new UnsupportedOperationException(); + } + } + + @Override + protected List getComponents() { + return List.of(); + } + + /** Any time {@link ReadEverythingPageToken} is updated, everything has been read */ + @Override + public PageToken updated(List newData) { + return PageToken.DONE; + } + + /** {@link ReadEverythingPageToken} does not support page size */ + @Override + public PageToken withPageSize(Integer pageSize) { + if (pageSize == null || pageSize == Integer.MAX_VALUE) { + return ReadEverythingPageToken.get(); + } else { + throw new UnsupportedOperationException(); + } + } +} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/AbstractTransactionalPersistence.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/AbstractTransactionalPersistence.java index e949b33fe2..3cda5cf5ed 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/AbstractTransactionalPersistence.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/AbstractTransactionalPersistence.java @@ -37,6 +37,8 @@ import org.apache.polaris.core.persistence.EntityAlreadyExistsException; import org.apache.polaris.core.persistence.PolicyMappingAlreadyExistsException; import org.apache.polaris.core.persistence.RetryOnConcurrencyException; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; import org.apache.polaris.core.policy.PolicyType; import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; @@ -352,46 +354,50 @@ public List lookupEntityVersions( /** {@inheritDoc} */ @Override @Nonnull - public List listEntities( + public PolarisPage listEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, - @Nonnull PolarisEntityType entityType) { + @Nonnull PolarisEntityType entityType, + @Nonnull PageToken pageToken) { return runInReadTransaction( - callCtx, () -> this.listEntitiesInCurrentTxn(callCtx, catalogId, parentId, entityType)); + callCtx, + () -> this.listEntitiesInCurrentTxn(callCtx, catalogId, parentId, entityType, pageToken)); } /** {@inheritDoc} */ @Override @Nonnull - public List listEntities( + public PolarisPage listEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, - @Nonnull Predicate entityFilter) { + @Nonnull Predicate entityFilter, + @Nonnull PageToken pageToken) { return runInReadTransaction( callCtx, () -> - this.listEntitiesInCurrentTxn(callCtx, catalogId, parentId, entityType, entityFilter)); + this.listEntitiesInCurrentTxn( + callCtx, catalogId, parentId, entityType, entityFilter, pageToken)); } /** {@inheritDoc} */ @Override @Nonnull - public List listEntities( + public PolarisPage listEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, - int limit, @Nonnull Predicate entityFilter, - @Nonnull Function transformer) { + @Nonnull Function transformer, + @Nonnull PageToken pageToken) { return runInReadTransaction( callCtx, () -> this.listEntitiesInCurrentTxn( - callCtx, catalogId, parentId, entityType, limit, entityFilter, transformer)); + callCtx, catalogId, parentId, entityType, entityFilter, transformer, pageToken)); } /** {@inheritDoc} */ diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java index 22120c7a3d..a38a3182c8 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java @@ -28,6 +28,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.Set; import java.util.function.Function; import java.util.stream.Collectors; @@ -63,6 +64,8 @@ import org.apache.polaris.core.persistence.dao.entity.PrivilegeResult; import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult; import org.apache.polaris.core.persistence.dao.entity.ScopedCredentialsResult; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; import org.apache.polaris.core.policy.PolicyEntity; import org.apache.polaris.core.policy.PolicyType; @@ -654,37 +657,41 @@ private void bootstrapPolarisService( } /** - * See {@link #listEntities(PolarisCallContext, List, PolarisEntityType, PolarisEntitySubType)} + * See {@link #listEntities(PolarisCallContext, List, PolarisEntityType, PolarisEntitySubType, + * PageToken)} */ private @Nonnull ListEntitiesResult listEntities( @Nonnull PolarisCallContext callCtx, @Nonnull TransactionalPersistence ms, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, - @Nonnull PolarisEntitySubType entitySubType) { + @Nonnull PolarisEntitySubType entitySubType, + @Nonnull PageToken pageToken) { // first resolve again the catalogPath to that entity PolarisEntityResolver resolver = new PolarisEntityResolver(callCtx, ms, catalogPath); // return if we failed to resolve if (resolver.isFailure()) { - return new ListEntitiesResult(BaseResult.ReturnStatus.CATALOG_PATH_CANNOT_BE_RESOLVED, null); + return new ListEntitiesResult( + BaseResult.ReturnStatus.CATALOG_PATH_CANNOT_BE_RESOLVED, null, Optional.empty()); } // return list of active entities - List toreturnList = + PolarisPage resultPage = ms.listEntitiesInCurrentTxn( - callCtx, resolver.getCatalogIdOrNull(), resolver.getParentId(), entityType); + callCtx, resolver.getCatalogIdOrNull(), resolver.getParentId(), entityType, pageToken); // prune the returned list with only entities matching the entity subtype if (entitySubType != PolarisEntitySubType.ANY_SUBTYPE) { - toreturnList = - toreturnList.stream() - .filter(rec -> rec.getSubTypeCode() == entitySubType.getCode()) - .collect(Collectors.toList()); + resultPage = + pageToken.buildNextPage( + resultPage.data.stream() + .filter(rec -> rec.getSubTypeCode() == entitySubType.getCode()) + .collect(Collectors.toList())); } // done - return new ListEntitiesResult(toreturnList); + return ListEntitiesResult.fromPolarisPage(resultPage); } /** {@inheritDoc} */ @@ -693,13 +700,15 @@ private void bootstrapPolarisService( @Nonnull PolarisCallContext callCtx, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, - @Nonnull PolarisEntitySubType entitySubType) { + @Nonnull PolarisEntitySubType entitySubType, + @Nonnull PageToken pageToken) { // get meta store we should be using TransactionalPersistence ms = ((TransactionalPersistence) callCtx.getMetaStore()); // run operation in a read transaction return ms.runInReadTransaction( - callCtx, () -> listEntities(callCtx, ms, catalogPath, entityType, entitySubType)); + callCtx, + () -> listEntities(callCtx, ms, catalogPath, entityType, entitySubType, pageToken)); } /** {@link #createPrincipal(PolarisCallContext, PolarisBaseEntity)} */ @@ -1336,13 +1345,14 @@ private void bootstrapPolarisService( // get the list of catalog roles, at most 2 List catalogRoles = ms.listEntitiesInCurrentTxn( - callCtx, - catalogId, - catalogId, - PolarisEntityType.CATALOG_ROLE, - 2, - entity -> true, - Function.identity()); + callCtx, + catalogId, + catalogId, + PolarisEntityType.CATALOG_ROLE, + entity -> true, + Function.identity(), + ms.pageTokenBuilder().fromLimit(2)) + .data; // if we have 2, we cannot drop the catalog. If only one left, better be the admin role if (catalogRoles.size() > 1) { @@ -1884,21 +1894,20 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( () -> this.loadEntity(callCtx, ms, entityCatalogId, entityId, entityType.getCode())); } - /** Refer to {@link #loadTasks(PolarisCallContext, String, int)} */ + /** Refer to {@link #loadTasks(PolarisCallContext, String, PageToken)} */ private @Nonnull EntitiesResult loadTasks( @Nonnull PolarisCallContext callCtx, @Nonnull TransactionalPersistence ms, String executorId, - int limit) { + PageToken pageToken) { // find all available tasks - List availableTasks = + PolarisPage availableTasks = ms.listEntitiesInCurrentTxn( callCtx, PolarisEntityConstants.getRootEntityId(), PolarisEntityConstants.getRootEntityId(), PolarisEntityType.TASK, - limit, entity -> { PolarisObjectMapperUtil.TaskExecutionState taskState = PolarisObjectMapperUtil.parseTaskState(entity); @@ -1913,10 +1922,11 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( || taskState.executor == null || callCtx.getClock().millis() - taskState.lastAttemptStartTime > taskAgeTimeout; }, - Function.identity()); + Function.identity(), + pageToken); List loadedTasks = new ArrayList<>(); - availableTasks.forEach( + availableTasks.data.forEach( task -> { // Make a copy to avoid mutating someone else's reference. // TODO: Refactor into immutable/Builder pattern. @@ -1947,14 +1957,14 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( result.getReturnStatus(), result.getExtraInformation()); } }); - return new EntitiesResult(loadedTasks); + return EntitiesResult.fromPolarisPage(PolarisPage.fromData(loadedTasks)); } @Override public @Nonnull EntitiesResult loadTasks( - @Nonnull PolarisCallContext callCtx, String executorId, int limit) { + @Nonnull PolarisCallContext callCtx, String executorId, PageToken pageToken) { TransactionalPersistence ms = ((TransactionalPersistence) callCtx.getMetaStore()); - return ms.runInTransaction(callCtx, () -> this.loadTasks(callCtx, ms, executorId, limit)); + return ms.runInTransaction(callCtx, () -> this.loadTasks(callCtx, ms, executorId, pageToken)); } /** {@inheritDoc} */ diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalPersistence.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalPersistence.java index 2057991db0..6835caaccb 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalPersistence.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalPersistence.java @@ -36,6 +36,8 @@ import org.apache.polaris.core.entity.PolarisPrincipalSecrets; import org.apache.polaris.core.persistence.BasePersistence; import org.apache.polaris.core.persistence.IntegrationPersistence; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; import org.apache.polaris.core.policy.TransactionalPolicyMappingPersistence; import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; import org.apache.polaris.core.storage.PolarisStorageIntegration; @@ -201,31 +203,33 @@ List lookupEntityVersionsInCurrentTxn( /** See {@link org.apache.polaris.core.persistence.BasePersistence#listEntities} */ @Nonnull - List listEntitiesInCurrentTxn( + PolarisPage listEntitiesInCurrentTxn( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, - @Nonnull PolarisEntityType entityType); + @Nonnull PolarisEntityType entityType, + @Nonnull PageToken pageToken); /** See {@link org.apache.polaris.core.persistence.BasePersistence#listEntities} */ @Nonnull - List listEntitiesInCurrentTxn( + PolarisPage listEntitiesInCurrentTxn( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, - @Nonnull Predicate entityFilter); + @Nonnull Predicate entityFilter, + @Nonnull PageToken pageToken); /** See {@link org.apache.polaris.core.persistence.BasePersistence#listEntities} */ @Nonnull - List listEntitiesInCurrentTxn( + PolarisPage listEntitiesInCurrentTxn( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, - int limit, @Nonnull Predicate entityFilter, - @Nonnull Function transformer); + @Nonnull Function transformer, + @Nonnull PageToken pageToken); /** * See {@link org.apache.polaris.core.persistence.BasePersistence#lookupEntityGrantRecordsVersion} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TreeMapTransactionalPersistenceImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TreeMapTransactionalPersistenceImpl.java index e9cbd53f39..a2a92d7ad7 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TreeMapTransactionalPersistenceImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TreeMapTransactionalPersistenceImpl.java @@ -21,11 +21,13 @@ import com.google.common.base.Predicates; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; +import java.util.Comparator; import java.util.List; import java.util.function.Function; import java.util.function.Predicate; import java.util.function.Supplier; import java.util.stream.Collectors; +import java.util.stream.Stream; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.PolarisBaseEntity; @@ -38,6 +40,10 @@ import org.apache.polaris.core.entity.PolarisPrincipalSecrets; import org.apache.polaris.core.persistence.BaseMetaStoreManager; import org.apache.polaris.core.persistence.PrincipalSecretsGenerator; +import org.apache.polaris.core.persistence.pagination.OffsetPageToken; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; import org.apache.polaris.core.storage.PolarisStorageIntegration; @@ -301,29 +307,30 @@ public List lookupEntityActiveBatchInCurrentTxn( /** {@inheritDoc} */ @Override - public @Nonnull List listEntitiesInCurrentTxn( + public @Nonnull PolarisPage listEntitiesInCurrentTxn( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, - @Nonnull PolarisEntityType entityType) { + @Nonnull PolarisEntityType entityType, + @Nonnull PageToken pageToken) { return this.listEntitiesInCurrentTxn( - callCtx, catalogId, parentId, entityType, Predicates.alwaysTrue()); + callCtx, catalogId, parentId, entityType, Predicates.alwaysTrue(), pageToken); } @Override - public @Nonnull List listEntitiesInCurrentTxn( + public @Nonnull PolarisPage listEntitiesInCurrentTxn( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, - @Nonnull Predicate entityFilter) { + @Nonnull Predicate entityFilter, + @Nonnull PageToken pageToken) { // full range scan under the parent for that type return this.listEntitiesInCurrentTxn( callCtx, catalogId, parentId, entityType, - Integer.MAX_VALUE, entityFilter, entity -> new EntityNameLookupRecord( @@ -332,31 +339,46 @@ public List lookupEntityActiveBatchInCurrentTxn( entity.getParentId(), entity.getName(), entity.getTypeCode(), - entity.getSubTypeCode())); + entity.getSubTypeCode()), + pageToken); } @Override - public @Nonnull List listEntitiesInCurrentTxn( + public @Nonnull PolarisPage listEntitiesInCurrentTxn( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, - int limit, @Nonnull Predicate entityFilter, - @Nonnull Function transformer) { - // full range scan under the parent for that type - return this.store - .getSliceEntitiesActive() - .readRange(this.store.buildPrefixKeyComposite(catalogId, parentId, entityType.getCode())) - .stream() - .map( - nameRecord -> - this.lookupEntityInCurrentTxn( - callCtx, catalogId, nameRecord.getId(), entityType.getCode())) - .filter(entityFilter) - .limit(limit) - .map(transformer) - .collect(Collectors.toList()); + @Nonnull Function transformer, + @Nonnull PageToken pageToken) { + if (!(pageToken instanceof ReadEverythingPageToken) + && !(pageToken instanceof OffsetPageToken)) { + throw new IllegalArgumentException("Unexpected pageToken: " + pageToken); + } + + Stream partialResults = + this.store + .getSliceEntitiesActive() + .readRange( + this.store.buildPrefixKeyComposite(catalogId, parentId, entityType.getCode())) + .stream() + .map( + nameRecord -> + this.lookupEntityInCurrentTxn( + callCtx, catalogId, nameRecord.getId(), entityType.getCode())) + .filter(entityFilter); + + if (pageToken instanceof OffsetPageToken) { + partialResults = + partialResults + .sorted(Comparator.comparingLong(PolarisEntityCore::getId)) + .skip(((OffsetPageToken) pageToken).offset) + .limit(pageToken.pageSize); + } + + List entities = partialResults.map(transformer).collect(Collectors.toList()); + return pageToken.buildNextPage(entities); } /** {@inheritDoc} */ @@ -558,6 +580,11 @@ public void rollback() { this.store.rollback(); } + @Override + public @Nonnull PageToken.PageTokenBuilder pageTokenBuilder() { + return OffsetPageToken.builder(); + } + /** {@inheritDoc} */ @Override public void writeToPolicyMappingRecordsInCurrentTxn( diff --git a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java index d11bed82d8..e8029bc4a8 100644 --- a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java +++ b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java @@ -43,6 +43,7 @@ import org.apache.polaris.core.entity.PolarisEntitySubType; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.TaskEntity; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; import org.assertj.core.api.Assertions; import org.assertj.core.api.InstanceOfAssertFactories; import org.junit.jupiter.api.BeforeEach; @@ -128,7 +129,8 @@ protected void testCreateEntities() { polarisTestMetaStoreManager.polarisCallContext, null, PolarisEntityType.TASK, - PolarisEntitySubType.NULL_SUBTYPE) + PolarisEntitySubType.NULL_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities(); Assertions.assertThat(listedEntities) .isNotNull() @@ -302,7 +304,9 @@ protected void testLoadTasks() { PolarisMetaStoreManager metaStoreManager = polarisTestMetaStoreManager.polarisMetaStoreManager; PolarisCallContext callCtx = polarisTestMetaStoreManager.polarisCallContext; List taskList = - metaStoreManager.loadTasks(callCtx, executorId, 5).getEntities(); + metaStoreManager + .loadTasks(callCtx, executorId, callCtx.getMetaStore().pageTokenBuilder().fromLimit(5)) + .getEntities(); Assertions.assertThat(taskList) .isNotNull() .isNotEmpty() @@ -322,7 +326,9 @@ protected void testLoadTasks() { // grab a second round of tasks. Assert that none of the original 5 are in the list List newTaskList = - metaStoreManager.loadTasks(callCtx, executorId, 5).getEntities(); + metaStoreManager + .loadTasks(callCtx, executorId, callCtx.getMetaStore().pageTokenBuilder().fromLimit(5)) + .getEntities(); Assertions.assertThat(newTaskList) .isNotNull() .isNotEmpty() @@ -336,7 +342,9 @@ protected void testLoadTasks() { // only 10 tasks are unassigned. Requesting 20, we should only receive those 10 List lastTen = - metaStoreManager.loadTasks(callCtx, executorId, 20).getEntities(); + metaStoreManager + .loadTasks(callCtx, executorId, callCtx.getMetaStore().pageTokenBuilder().fromLimit(20)) + .getEntities(); Assertions.assertThat(lastTen) .isNotNull() @@ -350,7 +358,9 @@ protected void testLoadTasks() { .collect(Collectors.toSet()); List emtpyList = - metaStoreManager.loadTasks(callCtx, executorId, 20).getEntities(); + metaStoreManager + .loadTasks(callCtx, executorId, callCtx.getMetaStore().pageTokenBuilder().fromLimit(20)) + .getEntities(); Assertions.assertThat(emtpyList).isNotNull().isEmpty(); @@ -358,7 +368,9 @@ protected void testLoadTasks() { // all the tasks are unassigned. Fetch them all List allTasks = - metaStoreManager.loadTasks(callCtx, executorId, 20).getEntities(); + metaStoreManager + .loadTasks(callCtx, executorId, callCtx.getMetaStore().pageTokenBuilder().fromLimit(20)) + .getEntities(); Assertions.assertThat(allTasks) .isNotNull() @@ -373,7 +385,9 @@ protected void testLoadTasks() { timeSource.add(Duration.ofMinutes(10)); List finalList = - metaStoreManager.loadTasks(callCtx, executorId, 20).getEntities(); + metaStoreManager + .loadTasks(callCtx, executorId, callCtx.getMetaStore().pageTokenBuilder().fromLimit(20)) + .getEntities(); Assertions.assertThat(finalList).isNotNull().isEmpty(); } @@ -401,7 +415,13 @@ protected void testLoadTasksInParallel() throws Exception { do { retry = false; try { - taskList = metaStoreManager.loadTasks(callCtx, executorId, 5).getEntities(); + taskList = + metaStoreManager + .loadTasks( + callCtx, + executorId, + callCtx.getMetaStore().pageTokenBuilder().fromLimit(5)) + .getEntities(); taskList.stream().map(PolarisBaseEntity::getName).forEach(taskNames::add); } catch (RetryOnConcurrencyException e) { retry = true; diff --git a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java index b401d4efb3..86596df02a 100644 --- a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java +++ b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java @@ -51,6 +51,7 @@ import org.apache.polaris.core.persistence.dao.entity.LoadPolicyMappingsResult; import org.apache.polaris.core.persistence.dao.entity.PolicyAttachmentResult; import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; import org.apache.polaris.core.policy.PolicyEntity; import org.apache.polaris.core.policy.PolicyType; @@ -766,7 +767,8 @@ void dropEntity(List catalogPath, PolarisBaseEntity entityToD this.polarisCallContext, path, PolarisEntityType.NAMESPACE, - PolarisEntitySubType.NULL_SUBTYPE) + PolarisEntitySubType.NULL_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities(); Assertions.assertThat(children).isNotNull(); if (children.isEmpty() && entity.getType() == PolarisEntityType.NAMESPACE) { @@ -776,7 +778,8 @@ void dropEntity(List catalogPath, PolarisBaseEntity entityToD this.polarisCallContext, path, PolarisEntityType.TABLE_LIKE, - PolarisEntitySubType.ANY_SUBTYPE) + PolarisEntitySubType.ANY_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities(); Assertions.assertThat(children).isNotNull(); } else if (children.isEmpty()) { @@ -786,7 +789,8 @@ void dropEntity(List catalogPath, PolarisBaseEntity entityToD this.polarisCallContext, path, PolarisEntityType.CATALOG_ROLE, - PolarisEntitySubType.ANY_SUBTYPE) + PolarisEntitySubType.ANY_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities(); Assertions.assertThat(children).isNotNull(); // if only one left, it can be dropped. @@ -1510,7 +1514,12 @@ private void validateListReturn( // list the entities under the specified path List result = polarisMetaStoreManager - .listEntities(this.polarisCallContext, path, entityType, entitySubType) + .listEntities( + this.polarisCallContext, + path, + entityType, + entitySubType, + ReadEverythingPageToken.get()) .getEntities(); Assertions.assertThat(result).isNotNull(); @@ -1827,7 +1836,8 @@ void validateBootstrap() { this.polarisCallContext, null, PolarisEntityType.PRINCIPAL, - PolarisEntitySubType.NULL_SUBTYPE) + PolarisEntitySubType.NULL_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities(); // ensure not null, one element only @@ -1853,7 +1863,8 @@ void validateBootstrap() { this.polarisCallContext, null, PolarisEntityType.PRINCIPAL_ROLE, - PolarisEntitySubType.NULL_SUBTYPE) + PolarisEntitySubType.NULL_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities(); // ensure not null, one element only @@ -2591,7 +2602,8 @@ public void testLookup() { this.polarisCallContext, null, PolarisEntityType.PRINCIPAL, - PolarisEntitySubType.NULL_SUBTYPE) + PolarisEntitySubType.NULL_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities(); // ensure not null, one element only diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/IcebergCatalogTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/IcebergCatalogTest.java index 5a83330a5b..b6f5f460c1 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/IcebergCatalogTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/IcebergCatalogTest.java @@ -62,6 +62,7 @@ import org.apache.iceberg.UpdateSchema; import org.apache.iceberg.catalog.CatalogTests; import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.SupportsNamespaces; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.exceptions.AlreadyExistsException; import org.apache.iceberg.exceptions.BadRequestException; @@ -98,6 +99,7 @@ import org.apache.polaris.core.persistence.dao.entity.BaseResult; import org.apache.polaris.core.persistence.dao.entity.EntityResult; import org.apache.polaris.core.persistence.dao.entity.PrincipalSecretsResult; +import org.apache.polaris.core.persistence.pagination.PolarisPage; import org.apache.polaris.core.persistence.transactional.TransactionalPersistence; import org.apache.polaris.core.secrets.UserSecretsManager; import org.apache.polaris.core.secrets.UserSecretsManagerFactory; @@ -164,7 +166,9 @@ public Map getConfigOverrides() { "polaris.features.defaults.\"INITIALIZE_DEFAULT_CATALOG_FILEIO_FOR_TEST\"", "true", "polaris.features.defaults.\"SUPPORTED_CATALOG_STORAGE_TYPES\"", - "[\"FILE\"]"); + "[\"FILE\"]", + "polaris.features.defaults.\"LIST_PAGINATION_ENABLED\"", + "true"); } } @@ -1517,7 +1521,12 @@ public void testDropTableWithPurge() { .as("Table should not exist after drop") .rejects(TABLE); List tasks = - metaStoreManager.loadTasks(polarisContext, "testExecutor", 1).getEntities(); + metaStoreManager + .loadTasks( + polarisContext, + "testExecutor", + callContext.getPolarisCallContext().getMetaStore().pageTokenBuilder().fromLimit(1)) + .getEntities(); Assertions.assertThat(tasks).hasSize(1); TaskEntity taskEntity = TaskEntity.of(tasks.get(0)); EnumMap credentials = @@ -1722,7 +1731,14 @@ public void testFileIOWrapper() { TaskEntity taskEntity = TaskEntity.of( metaStoreManager - .loadTasks(callContext.getPolarisCallContext(), "testExecutor", 1) + .loadTasks( + callContext.getPolarisCallContext(), + "testExecutor", + callContext + .getPolarisCallContext() + .getMetaStore() + .pageTokenBuilder() + .fromLimit(1)) .getEntities() .getFirst()); Map properties = taskEntity.getInternalPropertiesAsMap(); @@ -1918,4 +1934,130 @@ public void testTableOperationsDoesNotRefreshAfterCommit(boolean updateMetadataO private static InMemoryFileIO getInMemoryIo(IcebergCatalog catalog) { return (InMemoryFileIO) ((ExceptionMappingFileIO) catalog.getIo()).getInnerIo(); } + + @Test + public void testPaginatedListTables() { + if (this.requiresNamespaceCreate()) { + ((SupportsNamespaces) catalog).createNamespace(NS); + } + + for (int i = 0; i < 5; i++) { + catalog.buildTable(TableIdentifier.of(NS, "pagination_table_" + i), SCHEMA).create(); + } + + try { + // List without pagination + Assertions.assertThat(catalog.listTables(NS)).isNotNull().hasSize(5); + + // List with a limit: + PolarisPage firstListResult = + catalog.listTables(NS, polarisContext.getMetaStore().pageTokenBuilder().fromLimit(2)); + Assertions.assertThat(firstListResult.data.size()).isEqualTo(2); + Assertions.assertThat(firstListResult.pageToken.toString()).isNotNull().isNotEmpty(); + + // List using the previously obtained token: + PolarisPage secondListResult = catalog.listTables(NS, firstListResult.pageToken); + Assertions.assertThat(secondListResult.data.size()).isEqualTo(2); + Assertions.assertThat(secondListResult.pageToken.toString()).isNotNull().isNotEmpty(); + + // List using the final token: + PolarisPage finalListResult = catalog.listTables(NS, secondListResult.pageToken); + Assertions.assertThat(finalListResult.data.size()).isEqualTo(1); + Assertions.assertThat(finalListResult.pageToken).isNull(); + } finally { + for (int i = 0; i < 5; i++) { + catalog.dropTable(TableIdentifier.of(NS, "pagination_table_" + i)); + } + } + } + + @Test + public void testPaginatedListViews() { + if (this.requiresNamespaceCreate()) { + ((SupportsNamespaces) catalog).createNamespace(NS); + } + + for (int i = 0; i < 5; i++) { + catalog + .buildView(TableIdentifier.of(NS, "pagination_view_" + i)) + .withQuery("a_" + i, "SELECT 1 id") + .withSchema(SCHEMA) + .withDefaultNamespace(NS) + .create(); + } + + try { + // List without pagination + Assertions.assertThat(catalog.listViews(NS)).isNotNull().hasSize(5); + + // List with a limit: + PolarisPage firstListResult = + catalog.listViews(NS, polarisContext.getMetaStore().pageTokenBuilder().fromLimit(2)); + Assertions.assertThat(firstListResult.data.size()).isEqualTo(2); + Assertions.assertThat(firstListResult.pageToken.toString()).isNotNull().isNotEmpty(); + + // List using the previously obtained token: + PolarisPage secondListResult = catalog.listViews(NS, firstListResult.pageToken); + Assertions.assertThat(secondListResult.data.size()).isEqualTo(2); + Assertions.assertThat(secondListResult.pageToken.toString()).isNotNull().isNotEmpty(); + + // List using the final token: + PolarisPage finalListResult = catalog.listViews(NS, secondListResult.pageToken); + Assertions.assertThat(finalListResult.data.size()).isEqualTo(1); + Assertions.assertThat(finalListResult.pageToken).isNull(); + } finally { + for (int i = 0; i < 5; i++) { + catalog.dropTable(TableIdentifier.of(NS, "pagination_view_" + i)); + } + } + } + + @Test + public void testPaginatedListNamespaces() { + for (int i = 0; i < 5; i++) { + catalog.createNamespace(Namespace.of("pagination_namespace_" + i)); + } + + try { + // List without pagination + Assertions.assertThat(catalog.listNamespaces()).isNotNull().hasSize(5); + + // List with a limit: + PolarisPage firstListResult = + catalog.listNamespaces(polarisContext.getMetaStore().pageTokenBuilder().fromLimit(2)); + Assertions.assertThat(firstListResult.data.size()).isEqualTo(2); + Assertions.assertThat(firstListResult.pageToken.toString()).isNotNull().isNotEmpty(); + + // List using the previously obtained token: + PolarisPage secondListResult = catalog.listNamespaces(firstListResult.pageToken); + Assertions.assertThat(secondListResult.data.size()).isEqualTo(2); + Assertions.assertThat(secondListResult.pageToken.toString()).isNotNull().isNotEmpty(); + + // List using the final token: + PolarisPage finalListResult = catalog.listNamespaces(secondListResult.pageToken); + Assertions.assertThat(finalListResult.data.size()).isEqualTo(1); + Assertions.assertThat(finalListResult.pageToken).isNull(); + + // List with page size matching the amount of data + PolarisPage firstExactListResult = + catalog.listNamespaces(polarisContext.getMetaStore().pageTokenBuilder().fromLimit(5)); + Assertions.assertThat(firstExactListResult.data.size()).isEqualTo(5); + Assertions.assertThat(firstExactListResult.pageToken.toString()).isNotNull().isNotEmpty(); + + // Again list with matching page size + PolarisPage secondExactListResult = catalog.listNamespaces(firstExactListResult.pageToken); + Assertions.assertThat(secondExactListResult.data).isEmpty(); + Assertions.assertThat(secondExactListResult.pageToken).isNull(); + + // List with huge page size: + PolarisPage bigListResult = + catalog.listNamespaces(polarisContext.getMetaStore().pageTokenBuilder().fromLimit(9999)); + Assertions.assertThat(bigListResult.data.size()).isEqualTo(5); + Assertions.assertThat(bigListResult.pageToken).isNull(); + } finally { + for (int i = 0; i < 5; i++) { + catalog.dropNamespace(Namespace.of("pagination_namespace_" + i)); + } + } + } } diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/TableCleanupTaskHandlerTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/TableCleanupTaskHandlerTest.java index 33e96e2a16..21d27decc9 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/TableCleanupTaskHandlerTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/TableCleanupTaskHandlerTest.java @@ -141,7 +141,14 @@ public void testTableCleanup() throws IOException { assertThat( metaStoreManagerFactory .getOrCreateMetaStoreManager(realmContext) - .loadTasks(callContext.getPolarisCallContext(), "test", 2) + .loadTasks( + callContext.getPolarisCallContext(), + "test", + callContext + .getPolarisCallContext() + .getMetaStore() + .pageTokenBuilder() + .fromLimit(2)) .getEntities()) .hasSize(2) .satisfiesExactlyInAnyOrder( @@ -221,7 +228,14 @@ public void close() { assertThat( metaStoreManagerFactory .getOrCreateMetaStoreManager(realmContext) - .loadTasks(callContext.getPolarisCallContext(), "test", 5) + .loadTasks( + callContext.getPolarisCallContext(), + "test", + callContext + .getPolarisCallContext() + .getMetaStore() + .pageTokenBuilder() + .fromLimit(5)) .getEntities()) .hasSize(2); } @@ -282,10 +296,17 @@ public void close() { assertThat( metaStoreManagerFactory .getOrCreateMetaStoreManager(realmContext) - .loadTasks(callContext.getPolarisCallContext(), "test", 5) + .loadTasks( + callContext.getPolarisCallContext(), + "test", + callContext + .getPolarisCallContext() + .getMetaStore() + .pageTokenBuilder() + .fromLimit(5)) .getEntities()) .hasSize(4) - .satisfiesExactly( + .satisfiesExactlyInAnyOrder( taskEntity -> assertThat(taskEntity) .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) @@ -402,7 +423,10 @@ public void testTableCleanupMultipleSnapshots() throws IOException { List entities = metaStoreManagerFactory .getOrCreateMetaStoreManager(realmContext) - .loadTasks(callContext.getPolarisCallContext(), "test", 5) + .loadTasks( + callContext.getPolarisCallContext(), + "test", + callContext.getPolarisCallContext().getMetaStore().pageTokenBuilder().fromLimit(5)) .getEntities(); List manifestCleanupTasks = @@ -561,7 +585,10 @@ public void testTableCleanupMultipleMetadata() throws IOException { List entities = metaStoreManagerFactory .getOrCreateMetaStoreManager(callContext.getRealmContext()) - .loadTasks(callContext.getPolarisCallContext(), "test", 6) + .loadTasks( + callContext.getPolarisCallContext(), + "test", + callContext.getPolarisCallContext().getMetaStore().pageTokenBuilder().fromLimit(6)) .getEntities(); List manifestCleanupTasks = diff --git a/service/common/build.gradle.kts b/service/common/build.gradle.kts index d38838a85c..310544b181 100644 --- a/service/common/build.gradle.kts +++ b/service/common/build.gradle.kts @@ -29,6 +29,7 @@ dependencies { implementation(project(":polaris-api-management-service")) implementation(project(":polaris-api-iceberg-service")) implementation(project(":polaris-api-catalog-service")) + implementation(project(":polaris-jpa-model")) implementation(platform(libs.iceberg.bom)) implementation("org.apache.iceberg:iceberg-api") @@ -104,6 +105,7 @@ dependencies { testFixturesImplementation(project(":polaris-api-management-model")) testFixturesImplementation(project(":polaris-api-management-service")) testFixturesImplementation(project(":polaris-api-iceberg-service")) + testFixturesImplementation(project(":polaris-jpa-model")) testFixturesImplementation(project(":polaris-api-catalog-service")) testFixturesImplementation(libs.jakarta.enterprise.cdi.api) diff --git a/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java b/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java index 9b40f4228f..81cc817016 100644 --- a/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java +++ b/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java @@ -95,6 +95,7 @@ import org.apache.polaris.core.persistence.dao.entity.DropEntityResult; import org.apache.polaris.core.persistence.dao.entity.EntityResult; import org.apache.polaris.core.persistence.dao.entity.LoadGrantsResult; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest; import org.apache.polaris.core.persistence.resolver.ResolverPath; import org.apache.polaris.core.persistence.resolver.ResolverStatus; @@ -875,7 +876,8 @@ private List listCatalogsUnsafe() { getCurrentPolarisContext(), null, PolarisEntityType.CATALOG, - PolarisEntitySubType.ANY_SUBTYPE) + PolarisEntitySubType.ANY_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities() .stream() .map( @@ -1041,7 +1043,8 @@ public List listPrincipals() { getCurrentPolarisContext(), null, PolarisEntityType.PRINCIPAL, - PolarisEntitySubType.NULL_SUBTYPE) + PolarisEntitySubType.NULL_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities() .stream() .map( @@ -1150,7 +1153,8 @@ public List listPrincipalRoles() { getCurrentPolarisContext(), null, PolarisEntityType.PRINCIPAL_ROLE, - PolarisEntitySubType.NULL_SUBTYPE) + PolarisEntitySubType.NULL_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities() .stream() .map( @@ -1278,7 +1282,8 @@ public List listCatalogRoles(String catalogName) { getCurrentPolarisContext(), PolarisEntity.toCoreList(List.of(catalogEntity)), PolarisEntityType.CATALOG_ROLE, - PolarisEntitySubType.NULL_SUBTYPE) + PolarisEntitySubType.NULL_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities() .stream() .map( diff --git a/service/common/src/main/java/org/apache/polaris/service/catalog/generic/GenericTableCatalog.java b/service/common/src/main/java/org/apache/polaris/service/catalog/generic/GenericTableCatalog.java index b2fb31f67d..708bb85582 100644 --- a/service/common/src/main/java/org/apache/polaris/service/catalog/generic/GenericTableCatalog.java +++ b/service/common/src/main/java/org/apache/polaris/service/catalog/generic/GenericTableCatalog.java @@ -37,6 +37,7 @@ import org.apache.polaris.core.persistence.dao.entity.BaseResult; import org.apache.polaris.core.persistence.dao.entity.DropEntityResult; import org.apache.polaris.core.persistence.dao.entity.EntityResult; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifestCatalogView; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -175,7 +176,8 @@ public List listGenericTables(Namespace namespace) { this.callContext.getPolarisCallContext(), PolarisEntity.toCoreList(catalogPath), PolarisEntityType.TABLE_LIKE, - PolarisEntitySubType.GENERIC_TABLE) + PolarisEntitySubType.GENERIC_TABLE, + ReadEverythingPageToken.get()) .getEntities()); return PolarisCatalogHelpers.nameAndIdToTableIdentifiers(catalogPath, entities); } diff --git a/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalog.java b/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalog.java index 27ef12c3d7..5bdf9dcd6f 100644 --- a/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalog.java +++ b/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalog.java @@ -109,6 +109,9 @@ import org.apache.polaris.core.persistence.dao.entity.DropEntityResult; import org.apache.polaris.core.persistence.dao.entity.EntityResult; import org.apache.polaris.core.persistence.dao.entity.ListEntitiesResult; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifestCatalogView; import org.apache.polaris.core.persistence.resolver.ResolverPath; @@ -180,6 +183,7 @@ public class IcebergCatalog extends BaseMetastoreViewCatalog private Map tableDefaultProperties; private FileIOFactory fileIOFactory; private PolarisMetaStoreManager metaStoreManager; + private final PageToken.PageTokenBuilder pageTokenBuilder; /** * @param entityManager provides handle to underlying PolarisMetaStoreManager with which to @@ -208,6 +212,7 @@ public IcebergCatalog( this.catalogName = catalogEntity.getName(); this.fileIOFactory = fileIOFactory; this.metaStoreManager = metaStoreManager; + this.pageTokenBuilder = callContext.getPolarisCallContext().getMetaStore().pageTokenBuilder(); } @Override @@ -484,14 +489,32 @@ public boolean dropTable(TableIdentifier tableIdentifier, boolean purge) { return true; } + /** Check whether pagination is enabled for list operations */ + private boolean paginationEnabled() { + return callContext + .getPolarisCallContext() + .getConfigurationStore() + .getConfiguration( + callContext.getPolarisCallContext(), + catalogEntity, + FeatureConfiguration.LIST_PAGINATION_ENABLED); + } + @Override public List listTables(Namespace namespace) { + return listTables(namespace, ReadEverythingPageToken.get()).data; + } + + public PolarisPage listTables(Namespace namespace, PageToken pageToken) { if (!namespaceExists(namespace)) { throw new NoSuchNamespaceException( "Cannot list tables for namespace. Namespace does not exist: '%s'", namespace); } + if (!paginationEnabled()) { + pageToken = ReadEverythingPageToken.get(); + } - return listTableLike(PolarisEntitySubType.ICEBERG_TABLE, namespace); + return listTableLike(PolarisEntitySubType.ICEBERG_TABLE, namespace, pageToken); } @Override @@ -801,22 +824,40 @@ public List listNamespaces() { @Override public List listNamespaces(Namespace namespace) throws NoSuchNamespaceException { + return listNamespaces(namespace, ReadEverythingPageToken.get()).data; + } + + public PolarisPage listNamespaces(PageToken pageToken) + throws NoSuchNamespaceException { + return listNamespaces(Namespace.empty(), pageToken); + } + + public PolarisPage listNamespaces(Namespace namespace, PageToken pageToken) + throws NoSuchNamespaceException { PolarisResolvedPathWrapper resolvedEntities = resolvedEntityView.getResolvedPath(namespace); if (resolvedEntities == null) { throw new NoSuchNamespaceException("Namespace does not exist: %s", namespace); } + if (!paginationEnabled()) { + pageToken = ReadEverythingPageToken.get(); + } List catalogPath = resolvedEntities.getRawFullPath(); + ListEntitiesResult listResult = + getMetaStoreManager() + .listEntities( + getCurrentPolarisContext(), + PolarisEntity.toCoreList(catalogPath), + PolarisEntityType.NAMESPACE, + PolarisEntitySubType.NULL_SUBTYPE, + pageToken); List entities = - PolarisEntity.toNameAndIdList( - getMetaStoreManager() - .listEntities( - getCurrentPolarisContext(), - PolarisEntity.toCoreList(catalogPath), - PolarisEntityType.NAMESPACE, - PolarisEntitySubType.NULL_SUBTYPE) - .getEntities()); - return PolarisCatalogHelpers.nameAndIdToNamespaces(catalogPath, entities); + PolarisEntity.toNameAndIdList(listResult.getEntities()); + List namespaces = PolarisCatalogHelpers.nameAndIdToNamespaces(catalogPath, entities); + return listResult + .getPageToken() + .map(token -> new PolarisPage<>(token, namespaces)) + .orElseGet(() -> PolarisPage.fromData(namespaces)); } @Override @@ -828,12 +869,19 @@ public void close() throws IOException { @Override public List listViews(Namespace namespace) { + return listViews(namespace, ReadEverythingPageToken.get()).data; + } + + public PolarisPage listViews(Namespace namespace, PageToken pageToken) { if (!namespaceExists(namespace)) { throw new NoSuchNamespaceException( "Cannot list views for namespace. Namespace does not exist: '%s'", namespace); } + if (!paginationEnabled()) { + pageToken = ReadEverythingPageToken.get(); + } - return listTableLike(PolarisEntitySubType.ICEBERG_VIEW, namespace); + return listTableLike(PolarisEntitySubType.ICEBERG_VIEW, namespace, pageToken); } @Override @@ -1060,7 +1108,8 @@ private void validateNoLocationOverlap( callContext.getPolarisCallContext(), parentPath.stream().map(PolarisEntity::toCore).collect(Collectors.toList()), PolarisEntityType.NAMESPACE, - PolarisEntitySubType.ANY_SUBTYPE); + PolarisEntitySubType.ANY_SUBTYPE, + ReadEverythingPageToken.get()); if (!siblingNamespacesResult.isSuccess()) { throw new IllegalStateException( "Unable to resolve siblings entities to validate location - could not list namespaces"); @@ -1085,7 +1134,8 @@ private void validateNoLocationOverlap( .map(PolarisEntity::toCore) .collect(Collectors.toList()), PolarisEntityType.TABLE_LIKE, - PolarisEntitySubType.ANY_SUBTYPE); + PolarisEntitySubType.ANY_SUBTYPE, + ReadEverythingPageToken.get()); if (!siblingTablesResult.isSuccess()) { throw new IllegalStateException( "Unable to resolve siblings entities to validate location - could not list tables"); @@ -2428,7 +2478,8 @@ private void createNonExistingNamespaces(Namespace namespace) { } } - private List listTableLike(PolarisEntitySubType subType, Namespace namespace) { + private PolarisPage listTableLike( + PolarisEntitySubType subType, Namespace namespace, PageToken pageToken) { PolarisResolvedPathWrapper resolvedEntities = resolvedEntityView.getResolvedPath(namespace); if (resolvedEntities == null) { // Illegal state because the namespace should've already been in the static resolution set. @@ -2437,16 +2488,23 @@ private List listTableLike(PolarisEntitySubType subType, Namesp } List catalogPath = resolvedEntities.getRawFullPath(); + ListEntitiesResult listResult = + getMetaStoreManager() + .listEntities( + getCurrentPolarisContext(), + PolarisEntity.toCoreList(catalogPath), + PolarisEntityType.TABLE_LIKE, + subType, + pageToken); List entities = - PolarisEntity.toNameAndIdList( - getMetaStoreManager() - .listEntities( - getCurrentPolarisContext(), - PolarisEntity.toCoreList(catalogPath), - PolarisEntityType.TABLE_LIKE, - subType) - .getEntities()); - return PolarisCatalogHelpers.nameAndIdToTableIdentifiers(catalogPath, entities); + PolarisEntity.toNameAndIdList(listResult.getEntities()); + List identifiers = + PolarisCatalogHelpers.nameAndIdToTableIdentifiers(catalogPath, entities); + + return listResult + .getPageToken() + .map(token -> new PolarisPage<>(token, identifiers)) + .orElseGet(() -> PolarisPage.fromData(identifiers)); } /** diff --git a/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogAdapter.java b/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogAdapter.java index d1c930bf40..6cf361f95a 100644 --- a/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogAdapter.java +++ b/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogAdapter.java @@ -24,6 +24,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import jakarta.annotation.Nullable; import jakarta.enterprise.context.RequestScoped; import jakarta.inject.Inject; import jakarta.ws.rs.WebApplicationException; @@ -63,6 +64,7 @@ import org.apache.polaris.core.persistence.PolarisEntityManager; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; import org.apache.polaris.core.persistence.ResolvedPolarisEntity; +import org.apache.polaris.core.persistence.pagination.PageToken; import org.apache.polaris.core.persistence.resolver.Resolver; import org.apache.polaris.core.persistence.resolver.ResolverStatus; import org.apache.polaris.core.rest.PolarisEndpoints; @@ -195,6 +197,24 @@ private IcebergCatalogHandler newHandlerWrapper( polarisAuthorizer); } + /** Build a {@link PageToken} from a string and page size. */ + private PageToken buildPageToken(@Nullable String tokenString, @Nullable Integer pageSize) { + if (tokenString != null) { + return callContext + .getPolarisCallContext() + .getMetaStore() + .pageTokenBuilder() + .fromString(tokenString) + .withPageSize(pageSize); + } else { + return callContext + .getPolarisCallContext() + .getMetaStore() + .pageTokenBuilder() + .fromLimit(pageSize); + } + } + @Override public Response createNamespace( String prefix, @@ -216,11 +236,13 @@ public Response listNamespaces( RealmContext realmContext, SecurityContext securityContext) { Optional namespaceOptional = Optional.ofNullable(parent).map(this::decodeNamespace); + PageToken token = buildPageToken(pageToken, pageSize); return withCatalog( securityContext, prefix, catalog -> - Response.ok(catalog.listNamespaces(namespaceOptional.orElse(Namespace.of()))).build()); + Response.ok(catalog.listNamespaces(namespaceOptional.orElse(Namespace.of()), token)) + .build()); } @Override @@ -355,8 +377,9 @@ public Response listTables( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); + PageToken token = buildPageToken(pageToken, pageSize); return withCatalog( - securityContext, prefix, catalog -> Response.ok(catalog.listTables(ns)).build()); + securityContext, prefix, catalog -> Response.ok(catalog.listTables(ns, token)).build()); } @Override @@ -525,8 +548,9 @@ public Response listViews( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); + PageToken token = buildPageToken(pageToken, pageSize); return withCatalog( - securityContext, prefix, catalog -> Response.ok(catalog.listViews(ns)).build()); + securityContext, prefix, catalog -> Response.ok(catalog.listViews(ns, token)).build()); } @Override diff --git a/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogHandler.java b/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogHandler.java index 2785e004cf..cf9a21a724 100644 --- a/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogHandler.java +++ b/service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogHandler.java @@ -90,6 +90,8 @@ import org.apache.polaris.core.persistence.TransactionWorkspaceMetaStoreManager; import org.apache.polaris.core.persistence.dao.entity.EntitiesResult; import org.apache.polaris.core.persistence.dao.entity.EntityWithPath; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; import org.apache.polaris.core.secrets.UserSecretsManager; import org.apache.polaris.core.storage.PolarisStorageActions; import org.apache.polaris.service.catalog.SupportsNotifications; @@ -97,6 +99,8 @@ import org.apache.polaris.service.context.CallContextCatalogFactory; import org.apache.polaris.service.http.IcebergHttpUtil; import org.apache.polaris.service.http.IfNoneMatch; +import org.apache.polaris.service.types.ListNamespacesResponseWithPageToken; +import org.apache.polaris.service.types.ListTablesResponseWithPageToken; import org.apache.polaris.service.types.NotificationRequest; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -168,6 +172,20 @@ public static boolean isCreate(UpdateTableRequest request) { return isCreate; } + public ListNamespacesResponseWithPageToken listNamespaces(Namespace parent, PageToken pageToken) { + PolarisAuthorizableOperation op = PolarisAuthorizableOperation.LIST_NAMESPACES; + authorizeBasicNamespaceOperationOrThrow(op, parent); + + if (baseCatalog instanceof IcebergCatalog polarisCatalog) { + return ListNamespacesResponseWithPageToken.fromPolarisPage( + polarisCatalog.listNamespaces(parent, pageToken)); + } else { + return ListNamespacesResponseWithPageToken.fromPolarisPage( + PolarisPage.fromData( + CatalogHandlers.listNamespaces(namespaceCatalog, parent).namespaces())); + } + } + private UserSecretsManager getUserSecretsManager() { return userSecretsManager; } @@ -303,6 +321,19 @@ public UpdateNamespacePropertiesResponse updateNamespaceProperties( return CatalogHandlers.updateNamespaceProperties(namespaceCatalog, namespace, request); } + public ListTablesResponseWithPageToken listTables(Namespace namespace, PageToken pageToken) { + PolarisAuthorizableOperation op = PolarisAuthorizableOperation.LIST_TABLES; + authorizeBasicNamespaceOperationOrThrow(op, namespace); + + if (baseCatalog instanceof IcebergCatalog polarisCatalog) { + return ListTablesResponseWithPageToken.fromPolarisPage( + polarisCatalog.listTables(namespace, pageToken)); + } else { + return ListTablesResponseWithPageToken.fromPolarisPage( + PolarisPage.fromData(CatalogHandlers.listTables(baseCatalog, namespace).identifiers())); + } + } + public ListTablesResponse listTables(Namespace namespace) { PolarisAuthorizableOperation op = PolarisAuthorizableOperation.LIST_TABLES; authorizeBasicNamespaceOperationOrThrow(op, namespace); @@ -939,6 +970,19 @@ public void commitTransaction(CommitTransactionRequest commitTransactionRequest) } } + public ListTablesResponseWithPageToken listViews(Namespace namespace, PageToken pageToken) { + PolarisAuthorizableOperation op = PolarisAuthorizableOperation.LIST_VIEWS; + authorizeBasicNamespaceOperationOrThrow(op, namespace); + + if (baseCatalog instanceof IcebergCatalog polarisCatalog) { + return ListTablesResponseWithPageToken.fromPolarisPage( + polarisCatalog.listViews(namespace, pageToken)); + } else { + return ListTablesResponseWithPageToken.fromPolarisPage( + PolarisPage.fromData(CatalogHandlers.listTables(baseCatalog, namespace).identifiers())); + } + } + public ListTablesResponse listViews(Namespace namespace) { PolarisAuthorizableOperation op = PolarisAuthorizableOperation.LIST_VIEWS; authorizeBasicNamespaceOperationOrThrow(op, namespace); diff --git a/service/common/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java b/service/common/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java index 45e5f17e6d..3e2581298a 100644 --- a/service/common/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java +++ b/service/common/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java @@ -49,6 +49,7 @@ import org.apache.polaris.core.persistence.PolicyMappingAlreadyExistsException; import org.apache.polaris.core.persistence.dao.entity.EntityResult; import org.apache.polaris.core.persistence.dao.entity.LoadPolicyMappingsResult; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifestCatalogView; import org.apache.polaris.core.policy.PolicyEntity; import org.apache.polaris.core.policy.PolicyType; @@ -165,7 +166,8 @@ public List listPolicies(Namespace namespace, PolicyType polic callContext.getPolarisCallContext(), PolarisEntity.toCoreList(catalogPath), PolarisEntityType.POLICY, - PolarisEntitySubType.NULL_SUBTYPE) + PolarisEntitySubType.NULL_SUBTYPE, + ReadEverythingPageToken.get()) .getEntities() .stream() .map( diff --git a/service/common/src/main/java/org/apache/polaris/service/types/ListNamespacesResponseWithPageToken.java b/service/common/src/main/java/org/apache/polaris/service/types/ListNamespacesResponseWithPageToken.java new file mode 100644 index 0000000000..cbfaed2998 --- /dev/null +++ b/service/common/src/main/java/org/apache/polaris/service/types/ListNamespacesResponseWithPageToken.java @@ -0,0 +1,72 @@ +/* + * 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.service.types; + +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.MoreObjects; +import com.google.common.base.Preconditions; +import java.util.List; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.rest.responses.ListNamespacesResponse; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; + +/** + * Used in lieu of {@link ListNamespacesResponse} when there may be a {@link PageToken} associated + * with the response. Callers can use this {@link PageToken} to continue the listing operation and + * obtain more results. + */ +public class ListNamespacesResponseWithPageToken extends ListNamespacesResponse { + private final PageToken pageToken; + + private final List namespaces; + + public ListNamespacesResponseWithPageToken(PageToken pageToken, List namespaces) { + this.pageToken = pageToken; + this.namespaces = namespaces; + Preconditions.checkArgument(this.namespaces != null, "Invalid namespace: null"); + } + + public static ListNamespacesResponseWithPageToken fromPolarisPage( + PolarisPage polarisPage) { + return new ListNamespacesResponseWithPageToken(polarisPage.pageToken, polarisPage.data); + } + + @JsonProperty("next-page-token") + public String getPageToken() { + if (pageToken == null) { + return null; + } else { + return pageToken.toString(); + } + } + + @Override + public List namespaces() { + return this.namespaces != null ? this.namespaces : List.of(); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("namespaces", this.namespaces) + .add("pageToken", this.pageToken.toString()) + .toString(); + } +} diff --git a/service/common/src/main/java/org/apache/polaris/service/types/ListTablesResponseWithPageToken.java b/service/common/src/main/java/org/apache/polaris/service/types/ListTablesResponseWithPageToken.java new file mode 100644 index 0000000000..0c4dbaf009 --- /dev/null +++ b/service/common/src/main/java/org/apache/polaris/service/types/ListTablesResponseWithPageToken.java @@ -0,0 +1,72 @@ +/* + * 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.service.types; + +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.MoreObjects; +import com.google.common.base.Preconditions; +import java.util.List; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.rest.responses.ListTablesResponse; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.PolarisPage; + +/** + * Used in lieu of {@link ListTablesResponse} when there may be a {@link PageToken} associated with + * the response. Callers can use this {@link PageToken} to continue the listing operation and obtain + * more results. + */ +public class ListTablesResponseWithPageToken extends ListTablesResponse { + private final PageToken pageToken; + + private final List identifiers; + + public ListTablesResponseWithPageToken(PageToken pageToken, List identifiers) { + this.pageToken = pageToken; + this.identifiers = identifiers; + Preconditions.checkArgument(this.identifiers != null, "Invalid identifier list: null"); + } + + public static ListTablesResponseWithPageToken fromPolarisPage( + PolarisPage polarisPage) { + return new ListTablesResponseWithPageToken(polarisPage.pageToken, polarisPage.data); + } + + @JsonProperty("next-page-token") + public String getPageToken() { + if (pageToken == null) { + return null; + } else { + return pageToken.toString(); + } + } + + @Override + public List identifiers() { + return this.identifiers != null ? this.identifiers : List.of(); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("identifiers", this.identifiers) + .add("pageToken", this.pageToken) + .toString(); + } +} diff --git a/service/common/src/test/java/org/apache/polaris/service/catalog/io/FileIOFactoryTest.java b/service/common/src/test/java/org/apache/polaris/service/catalog/io/FileIOFactoryTest.java index 9a070273b8..2afceead8e 100644 --- a/service/common/src/test/java/org/apache/polaris/service/catalog/io/FileIOFactoryTest.java +++ b/service/common/src/test/java/org/apache/polaris/service/catalog/io/FileIOFactoryTest.java @@ -184,7 +184,10 @@ public void testLoadFileIOForCleanupTask() { testServices .metaStoreManagerFactory() .getOrCreateMetaStoreManager(realmContext) - .loadTasks(callContext.getPolarisCallContext(), "testExecutor", 1) + .loadTasks( + callContext.getPolarisCallContext(), + "testExecutor", + callContext.getPolarisCallContext().getMetaStore().pageTokenBuilder().fromLimit(1)) .getEntities(); Assertions.assertThat(tasks).hasSize(1); TaskEntity taskEntity = TaskEntity.of(tasks.get(0)); diff --git a/service/common/src/test/java/org/apache/polaris/service/persistence/pagination/PageTokenTest.java b/service/common/src/test/java/org/apache/polaris/service/persistence/pagination/PageTokenTest.java new file mode 100644 index 0000000000..7581430bc0 --- /dev/null +++ b/service/common/src/test/java/org/apache/polaris/service/persistence/pagination/PageTokenTest.java @@ -0,0 +1,207 @@ +/* + * 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.service.persistence.pagination; + +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import org.apache.polaris.core.persistence.pagination.OffsetPageToken; +import org.apache.polaris.core.persistence.pagination.PageToken; +import org.apache.polaris.core.persistence.pagination.ReadEverythingPageToken; +import org.apache.polaris.jpa.models.EntityIdPageToken; +import org.apache.polaris.jpa.models.ModelEntity; +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class PageTokenTest { + private static final Logger LOGGER = LoggerFactory.getLogger(PageTokenTest.class); + + static Stream> getPageTokenBuilders() { + return Stream.of( + OffsetPageToken.builder(), EntityIdPageToken.builder(), ReadEverythingPageToken.builder()); + } + + @Test + void testDoneToken() { + Assertions.assertThat(PageToken.DONE).isNull(); + } + + @ParameterizedTest + @MethodSource("getPageTokenBuilders") + void testRoundTrips(PageToken.PageTokenBuilder builder) { + if (builder instanceof ReadEverythingPageToken.ReadEverythingPageTokenBuilder) { + // Skip ReadEverythingPageToken + return; + } + + for (int limit : List.of(1, 10, 100, Integer.MAX_VALUE)) { + PageToken token = builder.fromLimit(limit); + Assertions.assertThat(token.pageSize).isEqualTo(limit); + Assertions.assertThat(builder.fromString(token.toString())).isEqualTo(token); + } + } + + @ParameterizedTest + @MethodSource("getPageTokenBuilders") + void testInvalidLimits(PageToken.PageTokenBuilder builder) { + if (builder instanceof ReadEverythingPageToken.ReadEverythingPageTokenBuilder) { + // Skip ReadEverythingPageToken + return; + } + + for (int limit : List.of(-1, 0)) { + Assertions.assertThatThrownBy(() -> builder.fromLimit(limit)) + .isInstanceOf(IllegalArgumentException.class); + } + + Assertions.assertThat(builder.fromLimit(null)).isInstanceOf(ReadEverythingPageToken.class); + } + + @ParameterizedTest + @MethodSource("getPageTokenBuilders") + void testStartingTokens(PageToken.PageTokenBuilder builder) { + Assertions.assertThat(builder.fromString("")).isNotNull(); + if (!(builder instanceof ReadEverythingPageToken.ReadEverythingPageTokenBuilder)) { + Assertions.assertThat(builder.fromString("")).isNotEqualTo(ReadEverythingPageToken.get()); + } + + Assertions.assertThatThrownBy(() -> builder.fromString(null)) + .isInstanceOf(IllegalArgumentException.class); + } + + @ParameterizedTest + @MethodSource("getPageTokenBuilders") + void testPageBuilding(PageToken.PageTokenBuilder builder) { + if (builder instanceof ReadEverythingPageToken.ReadEverythingPageTokenBuilder) { + // Skip ReadEverythingPageToken + return; + } + + List data = + List.of(ModelEntity.builder().id(1).build(), ModelEntity.builder().id(2).build()); + + PageToken token = builder.fromLimit(1000); + Assertions.assertThat(token.buildNextPage(data).data).isEqualTo(data); + Assertions.assertThat(token.buildNextPage(data).pageToken).isNull(); + } + + @Test + void testUniquePrefixes() { + Stream> builders = getPageTokenBuilders(); + List prefixes = + builders.map(PageToken.PageTokenBuilder::tokenPrefix).collect(Collectors.toList()); + Assertions.assertThat(prefixes.size()).isEqualTo(prefixes.stream().distinct().count()); + } + + @ParameterizedTest + @MethodSource("getPageTokenBuilders") + void testCrossTokenParsing(PageToken.PageTokenBuilder builder) { + var otherBuilders = getPageTokenBuilders().collect(Collectors.toList()); + for (var otherBuilder : otherBuilders) { + LOGGER.info( + "Testing {} being parsed by {}", + builder.getClass().getSimpleName(), + otherBuilder.getClass().getSimpleName()); + + final PageToken token; + if (builder instanceof ReadEverythingPageToken.ReadEverythingPageTokenBuilder) { + token = ReadEverythingPageToken.get(); + } else { + token = builder.fromLimit(1234); + } + if (otherBuilder.getClass().equals(builder.getClass())) { + Assertions.assertThat(otherBuilder.fromString(token.toString())).isEqualTo(token); + } else { + Assertions.assertThatThrownBy(() -> otherBuilder.fromString(token.toString())) + .isInstanceOf(IllegalArgumentException.class); + } + } + } + + @ParameterizedTest + @MethodSource("getPageTokenBuilders") + void testDefaultTokens(PageToken.PageTokenBuilder builder) { + if (builder instanceof ReadEverythingPageToken.ReadEverythingPageTokenBuilder) { + // Skip ReadEverythingPageToken + return; + } + + PageToken token = builder.fromString(""); + Assertions.assertThat(token.toString()).isNotNull(); + Assertions.assertThat(token.pageSize).isEqualTo(PageToken.DEFAULT_PAGE_SIZE); + } + + @Test + void testReadEverythingPageToken() { + PageToken token = ReadEverythingPageToken.get(); + + Assertions.assertThat(token.toString()).isNotNull(); + Assertions.assertThat(token.buildNextPage(List.of("anything")).pageToken) + .isEqualTo(PageToken.DONE); + Assertions.assertThat(token.pageSize).isEqualTo(Integer.MAX_VALUE); + } + + @Test + void testOffsetPageToken() { + OffsetPageToken token = (OffsetPageToken) OffsetPageToken.builder().fromLimit(2); + + Assertions.assertThat(token).isInstanceOf(OffsetPageToken.class); + Assertions.assertThat(token.offset).isEqualTo(0); + + List data = List.of("some", "data"); + var page = token.buildNextPage(data); + Assertions.assertThat(page.pageToken).isNotNull(); + Assertions.assertThat(page.pageToken).isInstanceOf(OffsetPageToken.class); + Assertions.assertThat(page.pageToken.pageSize).isEqualTo(2); + Assertions.assertThat(((OffsetPageToken) page.pageToken).offset).isEqualTo(2); + Assertions.assertThat(page.data).isEqualTo(data); + + Assertions.assertThat(OffsetPageToken.builder().fromString(page.pageToken.toString())) + .isEqualTo(page.pageToken); + } + + @Test + void testEntityIdPageToken() { + EntityIdPageToken token = (EntityIdPageToken) EntityIdPageToken.builder().fromLimit(2); + + Assertions.assertThat(token).isInstanceOf(EntityIdPageToken.class); + Assertions.assertThat(token.id).isEqualTo(-1L); + + List badData = List.of("some", "data"); + Assertions.assertThatThrownBy(() -> token.buildNextPage(badData)) + .isInstanceOf(IllegalArgumentException.class); + + List data = + List.of(ModelEntity.builder().id(101).build(), ModelEntity.builder().id(102).build()); + var page = token.buildNextPage(data); + + Assertions.assertThat(page.pageToken).isNotNull(); + Assertions.assertThat(page.pageToken).isInstanceOf(EntityIdPageToken.class); + Assertions.assertThat(page.pageToken.pageSize).isEqualTo(2); + Assertions.assertThat(((EntityIdPageToken) page.pageToken).id).isEqualTo(102); + Assertions.assertThat(page.data).isEqualTo(data); + + Assertions.assertThat(EntityIdPageToken.builder().fromString(page.pageToken.toString())) + .isEqualTo(page.pageToken); + } +}