[flink] Add describe buckets procedure

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/DescribeBucketsProcedure.java

@@ -0,0 +1,167 @@
+/*
+ * 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.flink.procedure;
+
+import org.apache.fluss.metadata.BucketInfo;
+import org.apache.fluss.metadata.PartitionSpec;
+import org.apache.fluss.metadata.ResolvedPartitionSpec;
+import org.apache.fluss.metadata.TablePath;
+
+import org.apache.flink.table.annotation.ArgumentHint;
+import org.apache.flink.table.annotation.DataTypeHint;
+import org.apache.flink.table.annotation.ProcedureHint;
+import org.apache.flink.table.procedure.ProcedureContext;
+import org.apache.flink.types.Row;
+
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.OptionalInt;
+import java.util.OptionalLong;
+
+/**
+ * Procedure to describe bucket metadata for a table.
+ *
+ * <p>Usage examples:
+ *
+ * <pre>
+ * CALL sys.describe_buckets('db.table');
+ * CALL sys.describe_buckets('db.table', 'partition_key=partition_value');
+ * </pre>
+ */
+public class DescribeBucketsProcedure extends ProcedureBase {
+
+    private static final String OUTPUT_TYPE =
+            "ROW<table_path STRING, table_id BIGINT, partition_id BIGINT, "
+                    + "partition_name STRING, bucket_id INT, leader_id INT, leader_epoch INT, "
+                    + "bucket_epoch INT, replicas ARRAY<INT>, isr ARRAY<INT>>";
+
+    @ProcedureHint(
+            argument = {@ArgumentHint(name = "table_path", type = @DataTypeHint("STRING"))},
+            output = @DataTypeHint(OUTPUT_TYPE))
+    public Row[] call(ProcedureContext context, String tablePath) throws Exception {
+        TablePath parsedTablePath = parseTablePath(tablePath);
+        return toRows(admin.describeBuckets(parsedTablePath).get());
+    }
+
+    @ProcedureHint(
+            argument = {
+                @ArgumentHint(name = "table_path", type = @DataTypeHint("STRING")),
+                @ArgumentHint(name = "partition_spec", type = @DataTypeHint("STRING"))
+            },
+            output = @DataTypeHint(OUTPUT_TYPE))
+    public Row[] call(ProcedureContext context, String tablePath, String partitionSpec)
+            throws Exception {
+        TablePath parsedTablePath = parseTablePath(tablePath);
+        return toRows(
+                admin.describeBuckets(parsedTablePath, parsePartitionSpec(partitionSpec)).get());
+    }
+
+    private static Row[] toRows(List<BucketInfo> bucketInfos) {
+        return bucketInfos.stream().map(DescribeBucketsProcedure::toRow).toArray(Row[]::new);
+    }
+
+    private static Row toRow(BucketInfo bucketInfo) {
+        return Row.of(
+                bucketInfo.getTablePath().toString(),
+                bucketInfo.getTableId(),
+                optionalLong(bucketInfo.getPartitionId()),
+                bucketInfo.getPartitionName(),
+                bucketInfo.getBucketId(),
+                optionalInt(bucketInfo.getLeaderId()),
+                optionalInt(bucketInfo.getLeaderEpoch()),
+                optionalInt(bucketInfo.getBucketEpoch()),
+                bucketInfo.getReplicas().toArray(new Integer[0]),
+                bucketInfo.getIsr().toArray(new Integer[0]));
+    }
+
+    private static Long optionalLong(OptionalLong value) {
+        return value.isPresent() ? value.getAsLong() : null;
+    }
+
+    private static Integer optionalInt(OptionalInt value) {
+        return value.isPresent() ? value.getAsInt() : null;
+    }
+
+    private static TablePath parseTablePath(String tablePath) {
+        if (tablePath == null || tablePath.trim().isEmpty()) {
+            throw new IllegalArgumentException(
+                    "table_path cannot be null or empty. Expected format is 'database.table'.");
+        }
+
+        String normalizedTablePath = tablePath.trim();
+        String[] parts = normalizedTablePath.split("\\.", -1);
+        if (parts.length != 2 || parts[0].isEmpty() || parts[1].isEmpty()) {
+            throw new IllegalArgumentException(
+                    "Invalid table_path '" + tablePath + "'. Expected format is 'database.table'.");
+        }
+        TablePath parsedTablePath = TablePath.of(parts[0], parts[1]);
+        parsedTablePath.validate();
+        return parsedTablePath;
+    }
+
+    private static PartitionSpec parsePartitionSpec(String partitionSpec) {
+        if (partitionSpec == null || partitionSpec.trim().isEmpty()) {
+            throw new IllegalArgumentException(
+                    "partition_spec cannot be null or empty. Expected format is "
+                            + "'key=value[/key=value...]'.");
+        }
+
+        String normalizedPartitionSpec = partitionSpec.trim();
+        ResolvedPartitionSpec resolvedPartitionSpec;
+        try {
+            String[] keyValuePairs = normalizedPartitionSpec.split("/", -1);
+            for (String keyValuePair : keyValuePairs) {
+                if (keyValuePair.isEmpty()) {
+                    throw new IllegalArgumentException("Empty partition key-value pair.");
+                }
+            }
+            resolvedPartitionSpec =
+                    ResolvedPartitionSpec.fromPartitionQualifiedName(normalizedPartitionSpec);
+        } catch (IllegalArgumentException e) {
+            throw new IllegalArgumentException(
+                    "Invalid partition_spec '"
+                            + partitionSpec
+                            + "'. Expected format is 'key=value[/key=value...]'.",
+                    e);
+        }
+
+        Map<String, String> spec = new LinkedHashMap<>();
+        List<String> partitionKeys = resolvedPartitionSpec.getPartitionKeys();
+        List<String> partitionValues = resolvedPartitionSpec.getPartitionValues();
+        for (int i = 0; i < partitionKeys.size(); i++) {
+            String partitionKey = partitionKeys.get(i);
+            if (partitionKey.trim().isEmpty()) {
+                throw new IllegalArgumentException(
+                        "Invalid partition_spec '"
+                                + partitionSpec
+                                + "': partition key cannot be empty.");
+            }
+            if (spec.containsKey(partitionKey)) {
+                throw new IllegalArgumentException(
+                        "Duplicate partition key '"
+                                + partitionKey
+                                + "' in partition_spec '"
+                                + partitionSpec
+                                + "'.");
+            }
+            spec.put(partitionKey, partitionValues.get(i));
+        }
+        return new PartitionSpec(spec);
+    }
+}

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/ProcedureManager.java

@@ -68,6 +68,7 @@ private static Map<String, Class<? extends ProcedureBase>> initProcedureMap() {
 
     private enum ProcedureEnum {
         ADD_ACL("sys.add_acl", AddAclProcedure.class),
+        DESCRIBE_BUCKETS("sys.describe_buckets", DescribeBucketsProcedure.class),
         DROP_ACL("sys.drop_acl", DropAclProcedure.class),
         List_ACL("sys.list_acl", ListAclProcedure.class),
         SET_CLUSTER_CONFIGS("sys.set_cluster_configs", SetClusterConfigsProcedure.class),

fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/procedure/DescribeBucketsProcedureTest.java

@@ -0,0 +1,183 @@
+/*
+ * 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.flink.procedure;
+
+import org.apache.fluss.flink.sink.testutils.TestAdminAdapter;
+import org.apache.fluss.metadata.BucketInfo;
+import org.apache.fluss.metadata.PartitionSpec;
+import org.apache.fluss.metadata.TablePath;
+
+import org.apache.flink.types.Row;
+import org.junit.jupiter.api.Test;
+
+import javax.annotation.Nullable;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Test for {@link DescribeBucketsProcedure}. */
+class DescribeBucketsProcedureTest {
+
+    private static final TablePath TABLE_PATH = TablePath.of("test_db", "test_table");
+
+    @Test
+    void testDescribeAllBucketsAndConvertNullableFields() throws Exception {
+        BucketInfo bucketInfo =
+                new BucketInfo(
+                        TABLE_PATH,
+                        10L,
+                        null,
+                        null,
+                        0,
+                        null,
+                        null,
+                        null,
+                        Arrays.asList(1, 2, 3),
+                        Collections.emptyList());
+        TestingAdmin admin = new TestingAdmin(Collections.singletonList(bucketInfo));
+        DescribeBucketsProcedure procedure = newProcedure(admin);
+
+        Row[] rows = procedure.call(null, " test_db.test_table ");
+
+        assertThat(admin.requestedTablePath).isEqualTo(TABLE_PATH);
+        assertThat(admin.requestedPartitionSpec).isNull();
+        assertThat(rows).hasSize(1);
+        Row row = rows[0];
+        assertThat(row.getArity()).isEqualTo(10);
+        assertThat(row.getField(0)).isEqualTo("test_db.test_table");
+        assertThat(row.getField(1)).isEqualTo(10L);
+        assertThat(row.getField(2)).isNull();
+        assertThat(row.getField(3)).isNull();
+        assertThat(row.getField(4)).isEqualTo(0);
+        assertThat(row.getField(5)).isNull();
+        assertThat(row.getField(6)).isNull();
+        assertThat(row.getField(7)).isNull();
+        assertThat((Integer[]) row.getField(8)).containsExactly(1, 2, 3);
+        assertThat((Integer[]) row.getField(9)).isEmpty();
+    }
+
+    @Test
+    void testDescribeBucketsWithPartitionSpec() throws Exception {
+        BucketInfo bucketInfo =
+                new BucketInfo(
+                        TABLE_PATH,
+                        10L,
+                        100L,
+                        "cn$2026-09-16",
+                        1,
+                        2,
+                        3,
+                        -1,
+                        Arrays.asList(1, 2, 3),
+                        Arrays.asList(2, 3));
+        TestingAdmin admin = new TestingAdmin(Collections.singletonList(bucketInfo));
+        DescribeBucketsProcedure procedure = newProcedure(admin);
+
+        Row[] rows = procedure.call(null, "test_db.test_table", "region=cn/dt=2026-09-16");
+
+        assertThat(admin.requestedTablePath).isEqualTo(TABLE_PATH);
+        assertThat(admin.requestedPartitionSpec).isNotNull();
+        assertThat(admin.requestedPartitionSpec.getSpecMap())
+                .containsEntry("region", "cn")
+                .containsEntry("dt", "2026-09-16");
+        assertThat(rows).hasSize(1);
+        Row row = rows[0];
+        assertThat(row.getField(2)).isEqualTo(100L);
+        assertThat(row.getField(3)).isEqualTo("cn$2026-09-16");
+        assertThat(row.getField(4)).isEqualTo(1);
+        assertThat(row.getField(5)).isEqualTo(2);
+        assertThat(row.getField(6)).isEqualTo(3);
+        assertThat(row.getField(7)).isEqualTo(-1);
+        assertThat((Integer[]) row.getField(8)).containsExactly(1, 2, 3);
+        assertThat((Integer[]) row.getField(9)).containsExactly(2, 3);
+    }
+
+    @Test
+    void testEmptyAdminResult() throws Exception {
+        DescribeBucketsProcedure procedure =
+                newProcedure(new TestingAdmin(Collections.emptyList()));
+
+        assertThat(procedure.call(null, "test_db.test_table")).isEmpty();
+    }
+
+    @Test
+    void testInvalidTablePath() {
+        DescribeBucketsProcedure procedure =
+                newProcedure(new TestingAdmin(Collections.emptyList()));
+
+        for (String tablePath :
+                Arrays.asList(null, "", " ", "test_db", ".test_table", "test_db.", "a.b.c")) {
+            assertThatThrownBy(() -> procedure.call(null, tablePath))
+                    .isInstanceOf(IllegalArgumentException.class)
+                    .hasMessageContaining("table_path");
+        }
+    }
+
+    @Test
+    void testInvalidPartitionSpec() {
+        DescribeBucketsProcedure procedure =
+                newProcedure(new TestingAdmin(Collections.emptyList()));
+
+        for (String partitionSpec :
+                Arrays.asList(null, "", " ", "region", "=cn", "region=cn/", "region=cn//dt=d1")) {
+            assertThatThrownBy(() -> procedure.call(null, "test_db.test_table", partitionSpec))
+                    .isInstanceOf(IllegalArgumentException.class)
+                    .hasMessageContaining("partition_spec");
+        }
+
+        assertThatThrownBy(() -> procedure.call(null, "test_db.test_table", "region=cn/region=us"))
+                .isInstanceOf(IllegalArgumentException.class)
+                .hasMessageContaining("Duplicate partition key 'region'");
+    }
+
+    private static DescribeBucketsProcedure newProcedure(TestingAdmin admin) {
+        DescribeBucketsProcedure procedure = new DescribeBucketsProcedure();
+        procedure.withAdmin(admin);
+        return procedure;
+    }
+
+    private static class TestingAdmin extends TestAdminAdapter {
+        private final List<BucketInfo> bucketInfos;
+        private @Nullable TablePath requestedTablePath;
+        private @Nullable PartitionSpec requestedPartitionSpec;
+
+        private TestingAdmin(List<BucketInfo> bucketInfos) {
+            this.bucketInfos = bucketInfos;
+        }
+
+        @Override
+        public CompletableFuture<List<BucketInfo>> describeBuckets(TablePath tablePath) {
+            requestedTablePath = tablePath;
+            requestedPartitionSpec = null;
+            return CompletableFuture.completedFuture(bucketInfos);
+        }
+
+        @Override
+        public CompletableFuture<List<BucketInfo>> describeBuckets(
+                TablePath tablePath, PartitionSpec partitionSpec) {
+            requestedTablePath = tablePath;
+            requestedPartitionSpec = partitionSpec;
+            return CompletableFuture.completedFuture(bucketInfos);
+        }
+    }
+}