airbyte_ops_mcp.regression_tests

46class Command(Enum):
47    """Airbyte connector commands."""
48
49    CHECK = "check"
50    DISCOVER = "discover"
51    READ = "read"
52    READ_WITH_STATE = "read-with-state"
53    SPEC = "spec"
54
55    def needs_config(self) -> bool:
56        return self in {
57            Command.CHECK,
58            Command.DISCOVER,
59            Command.READ,
60            Command.READ_WITH_STATE,
61        }
62
63    def needs_catalog(self) -> bool:
64        return self in {Command.READ, Command.READ_WITH_STATE}
65
66    def needs_state(self) -> bool:
67        return self in {Command.READ_WITH_STATE}

27@dataclass
28class ConnectionData:
29    """Data fetched from an Airbyte Cloud connection."""
30
31    connection_id: str
32    source_id: str
33    source_name: str
34    source_definition_id: str
35    config: dict[str, Any]
36    catalog: dict[str, Any]
37    stream_names: list[str]
38    workspace_id: str | None = None
39    docker_repository: str | None = None
40    docker_image_tag: str | None = None
41    state: list[dict[str, Any]] | None = field(default=None)
42
43    @property
44    def connector_image(self) -> str | None:
45        """Get the full connector image name with tag."""
46        if self.docker_repository and self.docker_image_tag:
47            return f"{self.docker_repository}:{self.docker_image_tag}"
48        return None

 84@dataclass
 85class ConnectorUnderTest:
 86    """Represents a connector being tested.
 87
 88    In validation tests, there would be one connector under test.
 89    When running regression tests, there would be two connectors under test:
 90    the target and the control versions of the same connector.
 91    """
 92
 93    image_name: str
 94    target_or_control: TargetOrControl
 95
 96    @property
 97    def name(self) -> str:
 98        """Get connector name without registry prefix."""
 99        return self.image_name.replace("airbyte/", "").split(":")[0]
100
101    @property
102    def name_without_type_prefix(self) -> str:
103        """Get connector name without actor type prefix."""
104        return self.name.replace(f"{self.actor_type.value}-", "")
105
106    @property
107    def version(self) -> str:
108        """Get connector version from image tag."""
109        return self.image_name.replace("airbyte/", "").split(":")[1]
110
111    @property
112    def actor_type(self) -> ActorType:
113        """Infer actor type from image name."""
114        if "airbyte/destination-" in self.image_name:
115            return ActorType.DESTINATION
116        elif "airbyte/source-" in self.image_name:
117            return ActorType.SOURCE
118        else:
119            raise ValueError(
120                f"Can't infer the actor type. Connector image name {self.image_name} "
121                "does not contain 'airbyte/source' or 'airbyte/destination'"
122            )
123
124    @classmethod
125    def from_image_name(
126        cls,
127        image_name: str,
128        target_or_control: TargetOrControl,
129    ) -> ConnectorUnderTest:
130        """Create a ConnectorUnderTest from an image name."""
131        return cls(image_name, target_or_control)

156@dataclass
157class ExecutionResult:
158    """Result of executing a connector command."""
159
160    connector_under_test: ConnectorUnderTest
161    command: Command
162    stdout_file_path: Path
163    stderr_file_path: Path
164    success: bool
165    exit_code: int
166    configured_catalog: ConfiguredAirbyteCatalog | None = None
167    config: dict[str, Any] | None = None
168    _airbyte_messages: list[AirbyteMessage] = field(default_factory=list)
169    _messages_loaded: bool = field(default=False, repr=False)
170
171    @property
172    def logger(self) -> logging.Logger:
173        return logging.getLogger(
174            f"{self.connector_under_test.target_or_control.value}-{self.command.value}"
175        )
176
177    @cached_property
178    def airbyte_messages(self) -> list[AirbyteMessage]:
179        """Parse and return all Airbyte messages from stdout."""
180        if self._messages_loaded:
181            return self._airbyte_messages
182
183        messages = []
184        for line in self.stdout_file_path.read_text().splitlines():
185            line = line.strip()
186            if not line:
187                continue
188            with contextlib.suppress(ValidationError):
189                messages.append(AirbyteMessage.parse_raw(line))
190        self._airbyte_messages = messages
191        self._messages_loaded = True
192        return messages
193
194    @property
195    def configured_streams(self) -> list[str]:
196        """Get list of configured stream names."""
197        if not self.configured_catalog:
198            return []
199        return [stream.stream.name for stream in self.configured_catalog.streams]
200
201    def get_records(self) -> Iterator[AirbyteMessage]:
202        """Iterate over record messages."""
203        for message in self.airbyte_messages:
204            if message.type is AirbyteMessageType.RECORD:
205                yield message
206
207    def get_records_per_stream(self, stream: str) -> Iterator[AirbyteMessage]:
208        """Get records for a specific stream."""
209        for message in self.get_records():
210            if message.record.stream == stream:
211                yield message
212
213    def get_states(self) -> Iterator[AirbyteMessage]:
214        """Iterate over state messages."""
215        for message in self.airbyte_messages:
216            if message.type is AirbyteMessageType.STATE:
217                yield message
218
219    def get_final_state_per_stream(self) -> dict[tuple[str | None, str], Any]:
220        """The state each stream would resume from, keyed by stream identity.
221
222        A connector emits state as it goes, so only the last message for a
223        stream describes where the next sync starts -- earlier ones are
224        checkpoints it has already moved past. Messages are walked in emission
225        order and later ones overwrite earlier ones, which is exactly how the
226        platform stores them.
227
228        Keyed by `(namespace, name)` rather than by the `namespace.name` label it
229        displays as, for the reason the catalog index is: a stream literally
230        called `public.users` must not share a key with `users` in the `public`
231        namespace and silently take the other's state with it. The comparison
232        turns these into the labels a reviewer reads, qualifying the collision
233        instead of collapsing it.
234
235        The three state types are all kept, because all three are what a
236        connector resumes from:
237
238        - `STREAM` -- one entry per stream, under its own identity.
239        - `GLOBAL` -- unpacked into one entry per stream, plus the shared state
240          under `SHARED_STATE_KEY`.
241        - `LEGACY` -- the protocol carries no stream identity here, so the whole
242          blob is one entry under `LEGACY_STATE_KEY`. Skipping it would let a
243          comparison report "state unchanged" over a state nothing looked at.
244
245        `sourceStats` and `destinationStats` are left out: they live beside the
246        state rather than in it, and `recordCount` legitimately differs between
247        two runs -- record counts are compared in their own table.
248
249        Returns:
250            Stream identity to state value, empty when the run emitted no state.
251        """
252        final_states: dict[tuple[str | None, str], Any] = {}
253
254        for message in self.get_states():
255            state = message.state
256            if state is None:
257                continue
258
259            if state.type is AirbyteStateType.GLOBAL and state.global_ is not None:
260                for stream_state in state.global_.stream_states or []:
261                    stream_id, value = self._stream_state_entry(stream_state)
262                    final_states[stream_id] = value
263                if state.global_.shared_state is not None:
264                    final_states[SHARED_STATE_KEY] = to_plain_dict(
265                        state.global_.shared_state, exclude_none=False
266                    )
267            elif state.type is AirbyteStateType.STREAM and state.stream is not None:
268                stream_id, value = self._stream_state_entry(state.stream)
269                final_states[stream_id] = value
270            elif state.data is not None:
271                # LEGACY, and a message that declared no type at all: the
272                # protocol's own default, and still what the connector resumes
273                # from.
274                final_states[LEGACY_STATE_KEY] = state.data
275
276        return final_states
277
278    @staticmethod
279    def _stream_state_entry(
280        stream_state: AirbyteStreamState,
281    ) -> tuple[tuple[str | None, str], Any]:
282        """One stream's state, under the identity comparisons key streams by."""
283        descriptor = stream_state.stream_descriptor
284
285        return (
286            to_stream_id(descriptor.namespace, descriptor.name),
287            # `exclude_none=False`: a `null` inside a state blob is the
288            # connector's own data, not an unset protocol optional. Dropping the
289            # key would report a cursor one version emits as `null` and the
290            # other omits as an unchanged state.
291            to_plain_dict(stream_state.stream_state, exclude_none=False),
292        )
293
294    def get_message_count_per_type(self) -> dict[AirbyteMessageType, int]:
295        """Count messages by type."""
296        counts: dict[AirbyteMessageType, int] = defaultdict(int)
297        for message in self.airbyte_messages:
298            counts[message.type] += 1
299        return dict(counts)
300
301    def get_record_count_per_stream(self) -> dict[str, int]:
302        """Count records by stream name.
303
304        Returns:
305            Dictionary mapping stream names to record counts.
306        """
307        counts: dict[str, int] = defaultdict(int)
308        for message in self.get_records():
309            counts[message.record.stream] += 1
310        return dict(counts)
311
312    def get_catalog(self) -> AirbyteCatalog | None:
313        """Get discovered catalog from messages."""
314        for message in self.airbyte_messages:
315            if message.type is AirbyteMessageType.CATALOG:
316                return message.catalog
317        return None
318
319    def get_spec(self) -> Any | None:
320        """Get connector spec from messages."""
321        for message in self.airbyte_messages:
322            if message.type is AirbyteMessageType.SPEC:
323                return message.spec
324        return None
325
326    def get_connection_status(self) -> Any | None:
327        """Get connection status from check command."""
328        for message in self.airbyte_messages:
329            if message.type is AirbyteMessageType.CONNECTION_STATUS:
330                return message.connectionStatus
331        return None
332
333    def is_check_successful(self) -> bool:
334        """Check if the check command was successful."""
335        status = self.get_connection_status()
336        if status is None:
337            return False
338        return status.status.value == "SUCCEEDED"
339
340    def save_artifacts(self, output_dir: Path) -> None:
341        """Save execution artifacts to the output directory."""
342        output_dir.mkdir(parents=True, exist_ok=True)
343
344        airbyte_messages_dir = output_dir / "airbyte_messages"
345        airbyte_messages_dir.mkdir(parents=True, exist_ok=True)
346
347        messages_by_type: dict[str, list[str]] = defaultdict(list)
348        for message in self.airbyte_messages:
349            type_name = message.type.value.lower()
350            messages_by_type[type_name].append(message.model_dump_json())
351
352        for type_name, messages in messages_by_type.items():
353            file_path = airbyte_messages_dir / f"{type_name}.jsonl"
354            file_path.write_text("\n".join(messages))
355
356        # Save configured catalog (input) if available
357        if self.configured_catalog is not None:
358            catalog_path = output_dir / "configured_catalog.json"
359            catalog_path.write_text(self.configured_catalog.model_dump_json(indent=2))
360            self.logger.info(f"Saved configured catalog to {catalog_path}")
361
362        self.logger.info(f"Artifacts saved to {output_dir}")

195class SecretRetrievalError(Exception):
196    """Raised when secret retrieval fails.
197
198    This exception is raised when USE_CONNECTION_SECRET_RETRIEVER is enabled
199    but secrets cannot be retrieved (e.g., EU data residency restrictions).
200    """

70class TargetOrControl(Enum):
71    """Indicates whether a connector is the target (new) or control (baseline) version."""
72
73    TARGET = "target"
74    CONTROL = "control"

203def enrich_config_with_secrets(
204    connection_data: ConnectionData,
205    retrieval_reason: str = "MCP live tests",
206    raise_on_failure: bool = True,
207) -> ConnectionData:
208    """Enrich connection data with unmasked secrets from internal retriever.
209
210    This function takes a ConnectionData object (typically from the public
211    Cloud API with masked secrets) and replaces the config with unmasked
212    secrets from the internal connection-retriever.
213
214    Args:
215        connection_data: The connection data to enrich.
216        retrieval_reason: Reason for retrieval (for audit logging).
217        raise_on_failure: If True (default), raise SecretRetrievalError when
218            secrets cannot be retrieved. If False, return the original
219            connection_data with masked secrets (legacy behavior).
220
221    Returns:
222        A new ConnectionData with unmasked config.
223
224    Raises:
225        SecretRetrievalError: If raise_on_failure is True and secrets cannot
226            be retrieved (e.g., due to EU data residency restrictions).
227    """
228    unmasked_config = retrieve_unmasked_config(
229        connection_id=connection_data.connection_id,
230        retrieval_reason=retrieval_reason,
231    )
232
233    if unmasked_config is None:
234        error_msg = (
235            "Could not retrieve unmasked secrets for connection "
236            f"{connection_data.connection_id}. This may be due to EU data "
237            "residency restrictions or database connectivity issues. "
238            "The connection's credentials cannot be used for regression testing."
239        )
240        logger.warning(error_msg)
241        if raise_on_failure:
242            raise SecretRetrievalError(error_msg)
243        return connection_data
244
245    logger.info(
246        f"Successfully enriched config with unmasked secrets for "
247        f"{connection_data.connection_id}"
248    )
249
250    # Return a new ConnectionData with the unmasked config
251    return replace(connection_data, config=unmasked_config)

 75def fetch_connection_data(
 76    connection_id: str,
 77    client_id: str | None = None,
 78    client_secret: str | None = None,
 79) -> ConnectionData:
 80    """Fetch connection configuration and catalog from Airbyte Cloud.
 81
 82    Args:
 83        connection_id: The connection ID to fetch data for.
 84        client_id: Airbyte Cloud client ID (defaults to env var).
 85        client_secret: Airbyte Cloud client secret (defaults to env var).
 86
 87    Returns:
 88        ConnectionData with config and catalog.
 89
 90    Raises:
 91        PyAirbyteInputError: If the API request fails.
 92    """
 93    client_id = client_id or os.getenv("AIRBYTE_CLOUD_CLIENT_ID")
 94    client_secret = client_secret or os.getenv("AIRBYTE_CLOUD_CLIENT_SECRET")
 95
 96    if not client_id or not client_secret:
 97        raise PyAirbyteInputError(
 98            message="Missing Airbyte Cloud credentials",
 99            context={
100                "hint": "Set AIRBYTE_CLOUD_CLIENT_ID and AIRBYTE_CLOUD_CLIENT_SECRET env vars"
101            },
102        )
103
104    access_token = _get_access_token(client_id, client_secret)
105    public_api_root = constants.CLOUD_API_ROOT
106    headers = {
107        "Authorization": f"Bearer {access_token}",
108        "Content-Type": "application/json",
109    }
110
111    # Get connection details
112    conn_response = requests.get(
113        f"{public_api_root}/connections/{connection_id}",
114        headers=headers,
115        timeout=30,
116    )
117
118    if conn_response.status_code != 200:
119        raise PyAirbyteInputError(
120            message=f"Failed to get connection: {conn_response.status_code}",
121            context={"connection_id": connection_id, "response": conn_response.text},
122        )
123
124    conn_data = conn_response.json()
125    source_id = conn_data["sourceId"]
126
127    # Get source details (includes config)
128    source_response = requests.get(
129        f"{public_api_root}/sources/{source_id}",
130        headers=headers,
131        timeout=30,
132    )
133
134    if source_response.status_code != 200:
135        raise PyAirbyteInputError(
136            message=f"Failed to get source: {source_response.status_code}",
137            context={"source_id": source_id, "response": source_response.text},
138        )
139
140    source_data = source_response.json()
141    source_definition_id = source_data.get("definitionId", "")
142
143    # Try to get docker repository and image tag from source definition version
144    docker_repository = None
145    docker_image_tag = None
146    if source_definition_id:
147        try:
148            # Use the Config API to get version info for the source
149            config_api_root = constants.CLOUD_CONFIG_API_ROOT
150            version_response = requests.post(
151                f"{config_api_root}/actor_definition_versions/get_for_source",
152                json={"sourceId": source_id},
153                headers=headers,
154                timeout=30,
155            )
156            if version_response.status_code == 200:
157                version_data = version_response.json()
158                docker_repository = version_data.get("dockerRepository")
159                docker_image_tag = version_data.get("dockerImageTag")
160        except Exception:
161            # Non-fatal: we can still proceed without docker info
162            pass
163
164    # Build configured catalog from connection streams
165    streams_config = conn_data.get("configurations", {}).get("streams", [])
166    stream_names = [s["name"] for s in streams_config]
167
168    # Build Airbyte protocol catalog format
169    catalog = _build_configured_catalog(
170        streams_config, source_id, headers, public_api_root
171    )
172
173    return ConnectionData(
174        connection_id=connection_id,
175        source_id=source_id,
176        source_name=source_data.get("name", ""),
177        source_definition_id=source_definition_id,
178        config=source_data.get("configuration", {}),
179        catalog=catalog,
180        stream_names=stream_names,
181        workspace_id=conn_data.get("workspaceId"),
182        docker_repository=docker_repository,
183        docker_image_tag=docker_image_tag,
184    )

54def is_secret_retriever_enabled() -> bool:
55    """Check if secret retrieval is enabled via environment variable.
56
57    Returns:
58        True if USE_CONNECTION_SECRET_RETRIEVER is set to a truthy value.
59    """
60    value = os.getenv(ENV_USE_SECRET_RETRIEVER, "").lower()
61    return value in ("true", "1", "yes")

 73def retrieve_unmasked_config(
 74    connection_id: str,
 75    retrieval_reason: str = "MCP live tests",
 76) -> dict | None:
 77    """Retrieve unmasked source config from vendored connection-retriever.
 78
 79    This function directly queries the internal Postgres database to get
 80    the source configuration with unmasked secrets.
 81
 82    Args:
 83        connection_id: The Airbyte Cloud connection ID.
 84        retrieval_reason: Reason for retrieval (for audit logging).
 85
 86    Returns:
 87        The unmasked source config dict, or None if retrieval fails.
 88    """
 89    # Only request the source config - that's all we need for secrets
 90    requested_objects = [ConnectionObject.SOURCE_CONFIG]
 91
 92    candidates = retrieve_objects(
 93        connection_objects=requested_objects,
 94        retrieval_reason=retrieval_reason,
 95        connection_id=connection_id,
 96    )
 97
 98    if not candidates:
 99        logger.warning(
100            f"No connection data found for connection ID {connection_id} "
101            "via connection-retriever"
102        )
103        return None
104
105    candidate = candidates[0]
106    if candidate.source_config:
107        return dict(candidate.source_config)
108
109    return None

64def should_use_secret_retriever() -> bool:
65    """Check if secret retrieval should be used.
66
67    Returns:
68        True if USE_CONNECTION_SECRET_RETRIEVER env var is set to a truthy value.
69    """
70    return is_secret_retriever_enabled()