connections.py

# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
"""Cloud Connections."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from airbyte._util import api_util
from airbyte.cloud.connectors import CloudDestination, CloudSource
from airbyte.cloud.sync_results import SyncResult


if TYPE_CHECKING:
    from airbyte_api.models import ConnectionResponse, JobResponse

    from airbyte.cloud.workspaces import CloudWorkspace


class CloudConnection:
    """A connection is an extract-load (EL) pairing of a source and destination in Airbyte Cloud.

    You can use a connection object to run sync jobs, retrieve logs, and manage the connection.
    """

    def __init__(
        self,
        workspace: CloudWorkspace,
        connection_id: str,
        source: str | None = None,
        destination: str | None = None,
    ) -> None:
        """It is not recommended to create a `CloudConnection` object directly.

        Instead, use `CloudWorkspace.get_connection()` to create a connection object.
        """
        self.connection_id = connection_id
        """The ID of the connection."""

        self.workspace = workspace
        """The workspace that the connection belongs to."""

        self._source_id = source
        """The ID of the source."""

        self._destination_id = destination
        """The ID of the destination."""

        self._connection_info: ConnectionResponse | None = None
        """The connection info object. (Cached.)"""

        self._cloud_source_object: CloudSource | None = None
        """The source object. (Cached.)"""

        self._cloud_destination_object: CloudDestination | None = None
        """The destination object. (Cached.)"""

    def _fetch_connection_info(self) -> ConnectionResponse:
        """Populate the connection with data from the API."""
        return api_util.get_connection(
            workspace_id=self.workspace.workspace_id,
            connection_id=self.connection_id,
            api_root=self.workspace.api_root,
            client_id=self.workspace.client_id,
            client_secret=self.workspace.client_secret,
        )

    @classmethod
    def _from_connection_response(
        cls,
        workspace: CloudWorkspace,
        connection_response: ConnectionResponse,
    ) -> CloudConnection:
        """Create a CloudConnection from a ConnectionResponse."""
        result = cls(
            workspace=workspace,
            connection_id=connection_response.connection_id,
            source=connection_response.source_id,
            destination=connection_response.destination_id,
        )
        result._connection_info = connection_response  # noqa: SLF001 # Accessing Non-Public API
        return result

    # Properties

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

        E.g. "My Postgres to Snowflake", not the connection ID.
        """
        if not self._connection_info:
            self._connection_info = self._fetch_connection_info()

        return self._connection_info.name

    @property
    def source_id(self) -> str:
        """The ID of the source."""
        if not self._source_id:
            if not self._connection_info:
                self._connection_info = self._fetch_connection_info()

            self._source_id = self._connection_info.source_id

        return self._source_id

    @property
    def source(self) -> CloudSource:
        """Get the source object."""
        if self._cloud_source_object:
            return self._cloud_source_object

        self._cloud_source_object = CloudSource(
            workspace=self.workspace,
            connector_id=self.source_id,
        )
        return self._cloud_source_object

    @property
    def destination_id(self) -> str:
        """The ID of the destination."""
        if not self._destination_id:
            if not self._connection_info:
                self._connection_info = self._fetch_connection_info()

            self._destination_id = self._connection_info.destination_id

        return self._destination_id

    @property
    def destination(self) -> CloudDestination:
        """Get the destination object."""
        if self._cloud_destination_object:
            return self._cloud_destination_object

        self._cloud_destination_object = CloudDestination(
            workspace=self.workspace,
            connector_id=self.destination_id,
        )
        return self._cloud_destination_object

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

        return [stream.name for stream in self._connection_info.configurations.streams or []]

    @property
    def table_prefix(self) -> str:
        """The table prefix."""
        if not self._connection_info:
            self._connection_info = self._fetch_connection_info()

        return self._connection_info.prefix or ""

    @property
    def connection_url(self) -> str | None:
        """The web URL to the connection."""
        return f"{self.workspace.workspace_url}/connections/{self.connection_id}"

    @property
    def job_history_url(self) -> str | None:
        """The URL to the job history for the connection."""
        return f"{self.connection_url}/timeline"

    # Run Sync

    def run_sync(
        self,
        *,
        wait: bool = True,
        wait_timeout: int = 300,
    ) -> SyncResult:
        """Run a sync."""
        connection_response = api_util.run_connection(
            connection_id=self.connection_id,
            api_root=self.workspace.api_root,
            workspace_id=self.workspace.workspace_id,
            client_id=self.workspace.client_id,
            client_secret=self.workspace.client_secret,
        )
        sync_result = SyncResult(
            workspace=self.workspace,
            connection=self,
            job_id=connection_response.job_id,
        )

        if wait:
            sync_result.wait_for_completion(
                wait_timeout=wait_timeout,
                raise_failure=True,
                raise_timeout=True,
            )

        return sync_result

    def __repr__(self) -> str:
        """String representation of the connection."""
        return (
            f"CloudConnection(connection_id={self.connection_id}, source_id={self.source_id}, "
            f"destination_id={self.destination_id}, connection_url={self.connection_url})"
        )

    # Logs

    def get_previous_sync_logs(
        self,
        *,
        limit: int = 20,
        offset: int | None = None,
        from_tail: bool = True,
    ) -> list[SyncResult]:
        """Get previous sync jobs for a connection with pagination support.

        Returns SyncResult objects containing job metadata (job_id, status, bytes_synced,
        rows_synced, start_time). Full log text can be fetched lazily via
        `SyncResult.get_full_log_text()`.

        Args:
            limit: Maximum number of jobs to return. Defaults to 20.
            offset: Number of jobs to skip from the beginning. Defaults to None (0).
            from_tail: If True, returns jobs ordered newest-first (createdAt DESC).
                If False, returns jobs ordered oldest-first (createdAt ASC).
                Defaults to True.

        Returns:
            A list of SyncResult objects representing the sync jobs.
        """
        order_by = (
            api_util.JOB_ORDER_BY_CREATED_AT_DESC
            if from_tail
            else api_util.JOB_ORDER_BY_CREATED_AT_ASC
        )
        sync_logs: list[JobResponse] = api_util.get_job_logs(
            connection_id=self.connection_id,
            api_root=self.workspace.api_root,
            workspace_id=self.workspace.workspace_id,
            limit=limit,
            offset=offset,
            order_by=order_by,
            client_id=self.workspace.client_id,
            client_secret=self.workspace.client_secret,
        )
        return [
            SyncResult(
                workspace=self.workspace,
                connection=self,
                job_id=sync_log.job_id,
                _latest_job_info=sync_log,
            )
            for sync_log in sync_logs
        ]

    def get_sync_result(
        self,
        job_id: int | None = None,
    ) -> SyncResult | None:
        """Get the sync result for the connection.

        If `job_id` is not provided, the most recent sync job will be used.

        Returns `None` if job_id is omitted and no previous jobs are found.
        """
        if job_id is None:
            # Get the most recent sync job
            results = self.get_previous_sync_logs(
                limit=1,
            )
            if results:
                return results[0]

            return None

        # Get the sync job by ID (lazy loaded)
        return SyncResult(
            workspace=self.workspace,
            connection=self,
            job_id=job_id,
        )

    # Artifacts

    def get_state_artifacts(self) -> list[dict[str, Any]] | None:
        """Get the connection state artifacts.

        Returns the persisted state for this connection, which can be used
        when debugging incremental syncs.

        Uses the Config API endpoint: POST /v1/state/get

        Returns:
            List of state objects for each stream, or None if no state is set.
        """
        state_response = api_util.get_connection_state(
            connection_id=self.connection_id,
            api_root=self.workspace.api_root,
            client_id=self.workspace.client_id,
            client_secret=self.workspace.client_secret,
            bearer_token=self.workspace.bearer_token,
        )
        if state_response.get("stateType") == "not_set":
            return None
        return state_response.get("streamState", [])

    def get_catalog_artifact(self) -> dict[str, Any] | None:
        """Get the configured catalog for this connection.

        Returns the full configured catalog (syncCatalog) for this connection,
        including stream schemas, sync modes, cursor fields, and primary keys.

        Uses the Config API endpoint: POST /v1/web_backend/connections/get

        Returns:
            Dictionary containing the configured catalog, or `None` if not found.
        """
        connection_response = api_util.get_connection_catalog(
            connection_id=self.connection_id,
            api_root=self.workspace.api_root,
            client_id=self.workspace.client_id,
            client_secret=self.workspace.client_secret,
            bearer_token=self.workspace.bearer_token,
        )
        return connection_response.get("syncCatalog")

    def rename(self, name: str) -> CloudConnection:
        """Rename the connection.

        Args:
            name: New name for the connection

        Returns:
            Updated CloudConnection object with refreshed info
        """
        updated_response = api_util.patch_connection(
            connection_id=self.connection_id,
            api_root=self.workspace.api_root,
            client_id=self.workspace.client_id,
            client_secret=self.workspace.client_secret,
            bearer_token=self.workspace.bearer_token,
            name=name,
        )
        self._connection_info = updated_response
        return self

    def set_table_prefix(self, prefix: str) -> CloudConnection:
        """Set the table prefix for the connection.

        Args:
            prefix: New table prefix to use when syncing to the destination

        Returns:
            Updated CloudConnection object with refreshed info
        """
        updated_response = api_util.patch_connection(
            connection_id=self.connection_id,
            api_root=self.workspace.api_root,
            client_id=self.workspace.client_id,
            client_secret=self.workspace.client_secret,
            bearer_token=self.workspace.bearer_token,
            prefix=prefix,
        )
        self._connection_info = updated_response
        return self

    def set_selected_streams(self, stream_names: list[str]) -> CloudConnection:
        """Set the selected streams for the connection.

        This is a destructive operation that can break existing connections if the
        stream selection is changed incorrectly. Use with caution.

        Args:
            stream_names: List of stream names to sync

        Returns:
            Updated CloudConnection object with refreshed info
        """
        configurations = api_util.build_stream_configurations(stream_names)

        updated_response = api_util.patch_connection(
            connection_id=self.connection_id,
            api_root=self.workspace.api_root,
            client_id=self.workspace.client_id,
            client_secret=self.workspace.client_secret,
            bearer_token=self.workspace.bearer_token,
            configurations=configurations,
        )
        self._connection_info = updated_response
        return self

    # Deletions

    def permanently_delete(
        self,
        *,
        cascade_delete_source: bool = False,
        cascade_delete_destination: bool = False,
    ) -> None:
        """Delete the connection.

        Args:
            cascade_delete_source: Whether to also delete the source.
            cascade_delete_destination: Whether to also delete the destination.
        """
        self.workspace.permanently_delete_connection(self)

        if cascade_delete_source:
            self.workspace.permanently_delete_source(self.source_id)

        if cascade_delete_destination:
            self.workspace.permanently_delete_destination(self.destination_id)