airbyte.cloud.connections

  1# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
  2"""Cloud Connections."""
  3
  4from __future__ import annotations
  5
  6from typing import TYPE_CHECKING, cast
  7
  8from airbyte._util import api_util
  9from airbyte.cloud.sync_results import SyncResult
 10
 11
 12if TYPE_CHECKING:
 13    from airbyte_api.models import ConnectionResponse, JobResponse
 14
 15    from airbyte.cloud.workspaces import CloudWorkspace
 16
 17
 18class CloudConnection:
 19    """A connection is an extract-load (EL) pairing of a source and destination in Airbyte Cloud.
 20
 21    You can use a connection object to run sync jobs, retrieve logs, and manage the connection.
 22    """
 23
 24    def __init__(
 25        self,
 26        workspace: CloudWorkspace,
 27        connection_id: str,
 28        source: str | None = None,
 29        destination: str | None = None,
 30    ) -> None:
 31        """It is not recommended to create a `CloudConnection` object directly.
 32
 33        Instead, use `CloudWorkspace.get_connection()` to create a connection object.
 34        """
 35        self.connection_id = connection_id
 36        """The ID of the connection."""
 37
 38        self.workspace = workspace
 39        """The workspace that the connection belongs to."""
 40
 41        self._source_id = source
 42        """The ID of the source."""
 43
 44        self._destination_id = destination
 45        """The ID of the destination."""
 46
 47        self._connection_info: ConnectionResponse | None = None
 48
 49    def _fetch_connection_info(self) -> ConnectionResponse:
 50        """Populate the connection with data from the API."""
 51        return api_util.get_connection(
 52            workspace_id=self.workspace.workspace_id,
 53            connection_id=self.connection_id,
 54            api_root=self.workspace.api_root,
 55            api_key=self.workspace.api_key,
 56        )
 57
 58    # Properties
 59
 60    @property
 61    def source_id(self) -> str:
 62        """The ID of the source."""
 63        if not self._source_id:
 64            if not self._connection_info:
 65                self._connection_info = self._fetch_connection_info()
 66
 67            self._source_id = self._connection_info.source_id
 68
 69        return cast(str, self._source_id)
 70
 71    @property
 72    def destination_id(self) -> str:
 73        """The ID of the destination."""
 74        if not self._destination_id:
 75            if not self._connection_info:
 76                self._connection_info = self._fetch_connection_info()
 77
 78            self._destination_id = self._connection_info.source_id
 79
 80        return cast(str, self._destination_id)
 81
 82    @property
 83    def stream_names(self) -> list[str]:
 84        """The stream names."""
 85        if not self._connection_info:
 86            self._connection_info = self._fetch_connection_info()
 87
 88        return [stream.name for stream in self._connection_info.configurations.streams]
 89
 90    @property
 91    def table_prefix(self) -> str:
 92        """The table prefix."""
 93        if not self._connection_info:
 94            self._connection_info = self._fetch_connection_info()
 95
 96        return self._connection_info.prefix
 97
 98    @property
 99    def connection_url(self) -> str | None:
100        """The URL to the connection."""
101        return f"{self.workspace.workspace_url}/connections/{self.connection_id}"
102
103    @property
104    def job_history_url(self) -> str | None:
105        """The URL to the job history for the connection."""
106        return f"{self.connection_url}/job-history"
107
108    # Run Sync
109
110    def run_sync(
111        self,
112        *,
113        wait: bool = True,
114        wait_timeout: int = 300,
115    ) -> SyncResult:
116        """Run a sync."""
117        connection_response = api_util.run_connection(
118            connection_id=self.connection_id,
119            api_root=self.workspace.api_root,
120            api_key=self.workspace.api_key,
121            workspace_id=self.workspace.workspace_id,
122        )
123        sync_result = SyncResult(
124            workspace=self.workspace,
125            connection=self,
126            job_id=connection_response.job_id,
127        )
128
129        if wait:
130            sync_result.wait_for_completion(
131                wait_timeout=wait_timeout,
132                raise_failure=True,
133                raise_timeout=True,
134            )
135
136        return sync_result
137
138    # Logs
139
140    def get_previous_sync_logs(
141        self,
142        *,
143        limit: int = 10,
144    ) -> list[SyncResult]:
145        """Get the previous sync logs for a connection."""
146        sync_logs: list[JobResponse] = api_util.get_job_logs(
147            connection_id=self.connection_id,
148            api_root=self.workspace.api_root,
149            api_key=self.workspace.api_key,
150            workspace_id=self.workspace.workspace_id,
151            limit=limit,
152        )
153        return [
154            SyncResult(
155                workspace=self.workspace,
156                connection=self,
157                job_id=sync_log.job_id,
158                _latest_job_info=sync_log,
159            )
160            for sync_log in sync_logs
161        ]
162
163    def get_sync_result(
164        self,
165        job_id: str | None = None,
166    ) -> SyncResult | None:
167        """Get the sync result for the connection.
168
169        If `job_id` is not provided, the most recent sync job will be used.
170
171        Returns `None` if job_id is omitted and no previous jobs are found.
172        """
173        if job_id is None:
174            # Get the most recent sync job
175            results = self.get_previous_sync_logs(
176                limit=1,
177            )
178            if results:
179                return results[0]
180
181            return None
182
183        # Get the sync job by ID (lazy loaded)
184        return SyncResult(
185            workspace=self.workspace,
186            connection=self,
187            job_id=job_id,
188        )
189
190    # Deletions
191
192    def _permanently_delete(
193        self,
194        *,
195        delete_source: bool = False,
196        delete_destination: bool = False,
197    ) -> None:
198        """Delete the connection.
199
200        Args:
201            delete_source: Whether to also delete the source.
202            delete_destination: Whether to also delete the destination.
203        """
204        self.workspace._permanently_delete_connection(  # noqa: SLF001  # Non-public API (for now)
205            connection=self
206        )
207
208        if delete_source:
209            self.workspace._permanently_delete_source(  # noqa: SLF001  # Non-public API (for now)
210                source=self.source_id
211            )
212
213        if delete_destination:
214            self.workspace._permanently_delete_destination(  # noqa: SLF001  # Non-public API
215                destination=self.destination_id,
216            )