From f32277b8b7dc363efb96453d542ec1b494edd7f5 Mon Sep 17 00:00:00 2001 From: fhan Date: Tue, 18 Aug 2026 16:23:28 +0800 Subject: [PATCH 1/5] [client] Add Admin API to describe buckets --- .../org/apache/fluss/client/admin/Admin.java | 41 ++++ .../apache/fluss/client/admin/FlussAdmin.java | 29 +++ .../client/utils/ClientRpcMessageUtils.java | 24 ++ .../client/admin/DescribeBucketsITCase.java | 215 ++++++++++++++++++ .../fluss/client/admin/FlussAdminITCase.java | 3 +- .../acl/FlussAuthorizationITCase.java | 12 +- .../utils/ClientRpcMessageUtilsTest.java | 48 ++++ .../org/apache/fluss/metadata/BucketInfo.java | 180 +++++++++++++++ .../apache/fluss/metadata/BucketInfoTest.java | 130 +++++++++++ .../sink/testutils/TestAdminAdapter.java | 12 + .../rpc/gateway/AdminReadOnlyGateway.java | 11 + .../apache/fluss/rpc/protocol/ApiKeys.java | 3 +- fluss-rpc/src/main/proto/FlussApi.proto | 23 ++ .../rpc/TestingTabletGatewayService.java | 8 + .../apache/fluss/server/RpcServiceBase.java | 131 +++++++++++ .../server/coordinator/MetadataManager.java | 3 + .../fluss/server/zk/ZooKeeperClient.java | 53 +++-- .../coordinator/TestCoordinatorGateway.java | 8 + .../tablet/TestTabletServerGateway.java | 8 + .../fluss/server/zk/ZooKeeperClientTest.java | 5 + 20 files changed, 922 insertions(+), 25 deletions(-) create mode 100644 fluss-client/src/test/java/org/apache/fluss/client/admin/DescribeBucketsITCase.java create mode 100644 fluss-common/src/main/java/org/apache/fluss/metadata/BucketInfo.java create mode 100644 fluss-common/src/test/java/org/apache/fluss/metadata/BucketInfoTest.java diff --git a/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java b/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java index a99423d79ba..a43babfacca 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java @@ -55,6 +55,7 @@ import org.apache.fluss.exception.TableNotPartitionedException; import org.apache.fluss.exception.TooManyBucketsException; import org.apache.fluss.exception.TooManyPartitionsException; +import org.apache.fluss.metadata.BucketInfo; import org.apache.fluss.metadata.DatabaseChange; import org.apache.fluss.metadata.DatabaseDescriptor; import org.apache.fluss.metadata.DatabaseInfo; @@ -253,6 +254,46 @@ CompletableFuture createTable( */ CompletableFuture getTableInfo(TablePath tablePath); + /** + * Describes the buckets of the given table asynchronously. + * + *

For a non-partitioned table, this returns the table buckets. For a partitioned table, this + * returns the buckets of all partitions. For a partitioned table with many partitions, prefer + * {@link #describeBuckets(TablePath, PartitionSpec)} to limit the result. + * + *

The following exceptions can be anticipated when calling {@code get()} on the returned + * future. + * + *

    + *
  • {@link TableNotExistException} if the table does not exist. + *
+ * + * @param tablePath The table path of the table. + * @since 1.0 + */ + CompletableFuture> describeBuckets(TablePath tablePath); + + /** + * Describes the buckets matching the given partition spec asynchronously. + * + *

The partition spec may contain all partition keys or a subset of them. + * + *

The following exceptions can be anticipated when calling {@code get()} on the returned + * future. + * + *

    + *
  • {@link TableNotExistException} if the table does not exist. + *
  • {@link TableNotPartitionedException} if the table is not partitioned. + *
  • {@link InvalidPartitionException} if the partition spec is invalid. + *
+ * + * @param tablePath The table path of the table. + * @param partitionSpec The complete or partial partition spec. + * @since 1.0 + */ + CompletableFuture> describeBuckets( + TablePath tablePath, PartitionSpec partitionSpec); + /** * Drop the table with the given table path asynchronously. * diff --git a/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java b/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java index 143eaeecc42..a340802b49e 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java @@ -36,6 +36,7 @@ import org.apache.fluss.exception.FlussRuntimeException; import org.apache.fluss.exception.LeaderNotAvailableException; import org.apache.fluss.exception.PartitionNotExistException; +import org.apache.fluss.metadata.BucketInfo; import org.apache.fluss.metadata.DatabaseChange; import org.apache.fluss.metadata.DatabaseDescriptor; import org.apache.fluss.metadata.DatabaseInfo; @@ -70,6 +71,7 @@ import org.apache.fluss.rpc.messages.DatabaseExistsRequest; import org.apache.fluss.rpc.messages.DatabaseExistsResponse; import org.apache.fluss.rpc.messages.DeleteProducerOffsetsRequest; +import org.apache.fluss.rpc.messages.DescribeBucketsRequest; import org.apache.fluss.rpc.messages.DescribeClusterConfigsRequest; import org.apache.fluss.rpc.messages.DropAclsRequest; import org.apache.fluss.rpc.messages.DropDatabaseRequest; @@ -358,6 +360,33 @@ public CompletableFuture getTableInfo(TablePath tablePath) { r.hasBucketCountEpoch() ? r.getBucketCountEpoch() : 0L)); } + @Override + public CompletableFuture> describeBuckets(TablePath tablePath) { + tablePath.validate(); + DescribeBucketsRequest request = new DescribeBucketsRequest(); + request.setTablePath() + .setDatabaseName(tablePath.getDatabaseName()) + .setTableName(tablePath.getTableName()); + return readOnlyGateway + .describeBuckets(request) + .thenApply(ClientRpcMessageUtils::toBucketInfos); + } + + @Override + public CompletableFuture> describeBuckets( + TablePath tablePath, PartitionSpec partitionSpec) { + tablePath.validate(); + checkNotNull(partitionSpec, "partitionSpec must not be null"); + DescribeBucketsRequest request = new DescribeBucketsRequest(); + request.setTablePath() + .setDatabaseName(tablePath.getDatabaseName()) + .setTableName(tablePath.getTableName()); + request.setPartitionSpec(makePbPartitionSpec(partitionSpec)); + return readOnlyGateway + .describeBuckets(request) + .thenApply(ClientRpcMessageUtils::toBucketInfos); + } + @Override public CompletableFuture dropTable(TablePath tablePath, boolean ignoreIfNotExists) { DropTableRequest request = new DropTableRequest(); diff --git a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java index ee535f2b0d7..9c9a5be3a02 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java @@ -43,6 +43,7 @@ import org.apache.fluss.fs.FsPathAndFileName; import org.apache.fluss.fs.token.ObtainedSecurityToken; import org.apache.fluss.metadata.AggFunction; +import org.apache.fluss.metadata.BucketInfo; import org.apache.fluss.metadata.DatabaseChange; import org.apache.fluss.metadata.DatabaseSummary; import org.apache.fluss.metadata.PartitionInfo; @@ -57,6 +58,7 @@ import org.apache.fluss.rpc.messages.AlterDatabaseRequest; import org.apache.fluss.rpc.messages.AlterTableRequest; import org.apache.fluss.rpc.messages.CreatePartitionRequest; +import org.apache.fluss.rpc.messages.DescribeBucketsResponse; import org.apache.fluss.rpc.messages.DropPartitionRequest; import org.apache.fluss.rpc.messages.GetClusterHealthResponse; import org.apache.fluss.rpc.messages.GetFileSystemSecurityTokenResponse; @@ -75,6 +77,7 @@ import org.apache.fluss.rpc.messages.MetadataRequest; import org.apache.fluss.rpc.messages.PbAddColumn; import org.apache.fluss.rpc.messages.PbAlterConfig; +import org.apache.fluss.rpc.messages.PbBucketInfo; import org.apache.fluss.rpc.messages.PbBucketOffset; import org.apache.fluss.rpc.messages.PbDatabaseSummary; import org.apache.fluss.rpc.messages.PbDescribeConfig; @@ -696,6 +699,27 @@ public static List toPartitionInfos( .collect(Collectors.toList()); } + public static List toBucketInfos(DescribeBucketsResponse response) { + return response.getBucketInfosList().stream() + .map(ClientRpcMessageUtils::toBucketInfo) + .collect(Collectors.toList()); + } + + private static BucketInfo toBucketInfo(PbBucketInfo pbBucketInfo) { + return new BucketInfo( + TablePath.of( + pbBucketInfo.getTablePath().getDatabaseName(), + pbBucketInfo.getTablePath().getTableName()), + pbBucketInfo.getTableId(), + pbBucketInfo.hasPartitionId() ? pbBucketInfo.getPartitionId() : null, + pbBucketInfo.hasPartitionName() ? pbBucketInfo.getPartitionName() : null, + pbBucketInfo.getBucketId(), + pbBucketInfo.hasLeaderId() ? pbBucketInfo.getLeaderId() : null, + pbBucketInfo.hasLeaderEpoch() ? pbBucketInfo.getLeaderEpoch() : null, + Arrays.stream(pbBucketInfo.getReplicaIds()).boxed().collect(Collectors.toList()), + Arrays.stream(pbBucketInfo.getIsrIds()).boxed().collect(Collectors.toList())); + } + public static Map toKeyValueMap(List pbKeyValues) { return pbKeyValues.stream() .collect( diff --git a/fluss-client/src/test/java/org/apache/fluss/client/admin/DescribeBucketsITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/admin/DescribeBucketsITCase.java new file mode 100644 index 00000000000..7d0ade2e1a8 --- /dev/null +++ b/fluss-client/src/test/java/org/apache/fluss/client/admin/DescribeBucketsITCase.java @@ -0,0 +1,215 @@ +/* + * 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.fluss.client.admin; + +import org.apache.fluss.exception.InvalidPartitionException; +import org.apache.fluss.exception.TableNotExistException; +import org.apache.fluss.exception.TableNotPartitionedException; +import org.apache.fluss.metadata.BucketInfo; +import org.apache.fluss.metadata.PartitionInfo; +import org.apache.fluss.metadata.PartitionSpec; +import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.types.DataTypes; + +import org.junit.jupiter.api.Test; + +import javax.annotation.Nullable; + +import java.time.Duration; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.apache.fluss.testutils.common.CommonTestUtils.waitUntil; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Integration test for describing table buckets through {@link Admin}. */ +class DescribeBucketsITCase extends ClientToServerITCaseBase { + + private static final TablePath NON_PARTITIONED_TABLE_PATH = + TablePath.of("test_db", "non_partitioned_table"); + private static final TablePath PARTITIONED_TABLE_PATH = + TablePath.of("test_db", "partitioned_table"); + + @Test + void testDescribeBucketsForNonPartitionedTable() throws Exception { + TableDescriptor tableDescriptor = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("name", DataTypes.STRING()) + .primaryKey("id") + .build()) + .distributedBy(3, "id") + .build(); + long tableId = createTable(NON_PARTITIONED_TABLE_PATH, tableDescriptor, false); + + List bucketInfos = waitAndDescribeBuckets(NON_PARTITIONED_TABLE_PATH, null, 3); + assertThat(bucketInfos).extracting(BucketInfo::getBucketId).containsExactly(0, 1, 2); + bucketInfos.forEach( + bucketInfo -> { + assertBucketInfo(bucketInfo, NON_PARTITIONED_TABLE_PATH, tableId, null); + assertThat(bucketInfo.getPartitionName()).isNull(); + }); + + assertThatThrownBy( + () -> + admin.describeBuckets( + NON_PARTITIONED_TABLE_PATH, + newPartitionSpec("pt", "2025")) + .get()) + .cause() + .isInstanceOf(TableNotPartitionedException.class); + assertThatThrownBy( + () -> admin.describeBuckets(TablePath.of("test_db", "missing_table")).get()) + .cause() + .isInstanceOf(TableNotExistException.class); + } + + @Test + void testDescribeBucketsForPartitionedTable() throws Exception { + TableDescriptor tableDescriptor = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("id", DataTypes.STRING()) + .column("pt", DataTypes.STRING()) + .column("region", DataTypes.STRING()) + .build()) + .distributedBy(2, "id") + .partitionedBy("pt", "region") + .build(); + long tableId = createTable(PARTITIONED_TABLE_PATH, tableDescriptor, false); + PartitionSpec p2025Cn = + newPartitionSpec(Arrays.asList("pt", "region"), Arrays.asList("2025", "cn")); + PartitionSpec p2025Us = + newPartitionSpec(Arrays.asList("pt", "region"), Arrays.asList("2025", "us")); + PartitionSpec p2026Cn = + newPartitionSpec(Arrays.asList("pt", "region"), Arrays.asList("2026", "cn")); + admin.createPartition(PARTITIONED_TABLE_PATH, p2025Cn, false).get(); + admin.createPartition(PARTITIONED_TABLE_PATH, p2025Us, false).get(); + admin.createPartition(PARTITIONED_TABLE_PATH, p2026Cn, false).get(); + + Map partitionIds = + admin.listPartitionInfos(PARTITIONED_TABLE_PATH).get().stream() + .collect( + Collectors.toMap( + PartitionInfo::getPartitionName, + PartitionInfo::getPartitionId)); + + List allPartitionBuckets = + waitAndDescribeBuckets(PARTITIONED_TABLE_PATH, null, 6); + assertThat(allPartitionBuckets) + .extracting( + bucketInfo -> + bucketInfo.getPartitionName() + ":" + bucketInfo.getBucketId()) + .containsExactly( + "2025$cn:0", + "2025$cn:1", + "2025$us:0", + "2025$us:1", + "2026$cn:0", + "2026$cn:1"); + allPartitionBuckets.forEach( + bucketInfo -> + assertBucketInfo( + bucketInfo, + PARTITIONED_TABLE_PATH, + tableId, + partitionIds.get(bucketInfo.getPartitionName()))); + + List partialPartitionBuckets = + waitAndDescribeBuckets(PARTITIONED_TABLE_PATH, newPartitionSpec("pt", "2025"), 4); + assertThat(partialPartitionBuckets) + .extracting(BucketInfo::getPartitionName) + .containsExactly("2025$cn", "2025$cn", "2025$us", "2025$us"); + + List exactPartitionBuckets = + waitAndDescribeBuckets(PARTITIONED_TABLE_PATH, p2025Cn, 2); + assertThat(exactPartitionBuckets) + .extracting(BucketInfo::getPartitionName) + .containsOnly("2025$cn"); + assertThat(exactPartitionBuckets).extracting(BucketInfo::getBucketId).containsExactly(0, 1); + + assertThat( + admin.describeBuckets( + PARTITIONED_TABLE_PATH, newPartitionSpec("pt", "missing")) + .get()) + .isEmpty(); + assertThatThrownBy( + () -> + admin.describeBuckets( + PARTITIONED_TABLE_PATH, + newPartitionSpec("unknown", "2025")) + .get()) + .cause() + .isInstanceOf(InvalidPartitionException.class) + .hasMessageContaining("unknown"); + } + + private List waitAndDescribeBuckets( + TablePath tablePath, @Nullable PartitionSpec partitionSpec, int expectedBucketCount) + throws Exception { + waitUntil( + () -> { + List bucketInfos = describeBuckets(tablePath, partitionSpec); + return bucketInfos.size() == expectedBucketCount + && bucketInfos.stream() + .allMatch( + bucketInfo -> + bucketInfo.getLeaderId().isPresent() + && bucketInfo + .getLeaderEpoch() + .isPresent() + && !bucketInfo.getIsr().isEmpty()); + }, + Duration.ofMinutes(1), + "Waiting for bucket metadata"); + return describeBuckets(tablePath, partitionSpec); + } + + private List describeBuckets( + TablePath tablePath, @Nullable PartitionSpec partitionSpec) throws Exception { + return partitionSpec == null + ? admin.describeBuckets(tablePath).get() + : admin.describeBuckets(tablePath, partitionSpec).get(); + } + + private static void assertBucketInfo( + BucketInfo bucketInfo, + TablePath tablePath, + long tableId, + @Nullable Long expectedPartitionId) { + assertThat(bucketInfo.getTablePath()).isEqualTo(tablePath); + assertThat(bucketInfo.getTableId()).isEqualTo(tableId); + if (expectedPartitionId == null) { + assertThat(bucketInfo.getPartitionId()).isEmpty(); + } else { + assertThat(bucketInfo.getPartitionId()).hasValue(expectedPartitionId); + } + assertThat(bucketInfo.getReplicas()).hasSize(3); + assertThat(bucketInfo.getIsr()).isNotEmpty(); + assertThat(bucketInfo.getReplicas()).containsAll(bucketInfo.getIsr()); + assertThat(bucketInfo.getIsr()).contains(bucketInfo.getLeaderId().getAsInt()); + } +} diff --git a/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java index 7a0a9c1e247..e67faf84b24 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java @@ -38,7 +38,6 @@ import org.apache.fluss.exception.DatabaseAlreadyExistException; import org.apache.fluss.exception.DatabaseNotEmptyException; import org.apache.fluss.exception.DatabaseNotExistException; -import org.apache.fluss.exception.FlussRuntimeException; import org.apache.fluss.exception.InvalidAlterTableException; import org.apache.fluss.exception.InvalidConfigException; import org.apache.fluss.exception.InvalidDatabaseException; @@ -1295,7 +1294,7 @@ void testListPartitionInfosByPartitionSpec() throws Exception { admin.listPartitionInfos(partitionedTablePath, invalidPartitionSpec) .get()) .cause() - .isInstanceOf(FlussRuntimeException.class) + .isInstanceOf(InvalidPartitionException.class) .hasMessageContaining("table don't contains this partitionKey: pt1"); } diff --git a/fluss-client/src/test/java/org/apache/fluss/client/security/acl/FlussAuthorizationITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/security/acl/FlussAuthorizationITCase.java index be862d0e564..2cdaef69d15 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/security/acl/FlussAuthorizationITCase.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/security/acl/FlussAuthorizationITCase.java @@ -422,8 +422,9 @@ void testDescribeTableOperation() throws Exception { // 3. getTableSchema // 4. getLatestKvSnapshots // 5. listPartitionInfos - // 6. getLatestLakeSnapshot - // 7. listOffsets + // 6. describeBuckets + // 7. getLatestLakeSnapshot + // 8. listOffsets // first check call these methods without authorization. assertThat(guestAdmin.listTables(DATA1_TABLE_PATH_PK.getDatabaseName()).get()) @@ -432,6 +433,7 @@ void testDescribeTableOperation() throws Exception { assertNoTableDescribeAuth(() -> guestAdmin.getTableSchema(DATA1_TABLE_PATH_PK).get()); assertNoTableDescribeAuth(() -> guestAdmin.getLatestKvSnapshots(DATA1_TABLE_PATH_PK).get()); assertNoTableDescribeAuth(() -> guestAdmin.listPartitionInfos(DATA1_TABLE_PATH_PK).get()); + assertNoTableDescribeAuth(() -> guestAdmin.describeBuckets(DATA1_TABLE_PATH_PK).get()); assertNoTableDescribeAuth( () -> guestAdmin.getLatestLakeSnapshot(DATA1_TABLE_PATH_PK).get()); assertNoTableDescribeAuth( @@ -467,6 +469,12 @@ void testDescribeTableOperation() throws Exception { assertThat(guestAdmin.tableExists(DATA1_TABLE_PATH_PK).get()).isTrue(); assertThat(guestAdmin.getLatestKvSnapshots(DATA1_TABLE_PATH_PK).get().getBucketIds()) .containsExactlyInAnyOrder(0, 1, 2); + assertThat(guestAdmin.describeBuckets(DATA1_TABLE_PATH_PK).get()) + .hasSize(3) + .allSatisfy( + bucketInfo -> + assertThat(bucketInfo.getTablePath()) + .isEqualTo(DATA1_TABLE_PATH_PK)); assertThatThrownBy(() -> guestAdmin.listPartitionInfos(DATA1_TABLE_PATH_PK).get()) .rootCause() .isInstanceOf(TableNotPartitionedException.class) diff --git a/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java b/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java index 07c63e464e4..fbc016e2f5c 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java @@ -21,13 +21,17 @@ import org.apache.fluss.client.write.ReadyWriteBatch; import org.apache.fluss.memory.MemorySegment; import org.apache.fluss.memory.PreAllocatedPagedOutputView; +import org.apache.fluss.metadata.BucketInfo; import org.apache.fluss.metadata.KvFormat; import org.apache.fluss.metadata.PartitionInfo; import org.apache.fluss.metadata.PhysicalTablePath; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableChange; +import org.apache.fluss.metadata.TablePath; import org.apache.fluss.rpc.messages.AlterTableRequest; +import org.apache.fluss.rpc.messages.DescribeBucketsResponse; import org.apache.fluss.rpc.messages.ListPartitionInfosResponse; +import org.apache.fluss.rpc.messages.PbBucketInfo; import org.apache.fluss.rpc.messages.PbKeyValue; import org.apache.fluss.rpc.messages.PbPartitionInfo; import org.apache.fluss.rpc.messages.PbPartitionSpec; @@ -198,6 +202,50 @@ private static PbPartitionInfo makePbPartitionInfo( return pbPartitionInfo; } + @Test + void testToBucketInfos() { + DescribeBucketsResponse response = new DescribeBucketsResponse(); + PbBucketInfo tableBucket = + response.addBucketInfo().setTableId(10L).setBucketId(0).setLeaderId(1); + tableBucket.setTablePath().setDatabaseName("db").setTableName("table"); + tableBucket.setLeaderEpoch(7); + tableBucket.addReplicaId(1); + tableBucket.addReplicaId(2); + tableBucket.addReplicaId(3); + tableBucket.addIsrId(1); + tableBucket.addIsrId(3); + + PbBucketInfo partitionBucket = response.addBucketInfo().setTableId(10L).setBucketId(1); + partitionBucket.setTablePath().setDatabaseName("db").setTableName("table"); + partitionBucket.setPartitionId(100L).setPartitionName("p1"); + partitionBucket.addReplicaId(2); + partitionBucket.addReplicaId(3); + + List bucketInfos = ClientRpcMessageUtils.toBucketInfos(response); + + assertThat(bucketInfos).hasSize(2); + BucketInfo tableBucketInfo = bucketInfos.get(0); + assertThat(tableBucketInfo.getTablePath()).isEqualTo(TablePath.of("db", "table")); + assertThat(tableBucketInfo.getTableId()).isEqualTo(10L); + assertThat(tableBucketInfo.getPartitionId()).isEmpty(); + assertThat(tableBucketInfo.getPartitionName()).isNull(); + assertThat(tableBucketInfo.getBucketId()).isEqualTo(0); + assertThat(tableBucketInfo.getLeaderId()).hasValue(1); + assertThat(tableBucketInfo.getLeaderEpoch()).hasValue(7); + assertThat(tableBucketInfo.getReplicas()).containsExactly(1, 2, 3); + assertThat(tableBucketInfo.getIsr()).containsExactly(1, 3); + + BucketInfo partitionBucketInfo = bucketInfos.get(1); + assertThat(partitionBucketInfo.getTablePath()).isEqualTo(TablePath.of("db", "table")); + assertThat(partitionBucketInfo.getPartitionId()).hasValue(100L); + assertThat(partitionBucketInfo.getPartitionName()).isEqualTo("p1"); + assertThat(partitionBucketInfo.getBucketId()).isEqualTo(1); + assertThat(partitionBucketInfo.getLeaderId()).isEmpty(); + assertThat(partitionBucketInfo.getLeaderEpoch()).isEmpty(); + assertThat(partitionBucketInfo.getReplicas()).containsExactly(2, 3); + assertThat(partitionBucketInfo.getIsr()).isEmpty(); + } + private KvWriteBatch createKvWriteBatch(int bucketId, MergeMode mergeMode) throws Exception { MemorySegment segment = MemorySegment.allocateHeapMemory(1024); PreAllocatedPagedOutputView outputView = diff --git a/fluss-common/src/main/java/org/apache/fluss/metadata/BucketInfo.java b/fluss-common/src/main/java/org/apache/fluss/metadata/BucketInfo.java new file mode 100644 index 00000000000..e2d6b67a095 --- /dev/null +++ b/fluss-common/src/main/java/org/apache/fluss/metadata/BucketInfo.java @@ -0,0 +1,180 @@ +/* + * 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.fluss.metadata; + +import org.apache.fluss.annotation.PublicEvolving; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.OptionalInt; +import java.util.OptionalLong; + +import static org.apache.fluss.utils.Preconditions.checkNotNull; + +/** + * Information about a physical table bucket, including its replicas and leader/ISR state. + * + * @since 1.0 + */ +@PublicEvolving +public final class BucketInfo { + private final TablePath tablePath; + private final long tableId; + private final @Nullable Long partitionId; + private final @Nullable String partitionName; + private final int bucketId; + private final @Nullable Integer leaderId; + private final @Nullable Integer leaderEpoch; + private final List replicas; + private final List isr; + + /** Creates bucket information. */ + public BucketInfo( + TablePath tablePath, + long tableId, + @Nullable Long partitionId, + @Nullable String partitionName, + int bucketId, + @Nullable Integer leaderId, + @Nullable Integer leaderEpoch, + List replicas, + List isr) { + this.tablePath = checkNotNull(tablePath, "tablePath should not be null."); + this.tableId = tableId; + this.partitionId = partitionId; + this.partitionName = partitionName; + this.bucketId = bucketId; + this.leaderId = leaderId; + this.leaderEpoch = leaderEpoch; + this.replicas = + Collections.unmodifiableList( + new ArrayList<>(checkNotNull(replicas, "replicas should not be null."))); + this.isr = + Collections.unmodifiableList( + new ArrayList<>(checkNotNull(isr, "isr should not be null."))); + } + + /** Returns the table path. */ + public TablePath getTablePath() { + return tablePath; + } + + /** Returns the table ID. */ + public long getTableId() { + return tableId; + } + + /** Returns the partition ID, or an empty optional for a non-partitioned table. */ + public OptionalLong getPartitionId() { + return partitionId == null ? OptionalLong.empty() : OptionalLong.of(partitionId); + } + + /** Returns the partition name, or {@code null} for a non-partitioned table. */ + @Nullable + public String getPartitionName() { + return partitionName; + } + + /** Returns the bucket ID. */ + public int getBucketId() { + return bucketId; + } + + /** Returns the leader ID, or an empty optional if no leader has been elected. */ + public OptionalInt getLeaderId() { + return leaderId == null ? OptionalInt.empty() : OptionalInt.of(leaderId); + } + + /** Returns the leader epoch, or an empty optional if no leader has been elected. */ + public OptionalInt getLeaderEpoch() { + return leaderEpoch == null ? OptionalInt.empty() : OptionalInt.of(leaderEpoch); + } + + /** Returns the replica IDs. */ + public List getReplicas() { + return replicas; + } + + /** Returns the in-sync replica IDs. */ + public List getIsr() { + return isr; + } + + @Override + public String toString() { + return "BucketInfo{" + + "tablePath=" + + tablePath + + ", tableId=" + + tableId + + ", partitionId=" + + partitionId + + ", partitionName='" + + partitionName + + '\'' + + ", bucketId=" + + bucketId + + ", leaderId=" + + leaderId + + ", leaderEpoch=" + + leaderEpoch + + ", replicas=" + + replicas + + ", isr=" + + isr + + '}'; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof BucketInfo)) { + return false; + } + BucketInfo that = (BucketInfo) o; + return tableId == that.tableId + && bucketId == that.bucketId + && Objects.equals(tablePath, that.tablePath) + && Objects.equals(partitionId, that.partitionId) + && Objects.equals(partitionName, that.partitionName) + && Objects.equals(leaderId, that.leaderId) + && Objects.equals(leaderEpoch, that.leaderEpoch) + && replicas.equals(that.replicas) + && isr.equals(that.isr); + } + + @Override + public int hashCode() { + return Objects.hash( + tablePath, + tableId, + partitionId, + partitionName, + bucketId, + leaderId, + leaderEpoch, + replicas, + isr); + } +} diff --git a/fluss-common/src/test/java/org/apache/fluss/metadata/BucketInfoTest.java b/fluss-common/src/test/java/org/apache/fluss/metadata/BucketInfoTest.java new file mode 100644 index 00000000000..a931b01af85 --- /dev/null +++ b/fluss-common/src/test/java/org/apache/fluss/metadata/BucketInfoTest.java @@ -0,0 +1,130 @@ +/* + * 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.fluss.metadata; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Test for {@link BucketInfo}. */ +class BucketInfoTest { + + @Test + void testBucketInfoWithPartitionAndLeader() { + TablePath tablePath = TablePath.of("db", "table"); + List replicas = new ArrayList<>(Arrays.asList(1, 2, 3)); + List isr = new ArrayList<>(Arrays.asList(1, 3)); + + BucketInfo bucketInfo = new BucketInfo(tablePath, 10L, 100L, "p1", 0, 1, 7, replicas, isr); + + assertThat(bucketInfo.getTablePath()).isEqualTo(tablePath); + assertThat(bucketInfo.getTableId()).isEqualTo(10L); + assertThat(bucketInfo.getPartitionId()).hasValue(100L); + assertThat(bucketInfo.getPartitionName()).isEqualTo("p1"); + assertThat(bucketInfo.getBucketId()).isEqualTo(0); + assertThat(bucketInfo.getLeaderId()).hasValue(1); + assertThat(bucketInfo.getLeaderEpoch()).hasValue(7); + assertThat(bucketInfo.getReplicas()).containsExactly(1, 2, 3); + assertThat(bucketInfo.getIsr()).containsExactly(1, 3); + + replicas.add(4); + isr.clear(); + assertThat(bucketInfo.getReplicas()).containsExactly(1, 2, 3); + assertThat(bucketInfo.getIsr()).containsExactly(1, 3); + assertThatThrownBy(() -> bucketInfo.getReplicas().add(4)) + .isInstanceOf(UnsupportedOperationException.class); + assertThatThrownBy(() -> bucketInfo.getIsr().clear()) + .isInstanceOf(UnsupportedOperationException.class); + } + + @Test + void testBucketInfoWithoutPartitionAndLeader() { + BucketInfo bucketInfo = + new BucketInfo( + TablePath.of("db", "table"), + 10L, + null, + null, + 0, + null, + null, + Collections.singletonList(1), + Collections.emptyList()); + + assertThat(bucketInfo.getPartitionId()).isEmpty(); + assertThat(bucketInfo.getPartitionName()).isNull(); + assertThat(bucketInfo.getLeaderId()).isEmpty(); + assertThat(bucketInfo.getLeaderEpoch()).isEmpty(); + assertThat(bucketInfo.getReplicas()).containsExactly(1); + assertThat(bucketInfo.getIsr()).isEmpty(); + } + + @Test + void testBucketInfoRejectsNullRequiredFields() { + TablePath tablePath = TablePath.of("db", "table"); + + assertThatThrownBy( + () -> + new BucketInfo( + null, + 10L, + null, + null, + 0, + null, + null, + Collections.singletonList(1), + Collections.emptyList())) + .isInstanceOf(NullPointerException.class) + .hasMessageContaining("tablePath should not be null"); + assertThatThrownBy( + () -> + new BucketInfo( + tablePath, + 10L, + null, + null, + 0, + null, + null, + null, + Collections.emptyList())) + .isInstanceOf(NullPointerException.class) + .hasMessageContaining("replicas should not be null"); + assertThatThrownBy( + () -> + new BucketInfo( + tablePath, + 10L, + null, + null, + 0, + null, + null, + Collections.singletonList(1), + null)) + .isInstanceOf(NullPointerException.class) + .hasMessageContaining("isr should not be null"); + } +} diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/testutils/TestAdminAdapter.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/testutils/TestAdminAdapter.java index a9f0de762df..5979fb2bd80 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/testutils/TestAdminAdapter.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/testutils/TestAdminAdapter.java @@ -37,6 +37,7 @@ import org.apache.fluss.cluster.rebalance.ServerTag; import org.apache.fluss.config.cluster.AlterConfig; import org.apache.fluss.config.cluster.ConfigEntry; +import org.apache.fluss.metadata.BucketInfo; import org.apache.fluss.metadata.DatabaseChange; import org.apache.fluss.metadata.DatabaseDescriptor; import org.apache.fluss.metadata.DatabaseInfo; @@ -147,6 +148,17 @@ public CompletableFuture getTableInfo(TablePath tablePath) { throw new UnsupportedOperationException("Not implemented in TestAdminAdapter"); } + @Override + public CompletableFuture> describeBuckets(TablePath tablePath) { + throw new UnsupportedOperationException("Not implemented in TestAdminAdapter"); + } + + @Override + public CompletableFuture> describeBuckets( + TablePath tablePath, PartitionSpec partitionSpec) { + throw new UnsupportedOperationException("Not implemented in TestAdminAdapter"); + } + @Override public CompletableFuture dropTable(TablePath tablePath, boolean ignoreIfNotExists) { throw new UnsupportedOperationException("Not implemented in TestAdminAdapter"); diff --git a/fluss-rpc/src/main/java/org/apache/fluss/rpc/gateway/AdminReadOnlyGateway.java b/fluss-rpc/src/main/java/org/apache/fluss/rpc/gateway/AdminReadOnlyGateway.java index 574e1a510dd..639d6de8a03 100644 --- a/fluss-rpc/src/main/java/org/apache/fluss/rpc/gateway/AdminReadOnlyGateway.java +++ b/fluss-rpc/src/main/java/org/apache/fluss/rpc/gateway/AdminReadOnlyGateway.java @@ -20,6 +20,8 @@ import org.apache.fluss.rpc.RpcGateway; import org.apache.fluss.rpc.messages.DatabaseExistsRequest; import org.apache.fluss.rpc.messages.DatabaseExistsResponse; +import org.apache.fluss.rpc.messages.DescribeBucketsRequest; +import org.apache.fluss.rpc.messages.DescribeBucketsResponse; import org.apache.fluss.rpc.messages.DescribeClusterConfigsRequest; import org.apache.fluss.rpc.messages.DescribeClusterConfigsResponse; import org.apache.fluss.rpc.messages.GetClusterHealthRequest; @@ -104,6 +106,15 @@ public interface AdminReadOnlyGateway extends RpcGateway { @RPC(api = ApiKeys.GET_TABLE_INFO) CompletableFuture getTableInfo(GetTableInfoRequest request); + /** + * Describes bucket metadata for a table, optionally filtered by a partition spec. + * + * @param request Request containing the table path and optional partition spec + * @return The bucket metadata response + */ + @RPC(api = ApiKeys.DESCRIBE_BUCKETS) + CompletableFuture describeBuckets(DescribeBucketsRequest request); + /** * Return a {@link GetTableSchemaResponse} identified by the given {@link * GetTableSchemaRequest}. diff --git a/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java b/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java index 9356f35f224..fd08a540a4a 100644 --- a/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java +++ b/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java @@ -114,7 +114,8 @@ public enum ApiKeys { LIST_REMOTE_LOG_MANIFESTS(1063, 0, 0, PUBLIC), LIST_KV_SNAPSHOTS(1064, 0, 0, PUBLIC), ADD_SERVER_TAG_BY_RACK(1065, 0, 0, PUBLIC), - REMOVE_SERVER_TAG_BY_RACK(1066, 0, 0, PUBLIC); + REMOVE_SERVER_TAG_BY_RACK(1066, 0, 0, PUBLIC), + DESCRIBE_BUCKETS(1067, 0, 0, PUBLIC); private static final Map ID_TO_TYPE = Arrays.stream(ApiKeys.values()) diff --git a/fluss-rpc/src/main/proto/FlussApi.proto b/fluss-rpc/src/main/proto/FlussApi.proto index d35d235fe8e..56eb5ead360 100644 --- a/fluss-rpc/src/main/proto/FlussApi.proto +++ b/fluss-rpc/src/main/proto/FlussApi.proto @@ -165,6 +165,16 @@ message GetTableInfoResponse { optional int64 bucket_count_epoch = 7; } +// describe buckets request and response +message DescribeBucketsRequest { + required PbTablePath table_path = 1; + optional PbPartitionSpec partition_spec = 2; +} + +message DescribeBucketsResponse { + repeated PbBucketInfo bucket_info = 1; +} + // list tables request and response message ListTablesRequest { required string database_name = 1; @@ -916,6 +926,19 @@ message PbBucketMetadata { repeated int32 isr = 6 [packed = true]; } +message PbBucketInfo { + required PbTablePath table_path = 1; + required int64 table_id = 2; + optional int64 partition_id = 3; + optional string partition_name = 4; + required int32 bucket_id = 5; + // optional as the leader may not be elected yet + optional int32 leader_id = 6; + repeated int32 replica_id = 7 [packed = true]; + optional int32 leader_epoch = 8; + repeated int32 isr_id = 9 [packed = true]; +} + message PbProduceLogReqForBucket { optional int64 partition_id = 1; required int32 bucket_id = 2; diff --git a/fluss-rpc/src/test/java/org/apache/fluss/rpc/TestingTabletGatewayService.java b/fluss-rpc/src/test/java/org/apache/fluss/rpc/TestingTabletGatewayService.java index f465bb4a69a..3a71bd7e21d 100644 --- a/fluss-rpc/src/test/java/org/apache/fluss/rpc/TestingTabletGatewayService.java +++ b/fluss-rpc/src/test/java/org/apache/fluss/rpc/TestingTabletGatewayService.java @@ -21,6 +21,8 @@ import org.apache.fluss.rpc.gateway.TabletServerGateway; import org.apache.fluss.rpc.messages.DatabaseExistsRequest; import org.apache.fluss.rpc.messages.DatabaseExistsResponse; +import org.apache.fluss.rpc.messages.DescribeBucketsRequest; +import org.apache.fluss.rpc.messages.DescribeBucketsResponse; import org.apache.fluss.rpc.messages.DescribeClusterConfigsRequest; import org.apache.fluss.rpc.messages.DescribeClusterConfigsResponse; import org.apache.fluss.rpc.messages.FetchLogRequest; @@ -201,6 +203,12 @@ public CompletableFuture getTableInfo(GetTableInfoRequest return null; } + @Override + public CompletableFuture describeBuckets( + DescribeBucketsRequest request) { + return null; + } + @Override public CompletableFuture getTableSchema(GetTableSchemaRequest request) { return null; diff --git a/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java b/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java index 782c0629c28..0d676574d10 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java @@ -46,6 +46,8 @@ import org.apache.fluss.rpc.messages.ApiVersionsResponse; import org.apache.fluss.rpc.messages.DatabaseExistsRequest; import org.apache.fluss.rpc.messages.DatabaseExistsResponse; +import org.apache.fluss.rpc.messages.DescribeBucketsRequest; +import org.apache.fluss.rpc.messages.DescribeBucketsResponse; import org.apache.fluss.rpc.messages.DescribeClusterConfigsRequest; import org.apache.fluss.rpc.messages.DescribeClusterConfigsResponse; import org.apache.fluss.rpc.messages.GetDatabaseInfoRequest; @@ -73,6 +75,7 @@ import org.apache.fluss.rpc.messages.MetadataRequest; import org.apache.fluss.rpc.messages.MetadataResponse; import org.apache.fluss.rpc.messages.PbApiVersion; +import org.apache.fluss.rpc.messages.PbBucketInfo; import org.apache.fluss.rpc.messages.PbTablePath; import org.apache.fluss.rpc.messages.TableExistsRequest; import org.apache.fluss.rpc.messages.TableExistsResponse; @@ -87,6 +90,7 @@ import org.apache.fluss.server.coordinator.CoordinatorService; import org.apache.fluss.server.coordinator.MetadataManager; import org.apache.fluss.server.kv.snapshot.CompletedSnapshot; +import org.apache.fluss.server.metadata.BucketMetadata; import org.apache.fluss.server.metadata.MetadataProvider; import org.apache.fluss.server.metadata.PartitionMetadata; import org.apache.fluss.server.metadata.PartitionNegativeCache; @@ -106,6 +110,8 @@ import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; +import java.util.Comparator; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -324,6 +330,131 @@ public CompletableFuture getTableInfo(GetTableInfoRequest return CompletableFuture.completedFuture(response); } + @Override + public CompletableFuture describeBuckets( + DescribeBucketsRequest request) { + TablePath tablePath = toTablePath(request.getTablePath()); + authorizeTable(OperationType.DESCRIBE, tablePath); + + TableInfo tableInfo = metadataManager.getTable(tablePath); + DescribeBucketsResponse response = new DescribeBucketsResponse(); + if (tableInfo.isPartitioned()) { + Map partitionRegistrations = + listPartitionsForDescribeBuckets(request, tablePath); + partitionRegistrations.remove(HISTORICAL_PARTITION_VALUE); + Map> partitionBucketMetadata = + getPartitionBucketMetadataForDescribeBuckets( + tablePath, + partitionRegistrations.values().stream() + .map(PartitionRegistration::getPartitionId) + .collect(Collectors.toList())); + partitionRegistrations.entrySet().stream() + .sorted(Map.Entry.comparingByKey()) + .forEach( + entry -> { + long partitionId = entry.getValue().getPartitionId(); + addBucketInfos( + response, + tablePath, + tableInfo.getTableId(), + partitionId, + entry.getKey(), + partitionBucketMetadata.getOrDefault( + partitionId, Collections.emptyList())); + }); + } else { + if (request.hasPartitionSpec()) { + throw new TableNotPartitionedException( + "Table '" + tablePath + "' is not a partitioned table."); + } + addBucketInfos( + response, + tablePath, + tableInfo.getTableId(), + null, + null, + getTableBucketMetadataForDescribeBuckets(tablePath, tableInfo.getTableId())); + } + return CompletableFuture.completedFuture(response); + } + + private Map listPartitionsForDescribeBuckets( + DescribeBucketsRequest request, TablePath tablePath) { + if (request.hasPartitionSpec()) { + return metadataManager.listPartitions( + tablePath, toResolvedPartitionSpec(request.getPartitionSpec())); + } + return metadataManager.listPartitions(tablePath); + } + + private Map> getPartitionBucketMetadataForDescribeBuckets( + TablePath tablePath, Collection partitionIds) { + try { + return zkClient.getBucketMetadataForPartitions(partitionIds); + } catch (Exception e) { + throw new FlussRuntimeException( + String.format("Failed to describe buckets for table '%s'.", tablePath), e); + } + } + + private List getTableBucketMetadataForDescribeBuckets( + TablePath tablePath, long tableId) { + try { + return zkClient.getBucketMetadataForTables(Collections.singleton(tableId)) + .getOrDefault(tableId, Collections.emptyList()); + } catch (Exception e) { + throw new FlussRuntimeException( + String.format("Failed to describe buckets for table '%s'.", tablePath), e); + } + } + + private static void addBucketInfos( + DescribeBucketsResponse response, + TablePath tablePath, + long tableId, + @Nullable Long partitionId, + @Nullable String partitionName, + List bucketMetadataList) { + bucketMetadataList.stream() + .sorted(Comparator.comparingInt(BucketMetadata::getBucketId)) + .forEach( + bucketMetadata -> + addBucketInfo( + response, + tablePath, + tableId, + partitionId, + partitionName, + bucketMetadata)); + } + + private static void addBucketInfo( + DescribeBucketsResponse response, + TablePath tablePath, + long tableId, + @Nullable Long partitionId, + @Nullable String partitionName, + BucketMetadata bucketMetadata) { + PbBucketInfo pbBucketInfo = + response.addBucketInfo() + .setTableId(tableId) + .setBucketId(bucketMetadata.getBucketId()); + pbBucketInfo + .setTablePath() + .setDatabaseName(tablePath.getDatabaseName()) + .setTableName(tablePath.getTableName()); + if (partitionId != null) { + pbBucketInfo.setPartitionId(partitionId); + } + if (partitionName != null) { + pbBucketInfo.setPartitionName(partitionName); + } + bucketMetadata.getLeaderId().ifPresent(pbBucketInfo::setLeaderId); + bucketMetadata.getLeaderEpoch().ifPresent(pbBucketInfo::setLeaderEpoch); + bucketMetadata.getReplicas().forEach(pbBucketInfo::addReplicaId); + bucketMetadata.getIsr().forEach(pbBucketInfo::addIsrId); + } + @Override public CompletableFuture getTableSchema(GetTableSchemaRequest request) { TablePath tablePath = toTablePath(request.getTablePath()); diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java index 1ce16a9627b..833c16825dd 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java @@ -19,6 +19,7 @@ import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; +import org.apache.fluss.exception.ApiException; import org.apache.fluss.exception.DatabaseAlreadyExistException; import org.apache.fluss.exception.DatabaseNotEmptyException; import org.apache.fluss.exception.DatabaseNotExistException; @@ -315,6 +316,8 @@ public Map listPartitions( return zookeeperClient.getPartitionRegistrations( tablePath, tableInfo.getPartitionKeys(), partitionFilter); } + } catch (ApiException e) { + throw e; } catch (Exception e) { throw new FlussRuntimeException( String.format( diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java index 010b2e840ee..17da18dbdd6 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java @@ -941,12 +941,7 @@ public Map> getPartitionsForTables(Collection /** Get the partition registrations of a table in ZK. */ public Map getPartitionRegistrations(TablePath tablePath) throws Exception { - Map partitions = new HashMap<>(); - for (String partitionName : getPartitions(tablePath)) { - Optional optPartition = getPartition(tablePath, partitionName); - optPartition.ifPresent(partition -> partitions.put(partitionName, partition)); - } - return partitions; + return getPartitionRegistrations(tablePath, getPartitions(tablePath)); } /** Get the partition and the id for the partitions of tables in ZK. */ @@ -994,20 +989,38 @@ public Map getPartitionRegistrations( List partitionKeys, ResolvedPartitionSpec partialPartitionSpec) throws Exception { - Map partitions = new HashMap<>(); - - for (String partitionName : getPartitions(tablePath)) { - ResolvedPartitionSpec resolvedPartitionSpec = - fromPartitionName(partitionKeys, partitionName); - boolean contains = resolvedPartitionSpec.contains(partialPartitionSpec); - if (contains) { - Optional optPartition = - getPartition(tablePath, partitionName); - optPartition.ifPresent(partition -> partitions.put(partitionName, partition)); - } - } - - return partitions; + List matchedPartitionNames = + getPartitions(tablePath).stream() + .filter( + partitionName -> + fromPartitionName(partitionKeys, partitionName) + .contains(partialPartitionSpec)) + .collect(Collectors.toList()); + return getPartitionRegistrations(tablePath, matchedPartitionNames); + } + + private Map getPartitionRegistrations( + TablePath tablePath, Collection partitionNames) throws Exception { + Map path2PartitionName = + partitionNames.stream() + .collect( + toMap( + partitionName -> + PartitionZNode.path(tablePath, partitionName), + partitionName -> partitionName)); + List responses = getDataInBackground(path2PartitionName.keySet()); + return processGetDataResponses( + responses, + response -> path2PartitionName.get(response.getPath()), + data -> { + PartitionRegistration partitionRegistration = PartitionZNode.decode(data); + if (partitionRegistration.getRemoteDataDir() == null) { + partitionRegistration = + partitionRegistration.newRemoteDataDir(defaultRemoteDataDir); + } + return partitionRegistration; + }, + "partition registrations"); } /** Get the id and name for the partitions of a table in ZK. */ diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestCoordinatorGateway.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestCoordinatorGateway.java index 3d1beeea910..e553d625499 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestCoordinatorGateway.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestCoordinatorGateway.java @@ -60,6 +60,8 @@ import org.apache.fluss.rpc.messages.DatabaseExistsResponse; import org.apache.fluss.rpc.messages.DeleteProducerOffsetsRequest; import org.apache.fluss.rpc.messages.DeleteProducerOffsetsResponse; +import org.apache.fluss.rpc.messages.DescribeBucketsRequest; +import org.apache.fluss.rpc.messages.DescribeBucketsResponse; import org.apache.fluss.rpc.messages.DescribeClusterConfigsRequest; import org.apache.fluss.rpc.messages.DescribeClusterConfigsResponse; import org.apache.fluss.rpc.messages.DropAclsRequest; @@ -236,6 +238,12 @@ public CompletableFuture getTableInfo(GetTableInfoRequest throw new UnsupportedOperationException(); } + @Override + public CompletableFuture describeBuckets( + DescribeBucketsRequest request) { + throw new UnsupportedOperationException(); + } + @Override public CompletableFuture getTableSchema(GetTableSchemaRequest request) { throw new UnsupportedOperationException(); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/tablet/TestTabletServerGateway.java b/fluss-server/src/test/java/org/apache/fluss/server/tablet/TestTabletServerGateway.java index cb81cee0462..0e0a96b2030 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/tablet/TestTabletServerGateway.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/tablet/TestTabletServerGateway.java @@ -26,6 +26,8 @@ import org.apache.fluss.rpc.messages.ApiVersionsResponse; import org.apache.fluss.rpc.messages.DatabaseExistsRequest; import org.apache.fluss.rpc.messages.DatabaseExistsResponse; +import org.apache.fluss.rpc.messages.DescribeBucketsRequest; +import org.apache.fluss.rpc.messages.DescribeBucketsResponse; import org.apache.fluss.rpc.messages.DescribeClusterConfigsRequest; import org.apache.fluss.rpc.messages.DescribeClusterConfigsResponse; import org.apache.fluss.rpc.messages.FetchLogRequest; @@ -258,6 +260,12 @@ public CompletableFuture getTableInfo(GetTableInfoRequest throw new UnsupportedOperationException(); } + @Override + public CompletableFuture describeBuckets( + DescribeBucketsRequest request) { + throw new UnsupportedOperationException(); + } + @Override public CompletableFuture getTableSchema(GetTableSchemaRequest request) { throw new UnsupportedOperationException(); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java index 60020824a53..367c5f785df 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java @@ -759,6 +759,11 @@ void testPartition() throws Exception { assertThat(partition.getPartitionId()).isEqualTo(1L); partition = zookeeperClient.getPartition(tablePath, "p2").get(); assertThat(partition.getPartitionId()).isEqualTo(2L); + Map partitionRegistrations = + zookeeperClient.getPartitionRegistrations(tablePath); + assertThat(partitionRegistrations).containsOnlyKeys("p1", "p2"); + assertThat(partitionRegistrations.get("p1").getPartitionId()).isEqualTo(1L); + assertThat(partitionRegistrations.get("p2").getPartitionId()).isEqualTo(2L); assertThat(zookeeperClient.getPartitionsForTables(Arrays.asList(tablePath))) .containsValues(new ArrayList<>(partitions)); From 6220764d20997740db096aaf23e3a817bc7f24e4 Mon Sep 17 00:00:00 2001 From: fhan Date: Tue, 18 Aug 2026 19:32:14 +0800 Subject: [PATCH 2/5] [client] fix test failure in flink2 module --- .../org/apache/fluss/flink/catalog/FlinkCatalogTest.java | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkCatalogTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkCatalogTest.java index a5e348c7af7..cdcee44d372 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkCatalogTest.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkCatalogTest.java @@ -54,6 +54,7 @@ import org.apache.flink.table.catalog.exceptions.DatabaseNotExistException; import org.apache.flink.table.catalog.exceptions.FunctionNotExistException; import org.apache.flink.table.catalog.exceptions.PartitionAlreadyExistsException; +import org.apache.flink.table.catalog.exceptions.PartitionSpecInvalidException; import org.apache.flink.table.catalog.exceptions.TableAlreadyExistException; import org.apache.flink.table.catalog.exceptions.TableNotExistException; import org.apache.flink.table.catalog.exceptions.TableNotPartitionedException; @@ -832,9 +833,9 @@ void testOperatePartitions() throws Exception { CatalogPartitionSpec invalidTestSpec = new CatalogPartitionSpec(Collections.singletonMap("second", "")); assertThatThrownBy(() -> catalog.listPartitions(path2, invalidTestSpec)) - .isInstanceOf(CatalogException.class) - .hasMessage( - "Failed to list partitions of table fluss.partitioned_t1 in test-catalog, by partitionSpec CatalogPartitionSpec{{second=}}"); + .isInstanceOf(PartitionSpecInvalidException.class) + .hasMessageContaining( + "PartitionSpec CatalogPartitionSpec{{second=}} does not match"); // NEW: Test dropPartition functionality CatalogPartitionSpec firstPartSpec = catalogPartitionSpecs.get(0); From 505689807d41d0cd31be5dfc4448e3a34e419c71 Mon Sep 17 00:00:00 2001 From: fhan Date: Thu, 20 Aug 2026 19:50:10 +0800 Subject: [PATCH 3/5] [client] refine code impl according to review comments --- .../client/utils/ClientRpcMessageUtils.java | 16 ++++--- .../utils/ClientRpcMessageUtilsTest.java | 10 ++--- fluss-rpc/src/main/proto/FlussApi.proto | 20 ++++----- .../apache/fluss/server/RpcServiceBase.java | 41 ++++++------------ .../server/coordinator/MetadataManager.java | 16 ++++++- .../fluss/server/zk/ZooKeeperClient.java | 30 +++++++------ .../fluss/server/zk/ZooKeeperClientTest.java | 42 +++++++++++++++++++ 7 files changed, 111 insertions(+), 64 deletions(-) diff --git a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java index 9c9a5be3a02..1dc3d02a87b 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java @@ -700,17 +700,21 @@ public static List toPartitionInfos( } public static List toBucketInfos(DescribeBucketsResponse response) { + TablePath tablePath = + TablePath.of( + response.getTablePath().getDatabaseName(), + response.getTablePath().getTableName()); + long tableId = response.getTableId(); return response.getBucketInfosList().stream() - .map(ClientRpcMessageUtils::toBucketInfo) + .map(pbBucketInfo -> toBucketInfo(tablePath, tableId, pbBucketInfo)) .collect(Collectors.toList()); } - private static BucketInfo toBucketInfo(PbBucketInfo pbBucketInfo) { + private static BucketInfo toBucketInfo( + TablePath tablePath, long tableId, PbBucketInfo pbBucketInfo) { return new BucketInfo( - TablePath.of( - pbBucketInfo.getTablePath().getDatabaseName(), - pbBucketInfo.getTablePath().getTableName()), - pbBucketInfo.getTableId(), + tablePath, + tableId, pbBucketInfo.hasPartitionId() ? pbBucketInfo.getPartitionId() : null, pbBucketInfo.hasPartitionName() ? pbBucketInfo.getPartitionName() : null, pbBucketInfo.getBucketId(), diff --git a/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java b/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java index fbc016e2f5c..392f5f02d24 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java @@ -204,10 +204,9 @@ private static PbPartitionInfo makePbPartitionInfo( @Test void testToBucketInfos() { - DescribeBucketsResponse response = new DescribeBucketsResponse(); - PbBucketInfo tableBucket = - response.addBucketInfo().setTableId(10L).setBucketId(0).setLeaderId(1); - tableBucket.setTablePath().setDatabaseName("db").setTableName("table"); + DescribeBucketsResponse response = new DescribeBucketsResponse().setTableId(10L); + response.setTablePath().setDatabaseName("db").setTableName("table"); + PbBucketInfo tableBucket = response.addBucketInfo().setBucketId(0).setLeaderId(1); tableBucket.setLeaderEpoch(7); tableBucket.addReplicaId(1); tableBucket.addReplicaId(2); @@ -215,8 +214,7 @@ void testToBucketInfos() { tableBucket.addIsrId(1); tableBucket.addIsrId(3); - PbBucketInfo partitionBucket = response.addBucketInfo().setTableId(10L).setBucketId(1); - partitionBucket.setTablePath().setDatabaseName("db").setTableName("table"); + PbBucketInfo partitionBucket = response.addBucketInfo().setBucketId(1); partitionBucket.setPartitionId(100L).setPartitionName("p1"); partitionBucket.addReplicaId(2); partitionBucket.addReplicaId(3); diff --git a/fluss-rpc/src/main/proto/FlussApi.proto b/fluss-rpc/src/main/proto/FlussApi.proto index 56eb5ead360..10f6e9ae2c3 100644 --- a/fluss-rpc/src/main/proto/FlussApi.proto +++ b/fluss-rpc/src/main/proto/FlussApi.proto @@ -172,7 +172,9 @@ message DescribeBucketsRequest { } message DescribeBucketsResponse { - repeated PbBucketInfo bucket_info = 1; + required PbTablePath table_path = 1; + required int64 table_id = 2; + repeated PbBucketInfo bucket_info = 3; } // list tables request and response @@ -927,16 +929,14 @@ message PbBucketMetadata { } message PbBucketInfo { - required PbTablePath table_path = 1; - required int64 table_id = 2; - optional int64 partition_id = 3; - optional string partition_name = 4; - required int32 bucket_id = 5; + optional int64 partition_id = 1; + optional string partition_name = 2; + required int32 bucket_id = 3; // optional as the leader may not be elected yet - optional int32 leader_id = 6; - repeated int32 replica_id = 7 [packed = true]; - optional int32 leader_epoch = 8; - repeated int32 isr_id = 9 [packed = true]; + optional int32 leader_id = 4; + repeated int32 replica_id = 5 [packed = true]; + optional int32 leader_epoch = 6; + repeated int32 isr_id = 7 [packed = true]; } message PbProduceLogReqForBucket { diff --git a/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java b/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java index 0d676574d10..a41e32370a6 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java @@ -337,10 +337,14 @@ public CompletableFuture describeBuckets( authorizeTable(OperationType.DESCRIBE, tablePath); TableInfo tableInfo = metadataManager.getTable(tablePath); - DescribeBucketsResponse response = new DescribeBucketsResponse(); + DescribeBucketsResponse response = + new DescribeBucketsResponse().setTableId(tableInfo.getTableId()); + response.setTablePath() + .setDatabaseName(tablePath.getDatabaseName()) + .setTableName(tablePath.getTableName()); if (tableInfo.isPartitioned()) { Map partitionRegistrations = - listPartitionsForDescribeBuckets(request, tablePath); + listPartitionsForDescribeBuckets(request, tablePath, tableInfo); partitionRegistrations.remove(HISTORICAL_PARTITION_VALUE); Map> partitionBucketMetadata = getPartitionBucketMetadataForDescribeBuckets( @@ -355,8 +359,6 @@ public CompletableFuture describeBuckets( long partitionId = entry.getValue().getPartitionId(); addBucketInfos( response, - tablePath, - tableInfo.getTableId(), partitionId, entry.getKey(), partitionBucketMetadata.getOrDefault( @@ -369,8 +371,6 @@ public CompletableFuture describeBuckets( } addBucketInfos( response, - tablePath, - tableInfo.getTableId(), null, null, getTableBucketMetadataForDescribeBuckets(tablePath, tableInfo.getTableId())); @@ -379,12 +379,12 @@ public CompletableFuture describeBuckets( } private Map listPartitionsForDescribeBuckets( - DescribeBucketsRequest request, TablePath tablePath) { + DescribeBucketsRequest request, TablePath tablePath, TableInfo tableInfo) { if (request.hasPartitionSpec()) { return metadataManager.listPartitions( - tablePath, toResolvedPartitionSpec(request.getPartitionSpec())); + tablePath, tableInfo, toResolvedPartitionSpec(request.getPartitionSpec())); } - return metadataManager.listPartitions(tablePath); + return metadataManager.listPartitions(tablePath, tableInfo, null); } private Map> getPartitionBucketMetadataForDescribeBuckets( @@ -410,8 +410,6 @@ private List getTableBucketMetadataForDescribeBuckets( private static void addBucketInfos( DescribeBucketsResponse response, - TablePath tablePath, - long tableId, @Nullable Long partitionId, @Nullable String partitionName, List bucketMetadataList) { @@ -420,29 +418,16 @@ private static void addBucketInfos( .forEach( bucketMetadata -> addBucketInfo( - response, - tablePath, - tableId, - partitionId, - partitionName, - bucketMetadata)); + response, partitionId, partitionName, bucketMetadata)); } private static void addBucketInfo( DescribeBucketsResponse response, - TablePath tablePath, - long tableId, @Nullable Long partitionId, @Nullable String partitionName, BucketMetadata bucketMetadata) { PbBucketInfo pbBucketInfo = - response.addBucketInfo() - .setTableId(tableId) - .setBucketId(bucketMetadata.getBucketId()); - pbBucketInfo - .setTablePath() - .setDatabaseName(tablePath.getDatabaseName()) - .setTableName(tablePath.getTableName()); + response.addBucketInfo().setBucketId(bucketMetadata.getBucketId()); if (partitionId != null) { pbBucketInfo.setPartitionId(partitionId); } @@ -645,9 +630,9 @@ public CompletableFuture listPartitionInfos( ResolvedPartitionSpec partitionSpecFromRequest = toResolvedPartitionSpec(request.getPartialPartitionSpec()); partitionRegistrations = - metadataManager.listPartitions(tablePath, partitionSpecFromRequest); + metadataManager.listPartitions(tablePath, tableInfo, partitionSpecFromRequest); } else { - partitionRegistrations = metadataManager.listPartitions(tablePath); + partitionRegistrations = metadataManager.listPartitions(tablePath, tableInfo, null); } boolean includeSystemPartitions = request.hasIncludeSystemPartitions() && request.isIncludeSystemPartitions(); diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java index 833c16825dd..14b97c927eb 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java @@ -19,7 +19,6 @@ import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; -import org.apache.fluss.exception.ApiException; import org.apache.fluss.exception.DatabaseAlreadyExistException; import org.apache.fluss.exception.DatabaseNotEmptyException; import org.apache.fluss.exception.DatabaseNotExistException; @@ -304,6 +303,19 @@ public Map listPartitions( TablePath tablePath, ResolvedPartitionSpec partitionFilter) throws TableNotExistException, TableNotPartitionedException, InvalidPartitionException { TableInfo tableInfo = getTable(tablePath); + return listPartitions(tablePath, tableInfo, partitionFilter); + } + + /** + * List the partitions of the given table and partition spec using the provided table metadata. + * + * @return a map from partition name to partition registration. + */ + public Map listPartitions( + TablePath tablePath, + TableInfo tableInfo, + @Nullable ResolvedPartitionSpec partitionFilter) + throws TableNotExistException, TableNotPartitionedException, InvalidPartitionException { if (!tableInfo.isPartitioned()) { throw new TableNotPartitionedException( "Table '" + tablePath + "' is not a partitioned table."); @@ -316,7 +328,7 @@ public Map listPartitions( return zookeeperClient.getPartitionRegistrations( tablePath, tableInfo.getPartitionKeys(), partitionFilter); } - } catch (ApiException e) { + } catch (InvalidPartitionException e) { throw e; } catch (Exception e) { throw new FlussRuntimeException( diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java index 17da18dbdd6..cdb8c90c2ea 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java @@ -1009,18 +1009,24 @@ private Map getPartitionRegistrations( PartitionZNode.path(tablePath, partitionName), partitionName -> partitionName)); List responses = getDataInBackground(path2PartitionName.keySet()); - return processGetDataResponses( - responses, - response -> path2PartitionName.get(response.getPath()), - data -> { - PartitionRegistration partitionRegistration = PartitionZNode.decode(data); - if (partitionRegistration.getRemoteDataDir() == null) { - partitionRegistration = - partitionRegistration.newRemoteDataDir(defaultRemoteDataDir); - } - return partitionRegistration; - }, - "partition registrations"); + Map partitionRegistrations = new HashMap<>(); + for (ZkGetDataResponse response : responses) { + if (response.getResultCode() == KeeperException.Code.NONODE) { + continue; + } + if (response.getResultCode() != KeeperException.Code.OK) { + throw KeeperException.create(response.getResultCode(), response.getPath()); + } + + PartitionRegistration partitionRegistration = PartitionZNode.decode(response.getData()); + if (partitionRegistration.getRemoteDataDir() == null) { + partitionRegistration = + partitionRegistration.newRemoteDataDir(defaultRemoteDataDir); + } + partitionRegistrations.put( + path2PartitionName.get(response.getPath()), partitionRegistration); + } + return partitionRegistrations; } /** Get the id and name for the partitions of a table in ZK. */ diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java index 367c5f785df..c23e11238c3 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java @@ -31,6 +31,7 @@ import org.apache.fluss.metadata.TableDescriptor; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.server.entity.RegisterTableBucketLeadAndIsrInfo; +import org.apache.fluss.server.zk.ZkAsyncResponse.ZkGetDataResponse; import org.apache.fluss.server.zk.data.BucketAssignment; import org.apache.fluss.server.zk.data.BucketSnapshot; import org.apache.fluss.server.zk.data.CoordinatorAddress; @@ -43,6 +44,7 @@ import org.apache.fluss.server.zk.data.TableRegistration; import org.apache.fluss.server.zk.data.TabletServerRegistration; import org.apache.fluss.server.zk.data.ZkData.BucketIdZNode; +import org.apache.fluss.server.zk.data.ZkData.PartitionZNode; import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadata; import org.apache.fluss.shaded.curator5.org.apache.curator.CuratorZookeeperClient; import org.apache.fluss.shaded.curator5.org.apache.curator.framework.CuratorFramework; @@ -66,6 +68,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; @@ -80,6 +83,7 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.anyCollection; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; @@ -795,6 +799,44 @@ void testPartition() throws Exception { assertThat(partitions).containsExactly("p2"); } + @Test + void testGetPartitionRegistrationsPreservesZookeeperErrors() throws Exception { + TablePath tablePath = TablePath.of("db", "tb"); + String partition1Path = PartitionZNode.path(tablePath, "p1"); + String partition2Path = PartitionZNode.path(tablePath, "p2"); + PartitionRegistration partitionRegistration = + new PartitionRegistration(1L, 2L, remoteDataDir); + + ZooKeeperClient testingClient = spy(zookeeperClient); + doReturn(new HashSet<>(Arrays.asList("p1", "p2"))) + .when(testingClient) + .getPartitions(tablePath); + doReturn( + Arrays.asList( + new ZkGetDataResponse( + partition1Path, + KeeperException.Code.OK, + PartitionZNode.encode(partitionRegistration)), + new ZkGetDataResponse( + partition2Path, KeeperException.Code.NONODE, null))) + .when(testingClient) + .getDataInBackground(anyCollection()); + + assertThat(testingClient.getPartitionRegistrations(tablePath)) + .containsOnlyKeys("p1") + .containsValue(partitionRegistration); + + doReturn( + Collections.singletonList( + new ZkGetDataResponse( + partition1Path, KeeperException.Code.CONNECTIONLOSS, null))) + .when(testingClient) + .getDataInBackground(anyCollection()); + + assertThatThrownBy(() -> testingClient.getPartitionRegistrations(tablePath)) + .isInstanceOf(KeeperException.ConnectionLossException.class); + } + @Test void testServerTag() throws Exception { Map serverTags = new HashMap<>(); From 18a232202581d1aaf2f216b092bfafa9c86681fb Mon Sep 17 00:00:00 2001 From: fhan Date: Wed, 2 Sep 2026 20:41:54 +0800 Subject: [PATCH 4/5] [client] Complete describe buckets metadata semantics --- .../client/utils/ClientRpcMessageUtils.java | 3 +- .../client/admin/DescribeBucketsITCase.java | 4 + .../utils/ClientRpcMessageUtilsTest.java | 7 +- .../org/apache/fluss/metadata/BucketInfo.java | 15 +++ .../apache/fluss/metadata/BucketInfoTest.java | 9 +- fluss-rpc/src/main/proto/FlussApi.proto | 5 +- fluss-rust/crates/fluss/proto/FlussApi.proto | 26 +++++ fluss-rust/crates/fluss/src/proto/fluss.rs | 39 ++++++++ .../apache/fluss/server/RpcServiceBase.java | 16 ++- .../fluss/server/zk/ZooKeeperClient.java | 76 ++++++++------- .../fluss/server/RpcServiceBaseTest.java | 25 ++++- .../fluss/server/zk/ZooKeeperClientTest.java | 97 +++++++++++++++++++ 12 files changed, 279 insertions(+), 43 deletions(-) diff --git a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java index 1dc3d02a87b..856c426f08f 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java @@ -720,8 +720,9 @@ private static BucketInfo toBucketInfo( pbBucketInfo.getBucketId(), pbBucketInfo.hasLeaderId() ? pbBucketInfo.getLeaderId() : null, pbBucketInfo.hasLeaderEpoch() ? pbBucketInfo.getLeaderEpoch() : null, + pbBucketInfo.hasBucketEpoch() ? pbBucketInfo.getBucketEpoch() : null, Arrays.stream(pbBucketInfo.getReplicaIds()).boxed().collect(Collectors.toList()), - Arrays.stream(pbBucketInfo.getIsrIds()).boxed().collect(Collectors.toList())); + Arrays.stream(pbBucketInfo.getIsrs()).boxed().collect(Collectors.toList())); } public static Map toKeyValueMap(List pbKeyValues) { diff --git a/fluss-client/src/test/java/org/apache/fluss/client/admin/DescribeBucketsITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/admin/DescribeBucketsITCase.java index 7d0ade2e1a8..3d246046b3d 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/admin/DescribeBucketsITCase.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/admin/DescribeBucketsITCase.java @@ -181,6 +181,9 @@ private List waitAndDescribeBuckets( && bucketInfo .getLeaderEpoch() .isPresent() + && bucketInfo + .getBucketEpoch() + .isPresent() && !bucketInfo.getIsr().isEmpty()); }, Duration.ofMinutes(1), @@ -208,6 +211,7 @@ private static void assertBucketInfo( assertThat(bucketInfo.getPartitionId()).hasValue(expectedPartitionId); } assertThat(bucketInfo.getReplicas()).hasSize(3); + assertThat(bucketInfo.getBucketEpoch()).isPresent(); assertThat(bucketInfo.getIsr()).isNotEmpty(); assertThat(bucketInfo.getReplicas()).containsAll(bucketInfo.getIsr()); assertThat(bucketInfo.getIsr()).contains(bucketInfo.getLeaderId().getAsInt()); diff --git a/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java b/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java index 392f5f02d24..0c1a24a65af 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java @@ -208,11 +208,12 @@ void testToBucketInfos() { response.setTablePath().setDatabaseName("db").setTableName("table"); PbBucketInfo tableBucket = response.addBucketInfo().setBucketId(0).setLeaderId(1); tableBucket.setLeaderEpoch(7); + tableBucket.setBucketEpoch(8); tableBucket.addReplicaId(1); tableBucket.addReplicaId(2); tableBucket.addReplicaId(3); - tableBucket.addIsrId(1); - tableBucket.addIsrId(3); + tableBucket.addIsr(1); + tableBucket.addIsr(3); PbBucketInfo partitionBucket = response.addBucketInfo().setBucketId(1); partitionBucket.setPartitionId(100L).setPartitionName("p1"); @@ -230,6 +231,7 @@ void testToBucketInfos() { assertThat(tableBucketInfo.getBucketId()).isEqualTo(0); assertThat(tableBucketInfo.getLeaderId()).hasValue(1); assertThat(tableBucketInfo.getLeaderEpoch()).hasValue(7); + assertThat(tableBucketInfo.getBucketEpoch()).hasValue(8); assertThat(tableBucketInfo.getReplicas()).containsExactly(1, 2, 3); assertThat(tableBucketInfo.getIsr()).containsExactly(1, 3); @@ -240,6 +242,7 @@ void testToBucketInfos() { assertThat(partitionBucketInfo.getBucketId()).isEqualTo(1); assertThat(partitionBucketInfo.getLeaderId()).isEmpty(); assertThat(partitionBucketInfo.getLeaderEpoch()).isEmpty(); + assertThat(partitionBucketInfo.getBucketEpoch()).isEmpty(); assertThat(partitionBucketInfo.getReplicas()).containsExactly(2, 3); assertThat(partitionBucketInfo.getIsr()).isEmpty(); } diff --git a/fluss-common/src/main/java/org/apache/fluss/metadata/BucketInfo.java b/fluss-common/src/main/java/org/apache/fluss/metadata/BucketInfo.java index e2d6b67a095..c59d023caf4 100644 --- a/fluss-common/src/main/java/org/apache/fluss/metadata/BucketInfo.java +++ b/fluss-common/src/main/java/org/apache/fluss/metadata/BucketInfo.java @@ -44,6 +44,7 @@ public final class BucketInfo { private final int bucketId; private final @Nullable Integer leaderId; private final @Nullable Integer leaderEpoch; + private final @Nullable Integer bucketEpoch; private final List replicas; private final List isr; @@ -56,6 +57,7 @@ public BucketInfo( int bucketId, @Nullable Integer leaderId, @Nullable Integer leaderEpoch, + @Nullable Integer bucketEpoch, List replicas, List isr) { this.tablePath = checkNotNull(tablePath, "tablePath should not be null."); @@ -65,6 +67,7 @@ public BucketInfo( this.bucketId = bucketId; this.leaderId = leaderId; this.leaderEpoch = leaderEpoch; + this.bucketEpoch = bucketEpoch; this.replicas = Collections.unmodifiableList( new ArrayList<>(checkNotNull(replicas, "replicas should not be null."))); @@ -109,6 +112,14 @@ public OptionalInt getLeaderEpoch() { return leaderEpoch == null ? OptionalInt.empty() : OptionalInt.of(leaderEpoch); } + /** + * Returns the generation of the leader/ISR state, or an empty optional for legacy metadata. The + * value {@code -1} indicates that no leader/ISR state exists. + */ + public OptionalInt getBucketEpoch() { + return bucketEpoch == null ? OptionalInt.empty() : OptionalInt.of(bucketEpoch); + } + /** Returns the replica IDs. */ public List getReplicas() { return replicas; @@ -137,6 +148,8 @@ public String toString() { + leaderId + ", leaderEpoch=" + leaderEpoch + + ", bucketEpoch=" + + bucketEpoch + ", replicas=" + replicas + ", isr=" @@ -160,6 +173,7 @@ public boolean equals(Object o) { && Objects.equals(partitionName, that.partitionName) && Objects.equals(leaderId, that.leaderId) && Objects.equals(leaderEpoch, that.leaderEpoch) + && Objects.equals(bucketEpoch, that.bucketEpoch) && replicas.equals(that.replicas) && isr.equals(that.isr); } @@ -174,6 +188,7 @@ public int hashCode() { bucketId, leaderId, leaderEpoch, + bucketEpoch, replicas, isr); } diff --git a/fluss-common/src/test/java/org/apache/fluss/metadata/BucketInfoTest.java b/fluss-common/src/test/java/org/apache/fluss/metadata/BucketInfoTest.java index a931b01af85..2d3a3a87d32 100644 --- a/fluss-common/src/test/java/org/apache/fluss/metadata/BucketInfoTest.java +++ b/fluss-common/src/test/java/org/apache/fluss/metadata/BucketInfoTest.java @@ -36,7 +36,8 @@ void testBucketInfoWithPartitionAndLeader() { List replicas = new ArrayList<>(Arrays.asList(1, 2, 3)); List isr = new ArrayList<>(Arrays.asList(1, 3)); - BucketInfo bucketInfo = new BucketInfo(tablePath, 10L, 100L, "p1", 0, 1, 7, replicas, isr); + BucketInfo bucketInfo = + new BucketInfo(tablePath, 10L, 100L, "p1", 0, 1, 7, 8, replicas, isr); assertThat(bucketInfo.getTablePath()).isEqualTo(tablePath); assertThat(bucketInfo.getTableId()).isEqualTo(10L); @@ -45,6 +46,7 @@ void testBucketInfoWithPartitionAndLeader() { assertThat(bucketInfo.getBucketId()).isEqualTo(0); assertThat(bucketInfo.getLeaderId()).hasValue(1); assertThat(bucketInfo.getLeaderEpoch()).hasValue(7); + assertThat(bucketInfo.getBucketEpoch()).hasValue(8); assertThat(bucketInfo.getReplicas()).containsExactly(1, 2, 3); assertThat(bucketInfo.getIsr()).containsExactly(1, 3); @@ -69,6 +71,7 @@ void testBucketInfoWithoutPartitionAndLeader() { 0, null, null, + null, Collections.singletonList(1), Collections.emptyList()); @@ -76,6 +79,7 @@ void testBucketInfoWithoutPartitionAndLeader() { assertThat(bucketInfo.getPartitionName()).isNull(); assertThat(bucketInfo.getLeaderId()).isEmpty(); assertThat(bucketInfo.getLeaderEpoch()).isEmpty(); + assertThat(bucketInfo.getBucketEpoch()).isEmpty(); assertThat(bucketInfo.getReplicas()).containsExactly(1); assertThat(bucketInfo.getIsr()).isEmpty(); } @@ -94,6 +98,7 @@ void testBucketInfoRejectsNullRequiredFields() { 0, null, null, + null, Collections.singletonList(1), Collections.emptyList())) .isInstanceOf(NullPointerException.class) @@ -109,6 +114,7 @@ void testBucketInfoRejectsNullRequiredFields() { null, null, null, + null, Collections.emptyList())) .isInstanceOf(NullPointerException.class) .hasMessageContaining("replicas should not be null"); @@ -122,6 +128,7 @@ void testBucketInfoRejectsNullRequiredFields() { 0, null, null, + null, Collections.singletonList(1), null)) .isInstanceOf(NullPointerException.class) diff --git a/fluss-rpc/src/main/proto/FlussApi.proto b/fluss-rpc/src/main/proto/FlussApi.proto index 10f6e9ae2c3..08dcfab9e78 100644 --- a/fluss-rpc/src/main/proto/FlussApi.proto +++ b/fluss-rpc/src/main/proto/FlussApi.proto @@ -936,7 +936,10 @@ message PbBucketInfo { optional int32 leader_id = 4; repeated int32 replica_id = 5 [packed = true]; optional int32 leader_epoch = 6; - repeated int32 isr_id = 7 [packed = true]; + // Generation of the complete leader/ISR state. + // Absence indicates legacy metadata; -1 indicates no leader/ISR state exists. + optional int32 bucket_epoch = 7; + repeated int32 isr = 8 [packed = true]; } message PbProduceLogReqForBucket { diff --git a/fluss-rust/crates/fluss/proto/FlussApi.proto b/fluss-rust/crates/fluss/proto/FlussApi.proto index d35d235fe8e..08dcfab9e78 100644 --- a/fluss-rust/crates/fluss/proto/FlussApi.proto +++ b/fluss-rust/crates/fluss/proto/FlussApi.proto @@ -165,6 +165,18 @@ message GetTableInfoResponse { optional int64 bucket_count_epoch = 7; } +// describe buckets request and response +message DescribeBucketsRequest { + required PbTablePath table_path = 1; + optional PbPartitionSpec partition_spec = 2; +} + +message DescribeBucketsResponse { + required PbTablePath table_path = 1; + required int64 table_id = 2; + repeated PbBucketInfo bucket_info = 3; +} + // list tables request and response message ListTablesRequest { required string database_name = 1; @@ -916,6 +928,20 @@ message PbBucketMetadata { repeated int32 isr = 6 [packed = true]; } +message PbBucketInfo { + optional int64 partition_id = 1; + optional string partition_name = 2; + required int32 bucket_id = 3; + // optional as the leader may not be elected yet + optional int32 leader_id = 4; + repeated int32 replica_id = 5 [packed = true]; + optional int32 leader_epoch = 6; + // Generation of the complete leader/ISR state. + // Absence indicates legacy metadata; -1 indicates no leader/ISR state exists. + optional int32 bucket_epoch = 7; + repeated int32 isr = 8 [packed = true]; +} + message PbProduceLogReqForBucket { optional int64 partition_id = 1; required int32 bucket_id = 2; diff --git a/fluss-rust/crates/fluss/src/proto/fluss.rs b/fluss-rust/crates/fluss/src/proto/fluss.rs index bc6578ac845..8121240df8f 100644 --- a/fluss-rust/crates/fluss/src/proto/fluss.rs +++ b/fluss-rust/crates/fluss/src/proto/fluss.rs @@ -195,6 +195,23 @@ pub struct GetTableInfoResponse { #[prost(int64, optional, tag = "7")] pub bucket_count_epoch: ::core::option::Option, } +/// describe buckets request and response +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct DescribeBucketsRequest { + #[prost(message, required, tag = "1")] + pub table_path: PbTablePath, + #[prost(message, optional, tag = "2")] + pub partition_spec: ::core::option::Option, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct DescribeBucketsResponse { + #[prost(message, required, tag = "1")] + pub table_path: PbTablePath, + #[prost(int64, required, tag = "2")] + pub table_id: i64, + #[prost(message, repeated, tag = "3")] + pub bucket_info: ::prost::alloc::vec::Vec, +} /// list tables request and response #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct ListTablesRequest { @@ -1206,6 +1223,28 @@ pub struct PbBucketMetadata { pub isr: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct PbBucketInfo { + #[prost(int64, optional, tag = "1")] + pub partition_id: ::core::option::Option, + #[prost(string, optional, tag = "2")] + pub partition_name: ::core::option::Option<::prost::alloc::string::String>, + #[prost(int32, required, tag = "3")] + pub bucket_id: i32, + /// optional as the leader may not be elected yet + #[prost(int32, optional, tag = "4")] + pub leader_id: ::core::option::Option, + #[prost(int32, repeated, tag = "5")] + pub replica_id: ::prost::alloc::vec::Vec, + #[prost(int32, optional, tag = "6")] + pub leader_epoch: ::core::option::Option, + /// Generation of the complete leader/ISR state. + /// Absence indicates legacy metadata; -1 indicates no leader/ISR state exists. + #[prost(int32, optional, tag = "7")] + pub bucket_epoch: ::core::option::Option, + #[prost(int32, repeated, tag = "8")] + pub isr: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct PbProduceLogReqForBucket { #[prost(int64, optional, tag = "1")] pub partition_id: ::core::option::Option, diff --git a/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java b/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java index a41e32370a6..8be94252d8f 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java @@ -135,6 +135,7 @@ import static org.apache.fluss.server.utils.ServerRpcMessageUtils.toPbConfigEntries; import static org.apache.fluss.server.utils.ServerRpcMessageUtils.toPbDatabaseSummary; import static org.apache.fluss.server.utils.ServerRpcMessageUtils.toTablePath; +import static org.apache.fluss.server.zk.data.LeaderAndIsr.NO_LEADER; import static org.apache.fluss.utils.PartitionUtils.HISTORICAL_PARTITION_VALUE; import static org.apache.fluss.utils.Preconditions.checkState; @@ -421,7 +422,8 @@ private static void addBucketInfos( response, partitionId, partitionName, bucketMetadata)); } - private static void addBucketInfo( + @VisibleForTesting + static void addBucketInfo( DescribeBucketsResponse response, @Nullable Long partitionId, @Nullable String partitionName, @@ -434,10 +436,16 @@ private static void addBucketInfo( if (partitionName != null) { pbBucketInfo.setPartitionName(partitionName); } - bucketMetadata.getLeaderId().ifPresent(pbBucketInfo::setLeaderId); - bucketMetadata.getLeaderEpoch().ifPresent(pbBucketInfo::setLeaderEpoch); + if (bucketMetadata.getLeaderId().isPresent() + && bucketMetadata.getLeaderId().getAsInt() != NO_LEADER) { + pbBucketInfo.setLeaderId(bucketMetadata.getLeaderId().getAsInt()); + bucketMetadata.getLeaderEpoch().ifPresent(pbBucketInfo::setLeaderEpoch); + } + if (bucketMetadata.getBucketEpoch() != null) { + pbBucketInfo.setBucketEpoch(bucketMetadata.getBucketEpoch()); + } bucketMetadata.getReplicas().forEach(pbBucketInfo::addReplicaId); - bucketMetadata.getIsr().forEach(pbBucketInfo::addIsrId); + bucketMetadata.getIsr().forEach(pbBucketInfo::addIsr); } @Override diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java index cdb8c90c2ea..ced25f3bff9 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java @@ -446,11 +446,10 @@ public Map getTablesAssignments(Collection tableIds tableIds.stream().collect(toMap(TableIdZNode::path, id -> id)); List responses = getDataInBackground(path2TableIdMap.keySet()); - return processGetDataResponses( + return processGetDataResponsesOrThrow( responses, response -> path2TableIdMap.get(response.getPath()), - TableIdZNode::decode, - "table assignment"); + data -> data == null || data.length == 0 ? null : TableIdZNode.decode(data)); } /** Get the partition assignment in ZK. */ @@ -466,11 +465,10 @@ public Map getPartitionsAssignments(Collection partitionIds.stream().collect(toMap(PartitionIdZNode::path, id -> id)); List responses = getDataInBackground(path2PartitionIdMap.keySet()); - return processGetDataResponses( + return processGetDataResponsesOrThrow( responses, response -> path2PartitionIdMap.get(response.getPath()), - PartitionIdZNode::decode, - "partition assignment"); + PartitionIdZNode::decode); } public void updateTableAssignment( @@ -597,11 +595,10 @@ public Map getLeaderAndIsrs(Collection t tableBuckets.stream().collect(toMap(LeaderAndIsrZNode::path, bucket -> bucket)); List responses = getDataInBackground(path2TableBucketMap.keySet()); - return processGetDataResponses( + return processGetDataResponsesOrThrow( responses, response -> path2TableBucketMap.get(response.getPath()), - LeaderAndIsrZNode::decode, - "leader and isr"); + LeaderAndIsrZNode::decode); } public void updateLeaderAndIsr( @@ -1008,25 +1005,12 @@ private Map getPartitionRegistrations( partitionName -> PartitionZNode.path(tablePath, partitionName), partitionName -> partitionName)); - List responses = getDataInBackground(path2PartitionName.keySet()); - Map partitionRegistrations = new HashMap<>(); - for (ZkGetDataResponse response : responses) { - if (response.getResultCode() == KeeperException.Code.NONODE) { - continue; - } - if (response.getResultCode() != KeeperException.Code.OK) { - throw KeeperException.create(response.getResultCode(), response.getPath()); - } - - PartitionRegistration partitionRegistration = PartitionZNode.decode(response.getData()); - if (partitionRegistration.getRemoteDataDir() == null) { - partitionRegistration = - partitionRegistration.newRemoteDataDir(defaultRemoteDataDir); - } - partitionRegistrations.put( - path2PartitionName.get(response.getPath()), partitionRegistration); - } - return partitionRegistrations; + return getPartitionZNodeData( + path2PartitionName, + partitionRegistration -> + partitionRegistration.getRemoteDataDir() == null + ? partitionRegistration.newRemoteDataDir(defaultRemoteDataDir) + : partitionRegistration); } /** Get the id and name for the partitions of a table in ZK. */ @@ -1218,12 +1202,18 @@ public Map getPartitionIds( checkNotNull(p.getPartitionName())), path -> path)); - List responses = getDataInBackground(path2PartitionPathMap.keySet()); - return processGetDataResponses( + return getPartitionZNodeData( + path2PartitionPathMap, PartitionRegistration::toTablePartition); + } + + private Map getPartitionZNodeData( + Map path2Key, Function partitionRegistrationMapper) + throws Exception { + List responses = getDataInBackground(path2Key.keySet()); + return processGetDataResponsesOrThrow( responses, - response -> path2PartitionPathMap.get(response.getPath()), - (byte[] data) -> PartitionZNode.decode(data).toTablePartition(), - "partition"); + response -> path2Key.get(response.getPath()), + data -> partitionRegistrationMapper.apply(PartitionZNode.decode(data))); } /** Get partition num of a table in ZK. */ @@ -2211,6 +2201,26 @@ public static Map processGetDataResponses( return result; } + private static Map processGetDataResponsesOrThrow( + List responses, + Function keyExtractor, + Function decoder) + throws KeeperException { + Map result = new HashMap<>(); + for (ZkGetDataResponse response : responses) { + if (response.getResultCode() == KeeperException.Code.NONODE) { + continue; + } + response.maybeThrow(); + + V value = decoder.apply(response.getData()); + if (value != null) { + result.put(keyExtractor.apply(response), value); + } + } + return result; + } + /** * Template method to process multiple ZooKeeper children responses with decoder. * diff --git a/fluss-server/src/test/java/org/apache/fluss/server/RpcServiceBaseTest.java b/fluss-server/src/test/java/org/apache/fluss/server/RpcServiceBaseTest.java index 645853776e7..d2180e9deb6 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/RpcServiceBaseTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/RpcServiceBaseTest.java @@ -24,16 +24,22 @@ import org.apache.fluss.metadata.TablePath; import org.apache.fluss.record.TestData; import org.apache.fluss.row.encode.KvValueLayout; +import org.apache.fluss.rpc.messages.DescribeBucketsResponse; +import org.apache.fluss.rpc.messages.PbBucketInfo; +import org.apache.fluss.server.metadata.BucketMetadata; import org.junit.jupiter.api.Test; +import java.util.Arrays; import java.util.HashMap; import java.util.Map; import static org.apache.fluss.server.RpcServiceBase.validateKvSnapshotMetadataVersion; +import static org.apache.fluss.server.zk.data.LeaderAndIsr.NO_LEADER; +import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; -/** Tests KV snapshot metadata API compatibility. */ +/** Tests for {@link RpcServiceBase}. */ class RpcServiceBaseTest { @Test @@ -47,6 +53,23 @@ void testValidateKvSnapshotMetadataVersion() { validateKvSnapshotMetadataVersion((short) 0, plainTableInfo); } + @Test + void testAddBucketInfoNormalizesNoLeader() { + DescribeBucketsResponse response = new DescribeBucketsResponse(); + BucketMetadata bucketMetadata = + new BucketMetadata(0, NO_LEADER, 3, Arrays.asList(1, 2, 3), Arrays.asList(1, 2), 4); + + RpcServiceBase.addBucketInfo(response, null, null, bucketMetadata); + + PbBucketInfo bucketInfo = response.getBucketInfosList().get(0); + assertThat(bucketInfo.hasLeaderId()).isFalse(); + assertThat(bucketInfo.hasLeaderEpoch()).isFalse(); + assertThat(bucketInfo.hasBucketEpoch()).isTrue(); + assertThat(bucketInfo.getBucketEpoch()).isEqualTo(4); + assertThat(bucketInfo.getReplicaIds()).containsExactly(1, 2, 3); + assertThat(bucketInfo.getIsrs()).containsExactly(1, 2); + } + private static TableInfo tableInfo(KvValueLayout layout) { Map properties = new HashMap<>(TestData.DATA1_TABLE_DESCRIPTOR_PK.getProperties()); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java index c23e11238c3..ab547328c74 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java @@ -25,10 +25,12 @@ import org.apache.fluss.config.Configuration; import org.apache.fluss.fs.FsPath; import org.apache.fluss.metadata.DatabaseSummary; +import org.apache.fluss.metadata.PhysicalTablePath; import org.apache.fluss.metadata.Schema; import org.apache.fluss.metadata.SchemaInfo; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TablePartition; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.server.entity.RegisterTableBucketLeadAndIsrInfo; import org.apache.fluss.server.zk.ZkAsyncResponse.ZkGetDataResponse; @@ -44,7 +46,10 @@ import org.apache.fluss.server.zk.data.TableRegistration; import org.apache.fluss.server.zk.data.TabletServerRegistration; import org.apache.fluss.server.zk.data.ZkData.BucketIdZNode; +import org.apache.fluss.server.zk.data.ZkData.LeaderAndIsrZNode; +import org.apache.fluss.server.zk.data.ZkData.PartitionIdZNode; import org.apache.fluss.server.zk.data.ZkData.PartitionZNode; +import org.apache.fluss.server.zk.data.ZkData.TableIdZNode; import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadata; import org.apache.fluss.shaded.curator5.org.apache.curator.CuratorZookeeperClient; import org.apache.fluss.shaded.curator5.org.apache.curator.framework.CuratorFramework; @@ -203,6 +208,62 @@ void testTabletAssignments() throws Exception { assertThat(zookeeperClient.getTableAssignment(tableId1)).isEmpty(); } + @Test + void testGetTablesAssignmentsPreservesZookeeperErrors() throws Exception { + long tableId = 1L; + String tablePath = TableIdZNode.path(tableId); + ZooKeeperClient testingClient = spy(zookeeperClient); + doReturn( + Arrays.asList( + new ZkGetDataResponse( + tablePath, KeeperException.Code.OK, new byte[0]), + new ZkGetDataResponse( + TableIdZNode.path(2L), KeeperException.Code.NONODE, null))) + .when(testingClient) + .getDataInBackground(anyCollection()); + + assertThat(testingClient.getTablesAssignments(Arrays.asList(tableId, 2L))).isEmpty(); + + doReturn( + Collections.singletonList( + new ZkGetDataResponse( + tablePath, KeeperException.Code.CONNECTIONLOSS, null))) + .when(testingClient) + .getDataInBackground(anyCollection()); + + assertThatThrownBy(() -> testingClient.getTablesAssignments(Collections.singleton(tableId))) + .isInstanceOf(KeeperException.ConnectionLossException.class); + } + + @Test + void testGetPartitionsAssignmentsPreservesZookeeperErrors() throws Exception { + long partitionId = 1L; + String partitionPath = PartitionIdZNode.path(partitionId); + ZooKeeperClient testingClient = spy(zookeeperClient); + doReturn( + Collections.singletonList( + new ZkGetDataResponse( + partitionPath, KeeperException.Code.NONODE, null))) + .when(testingClient) + .getDataInBackground(anyCollection()); + + assertThat(testingClient.getPartitionsAssignments(Collections.singleton(partitionId))) + .isEmpty(); + + doReturn( + Collections.singletonList( + new ZkGetDataResponse( + partitionPath, KeeperException.Code.CONNECTIONLOSS, null))) + .when(testingClient) + .getDataInBackground(anyCollection()); + + assertThatThrownBy( + () -> + testingClient.getPartitionsAssignments( + Collections.singleton(partitionId))) + .isInstanceOf(KeeperException.ConnectionLossException.class); + } + @Test void testLeaderAndIsr() throws Exception { // try to get bucket leadership, should return empty @@ -240,6 +301,33 @@ void testLeaderAndIsr() throws Exception { assertThat(zookeeperClient.getLeaderAndIsr(tableBucket1)).isEmpty(); } + @Test + void testGetLeaderAndIsrsPreservesZookeeperErrors() throws Exception { + TableBucket tableBucket = new TableBucket(1L, 0); + String leaderAndIsrPath = LeaderAndIsrZNode.path(tableBucket); + ZooKeeperClient testingClient = spy(zookeeperClient); + doReturn( + Collections.singletonList( + new ZkGetDataResponse( + leaderAndIsrPath, KeeperException.Code.NONODE, null))) + .when(testingClient) + .getDataInBackground(anyCollection()); + + assertThat(testingClient.getLeaderAndIsrs(Collections.singleton(tableBucket))).isEmpty(); + + doReturn( + Collections.singletonList( + new ZkGetDataResponse( + leaderAndIsrPath, + KeeperException.Code.CONNECTIONLOSS, + null))) + .when(testingClient) + .getDataInBackground(anyCollection()); + + assertThatThrownBy(() -> testingClient.getLeaderAndIsrs(Collections.singleton(tableBucket))) + .isInstanceOf(KeeperException.ConnectionLossException.class); + } + @ParameterizedTest @ValueSource(booleans = {true, false}) void testBatchCreateAndUpdateLeaderAndIsr(boolean isPartitionTable) throws Exception { @@ -825,6 +913,12 @@ void testGetPartitionRegistrationsPreservesZookeeperErrors() throws Exception { assertThat(testingClient.getPartitionRegistrations(tablePath)) .containsOnlyKeys("p1") .containsValue(partitionRegistration); + PhysicalTablePath partitionPath = PhysicalTablePath.of(tablePath, "p1"); + Map partitionIds = + testingClient.getPartitionIds(Collections.singleton(partitionPath)); + assertThat(partitionIds).containsOnlyKeys(partitionPath); + assertThat(partitionIds.get(partitionPath).getTableId()).isEqualTo(1L); + assertThat(partitionIds.get(partitionPath).getPartitionId()).isEqualTo(2L); doReturn( Collections.singletonList( @@ -835,6 +929,9 @@ void testGetPartitionRegistrationsPreservesZookeeperErrors() throws Exception { assertThatThrownBy(() -> testingClient.getPartitionRegistrations(tablePath)) .isInstanceOf(KeeperException.ConnectionLossException.class); + assertThatThrownBy( + () -> testingClient.getPartitionIds(Collections.singleton(partitionPath))) + .isInstanceOf(KeeperException.ConnectionLossException.class); } @Test From 6656bb7a1bd6cf7b65c5d6effa34385ad998659e Mon Sep 17 00:00:00 2001 From: fhan Date: Tue, 8 Sep 2026 17:08:52 +0800 Subject: [PATCH 5/5] [test] Avoid Mockito in ZooKeeper response handling tests --- .../fluss/server/zk/ZooKeeperClient.java | 3 +- .../fluss/server/zk/ZooKeeperClientTest.java | 200 ++++++------------ 2 files changed, 69 insertions(+), 134 deletions(-) diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java index ced25f3bff9..f2906ceb245 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java @@ -2201,7 +2201,8 @@ public static Map processGetDataResponses( return result; } - private static Map processGetDataResponsesOrThrow( + @VisibleForTesting + static Map processGetDataResponsesOrThrow( List responses, Function keyExtractor, Function decoder) diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java index ab547328c74..65c36451e63 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java @@ -46,9 +46,6 @@ import org.apache.fluss.server.zk.data.TableRegistration; import org.apache.fluss.server.zk.data.TabletServerRegistration; import org.apache.fluss.server.zk.data.ZkData.BucketIdZNode; -import org.apache.fluss.server.zk.data.ZkData.LeaderAndIsrZNode; -import org.apache.fluss.server.zk.data.ZkData.PartitionIdZNode; -import org.apache.fluss.server.zk.data.ZkData.PartitionZNode; import org.apache.fluss.server.zk.data.ZkData.TableIdZNode; import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadata; import org.apache.fluss.shaded.curator5.org.apache.curator.CuratorZookeeperClient; @@ -73,7 +70,6 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; -import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; @@ -88,7 +84,6 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.anyCollection; import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; @@ -196,6 +191,19 @@ void testTabletAssignments() throws Exception { assertThat(zookeeperClient.getTablesAssignments(Arrays.asList(tableId1, tableId2))) .containsValues(tableAssignment1, tableAssignment2); + // An existing table id node may have no assignment when it only contains child nodes. + long tableIdWithoutAssignment = 3L; + zookeeperClient + .getCuratorClient() + .create() + .forPath(TableIdZNode.path(tableIdWithoutAssignment), new byte[0]); + assertThat( + zookeeperClient.getTablesAssignments( + Arrays.asList(tableId1, tableId2, tableIdWithoutAssignment))) + .containsOnlyKeys(tableId1, tableId2) + .containsEntry(tableId1, tableAssignment1) + .containsEntry(tableId2, tableAssignment2); + // test update TableAssignment tableAssignment3 = TableAssignment.builder().add(3, BucketAssignment.of(1, 5)).build(); @@ -206,62 +214,11 @@ void testTabletAssignments() throws Exception { // test delete zookeeperClient.deleteTableAssignment(tableId1); assertThat(zookeeperClient.getTableAssignment(tableId1)).isEmpty(); - } - - @Test - void testGetTablesAssignmentsPreservesZookeeperErrors() throws Exception { - long tableId = 1L; - String tablePath = TableIdZNode.path(tableId); - ZooKeeperClient testingClient = spy(zookeeperClient); - doReturn( - Arrays.asList( - new ZkGetDataResponse( - tablePath, KeeperException.Code.OK, new byte[0]), - new ZkGetDataResponse( - TableIdZNode.path(2L), KeeperException.Code.NONODE, null))) - .when(testingClient) - .getDataInBackground(anyCollection()); - - assertThat(testingClient.getTablesAssignments(Arrays.asList(tableId, 2L))).isEmpty(); - - doReturn( - Collections.singletonList( - new ZkGetDataResponse( - tablePath, KeeperException.Code.CONNECTIONLOSS, null))) - .when(testingClient) - .getDataInBackground(anyCollection()); - - assertThatThrownBy(() -> testingClient.getTablesAssignments(Collections.singleton(tableId))) - .isInstanceOf(KeeperException.ConnectionLossException.class); - } - - @Test - void testGetPartitionsAssignmentsPreservesZookeeperErrors() throws Exception { - long partitionId = 1L; - String partitionPath = PartitionIdZNode.path(partitionId); - ZooKeeperClient testingClient = spy(zookeeperClient); - doReturn( - Collections.singletonList( - new ZkGetDataResponse( - partitionPath, KeeperException.Code.NONODE, null))) - .when(testingClient) - .getDataInBackground(anyCollection()); - - assertThat(testingClient.getPartitionsAssignments(Collections.singleton(partitionId))) - .isEmpty(); - - doReturn( - Collections.singletonList( - new ZkGetDataResponse( - partitionPath, KeeperException.Code.CONNECTIONLOSS, null))) - .when(testingClient) - .getDataInBackground(anyCollection()); - - assertThatThrownBy( - () -> - testingClient.getPartitionsAssignments( - Collections.singleton(partitionId))) - .isInstanceOf(KeeperException.ConnectionLossException.class); + assertThat( + zookeeperClient.getTablesAssignments( + Arrays.asList(tableId1, tableId2, tableIdWithoutAssignment))) + .containsOnlyKeys(tableId2) + .containsEntry(tableId2, tableAssignment2); } @Test @@ -285,8 +242,13 @@ void testLeaderAndIsr() throws Exception { tableBucket2, leaderAndIsr2, zkEpoch.getCoordinatorEpochZkVersion()); assertThat(zookeeperClient.getLeaderAndIsr(tableBucket1)).hasValue(leaderAndIsr1); assertThat(zookeeperClient.getLeaderAndIsr(tableBucket2)).hasValue(leaderAndIsr2); - assertThat(zookeeperClient.getLeaderAndIsrs(Arrays.asList(tableBucket1, tableBucket2))) - .containsValues(leaderAndIsr1, leaderAndIsr2); + Map leaderAndIsrs = + zookeeperClient.getLeaderAndIsrs( + Arrays.asList(tableBucket1, tableBucket2, new TableBucket(1, 3))); + assertThat(leaderAndIsrs) + .containsOnlyKeys(tableBucket1, tableBucket2) + .containsEntry(tableBucket1, leaderAndIsr1) + .containsEntry(tableBucket2, leaderAndIsr2); // test update leaderAndIsr1 = @@ -302,29 +264,31 @@ void testLeaderAndIsr() throws Exception { } @Test - void testGetLeaderAndIsrsPreservesZookeeperErrors() throws Exception { - TableBucket tableBucket = new TableBucket(1L, 0); - String leaderAndIsrPath = LeaderAndIsrZNode.path(tableBucket); - ZooKeeperClient testingClient = spy(zookeeperClient); - doReturn( - Collections.singletonList( - new ZkGetDataResponse( - leaderAndIsrPath, KeeperException.Code.NONODE, null))) - .when(testingClient) - .getDataInBackground(anyCollection()); - - assertThat(testingClient.getLeaderAndIsrs(Collections.singleton(tableBucket))).isEmpty(); - - doReturn( - Collections.singletonList( - new ZkGetDataResponse( - leaderAndIsrPath, - KeeperException.Code.CONNECTIONLOSS, - null))) - .when(testingClient) - .getDataInBackground(anyCollection()); - - assertThatThrownBy(() -> testingClient.getLeaderAndIsrs(Collections.singleton(tableBucket))) + void testProcessGetDataResponsesOrThrow() throws Exception { + String existingPath = "/existing"; + List responses = + Arrays.asList( + new ZkGetDataResponse( + existingPath, KeeperException.Code.OK, new byte[] {1}, null), + new ZkGetDataResponse( + "/without-value", KeeperException.Code.OK, new byte[0], null), + new ZkGetDataResponse("/missing", KeeperException.Code.NONODE, null, null)); + + Map result = + ZooKeeperClient.processGetDataResponsesOrThrow( + responses, + ZkGetDataResponse::getPath, + data -> data.length == 0 ? null : (int) data[0]); + assertThat(result).containsOnlyKeys(existingPath).containsEntry(existingPath, 1); + + ZkGetDataResponse failedResponse = + new ZkGetDataResponse("/failed", KeeperException.Code.CONNECTIONLOSS, null, null); + assertThatThrownBy( + () -> + ZooKeeperClient.processGetDataResponsesOrThrow( + Collections.singletonList(failedResponse), + ZkGetDataResponse::getPath, + data -> data)) .isInstanceOf(KeeperException.ConnectionLossException.class); } @@ -844,6 +808,11 @@ void testPartition() throws Exception { tableId, partitionAssignment.getBucketAssignments().size()); + assertThat(zookeeperClient.getPartitionsAssignments(Arrays.asList(1L, 2L, 3L))) + .containsOnlyKeys(1L, 2L) + .containsEntry(1L, partitionAssignment) + .containsEntry(2L, partitionAssignment); + // check created partitions partitions = zookeeperClient.getPartitions(tablePath); assertThat(partitions).containsExactly("p1", "p2"); @@ -859,6 +828,18 @@ void testPartition() throws Exception { assertThat(zookeeperClient.getPartitionsForTables(Arrays.asList(tablePath))) .containsValues(new ArrayList<>(partitions)); + PhysicalTablePath partitionPath1 = PhysicalTablePath.of(tablePath, "p1"); + PhysicalTablePath partitionPath2 = PhysicalTablePath.of(tablePath, "p2"); + PhysicalTablePath missingPartitionPath = PhysicalTablePath.of(tablePath, "p3"); + Map partitionIds = + zookeeperClient.getPartitionIds( + Arrays.asList(partitionPath1, partitionPath2, missingPartitionPath)); + assertThat(partitionIds).containsOnlyKeys(partitionPath1, partitionPath2); + assertThat(partitionIds.get(partitionPath1).getTableId()).isEqualTo(tableId); + assertThat(partitionIds.get(partitionPath1).getPartitionId()).isEqualTo(1L); + assertThat(partitionIds.get(partitionPath2).getTableId()).isEqualTo(tableId); + assertThat(partitionIds.get(partitionPath2).getPartitionId()).isEqualTo(2L); + // A batch read returns every registration and preserves the version needed by CAS updates. PartitionRegistration p1Registration = zookeeperClient.getPartition(tablePath, "p1").get(); zookeeperClient.updatePartitionRegistration(tablePath, "p1", p1Registration); @@ -887,53 +868,6 @@ void testPartition() throws Exception { assertThat(partitions).containsExactly("p2"); } - @Test - void testGetPartitionRegistrationsPreservesZookeeperErrors() throws Exception { - TablePath tablePath = TablePath.of("db", "tb"); - String partition1Path = PartitionZNode.path(tablePath, "p1"); - String partition2Path = PartitionZNode.path(tablePath, "p2"); - PartitionRegistration partitionRegistration = - new PartitionRegistration(1L, 2L, remoteDataDir); - - ZooKeeperClient testingClient = spy(zookeeperClient); - doReturn(new HashSet<>(Arrays.asList("p1", "p2"))) - .when(testingClient) - .getPartitions(tablePath); - doReturn( - Arrays.asList( - new ZkGetDataResponse( - partition1Path, - KeeperException.Code.OK, - PartitionZNode.encode(partitionRegistration)), - new ZkGetDataResponse( - partition2Path, KeeperException.Code.NONODE, null))) - .when(testingClient) - .getDataInBackground(anyCollection()); - - assertThat(testingClient.getPartitionRegistrations(tablePath)) - .containsOnlyKeys("p1") - .containsValue(partitionRegistration); - PhysicalTablePath partitionPath = PhysicalTablePath.of(tablePath, "p1"); - Map partitionIds = - testingClient.getPartitionIds(Collections.singleton(partitionPath)); - assertThat(partitionIds).containsOnlyKeys(partitionPath); - assertThat(partitionIds.get(partitionPath).getTableId()).isEqualTo(1L); - assertThat(partitionIds.get(partitionPath).getPartitionId()).isEqualTo(2L); - - doReturn( - Collections.singletonList( - new ZkGetDataResponse( - partition1Path, KeeperException.Code.CONNECTIONLOSS, null))) - .when(testingClient) - .getDataInBackground(anyCollection()); - - assertThatThrownBy(() -> testingClient.getPartitionRegistrations(tablePath)) - .isInstanceOf(KeeperException.ConnectionLossException.class); - assertThatThrownBy( - () -> testingClient.getPartitionIds(Collections.singleton(partitionPath))) - .isInstanceOf(KeeperException.ConnectionLossException.class); - } - @Test void testServerTag() throws Exception { Map serverTags = new HashMap<>();