diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index bb6ebc4da..910896771 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -94,6 +94,10 @@ _T = TypeVar("_T") +DEV_NULL_DESTINATION_DEFINITION_ID = "a7bcc9d8-13b3-4e49-b80d-d020b90045e3" +"""Definition ID of the `destination-dev-null` (no-op) connector on Airbyte Cloud.""" + + def _register_queued_job_status() -> None: """Register the `queued` job status on the `airbyte-api` SDK enum. @@ -922,6 +926,34 @@ def list_destinations( # Get and run connections +def get_source_stream_names( + source_id: str, + *, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, +) -> list[str]: + """Discover a source's streams (bypassing the cached catalog) and return their names.""" + airbyte_instance = get_airbyte_server_instance( + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + api_root=api_root, + ) + try: + response = airbyte_instance.streams.get_stream_properties( + api.GetStreamPropertiesRequest(source_id=source_id, ignore_cache=True), + ) + except SDKError as e: + raise _wrap_sdk_error(e, {"source_id": source_id, "api_root": api_root}) from e + return [ + stream.stream_name + for stream in response.stream_properties_response or [] + if stream.stream_name + ] + + def get_connection( workspace_id: str, connection_id: str, @@ -1583,9 +1615,7 @@ def create_destination( ) definition_id_override: str | None = None if _get_destination_type_str(config) == "dev-null": - # TODO: We have to hard-code the definition ID for dev-null destination. - # https://github.com/airbytehq/PyAirbyte/issues/743 - definition_id_override = "a7bcc9d8-13b3-4e49-b80d-d020b90045e3" + definition_id_override = DEV_NULL_DESTINATION_DEFINITION_ID try: response: api.CreateDestinationResponse = airbyte_instance.destinations.create_destination( models.DestinationCreateRequest( @@ -3047,6 +3077,7 @@ def get_connector_builder_project( client_secret: SecretString | None, bearer_token: SecretString | None, config_api_root: str | None = None, + version: int | None = None, ) -> dict[str, Any]: """Get a connector builder project, including the draft manifest if one exists. @@ -3063,6 +3094,8 @@ def get_connector_builder_project( client_secret: OAuth client secret bearer_token: Bearer token for authentication (alternative to client credentials). config_api_root: Optional explicit Config API root URL. + version: Optional published manifest version. When set, `declarativeManifest` + holds that published version instead of the draft. Returns: A dictionary containing the builder project details. Key fields include: @@ -3070,11 +3103,43 @@ def get_connector_builder_project( - declarativeManifest: The draft manifest data (if hasDraft is True), which contains a 'manifest' field with the actual YAML manifest dict. """ + request_body: dict[str, Any] = { + "workspaceId": workspace_id, + "builderProjectId": builder_project_id, + } + if version is not None: + request_body["version"] = version return _make_config_api_request( path="/connector_builder_projects/get_with_manifest", + json=request_body, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + + +def list_declarative_manifest_versions( + *, + workspace_id: str, + source_definition_id: str, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, +) -> list[dict[str, Any]]: + """List the published manifest versions of a declarative custom source definition. + + Calls `POST /v1/declarative_source_definitions/list_manifests`. Each item has `version`, + `isActive` and `description`, but not the manifest. + """ + response = _make_config_api_request( + path="/declarative_source_definitions/list_manifests", json={ "workspaceId": workspace_id, - "builderProjectId": builder_project_id, + "sourceDefinitionId": source_definition_id, }, api_root=api_root, config_api_root=config_api_root, @@ -3082,6 +3147,7 @@ def get_connector_builder_project( client_secret=client_secret, bearer_token=bearer_token, ) + return response.get("manifestVersions") or [] def update_connector_builder_project( # noqa: PLR0913 @@ -3175,6 +3241,210 @@ def update_connector_builder_project_testing_values( # noqa: PLR0913 ) +def create_connector_builder_project( + *, + workspace_id: str, + name: str, + draft_manifest: dict[str, Any] | None, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, +) -> str: + """Create an unpublished connector builder project and return its ID. + + Calls `POST /v1/connector_builder_projects/create`. + """ + builder_project: dict[str, Any] = {"name": name} + if draft_manifest is not None: + builder_project["draftManifest"] = draft_manifest + + response = _make_config_api_request( + path="/connector_builder_projects/create", + json={"workspaceId": workspace_id, "builderProject": builder_project}, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + builder_project_id = response.get("builderProjectId") + if not isinstance(builder_project_id, str): + raise AirbyteCloudError( + message="Connector builder project API returned an unexpected response.", + context={"workspace_id": workspace_id}, + ) + return builder_project_id + + +def read_connector_builder_project_stream( # noqa: PLR0913 + *, + workspace_id: str, + builder_project_id: str, + manifest: dict[str, Any], + stream_name: str, + record_limit: int, + page_limit: int, + slice_limit: int, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, + timeout_seconds: float | None = None, +) -> dict[str, Any]: + """Run a test read of one stream on Airbyte Cloud, using the project's testing values. + + Calls `POST /v1/connector_builder_projects/read_stream`. `timeout_seconds` bounds the HTTP + request only; it doesn't cancel the read on the Cloud side. + """ + return _make_config_api_request( + path="/connector_builder_projects/read_stream", + json={ + "workspaceId": workspace_id, + "builderProjectId": builder_project_id, + "manifest": manifest, + "streamName": stream_name, + "recordLimit": record_limit, + "pageLimit": page_limit, + "sliceLimit": slice_limit, + }, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + timeout=None if timeout_seconds is None else (timeout_seconds, timeout_seconds), + ) + + +def full_resolve_connector_builder_project_manifest( # noqa: PLR0913 + *, + workspace_id: str, + builder_project_id: str, + manifest: dict[str, Any], + stream_limit: int, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, +) -> dict[str, Any]: + """Resolve a manifest, including dynamic streams, using the project's testing values. + + Calls `POST /v1/connector_builder_projects/full_resolve` and returns the resolved manifest. + """ + response = _make_config_api_request( + path="/connector_builder_projects/full_resolve", + json={ + "workspaceId": workspace_id, + "builderProjectId": builder_project_id, + "manifest": manifest, + "streamLimit": stream_limit, + }, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + resolved_manifest = response.get("manifest") + if not isinstance(resolved_manifest, dict): + raise AirbyteCloudError( + message="Connector builder resolve API returned an unexpected response.", + context={"workspace_id": workspace_id, "builder_project_id": builder_project_id}, + ) + return resolved_manifest + + +def publish_connector_builder_project( # noqa: PLR0913 + *, + workspace_id: str, + builder_project_id: str, + name: str, + manifest: dict[str, Any], + spec: dict[str, Any], + description: str, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, +) -> str: + """Publish a connector builder project as version 1 of a new custom source definition. + + Calls `POST /v1/connector_builder_projects/publish`, which also clears the project's + draft, and returns the new source definition ID. + """ + response = _make_config_api_request( + path="/connector_builder_projects/publish", + json={ + "workspaceId": workspace_id, + "builderProjectId": builder_project_id, + "name": name, + "initialDeclarativeManifest": { + "manifest": manifest, + "spec": spec, + "version": 1, + "description": description, + }, + }, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + source_definition_id = response.get("sourceDefinitionId") + if not isinstance(source_definition_id, str): + raise AirbyteCloudError( + message="Connector builder publish API returned an unexpected response.", + context={"workspace_id": workspace_id, "builder_project_id": builder_project_id}, + ) + return source_definition_id + + +def create_declarative_manifest_version( # noqa: PLR0913 + *, + workspace_id: str, + source_definition_id: str, + manifest: dict[str, Any], + spec: dict[str, Any], + version: int, + description: str, + api_root: str, + client_id: SecretString | None, + client_secret: SecretString | None, + bearer_token: SecretString | None, + config_api_root: str | None = None, +) -> None: + """Publish a manifest as a new active version of a declarative custom source definition. + + Calls `POST /v1/declarative_source_definitions/create_manifest`, which also deletes the + definition's Connector Builder draft. + """ + _make_config_api_request( + path="/declarative_source_definitions/create_manifest", + json={ + "workspaceId": workspace_id, + "sourceDefinitionId": source_definition_id, + "setAsActiveManifest": True, + "declarativeManifest": { + "manifest": manifest, + "spec": spec, + "version": version, + "description": description, + }, + }, + api_root=api_root, + config_api_root=config_api_root, + client_id=client_id, + client_secret=client_secret, + bearer_token=bearer_token, + ) + + # Organization and workspace listing diff --git a/airbyte/cloud/__init__.py b/airbyte/cloud/__init__.py index 2a2043f35..9b39ce135 100644 --- a/airbyte/cloud/__init__.py +++ b/airbyte/cloud/__init__.py @@ -112,6 +112,7 @@ client, client_config, connections, + connector_builder, constants, organizations, sync_results, @@ -125,6 +126,7 @@ "client", "organizations", "connections", + "connector_builder", "constants", "client_config", "sync_results", diff --git a/airbyte/cloud/connector_builder/__init__.py b/airbyte/cloud/connector_builder/__init__.py new file mode 100644 index 000000000..edee1145c --- /dev/null +++ b/airbyte/cloud/connector_builder/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Connector Builder projects in Airbyte Cloud. + +A Connector Builder project holds an unpublished draft manifest and the testing values +(including secrets) used to test it. Test reads and manifest resolution run on Airbyte Cloud, +so no connector code runs on the caller's machine. +""" + +from __future__ import annotations + +from airbyte.cloud.connector_builder._models import ( + ConnectorBuilderPageSummary, + ConnectorBuilderPublishedVersion, + ConnectorBuilderResolvedManifest, + ConnectorBuilderStreamReadResult, +) +from airbyte.cloud.connector_builder.project import ( + CloudConnectorBuilderProject, +) + + +__all__ = [ + "CloudConnectorBuilderProject", + "ConnectorBuilderPageSummary", + "ConnectorBuilderPublishedVersion", + "ConnectorBuilderResolvedManifest", + "ConnectorBuilderStreamReadResult", +] diff --git a/airbyte/cloud/connector_builder/_models.py b/airbyte/cloud/connector_builder/_models.py new file mode 100644 index 000000000..4b4926c90 --- /dev/null +++ b/airbyte/cloud/connector_builder/_models.py @@ -0,0 +1,293 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Result models for Connector Builder projects in Airbyte Cloud.""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, Field + + +_HTTP_ERROR_STATUS_MIN = 400 +_PROBLEM_LOG_LEVELS = frozenset({"WARN", "ERROR", "FATAL"}) +_FAILING_LOG_LEVELS = frozenset({"ERROR", "FATAL"}) +ManifestTextEditMode = Literal["replace_full", "replace_text", "replace_all"] +"""How `_update_draft_manifest_yaml()` edits the manifest text. + +- `replace_full`: `new_text` is the whole manifest. +- `replace_text`: replace the single exact match of `old_text`; fails on zero or several matches. +- `replace_all`: replace every exact match of `old_text`; fails on zero matches. +""" + + +class ConnectorBuilderPageSummary(BaseModel): + """One HTTP page fetched during a Connector Builder test read.""" + + url: str | None = None + http_method: str | None = None + status: int | None = None + record_count: int = 0 + response_body: str | None = Field( + default=None, + description="The response body, included only for HTTP error responses.", + ) + + +def _record_count_warnings( + record_count: int, + pages: list[ConnectorBuilderPageSummary], + *, + test_read_limit_reached: bool, +) -> list[str]: + """Return informational hints for record counts that may indicate a partial read.""" + if test_read_limit_reached: + return [] + if record_count == 0: + return ["Read 0 records. Check the request, record selector, and testing values."] + + warnings = [] + if record_count % 10 == 0: + warnings.append( + f"Read {record_count:,} records, a multiple of 10. A page size or API limit may " + "have cut the read short; check the paginator." + ) + + record_pages = [page for page in pages if page.record_count > 0] + if ( + record_pages + and record_pages[0].record_count > 1 + and pages[-1].record_count > 0 + and pages[-1].record_count == record_pages[0].record_count + ): + record_count_per_page = record_pages[0].record_count + warnings.append( + f"The last page returned a full page of {record_count_per_page:,} records, the " + "same as the first page, so more pages may exist; check the paginator." + ) + return warnings + + +class ConnectorBuilderStreamReadResult(BaseModel): + """Result of a Connector Builder test read of one stream.""" + + stream_name: str + record_count: int + slice_count: int + page_count: int + test_read_limit_reached: bool + records: list[dict[str, Any]] + pages: list[ConnectorBuilderPageSummary] + logs: list[dict[str, Any]] + inferred_schema: dict[str, Any] | None = None + warnings: list[str] = Field( + default_factory=list, + description=( + "Hints that the read may be incomplete, such as round record counts. " + "They don't fail the read." + ), + ) + + @classmethod + def _from_api_response( + cls, + stream_name: str, + response: dict[str, Any], + ) -> ConnectorBuilderStreamReadResult: + """Build a result from a `read_stream` API response.""" + records: list[dict[str, Any]] = [] + pages: list[ConnectorBuilderPageSummary] = [] + slices = response.get("slices") or [] + for slice_ in slices: + for page in slice_.get("pages") or []: + page_records = page.get("records") or [] + records.extend(page_records) + request = page.get("request") or {} + http_response = page.get("response") or {} + status = http_response.get("status") + pages.append( + ConnectorBuilderPageSummary( + url=request.get("url"), + http_method=request.get("http_method"), + status=status, + record_count=len(page_records), + response_body=( + http_response.get("body") + if isinstance(status, int) and status >= _HTTP_ERROR_STATUS_MIN + else None + ), + ) + ) + + test_read_limit_reached = bool(response.get("test_read_limit_reached")) + return cls( + stream_name=stream_name, + record_count=len(records), + slice_count=len(slices), + page_count=len(pages), + test_read_limit_reached=test_read_limit_reached, + records=records, + pages=pages, + logs=[ + {"level": log.get("level"), "message": log.get("message")} + for log in response.get("logs") or [] + ], + inferred_schema=response.get("inferred_schema"), + warnings=_record_count_warnings( + len(records), + pages, + test_read_limit_reached=test_read_limit_reached, + ), + ) + + +class ConnectorBuilderResolvedManifest(BaseModel): + """A manifest resolved by Airbyte Cloud, with dynamic streams expanded.""" + + manifest_dict: dict[str, Any] + stream_names: list[str] + + @classmethod + def _from_manifest( + cls, + manifest: dict[str, Any], + ) -> ConnectorBuilderResolvedManifest: + """Build a result from a resolved manifest.""" + return cls( + manifest_dict=manifest, + stream_names=[ + stream["name"] + for stream in manifest.get("streams") or [] + if isinstance(stream, dict) and isinstance(stream.get("name"), str) + ], + ) + + +class ConnectorBuilderPublishedVersion(BaseModel): + """One published manifest version of a Connector Builder project's definition.""" + + version: int + is_active: bool + description: str + + +class ConnectorBuilderProjectInfo(BaseModel): + """A snapshot of a Connector Builder project, read with a single project request.""" + + builder_project_id: str + definition_id: str | None + connector_builder_project_url: str + name: str + has_draft: bool + draft_manifest_dict: dict[str, Any] | None = Field(serialization_alias="draft_manifest") + draft_manifest_yaml: str | None + published_manifest_dict: dict[str, Any] | None = Field(serialization_alias="published_manifest") + + +class ConnectorBuilderStreamReadiness(BaseModel): + """Readiness of one stream, from a Connector Builder test read.""" + + stream_name: str + failed: bool = Field( + description=( + "True when the test read raised an error, returned HTTP errors, or logged errors." + ), + ) + record_count: int = 0 + page_count: int = 0 + read_to_end: bool = Field( + default=False, + description=( + "True when the stream ended within the test-read caps without HTTP errors. False " + "means a record, page, or slice cap was hit, an HTTP error occurred, or the read " + "did not run (see `error`)." + ), + ) + http_errors: list[ConnectorBuilderPageSummary] = Field(default_factory=list) + problem_logs: list[dict[str, Any]] = Field(default_factory=list) + inferred_schema: dict[str, Any] | None = None + warnings: list[str] = Field( + default_factory=list, + description=( + "Hints that the read may be incomplete, such as round record counts. " + "They don't fail the read." + ), + ) + error: str | None = Field( + default=None, + description="The error raised by the test read itself, if it failed to run.", + ) + + @classmethod + def _from_stream_read( + cls, + result: ConnectorBuilderStreamReadResult, + ) -> ConnectorBuilderStreamReadiness: + """Summarize errors and results from a stream test read.""" + http_errors = [ + page + for page in result.pages + if page.status is not None and page.status >= _HTTP_ERROR_STATUS_MIN + ] + problem_logs = [log for log in result.logs if log.get("level") in _PROBLEM_LOG_LEVELS] + has_errors = bool(http_errors) or any( + log.get("level") in _FAILING_LOG_LEVELS for log in problem_logs + ) + return cls( + stream_name=result.stream_name, + failed=has_errors, + record_count=result.record_count, + page_count=result.page_count, + read_to_end=not result.test_read_limit_reached and not http_errors, + http_errors=http_errors, + problem_logs=problem_logs, + inferred_schema=result.inferred_schema, + warnings=result.warnings, + ) + + +ReadinessCheckResult = Literal["passed", "failed", "unknown"] +"""Outcome of a readiness check; `unknown` when the report can't tell.""" + + +class ConnectorBuilderReadinessCheck(BaseModel): + """One readiness check toward a finished custom source.""" + + description: str + result: ReadinessCheckResult + detail: str | None = None + + +class ConnectorBuilderFullSync(BaseModel): + """A full sync started for a published Connector Builder project.""" + + source_id: str + destination_id: str + destination_created: bool + connection_id: str + connection_created: bool + job_id: int + job_url: str + + +class ConnectorBuilderReadinessReport(BaseModel): + """Readiness checks and per-stream test reads for a Connector Builder draft.""" + + failed_streams: list[str] = Field(default_factory=list) + pending_streams: list[str] = Field( + default_factory=list, + description=( + "Streams still reading when `timeout_seconds` ran out. Their results are dropped, " + "and Cloud may still finish those reads." + ), + ) + not_attempted_streams: list[str] = Field( + default_factory=list, + description="Streams not started before `timeout_seconds` ran out.", + ) + checks: list[ConnectorBuilderReadinessCheck] + next_steps: list[str] + streams: list[ConnectorBuilderStreamReadiness] + full_sync: ConnectorBuilderFullSync | None = Field( + default=None, + description="The full sync started with `run_full_sync=True`, if any.", + ) diff --git a/airbyte/cloud/connector_builder/_util.py b/airbyte/cloud/connector_builder/_util.py new file mode 100644 index 000000000..ee772d2e3 --- /dev/null +++ b/airbyte/cloud/connector_builder/_util.py @@ -0,0 +1,297 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Internal helpers for Connector Builder projects: manifest text, specs, and readiness.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import yaml + +from airbyte import exceptions as exc +from airbyte.cloud.connector_builder._models import ( + ConnectorBuilderReadinessCheck, + ConnectorBuilderStreamReadiness, + ManifestTextEditMode, + ReadinessCheckResult, +) + + +_MANIFEST_KEYS = frozenset({"streams", "definitions", "check"}) +_YAML_LINE_WIDTH = 2**31 - 1 +_SECRETS_MASK = "**********" +"""Placeholder Cloud returns instead of a stored secret value.""" +CONNECTOR_BUILDER_TEST_DESTINATION_NAME = "Connector Builder Test (No-op)" +"""Name of the no-op destination that full syncs find or create by default.""" +TEST_READ_MAX_RECORDS = 5_000 +"""Most records read per stream by `_run_readiness_checks()`.""" +TEST_READ_MAX_PAGES = 20 +TEST_READ_MAX_SLICES = 20 + + +def unwrap_builder_manifest(value: object) -> dict[str, Any] | None: + """Return the connector manifest from a builder manifest payload. + + Some builder projects store their draft as a `DeclarativeSourceManifest` + envelope (`{"manifest": {...}, "spec": ..., "version": ...}`) rather than + the bare manifest; unwrap that envelope so the result matches the shape of + `manifest`. + """ + if not isinstance(value, dict): + return None + nested = value.get("manifest") + if isinstance(nested, dict) and not any(key in value for key in _MANIFEST_KEYS): + return nested + return value + + +def render_manifest_text(manifest: dict[str, Any]) -> str: + """Render a manifest as deterministic YAML text, without line wrapping. + + Airbyte Cloud stores the parsed manifest, not its original text, so comments and + formatting are not preserved. Text edits match against this rendering. + """ + return yaml.safe_dump( + manifest, + sort_keys=False, + allow_unicode=True, + width=_YAML_LINE_WIDTH, + ) + + +def apply_manifest_text_edit( + current_text: str, + *, + edit_mode: ManifestTextEditMode, + new_text: str, + old_text: str | None = None, +) -> str: + """Return `current_text` after applying one text edit. + + Args: + current_text: The text to edit. + edit_mode: See `ManifestTextEditMode`. + new_text: The whole new text for `replace_full`, else the replacement for `old_text`. + old_text: The exact text to replace. Required unless `edit_mode` is `replace_full`. + """ + if edit_mode == "replace_full": + if old_text is not None: + raise exc.AirbyteLibInputError( + message="`old_text` is not allowed when `edit_mode` is `replace_full`.", + ) + return new_text + + if not old_text: + raise exc.AirbyteLibInputError( + message=f"`old_text` is required when `edit_mode` is `{edit_mode}`.", + ) + match_count = current_text.count(old_text) + if match_count == 0: + raise exc.AirbyteLibInputError( + message="`old_text` was not found in the draft manifest text.", + guidance="Copy `old_text` verbatim from the latest draft manifest text.", + ) + if edit_mode == "replace_text" and match_count > 1: + raise exc.AirbyteLibInputError( + message=f"`old_text` matches {match_count} places in the draft manifest text.", + guidance=( + "Include more surrounding text in `old_text` so it matches once, " + "or use `replace_all` to replace every match." + ), + context={"match_count": match_count}, + ) + return current_text.replace(old_text, new_text) + + +def parse_manifest_yaml(manifest_yaml: dict[str, Any] | Path | str) -> dict[str, Any]: + """Return a manifest dict from a dict, a YAML file path, or YAML text.""" + if isinstance(manifest_yaml, Path): + manifest_yaml = manifest_yaml.read_text() + if isinstance(manifest_yaml, str): + manifest_yaml = yaml.safe_load(manifest_yaml) + if not isinstance(manifest_yaml, dict): + raise exc.AirbyteLibInputError(message="The manifest must be a YAML mapping.") + return manifest_yaml + + +def builder_testing_values_spec(manifest: dict[str, Any]) -> dict[str, Any]: + """Build the `spec` object that the testing-values API expects, as the Builder UI does.""" + manifest_spec = manifest.get("spec") or {} + return { + "documentationUrl": manifest_spec.get("documentation_url"), + "connectionSpecification": manifest_spec.get("connection_specification") or {}, + "advancedAuth": manifest_spec.get("advanced_auth"), + } + + +def publish_spec(manifest: dict[str, Any]) -> dict[str, Any]: + """Build the connector spec for publishing, with the API's camelCase field names.""" + return {**(manifest.get("spec") or {}), **builder_testing_values_spec(manifest)} + + +def draft_from_project_data(project_data: dict[str, Any]) -> dict[str, Any] | None: + if not project_data["builderProject"].get("hasDraft"): + return None + declarative_manifest = project_data.get("declarativeManifest") or {} + return unwrap_builder_manifest(declarative_manifest.get("manifest")) + + +def _result(passed: bool) -> ReadinessCheckResult: # noqa: FBT001 + return "passed" if passed else "failed" + + +def readiness_next_steps( + checks: list[ConnectorBuilderReadinessCheck], + streams: list[ConnectorBuilderStreamReadiness], + pending_streams: list[str], + not_attempted_streams: list[str], +) -> list[str]: + steps = [check.description for check in checks[:4] if check.result != "passed"] + failed = [stream.stream_name for stream in streams if stream.failed] + capped = [ + stream.stream_name + for stream in streams + if stream.error is None and not stream.http_errors and not stream.read_to_end + ] + if failed: + steps.append(f"Fix the failed streams and test them again: {', '.join(failed)}") + if pending_streams: + steps.append( + "These streams were still reading at the timeout, so their results were dropped. " + "Retry them with a longer `timeout_seconds`, or check them with the full sync: " + f"{', '.join(pending_streams)}" + ) + if not_attempted_streams and not failed: + steps.append( + "Run the report again with the not-attempted streams, skipping any already " + f"validated or unchanged since: {', '.join(not_attempted_streams)}" + ) + if capped: + steps.append( + "These streams hit a test-read cap, so later pages weren't read; " + f"check them with a full sync: {', '.join(capped)}" + ) + warning_streams = [stream.stream_name for stream in streams if stream.warnings] + if warning_streams: + steps.append(f"Report these streams' `warnings` to the user: {', '.join(warning_streams)}") + return steps + + +def readiness_checks( + project_data: dict[str, Any], + draft: dict[str, Any], + streams: list[ConnectorBuilderStreamReadiness], + pending_streams: list[str], + url: str, +) -> list[ConnectorBuilderReadinessCheck]: + spec = (draft.get("spec") or {}).get("connection_specification") + takes_no_inputs = isinstance(spec, dict) and not spec.get("properties") + testing_values = project_data.get("testingValues") or {} + source_definition_id = project_data["builderProject"].get("sourceDefinitionId") + failed = [stream.stream_name for stream in streams if stream.failed] + capped = [ + stream.stream_name + for stream in streams + if stream.error is None and not stream.http_errors and not stream.read_to_end + ] + not_read = [ + stream.stream_name for stream in streams if stream.error is not None or stream.http_errors + ] + return [ + ConnectorBuilderReadinessCheck( + description="Save a draft manifest", + result=_result(bool(project_data["builderProject"].get("hasDraft"))), + ), + ConnectorBuilderReadinessCheck( + description="Define the connection spec (user inputs)", + result=_result(isinstance(spec, dict)), + ), + ConnectorBuilderReadinessCheck( + description="Define at least one stream", + result=_result(bool(draft.get("streams")) or bool(draft.get("dynamic_streams"))), + ), + ConnectorBuilderReadinessCheck( + description="Set testing values (enter secrets in the Connector Builder UI)", + result=_result(bool(testing_values) or takes_no_inputs), + detail=url, + ), + ConnectorBuilderReadinessCheck( + description="Test-read the named streams without errors", + result=_result(not failed and not pending_streams), + detail="; ".join( + part + for part in ( + f"failed: {', '.join(failed)}" if failed else "", + f"pending: {', '.join(pending_streams)}" if pending_streams else "", + ) + if part + ) + or None, + ), + ConnectorBuilderReadinessCheck( + description="Read the named streams to the end within the test-read caps", + result=_result(not pending_streams and all(stream.read_to_end for stream in streams)), + detail="; ".join( + part + for part in ( + ( + f"hit a cap: {', '.join(capped)}. Check these with a full sync " + "after publishing." + ) + if capped + else "", + f"not read: {', '.join(not_read)} (test read failed)" if not_read else "", + ) + if part + ) + or None, + ), + ConnectorBuilderReadinessCheck( + description="Publish the custom source definition", + result=_result(source_definition_id is not None), + detail=source_definition_id, + ), + ConnectorBuilderReadinessCheck( + description="Run a full sync of every stream to a no-op destination", + result="unknown", + detail="Not tracked by this report. Run it after publishing.", + ), + ] + + +def drop_masked_secrets(value: Any) -> Any: # noqa: ANN401 + """Return a copy of `value` without the masked secret fields Cloud returns.""" + if isinstance(value, dict): + return { + key: drop_masked_secrets(nested) + for key, nested in value.items() + if nested != _SECRETS_MASK + } + if isinstance(value, list): + return [drop_masked_secrets(nested) for nested in value] + return value + + +def spec_missing_required( + schema: Any, # noqa: ANN401 + config: dict[str, Any], +) -> bool: + """Return whether `config` lacks a required top-level field that has no default.""" + if not isinstance(schema, dict): + return False + properties = schema.get("properties") or {} + return any( + name not in config and "default" not in (properties.get(name) or {}) + for name in schema.get("required") or [] + ) + + +def spec_has_secrets(schema: Any) -> bool: # noqa: ANN401 + """Return whether a connection spec marks any field `airbyte_secret`.""" + if isinstance(schema, dict): + return schema.get("airbyte_secret") is True or any( + spec_has_secrets(nested) for nested in schema.values() + ) + if isinstance(schema, list): + return any(spec_has_secrets(nested) for nested in schema) + return False diff --git a/airbyte/cloud/connector_builder/project.py b/airbyte/cloud/connector_builder/project.py new file mode 100644 index 000000000..227900b7b --- /dev/null +++ b/airbyte/cloud/connector_builder/project.py @@ -0,0 +1,929 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""The `CloudConnectorBuilderProject` class.""" + +from __future__ import annotations + +import difflib +import time +from concurrent.futures import ThreadPoolExecutor, wait +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any, Literal, overload + +import requests + +from airbyte import exceptions as exc +from airbyte._util import api_util +from airbyte._util.api_util import get_web_url_root +from airbyte.cloud.connector_builder._models import ( + ConnectorBuilderFullSync, + ConnectorBuilderProjectInfo, + ConnectorBuilderPublishedVersion, + ConnectorBuilderReadinessReport, + ConnectorBuilderResolvedManifest, + ConnectorBuilderStreamReadiness, + ConnectorBuilderStreamReadResult, + ManifestTextEditMode, +) +from airbyte.cloud.connector_builder._util import ( + CONNECTOR_BUILDER_TEST_DESTINATION_NAME, + TEST_READ_MAX_PAGES, + TEST_READ_MAX_RECORDS, + TEST_READ_MAX_SLICES, + apply_manifest_text_edit, + builder_testing_values_spec, + draft_from_project_data, + drop_masked_secrets, + parse_manifest_yaml, + publish_spec, + readiness_checks, + readiness_next_steps, + render_manifest_text, + spec_has_secrets, + spec_missing_required, + unwrap_builder_manifest, +) + + +if TYPE_CHECKING: + from collections.abc import Callable + from pathlib import Path + + from airbyte.cloud.connectors import CloudSource, CustomCloudSourceDefinition + from airbyte.cloud.workspaces import CloudWorkspace + + +_READINESS_MAX_WORKERS = 5 +"""Most streams `_run_readiness_checks()` reads at once.""" +_MIN_READ_TIMEOUT = 1.0 +"""Shortest HTTP timeout for a test read started just before the readiness deadline.""" + + +def _default_version_description(action: str) -> str: + """Return a generic release note with the current UTC time.""" + return f"{action} by PyAirbyte at {datetime.now(UTC):%Y-%m-%dT%H:%M:%SZ}" + + +def _check_test_read_limits( + record_limit: int, + page_limit: int, + slice_limit: int, +) -> None: + """Reject test-read limits outside Airbyte Cloud's supported range.""" + limits = ( + ("record_limit", record_limit, TEST_READ_MAX_RECORDS), + ("page_limit", page_limit, TEST_READ_MAX_PAGES), + ("slice_limit", slice_limit, TEST_READ_MAX_SLICES), + ) + if any(not 1 <= value <= maximum for _, value, maximum in limits): + raise exc.AirbyteLibInputError( + message="Test-read limits are out of range.", + guidance=( + f"Airbyte Cloud accepts record_limit 1 to {TEST_READ_MAX_RECORDS:,}, " + f"page_limit 1 to {TEST_READ_MAX_PAGES}, and " + f"slice_limit 1 to {TEST_READ_MAX_SLICES}." + ), + context={name: value for name, value, _ in limits}, + ) + + +class CloudConnectorBuilderProject: + """A Connector Builder project in an Airbyte Cloud workspace.""" + + def __init__( + self, + *, + workspace: CloudWorkspace, + builder_project_id: str, + owner_workspace_id: str | None = None, + ) -> None: + """Initialize a Connector Builder project object. Does not call the API. + + Args: + workspace: The workspace whose credentials are used. + builder_project_id: The Connector Builder project ID. + owner_workspace_id: The workspace that owns the project, if it differs from + `workspace` (for example, a definition shared from another workspace). + """ + self.workspace = workspace + self.builder_project_id = builder_project_id + self.owner_workspace_id = owner_workspace_id or workspace.workspace_id + + @property + def connector_builder_project_url(self) -> str: + """The project's Connector Builder URL, where users enter testing values and secrets.""" + return ( + f"{get_web_url_root(self.workspace.api_root)}/workspaces/" + f"{self.owner_workspace_id}/connector-builder/edit/{self.builder_project_id}" + ) + + def _get_project_data(self) -> dict[str, Any]: + """Fetch the project metadata and draft manifest from the API.""" + return api_util.get_connector_builder_project( + workspace_id=self.owner_workspace_id, + builder_project_id=self.builder_project_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, + ) + + @property + def name(self) -> str: + """The project name.""" + return self._get_project_data()["builderProject"]["name"] + + @property + def source_definition_id(self) -> str | None: + """The ID of the custom source definition, or `None` if never published.""" + return self._get_project_data()["builderProject"].get("sourceDefinitionId") + + def _published_manifest( + self, + source_definition_id: str | None, + ) -> dict[str, Any] | None: + if source_definition_id is None: + return None + return self.workspace.get_custom_source_definition( + source_definition_id, + definition_type="yaml", + ).manifest + + @property + def published_manifest_dict(self) -> dict[str, Any] | None: + """The active manifest of the published definition, or `None` if never published.""" + return self._published_manifest(self.source_definition_id) + + def _require_source_definition_id(self) -> str: + source_definition_id = self.source_definition_id + if source_definition_id is None: + raise exc.AirbyteLibInputError( + message="This Connector Builder project has never been published.", + guidance="Publish the project first.", + context={"builder_project_id": self.builder_project_id}, + ) + return source_definition_id + + def list_published_versions(self) -> list[ConnectorBuilderPublishedVersion]: + """List the published manifest versions, oldest first. Empty if never published.""" + source_definition_id = self.source_definition_id + if source_definition_id is None: + return [] + versions = api_util.list_declarative_manifest_versions( + workspace_id=self.owner_workspace_id, + source_definition_id=source_definition_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, + ) + return sorted( + ( + ConnectorBuilderPublishedVersion( + version=item["version"], + is_active=item["isActive"], + description=item.get("description") or "", + ) + for item in versions + ), + key=lambda item: item.version, + ) + + def get_published_manifest_dict( + self, + version: int | None = None, + ) -> dict[str, Any] | None: + """Return a published manifest version, or `None` if never published. + + Args: + version: The published version to return. Defaults to the active version. + """ + if version is None: + return self.published_manifest_dict + self._require_source_definition_id() + project_data = api_util.get_connector_builder_project( + workspace_id=self.owner_workspace_id, + builder_project_id=self.builder_project_id, + version=version, + 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, + ) + declarative_manifest = project_data.get("declarativeManifest") or {} + return unwrap_builder_manifest(declarative_manifest.get("manifest")) + + def get_published_manifest_yaml( + self, + version: int | None = None, + ) -> str | None: + """Return a published manifest version as YAML text, or `None` if never published. + + Args: + version: The published version to return. Defaults to the active version. + """ + manifest = self.get_published_manifest_dict(version) + return render_manifest_text(manifest) if manifest is not None else None + + def revert_to_published_version( + self, + version: int, + *, + as_draft: bool = True, + description: str | None = None, + ) -> None: + """Revert the project to a published manifest version. + + With `as_draft=True`, the version's manifest replaces the draft (overwriting any + unpublished edits) and the active published version doesn't change. With + `as_draft=False`, the version's manifest is published as a new active version. That + fails if the draft has unpublished edits, because publishing deletes the draft. + + Args: + version: The published version to revert to. + as_draft: Whether to save the version as the draft instead of publishing it. + description: Release note for the new version when `as_draft=False`. Defaults + to a generic note with the reverted version and the UTC publish time. + """ + self._require_source_definition_id() + manifest = self.get_published_manifest_dict(version) + if manifest is None: + raise exc.AirbyteLibInputError( + message="The published version has no manifest.", + context={"builder_project_id": self.builder_project_id, "version": version}, + ) + if as_draft: + self.replace_draft_manifest(manifest, pre_validate=False) + return + + draft_manifest = self.draft_manifest_dict + if draft_manifest is not None and draft_manifest != self.published_manifest_dict: + raise exc.AirbyteLibInputError( + message="This Connector Builder project has unpublished draft edits.", + guidance=( + "Revert with `as_draft=True` to review the version as the draft, " + "or publish the draft first." + ), + context={"builder_project_id": self.builder_project_id, "version": version}, + ) + self._publish_manifest( + manifest, + name=self.name, + description=description + or _default_version_description(f"Reverted to version {version}"), + pre_validate=False, + ) + + def _get_info(self) -> ConnectorBuilderProjectInfo: + """Read the project once and return its IDs, URL, draft, and published manifest.""" + project_data = self._get_project_data() + draft_manifest = draft_from_project_data(project_data) + source_definition_id = project_data["builderProject"].get("sourceDefinitionId") + return ConnectorBuilderProjectInfo( + builder_project_id=self.builder_project_id, + definition_id=source_definition_id, + connector_builder_project_url=self.connector_builder_project_url, + name=project_data["builderProject"]["name"], + has_draft=draft_manifest is not None, + draft_manifest_dict=draft_manifest, + draft_manifest_yaml=( + render_manifest_text(draft_manifest) if draft_manifest is not None else None + ), + published_manifest_dict=self._published_manifest(source_definition_id), + ) + + @property + def draft_manifest_dict(self) -> dict[str, Any] | None: + """The unpublished draft manifest, or `None` if the project has no draft.""" + return draft_from_project_data(self._get_project_data()) + + def _require_draft_manifest( + self, + project_data: dict[str, Any] | None = None, + ) -> dict[str, Any]: + manifest = draft_from_project_data(project_data or self._get_project_data()) + if manifest is None: + raise exc.AirbyteLibInputError( + message="This Connector Builder project has no draft manifest.", + guidance="Save a draft manifest to the project first.", + context={"builder_project_id": self.builder_project_id}, + ) + return manifest + + def replace_draft_manifest( + self, + manifest_yaml: dict[str, Any] | Path | str, + *, + pre_validate: bool = True, + ) -> None: + """Replace the project's draft manifest. Does not publish anything. + + Args: + manifest_yaml: The manifest, as a dict, a YAML file path, or YAML text. + pre_validate: Whether to validate the manifest client-side before saving. + """ + manifest = parse_manifest_yaml(manifest_yaml) + if pre_validate: + api_util.validate_yaml_manifest(manifest, raise_on_error=True) + + api_util.update_connector_builder_project( + workspace_id=self.owner_workspace_id, + builder_project_id=self.builder_project_id, + name=self.name, + draft_manifest=manifest, + components_file_content=None, + 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, + ) + + @property + def draft_manifest_yaml(self) -> str | None: + """The draft manifest rendered by `render_manifest_text()`, or `None` if no draft.""" + manifest = self.draft_manifest_dict + return render_manifest_text(manifest) if manifest is not None else None + + def _update_draft_manifest_yaml( + self, + *, + edit_mode: ManifestTextEditMode, + new_text: str, + old_text: str | None = None, + pre_validate: bool = True, + ) -> str: + """Edit the draft manifest as text and save it. Does not publish anything. + + Edits match against `draft_manifest_yaml`. Returns a unified diff from the previous + rendering to the rendering of the saved draft. + + Args: + edit_mode: See `ManifestTextEditMode`. + new_text: The whole new manifest for `replace_full`, else the replacement text. + old_text: The exact text to replace. Required unless `edit_mode` is `replace_full`. + pre_validate: Whether to validate the edited manifest client-side before saving. + """ + current_text = self.draft_manifest_yaml + if current_text is None and edit_mode != "replace_full": + raise exc.AirbyteLibInputError( + message="This Connector Builder project has no draft manifest to edit.", + guidance="Use `replace_full` to save a whole manifest first.", + context={"builder_project_id": self.builder_project_id}, + ) + edited_text = apply_manifest_text_edit( + current_text or "", + edit_mode=edit_mode, + new_text=new_text, + old_text=old_text, + ) + self.replace_draft_manifest(edited_text, pre_validate=pre_validate) + + saved_text = self.draft_manifest_yaml or "" + return "".join( + difflib.unified_diff( + (current_text or "").splitlines(keepends=True), + saved_text.splitlines(keepends=True), + fromfile="previous", + tofile="saved", + ) + ) + + def replace_testing_values( + self, + testing_values: dict[str, Any], + ) -> None: + """Replace the project's testing values, which Cloud stores in its secret storage. + + The spec comes from the draft manifest, or the published one when there's no draft. + Fields marked `airbyte_secret` in the spec are stored as secrets. + """ + project_data = self._get_project_data() + manifest = draft_from_project_data(project_data) or self._published_manifest( + project_data["builderProject"].get("sourceDefinitionId") + ) + if manifest is None: + manifest = self._require_draft_manifest(project_data) + api_util.update_connector_builder_project_testing_values( + workspace_id=self.owner_workspace_id, + builder_project_id=self.builder_project_id, + testing_values=testing_values, + spec=builder_testing_values_spec(manifest), + 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, + ) + + def _testing_values_connection_spec(self) -> dict[str, Any]: + """Return the JSON schema of the testing values, from the draft or published spec.""" + project_data = self._get_project_data() + manifest = draft_from_project_data(project_data) or self._published_manifest( + project_data["builderProject"].get("sourceDefinitionId") + ) + if manifest is None: + manifest = self._require_draft_manifest(project_data) + return builder_testing_values_spec(manifest).get("connectionSpecification") or {} + + def _merge_testing_values( + self, + testing_values: dict[str, Any], + ) -> None: + """Merge `testing_values` onto the saved ones, keeping secrets stored on Cloud. + + Cloud returns saved secrets masked; sending the mask back keeps the stored secret. + """ + saved = self._get_project_data().get("testingValues") or {} + self.replace_testing_values({**saved, **testing_values}) + + def run_stream_test_read( + self, + stream_name: str, + *, + record_limit: int = 100, + page_limit: int = 5, + slice_limit: int = 5, + ) -> ConnectorBuilderStreamReadResult: + """Test-read one stream of the draft manifest on Airbyte Cloud. + + Uses the project's saved testing values as the connector config. + + Args: + stream_name: The stream to read. + record_limit: Maximum records to return (1 to 5,000). + page_limit: Maximum pages to fetch per slice (1 to 20). + slice_limit: Maximum slices to read (1 to 20). + """ + _check_test_read_limits(record_limit, page_limit, slice_limit) + return self._run_stream_test_read( + self._require_draft_manifest(), + stream_name, + record_limit=record_limit, + page_limit=page_limit, + slice_limit=slice_limit, + ) + + def _run_stream_test_read( + self, + manifest: dict[str, Any], + stream_name: str, + *, + record_limit: int, + page_limit: int, + slice_limit: int, + timeout_seconds: float | None = None, + ) -> ConnectorBuilderStreamReadResult: + response = api_util.read_connector_builder_project_stream( + workspace_id=self.owner_workspace_id, + builder_project_id=self.builder_project_id, + manifest=manifest, + stream_name=stream_name, + record_limit=record_limit, + page_limit=page_limit, + slice_limit=slice_limit, + 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, + timeout_seconds=timeout_seconds, + ) + return ConnectorBuilderStreamReadResult._from_api_response(stream_name, response) # noqa: SLF001 + + def get_resolved_manifest( + self, + *, + stream_limit: int = 5, + ) -> ConnectorBuilderResolvedManifest: + """Resolve the draft manifest on Airbyte Cloud, including dynamic streams. + + Uses the project's saved testing values to generate dynamic streams. + + Args: + stream_limit: Maximum generated streams to return per dynamic stream (1 to 100). + """ + resolved_manifest = api_util.full_resolve_connector_builder_project_manifest( + workspace_id=self.owner_workspace_id, + builder_project_id=self.builder_project_id, + manifest=self._require_draft_manifest(), + stream_limit=stream_limit, + 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, + ) + return ConnectorBuilderResolvedManifest._from_manifest(resolved_manifest) # noqa: SLF001 + + def _run_readiness_checks( + self, + stream_names: list[str], + *, + record_limit: int = TEST_READ_MAX_RECORDS, + page_limit: int = TEST_READ_MAX_PAGES, + slice_limit: int = TEST_READ_MAX_SLICES, + timeout_seconds: float | None = None, + run_full_sync: bool = False, + test_source_id: str | None = None, + on_resource_created: Callable[[str], None] | None = None, + ) -> ConnectorBuilderReadinessReport: + """Check the draft's progress and test-read the given streams on Airbyte Cloud. + + Reads run in parallel. A stream is `failed` when its read errored, returned HTTP + errors, or logged errors. The report lists failures and warnings without judging + readiness; `read_to_end` indicates whether the read ended within the test-read caps. + The readiness checks are derived from the project, so nothing is stored. A full sync + after publishing validates the whole connector. + + Args: + stream_names: Streams to test-read (at least 1). For + dynamic streams, get the generated names from `get_resolved_manifest()`. + record_limit: Maximum records to read per stream (1 to 5,000). + page_limit: Maximum pages to fetch per slice (1 to 20). + slice_limit: Maximum slices to read (1 to 20). + timeout_seconds: Return after this many seconds with the streams finished so far; + the rest are listed in `pending_streams` and `not_attempted_streams`. No limit + by default. + run_full_sync: If no named stream failed or is still pending (or none are named), + publish the draft and start a full sync with `_run_full_sync()`. + test_source_id: The source for the full sync. See `_run_full_sync()`. + on_resource_created: Passed to `_run_full_sync()`. + """ + _check_test_read_limits(record_limit, page_limit, slice_limit) + stream_names = list(dict.fromkeys(stream_names)) + if not stream_names and not run_full_sync: + raise exc.AirbyteLibInputError( + message="Pass at least 1 stream name, or set `run_full_sync`.", + ) + project_data = self._get_project_data() + if not stream_names: + # A full sync alone can run on the published manifest; the first publish clears + # the draft. + manifest = draft_from_project_data(project_data) or self._published_manifest( + project_data["builderProject"].get("sourceDefinitionId") + ) + if manifest is None: + manifest = self._require_draft_manifest(project_data) + return self._readiness_report(project_data, manifest, [], [], []).model_copy( + update={ + "full_sync": self._run_full_sync( + test_source_id=test_source_id, + on_resource_created=on_resource_created, + ), + }, + ) + + draft = self._require_draft_manifest(project_data) + deadline = None if timeout_seconds is None else time.monotonic() + timeout_seconds + + def read(stream_name: str) -> ConnectorBuilderStreamReadiness: + # Bound each HTTP request by the deadline so worker threads end with the call. + request_timeout = ( + None if deadline is None else max(deadline - time.monotonic(), _MIN_READ_TIMEOUT) + ) + try: + result = self._run_stream_test_read( + draft, + stream_name, + record_limit=record_limit, + page_limit=page_limit, + slice_limit=slice_limit, + timeout_seconds=request_timeout, + ) + except exc.AirbyteLibError as ex: + error = ex.get_message() + if ex.guidance: + error += f" {ex.guidance}" + return ConnectorBuilderStreamReadiness( + stream_name=stream_name, failed=True, error=error + ) + except requests.RequestException as ex: + return ConnectorBuilderStreamReadiness( + stream_name=stream_name, + failed=True, + error=f"Request to Airbyte Cloud failed: {type(ex).__name__}", + ) + return ConnectorBuilderStreamReadiness._from_stream_read(result) # noqa: SLF001 + + pool = ThreadPoolExecutor(max_workers=min(_READINESS_MAX_WORKERS, len(stream_names))) + try: + futures = [pool.submit(read, name) for name in stream_names] + done, _ = wait(futures, timeout=timeout_seconds) + finally: + pool.shutdown(wait=False, cancel_futures=True) + streams = [future.result() for future in futures if future in done] + unfinished = [ + (name, future) + for name, future in zip(stream_names, futures, strict=True) + if future not in done + ] + pending_streams = [name for name, future in unfinished if not future.cancelled()] + not_attempted_streams = [name for name, future in unfinished if future.cancelled()] + report = self._readiness_report( + project_data, draft, streams, pending_streams, not_attempted_streams + ) + if run_full_sync and not ( + report.failed_streams or report.pending_streams or report.not_attempted_streams + ): + report.full_sync = self._run_full_sync( + test_source_id=test_source_id, + on_resource_created=on_resource_created, + ) + return report + + def _readiness_report( + self, + project_data: dict[str, Any], + draft: dict[str, Any], + streams: list[ConnectorBuilderStreamReadiness], + pending_streams: list[str], + not_attempted_streams: list[str], + ) -> ConnectorBuilderReadinessReport: + remaining = pending_streams + not_attempted_streams + checks = readiness_checks( + project_data, draft, streams, remaining, self.connector_builder_project_url + ) + return ConnectorBuilderReadinessReport( + failed_streams=[stream.stream_name for stream in streams if stream.failed], + pending_streams=pending_streams, + not_attempted_streams=not_attempted_streams, + checks=checks, + next_steps=readiness_next_steps( + checks, streams, pending_streams, not_attempted_streams + ), + streams=streams, + ) + + def publish( + self, + name: str, + *, + pre_validate: bool = True, + description: str | None = None, + ) -> CustomCloudSourceDefinition: + """Publish the draft manifest as a custom source definition. + + The first publish creates the definition with `name`. If the project is already + published, the draft becomes a new active version of the existing definition. + + Args: + name: The definition name, used only by the first publish. + pre_validate: Whether to validate the manifest client-side before publishing. + description: Release note for the new version. Defaults to a generic note with + the UTC publish time. + """ + return self._publish_manifest( + self._require_draft_manifest(), + name=name, + description=description or _default_version_description("Published"), + pre_validate=pre_validate, + ) + + def _publish_manifest( + self, + manifest: dict[str, Any], + *, + name: str, + description: str, + pre_validate: bool, + ) -> CustomCloudSourceDefinition: + """Publish `manifest` as version 1 or as the next active version of the definition.""" + if pre_validate: + api_util.validate_yaml_manifest(manifest, raise_on_error=True) + + source_definition_id = self.source_definition_id + if source_definition_id is None: + source_definition_id = api_util.publish_connector_builder_project( + workspace_id=self.owner_workspace_id, + builder_project_id=self.builder_project_id, + name=name, + manifest=manifest, + spec=publish_spec(manifest), + description=description, + 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, + ) + else: + next_version = ( + max( + (item.version for item in self.list_published_versions()), + default=0, + ) + + 1 + ) + api_util.create_declarative_manifest_version( + workspace_id=self.owner_workspace_id, + source_definition_id=source_definition_id, + manifest=manifest, + spec=publish_spec(manifest), + version=next_version, + description=description, + 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, + ) + return self.workspace.get_custom_source_definition( + source_definition_id, + definition_type="yaml", + ) + + def _get_default_test_source_name(self) -> str: + """Default name of the source `_get_or_create_default_test_source()` finds or creates.""" + return f"{self.name} (Connector Builder test)" + + @overload + def _get_or_create_default_test_source( + self, + *, + allow_create: Literal[True] = True, + default_source_name: str | Callable[[CloudConnectorBuilderProject], str] | None = None, + ) -> CloudSource: ... + + @overload + def _get_or_create_default_test_source( + self, + *, + allow_create: bool, + default_source_name: str | Callable[[CloudConnectorBuilderProject], str] | None = None, + ) -> CloudSource | None: ... + + def _get_or_create_default_test_source( + self, + *, + allow_create: bool = True, + default_source_name: str | Callable[[CloudConnectorBuilderProject], str] | None = None, + ) -> CloudSource | None: + """Return the project's test source on the published definition, creating it if needed. + + A new source gets the project's non-secret testing values. When the spec has secret + fields, the source is a draft (deferred credentials): a person enters the secrets at + its `connector_url`, then tests and saves. Publish the project first. + + Args: + allow_create: Whether to create the source when it doesn't exist. When False, + returns `None` instead. + default_source_name: The source name, or a callable that takes this project and + returns it. Defaults to " (Connector Builder test)". + """ + source_name = ( + default_source_name(self) if callable(default_source_name) else default_source_name + ) or self._get_default_test_source_name() + project_data = self._get_project_data() + source_definition_id = project_data["builderProject"].get("sourceDefinitionId") + if source_definition_id is None: + if not allow_create: + return None + raise exc.AirbyteLibInputError( + message="Publish the Connector Builder project before creating a test source.", + ) + existing = next( + ( + source + for source in self.workspace.list_sources(name=source_name) + if source.definition_id == source_definition_id + ), + None, + ) + if existing is not None or not allow_create: + return existing + # The source runs the published version, so its spec decides the setup. + manifest = self._published_manifest(source_definition_id) + spec = ((manifest or {}).get("spec") or {}).get("connection_specification") + config = drop_masked_secrets(project_data.get("testingValues") or {}) + return self.workspace.deploy_source( + name=source_name, + source=config, + definition_id=source_definition_id, + defer_credentials=spec_has_secrets(spec) or spec_missing_required(spec, config), + ) + + def _run_full_sync( + self, + *, + test_source_id: str | None = None, + pre_validate: bool = True, + on_resource_created: Callable[[str], None] | None = None, + ) -> ConnectorBuilderFullSync: + """Publish the draft and start a full sync of every stream to a no-op destination. + + Without a draft, syncs the published version. Selects every stream the test source + discovers, so dynamic streams follow the source's own credentials. Returns without + waiting for the sync to finish. + + Args: + test_source_id: Source on the published definition to sync from. Defaults to the + source from `_get_or_create_default_test_source()`. + pre_validate: Whether to validate the manifest client-side before publishing. + on_resource_created: Called with the ID of each destination or connection this + call creates, as soon as it exists. + """ + project_data = self._get_project_data() + has_draft = draft_from_project_data(project_data) is not None + source_definition_id = project_data["builderProject"].get("sourceDefinitionId") + if not has_draft and source_definition_id is None: + self._require_draft_manifest(project_data) + # Check the source and destination first: the first publish clears the draft. + if test_source_id is None: + test_source = self._get_or_create_default_test_source(allow_create=False) + if test_source is None: + raise exc.AirbyteLibInputError( + message="No test source found for this Connector Builder project.", + guidance=( + "Pass `test_source_id`, or create a source on this project's published " + "definition and have the user complete its settings." + ), + ) + test_source_id = test_source.connector_id + elif self.workspace.get_source(test_source_id).definition_id != source_definition_id: + raise exc.AirbyteLibInputError( + message="`test_source_id` isn't a source of this project's published definition.", + context={"test_source_id": test_source_id, "definition_id": source_definition_id}, + ) + test_destination_id, destination_created = self._get_or_create_test_destination() + if destination_created and on_resource_created is not None: + on_resource_created(test_destination_id) + if has_draft: + self.publish(self.name, pre_validate=pre_validate) + + stream_names = api_util.get_source_stream_names( + test_source_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, + ) + if not stream_names: + raise exc.AirbyteLibInputError( + message="The test source discovered no streams.", + context={"test_source_id": test_source_id}, + ) + connection_name = f"{self.name} (Connector Builder full sync)" + connection = next( + ( + connection + for connection in self.workspace.list_connections(name=connection_name) + if connection.source_id == test_source_id + and connection.destination_id == test_destination_id + ), + None, + ) + connection_created = connection is None + if connection is None: + connection = self.workspace.deploy_connection( + connection_name, + source=test_source_id, + destination=test_destination_id, + selected_streams=stream_names, + ) + if on_resource_created is not None: + on_resource_created(connection.connection_id) + else: + connection.set_selected_streams(stream_names) + sync_result = connection.run_sync(wait=False) + return ConnectorBuilderFullSync( + source_id=test_source_id, + destination_id=test_destination_id, + destination_created=destination_created, + connection_id=connection.connection_id, + connection_created=connection_created, + job_id=sync_result.job_id, + job_url=sync_result.job_url, + ) + + def _get_or_create_test_destination(self) -> tuple[str, bool]: + """Return the test destination's ID and whether this call created it.""" + existing = self.workspace.list_destinations(name=CONNECTOR_BUILDER_TEST_DESTINATION_NAME) + noop = [ + destination + for destination in existing + if destination.definition_id == api_util.DEV_NULL_DESTINATION_DEFINITION_ID + ] + if noop: + return noop[0].connector_id, False + if existing: + raise exc.AirbyteLibInputError( + message=( + f"Destination `{CONNECTOR_BUILDER_TEST_DESTINATION_NAME}` exists " + "but isn't a no-op (dev-null) destination." + ), + guidance="Rename or delete that destination.", + context={"destination_ids": [destination.connector_id for destination in existing]}, + ) + from airbyte.destinations.util import get_noop_destination # noqa: PLC0415 + + return self.workspace.deploy_destination( + name=CONNECTOR_BUILDER_TEST_DESTINATION_NAME, + destination=get_noop_destination(), + ).connector_id, True + + def __repr__(self) -> str: + """Return a string representation of the project.""" + return f"CloudConnectorBuilderProject(builder_project_id={self.builder_project_id!r})" diff --git a/airbyte/cloud/connectors.py b/airbyte/cloud/connectors.py index c8eeee050..8db297888 100644 --- a/airbyte/cloud/connectors.py +++ b/airbyte/cloud/connectors.py @@ -88,6 +88,8 @@ ) from airbyte._util import api_util, text_util from airbyte._util.api_util import get_web_url_root +from airbyte.cloud.connector_builder import CloudConnectorBuilderProject +from airbyte.cloud.connector_builder._util import unwrap_builder_manifest from airbyte.cloud.models import ( CheckResult, CloudCustomSourceDefinitionInfo, @@ -1273,25 +1275,6 @@ def _from_destination_response( return result -_MANIFEST_KEYS = frozenset({"streams", "definitions", "check"}) - - -def _unwrap_builder_manifest(value: object) -> dict[str, Any] | None: - """Return the connector manifest from a builder manifest payload. - - Some builder projects store their draft as a `DeclarativeSourceManifest` - envelope (`{"manifest": {...}, "spec": ..., "version": ...}`) rather than - the bare manifest; unwrap that envelope so the result matches the shape of - `manifest`. - """ - if not isinstance(value, dict): - return None - nested = value.get("manifest") - if isinstance(nested, dict) and not any(key in value for key in _MANIFEST_KEYS): - return nested - return value - - class CustomCloudSourceDefinition: # noqa: PLR0904 """A custom source connector definition in Airbyte Cloud. @@ -1465,6 +1448,22 @@ def connector_builder_project_id(self) -> str | None: return self._connector_builder_project_id + @property + def _connector_builder_project(self) -> CloudConnectorBuilderProject | None: + """The Connector Builder project behind this definition. YAML connectors only.""" + if self.definition_type != "yaml": + return None + + project_id = self.connector_builder_project_id + if not project_id: + return None + + return CloudConnectorBuilderProject( + workspace=self.workspace, + builder_project_id=project_id, + owner_workspace_id=self._builder_project_workspace_id, + ) + @property def connector_builder_project_url(self) -> str | None: """Get the connector builder project URL. Only present for YAML connectors.""" @@ -1577,7 +1576,7 @@ def draft_manifest(self) -> dict[str, Any] | None: return None declarative_manifest = project_data.get("declarativeManifest", {}) - return _unwrap_builder_manifest(declarative_manifest.get("manifest")) + return unwrap_builder_manifest(declarative_manifest.get("manifest")) @property def definition_url(self) -> str: diff --git a/airbyte/cloud/workspaces.py b/airbyte/cloud/workspaces.py index 101e7d1eb..de241ca6e 100644 --- a/airbyte/cloud/workspaces.py +++ b/airbyte/cloud/workspaces.py @@ -39,11 +39,8 @@ import uuid from dataclasses import dataclass, field from functools import cached_property -from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, overload -import yaml - from airbyte import exceptions as exc from airbyte._direct_connectors import api_util as agents_api_util from airbyte._direct_connectors import connector_docs @@ -60,6 +57,8 @@ from airbyte.cloud._credentials import _AirbyteCredentials from airbyte.cloud.client_config import CloudClientConfig from airbyte.cloud.connections import CloudConnection +from airbyte.cloud.connector_builder import CloudConnectorBuilderProject +from airbyte.cloud.connector_builder._util import parse_manifest_yaml from airbyte.cloud.models import ( ConnectorFeature, ConnectorType, @@ -75,6 +74,7 @@ if TYPE_CHECKING: from collections.abc import Callable + from pathlib import Path from airbyte.cloud.organizations import CloudOrganization from airbyte.sources.base import Source @@ -1257,11 +1257,14 @@ def publish_custom_source_definition( unique: bool = True, pre_validate: bool = True, testing_values: dict[str, Any] | None = None, + builder_project_id: str | None = None, ) -> cloud_connectors.CustomCloudSourceDefinition: """Publish a custom source connector definition. You must specify EITHER manifest_yaml (for YAML connectors) OR both docker_image - and docker_tag (for Docker connectors), but not both. + and docker_tag (for Docker connectors), but not both. To publish an existing + Connector Builder project, pass `builder_project_id`; `manifest_yaml` is then optional + and replaces the project's draft before publishing. Args: name: Display name for the connector definition @@ -1274,6 +1277,9 @@ def publish_custom_source_definition( Connector Builder UI. If provided, these values are stored as the complete testing values object for the connector builder project (replaces any existing values), allowing immediate test read operations. + builder_project_id: Optional Connector Builder project to publish. If the project + is already published, its draft becomes a new active version of the existing + definition. Returns: CustomCloudSourceDefinition object representing the created definition @@ -1282,6 +1288,20 @@ def publish_custom_source_definition( AirbyteLibInputError: If both or neither of manifest_yaml and docker_image provided AirbyteDuplicateResourcesError: If unique=True and name already exists """ + if builder_project_id is not None: + if docker_image is not None: + raise exc.AirbyteLibInputError( + message="docker_image cannot be combined with builder_project_id.", + ) + return self._publish_connector_builder_project( + builder_project_id, + name=name, + manifest_yaml=manifest_yaml, + unique=unique, + pre_validate=pre_validate, + testing_values=testing_values, + ) + is_yaml = manifest_yaml is not None is_docker = docker_image is not None @@ -1314,18 +1334,12 @@ def publish_custom_source_definition( ) if is_yaml: - manifest_dict: dict[str, Any] - if isinstance(manifest_yaml, Path): - manifest_dict = yaml.safe_load(manifest_yaml.read_text()) - elif isinstance(manifest_yaml, str): - manifest_dict = yaml.safe_load(manifest_yaml) - elif manifest_yaml is not None: - manifest_dict = manifest_yaml - else: + if manifest_yaml is None: raise exc.AirbyteLibInputError( message="manifest_yaml is required for YAML connectors", context={"name": name}, ) + manifest_dict = parse_manifest_yaml(manifest_yaml) if pre_validate: api_util.validate_yaml_manifest(manifest_dict, raise_on_error=True) @@ -1354,6 +1368,104 @@ def publish_custom_source_definition( "Only YAML manifest-based custom sources are currently available." ) + def _publish_connector_builder_project( + self, + builder_project_id: str, + *, + name: str, + manifest_yaml: dict[str, Any] | Path | str | None, + unique: bool, + pre_validate: bool, + testing_values: dict[str, Any] | None, + ) -> cloud_connectors.CustomCloudSourceDefinition: + project = self.get_connector_builder_project(builder_project_id) + if unique and project.source_definition_id is None: + existing = self.list_custom_source_definitions(definition_type="yaml") + if any(d.name == name for d in existing): + raise exc.AirbyteDuplicateResourcesError( + resource_type="custom_source_definition", + resource_name=name, + ) + + if manifest_yaml is not None: + project.replace_draft_manifest(manifest_yaml, pre_validate=pre_validate) + + definition = project.publish(name, pre_validate=pre_validate) + if testing_values is not None: + definition.set_testing_values(testing_values) + return definition + + def create_connector_builder_project( + self, + name: str, + *, + manifest_yaml: dict[str, Any] | Path | str | None = None, + pre_validate: bool = True, + ) -> CloudConnectorBuilderProject: + """Create an unpublished Connector Builder project. + + Nothing appears in the workspace's source list until the project is published. + + Args: + name: Display name for the project. + manifest_yaml: Optional initial draft manifest (dict, YAML file path, or YAML text). + pre_validate: Whether to validate the manifest client-side before saving. + """ + manifest_dict: dict[str, Any] | None = None + if manifest_yaml is not None: + manifest_dict = parse_manifest_yaml(manifest_yaml) + if pre_validate: + api_util.validate_yaml_manifest(manifest_dict, raise_on_error=True) + + builder_project_id = api_util.create_connector_builder_project( + workspace_id=self.workspace_id, + name=name, + draft_manifest=manifest_dict, + 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 CloudConnectorBuilderProject(workspace=self, builder_project_id=builder_project_id) + + def get_connector_builder_project( + self, + builder_project_id: str | None = None, + *, + definition_id: str | None = None, + ) -> CloudConnectorBuilderProject: + """Get a Connector Builder project by its ID or by its published definition's ID. + + Args: + builder_project_id: The Connector Builder project ID. + definition_id: The ID of the custom source definition published from the project. + + Raises: + AirbyteLibInputError: If not exactly one ID is given, or the definition has no + Connector Builder project. + """ + if (builder_project_id is None) == (definition_id is None): + raise exc.AirbyteLibInputError( + message="Provide exactly one of `builder_project_id` or `definition_id`.", + ) + if builder_project_id is not None: + return CloudConnectorBuilderProject( + workspace=self, builder_project_id=builder_project_id + ) + + assert definition_id is not None + project = self.get_custom_source_definition( # noqa: SLF001 + definition_id, + definition_type="yaml", + )._connector_builder_project + if project is None: + raise exc.AirbyteLibInputError( + message="No Connector Builder project was found for this definition.", + context={"definition_id": definition_id}, + ) + return project + def list_custom_source_definitions( self, *, diff --git a/airbyte/mcp/__init__.py b/airbyte/mcp/__init__.py index f6536eed6..3834190de 100644 --- a/airbyte/mcp/__init__.py +++ b/airbyte/mcp/__init__.py @@ -414,11 +414,12 @@ """ # noqa: D415 -from airbyte.mcp import cloud, guidance, interactive, local, registry +from airbyte.mcp import cloud, connector_builder, guidance, interactive, local, registry __all__: list[str] = [ "cloud", + "connector_builder", "guidance", "interactive", "local", diff --git a/airbyte/mcp/_tool_utils.py b/airbyte/mcp/_tool_utils.py index f87bd88da..94f163b06 100644 --- a/airbyte/mcp/_tool_utils.py +++ b/airbyte/mcp/_tool_utils.py @@ -548,7 +548,7 @@ def _normalize_mcp_module(mcp_module: str) -> str: return mcp_module -_CLOUD_MCP_MODULES = frozenset({"cloud", "interactive"}) +_CLOUD_MCP_MODULES = frozenset({"cloud", "connector_builder", "interactive"}) """MCP tool modules backed by Airbyte Cloud APIs. Cloud-scoped policies (pipeline changes, read-only mode, destructive tools) diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 39c1d41b0..1723055f1 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -35,6 +35,7 @@ ) from airbyte._util import api_util from airbyte.cloud.client import _MAX_WORKSPACES_TO_VALIDATE, CloudClient +from airbyte.cloud.connector_builder import CloudConnectorBuilderProject from airbyte.cloud.connectors import ( CheckResult, CloudConnector, @@ -74,6 +75,7 @@ AirbyteExternalAccessNotEnabledError, AirbyteLibError, AirbyteLibInputError, + AirbyteLocalDemoModeError, AirbyteMissingResourceError, AirbyteMissingWorkspaceContextError, ) @@ -125,10 +127,25 @@ ) +TESTING_VALUES_SECRET_HINT = ( + "`testing_values` supports inline secret refs via `secret_reference::ENV_VAR_NAME`." +) + + def _manifest_yaml_source_hint() -> str | None: return MANIFEST_YAML_SOURCE_HINT if is_trusted_execution_enabled() else None +def _testing_values_secret_hint() -> str | None: + return TESTING_VALUES_SECRET_HINT if is_trusted_execution_enabled() else None + + +def _manifest_and_testing_values_hint() -> str | None: + if not is_trusted_execution_enabled(): + return None + return f"{MANIFEST_YAML_SOURCE_HINT} {TESTING_VALUES_SECRET_HINT}" + + SKILL_DOCS_SECTION_HINT = """\ Skill doc sections for direct API queries and actions are named with the pattern `actions..`, matching the `entity_type` and `action` arguments of the @@ -3211,6 +3228,15 @@ def get_cloud_organization_billing_status( ) +def _check_builder_project_created_in_session(project: CloudConnectorBuilderProject) -> None: + try: + check_guid_created_in_session(project.builder_project_id) + except AirbyteLocalDemoModeError: + if project.source_definition_id is None: + raise + check_guid_created_in_session(project.source_definition_id) + + def _get_custom_source_definition_description( custom_source: CustomCloudSourceDefinition, ) -> str: @@ -3225,7 +3251,7 @@ def _get_custom_source_definition_description( ) -@mcp_tool(extra_help_text=_manifest_yaml_source_hint) +@mcp_tool(extra_help_text=_manifest_and_testing_values_hint) def publish_custom_source_definition( # noqa: PLR0913 # Explicit args mirror MCP tool schema. ctx: Context, name: Annotated[ @@ -3281,7 +3307,6 @@ def publish_custom_source_definition( # noqa: PLR0913 # Explicit args mirror M description=( "Optional testing configuration values for the Builder UI. " "Can be provided as a JSON object or JSON string. " - "Supports inline secret refs via 'secret_reference::ENV_VAR_NAME' syntax. " "If provided, these values replace any existing testing values " "for the connector builder project, allowing immediate test read operations." ), @@ -3302,6 +3327,18 @@ def publish_custom_source_definition( # noqa: PLR0913 # Explicit args mirror M ), TraceArg.FINGERPRINT, ], + builder_project_id: Annotated[ + str | None, + Field( + description=( + "Optional. Publish this Connector Builder project's draft instead of creating " + "a new project. The manifest arguments are then optional and replace the draft " + "before publishing. If the project is already published, the draft becomes a " + "new active version of its existing definition." + ), + default=None, + ), + ] = None, ) -> str: """Publish a custom YAML source connector definition to Airbyte Cloud. @@ -3312,7 +3349,7 @@ def publish_custom_source_definition( # noqa: PLR0913 # Explicit args mirror M manifest_yaml_text=manifest_yaml_text, manifest_yaml_file=manifest_yaml_file, ) - if processed_manifest is None: + if processed_manifest is None and builder_project_id is None: error_message = ( "Provide the manifest in `manifest_yaml_text` (full YAML contents) or " "`manifest_yaml_file` (local path)." @@ -3335,12 +3372,17 @@ def publish_custom_source_definition( # noqa: PLR0913 # Explicit args mirror M ) workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + if builder_project_id is not None: + _check_builder_project_created_in_session( + workspace.get_connector_builder_project(builder_project_id) + ) custom_source = workspace.publish_custom_source_definition( name=name, manifest_yaml=processed_manifest, unique=unique, pre_validate=pre_validate, testing_values=testing_values_dict, + builder_project_id=builder_project_id, ) register_guid_created_in_session(custom_source.definition_id) return ( @@ -3451,53 +3493,7 @@ def get_custom_source_definition( return result -@mcp_tool( - read_only=True, - idempotent=True, -) -def get_connector_builder_draft_manifest( - ctx: Context, - definition_id: Annotated[ - str, - Field(description="The ID of the custom source definition to retrieve the draft for."), - ], - *, - workspace_id: Annotated[ - str | None, - Field( - description=WORKSPACE_ID_TIP_TEXT, - default=None, - ), - TraceArg.OMIT, - ], -) -> dict[str, Any]: - """Get the Connector Builder draft manifest for a custom source definition. - - Returns the working draft manifest that has been saved in the Connector Builder UI - but not yet published. This is useful for inspecting what a user is currently working - on before they publish their changes. - - If no draft exists, 'has_draft' will be False and 'draft_manifest' will be None. - The published manifest is always included for comparison. - """ - workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) - definition = workspace.get_custom_source_definition( - definition_id=definition_id, - definition_type="yaml", - ) - - return { - "definition_id": definition.definition_id, - "name": definition.name, - "connector_builder_project_id": definition.connector_builder_project_id, - "connector_builder_project_url": definition.connector_builder_project_url, - "has_draft": definition.has_draft, - "draft_manifest": definition.draft_manifest, - "published_manifest": definition.manifest, - } - - -@mcp_tool(destructive=True, extra_help_text=_manifest_yaml_source_hint) +@mcp_tool(destructive=True, extra_help_text=_manifest_and_testing_values_hint) def update_custom_source_definition( # noqa: PLR0913 # Explicit args mirror MCP tool schema. ctx: Context, definition_id: Annotated[ @@ -3557,7 +3553,6 @@ def update_custom_source_definition( # noqa: PLR0913 # Explicit args mirror MC description=( "Optional testing configuration values for the Builder UI. " "Can be provided as a JSON object or JSON string. " - "Supports inline secret refs via 'secret_reference::ENV_VAR_NAME' syntax. " "If provided, these values replace any existing testing values " "for the connector builder project. The entire testing values object " "is overwritten, so pass the full set of values you want to persist." @@ -3600,7 +3595,9 @@ def update_custom_source_definition( # noqa: PLR0913 # Explicit args mirror MC required_args = ["name", "manifest_yaml_text"] if is_trusted_execution_enabled(): required_args.append("manifest_yaml_file") - required_args.extend(("testing_values", "testing_values_secret_name")) + required_args.append("testing_values") + if is_trusted_execution_enabled(): + required_args.append("testing_values_secret_name") raise AirbyteLibInputError( message=( f"At least one of {', '.join(required_args)} must be provided to update " @@ -4330,7 +4327,7 @@ def register_cloud_tools(app: FastMCP) -> None: """ exclude_args = ["workspace_id"] if AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET else [] if not is_trusted_execution_enabled(): - exclude_args.append("manifest_yaml_file") + exclude_args.extend(("manifest_yaml_file", "testing_values_secret_name")) if exclude_args: _add_defaults_for_exclude_args(exclude_args) register_mcp_tools( diff --git a/airbyte/mcp/connector_builder.py b/airbyte/mcp/connector_builder.py new file mode 100644 index 000000000..12d1d2690 --- /dev/null +++ b/airbyte/mcp/connector_builder.py @@ -0,0 +1,969 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Airbyte Cloud Connector Builder MCP operations. + +.. include:: ../../docs/mcp-generated/connector_builder.md +""" + +# No public Python API — MCP primitives are registered via decorators and +# documented via the generated Markdown include above. +__all__: list[str] = [] + +from typing import TYPE_CHECKING, Annotated, Any + +from fastmcp import Context, FastMCP +from fastmcp_extensions import TraceArg, register_mcp_tools +from pydantic import BaseModel, Field + +from airbyte.cloud.connector_builder._models import ( + ConnectorBuilderPublishedVersion, + ConnectorBuilderReadinessReport, + ConnectorBuilderStreamReadResult, + ManifestTextEditMode, +) +from airbyte.cloud.connector_builder._util import ( + TEST_READ_MAX_PAGES, + TEST_READ_MAX_RECORDS, + TEST_READ_MAX_SLICES, +) +from airbyte.exceptions import ( + AirbyteLibInputError, +) +from airbyte.mcp._arg_resolvers import ( + resolve_connector_config, + resolve_list_of_strings, + resolve_manifest_yaml, +) +from airbyte.mcp._error_handling import redact_agent_text +from airbyte.mcp._guards import is_trusted_execution_enabled +from airbyte.mcp._tool_utils import ( + AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET, + mcp_tool, + register_guid_created_in_session, +) + + +if TYPE_CHECKING: + from airbyte.cloud.workspaces import CloudWorkspace + +from airbyte.mcp.cloud import ( + WORKSPACE_ID_TIP_TEXT, + _add_defaults_for_exclude_args, + _check_builder_project_created_in_session, + _get_cloud_workspace, + _manifest_yaml_source_hint, +) + + +_BUILDER_PROJECT_ID_DESCRIPTION = ( + "The Connector Builder project ID. Provide exactly one of `builder_project_id` or " + "`definition_id`." +) +_BUILDER_DEFINITION_ID_DESCRIPTION = ( + "The custom source definition ID, for a project that has been published. Provide exactly " + "one of `builder_project_id` or `definition_id`." +) +_TEST_READ_RECORDS_PREVIEW_LIMIT = 10 + + +class ConnectorBuilderDraftManifestResult(BaseModel): + """Draft and published manifests of a Connector Builder project.""" + + builder_project_id: str + """The Connector Builder project ID.""" + definition_id: str | None + """The published custom source definition ID, or `None` if never published.""" + connector_builder_project_url: str + """The project's Connector Builder URL.""" + name: str + """The project name.""" + has_draft: bool + """Whether the project has a saved draft manifest.""" + draft_manifest: dict[str, Any] | None + """The draft manifest, or `None` without a draft.""" + draft_manifest_yaml: str | None + """The draft manifest as YAML text, or `None` without a draft.""" + published_manifest: dict[str, Any] | None + """The published manifest, or `None` if never published.""" + + +class ConnectorBuilderProjectCreateResult(BaseModel): + """Result of creating a Connector Builder project.""" + + builder_project_id: str + """The new Connector Builder project ID.""" + definition_id: str | None + """The published custom source definition ID, set only with `create_test_source`.""" + connector_builder_project_url: str + """The project's Connector Builder URL.""" + name: str + """The project name.""" + test_source_id: str | None = None + """The test source ID, set only with `create_test_source`.""" + test_source_settings_url: str | None = None + """Cloud settings page of the test source, set only with `create_test_source`.""" + + +class ConnectorBuilderStreamTestReadToolResult(ConnectorBuilderStreamReadResult): + """Test read of one stream, with credentials redacted and a capped record preview.""" + + builder_project_id: str + """The Connector Builder project ID.""" + + +class ConnectorBuilderResolvedManifestResult(BaseModel): + """Resolved draft manifest of a Connector Builder project.""" + + builder_project_id: str + """The Connector Builder project ID.""" + stream_names: list[str] + """Stream names in the resolved manifest, including generated dynamic streams.""" + resolved_manifest: dict[str, Any] + """The resolved manifest.""" + + +class ConnectorBuilderPublishedVersionsResult(BaseModel): + """Published manifest versions of a Connector Builder project.""" + + builder_project_id: str + """The Connector Builder project ID.""" + definition_id: str | None + """The published custom source definition ID, or `None` if never published.""" + versions: list[ConnectorBuilderPublishedVersion] + """Published versions, oldest first. Empty if never published.""" + + +class ConnectorBuilderPublishedManifestResult(BaseModel): + """One published manifest version of a Connector Builder project.""" + + builder_project_id: str + """The Connector Builder project ID.""" + definition_id: str + """The published custom source definition ID.""" + version: int | None + """The returned version; `None` when the active version was requested.""" + manifest_yaml: str + """The published manifest as YAML text.""" + + +class ConnectorBuilderRevertResult(BaseModel): + """Result of reverting a Connector Builder project to a published version.""" + + builder_project_id: str + """The Connector Builder project ID.""" + definition_id: str + """The published custom source definition ID.""" + reverted_to_version: int + """The published version whose manifest was restored.""" + as_draft: bool + """Whether the version was saved as the draft instead of published.""" + active_version: int | None + """The active published version after the revert.""" + + +class ConnectorBuilderReadinessReportResult(ConnectorBuilderReadinessReport): + """Readiness report of a Connector Builder project, with credentials redacted.""" + + builder_project_id: str + """The Connector Builder project ID.""" + connector_builder_project_url: str + """The project's Connector Builder URL.""" + + +@mcp_tool( + read_only=True, + idempotent=True, +) +def get_connector_builder_draft_manifest( + ctx: Context, + definition_id: Annotated[ + str | None, + Field( + description=_BUILDER_DEFINITION_ID_DESCRIPTION, + default=None, + ), + ] = None, + *, + builder_project_id: Annotated[ + str | None, + Field( + description=_BUILDER_PROJECT_ID_DESCRIPTION, + default=None, + ), + ] = None, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], +) -> ConnectorBuilderDraftManifestResult: + """Get the Connector Builder draft manifest of a project or custom source definition. + + Returns the working draft manifest that has been saved but not yet published, both as + a dict and as `draft_manifest_yaml`. Airbyte Cloud stores the parsed manifest, so the text + is a normalized YAML rendering without comments; text edits with + `update_connector_builder_manifest_text` match against it. + + If no draft exists, 'has_draft' will be False and the draft fields will be None. + The published manifest is included for comparison, or None if never published. + """ + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + project = workspace.get_connector_builder_project( + builder_project_id, + definition_id=definition_id, + ) + return ConnectorBuilderDraftManifestResult.model_validate( + project._get_info().model_dump(by_alias=True), # noqa: SLF001 + ) + + +@mcp_tool( + extra_help_text=_manifest_yaml_source_hint, + insiders_only=True, +) +def create_connector_builder_project( + ctx: Context, + name: Annotated[ + str, + Field(description="The name for the Connector Builder project."), + TraceArg.FINGERPRINT, + ], + *, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], + manifest_yaml_text: Annotated[ + str | None, + Field( + description=( + "Optional initial draft. Full YAML contents of the declarative (low-code) " + "connector manifest: the file's text, not a filename or path." + ), + default=None, + ), + ] = None, + manifest_yaml_file: Annotated[ + str | None, + Field( + description=( + "Optional initial draft. Absolute path to a manifest YAML file on the machine " + "running this MCP server." + ), + default=None, + ), + ] = None, + pre_validate: Annotated[ + bool, + Field( + description="Whether to validate the manifest client-side before saving.", + default=True, + ), + ] = True, + create_test_source: Annotated[ + bool, + Field( + description=( + "Optional. Publish the initial manifest (which needs its spec) and create a " + "test source on it for a later full sync. Ask the user to enter their " + "credentials at both `connector_builder_project_url` (testing values) and " + "`test_source_settings_url` (then test and save) before you start building." + ), + default=False, + ), + ] = False, +) -> ConnectorBuilderProjectCreateResult: + """Create a Connector Builder project to develop a custom YAML source. + + The project holds the draft manifest and the testing values used for test reads. + Nothing appears in the workspace's source list until the project is published with + `publish_custom_source_definition`. + """ + processed_manifest = resolve_manifest_yaml( + manifest_yaml_text=manifest_yaml_text, + manifest_yaml_file=manifest_yaml_file, + ) + if create_test_source and processed_manifest is None: + raise AirbyteLibInputError( + message="`create_test_source` needs an initial manifest to publish.", + guidance="Pass `manifest_yaml_text` or `manifest_yaml_file`.", + ) + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + project = workspace.create_connector_builder_project( + name, + manifest_yaml=processed_manifest, + pre_validate=pre_validate, + ) + register_guid_created_in_session(project.builder_project_id) + result = ConnectorBuilderProjectCreateResult( + builder_project_id=project.builder_project_id, + definition_id=None, + connector_builder_project_url=project.connector_builder_project_url, + name=name, + ) + if create_test_source: + definition = project.publish(name, pre_validate=pre_validate) + register_guid_created_in_session(definition.definition_id) + test_source = project._get_or_create_default_test_source() # noqa: SLF001 + register_guid_created_in_session(test_source.connector_id) + result.definition_id = definition.definition_id + result.test_source_id = test_source.connector_id + result.test_source_settings_url = test_source.connector_url + return result + + +@mcp_tool( + destructive=True, + insiders_only=True, +) +def update_connector_builder_manifest_text( + ctx: Context, + edit_mode: Annotated[ + ManifestTextEditMode, + Field( + description=( + "`replace_full`: `new_text` is the whole manifest. " + "`replace_text`: replace the single exact match of `old_text`; fails if " + "`old_text` matches zero or several places (add surrounding text to make it " + "unique). `replace_all`: replace every exact match of `old_text`." + ), + ), + ], + new_text: Annotated[ + str, + Field( + description=( + "The whole manifest YAML for `replace_full`, else the text that replaces " + "`old_text`." + ), + ), + ], + *, + old_text: Annotated[ + str | None, + Field( + description=( + "The exact text to replace, copied verbatim from `draft_manifest_yaml`. " + "Required unless `edit_mode` is `replace_full`." + ), + default=None, + ), + ] = None, + builder_project_id: Annotated[ + str | None, + Field( + description=_BUILDER_PROJECT_ID_DESCRIPTION, + default=None, + ), + ] = None, + definition_id: Annotated[ + str | None, + Field( + description=_BUILDER_DEFINITION_ID_DESCRIPTION, + default=None, + ), + ] = None, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], + pre_validate: Annotated[ + bool, + Field( + description="Whether to validate the edited manifest client-side before saving.", + default=True, + ), + ] = True, +) -> str: + """Edit and save a Connector Builder project's draft manifest. Does not publish anything. + + Text edits match against `draft_manifest_yaml` from `get_connector_builder_draft_manifest`. + Returns a unified diff from the previous text to the saved text; base later edits on the + saved text. + """ + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + project = workspace.get_connector_builder_project( + builder_project_id, + definition_id=definition_id, + ) + _check_builder_project_created_in_session(project) + diff = project._update_draft_manifest_yaml( # noqa: SLF001 + edit_mode=edit_mode, + new_text=new_text, + old_text=old_text, + pre_validate=pre_validate, + ) + return ( + f"Saved the draft manifest of Connector Builder project {project.builder_project_id}." + "\n\n" + (diff or "No changes.") + ) + + +@mcp_tool( + destructive=True, + idempotent=True, + insiders_only=True, +) +def set_connector_builder_testing_values( + ctx: Context, + *, + builder_project_id: Annotated[ + str | None, + Field( + description=_BUILDER_PROJECT_ID_DESCRIPTION, + default=None, + ), + ] = None, + definition_id: Annotated[ + str | None, + Field( + description=_BUILDER_DEFINITION_ID_DESCRIPTION, + default=None, + ), + ] = None, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], + nonsecret_testing_values: Annotated[ + dict | str, + Field( + description=( + "Non-secret configuration values for test reads, as a JSON object or JSON " + "string. Merged onto the saved testing values, so secrets the user entered in " + "the Connector Builder are kept. Values for `airbyte_secret` fields are rejected." + ), + ), + ], +) -> str: + """Update the testing values (connector config) of a Connector Builder project. + + Test reads use these values. Airbyte Cloud stores fields marked `airbyte_secret` in the + draft manifest's spec in its secret storage. Requires a saved draft manifest. + Don't ask users to paste secrets into the chat: send them the returned project URL to + enter secrets in the Connector Builder's testing values form instead. + """ + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + project = workspace.get_connector_builder_project( + builder_project_id, + definition_id=definition_id, + ) + _check_builder_project_created_in_session(project) + testing_values_dict = resolve_connector_config( + config=nonsecret_testing_values, + config_spec_jsonschema=project._testing_values_connection_spec(), # noqa: SLF001 + ) + project._merge_testing_values(testing_values_dict) # noqa: SLF001 + return ( + f"Saved testing values for Connector Builder project {project.builder_project_id}. " + f"Users can review them or enter secrets at: {project.connector_builder_project_url}" + ) + + +def _redact_stream_diagnostics(stream_data: dict[str, Any]) -> dict[str, Any]: + """Redact credentials from a test read's page URLs, error bodies, logs, and error text.""" + for page in [*stream_data.get("pages", []), *stream_data.get("http_errors", [])]: + for key in ("url", "response_body"): + if page.get(key): + page[key] = redact_agent_text(page[key]) + for log in [*stream_data.get("logs", []), *stream_data.get("problem_logs", [])]: + if isinstance(log.get("message"), str): + log["message"] = redact_agent_text(log["message"]) + if stream_data.get("error"): + stream_data["error"] = redact_agent_text(stream_data["error"]) + return stream_data + + +@mcp_tool( + read_only=True, + external_access=True, + insiders_only=True, +) +def run_connector_builder_stream_test_read( + ctx: Context, + stream_name: Annotated[ + str, + Field(description="The name of the stream to test-read."), + TraceArg.FINGERPRINT, + ], + *, + builder_project_id: Annotated[ + str | None, + Field( + description=_BUILDER_PROJECT_ID_DESCRIPTION, + default=None, + ), + ] = None, + definition_id: Annotated[ + str | None, + Field( + description=_BUILDER_DEFINITION_ID_DESCRIPTION, + default=None, + ), + ] = None, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], + record_limit: Annotated[ + int, + Field( + description=f"Maximum records to read (1 to {TEST_READ_MAX_RECORDS:,}).", + default=100, + ge=1, + le=TEST_READ_MAX_RECORDS, + ), + TraceArg.VALUE, + ] = 100, + page_limit: Annotated[ + int, + Field( + description=f"Maximum pages to fetch per slice (1 to {TEST_READ_MAX_PAGES}).", + default=5, + ge=1, + le=TEST_READ_MAX_PAGES, + ), + TraceArg.VALUE, + ] = 5, + slice_limit: Annotated[ + int, + Field( + description=f"Maximum slices to read (1 to {TEST_READ_MAX_SLICES}).", + default=5, + ge=1, + le=TEST_READ_MAX_SLICES, + ), + TraceArg.VALUE, + ] = 5, +) -> ConnectorBuilderStreamTestReadToolResult: + """Test-read one stream of a Connector Builder project's draft manifest on Airbyte Cloud. + + Uses the project's saved testing values as the connector config. Returns counts, a + per-page summary of HTTP requests (with response bodies for HTTP errors), connector logs, + the inferred schema, and a preview of the first records (`records` holds at most + 10 records; `record_count` is the total read). `test_read_limit_reached` true means later + records weren't read. Cloud stops at page boundaries, so `record_count` can exceed + `record_limit`. + `warnings` flags reads that may be incomplete (round record counts, a full last page); + check them, but they don't fail the read. + """ + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + project = workspace.get_connector_builder_project( + builder_project_id, + definition_id=definition_id, + ) + result = project.run_stream_test_read( + stream_name, + record_limit=record_limit, + page_limit=page_limit, + slice_limit=slice_limit, + ) + return ConnectorBuilderStreamTestReadToolResult.model_validate( + { + **_redact_stream_diagnostics(result.model_dump()), + "records": result.records[:_TEST_READ_RECORDS_PREVIEW_LIMIT], + "builder_project_id": project.builder_project_id, + } + ) + + +@mcp_tool( + read_only=True, + idempotent=True, + external_access=True, + insiders_only=True, +) +def get_connector_builder_resolved_manifest( + ctx: Context, + *, + builder_project_id: Annotated[ + str | None, + Field( + description=_BUILDER_PROJECT_ID_DESCRIPTION, + default=None, + ), + ] = None, + definition_id: Annotated[ + str | None, + Field( + description=_BUILDER_DEFINITION_ID_DESCRIPTION, + default=None, + ), + ] = None, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], + stream_limit: Annotated[ + int, + Field( + description="Maximum generated streams per dynamic stream (1 to 100).", + default=5, + ge=1, + le=100, + ), + TraceArg.VALUE, + ] = 5, +) -> ConnectorBuilderResolvedManifestResult: + """Resolve a Connector Builder project's draft manifest on Airbyte Cloud. + + Expands `$ref`s, `$parameters`, and dynamic streams (using the project's testing values), + and lists the resulting stream names. + """ + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + project = workspace.get_connector_builder_project( + builder_project_id, + definition_id=definition_id, + ) + resolved = project.get_resolved_manifest(stream_limit=stream_limit) + return ConnectorBuilderResolvedManifestResult( + builder_project_id=project.builder_project_id, + stream_names=resolved.stream_names, + resolved_manifest=resolved.manifest_dict, + ) + + +@mcp_tool( + read_only=True, + idempotent=True, + insiders_only=True, +) +def list_connector_builder_published_versions( + ctx: Context, + *, + builder_project_id: Annotated[ + str | None, + Field( + description=_BUILDER_PROJECT_ID_DESCRIPTION, + default=None, + ), + ] = None, + definition_id: Annotated[ + str | None, + Field( + description=_BUILDER_DEFINITION_ID_DESCRIPTION, + default=None, + ), + ] = None, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], +) -> ConnectorBuilderPublishedVersionsResult: + """List the published manifest versions of a Connector Builder project, oldest first. + + Each version has `version`, `is_active` and `description`. Get a version's manifest with + `get_connector_builder_published_manifest`. + """ + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + project = workspace.get_connector_builder_project( + builder_project_id, + definition_id=definition_id, + ) + return ConnectorBuilderPublishedVersionsResult( + builder_project_id=project.builder_project_id, + definition_id=project.source_definition_id, + versions=project.list_published_versions(), + ) + + +@mcp_tool( + read_only=True, + idempotent=True, + insiders_only=True, +) +def get_connector_builder_published_manifest( + ctx: Context, + *, + builder_project_id: Annotated[ + str | None, + Field( + description=_BUILDER_PROJECT_ID_DESCRIPTION, + default=None, + ), + ] = None, + definition_id: Annotated[ + str | None, + Field( + description=_BUILDER_DEFINITION_ID_DESCRIPTION, + default=None, + ), + ] = None, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], + version: Annotated[ + int | None, + Field( + description="Optional published version. Defaults to the active version.", + default=None, + ge=1, + ), + TraceArg.VALUE, + ] = None, +) -> ConnectorBuilderPublishedManifestResult: + """Get a published manifest version of a Connector Builder project as YAML text.""" + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + project = workspace.get_connector_builder_project( + builder_project_id, + definition_id=definition_id, + ) + source_definition_id = project._require_source_definition_id() # noqa: SLF001 + return ConnectorBuilderPublishedManifestResult( + builder_project_id=project.builder_project_id, + definition_id=source_definition_id, + version=version, + manifest_yaml=project.get_published_manifest_yaml(version) or "", + ) + + +@mcp_tool( + destructive=True, + insiders_only=True, +) +def revert_connector_builder_published_version( + ctx: Context, + *, + version: Annotated[ + int, + Field( + description="The published version to revert to.", + ge=1, + ), + TraceArg.VALUE, + ], + builder_project_id: Annotated[ + str | None, + Field( + description=_BUILDER_PROJECT_ID_DESCRIPTION, + default=None, + ), + ] = None, + definition_id: Annotated[ + str | None, + Field( + description=_BUILDER_DEFINITION_ID_DESCRIPTION, + default=None, + ), + ] = None, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], + as_draft: Annotated[ + bool, + Field( + description=( + "True saves the version as the pending draft (overwriting unpublished edits) " + "and leaves the active version alone. False publishes it as a new active " + "version, which fails if the draft has unpublished edits." + ), + default=True, + ), + TraceArg.VALUE, + ] = True, + description: Annotated[ + str | None, + Field( + description=( + "Optional release note for the new version when `as_draft` is false. " + "Defaults to a generic note with the reverted version and the UTC time." + ), + default=None, + max_length=256, + ), + TraceArg.FINGERPRINT, + ] = None, +) -> ConnectorBuilderRevertResult: + """Revert a Connector Builder project to a published manifest version. + + List versions with `list_connector_builder_published_versions`. + """ + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + project = workspace.get_connector_builder_project( + builder_project_id, + definition_id=definition_id, + ) + _check_builder_project_created_in_session(project) + source_definition_id = project._require_source_definition_id() # noqa: SLF001 + project.revert_to_published_version( + version, + as_draft=as_draft, + description=description, + ) + active_versions = [item.version for item in project.list_published_versions() if item.is_active] + return ConnectorBuilderRevertResult( + builder_project_id=project.builder_project_id, + definition_id=source_definition_id, + reverted_to_version=version, + as_draft=as_draft, + active_version=active_versions[0] if active_versions else None, + ) + + +@mcp_tool( + external_access=True, + insiders_only=True, +) +def run_connector_builder_readiness_checks( + ctx: Context, + stream_names: Annotated[ + str | list[str] | None, + Field( + description=( + "Streams to test-read. The report covers only these streams. Required unless " + "`run_full_sync` is set. For dynamic streams, get the generated names from " + "`get_connector_builder_resolved_manifest`." + ), + default=None, + ), + TraceArg.FINGERPRINT, + ] = None, + *, + builder_project_id: Annotated[ + str | None, + Field( + description=_BUILDER_PROJECT_ID_DESCRIPTION, + default=None, + ), + ] = None, + definition_id: Annotated[ + str | None, + Field( + description=_BUILDER_DEFINITION_ID_DESCRIPTION, + default=None, + ), + ] = None, + workspace_id: Annotated[ + str | None, + Field( + description=WORKSPACE_ID_TIP_TEXT, + default=None, + ), + TraceArg.OMIT, + ], + timeout_seconds: Annotated[ + float | None, + Field( + description=( + "Optional. Return after this many seconds with the streams finished so far. " + "Streams still reading are dropped (`pending_streams`); streams not started are " + "listed in `not_attempted_streams` to pass in another call. No limit by default." + ), + default=None, + ), + ] = None, + run_full_sync: Annotated[ + bool, + Field( + description=( + "Optional. If no named stream failed or is still pending (or none are named), " + "publish the draft and start a full sync of every stream to a no-op destination. " + "Returns the job in `full_sync`; poll it with `get_cloud_sync_status`." + ), + default=False, + ), + ] = False, + test_source_id: Annotated[ + str | None, + Field( + description=( + "Optional. Source on the published definition for the full sync. Defaults to " + "the test source from `create_connector_builder_project(create_test_source=True)`." + ), + default=None, + ), + ] = None, +) -> ConnectorBuilderReadinessReportResult: + """Check a Connector Builder draft's progress and test-read its streams on Airbyte Cloud. + + Returns readiness checks for the steps to a finished custom source, with next steps, plus + per-stream results from parallel test reads using the project's testing values: `failed`, + `read_to_end`, record and page counts, HTTP errors, warning and error logs, and inferred + schemas. A stream is `failed` when its read errored, returned HTTP errors, or logged + errors. Cloud test reads stop at 5,000 records, 20 pages per slice, and 20 slices, so a + stream that hits a cap isn't `read_to_end`. The report calls out failures and warnings + without judging readiness; report every failure and warning to the user. A full sync + validates the whole connector. With `run_full_sync`, the tool publishes the draft and + starts that sync once no named stream failed or is pending. + `warnings` flags reads that may be incomplete (round record counts, a full last page); + check them, but they don't fail the read. + """ + workspace: CloudWorkspace = _get_cloud_workspace(ctx, workspace_id) + project = workspace.get_connector_builder_project( + builder_project_id, + definition_id=definition_id, + ) + if run_full_sync: + _check_builder_project_created_in_session(project) + report = project._run_readiness_checks( # noqa: SLF001 + resolve_list_of_strings(stream_names) or [], + timeout_seconds=timeout_seconds, + run_full_sync=run_full_sync, + test_source_id=test_source_id, + on_resource_created=register_guid_created_in_session, + ) + report_data = report.model_dump() + report_data["streams"] = [ + _redact_stream_diagnostics(stream) for stream in report_data["streams"] + ] + return ConnectorBuilderReadinessReportResult.model_validate( + { + "builder_project_id": project.builder_project_id, + "connector_builder_project_url": project.connector_builder_project_url, + **report_data, + } + ) + + +def register_connector_builder_tools(app: FastMCP) -> None: + """Register Connector Builder tools with the FastMCP app. + + Args: + app: FastMCP application instance + """ + exclude_args = ["workspace_id"] if AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET else [] + if not is_trusted_execution_enabled(): + exclude_args.extend(("manifest_yaml_file", "testing_values_secret_name")) + if exclude_args: + _add_defaults_for_exclude_args(exclude_args) + register_mcp_tools( + app, + mcp_module=__name__, + exclude_args=exclude_args or None, + ) diff --git a/airbyte/mcp/server.py b/airbyte/mcp/server.py index dd9c12d66..7ff5bd9de 100644 --- a/airbyte/mcp/server.py +++ b/airbyte/mcp/server.py @@ -121,6 +121,7 @@ current_airbyte_user_id, ) from airbyte.mcp.cloud import register_cloud_tools +from airbyte.mcp.connector_builder import register_connector_builder_tools from airbyte.mcp.guidance import ( KAPA_API_KEY_CONFIG_ARG, KAPA_RETRIEVAL_API_URL_CONFIG_ARG, @@ -602,6 +603,7 @@ async def _mcp_mode_lifespan( # noqa: RUF029 # Register tools from each module register_cloud_tools(app) +register_connector_builder_tools(app) register_local_tools(app) register_registry_tools(app) register_interactive_tools(app) diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md index 3513c3458..e2260f751 100644 --- a/docs/CONTRIBUTING.md +++ b/docs/CONTRIBUTING.md @@ -164,6 +164,7 @@ PyAirbyte server that is: - `index.md` — server overview (name, version, instructions, totals, module table) - `cloud.md` — tools registered by `airbyte.mcp.cloud` +- `connector_builder.md` — tools registered by `airbyte.mcp.connector_builder` - `local.md` — tools registered by `airbyte.mcp.local` - `registry.md` — tools registered by `airbyte.mcp.registry` - `guidance.md` — tools and prompts registered by `airbyte.mcp.guidance` diff --git a/docs/generate.py b/docs/generate.py index 197c65ce7..c28f879d8 100755 --- a/docs/generate.py +++ b/docs/generate.py @@ -22,7 +22,7 @@ def _regenerate_mcp_markdown() -> None: """Regenerate `docs/mcp-generated/` before pdoc runs. - The `airbyte.mcp.{cloud,local,interactive,registry,guidance}` modules pull the + The `airbyte.mcp.{cloud,connector_builder,local,interactive,registry,guidance}` modules pull the per-module Markdown files from `docs/mcp-generated/` via pdoc's `.. include::` directive. That directory is git-ignored, so on a clean checkout pdoc would fail to resolve the include unless we regenerate it diff --git a/tests/unit_tests/fixtures/arg_trace_plan.json b/tests/unit_tests/fixtures/arg_trace_plan.json index 758a80928..1a5b51832 100644 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ b/tests/unit_tests/fixtures/arg_trace_plan.json @@ -17,6 +17,14 @@ "table_prefix": "fingerprint", "workspace_id": "omit" }, + "create_connector_builder_project": { + "create_test_source": "value allowed=[false,true]", + "manifest_yaml_file": "hash", + "manifest_yaml_text": "hash", + "name": "fingerprint", + "pre_validate": "value allowed=[false,true]", + "workspace_id": "omit" + }, "deploy_connector_to_cloud": { "config": "presence", "config_secret_name": "fingerprint", @@ -143,7 +151,20 @@ "workspace_id": "omit" }, "get_connector_builder_draft_manifest": { + "builder_project_id": "hash", + "definition_id": "hash", + "workspace_id": "omit" + }, + "get_connector_builder_published_manifest": { + "builder_project_id": "hash", "definition_id": "hash", + "version": "value int", + "workspace_id": "omit" + }, + "get_connector_builder_resolved_manifest": { + "builder_project_id": "hash", + "definition_id": "hash", + "stream_limit": "value int", "workspace_id": "omit" }, "get_connector_info": { @@ -213,6 +234,11 @@ "organization_name": "fingerprint", "privilege_scope": "value allowed=[\"any\",\"instance_admin\",\"member_of\",\"organization_admin\"]" }, + "list_connector_builder_published_versions": { + "builder_project_id": "hash", + "definition_id": "hash", + "workspace_id": "omit" + }, "list_connector_config_secrets": { "connector_name": "fingerprint" }, @@ -252,6 +278,7 @@ "workspace_id": "omit" }, "publish_custom_source_definition": { + "builder_project_id": "hash", "manifest_yaml_file": "hash", "manifest_yaml_text": "hash", "name": "fingerprint", @@ -288,12 +315,38 @@ "streams": "hash list", "workspace_id": "omit" }, + "revert_connector_builder_published_version": { + "as_draft": "value allowed=[false,true]", + "builder_project_id": "hash", + "definition_id": "hash", + "description": "fingerprint", + "version": "value int", + "workspace_id": "omit" + }, "run_cloud_sync": { "connection_id": "hash", "wait": "value allowed=[false,true]", "wait_timeout": "hash", "workspace_id": "omit" }, + "run_connector_builder_readiness_checks": { + "builder_project_id": "hash", + "definition_id": "hash", + "run_full_sync": "value allowed=[false,true]", + "stream_names": "fingerprint list", + "test_source_id": "hash", + "timeout_seconds": "hash", + "workspace_id": "omit" + }, + "run_connector_builder_stream_test_read": { + "builder_project_id": "hash", + "definition_id": "hash", + "page_limit": "value int", + "record_limit": "value int", + "slice_limit": "value int", + "stream_name": "fingerprint", + "workspace_id": "omit" + }, "run_sql_query": { "max_records": "hash", "sql_query": "hash" @@ -311,6 +364,12 @@ "prefix": "fingerprint", "workspace_id": "omit" }, + "set_connector_builder_testing_values": { + "builder_project_id": "hash", + "definition_id": "hash", + "nonsecret_testing_values": "presence", + "workspace_id": "omit" + }, "set_default_cloud_workspace": { "user_email": "hash", "workspace_id": "omit" @@ -362,6 +421,15 @@ "connector_type": "value allowed=[\"destination\",\"source\"]", "workspace_id": "omit" }, + "update_connector_builder_manifest_text": { + "builder_project_id": "hash", + "definition_id": "hash", + "edit_mode": "value allowed=[\"replace_all\",\"replace_full\",\"replace_text\"]", + "new_text": "hash", + "old_text": "hash", + "pre_validate": "value allowed=[false,true]", + "workspace_id": "omit" + }, "update_custom_source_definition": { "definition_id": "hash", "manifest_yaml_file": "hash", diff --git a/tests/unit_tests/test_cloud_connector_builder.py b/tests/unit_tests/test_cloud_connector_builder.py new file mode 100644 index 000000000..126e6b05c --- /dev/null +++ b/tests/unit_tests/test_cloud_connector_builder.py @@ -0,0 +1,1731 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Unit tests for Connector Builder projects in Airbyte Cloud.""" + +from __future__ import annotations + +import re + +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import threading +import pytest +import requests + +from airbyte._util import api_util +from airbyte.cloud import connector_builder +from airbyte.cloud.connector_builder import _models +from airbyte.cloud.connector_builder import ( + CloudConnectorBuilderProject, + ConnectorBuilderPageSummary, + ConnectorBuilderResolvedManifest, + ConnectorBuilderStreamReadResult, +) +from airbyte.cloud.connector_builder._models import ( + ConnectorBuilderStreamReadiness, + _record_count_warnings, +) +from airbyte.cloud.connector_builder._util import ( + apply_manifest_text_edit, + drop_masked_secrets, + publish_spec, + render_manifest_text, + spec_has_secrets, +) +from airbyte.cloud.workspaces import CloudWorkspace +from airbyte.exceptions import ( + AirbyteCloudError, + AirbyteDuplicateResourcesError, + AirbyteLibInputError, +) +from airbyte.secrets.base import SecretString + +MANIFEST: dict[str, Any] = { + "version": "6.0.0", + "streams": [{"name": "users"}, {"name": "orders"}], + "spec": { + "connection_specification": {"type": "object", "properties": {}}, + "documentation_url": "https://example.com/docs", + }, +} + + +@pytest.mark.parametrize( + ("edit_mode", "new_text", "old_text", "expected"), + [ + pytest.param("replace_full", "b: 2\n", None, "b: 2\n", id="replace-full"), + pytest.param( + "replace_text", "y", "x", "a: y\nb: z\n", id="replace-text-single" + ), + pytest.param("replace_all", "q", "z", "a: x\nb: q\n", id="replace-all-single"), + ], +) +def test_apply_manifest_text_edit( + edit_mode: _models.ManifestTextEditMode, + new_text: str, + old_text: str | None, + expected: str, +) -> None: + result = apply_manifest_text_edit( + "a: x\nb: z\n", + edit_mode=edit_mode, + new_text=new_text, + old_text=old_text, + ) + assert result == expected + + +def test_apply_manifest_text_edit_replace_all_replaces_every_match() -> None: + result = apply_manifest_text_edit( + "a: x\nb: x\n", edit_mode="replace_all", new_text="y", old_text="x" + ) + assert result == "a: y\nb: y\n" + + +@pytest.mark.parametrize( + ("edit_mode", "old_text", "match"), + [ + pytest.param("replace_text", "x", "matches 2 places", id="multiple-matches"), + pytest.param("replace_text", "nope", "was not found", id="no-match"), + pytest.param("replace_all", "nope", "was not found", id="replace-all-no-match"), + pytest.param("replace_text", None, "is required", id="missing-old-text"), + pytest.param("replace_text", "", "is required", id="empty-old-text"), + pytest.param("replace_full", "x", "not allowed", id="old-text-with-full"), + ], +) +def test_apply_manifest_text_edit_errors( + edit_mode: _models.ManifestTextEditMode, + old_text: str | None, + match: str, +) -> None: + with pytest.raises(AirbyteLibInputError, match=match): + apply_manifest_text_edit( + "a: x\nb: x\n", edit_mode=edit_mode, new_text="y", old_text=old_text + ) + + +def test_render_manifest_text_keeps_order_and_does_not_wrap() -> None: + long_value = "word " * 40 + text = render_manifest_text({"z": 1, "a": long_value.strip()}) + assert text.splitlines()[0] == "z: 1" + assert len(text.splitlines()) == 2 + + +def test_stream_read_result__from_api_response() -> None: + response = { + "logs": [{"level": "ERROR", "message": "boom", "stacktrace": "..."}], + "slices": [ + { + "pages": [ + { + "records": [{"id": 1}, {"id": 2}], + "request": {"url": "https://api/x", "http_method": "GET"}, + "response": {"status": 200, "body": "ok"}, + }, + { + "records": [], + "request": { + "url": "https://api/x?page=2", + "http_method": "GET", + }, + "response": {"status": 401, "body": "unauthorized"}, + }, + ] + } + ], + "test_read_limit_reached": True, + "inferred_schema": {"type": "object"}, + } + + result = ConnectorBuilderStreamReadResult._from_api_response("users", response) + + assert result.record_count == 2 + assert result.slice_count == 1 + assert result.page_count == 2 + assert result.test_read_limit_reached is True + assert [page.status for page in result.pages] == [200, 401] + assert result.pages[0].response_body is None + assert result.pages[1].response_body == "unauthorized" + assert result.logs == [{"level": "ERROR", "message": "boom"}] + assert result.inferred_schema == {"type": "object"} + + +@pytest.mark.parametrize( + ("record_count", "page_counts", "limit_reached", "expected"), + [ + pytest.param( + 100, + [100], + False, + [ + "Read 100 records, a multiple of 10. A page size or API limit may have cut " + "the read short; check the paginator.", + "The last page returned a full page of 100 records, the same as the first " + "page, so more pages may exist; check the paginator.", + ], + id="round-single-page", + ), + pytest.param(7, [5, 2], False, [], id="non-round-partial-last-page"), + pytest.param(5, [5, 0], False, [], id="empty-final-page"), + pytest.param( + 10, + [5, 5], + False, + [ + "Read 10 records, a multiple of 10. A page size or API limit may have cut " + "the read short; check the paginator.", + "The last page returned a full page of 5 records, the same as the first " + "page, so more pages may exist; check the paginator.", + ], + id="full-final-page", + ), + pytest.param( + 0, + [0], + False, + ["Read 0 records. Check the request, record selector, and testing values."], + id="zero-records", + ), + pytest.param(1, [1], False, [], id="one-record"), + pytest.param( + 150, + [50, 50, 50], + False, + [ + "Read 150 records, a multiple of 10. A page size or API limit may have cut " + "the read short; check the paginator.", + "The last page returned a full page of 50 records, the same as the first " + "page, so more pages may exist; check the paginator.", + ], + id="multiple-full-pages", + ), + pytest.param(100, [100], True, [], id="reached-test-read-limit"), + ], +) +def test_record_count_warnings( + record_count: int, + page_counts: list[int], + limit_reached: bool, + expected: list[str], +) -> None: + pages = [ConnectorBuilderPageSummary(record_count=count) for count in page_counts] + + assert ( + _record_count_warnings( + record_count, + pages, + test_read_limit_reached=limit_reached, + ) + == expected + ) + + +@pytest.mark.parametrize("limit_reached", [False, True]) +def test_stream_read_warnings_propagate_to_readiness_without_failing( + limit_reached: bool, +) -> None: + response = { + "slices": [ + { + "pages": [ + {"records": [{"id": index} for index in range(100)]}, + ] + } + ], + "test_read_limit_reached": limit_reached, + } + + result = ConnectorBuilderStreamReadResult._from_api_response("users", response) + readiness = ConnectorBuilderStreamReadiness._from_stream_read(result) + + if limit_reached: + assert result.warnings == [] + else: + assert len(result.warnings) == 2 + assert readiness.warnings == result.warnings + assert readiness.failed is False + assert readiness.read_to_end is not limit_reached + + +def test_resolved_manifest_lists_stream_names() -> None: + resolved = ConnectorBuilderResolvedManifest._from_manifest(MANIFEST) + assert resolved.stream_names == ["users", "orders"] + + +@pytest.mark.parametrize( + ("helper", "kwargs", "expected_path", "expected_json", "response", "expected"), + [ + pytest.param( + api_util.create_connector_builder_project, + {"workspace_id": "ws", "name": "My API", "draft_manifest": MANIFEST}, + "/connector_builder_projects/create", + { + "workspaceId": "ws", + "builderProject": {"name": "My API", "draftManifest": MANIFEST}, + }, + {"builderProjectId": "bp", "workspaceId": "ws"}, + "bp", + id="create", + ), + pytest.param( + api_util.read_connector_builder_project_stream, + { + "workspace_id": "ws", + "builder_project_id": "bp", + "manifest": MANIFEST, + "stream_name": "users", + "record_limit": 10, + "page_limit": 2, + "slice_limit": 3, + }, + "/connector_builder_projects/read_stream", + { + "workspaceId": "ws", + "builderProjectId": "bp", + "manifest": MANIFEST, + "streamName": "users", + "recordLimit": 10, + "pageLimit": 2, + "sliceLimit": 3, + }, + {"logs": [], "slices": []}, + {"logs": [], "slices": []}, + id="read-stream", + ), + pytest.param( + api_util.full_resolve_connector_builder_project_manifest, + { + "workspace_id": "ws", + "builder_project_id": "bp", + "manifest": MANIFEST, + "stream_limit": 7, + }, + "/connector_builder_projects/full_resolve", + { + "workspaceId": "ws", + "builderProjectId": "bp", + "manifest": MANIFEST, + "streamLimit": 7, + }, + {"manifest": {"streams": []}}, + {"streams": []}, + id="full-resolve", + ), + pytest.param( + api_util.publish_connector_builder_project, + { + "workspace_id": "ws", + "builder_project_id": "bp", + "name": "My API", + "manifest": MANIFEST, + "spec": {"connectionSpecification": {}}, + "description": "First release", + }, + "/connector_builder_projects/publish", + { + "workspaceId": "ws", + "builderProjectId": "bp", + "name": "My API", + "initialDeclarativeManifest": { + "manifest": MANIFEST, + "spec": {"connectionSpecification": {}}, + "version": 1, + "description": "First release", + }, + }, + {"sourceDefinitionId": "def"}, + "def", + id="publish", + ), + pytest.param( + api_util.create_declarative_manifest_version, + { + "workspace_id": "ws", + "source_definition_id": "def", + "manifest": MANIFEST, + "spec": {"connectionSpecification": {}}, + "version": 3, + "description": "Fix paging", + }, + "/declarative_source_definitions/create_manifest", + { + "workspaceId": "ws", + "sourceDefinitionId": "def", + "setAsActiveManifest": True, + "declarativeManifest": { + "manifest": MANIFEST, + "spec": {"connectionSpecification": {}}, + "version": 3, + "description": "Fix paging", + }, + }, + {}, + None, + id="create-manifest-version", + ), + ], +) +def test_connector_builder_api_helpers_forward_requests( + monkeypatch: pytest.MonkeyPatch, + helper: Any, + kwargs: dict[str, Any], + expected_path: str, + expected_json: dict[str, Any], + response: dict[str, Any], + expected: object, +) -> None: + captured: dict[str, Any] = {} + + def fake_config_request(**request_kwargs: Any) -> dict[str, Any]: + captured.update(request_kwargs) + return response + + monkeypatch.setattr(api_util, "_make_config_api_request", fake_config_request) + + result = helper( + **kwargs, + api_root="https://api.example", + config_api_root="https://config.example", + client_id=None, + client_secret=None, + bearer_token=SecretString("token"), + ) + + assert result == expected + assert captured["path"] == expected_path + assert captured["json"] == expected_json + + +def _fake_workspace() -> Any: + return SimpleNamespace( + workspace_id="ws", + api_root="https://api.airbyte.com/v1", + config_api_root=None, + client_id=None, + client_secret=None, + bearer_token=SecretString("token"), + get_custom_source_definition=MagicMock(), + list_sources=MagicMock(return_value=[]), + ) + + +def _patch_project_store( + monkeypatch: pytest.MonkeyPatch, + *, + draft: dict[str, Any] | None, + source_definition_id: str | None = None, + testing_values: dict[str, Any] | None = None, +) -> dict[str, Any]: + store: dict[str, Any] = {"draft": draft, "saved": []} + + def fake_get(**_: Any) -> dict[str, Any]: + return { + "builderProject": { + "name": "My API", + "builderProjectId": "bp", + "hasDraft": store["draft"] is not None, + "sourceDefinitionId": source_definition_id, + }, + "declarativeManifest": {"manifest": store["draft"], "isDraft": True}, + "testingValues": testing_values, + } + + def fake_update(**kwargs: Any) -> None: + store["draft"] = kwargs["draft_manifest"] + store["saved"].append(kwargs) + + monkeypatch.setattr(api_util, "get_connector_builder_project", fake_get) + monkeypatch.setattr(api_util, "update_connector_builder_project", fake_update) + monkeypatch.setattr(api_util, "validate_yaml_manifest", lambda *_, **__: None) + return store + + +def test_project_url() -> None: + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + assert ( + project.connector_builder_project_url + == "https://cloud.airbyte.com/workspaces/ws/connector-builder/edit/bp" + ) + + +def test_update_draft_manifest_yaml_saves_and_returns_diff( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _patch_project_store(monkeypatch, draft={"version": "6.0.0", "streams": []}) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + + diff = project._update_draft_manifest_yaml( + edit_mode="replace_text", + old_text="version: 6.0.0", + new_text="version: 6.1.0", + ) + + assert store["draft"] == {"version": "6.1.0", "streams": []} + assert store["saved"][0]["name"] == "My API" + assert "-version: 6.0.0" in diff + assert "+version: 6.1.0" in diff + + +def test_update_draft_manifest_yaml_without_draft_requires_replace_full( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _patch_project_store(monkeypatch, draft=None) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + + with pytest.raises(AirbyteLibInputError, match="no draft manifest"): + project._update_draft_manifest_yaml( + edit_mode="replace_all", old_text="a", new_text="b" + ) + + project._update_draft_manifest_yaml( + edit_mode="replace_full", new_text="version: 6.0.0\n" + ) + assert store["draft"] == {"version": "6.0.0"} + + +def test_replace_testing_values_uses_draft_spec( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST) + update_testing_values = MagicMock() + monkeypatch.setattr( + api_util, + "update_connector_builder_project_testing_values", + update_testing_values, + ) + + CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ).replace_testing_values({"api_key": "k"}) + + kwargs = update_testing_values.call_args.kwargs + assert kwargs["testing_values"] == {"api_key": "k"} + assert kwargs["spec"] == { + "documentationUrl": "https://example.com/docs", + "connectionSpecification": {"type": "object", "properties": {}}, + "advancedAuth": None, + } + + +def test_replace_testing_values_without_draft_uses_published_spec( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft=None, source_definition_id="def") + update_testing_values = MagicMock() + monkeypatch.setattr( + api_util, + "update_connector_builder_project_testing_values", + update_testing_values, + ) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + monkeypatch.setattr(project, "_published_manifest", lambda _id: _secret_manifest()) + + project.replace_testing_values({"api_key": "k"}) + + spec = update_testing_values.call_args.kwargs["spec"] + assert spec["connectionSpecification"]["properties"]["api_key"]["airbyte_secret"] + + +def test_publish_unpublished_project(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST) + publish = MagicMock(return_value="def") + monkeypatch.setattr(api_util, "publish_connector_builder_project", publish) + workspace = _fake_workspace() + + CloudConnectorBuilderProject(workspace=workspace, builder_project_id="bp").publish( + "My API" + ) + + spec = publish.call_args.kwargs["spec"] + assert ( + spec["connectionSpecification"] == MANIFEST["spec"]["connection_specification"] + ) + workspace.get_custom_source_definition.assert_called_once_with( + "def", definition_type="yaml" + ) + + +def test_publish_already_published_project_adds_version( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST, source_definition_id="def") + publish = MagicMock() + create_version = MagicMock() + monkeypatch.setattr(api_util, "publish_connector_builder_project", publish) + monkeypatch.setattr(api_util, "create_declarative_manifest_version", create_version) + monkeypatch.setattr( + api_util, + "list_declarative_manifest_versions", + lambda **_: [ + {"version": 1, "isActive": False}, + {"version": 2, "isActive": True}, + ], + ) + workspace = _fake_workspace() + + CloudConnectorBuilderProject(workspace=workspace, builder_project_id="bp").publish( + "My API", + description="Fix paging", + ) + + publish.assert_not_called() + kwargs = create_version.call_args.kwargs + assert kwargs["manifest"] == MANIFEST + assert kwargs["version"] == 3 + assert kwargs["description"] == "Fix paging" + + +def test_publish_defaults_description_to_utc_timestamp( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST) + publish = MagicMock(return_value="def") + monkeypatch.setattr(api_util, "publish_connector_builder_project", publish) + + CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ).publish("My API") + + assert re.fullmatch( + r"Published by PyAirbyte at \d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z", + publish.call_args.kwargs["description"], + ) + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({}, id="neither"), + pytest.param({"builder_project_id": "bp", "definition_id": "d"}, id="both"), + ], +) +def test_get_connector_builder_project_requires_exactly_one_id( + kwargs: dict[str, str], +) -> None: + workspace = CloudWorkspace( + workspace_id="ws", + client_id=SecretString("id"), + client_secret=SecretString("secret"), + ) + with pytest.raises(AirbyteLibInputError, match="exactly one"): + workspace.get_connector_builder_project(**kwargs) + + +def test_project_uses_owner_workspace() -> None: + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp", owner_workspace_id="owner" + ) + assert project.connector_builder_project_url.endswith( + "/workspaces/owner/connector-builder/edit/bp" + ) + + +def test_get_info_reads_project_once(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_project_store(monkeypatch, draft={"version": "6.0.0"}) + calls: list[dict[str, Any]] = [] + original = api_util.get_connector_builder_project + + def counting_get(**kwargs: Any) -> dict[str, Any]: + calls.append(kwargs) + return original(**kwargs) + + monkeypatch.setattr(api_util, "get_connector_builder_project", counting_get) + + info = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + )._get_info() + + assert len(calls) == 1 + assert info.has_draft is True + assert info.draft_manifest_yaml == "version: 6.0.0\n" + assert info.definition_id is None + assert info.published_manifest_dict is None + dumped = info.model_dump(by_alias=True) + assert "draft_manifest" in dumped + assert dumped["published_manifest"] is None + + +def test_publish_validates_unpublished_draft_unless_disabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST) + validate = MagicMock() + monkeypatch.setattr(api_util, "validate_yaml_manifest", validate) + monkeypatch.setattr( + api_util, "publish_connector_builder_project", MagicMock(return_value="d") + ) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + + project.publish("My API", pre_validate=False) + validate.assert_not_called() + + project.publish("My API") + validate.assert_called_once() + + +def _stream_read( + stream_name: str, + *, + record_count: int, + status: int = 200, + logs: list[dict[str, Any]] | None = None, + limit_reached: bool = False, +) -> ConnectorBuilderStreamReadResult: + return ConnectorBuilderStreamReadResult( + stream_name=stream_name, + record_count=record_count, + slice_count=1, + page_count=1, + test_read_limit_reached=limit_reached, + records=[], + pages=[ConnectorBuilderPageSummary(status=status, record_count=record_count)], + logs=logs or [], + ) + + +def _patch_reads( + monkeypatch: pytest.MonkeyPatch, + project: CloudConnectorBuilderProject, + reads: dict[str, ConnectorBuilderStreamReadResult], +) -> None: + def fake_read( + _manifest: Any, stream_name: str, **_: Any + ) -> ConnectorBuilderStreamReadResult: + if stream_name == "dropped": + raise requests.ConnectionError("connection reset") + if stream_name not in reads: + raise AirbyteLibInputError( + message="read failed", guidance="Fix the stream." + ) + return reads[stream_name] + + monkeypatch.setattr(project, "_run_stream_test_read", fake_read) + + +def test_run_readiness_checks_summarizes_each_stream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + names = ["ok", "empty", "http_error", "error_log", "warn_log", "broken", "dropped"] + _patch_project_store( + monkeypatch, draft={**MANIFEST, "streams": [{"name": name} for name in names]} + ) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + _patch_reads( + monkeypatch, + project, + { + "ok": _stream_read("ok", record_count=3), + "empty": ConnectorBuilderStreamReadResult._from_api_response( + "empty", + {"slices": [{"pages": [{"records": []}]}]}, + ), + "http_error": _stream_read("http_error", record_count=1, status=422), + "error_log": _stream_read( + "error_log", + record_count=1, + logs=[{"level": "ERROR", "message": "boom"}], + ), + "warn_log": _stream_read( + "warn_log", record_count=1, logs=[{"level": "WARN", "message": "hmm"}] + ), + }, + ) + + report = project._run_readiness_checks(names) + + assert [stream.stream_name for stream in report.streams] == names + assert {stream.stream_name: stream.failed for stream in report.streams} == { + "ok": False, + "empty": False, + "http_error": True, + "error_log": True, + "warn_log": False, + "broken": True, + "dropped": True, + } + assert report.failed_streams == [ + "http_error", + "error_log", + "broken", + "dropped", + ] + by_name = {stream.stream_name: stream for stream in report.streams} + assert by_name["empty"].warnings == [ + "Read 0 records. Check the request, record selector, and testing values." + ] + assert "empty" not in report.failed_streams + assert by_name["http_error"].http_errors[0].status == 422 + assert by_name["http_error"].read_to_end is False + assert by_name["warn_log"].problem_logs == [{"level": "WARN", "message": "hmm"}] + assert by_name["broken"].error == "read failed Fix the stream." + assert ( + by_name["dropped"].error == "Request to Airbyte Cloud failed: ConnectionError" + ) + assert report.checks[4].detail == "failed: http_error, error_log, broken, dropped" + assert report.checks[5].result == "failed" + assert "not read: http_error, broken, dropped (test read failed)" in ( + report.checks[5].detail or "" + ) + assert ( + report.next_steps[-1] == "Report these streams' `warnings` to the user: empty" + ) + assert not any("All named streams passed" in step for step in report.next_steps) + assert not any("ready" in step.lower() for step in report.next_steps) + + +def test_run_readiness_checks_requires_stream_names() -> None: + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + with pytest.raises(AirbyteLibInputError): + project._run_readiness_checks([]) + + +@pytest.mark.parametrize("method", ["run_stream_test_read", "_run_readiness_checks"]) +@pytest.mark.parametrize( + ("limit_name", "limit_value"), + [ + ("record_limit", 5_001), + ("page_limit", 21), + ("slice_limit", 21), + ], +) +def test_test_read_limits_reject_out_of_range_without_api_calls( + monkeypatch: pytest.MonkeyPatch, + method: str, + limit_name: str, + limit_value: int, +) -> None: + def unexpected_api_call(*_: Any, **__: Any) -> Any: + pytest.fail("A test-read limit was rejected only after an API call.") + + monkeypatch.setattr(api_util, "get_connector_builder_project", unexpected_api_call) + monkeypatch.setattr( + api_util, "read_connector_builder_project_stream", unexpected_api_call + ) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + expected: dict[str, int] = ( + {"record_limit": 100, "page_limit": 5, "slice_limit": 5} + if method == "run_stream_test_read" + else { + "record_limit": connector_builder._util.TEST_READ_MAX_RECORDS, + "page_limit": connector_builder._util.TEST_READ_MAX_PAGES, + "slice_limit": connector_builder._util.TEST_READ_MAX_SLICES, + } + ) + + with pytest.raises(AirbyteLibInputError) as error: + if method == "run_stream_test_read": + project.run_stream_test_read("users", **{limit_name: limit_value}) + else: + project._run_readiness_checks(["users"], **{limit_name: limit_value}) + + expected[limit_name] = limit_value + assert error.value.get_message() == "Test-read limits are out of range." + assert error.value.context == expected + + +def test_readiness_report_defaults_forward_cloud_test_read_limits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft={"streams": [{"name": "users"}]}) + received: list[dict[str, Any]] = [] + + def fake_read(**kwargs: Any) -> dict[str, Any]: + received.append(kwargs) + return {"slices": [{"pages": [{"records": [{"id": 1}]}]}]} + + monkeypatch.setattr(api_util, "read_connector_builder_project_stream", fake_read) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + + project._run_readiness_checks(["users"]) + + assert len(received) == 1 + assert received[0]["record_limit"] == connector_builder._util.TEST_READ_MAX_RECORDS + assert received[0]["page_limit"] == connector_builder._util.TEST_READ_MAX_PAGES + assert received[0]["slice_limit"] == connector_builder._util.TEST_READ_MAX_SLICES + + +def test_readiness_checklist_does_not_call_failed_reads_capped( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft={"streams": [{"name": "users"}]}) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + + def fail_read(*_: Any, **__: Any) -> ConnectorBuilderStreamReadResult: + raise AirbyteCloudError(message="Cloud read failed.") + + monkeypatch.setattr(project, "_run_stream_test_read", fail_read) + + report = project._run_readiness_checks(["users"]) + + assert report.streams[0].read_to_end is False + assert report.checks[5].result == "failed" + assert report.checks[5].detail == "not read: users (test read failed)" + assert "hit a cap" not in (report.checks[5].detail or "") + + +def test_readiness_checklist_marks_failed_capped_stream_incomplete( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft={"streams": [{"name": "users"}]}) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + _patch_reads( + monkeypatch, + project, + { + "users": _stream_read( + "users", + record_count=1, + logs=[{"level": "ERROR", "message": "read failed"}], + limit_reached=True, + ) + }, + ) + + report = project._run_readiness_checks(["users"]) + + assert report.streams[0].failed is True + assert report.streams[0].error is None + assert report.streams[0].http_errors == [] + assert report.failed_streams == ["users"] + assert report.checks[5].result == "failed" + assert report.checks[5].detail == ( + "hit a cap: users. Check these with a full sync after publishing." + ) + assert any( + "Fix the failed streams and test them again: users" in step + for step in report.next_steps + ) + assert any( + "hit a test-read cap" in step and "users" in step for step in report.next_steps + ) + + +def test_run_readiness_checks_without_draft_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft=None) + with pytest.raises(AirbyteLibInputError): + CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + )._run_readiness_checks(["users"]) + + +def test_readiness_report_lists_caps_without_a_success_verdict( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + reads = { + "users": _stream_read("users", record_count=1), + "orders": _stream_read("orders", record_count=1, limit_reached=True), + } + _patch_reads(monkeypatch, project, reads) + + partial = project._run_readiness_checks(["users"]) + capped = project._run_readiness_checks(["users", "orders"]) + reads["orders"] = _stream_read("orders", record_count=1) + full = project._run_readiness_checks(["users", "orders", "users"]) + + assert partial.failed_streams == [] + assert capped.failed_streams == [] + assert "passed" not in partial.model_dump() + assert any( + "These streams hit a test-read cap, so later pages weren't read; " + "check them with a full sync: orders" == step + for step in capped.next_steps + ) + assert not any("Not ready yet" in step for step in capped.next_steps) + assert not any("All named streams passed" in step for step in capped.next_steps) + assert not any("ready" in step.lower() for step in capped.next_steps) + assert [s.read_to_end for s in capped.streams] == [True, False] + assert capped.checks[5].result == "failed" + assert full.failed_streams == [] + assert not any("All named streams passed" in step for step in full.next_steps) + assert not any("ready" in step.lower() for step in full.next_steps) + + +def test_readiness_checks_require_testing_values_when_spec_has_inputs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + draft = { + "version": "6.0.0", + "spec": {"connection_specification": {"properties": {"api_key": {}}}}, + "streams": [{"name": "users"}], + } + _patch_project_store(monkeypatch, draft=draft, source_definition_id="def") + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + _patch_reads(monkeypatch, project, {"users": _stream_read("users", record_count=1)}) + + report = project._run_readiness_checks(["users"]) + + assert [check.result for check in report.checks] == [ + "passed", + "passed", + "passed", + "failed", + "passed", + "passed", + "passed", + "unknown", + ] + assert report.failed_streams == [] + + +def test_publish_spec_uses_api_field_names() -> None: + manifest = { + "spec": { + "connection_specification": {"type": "object"}, + "documentation_url": "https://example.com/docs", + "advanced_auth": {"auth_flow_type": "oauth2.0"}, + } + } + + spec = publish_spec(manifest) + + assert spec["connectionSpecification"] == {"type": "object"} + assert spec["documentationUrl"] == "https://example.com/docs" + assert spec["advancedAuth"] == {"auth_flow_type": "oauth2.0"} + + +def test_publish_duplicate_name_leaves_draft_unchanged( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _patch_project_store(monkeypatch, draft={"version": "6.0.0"}) + workspace = CloudWorkspace( + workspace_id="ws", + client_id=SecretString("id"), + client_secret=SecretString("secret"), + ) + monkeypatch.setattr( + CloudWorkspace, + "list_custom_source_definitions", + lambda *_, **__: [SimpleNamespace(name="My API")], + ) + + with pytest.raises(AirbyteDuplicateResourcesError): + workspace.publish_custom_source_definition( + "My API", + builder_project_id="bp", + manifest_yaml=MANIFEST, + unique=True, + ) + + assert store["saved"] == [] + + +def test_readiness_timeout_splits_pending_and_not_attempted( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + release = threading.Event() + slow = {f"slow{i}" for i in range(5)} + + def fake_read( + _manifest: Any, stream_name: str, **_: Any + ) -> ConnectorBuilderStreamReadResult: + if stream_name in slow: + release.wait(5) + return _stream_read(stream_name, record_count=1) + + monkeypatch.setattr(project, "_run_stream_test_read", fake_read) + names = ["fast", *sorted(slow), "queued"] + try: + report = project._run_readiness_checks(names, timeout_seconds=0.5) + finally: + release.set() + assert report.failed_streams == [] + assert report.pending_streams == sorted(slow) + assert report.not_attempted_streams == ["queued"] + assert any( + "not-attempted streams" in step and "queued" in step + for step in report.next_steps + ) + assert any("still reading at the timeout" in step for step in report.next_steps) + + +def test_readiness_checklist_is_incomplete_when_a_stream_is_pending( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + release = threading.Event() + + def fake_read( + _manifest: Any, stream_name: str, **_: Any + ) -> ConnectorBuilderStreamReadResult: + if stream_name == "pending": + release.wait(5) + return _stream_read(stream_name, record_count=1) + + monkeypatch.setattr(project, "_run_stream_test_read", fake_read) + try: + report = project._run_readiness_checks( + ["ready", "pending"], + timeout_seconds=0.5, + ) + finally: + release.set() + + assert report.pending_streams == ["pending"] + assert [stream.stream_name for stream in report.streams] == ["ready"] + assert all(stream.read_to_end for stream in report.streams) + assert report.checks[5].result == "failed" + + +def test_readiness_bounds_read_requests_by_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + timeouts: list[float | None] = [] + + def fake_read( + _manifest: Any, stream_name: str, **kwargs: Any + ) -> ConnectorBuilderStreamReadResult: + timeouts.append(kwargs["timeout_seconds"]) + return _stream_read(stream_name, record_count=1) + + monkeypatch.setattr(project, "_run_stream_test_read", fake_read) + project._run_readiness_checks(["users"], timeout_seconds=30) + project._run_readiness_checks(["users"]) + + # Windows clock granularity can make the remaining time a hair over the timeout. + assert timeouts[0] is not None and 0 < timeouts[0] <= 30 + 1e-3 + assert timeouts[1] is None + + +def test_drop_masked_secrets() -> None: + values = { + "api_key": "**********", + "region": "us", + "auth": {"token": "**********", "kind": "bearer"}, + "items": [{"secret": "**********", "id": 1}], + } + + assert drop_masked_secrets(values) == { + "region": "us", + "auth": {"kind": "bearer"}, + "items": [{"id": 1}], + } + + +@pytest.mark.parametrize( + ("spec", "expected"), + [ + pytest.param({"type": "object", "properties": {}}, False, id="no-inputs"), + pytest.param( + {"properties": {"region": {"type": "string"}}}, False, id="no-secrets" + ), + pytest.param( + { + "properties": { + "auth": { + "oneOf": [{"properties": {"key": {"airbyte_secret": True}}}] + } + } + }, + True, + id="nested-secret", + ), + ], +) +def test_spec_has_secrets(spec: dict[str, Any], expected: bool) -> None: + assert spec_has_secrets(spec) is expected + + +def _secret_manifest() -> dict[str, Any]: + return { + **MANIFEST, + "spec": { + "connection_specification": { + "type": "object", + "properties": { + "api_key": {"type": "string", "airbyte_secret": True}, + "region": {"type": "string"}, + }, + }, + }, + } + + +@pytest.mark.parametrize( + ("manifest", "deferred"), + [ + pytest.param(_secret_manifest(), True, id="secret-spec-deferred"), + pytest.param(MANIFEST, False, id="no-secret-spec-ready"), + pytest.param( + { + **MANIFEST, + "spec": { + "type": "Spec", + "connection_specification": { + "type": "object", + "required": ["base_url"], + "properties": {"base_url": {"type": "string"}}, + }, + }, + }, + True, + id="missing-required-deferred", + ), + ], +) +def test_create_test_source_uses_non_secret_testing_values( + monkeypatch: pytest.MonkeyPatch, + manifest: dict[str, Any], + deferred: bool, +) -> None: + _patch_project_store( + monkeypatch, + draft=MANIFEST, + source_definition_id="def", + testing_values={"api_key": "**********", "region": "us"}, + ) + workspace = _fake_workspace() + workspace.deploy_source = MagicMock() + project = CloudConnectorBuilderProject(workspace=workspace, builder_project_id="bp") + # The source runs the published version, so its spec wins over the draft's. + monkeypatch.setattr(project, "_published_manifest", lambda _id: manifest) + + project._get_or_create_default_test_source() + + workspace.deploy_source.assert_called_once_with( + name="My API (Connector Builder test)", + source={"region": "us"}, + definition_id="def", + defer_credentials=deferred, + ) + + +def test_create_test_source_requires_publish(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + + with pytest.raises(AirbyteLibInputError, match="Publish"): + project._get_or_create_default_test_source() + + +def _full_sync_project( + monkeypatch: pytest.MonkeyPatch, + *, + destinations: list[Any], + connections: list[Any], + sources: list[Any] | None = None, + draft: dict[str, Any] | None = MANIFEST, + discovered_streams: list[str] | None = None, +) -> tuple[CloudConnectorBuilderProject, Any]: + _patch_project_store(monkeypatch, draft=draft, source_definition_id="def") + workspace = _fake_workspace() + workspace.get_source = MagicMock(return_value=SimpleNamespace(definition_id="def")) + workspace.list_sources = MagicMock(return_value=sources or []) + workspace.list_destinations = MagicMock(return_value=destinations) + workspace.list_connections = MagicMock(return_value=connections) + workspace.deploy_destination = MagicMock( + return_value=SimpleNamespace(connector_id="new-dest") + ) + new_connection = MagicMock(connection_id="conn") + new_connection.run_sync.return_value = SimpleNamespace( + job_id=7, job_url="https://job" + ) + workspace.deploy_connection = MagicMock(return_value=new_connection) + project = CloudConnectorBuilderProject(workspace=workspace, builder_project_id="bp") + monkeypatch.setattr(project, "publish", MagicMock()) + monkeypatch.setattr( + api_util, + "get_source_stream_names", + MagicMock( + return_value=["users", "orders"] + if discovered_streams is None + else discovered_streams + ), + ) + monkeypatch.setattr(project, "_published_manifest", lambda _id: MANIFEST) + return project, workspace + + +def test_run_full_sync_creates_noop_destination_and_connection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + project, workspace = _full_sync_project( + monkeypatch, + destinations=[], + connections=[], + sources=[ + SimpleNamespace(connector_id="other", definition_id="other-def"), + SimpleNamespace(connector_id="src", definition_id="def"), + ], + ) + + result = project._run_full_sync() + + project.publish.assert_called_once_with("My API", pre_validate=True) + workspace.list_sources.assert_called_once_with( + name="My API (Connector Builder test)" + ) + assert workspace.deploy_destination.call_args.kwargs["name"] == ( + connector_builder._util.CONNECTOR_BUILDER_TEST_DESTINATION_NAME + ) + workspace.deploy_connection.assert_called_once_with( + "My API (Connector Builder full sync)", + source="src", + destination="new-dest", + selected_streams=["users", "orders"], + ) + assert result.model_dump() == { + "source_id": "src", + "destination_id": "new-dest", + "connection_id": "conn", + "destination_created": True, + "connection_created": True, + "job_id": 7, + "job_url": "https://job", + } + + +def test_run_full_sync_reports_created_resources_before_sync_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + project, workspace = _full_sync_project( + monkeypatch, destinations=[], connections=[] + ) + workspace.deploy_connection.return_value.run_sync.side_effect = RuntimeError("boom") + created: list[str] = [] + + with pytest.raises(RuntimeError, match="boom"): + project._run_full_sync(test_source_id="src", on_resource_created=created.append) + + assert created == ["new-dest", "conn"] + + +def test_run_full_sync_reuses_destination_and_connection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + existing = MagicMock(connection_id="old", source_id="src", destination_id="dest") + existing.run_sync.return_value = SimpleNamespace(job_id=8, job_url="https://job8") + project, workspace = _full_sync_project( + monkeypatch, + destinations=[ + SimpleNamespace( + connector_id="dest", + definition_id=api_util.DEV_NULL_DESTINATION_DEFINITION_ID, + ) + ], + connections=[existing], + ) + + result = project._run_full_sync(test_source_id="src") + + workspace.deploy_destination.assert_not_called() + workspace.deploy_connection.assert_not_called() + existing.set_selected_streams.assert_called_once_with(["users", "orders"]) + existing.run_sync.assert_called_once_with(wait=False) + workspace.get_source.assert_called_once_with("src") + assert (result.connection_id, result.connection_created, result.job_id) == ( + "old", + False, + 8, + ) + + +def test_run_full_sync_without_test_source_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + project, workspace = _full_sync_project( + monkeypatch, destinations=[], connections=[] + ) + + with pytest.raises(AirbyteLibInputError, match="No test source"): + project._run_full_sync() + + project.publish.assert_not_called() + workspace.deploy_destination.assert_not_called() + + +@pytest.mark.parametrize( + ("stream_names", "record_count", "status", "expect_sync"), + [ + pytest.param([], 0, 200, True, id="no-streams"), + pytest.param(["users"], 1, 200, True, id="no-failed-streams"), + pytest.param(["users"], 0, 200, True, id="zero-records-no-error"), + pytest.param(["users"], 1, 422, False, id="http-error-skips-sync"), + ], +) +def test_readiness_full_sync_gating( + monkeypatch: pytest.MonkeyPatch, + stream_names: list[str], + record_count: int, + status: int, + expect_sync: bool, +) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + full_sync = _models.ConnectorBuilderFullSync( + source_id="src", + destination_id="dest", + connection_id="conn", + destination_created=False, + connection_created=True, + job_id=1, + job_url="u", + ) + run_full_sync = MagicMock(return_value=full_sync) + monkeypatch.setattr(project, "_run_full_sync", run_full_sync) + publish = MagicMock() + monkeypatch.setattr(project, "publish", publish) + monkeypatch.setattr( + project, + "_run_stream_test_read", + lambda _m, name, **_: _stream_read( + name, record_count=record_count, status=status + ), + ) + + report = project._run_readiness_checks( + stream_names, run_full_sync=True, test_source_id="src" + ) + + assert (report.full_sync is not None) is expect_sync + if expect_sync: + run_full_sync.assert_called_once_with( + test_source_id="src", on_resource_created=None + ) + else: + assert report.full_sync is None + assert report.failed_streams == ["users"] + run_full_sync.assert_not_called() + publish.assert_not_called() + + +def test_readiness_requires_streams_without_full_sync( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft=MANIFEST) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + + with pytest.raises(AirbyteLibInputError, match="run_full_sync"): + project._run_readiness_checks([]) + + +def test_run_full_sync_selects_streams_discovered_after_publish( + monkeypatch: pytest.MonkeyPatch, +) -> None: + project, workspace = _full_sync_project( + monkeypatch, + destinations=[ + SimpleNamespace( + connector_id="dest", + definition_id=api_util.DEV_NULL_DESTINATION_DEFINITION_ID, + ) + ], + connections=[], + discovered_streams=["users", "tables_a", "tables_b"], + ) + calls: list[str] = [] + project.publish.side_effect = lambda *_, **__: calls.append("publish") + api_util.get_source_stream_names.side_effect = lambda *_, **__: ( + calls.append("discover"), + ["users", "tables_a", "tables_b"], + )[1] + + project._run_full_sync(test_source_id="src") + + assert calls == ["publish", "discover"] + assert api_util.get_source_stream_names.call_args.args == ("src",) + assert workspace.deploy_connection.call_args.kwargs["selected_streams"] == [ + "users", + "tables_a", + "tables_b", + ] + + +def test_run_full_sync_rejects_source_without_streams( + monkeypatch: pytest.MonkeyPatch, +) -> None: + project, workspace = _full_sync_project( + monkeypatch, destinations=[], connections=[], discovered_streams=[] + ) + + with pytest.raises(AirbyteLibInputError, match="no streams"): + project._run_full_sync(test_source_id="src") + + workspace.deploy_connection.assert_not_called() + + +def test_run_full_sync_without_draft_syncs_published_version( + monkeypatch: pytest.MonkeyPatch, +) -> None: + project, workspace = _full_sync_project( + monkeypatch, + destinations=[ + SimpleNamespace( + connector_id="dest", + definition_id=api_util.DEV_NULL_DESTINATION_DEFINITION_ID, + ) + ], + connections=[], + draft=None, + ) + + project._run_full_sync(test_source_id="src") + + project.publish.assert_not_called() + workspace.deploy_connection.assert_called_once() + + +def test_run_full_sync_rejects_source_of_other_definition( + monkeypatch: pytest.MonkeyPatch, +) -> None: + project, workspace = _full_sync_project( + monkeypatch, destinations=[], connections=[] + ) + workspace.get_source.return_value = SimpleNamespace(definition_id="other-def") + + with pytest.raises(AirbyteLibInputError, match="published definition"): + project._run_full_sync(test_source_id="src") + + workspace.deploy_connection.assert_not_called() + + +def test_create_test_source_after_publish_cleared_draft( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store( + monkeypatch, + draft=None, + source_definition_id="def", + testing_values={"api_key": "**********", "region": "us"}, + ) + workspace = _fake_workspace() + workspace.deploy_source = MagicMock() + project = CloudConnectorBuilderProject(workspace=workspace, builder_project_id="bp") + monkeypatch.setattr(project, "_published_manifest", lambda _id: _secret_manifest()) + + project._get_or_create_default_test_source() + + assert workspace.deploy_source.call_args.kwargs["defer_credentials"] is True + + +def test_run_readiness_checks_full_sync_only_without_draft( + monkeypatch: pytest.MonkeyPatch, +) -> None: + project, workspace = _full_sync_project( + monkeypatch, + destinations=[ + SimpleNamespace( + connector_id="dest", + definition_id=api_util.DEV_NULL_DESTINATION_DEFINITION_ID, + ) + ], + connections=[], + draft=None, + ) + + report = project._run_readiness_checks([], run_full_sync=True, test_source_id="src") + + assert report.full_sync is not None + report_data = report.model_dump() + assert "passed" not in report_data + assert report_data["full_sync"] is not None + assert report.full_sync.destination_created is False + project.publish.assert_not_called() + workspace.deploy_connection.assert_called_once() + draft_check = next( + check for check in report.checks if check.description == "Save a draft manifest" + ) + assert draft_check.result == "failed" + + +def test_run_full_sync_rejects_non_noop_destination_with_default_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + project, workspace = _full_sync_project( + monkeypatch, + destinations=[SimpleNamespace(connector_id="s3", definition_id="s3-def")], + connections=[], + ) + + with pytest.raises(AirbyteLibInputError, match="isn't a no-op"): + project._run_full_sync(test_source_id="src") + + workspace.deploy_destination.assert_not_called() + workspace.deploy_connection.assert_not_called() + + +def _patch_published_versions( + monkeypatch: pytest.MonkeyPatch, + *, + draft: dict[str, Any] | None, +) -> dict[str, Any]: + store = _patch_project_store(monkeypatch, draft=draft, source_definition_id="def") + versions = { + 1: {"version": "6.0.0", "streams": []}, + 2: {"version": "6.1.0", "streams": []}, + } + unversioned_get = api_util.get_connector_builder_project + + def fake_get(**kwargs: Any) -> dict[str, Any]: + if kwargs.get("version") is None: + return unversioned_get(**kwargs) + return { + "builderProject": {"name": "My API", "sourceDefinitionId": "def"}, + "declarativeManifest": { + "manifest": versions[kwargs["version"]], + "isDraft": False, + "version": kwargs["version"], + }, + } + + monkeypatch.setattr(api_util, "get_connector_builder_project", fake_get) + monkeypatch.setattr( + api_util, + "list_declarative_manifest_versions", + lambda **_: [ + {"version": 2, "isActive": True, "description": ""}, + {"version": 1, "isActive": False, "description": "first"}, + ], + ) + store["versions"] = versions + return store + + +def test_list_published_versions_sorts_oldest_first( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_published_versions(monkeypatch, draft=None) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + + versions = project.list_published_versions() + + assert [(v.version, v.is_active, v.description) for v in versions] == [ + (1, False, "first"), + (2, True, ""), + ] + + +def test_list_published_versions_unpublished_is_empty( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store(monkeypatch, draft={"version": "6.0.0"}) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + + assert project.list_published_versions() == [] + with pytest.raises(AirbyteLibInputError, match="never been published"): + project.get_published_manifest_dict(1) + + +def test_get_published_manifest_by_version(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_published_versions(monkeypatch, draft=None) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + + assert project.get_published_manifest_dict(1) == {"version": "6.0.0", "streams": []} + assert "version: 6.0.0" in (project.get_published_manifest_yaml(1) or "") + + +def test_revert_as_draft_replaces_draft(monkeypatch: pytest.MonkeyPatch) -> None: + store = _patch_published_versions(monkeypatch, draft={"version": "7.0.0"}) + workspace = _fake_workspace() + project = CloudConnectorBuilderProject(workspace=workspace, builder_project_id="bp") + create_version = MagicMock() + monkeypatch.setattr(api_util, "create_declarative_manifest_version", create_version) + + project.revert_to_published_version(1) + + assert store["draft"] == store["versions"][1] + create_version.assert_not_called() + + +def test_revert_publish_refuses_unpublished_draft_edits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_published_versions(monkeypatch, draft={"version": "7.0.0"}) + workspace = _fake_workspace() + workspace.get_custom_source_definition.return_value.manifest = {"version": "6.1.0"} + project = CloudConnectorBuilderProject(workspace=workspace, builder_project_id="bp") + create_version = MagicMock() + monkeypatch.setattr(api_util, "create_declarative_manifest_version", create_version) + + with pytest.raises(AirbyteLibInputError, match="unpublished draft edits"): + project.revert_to_published_version(1, as_draft=False) + + create_version.assert_not_called() + + +def test_revert_publish_publishes_version(monkeypatch: pytest.MonkeyPatch) -> None: + store = _patch_published_versions(monkeypatch, draft=None) + workspace = _fake_workspace() + project = CloudConnectorBuilderProject(workspace=workspace, builder_project_id="bp") + + create_version = MagicMock() + monkeypatch.setattr(api_util, "create_declarative_manifest_version", create_version) + + project.revert_to_published_version(1, as_draft=False) + + kwargs = create_version.call_args.kwargs + assert kwargs["manifest"] == store["versions"][1] + assert kwargs["version"] == 3 + assert kwargs["description"].startswith("Reverted to version 1 by PyAirbyte at ") + + +def test_merge_testing_values_keeps_saved_secrets( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_project_store( + monkeypatch, + draft=MANIFEST, + testing_values={"api_key": "**********", "region": "us"}, + ) + project = CloudConnectorBuilderProject( + workspace=_fake_workspace(), builder_project_id="bp" + ) + replace = MagicMock() + monkeypatch.setattr(project, "replace_testing_values", replace) + + project._merge_testing_values({"region": "eu"}) + + replace.assert_called_once_with({"api_key": "**********", "region": "eu"}) diff --git a/tests/unit_tests/test_mcp_arg_trace_parity.py b/tests/unit_tests/test_mcp_arg_trace_parity.py index a30295b3a..0249924b8 100644 --- a/tests/unit_tests/test_mcp_arg_trace_parity.py +++ b/tests/unit_tests/test_mcp_arg_trace_parity.py @@ -70,6 +70,7 @@ def test_inventory_matches_fixture() -> None: assert all( "manifest_yaml_file" in plan[name] for name in ( + "create_connector_builder_project", "publish_custom_source_definition", "update_custom_source_definition", ) @@ -77,9 +78,19 @@ def test_inventory_matches_fixture() -> None: FIXTURE.write_text(json.dumps(plan, indent=2, sort_keys=True) + "\n") expected = json.loads(FIXTURE.read_text()) - for name in ("publish_custom_source_definition", "update_custom_source_definition"): + for name in ( + "create_connector_builder_project", + "publish_custom_source_definition", + "update_custom_source_definition", + ): if "manifest_yaml_file" not in plan[name]: expected[name].pop("manifest_yaml_file", None) + for name in ( + "publish_custom_source_definition", + "update_custom_source_definition", + ): + if "testing_values_secret_name" not in plan[name]: + expected[name].pop("testing_values_secret_name", None) assert plan == expected diff --git a/tests/unit_tests/test_mcp_cloud.py b/tests/unit_tests/test_mcp_cloud.py index b84ec280c..68efb00d1 100644 --- a/tests/unit_tests/test_mcp_cloud.py +++ b/tests/unit_tests/test_mcp_cloud.py @@ -38,7 +38,9 @@ JobStatusEnum, ) from airbyte.constants import MCP_DISABLE_DELETE_GUARD_ENV_VAR, MCP_INSIDERS_ENV_VAR +from airbyte.cloud.connector_builder import _models as _builder_models from airbyte.mcp import cloud as cloud_mcp +from airbyte.mcp import connector_builder as connector_builder_mcp from airbyte.mcp import guidance as guidance_mcp from airbyte.mcp._arg_resolvers import resolve_list_of_dicts from airbyte.mcp._tool_utils import _TOOL_POLICIES # noqa: PLC2701 @@ -3570,3 +3572,93 @@ def test_request_cloud_search_index_backfill_is_a_write_tool() -> None: policy = _TOOL_POLICIES["request_cloud_search_index_backfill"] assert policy.external_access is True assert policy.pipeline_change is False + + +def test_connector_builder_tools_return_read_warnings( + monkeypatch: pytest.MonkeyPatch, +) -> None: + warnings = ["Read 100 records, a multiple of 10. Check the paginator."] + read_result = _builder_models.ConnectorBuilderStreamReadResult( + stream_name="users", + record_count=0, + slice_count=0, + page_count=0, + test_read_limit_reached=False, + records=[], + pages=[], + logs=[], + warnings=warnings, + ) + readiness_report = _builder_models.ConnectorBuilderReadinessReport( + checks=[], + next_steps=[], + streams=[ + _builder_models.ConnectorBuilderStreamReadiness( + stream_name="users", + failed=False, + warnings=warnings, + ) + ], + ) + project_info = _builder_models.ConnectorBuilderProjectInfo( + builder_project_id="builder-project", + definition_id=None, + connector_builder_project_url="https://cloud.airbyte.com/workspaces/ws/connector-builder/edit/bp", + name="Users", + has_draft=True, + draft_manifest_dict={"version": "6.0.0"}, + draft_manifest_yaml="version: 6.0.0\n", + published_manifest_dict=None, + ) + project = SimpleNamespace( + builder_project_id="builder-project", + connector_builder_project_url="https://cloud.airbyte.com/workspaces/ws/connector-builder/edit/bp", + run_stream_test_read=lambda *_args, **_kwargs: read_result, + _run_readiness_checks=lambda *_args, **_kwargs: readiness_report, + _get_info=lambda: project_info, + ) + workspace = SimpleNamespace( + get_connector_builder_project=lambda *_args, **_kwargs: project + ) + monkeypatch.setattr( + connector_builder_mcp, "_get_cloud_workspace", lambda *_args: workspace + ) + ctx = cast(Context, object()) + + stream_result = connector_builder_mcp.run_connector_builder_stream_test_read( + ctx, + "users", + builder_project_id="builder-project", + workspace_id="ws", + ) + readiness_result = connector_builder_mcp.run_connector_builder_readiness_checks( + ctx, + ["users"], + builder_project_id="builder-project", + workspace_id="ws", + ) + + draft_result = connector_builder_mcp.get_connector_builder_draft_manifest( + ctx, + builder_project_id="builder-project", + workspace_id="ws", + ) + + assert stream_result.warnings == warnings + assert stream_result.builder_project_id == "builder-project" + assert readiness_result.streams[0].warnings == warnings + assert "passed" not in readiness_result.model_dump() + assert draft_result.draft_manifest == {"version": "6.0.0"} + assert draft_result.published_manifest is None + + +def test_redact_stream_diagnostics() -> None: + leaky = "GET https://api.example.com/data?api_key=s3cret failed: Bearer s3cret" + data = connector_builder_mcp._redact_stream_diagnostics({ + "pages": [{"url": leaky, "response_body": leaky}], + "http_errors": [{"url": leaky, "response_body": None}], + "logs": [{"level": "WARN", "message": leaky}], + "problem_logs": [{"level": "ERROR", "message": leaky}], + "error": leaky, + }) + assert "s3cret" not in str(data) diff --git a/tests/unit_tests/test_mcp_docs_includes.py b/tests/unit_tests/test_mcp_docs_includes.py index e8d311baa..d93dc5911 100644 --- a/tests/unit_tests/test_mcp_docs_includes.py +++ b/tests/unit_tests/test_mcp_docs_includes.py @@ -49,6 +49,7 @@ def test_existing_includes_name_generated_mcp_modules() -> None: repo_root = Path(__file__).parents[2] generated_modules = { "cloud", + "connector_builder", "local", "interactive", "registry", diff --git a/tests/unit_tests/test_mcp_tool_utils.py b/tests/unit_tests/test_mcp_tool_utils.py index 6970e6591..1087696bc 100644 --- a/tests/unit_tests/test_mcp_tool_utils.py +++ b/tests/unit_tests/test_mcp_tool_utils.py @@ -295,7 +295,17 @@ def test_server_tools_register_expected_policies() -> None: insiders_tools = { name for name, policy in _TOOL_POLICIES.items() if policy.insiders_only } - assert insiders_tools == set() + assert insiders_tools == { + "create_connector_builder_project", + "update_connector_builder_manifest_text", + "set_connector_builder_testing_values", + "run_connector_builder_stream_test_read", + "get_connector_builder_resolved_manifest", + "run_connector_builder_readiness_checks", + "list_connector_builder_published_versions", + "get_connector_builder_published_manifest", + "revert_connector_builder_published_version", + } @pytest.mark.parametrize( diff --git a/tests/unit_tests/test_mcp_trusted_execution.py b/tests/unit_tests/test_mcp_trusted_execution.py index 1dafff09a..acfa311d7 100644 --- a/tests/unit_tests/test_mcp_trusted_execution.py +++ b/tests/unit_tests/test_mcp_trusted_execution.py @@ -18,7 +18,7 @@ AirbyteLibInputError, AirbyteTrustedExecutionRequiredError, ) -from airbyte.mcp import cloud, local +from airbyte.mcp import cloud, connector_builder, local from airbyte.mcp._arg_resolvers import resolve_connector_config, resolve_manifest_yaml from airbyte.mcp._guards import ( is_trusted_execution_enabled, @@ -261,6 +261,7 @@ def test_manifest_file_schema_depends_on_trusted_execution( _set_trusted(monkeypatch, enabled=enabled) app = FastMCP(f"manifest-schema-{enabled}") cloud.register_cloud_tools(app) + connector_builder.register_connector_builder_tools(app) async def list_tools() -> dict[str, object]: async with Client(app) as client: @@ -269,6 +270,7 @@ async def list_tools() -> dict[str, object]: tools = asyncio.run(list_tools()) for tool_name in ( + "create_connector_builder_project", "publish_custom_source_definition", "update_custom_source_definition", ): @@ -287,6 +289,22 @@ async def list_tools() -> dict[str, object]: not in properties["manifest_yaml_file"]["description"] ) + for tool_name in ( + "publish_custom_source_definition", + "update_custom_source_definition", + ): + tool = tools[tool_name] + properties = tool.input_schema["properties"] # type: ignore[attr-defined] + assert "testing_values" in properties + assert ("testing_values_secret_name" in properties) is enabled + assert ( + "secret_reference::" not in properties["testing_values"]["description"] + ) + assert ( + cloud.TESTING_VALUES_SECRET_HINT + in (getattr(tool, "description", None) or "") + ) is enabled + @pytest.mark.parametrize("enabled", [False, True], ids=["untrusted", "trusted"]) def test_publish_missing_manifest_error_depends_on_trusted_execution(