From 6ee930662ec5d863bb661cc40c463f0fc9a6b0e4 Mon Sep 17 00:00:00 2001
From: Pierre Villard <pierre.villard.fr@gmail.com>
Date: Mon, 22 Jun 2026 20:54:20 +0200
Subject: [PATCH] NIFI-16040 - Correct branches paging in BitBucket registry
 client

---
 .../bitbucket/BitbucketRepositoryClient.java  | 30 +++++++------------
 .../BitbucketRepositoryClientTest.java        | 18 +++++++++++
 2 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/nifi-extension-bundles/nifi-atlassian-bundle/nifi-atlassian-extensions/src/main/java/org/apache/nifi/atlassian/bitbucket/BitbucketRepositoryClient.java b/nifi-extension-bundles/nifi-atlassian-bundle/nifi-atlassian-extensions/src/main/java/org/apache/nifi/atlassian/bitbucket/BitbucketRepositoryClient.java
index 2972f0b40a21..5a62df207e08 100644
--- a/nifi-extension-bundles/nifi-atlassian-bundle/nifi-atlassian-extensions/src/main/java/org/apache/nifi/atlassian/bitbucket/BitbucketRepositoryClient.java
+++ b/nifi-extension-bundles/nifi-atlassian-bundle/nifi-atlassian-extensions/src/main/java/org/apache/nifi/atlassian/bitbucket/BitbucketRepositoryClient.java
@@ -251,29 +251,19 @@ public Set<String> getBranches() throws FlowRegistryException {
 
     private Set<String> getBranchesCloud() throws FlowRegistryException {
         final URI uri = getRepositoryUriBuilder().addPathSegment("refs").addPathSegment("branches").build();
-        try (final HttpResponseEntity response = this.webClient.getWebClientService()
-                .get()
-                .uri(uri)
-                .header(AUTHORIZATION_HEADER, authToken.getAuthzHeaderValue())
-                .retrieve()) {
+        final String errorMessage = "Error while listing branches for repository [%s]".formatted(repoName);
+        final Iterator<JsonNode> branches = getPagedResponseValues(uri, errorMessage);
 
-            verifyStatusCode(response, "Error while listing branches for repository [%s]".formatted(repoName), HttpURLConnection.HTTP_OK);
-
-            final JsonNode jsonResponse = parseResponseBody(response, uri);
-            final JsonNode values = jsonResponse.get(FIELD_VALUES);
-            final Set<String> result = new HashSet<>();
-            if (values != null && values.isArray()) {
-                for (JsonNode branch : values) {
-                    final String branchName = branch.path(FIELD_NAME).asText(EMPTY_STRING);
-                    if (!branchName.isEmpty()) {
-                        result.add(branchName);
-                    }
-                }
+        final Set<String> result = new HashSet<>();
+        while (branches.hasNext()) {
+            final JsonNode branch = branches.next();
+            final String branchName = branch.path(FIELD_NAME).asText(EMPTY_STRING);
+            if (!branchName.isEmpty()) {
+                result.add(branchName);
             }
-            return result;
-        } catch (final IOException e) {
-            throw new FlowRegistryException("Failed closing Bitbucket branch listing response", e);
         }
+
+        return result;
     }
 
     private Set<String> getBranchesDataCenter() throws FlowRegistryException {
diff --git a/nifi-extension-bundles/nifi-atlassian-bundle/nifi-atlassian-extensions/src/test/java/org/apache/nifi/atlassian/bitbucket/BitbucketRepositoryClientTest.java b/nifi-extension-bundles/nifi-atlassian-bundle/nifi-atlassian-extensions/src/test/java/org/apache/nifi/atlassian/bitbucket/BitbucketRepositoryClientTest.java
index 962c48cd0083..8129d237ca19 100644
--- a/nifi-extension-bundles/nifi-atlassian-bundle/nifi-atlassian-extensions/src/test/java/org/apache/nifi/atlassian/bitbucket/BitbucketRepositoryClientTest.java
+++ b/nifi-extension-bundles/nifi-atlassian-bundle/nifi-atlassian-extensions/src/test/java/org/apache/nifi/atlassian/bitbucket/BitbucketRepositoryClientTest.java
@@ -42,6 +42,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.Optional;
 import java.util.OptionalLong;
+import java.util.Set;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -370,6 +371,23 @@ void testCreateBranchBlankSourceRejected() throws FlowRegistryException {
         assertThrows(IllegalArgumentException.class, () -> client.createBranch("feature", "  ", Optional.empty()));
     }
 
+    @Test
+    void testGetBranchesCloudFollowsPagination() throws FlowRegistryException {
+        final String firstPage = "{\"values\":[{\"name\":\"main\"},{\"name\":\"develop\"}],"
+                + "\"next\":\"https://api.bitbucket.org/2.0/repositories/test-workspace/test-repo/refs/branches?page=2\"}";
+        final String secondPage = "{\"values\":[{\"name\":\"feature-1\"},{\"name\":\"feature-2\"}]}";
+        stubGetChain(
+                branchListResponse(),
+                mockResponse(HttpURLConnection.HTTP_OK, firstPage),
+                mockResponse(HttpURLConnection.HTTP_OK, secondPage)
+        );
+
+        final BitbucketRepositoryClient client = buildCloudClient();
+        final Set<String> branches = client.getBranches();
+
+        assertEquals(Set.of("main", "develop", "feature-1", "feature-2"), branches);
+    }
+
     private BitbucketRepositoryClient buildDataCenterClient() throws FlowRegistryException {
         return BitbucketRepositoryClient.builder()
                 .clientId("test-client")