airbyte-internal-ops 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

airbyte_ops_mcp/_legacy/airbyte_ci/connector_live_tests/commons/mitm_addons.py

@@ -1,37 +0,0 @@
-# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
-from __future__ import annotations
-
-from urllib.parse import parse_qs, urlencode, urlparse
-
-from mitmproxy import http
-
-
-class SortQueryParams:
-    """This addon sorts query parameters in the request URL.
-    It is useful for testing purposes, as it makes it easier to compare requests and get cache hits.
-    """
-
-    def request(self, flow: http.HTTPFlow) -> None:
-        if url := flow.request.url:
-            parsed_url = urlparse(url)
-            # Get query parameters as dictionary
-            query_params = parse_qs(parsed_url.query)
-            # Sort query parameters alphabetically
-            sorted_params = {
-                key: query_params[key] for key in sorted(query_params.keys())
-            }
-            # Reconstruct the URL with sorted query parameters
-            sorted_url = parsed_url._replace(
-                query=urlencode(sorted_params, doseq=True)
-            ).geturl()
-
-            # Update the request URL
-            flow.request.url = sorted_url
-
-
-# Disabling the addon.
-# It can alter the request URL when some connector URL are already encoded.
-# See discussion here https://github.com/airbytehq/airbyte-internal-issues/issues/9302#issuecomment-2311854334
-
-# addons = [SortQueryParams()]
-addons = []

airbyte_ops_mcp/_legacy/airbyte_ci/connector_live_tests/commons/models.py

@@ -1,595 +0,0 @@
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-
-from __future__ import annotations
-
-import _collections_abc
-import json
-import logging
-import tempfile
-from collections import defaultdict
-from collections.abc import Iterable, Iterator, MutableMapping
-from dataclasses import dataclass, field
-from enum import Enum
-from functools import cache
-from pathlib import Path
-from typing import Any, Dict, List, Optional
-
-import dagger
-import requests
-from airbyte_protocol.models import (
-    AirbyteCatalog,  # type: ignore
-    AirbyteMessage,  # type: ignore
-    AirbyteStateMessage,  # type: ignore
-    AirbyteStreamStatusTraceMessage,  # type: ignore
-    ConfiguredAirbyteCatalog,  # type: ignore
-    TraceType,  # type: ignore
-)
-from airbyte_protocol.models import Type as AirbyteMessageType
-from genson import SchemaBuilder  # type: ignore
-from live_tests.commons.backends import DuckDbBackend, FileBackend
-from live_tests.commons.secret_access import get_airbyte_api_key
-from live_tests.commons.utils import (
-    get_connector_container,
-    get_http_flows_from_mitm_dump,
-    mitm_http_stream_to_har,
-    sanitize_stream_name,
-    sort_dict_keys,
-)
-from mitmproxy import http
-from pydantic import ValidationError
-
-
-class UserDict(_collections_abc.MutableMapping):  # type: ignore
-    # Start by filling-out the abstract methods
-    def __init__(self, _dict: Optional[MutableMapping] = None, **kwargs: Any):
-        self.data: MutableMapping = {}
-        if _dict is not None:
-            self.update(_dict)
-        if kwargs:
-            self.update(kwargs)
-
-    def __len__(self) -> int:
-        return len(self.data)
-
-    def __getitem__(self, key: Any) -> Any:
-        if key in self.data:
-            return self.data[key]
-        if hasattr(self.__class__, "__missing__"):
-            return self.__class__.__missing__(self, key)
-        raise KeyError(key)
-
-    def __setitem__(self, key: Any, item: Any) -> None:
-        self.data[key] = item
-
-    def __delitem__(self, key: Any) -> None:
-        del self.data[key]
-
-    def __iter__(self) -> Iterator:
-        return iter(self.data)
-
-    # Modify __contains__ to work correctly when __missing__ is present
-    def __contains__(self, key: Any) -> bool:
-        return key in self.data
-
-    # Now, add the methods in dicts but not in MutableMapping
-    def __repr__(self) -> str:
-        return repr(self.data)
-
-    def __or__(self, other: UserDict | dict) -> UserDict:
-        if isinstance(other, UserDict):
-            return self.__class__(self.data | other.data)  # type: ignore
-        if isinstance(other, dict):
-            return self.__class__(self.data | other)  # type: ignore
-        return NotImplemented
-
-    def __ror__(self, other: UserDict | dict) -> UserDict:
-        if isinstance(other, UserDict):
-            return self.__class__(other.data | self.data)  # type: ignore
-        if isinstance(other, dict):
-            return self.__class__(other | self.data)  # type: ignore
-        return NotImplemented
-
-    def __ior__(self, other: UserDict | dict) -> UserDict:
-        if isinstance(other, UserDict):
-            self.data |= other.data  # type: ignore
-        else:
-            self.data |= other  # type: ignore
-        return self
-
-    def __copy__(self) -> UserDict:
-        inst = self.__class__.__new__(self.__class__)
-        inst.__dict__.update(self.__dict__)
-        # Create a copy and avoid triggering descriptors
-        inst.__dict__["data"] = self.__dict__["data"].copy()
-        return inst
-
-    def copy(self) -> UserDict:
-        if self.__class__ is UserDict:
-            return UserDict(self.data.copy())  # type: ignore
-        import copy
-
-        data = self.data
-        try:
-            self.data = {}
-            c = copy.copy(self)
-        finally:
-            self.data = data
-        c.update(self)
-        return c
-
-    @classmethod
-    def fromkeys(cls, iterable: Iterable, value: Optional[Any] = None) -> UserDict:
-        d = cls()
-        for key in iterable:
-            d[key] = value
-        return d
-
-
-class SecretDict(UserDict):
-    def __str__(self) -> str:
-        return f"{self.__class__.__name__}(******)"
-
-    def __repr__(self) -> str:
-        return str(self)
-
-
-class Command(Enum):
-    CHECK = "check"
-    DISCOVER = "discover"
-    READ = "read"
-    READ_WITH_STATE = "read-with-state"
-    SPEC = "spec"
-
-    def needs_config(self) -> bool:
-        return self in {
-            Command.CHECK,
-            Command.DISCOVER,
-            Command.READ,
-            Command.READ_WITH_STATE,
-        }
-
-    def needs_catalog(self) -> bool:
-        return self in {Command.READ, Command.READ_WITH_STATE}
-
-    def needs_state(self) -> bool:
-        return self in {Command.READ_WITH_STATE}
-
-
-class TargetOrControl(Enum):
-    TARGET = "target"
-    CONTROL = "control"
-
-
-class ActorType(Enum):
-    SOURCE = "source"
-    DESTINATION = "destination"
-
-
-class ConnectionSubset(Enum):
-    """Signals which connection pool to consider for this live test — just the Airbyte sandboxes, or all possible connctions on Cloud."""
-
-    SANDBOXES = "sandboxes"
-    ALL = "all"
-
-
-@dataclass
-class ConnectorUnderTest:
-    """Represents a connector being tested.
-    In validation tests, there would be one connector under test.
-    When running regression tests, there would be two connectors under test: the target and the control versions of the same connector.
-    """
-
-    # connector image, assuming it's in the format "airbyte/{actor_type}-{connector_name}:{version}"
-    image_name: str
-    container: dagger.Container
-    target_or_control: TargetOrControl
-
-    @property
-    def name(self) -> str:
-        return self.image_name.replace("airbyte/", "").split(":")[0]
-
-    @property
-    def name_without_type_prefix(self) -> str:
-        return self.name.replace(f"{self.actor_type.value}-", "")
-
-    @property
-    def version(self) -> str:
-        return self.image_name.replace("airbyte/", "").split(":")[1]
-
-    @property
-    def actor_type(self) -> ActorType:
-        if "airbyte/destination-" in self.image_name:
-            return ActorType.DESTINATION
-        elif "airbyte/source-" in self.image_name:
-            return ActorType.SOURCE
-        else:
-            raise ValueError(
-                f"Can't infer the actor type. Connector image name {self.image_name} does not contain 'airbyte/source' or 'airbyte/destination'"
-            )
-
-    @classmethod
-    async def from_image_name(
-        cls: type[ConnectorUnderTest],
-        dagger_client: dagger.Client,
-        image_name: str,
-        target_or_control: TargetOrControl,
-    ) -> ConnectorUnderTest:
-        container = await get_connector_container(dagger_client, image_name)
-        return cls(image_name, container, target_or_control)
-
-
-@dataclass
-class ExecutionInputs:
-    hashed_connection_id: str
-    connector_under_test: ConnectorUnderTest
-    actor_id: str
-    global_output_dir: Path
-    command: Command
-    config: Optional[SecretDict] = None
-    configured_catalog: Optional[ConfiguredAirbyteCatalog] = None
-    state: Optional[dict] = None
-    environment_variables: Optional[dict] = None
-    duckdb_path: Optional[Path] = None
-
-    def raise_if_missing_attr_for_command(self, attribute: str) -> None:
-        if getattr(self, attribute) is None:
-            raise ValueError(
-                f"We need a {attribute} to run the {self.command.value} command"
-            )
-
-    def __post_init__(self) -> None:
-        if self.command is Command.CHECK:
-            self.raise_if_missing_attr_for_command("config")
-        if self.command is Command.DISCOVER:
-            self.raise_if_missing_attr_for_command("config")
-        if self.command is Command.READ:
-            self.raise_if_missing_attr_for_command("config")
-            self.raise_if_missing_attr_for_command("configured_catalog")
-        if self.command is Command.READ_WITH_STATE:
-            self.raise_if_missing_attr_for_command("config")
-            self.raise_if_missing_attr_for_command("configured_catalog")
-            self.raise_if_missing_attr_for_command("state")
-
-    @property
-    def output_dir(self) -> Path:
-        output_dir = (
-            self.global_output_dir
-            / f"command_execution_artifacts/{self.connector_under_test.name}/{self.command.value}/{self.connector_under_test.version}/{self.hashed_connection_id}"
-        )
-        output_dir.mkdir(parents=True, exist_ok=True)
-        return output_dir
-
-
-@dataclass
-class ExecutionResult:
-    hashed_connection_id: str
-    actor_id: str
-    configured_catalog: ConfiguredAirbyteCatalog
-    connector_under_test: ConnectorUnderTest
-    command: Command
-    stdout_file_path: Path
-    stderr_file_path: Path
-    success: bool
-    executed_container: Optional[dagger.Container]
-    config: Optional[SecretDict]
-    http_dump: Optional[dagger.File] = None
-    http_flows: list[http.HTTPFlow] = field(default_factory=list)
-    stream_schemas: Optional[dict[str, Any]] = None
-    backend: Optional[FileBackend] = None
-
-    HTTP_DUMP_FILE_NAME = "http_dump.mitm"
-    HAR_FILE_NAME = "http_dump.har"
-
-    @property
-    def logger(self) -> logging.Logger:
-        return logging.getLogger(
-            f"{self.connector_under_test.target_or_control.value}-{self.command.value}"
-        )
-
-    @property
-    def airbyte_messages(self) -> Iterable[AirbyteMessage]:
-        return self.parse_airbyte_messages_from_command_output(self.stdout_file_path)
-
-    @property
-    def duckdb_schema(self) -> Iterable[str]:
-        return (
-            self.connector_under_test.target_or_control.value,
-            self.command.value,
-            self.hashed_connection_id,
-        )
-
-    @property
-    def configured_streams(self) -> List[str]:
-        return [stream.stream.name for stream in self.configured_catalog.streams]
-
-    @property
-    def primary_keys_per_stream(self) -> Dict[str, List[str]]:
-        return {
-            stream.stream.name: stream.primary_key[0] if stream.primary_key else None
-            for stream in self.configured_catalog.streams
-        }
-
-    @classmethod
-    async def load(
-        cls: type[ExecutionResult],
-        connector_under_test: ConnectorUnderTest,
-        hashed_connection_id: str,
-        actor_id: str,
-        configured_catalog: ConfiguredAirbyteCatalog,
-        command: Command,
-        stdout_file_path: Path,
-        stderr_file_path: Path,
-        success: bool,
-        executed_container: Optional[dagger.Container],
-        config: Optional[SecretDict] = None,
-        http_dump: Optional[dagger.File] = None,
-    ) -> ExecutionResult:
-        execution_result = cls(
-            hashed_connection_id,
-            actor_id,
-            configured_catalog,
-            connector_under_test,
-            command,
-            stdout_file_path,
-            stderr_file_path,
-            success,
-            executed_container,
-            config,
-            http_dump,
-        )
-        await execution_result.load_http_flows()
-        return execution_result
-
-    async def load_http_flows(self) -> None:
-        if not self.http_dump:
-            return
-        with tempfile.NamedTemporaryFile() as temp_file:
-            await self.http_dump.export(temp_file.name)
-            self.http_flows = get_http_flows_from_mitm_dump(Path(temp_file.name))
-
-    def parse_airbyte_messages_from_command_output(
-        self, command_output_path: Path, log_validation_errors: bool = False
-    ) -> Iterable[AirbyteMessage]:
-        with open(command_output_path) as command_output:
-            for line in command_output:
-                try:
-                    yield AirbyteMessage.parse_raw(line)
-                except ValidationError as e:
-                    if log_validation_errors:
-                        self.logger.warn(f"Error parsing AirbyteMessage: {e}")
-
-    def get_records(self) -> Iterable[AirbyteMessage]:
-        self.logger.info(
-            f"Reading records all records for command {self.command.value} on {self.connector_under_test.target_or_control.value} version."
-        )
-        for message in self.airbyte_messages:
-            if message.type is AirbyteMessageType.RECORD:
-                yield message
-
-    def generate_stream_schemas(self) -> dict[str, Any]:
-        self.logger.info("Generating stream schemas")
-        stream_builders: dict[str, SchemaBuilder] = {}
-        for record in self.get_records():
-            stream = record.record.stream
-            if stream not in stream_builders:
-                stream_schema_builder = SchemaBuilder()
-                stream_schema_builder.add_schema({"type": "object", "properties": {}})
-                stream_builders[stream] = stream_schema_builder
-            stream_builders[stream].add_object(
-                self.get_obfuscated_types(record.record.data)
-            )
-        self.logger.info("Stream schemas generated")
-        return {
-            stream: sort_dict_keys(stream_builders[stream].to_schema())
-            for stream in stream_builders
-        }
-
-    @staticmethod
-    def get_obfuscated_types(data: dict[str, Any]) -> dict[str, Any]:
-        """
-        Convert obfuscated records into a record whose values have the same type as the original values.
-        """
-        types = {}
-        for k, v in data.items():
-            if v.startswith("string_"):
-                types[k] = "a"
-            elif v.startswith("integer_"):
-                types[k] = 0
-            elif v.startswith("number_"):
-                types[k] = 0.1
-            elif v.startswith("boolean_"):
-                types[k] = True
-            elif v.startswith("null_"):
-                types[k] = None
-            elif v.startswith("array_"):
-                types[k] = []
-            elif v.startswith("object_"):
-                types[k] = {}
-            else:
-                types[k] = v
-
-        return types
-
-    def get_records_per_stream(self, stream: str) -> Iterator[AirbyteMessage]:
-        assert self.backend is not None, "Backend must be set to get records per stream"
-        self.logger.info(f"Reading records for stream {stream}")
-        if stream not in self.backend.record_per_stream_paths:
-            self.logger.warning(f"No records found for stream {stream}")
-            yield from []
-        else:
-            for message in self.parse_airbyte_messages_from_command_output(
-                self.backend.record_per_stream_paths[stream], log_validation_errors=True
-            ):
-                if message.type is AirbyteMessageType.RECORD:
-                    yield message
-
-    def get_states_per_stream(
-        self, stream: str
-    ) -> Dict[str, List[AirbyteStateMessage]]:
-        self.logger.info(f"Reading state messages for stream {stream}")
-        states = defaultdict(list)
-        for message in self.airbyte_messages:
-            if message.type is AirbyteMessageType.STATE:
-                states[message.state.stream.stream_descriptor.name].append(
-                    message.state
-                )
-        return states
-
-    def get_status_messages_per_stream(
-        self, stream: str
-    ) -> Dict[str, List[AirbyteStreamStatusTraceMessage]]:
-        self.logger.info(f"Reading state messages for stream {stream}")
-        statuses = defaultdict(list)
-        for message in self.airbyte_messages:
-            if (
-                message.type is AirbyteMessageType.TRACE
-                and message.trace.type == TraceType.STREAM_STATUS
-            ):
-                statuses[message.trace.stream_status.stream_descriptor.name].append(
-                    message.trace.stream_status
-                )
-        return statuses
-
-    @cache
-    def get_message_count_per_type(self) -> dict[AirbyteMessageType, int]:
-        message_count: dict[AirbyteMessageType, int] = defaultdict(int)
-        for message in self.airbyte_messages:
-            message_count[message.type] += 1
-        return message_count
-
-    async def save_http_dump(self, output_dir: Path) -> None:
-        if self.http_dump:
-            self.logger.info(
-                "An http dump was captured during the execution of the command, saving it."
-            )
-            http_dump_file_path = (output_dir / self.HTTP_DUMP_FILE_NAME).resolve()
-            await self.http_dump.export(str(http_dump_file_path))
-            self.logger.info(f"Http dump saved to {http_dump_file_path}")
-
-            # Define where the har file will be saved
-            har_file_path = (output_dir / self.HAR_FILE_NAME).resolve()
-            # Convert the mitmproxy dump file to a har file
-            mitm_http_stream_to_har(http_dump_file_path, har_file_path)
-            self.logger.info(f"Har file saved to {har_file_path}")
-        else:
-            self.logger.warning("No http dump to save")
-
-    def save_airbyte_messages(
-        self, output_dir: Path, duckdb_path: Optional[Path] = None
-    ) -> None:
-        self.logger.info("Saving Airbyte messages to disk")
-        airbyte_messages_dir = output_dir / "airbyte_messages"
-        airbyte_messages_dir.mkdir(parents=True, exist_ok=True)
-        if duckdb_path:
-            self.backend = DuckDbBackend(
-                airbyte_messages_dir, duckdb_path, self.duckdb_schema
-            )
-        else:
-            self.backend = FileBackend(airbyte_messages_dir)
-        self.backend.write(self.airbyte_messages)
-        self.logger.info("Airbyte messages saved")
-
-    def save_stream_schemas(self, output_dir: Path) -> None:
-        self.stream_schemas = self.generate_stream_schemas()
-        stream_schemas_dir = output_dir / "stream_schemas"
-        stream_schemas_dir.mkdir(parents=True, exist_ok=True)
-        for stream_name, stream_schema in self.stream_schemas.items():
-            (
-                stream_schemas_dir / f"{sanitize_stream_name(stream_name)}.json"
-            ).write_text(json.dumps(stream_schema, sort_keys=True))
-        self.logger.info("Stream schemas saved to disk")
-
-    async def save_artifacts(
-        self, output_dir: Path, duckdb_path: Optional[Path] = None
-    ) -> None:
-        self.logger.info("Saving artifacts to disk")
-        self.save_airbyte_messages(output_dir, duckdb_path)
-        self.update_configuration()
-        await self.save_http_dump(output_dir)
-        self.save_stream_schemas(output_dir)
-        self.logger.info("All artifacts saved to disk")
-
-    def get_updated_configuration(
-        self, control_message_path: Path
-    ) -> Optional[dict[str, Any]]:
-        """Iterate through the control messages to find CONNECTOR_CONFIG message and return the last updated configuration."""
-        if not control_message_path.exists():
-            return None
-        updated_config = None
-        for line in control_message_path.read_text().splitlines():
-            if line.strip():
-                connector_config = json.loads(line.strip()).get("connectorConfig", {})
-                if connector_config:
-                    updated_config = connector_config
-        return updated_config
-
-    def update_configuration(self) -> None:
-        """This function checks if a configuration has to be updated by reading the control messages file.
-        If a configuration has to be updated, it updates the configuration on the actor using the Airbyte API.
-        """
-        assert self.backend is not None, (
-            "Backend must be set to update configuration in order to find the control messages path"
-        )
-        updated_configuration = self.get_updated_configuration(
-            self.backend.jsonl_controls_path
-        )
-        if updated_configuration is None:
-            return
-
-        self.logger.warning(
-            f"Updating configuration for {self.connector_under_test.name}, actor {self.actor_id}"
-        )
-        url = f"https://api.airbyte.com/v1/{self.connector_under_test.actor_type.value}s/{self.actor_id}"
-
-        payload = {
-            "configuration": {
-                **updated_configuration,
-                f"{self.connector_under_test.actor_type.value}Type": self.connector_under_test.name_without_type_prefix,
-            }
-        }
-        headers = {
-            "accept": "application/json",
-            "content-type": "application/json",
-            "authorization": f"Bearer {get_airbyte_api_key()}",
-        }
-
-        response = requests.patch(url, json=payload, headers=headers)
-        try:
-            response.raise_for_status()
-        except requests.HTTPError as e:
-            self.logger.error(
-                f"Failed to update {self.connector_under_test.name} configuration on actor {self.actor_id}: {e}"
-            )
-            self.logger.error(f"Response: {response.text}")
-        self.logger.info(
-            f"Updated configuration for {self.connector_under_test.name}, actor {self.actor_id}"
-        )
-
-    def __hash__(self):
-        return hash(self.connector_under_test.version)
-
-
-@dataclass(kw_only=True)
-class ConnectionObjects:
-    source_config: Optional[SecretDict]
-    destination_config: Optional[SecretDict]
-    configured_catalog: Optional[ConfiguredAirbyteCatalog]
-    catalog: Optional[AirbyteCatalog]
-    state: Optional[dict]
-    workspace_id: Optional[str]
-    source_id: Optional[str]
-    destination_id: Optional[str]
-    source_docker_image: Optional[str]
-    connection_id: Optional[str]
-
-    @property
-    def url(self) -> Optional[str]:
-        if not self.workspace_id or not self.connection_id:
-            return None
-        return f"https://cloud.airbyte.com/workspaces/{self.workspace_id}/connections/{self.connection_id}"
-
-    @property
-    def hashed_connection_id(self) -> Optional[str]:
-        if not self.connection_id:
-            return None
-        return self.connection_id[:8]