`airbyte_prefect.connections`

Tasks for connecting to Airbyte and triggering connection syncs

Classes

`AirbyteConnection`

Bases: JobBlock

A block representing an existing Airbyte connection.

Attributes:

Name	Type	Description
`airbyte_server`	`AirbyteServer`	`AirbyteServer` block representing the Airbyte instance where the Airbyte connection is defined.
`connection_id`	`UUID`	UUID of the Airbyte Connection to trigger.
`poll_interval_s`	`int`	Time in seconds between status checks of the Airbyte sync job.
`status_updates`	`bool`	Whether to log job status on each poll of the Airbyte sync job.
`timeout`	`int`	Timeout in seconds for requests made by `httpx.AsyncClient`.
`max_wait_seconds`	`Optional[int]`	Maximum time in seconds to wait for a sync job to reach a terminal status. Leave unset to wait indefinitely.

Examples:

Load an existing AirbyteConnection block:

from airbyte_prefect import AirbyteConnection

airbyte_connection = AirbyteConnection.load("BLOCK_NAME")

Run an Airbyte connection sync as a flow:

from prefect import flow
from airbyte_prefect import AirbyteConnection
from airbyte_prefect.flows import run_connection_sync # this is a flow

airbyte_connection = AirbyteConnection.load("BLOCK_NAME")

@flow
def airbyte_orchestrator():
    run_connection_sync(airbyte_connection) # now it's a subflow

Source code in airbyte_prefect/connections.py

class AirbyteConnection(JobBlock):
    """A block representing an existing Airbyte connection.

    Attributes:
        airbyte_server: `AirbyteServer` block representing the Airbyte instance
            where the Airbyte connection is defined.
        connection_id: UUID of the Airbyte Connection to trigger.
        poll_interval_s: Time in seconds between status checks of the Airbyte sync job.
        status_updates: Whether to log job status on each poll of the Airbyte sync job.
        timeout: Timeout in seconds for requests made by `httpx.AsyncClient`.
        max_wait_seconds: Maximum time in seconds to wait for a sync job to reach a
            terminal status. Leave unset to wait indefinitely.

    Examples:
        Load an existing `AirbyteConnection` block:
        ```python
        from airbyte_prefect import AirbyteConnection

        airbyte_connection = AirbyteConnection.load("BLOCK_NAME")
        ```

        Run an Airbyte connection sync as a flow:
        ```python
        from prefect import flow
        from airbyte_prefect import AirbyteConnection
        from airbyte_prefect.flows import run_connection_sync # this is a flow

        airbyte_connection = AirbyteConnection.load("BLOCK_NAME")

        @flow
        def airbyte_orchestrator():
            run_connection_sync(airbyte_connection) # now it's a subflow
        ```
    """

    _block_type_name = "Airbyte Connection"
    _logo_url = "https://avatars.githubusercontent.com/u/59758427?s=250&v=4"
    _documentation_url = (
        "https://github.com/haybankz/airbyte-prefect#airbyteconnection"  # noqa
    )

    airbyte_server: AirbyteServer = Field(
        default=...,
        description=(
            "AirbyteServer block representing the Airbyte instance "
            "where the Airbyte connection is defined."
        ),
    )

    connection_id: uuid.UUID = Field(
        default=...,
        description="UUID of the Airbyte Connection to trigger.",
    )

    poll_interval_s: int = Field(
        default=15,
        description="Time in seconds between status checks of the Airbyte sync job.",
    )

    status_updates: bool = Field(
        default=False,
        description="Whether to log job status on each poll of the Airbyte sync job.",
    )

    timeout: int = Field(
        default=5,
        description="Timeout in seconds for requests made by httpx.AsyncClient.",
    )

    max_wait_seconds: Optional[int] = Field(
        default=None,
        description=(
            "Maximum time in seconds to wait for a sync job to reach a terminal "
            "status. Leave unset to wait indefinitely."
        ),
    )

    @sync_compatible
    async def trigger(self) -> AirbyteSync:
        """Trigger a sync of the defined Airbyte connection.

        Returns:
            An `AirbyteSync` `JobRun` object representing the active sync job.

        Raises:
            AirbyteConnectionInactiveException: If the connection is inactive.
            AirbyteConnectionDeprecatedException: If the connection is deprecated.
        """
        str_connection_id = str(self.connection_id)

        async with self.airbyte_server.get_client(
            logger=self.logger, timeout=self.timeout
        ) as airbyte_client:

            self.logger.info(
                f"Triggering Airbyte Connection {self.connection_id}, "
                f"in workspace at {self.airbyte_server.base_url!r}"
            )

            connection_status = await airbyte_client.get_connection_status(
                str_connection_id
            )

            if connection_status == CONNECTION_STATUS_ACTIVE:
                (
                    job_id,
                    _,
                ) = await airbyte_client.trigger_manual_sync_connection(
                    str_connection_id
                )

                return AirbyteSync(
                    airbyte_connection=self,
                    job_id=job_id,
                )

            elif connection_status == CONNECTION_STATUS_INACTIVE:
                raise err.AirbyteConnectionInactiveException(
                    f"Connection: {self.connection_id!r} is inactive"
                    f"Please enable the connection {self.connection_id!r} "
                    "in your Airbyte instance."
                )
            elif connection_status == CONNECTION_STATUS_DEPRECATED:
                raise err.AirbyeConnectionDeprecatedException(
                    f"Connection {self.connection_id!r} is deprecated."
                )

Methods:

`trigger` `async`

Trigger a sync of the defined Airbyte connection.

Returns:

Type	Description
`AirbyteSync`	An `AirbyteSync` `JobRun` object representing the active sync job.

Raises:

Type	Description
`AirbyteConnectionInactiveException`	If the connection is inactive.
`AirbyteConnectionDeprecatedException`	If the connection is deprecated.

Source code in airbyte_prefect/connections.py

@sync_compatible
async def trigger(self) -> AirbyteSync:
    """Trigger a sync of the defined Airbyte connection.

    Returns:
        An `AirbyteSync` `JobRun` object representing the active sync job.

    Raises:
        AirbyteConnectionInactiveException: If the connection is inactive.
        AirbyteConnectionDeprecatedException: If the connection is deprecated.
    """
    str_connection_id = str(self.connection_id)

    async with self.airbyte_server.get_client(
        logger=self.logger, timeout=self.timeout
    ) as airbyte_client:

        self.logger.info(
            f"Triggering Airbyte Connection {self.connection_id}, "
            f"in workspace at {self.airbyte_server.base_url!r}"
        )

        connection_status = await airbyte_client.get_connection_status(
            str_connection_id
        )

        if connection_status == CONNECTION_STATUS_ACTIVE:
            (
                job_id,
                _,
            ) = await airbyte_client.trigger_manual_sync_connection(
                str_connection_id
            )

            return AirbyteSync(
                airbyte_connection=self,
                job_id=job_id,
            )

        elif connection_status == CONNECTION_STATUS_INACTIVE:
            raise err.AirbyteConnectionInactiveException(
                f"Connection: {self.connection_id!r} is inactive"
                f"Please enable the connection {self.connection_id!r} "
                "in your Airbyte instance."
            )
        elif connection_status == CONNECTION_STATUS_DEPRECATED:
            raise err.AirbyeConnectionDeprecatedException(
                f"Connection {self.connection_id!r} is deprecated."
            )

`AirbyteSync`

Bases: JobRun

A JobRun representing an Airbyte sync job.

Source code in airbyte_prefect/connections.py

class AirbyteSync(JobRun):
    """A `JobRun` representing an Airbyte sync job."""

    def __init__(self, airbyte_connection: "AirbyteConnection", job_id: int):
        self.airbyte_connection: "AirbyteConnection" = airbyte_connection
        self.job_id: int = job_id
        self._records_synced: int = 0

    @sync_compatible
    async def wait_for_completion(self):
        """Wait for the `AirbyteConnection` sync to reach a terminal state.

        Raises:
            AirbyteSyncJobFailed: If the sync job ends `failed`, `cancelled` or
                `incomplete`.
            AirbyteSyncJobTimeout: If the connection's `max_wait_seconds` elapses
                before the job reaches a terminal status.
        """
        async with self.airbyte_connection.airbyte_server.get_client(
            logger=self.airbyte_connection.logger,
            timeout=self.airbyte_connection.timeout,
        ) as airbyte_client:

            job_status = JOB_STATUS_PENDING
            max_wait_seconds = self.airbyte_connection.max_wait_seconds
            deadline = (
                None if max_wait_seconds is None else monotonic() + max_wait_seconds
            )

            while job_status not in terminal_job_statuses:
                if deadline is not None and monotonic() >= deadline:
                    raise err.AirbyteSyncJobTimeout(
                        f"Job {self.job_id} did not reach a terminal status within "
                        f"{max_wait_seconds} seconds. Last status: {job_status}."
                    )

                job_info = await airbyte_client.get_job_info(self.job_id)

                job_status = job_info["job"]["status"]

                self._records_synced = _records_synced_from(job_info)

                # pending┃running┃incomplete┃failed┃succeeded┃cancelled
                if job_status == JOB_STATUS_SUCCEEDED:
                    self.logger.info(f"Job {self.job_id} succeeded.")
                elif job_status in unsuccessful_job_statuses:
                    self.logger.error(f"Job {self.job_id} {job_status}.")
                    raise err.AirbyteSyncJobFailed(f"Job {self.job_id} {job_status}.")
                else:
                    if self.airbyte_connection.status_updates:
                        self.logger.info(job_status)
                    # wait for next poll interval
                    await sleep(self.airbyte_connection.poll_interval_s)

    @sync_compatible
    async def fetch_result(self) -> AirbyteSyncResult:
        """Fetch the result of the `AirbyteSync`.

        Returns:
            `AirbyteSyncResult`: object containing metadata for the `AirbyteSync`.
        """
        async with self.airbyte_connection.airbyte_server.get_client(
            logger=self.airbyte_connection.logger,
            timeout=self.airbyte_connection.timeout,
        ) as airbyte_client:
            job_info = await airbyte_client.get_job_info(self.job_id)

            job_status = job_info["job"]["status"]
            job_created_at = job_info["job"]["createdAt"]
            job_updated_at = job_info["job"]["updatedAt"]

            return AirbyteSyncResult(
                created_at=job_created_at,
                job_id=self.job_id,
                job_status=job_status,
                records_synced=_records_synced_from(job_info),
                updated_at=job_updated_at,
            )

Methods:

`fetch_result` `async`

Fetch the result of the AirbyteSync.

Returns:

Type	Description
`AirbyteSyncResult`	`AirbyteSyncResult`: object containing metadata for the `AirbyteSync`.

Source code in airbyte_prefect/connections.py

@sync_compatible
async def fetch_result(self) -> AirbyteSyncResult:
    """Fetch the result of the `AirbyteSync`.

    Returns:
        `AirbyteSyncResult`: object containing metadata for the `AirbyteSync`.
    """
    async with self.airbyte_connection.airbyte_server.get_client(
        logger=self.airbyte_connection.logger,
        timeout=self.airbyte_connection.timeout,
    ) as airbyte_client:
        job_info = await airbyte_client.get_job_info(self.job_id)

        job_status = job_info["job"]["status"]
        job_created_at = job_info["job"]["createdAt"]
        job_updated_at = job_info["job"]["updatedAt"]

        return AirbyteSyncResult(
            created_at=job_created_at,
            job_id=self.job_id,
            job_status=job_status,
            records_synced=_records_synced_from(job_info),
            updated_at=job_updated_at,
        )

`wait_for_completion` `async`

Wait for the AirbyteConnection sync to reach a terminal state.

Raises:

Type	Description
`AirbyteSyncJobFailed`	If the sync job ends `failed`, `cancelled` or `incomplete`.
`AirbyteSyncJobTimeout`	If the connection's `max_wait_seconds` elapses before the job reaches a terminal status.

Source code in airbyte_prefect/connections.py

@sync_compatible
async def wait_for_completion(self):
    """Wait for the `AirbyteConnection` sync to reach a terminal state.

    Raises:
        AirbyteSyncJobFailed: If the sync job ends `failed`, `cancelled` or
            `incomplete`.
        AirbyteSyncJobTimeout: If the connection's `max_wait_seconds` elapses
            before the job reaches a terminal status.
    """
    async with self.airbyte_connection.airbyte_server.get_client(
        logger=self.airbyte_connection.logger,
        timeout=self.airbyte_connection.timeout,
    ) as airbyte_client:

        job_status = JOB_STATUS_PENDING
        max_wait_seconds = self.airbyte_connection.max_wait_seconds
        deadline = (
            None if max_wait_seconds is None else monotonic() + max_wait_seconds
        )

        while job_status not in terminal_job_statuses:
            if deadline is not None and monotonic() >= deadline:
                raise err.AirbyteSyncJobTimeout(
                    f"Job {self.job_id} did not reach a terminal status within "
                    f"{max_wait_seconds} seconds. Last status: {job_status}."
                )

            job_info = await airbyte_client.get_job_info(self.job_id)

            job_status = job_info["job"]["status"]

            self._records_synced = _records_synced_from(job_info)

            # pending┃running┃incomplete┃failed┃succeeded┃cancelled
            if job_status == JOB_STATUS_SUCCEEDED:
                self.logger.info(f"Job {self.job_id} succeeded.")
            elif job_status in unsuccessful_job_statuses:
                self.logger.error(f"Job {self.job_id} {job_status}.")
                raise err.AirbyteSyncJobFailed(f"Job {self.job_id} {job_status}.")
            else:
                if self.airbyte_connection.status_updates:
                    self.logger.info(job_status)
                # wait for next poll interval
                await sleep(self.airbyte_connection.poll_interval_s)

`AirbyteSyncResult`

Bases: BaseModel

Model representing a result from an AirbyteSync job run.

Source code in airbyte_prefect/connections.py

class AirbyteSyncResult(BaseModel):
    """Model representing a result from an `AirbyteSync` job run."""

    created_at: datetime
    job_status: JobStatus
    job_id: int
    records_synced: int
    updated_at: datetime

`JobStatus`

Bases: str, Enum

The statuses Airbyte reports for a sync job.

Source code in airbyte_prefect/connections.py

class JobStatus(str, Enum):
    """The statuses Airbyte reports for a sync job."""

    CANCELLED = "cancelled"
    FAILED = "failed"
    INCOMPLETE = "incomplete"
    PENDING = "pending"
    RUNNING = "running"
    SUCCEEDED = "succeeded"

    # Without this, str() and f-strings render "JobStatus.SUCCEEDED" rather than
    # "succeeded", which would change every log line and any caller interpolating
    # `AirbyteSyncResult.job_status`.
    __str__ = str.__str__

Functions:

`trigger_sync` `async`

Prefect Task for triggering an Airbyte connection sync.

It is assumed that the user will have previously configured a Source & Destination into a Connection. e.g. MySql -> CSV

An invocation of trigger_sync will attempt to start a sync job for the specified connection_id representing the Connection in Airbyte.

trigger_sync will poll Airbyte Server for the Connection status and will only complete when the sync has completed or when it receives an error status code from an API call.

As of airbyte-prefect==0.1.3, the kwargs airbyte_server_host and airbyte_server_port can be replaced by passing an airbyte_server block instance to generate the AirbyteClient. Using the airbyte_server block is preferred, but the individual kwargs remain for backwards compatibility.

Parameters:

Name	Type	Description	Default
`connection_id`	`str`	Airbyte connection ID to trigger a sync for.	required
`airbyte_server`	`Optional[AirbyteServer]`	An `AirbyteServer` block to create an `AirbyteClient`.	`None`
`airbyte_server_host`	`Optional[str]`	Airbyte server host to connect to.	`None`
`airbyte_server_port`	`Optional[int]`	Airbyte server port to connect to.	`None`
`airbyte_api_version`	`Optional[str]`	Airbyte API version to use.	`None`
`poll_interval_s`	`int`	How often to poll Airbyte for sync status.	`15`
`status_updates`	`bool`	Whether to log sync job status while polling.	`False`
`timeout`	`int`	The POST request `timeout` for the `httpx.AsyncClient`.	`5`
`max_wait_seconds`	`Optional[int]`	Maximum time to wait for the sync job to reach a terminal status. `None`, the default, waits indefinitely.	`None`

Source code in airbyte_prefect/connections.py

@task
async def trigger_sync(
    connection_id: str,
    airbyte_server: Optional[AirbyteServer] = None,
    airbyte_server_host: Optional[str] = None,
    airbyte_server_port: Optional[int] = None,
    airbyte_api_version: Optional[str] = None,
    poll_interval_s: int = 15,
    status_updates: bool = False,
    timeout: int = 5,
    max_wait_seconds: Optional[int] = None,
) -> Dict[str, Any]:
    """Prefect Task for triggering an Airbyte connection sync.

    *It is assumed that the user will have previously configured
    a Source & Destination into a Connection.*
    e.g. MySql -> CSV

    An invocation of `trigger_sync` will attempt to start a sync job for
    the specified `connection_id` representing the Connection in
    Airbyte.

    `trigger_sync` will poll Airbyte Server for the Connection status and
    will only complete when the sync has completed or
    when it receives an error status code from an API call.

    As of `airbyte-prefect==0.1.3`, the kwargs `airbyte_server_host` and
    `airbyte_server_port` can be replaced by passing an `airbyte_server` block
    instance to generate the `AirbyteClient`. Using the `airbyte_server` block is
    preferred, but the individual kwargs remain for backwards compatibility.

    Args:
        connection_id: Airbyte connection ID to trigger a sync for.
        airbyte_server: An `AirbyteServer` block to create an `AirbyteClient`.
        airbyte_server_host: Airbyte server host to connect to.
        airbyte_server_port: Airbyte server port to connect to.
        airbyte_api_version: Airbyte API version to use.
        poll_interval_s: How often to poll Airbyte for sync status.
        status_updates: Whether to log sync job status while polling.
        timeout: The POST request `timeout` for the `httpx.AsyncClient`.
        max_wait_seconds: Maximum time to wait for the sync job to reach a
            terminal status. `None`, the default, waits indefinitely.
    Raises:
        ValueError: If `connection_id` is not a valid UUID.
        AirbyteSyncJobFailed: If the sync job ends `failed`, `cancelled` or
            `incomplete`.
        AirbyteSyncJobTimeout: If `max_wait_seconds` elapses before the job
            reaches a terminal status.
        AirbyteConnectionInactiveException: If a given connection is inactive.
        AirbyeConnectionDeprecatedException: If a given connection is deprecated.
    Returns:
        Job metadata, including the connection ID and final status of the sync.
    Examples:
        Flow that triggers an Airybte connection sync:
        ```python
        from prefect import flow
        from airbyte_prefect.connections import trigger_sync
        from airbyte_prefect.server import AirbyteServer
        @flow
        def example_trigger_sync_flow():
            # Run other tasks and subflows here
            trigger_sync(
                airbyte_server=AirbyteServer.load("oss-airbyte"),
                connection_id="your-connection-id-to-sync"
            )
        example_trigger_sync_flow()
        ```
    """
    logger = get_run_logger()

    if not airbyte_server:
        warn(
            "The use of `airbyte_server_host`, `airbyte_server_port`, and "
            "`airbyte_api_version` is deprecated and will be removed in a "
            "future release. Please pass an `airbyte_server` block to this "
            "task instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        if any([airbyte_server_host, airbyte_server_port, airbyte_api_version]):
            airbyte_server = AirbyteServer(
                server_host=airbyte_server_host or "localhost",
                server_port=airbyte_server_port or 8000,
                api_version=airbyte_api_version or "v1",
            )
        else:
            airbyte_server = AirbyteServer()
    else:
        if any([airbyte_server_host, airbyte_server_port, airbyte_api_version]):
            logger.info(
                "Ignoring `airbyte_server_host`, `airbyte_api_version`, "
                "and `airbyte_server_port` because `airbyte_server` block "
                " was passed. Using API URL from `airbyte_server` block: "
                f"{airbyte_server.base_url!r}."
            )

    try:
        uuid.UUID(connection_id)
    except (TypeError, ValueError):
        raise ValueError("Parameter `connection_id` *must* be a valid UUID \
            i.e. 32 hex characters, including hyphens.")

    async with airbyte_server.get_client(
        logger=logger, timeout=timeout
    ) as airbyte_client:

        logger.info(
            f"Getting Airbyte Connection {connection_id}, poll interval "
            f"{poll_interval_s} seconds, airbyte_base_url {airbyte_server.base_url}"
        )

        connection_status = await airbyte_client.get_connection_status(connection_id)

        if connection_status == CONNECTION_STATUS_ACTIVE:
            # Trigger manual sync on the Connection ...
            (
                job_id,
                job_created_at,
            ) = await airbyte_client.trigger_manual_sync_connection(connection_id)

            job_status = JOB_STATUS_PENDING
            deadline = (
                None if max_wait_seconds is None else monotonic() + max_wait_seconds
            )

            while job_status not in terminal_job_statuses:
                if deadline is not None and monotonic() >= deadline:
                    raise err.AirbyteSyncJobTimeout(
                        f"Job {job_id} did not reach a terminal status within "
                        f"{max_wait_seconds} seconds. Last status: {job_status}."
                    )

                (
                    job_status,
                    job_created_at,
                    job_updated_at,
                ) = await airbyte_client.get_job_status(job_id)

                # pending┃running┃incomplete┃failed┃succeeded┃cancelled
                if job_status == JOB_STATUS_SUCCEEDED:
                    logger.info(f"Job {job_id} succeeded.")
                elif job_status in unsuccessful_job_statuses:
                    logger.error(f"Job {job_id} {job_status}.")
                    raise err.AirbyteSyncJobFailed(f"Job {job_id} {job_status}.")
                else:
                    if status_updates:
                        logger.info(job_status)
                    # wait for next poll interval
                    await sleep(poll_interval_s)

            return {
                "connection_id": connection_id,
                "status": connection_status,
                "job_status": job_status,
                "job_created_at": job_created_at,
                "job_updated_at": job_updated_at,
            }
        elif connection_status == CONNECTION_STATUS_INACTIVE:
            logger.error(
                f"Connection: {connection_id} is inactive"
                " - you'll need to enable it in your Airbyte instance"
            )
            raise err.AirbyteConnectionInactiveException(
                f"Please enable the Connection {connection_id} in Airbyte instance."
            )
        elif connection_status == CONNECTION_STATUS_DEPRECATED:
            logger.error(f"Connection {connection_id} is deprecated.")
            raise err.AirbyeConnectionDeprecatedException(
                f"Connection {connection_id} is deprecated."
            )

airbyte_prefect.connections

Classes

AirbyteConnection

Methods:

trigger async

AirbyteSync

Methods:

fetch_result async

wait_for_completion async

AirbyteSyncResult

JobStatus

Functions:

trigger_sync async

`airbyte_prefect.connections`

`AirbyteConnection`

`trigger` `async`

`AirbyteSync`

`fetch_result` `async`

`wait_for_completion` `async`

`AirbyteSyncResult`

`JobStatus`

`trigger_sync` `async`