feat: bitbucket connector (#5294)

Author: danielkravets

backend/onyx/configs/constants.py

@@ -199,6 +199,7 @@ class DocumentSource(str, Enum):
     HIGHSPOT = "highspot"
 
     IMAP = "imap"
+    BITBUCKET = "bitbucket"
 
     # Special case just for integration tests
     MOCK_CONNECTOR = "mock_connector"
@@ -541,6 +542,7 @@ class OnyxCallTypes(str, Enum):
     DocumentSource.GITHUB: "github data (issues, PRs)",
     DocumentSource.GITBOOK: "gitbook data",
     DocumentSource.GITLAB: "gitlab data",
+    DocumentSource.BITBUCKET: "bitbucket data",
     DocumentSource.GURU: "guru data",
     DocumentSource.BOOKSTACK: "bookstack data",
     DocumentSource.OUTLINE: "outline data",

backend/onyx/connectors/bitbucket/__init__.py

@@ -0,0 +1,345 @@
+from __future__ import annotations
+
+import copy
+from collections.abc import Callable
+from collections.abc import Iterator
+from datetime import datetime
+from datetime import timezone
+from typing import Any
+from typing import TYPE_CHECKING
+
+from typing_extensions import override
+
+from onyx.configs.app_configs import INDEX_BATCH_SIZE
+from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS
+from onyx.configs.constants import DocumentSource
+from onyx.connectors.bitbucket.utils import build_auth_client
+from onyx.connectors.bitbucket.utils import list_repositories
+from onyx.connectors.bitbucket.utils import map_pr_to_document
+from onyx.connectors.bitbucket.utils import paginate
+from onyx.connectors.bitbucket.utils import PR_LIST_RESPONSE_FIELDS
+from onyx.connectors.bitbucket.utils import SLIM_PR_LIST_RESPONSE_FIELDS
+from onyx.connectors.exceptions import CredentialExpiredError
+from onyx.connectors.exceptions import InsufficientPermissionsError
+from onyx.connectors.exceptions import UnexpectedValidationError
+from onyx.connectors.interfaces import CheckpointedConnector
+from onyx.connectors.interfaces import CheckpointOutput
+from onyx.connectors.interfaces import SecondsSinceUnixEpoch
+from onyx.connectors.interfaces import SlimConnector
+from onyx.connectors.models import ConnectorCheckpoint
+from onyx.connectors.models import ConnectorFailure
+from onyx.connectors.models import ConnectorMissingCredentialError
+from onyx.connectors.models import DocumentFailure
+from onyx.connectors.models import SlimDocument
+from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface
+from onyx.utils.logger import setup_logger
+
+if TYPE_CHECKING:
+    import httpx
+
+logger = setup_logger()
+
+
+class BitbucketConnectorCheckpoint(ConnectorCheckpoint):
+    """Checkpoint state for resumable Bitbucket PR indexing.
+
+    Fields:
+        repos_queue: Materialized list of repository slugs to process.
+        current_repo_index: Index of the repository currently being processed.
+        next_url: Bitbucket "next" URL for continuing pagination within the current repo.
+    """
+
+    repos_queue: list[str] = []
+    current_repo_index: int = 0
+    next_url: str | None = None
+
+
+class BitbucketConnector(
+    CheckpointedConnector[BitbucketConnectorCheckpoint],
+    SlimConnector,
+):
+    """Connector for indexing Bitbucket Cloud pull requests.
+
+    Args:
+        workspace: Bitbucket workspace ID.
+        repositories: Comma-separated list of repository slugs to index.
+        projects: Comma-separated list of project keys to index all repositories within.
+        batch_size: Max number of documents to yield per batch.
+    """
+
+    def __init__(
+        self,
+        workspace: str,
+        repositories: str | None = None,
+        projects: str | None = None,
+        batch_size: int = INDEX_BATCH_SIZE,
+    ) -> None:
+        self.workspace = workspace
+        self._repositories = (
+            [s.strip() for s in repositories.split(",") if s.strip()]
+            if repositories
+            else None
+        )
+        self._projects: list[str] | None = (
+            [s.strip() for s in projects.split(",") if s.strip()] if projects else None
+        )
+        self.batch_size = batch_size
+        self.email: str | None = None
+        self.api_token: str | None = None
+
+    def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:
+        """Load API token-based credentials.
+
+        Expects a dict with keys: `bitbucket_email`, `bitbucket_api_token`.
+        """
+        self.email = credentials.get("bitbucket_email")
+        self.api_token = credentials.get("bitbucket_api_token")
+        if not self.email or not self.api_token:
+            raise ConnectorMissingCredentialError("Bitbucket")
+        return None
+
+    def _client(self) -> httpx.Client:
+        """Build an authenticated HTTP client or raise if credentials missing."""
+        if not self.email or not self.api_token:
+            raise ConnectorMissingCredentialError("Bitbucket")
+        return build_auth_client(self.email, self.api_token)
+
+    def _iter_pull_requests_for_repo(
+        self,
+        client: httpx.Client,
+        repo_slug: str,
+        params: dict[str, Any] | None = None,
+        start_url: str | None = None,
+        on_page: Callable[[str | None], None] | None = None,
+    ) -> Iterator[dict[str, Any]]:
+        base = f"https://api.bitbucket.org/2.0/repositories/{self.workspace}/{repo_slug}/pullrequests"
+        yield from paginate(
+            client,
+            base,
+            params,
+            start_url=start_url,
+            on_page=on_page,
+        )
+
+    def _build_params(
+        self,
+        fields: str = PR_LIST_RESPONSE_FIELDS,
+        start: SecondsSinceUnixEpoch | None = None,
+        end: SecondsSinceUnixEpoch | None = None,
+    ) -> dict[str, Any]:
+        """Build Bitbucket fetch params.
+
+        Always include OPEN, MERGED, and DECLINED PRs. If both ``start`` and
+        ``end`` are provided, apply a single updated_on time window.
+        """
+
+        def _iso(ts: SecondsSinceUnixEpoch) -> str:
+            return datetime.fromtimestamp(ts, tz=timezone.utc).isoformat()
+
+        def _tc_epoch(
+            lower_epoch: SecondsSinceUnixEpoch | None,
+            upper_epoch: SecondsSinceUnixEpoch | None,
+        ) -> str | None:
+            if lower_epoch is not None and upper_epoch is not None:
+                lower_iso = _iso(lower_epoch)
+                upper_iso = _iso(upper_epoch)
+                return f'(updated_on >= "{lower_iso}" AND updated_on <= "{upper_iso}")'
+            return None
+
+        params: dict[str, Any] = {"fields": fields, "pagelen": 50}
+        time_clause = _tc_epoch(start, end)
+        q = '(state = "OPEN" OR state = "MERGED" OR state = "DECLINED")'
+        if time_clause:
+            q = f"{q} AND {time_clause}"
+        params["q"] = q
+        return params
+
+    def _iter_target_repositories(self, client: httpx.Client) -> Iterator[str]:
+        """Yield repository slugs based on configuration.
+
+        Priority:
+        - repositories list
+        - projects list (list repos by project key)
+        - workspace (all repos)
+        """
+        if self._repositories:
+            for slug in self._repositories:
+                yield slug
+            return
+        if self._projects:
+            for project_key in self._projects:
+                for repo in list_repositories(client, self.workspace, project_key):
+                    slug_val = repo.get("slug")
+                    if isinstance(slug_val, str) and slug_val:
+                        yield slug_val
+            return
+        for repo in list_repositories(client, self.workspace, None):
+            slug_val = repo.get("slug")
+            if isinstance(slug_val, str) and slug_val:
+                yield slug_val
+
+    @override
+    def load_from_checkpoint(
+        self,
+        start: SecondsSinceUnixEpoch,
+        end: SecondsSinceUnixEpoch,
+        checkpoint: BitbucketConnectorCheckpoint,
+    ) -> CheckpointOutput[BitbucketConnectorCheckpoint]:
+        """Resumable PR ingestion across repos and pages within a time window.
+
+        Yields Documents (or ConnectorFailure for per-PR mapping failures) and returns
+        an updated checkpoint that records repo position and next page URL.
+        """
+        new_checkpoint = copy.deepcopy(checkpoint)
+
+        with self._client() as client:
+            # Materialize target repositories once
+            if not new_checkpoint.repos_queue:
+                # Preserve explicit order; otherwise ensure deterministic ordering
+                repos_list = list(self._iter_target_repositories(client))
+                new_checkpoint.repos_queue = sorted(set(repos_list))
+                new_checkpoint.current_repo_index = 0
+                new_checkpoint.next_url = None
+
+            repos = new_checkpoint.repos_queue
+            if not repos or new_checkpoint.current_repo_index >= len(repos):
+                new_checkpoint.has_more = False
+                return new_checkpoint
+
+            repo_slug = repos[new_checkpoint.current_repo_index]
+
+            first_page_params = self._build_params(
+                fields=PR_LIST_RESPONSE_FIELDS,
+                start=start,
+                end=end,
+            )
+
+            def _on_page(next_url: str | None) -> None:
+                new_checkpoint.next_url = next_url
+
+            for pr in self._iter_pull_requests_for_repo(
+                client,
+                repo_slug,
+                params=first_page_params,
+                start_url=new_checkpoint.next_url,
+                on_page=_on_page,
+            ):
+                try:
+                    document = map_pr_to_document(pr, self.workspace, repo_slug)
+                    yield document
+                except Exception as e:
+                    pr_id = pr.get("id")
+                    pr_link = (
+                        f"https://bitbucket.org/{self.workspace}/{repo_slug}/pull-requests/{pr_id}"
+                        if pr_id is not None
+                        else None
+                    )
+                    yield ConnectorFailure(
+                        failed_document=DocumentFailure(
+                            document_id=(
+                                f"{DocumentSource.BITBUCKET.value}:{self.workspace}:{repo_slug}:pr:{pr_id}"
+                                if pr_id is not None
+                                else f"{DocumentSource.BITBUCKET.value}:{self.workspace}:{repo_slug}:pr:unknown"
+                            ),
+                            document_link=pr_link,
+                        ),
+                        failure_message=f"Failed to process Bitbucket PR: {e}",
+                        exception=e,
+                    )
+
+            # Advance to next repository (if any) and set has_more accordingly
+            new_checkpoint.current_repo_index += 1
+            new_checkpoint.next_url = None
+            new_checkpoint.has_more = new_checkpoint.current_repo_index < len(repos)
+
+        return new_checkpoint
+
+    @override
+    def build_dummy_checkpoint(self) -> BitbucketConnectorCheckpoint:
+        """Create an initial checkpoint with work remaining."""
+        return BitbucketConnectorCheckpoint(has_more=True)
+
+    @override
+    def validate_checkpoint_json(
+        self, checkpoint_json: str
+    ) -> BitbucketConnectorCheckpoint:
+        """Validate and deserialize a checkpoint instance from JSON."""
+        return BitbucketConnectorCheckpoint.model_validate_json(checkpoint_json)
+
+    def retrieve_all_slim_documents(
+        self,
+        start: SecondsSinceUnixEpoch | None = None,
+        end: SecondsSinceUnixEpoch | None = None,
+        callback: IndexingHeartbeatInterface | None = None,
+    ) -> Iterator[list[SlimDocument]]:
+        """Return only document IDs for all existing pull requests."""
+        batch: list[SlimDocument] = []
+        params = self._build_params(
+            fields=SLIM_PR_LIST_RESPONSE_FIELDS,
+            start=start,
+            end=end,
+        )
+        with self._client() as client:
+            for slug in self._iter_target_repositories(client):
+                for pr in self._iter_pull_requests_for_repo(
+                    client, slug, params=params
+                ):
+                    pr_id = pr["id"]
+                    doc_id = f"{DocumentSource.BITBUCKET.value}:{self.workspace}:{slug}:pr:{pr_id}"
+                    batch.append(SlimDocument(id=doc_id))
+                    if len(batch) >= self.batch_size:
+                        yield batch
+                        batch = []
+                        if callback:
+                            if callback.should_stop():
+                                # Note: this is not actually used for permission sync yet, just pruning
+                                raise RuntimeError(
+                                    "bitbucket_pr_sync: Stop signal detected"
+                                )
+                            callback.progress("bitbucket_pr_sync", len(batch))
+        if batch:
+            yield batch
+
+    def validate_connector_settings(self) -> None:
+        """Validate Bitbucket credentials and workspace access by probing a lightweight endpoint.
+
+        Raises:
+            CredentialExpiredError: on HTTP 401
+            InsufficientPermissionsError: on HTTP 403
+            UnexpectedValidationError: on any other failure
+        """
+        try:
+            with self._client() as client:
+                url = f"https://api.bitbucket.org/2.0/repositories/{self.workspace}"
+                resp = client.get(
+                    url,
+                    params={"pagelen": 1, "fields": "pagelen"},
+                    timeout=REQUEST_TIMEOUT_SECONDS,
+                )
+                if resp.status_code == 401:
+                    raise CredentialExpiredError(
+                        "Invalid or expired Bitbucket credentials (HTTP 401)."
+                    )
+                if resp.status_code == 403:
+                    raise InsufficientPermissionsError(
+                        "Insufficient permissions to access Bitbucket workspace (HTTP 403)."
+                    )
+                if resp.status_code < 200 or resp.status_code >= 300:
+                    raise UnexpectedValidationError(
+                        f"Unexpected Bitbucket error (status={resp.status_code})."
+                    )
+        except Exception as e:
+            # Network or other unexpected errors
+            if isinstance(
+                e,
+                (
+                    CredentialExpiredError,
+                    InsufficientPermissionsError,
+                    UnexpectedValidationError,
+                    ConnectorMissingCredentialError,
+                ),
+            ):
+                raise
+            raise UnexpectedValidationError(
+                f"Unexpected error while validating Bitbucket settings: {e}"
+            )