Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 60 additions & 0 deletions airbyte/_util/api_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -285,6 +285,34 @@ def get_workspace(
)


def get_workspace_via_config_api(
workspace_id: str,
*,
api_root: str,
client_id: SecretString | None,
client_secret: SecretString | None,
bearer_token: SecretString | None,
config_api_root: str | None = None,
) -> dict[str, Any]:
"""Get a workspace via the Config API endpoint `POST /v1/workspaces/get`.

Unlike `get_workspace`, this uses the internal Config API rather than the public API.
The Config API accepts user-realm bearer tokens (from an interactive OIDC login),
which the public API rejects because it only accepts application-client tokens.
Comment on lines +297 to +301

Returns the raw `WorkspaceRead` response as a dictionary.
"""
return _make_config_api_request(
path="/workspaces/get",
json={"workspaceId": workspace_id},
api_root=api_root,
config_api_root=config_api_root,
client_id=client_id,
client_secret=client_secret,
bearer_token=bearer_token,
)


def create_workspace(
*,
name: str,
Expand Down Expand Up @@ -798,6 +826,38 @@ def get_connection(
)


def get_connection_via_config_api(
connection_id: str,
*,
api_root: str,
client_id: SecretString | None,
client_secret: SecretString | None,
bearer_token: SecretString | None,
config_api_root: str | None = None,
) -> dict[str, Any]:
"""Get a connection via the Config API endpoint `POST /v1/web_backend/connections/get`.

Unlike `get_connection`, this uses the internal Config API rather than the public API.
The Config API accepts user-realm bearer tokens (from an interactive OIDC login),
which the public API rejects because it only accepts application-client tokens.
Comment on lines +838 to +842

The `WebBackendConnectionRead` response embeds the full `source` and `destination`
objects (including their names), so callers can resolve connector names without
issuing separate public API lookups.

Returns the raw `WebBackendConnectionRead` response as a dictionary.
"""
return _make_config_api_request(
path="/web_backend/connections/get",
json={"connectionId": connection_id, "withRefreshedCatalog": False},
api_root=api_root,
config_api_root=config_api_root,
client_id=client_id,
client_secret=client_secret,
bearer_token=bearer_token,
)


def run_connection(
workspace_id: str,
connection_id: str,
Expand Down
74 changes: 63 additions & 11 deletions airbyte/cloud/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,16 +111,32 @@ def _fetch_connection_info(
self._verify_workspace_match(self._connection_info)
return self._connection_info

# Fetch from API
connection_info = api_util.get_connection(
workspace_id=self.workspace.workspace_id,
connection_id=self.connection_id,
api_root=self.workspace.api_root,
client_id=self.workspace.client_id,
client_secret=self.workspace.client_secret,
bearer_token=self.workspace.bearer_token,
)
result = CloudConnectionInfo.from_api_response(connection_info)
# Fetch from API. A bearer-only (interactive OIDC) token is rejected by the
# public API, so route those reads through the Config API, which also embeds
# the source/destination objects and sync catalog in a single response.
if self.workspace._uses_bearer_only_auth: # noqa: SLF001 # Internal auth check
raw = api_util.get_connection_via_config_api(
connection_id=self.connection_id,
api_root=self.workspace.api_root,
client_id=self.workspace.client_id,
client_secret=self.workspace.client_secret,
bearer_token=self.workspace.bearer_token,
config_api_root=self.workspace.config_api_root,
)
result = CloudConnectionInfo.from_config_api_response(
raw,
fallback_workspace_id=self.workspace.workspace_id,
)
else:
connection_info = api_util.get_connection(
workspace_id=self.workspace.workspace_id,
connection_id=self.connection_id,
api_root=self.workspace.api_root,
client_id=self.workspace.client_id,
client_secret=self.workspace.client_secret,
bearer_token=self.workspace.bearer_token,
)
result = CloudConnectionInfo.from_api_response(connection_info)

self._connection_info = result

Expand Down Expand Up @@ -242,13 +258,49 @@ def destination(self) -> CloudDestination:
)
return self._cloud_destination_object

@property
def source_name(self) -> str | None:
"""The display name of the source, if available.

Prefers a name embedded in the connection info (Config API reads). Falls back to
a separate source lookup for public API reads that do not embed the name.
"""
if not self._connection_info:
self._connection_info = self._fetch_connection_info()

if self._connection_info.source_name is not None:
return self._connection_info.source_name

return self.source.name

@property
def destination_name(self) -> str | None:
"""The display name of the destination, if available.

Prefers a name embedded in the connection info (Config API reads). Falls back to
a separate destination lookup for public API reads that do not embed the name.
"""
if not self._connection_info:
self._connection_info = self._fetch_connection_info()

if self._connection_info.destination_name is not None:
return self._connection_info.destination_name

return self.destination.name

@property
def stream_names(self) -> list[str]:
"""The stream names."""
if not self._connection_info:
self._connection_info = self._fetch_connection_info()

return [stream.name for stream in self._connection_info.configurations.streams or []]
if self._connection_info.stream_names is not None:
return self._connection_info.stream_names

configurations = self._connection_info.configurations
if configurations is None:
return []
return [stream.name for stream in configurations.streams or []]

@property
def table_prefix(self) -> str:
Expand Down
74 changes: 72 additions & 2 deletions airbyte/cloud/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,15 +133,28 @@ class CloudConnectionInfo(BaseModel):
name: str
"""The connection name."""

configurations: Any
"""Stream configuration details for the connection."""
configurations: Any = None
"""Stream configuration details for the connection.

Populated from the public API response. `None` when the connection was read via
the Config API, in which case `stream_names` carries the selected stream names.
"""

prefix: str | None = None
"""The destination table prefix."""

status: str
"""The connection status."""

source_name: str | None = None
"""The source display name, when embedded in the response (Config API only)."""

destination_name: str | None = None
"""The destination display name, when embedded in the response (Config API only)."""

stream_names: list[str] | None = None
"""Selected stream names, when the response embeds the sync catalog (Config API only)."""

@classmethod
def from_api_response(cls, connection: _ConnectionResponseLike) -> CloudConnectionInfo:
"""Create a public model from an internal API connection response."""
Expand All @@ -156,6 +169,52 @@ def from_api_response(cls, connection: _ConnectionResponseLike) -> CloudConnecti
status=_enum_value(connection.status),
)

@classmethod
def from_config_api_response(
cls,
connection: Mapping[str, Any],
*,
fallback_workspace_id: str | None = None,
) -> CloudConnectionInfo:
"""Create a public model from a Config API `WebBackendConnectionRead` response.

The Config API response embeds the full `source` and `destination` objects and
the `syncCatalog`, so connector names and selected streams are resolved from this
single response rather than via separate public API lookups. The `syncCatalog`
lists every discovered stream with a per-stream `config.selected` flag, so only
selected streams are returned, matching the public API's `stream_names` behavior.
"""
source: Mapping[str, Any] = connection.get("source") or {}
destination: Mapping[str, Any] = connection.get("destination") or {}
sync_catalog: Mapping[str, Any] = connection.get("syncCatalog") or {}
stream_entries: list[Mapping[str, Any]] = sync_catalog.get("streams") or []
stream_names = [
entry["stream"]["name"]
for entry in stream_entries
if isinstance(entry.get("stream"), Mapping)
and entry["stream"].get("name")
and _is_stream_selected(entry.get("config"))
]
Comment thread
aaronsteers marked this conversation as resolved.
Comment on lines +191 to +197

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Filter out deselected catalog streams?

Could we require entry["config"]["selected"] here, wdyt? This currently exposes all catalog streams as selected_streams, so disabled streams are reported by CloudConnection.stream_names and the MCP connection description.

Proposed fix
 stream_names = [
     entry["stream"]["name"]
     for entry in stream_entries
-    if isinstance(entry.get("stream"), Mapping) and entry["stream"].get("name")
+    if (
+        isinstance(entry.get("stream"), Mapping)
+        and entry["stream"].get("name")
+        and isinstance(entry.get("config"), Mapping)
+        and entry["config"].get("selected") is True
+    )
 ]

Please add a deselected stream to the fixture and assert it is excluded.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
stream_names = [
entry["stream"]["name"]
for entry in stream_entries
if isinstance(entry.get("stream"), Mapping) and entry["stream"].get("name")
]
stream_names = [
entry["stream"]["name"]
for entry in stream_entries
if (
isinstance(entry.get("stream"), Mapping)
and entry["stream"].get("name")
and isinstance(entry.get("config"), Mapping)
and entry["config"].get("selected") is True
)
]
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@airbyte/cloud/models.py` around lines 189 - 193, Update the stream_names
comprehension in CloudConnection to include only entries whose config indicates
selected=True, while preserving the existing stream mapping and nonempty-name
checks. Extend the relevant fixture with a deselected stream and assert it is
excluded from selected_streams and the MCP connection description.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

☑️ Resolved in 8dd2ebf. Agreed — filtered stream_names to selected streams. I use a _is_stream_selected helper rather than a strict config.selected is True: when config is present I honor selected, but a missing config defaults to selected (matching the platform default) so unusual payloads don't silently drop streams. Added a deselected stream to the fixture and assert it's excluded.

workspace_id = (
source.get("workspaceId")
or destination.get("workspaceId")
or fallback_workspace_id
or ""
)
return cls(
connection_id=connection["connectionId"],
workspace_id=workspace_id,
source_id=connection["sourceId"],
destination_id=connection["destinationId"],
name=connection["name"],
configurations=None,
prefix=connection.get("prefix"),
status=_enum_value(connection["status"]),
source_name=source.get("name"),
destination_name=destination.get("name"),
stream_names=stream_names,
Comment on lines +210 to +215
)


class CloudJobInfo(BaseModel):
"""Information about an Airbyte Cloud job."""
Expand Down Expand Up @@ -260,6 +319,17 @@ def _notifications_to_dict(notifications: object) -> dict[str, object | None]:
return {}


def _is_stream_selected(config: object) -> bool:
"""Return whether a `syncCatalog` stream entry is selected for sync.

A stream is selected when its `config.selected` flag is truthy. A missing `config` or
missing `selected` flag defaults to selected, matching the platform default.
"""
if not isinstance(config, Mapping):
return True
return bool(config.get("selected", True))


def _enum_value(value: object) -> str:
"""Return the string value for an enum-like object."""
if isinstance(value, Enum):
Expand Down
47 changes: 40 additions & 7 deletions airbyte/cloud/workspaces.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,45 @@ def workspace_url(self) -> str | None:
"""The web URL of the workspace."""
return f"{get_web_url_root(self.api_root)}/workspaces/{self.workspace_id}"

@property
def _uses_bearer_only_auth(self) -> bool:
"""Whether this workspace authenticates with a bearer token and no client credentials.

An interactive OIDC login produces a user-realm bearer token with no client
credentials. The public API rejects such tokens (it only accepts application-client
tokens), while the Config API accepts them, so metadata reads must route through the
Config API in this case.
"""
return (
self.bearer_token is not None and self.client_id is None and self.client_secret is None
)

def get_workspace_info(self) -> CloudWorkspaceInfo:
"""Get metadata about this workspace.

Routes bearer-only (interactive OIDC) reads through the Config API and all other
reads through the public API.
"""
if self._uses_bearer_only_auth:
raw = api_util.get_workspace_via_config_api(
workspace_id=self.workspace_id,
api_root=self.api_root,
client_id=self.client_id,
client_secret=self.client_secret,
bearer_token=self.bearer_token,
config_api_root=self.config_api_root,
)
return CloudWorkspaceInfo.from_mapping(raw)

response = api_util.get_workspace(
workspace_id=self.workspace_id,
api_root=self.api_root,
client_id=self.client_id,
client_secret=self.client_secret,
bearer_token=self.bearer_token,
)
return CloudWorkspaceInfo.from_api_response(response)
Comment thread
coderabbitai[bot] marked this conversation as resolved.

@cached_property
def _organization_info(self) -> dict[str, Any]:
"""Fetch and cache organization info for this workspace.
Expand Down Expand Up @@ -310,13 +349,7 @@ def connect(self) -> None:
serves primarily as a simple check to ensure that the workspace is reachable
and credentials are correct.
"""
_ = api_util.get_workspace(
api_root=self.api_root,
workspace_id=self.workspace_id,
client_id=self.client_id,
client_secret=self.client_secret,
bearer_token=self.bearer_token,
)
_ = self.get_workspace_info()
print(f"Successfully connected to workspace: {self.workspace_url}")

# Get sources, destinations, and connections
Expand Down
16 changes: 5 additions & 11 deletions airbyte/mcp/cloud.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
from pydantic import BaseModel, Field

from airbyte import cloud, get_destination, get_source
from airbyte._util import api_util
from airbyte.cloud.client import CloudClient
from airbyte.cloud.connectors import CustomCloudSourceDefinition
from airbyte.cloud.constants import FAILED_STATUSES
Expand Down Expand Up @@ -567,14 +566,9 @@ def check_airbyte_cloud_workspace(
"""
workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id)

# Get workspace details from the public API using workspace's credentials
workspace_response = api_util.get_workspace(
workspace_id=workspace.workspace_id,
api_root=workspace.api_root,
client_id=workspace.client_id,
client_secret=workspace.client_secret,
bearer_token=workspace.bearer_token,
)
# Get workspace details. Reads are routed through the Config API for bearer-only
# (interactive OIDC) credentials, which the public API rejects.
workspace_response = workspace.get_workspace_info()

# Try to get organization info (including billing), but fail gracefully if we don't have
# permissions. Fetching organization info requires ORGANIZATION_READER permissions on the
Expand Down Expand Up @@ -1006,9 +1000,9 @@ def describe_cloud_connection(
connection_name=cast(str, connection.name),
connection_url=cast(str, connection.connection_url),
source_id=connection.source_id,
source_name=cast(str, connection.source.name),
source_name=cast(str, connection.source_name),
destination_id=connection.destination_id,
destination_name=cast(str, connection.destination.name),
destination_name=cast(str, connection.destination_name),
selected_streams=connection.stream_names,
table_prefix=connection.table_prefix,
)
Expand Down
Loading
Loading