[client] Add Admin API to describe buckets

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<Void> createTable(
      */
     CompletableFuture<TableInfo> getTableInfo(TablePath tablePath);
 
+    /**
+     * Describes the buckets of the given table asynchronously.
+     *
+     * <p>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.
+     *
+     * <p>The following exceptions can be anticipated when calling {@code get()} on the returned
+     * future.
+     *
+     * <ul>
+     *   <li>{@link TableNotExistException} if the table does not exist.
+     * </ul>
+     *
+     * @param tablePath The table path of the table.
+     * @since 1.0
+     */
+    CompletableFuture<List<BucketInfo>> describeBuckets(TablePath tablePath);
+
+    /**
+     * Describes the buckets matching the given partition spec asynchronously.
+     *
+     * <p>The partition spec may contain all partition keys or a subset of them.
+     *
+     * <p>The following exceptions can be anticipated when calling {@code get()} on the returned
+     * future.
+     *
+     * <ul>
+     *   <li>{@link TableNotExistException} if the table does not exist.
+     *   <li>{@link TableNotPartitionedException} if the table is not partitioned.
+     *   <li>{@link InvalidPartitionException} if the partition spec is invalid.
+     * </ul>
+     *
+     * @param tablePath The table path of the table.
+     * @param partitionSpec The complete or partial partition spec.
+     * @since 1.0
+     */
+    CompletableFuture<List<BucketInfo>> describeBuckets(
+            TablePath tablePath, PartitionSpec partitionSpec);
+
     /**
      * Drop the table with the given table path asynchronously.
      *

fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java

@@ -34,6 +34,7 @@
 import org.apache.fluss.config.cluster.ConfigEntry;
 import org.apache.fluss.exception.FlussRuntimeException;
 import org.apache.fluss.exception.LeaderNotAvailableException;
+import org.apache.fluss.metadata.BucketInfo;
 import org.apache.fluss.metadata.DatabaseChange;
 import org.apache.fluss.metadata.DatabaseDescriptor;
 import org.apache.fluss.metadata.DatabaseInfo;
@@ -66,6 +67,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;
@@ -347,6 +349,33 @@ public CompletableFuture<TableInfo> getTableInfo(TablePath tablePath) {
                                         r.getModifiedTime()));
     }
 
+    @Override
+    public CompletableFuture<List<BucketInfo>> 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<List<BucketInfo>> 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<Void> dropTable(TablePath tablePath, boolean ignoreIfNotExists) {
         DropTableRequest request = new DropTableRequest();

fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java

@@ -42,6 +42,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;
@@ -55,6 +56,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;
@@ -73,6 +75,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;
@@ -649,6 +652,27 @@ public static List<PartitionInfo> toPartitionInfos(ListPartitionInfosResponse re
                 .collect(Collectors.toList());
     }
 
+    public static List<BucketInfo> 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<String, String> toKeyValueMap(List<PbKeyValue> pbKeyValues) {
         return pbKeyValues.stream()
                 .collect(

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<BucketInfo> 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<String, Long> partitionIds =
+                admin.listPartitionInfos(PARTITIONED_TABLE_PATH).get().stream()
+                        .collect(
+                                Collectors.toMap(
+                                        PartitionInfo::getPartitionName,
+                                        PartitionInfo::getPartitionId));
+
+        List<BucketInfo> 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<BucketInfo> partialPartitionBuckets =
+                waitAndDescribeBuckets(PARTITIONED_TABLE_PATH, newPartitionSpec("pt", "2025"), 4);
+        assertThat(partialPartitionBuckets)
+                .extracting(BucketInfo::getPartitionName)
+                .containsExactly("2025$cn", "2025$cn", "2025$us", "2025$us");
+
+        List<BucketInfo> 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<BucketInfo> waitAndDescribeBuckets(
+            TablePath tablePath, @Nullable PartitionSpec partitionSpec, int expectedBucketCount)
+            throws Exception {
+        waitUntil(
+                () -> {
+                    List<BucketInfo> 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<BucketInfo> 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());
+    }
+}

fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java

@@ -37,7 +37,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;
@@ -1270,7 +1269,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");
     }