prefect_census.syncs

Module containing tasks and flows for interacting with Census syncs.

Classes

CensusSync

Bases: JobBlock

A Job Block for triggering a Census sync run and waiting for completion.

Attributes:

Name	Type	Description
credentials	CensusCredentials	Credentials for authenticating with Census.
sync_id	int	The ID of the sync to trigger.
force_full_sync	bool	If True, a full sync will be triggered.
max_wait_seconds	int	Maximum number of seconds to wait for sync to complete.
poll_frequency_seconds	int	Number of seconds to wait in between checks for run completion.

Example

Trigger a Census sync and wait for completion as a subflow:

from prefect import flow
from prefect_census import CensusSync, run_census_sync

@flow
def my_census_flow():
    census_sync = CensusSync.load("BLOCK_NAME")

    # do some setup

    run_census_sync(census_sync)

    # do some cleanup

Source code in prefect_census/syncs.py

class CensusSync(JobBlock):
    """
    A Job Block for triggering a Census sync run and waiting for completion.

    Attributes:
        credentials: Credentials for authenticating with Census.
        sync_id: The ID of the sync to trigger.
        force_full_sync: If `True`, a full sync will be triggered.
        max_wait_seconds: Maximum number of seconds to wait for sync to complete.
        poll_frequency_seconds: Number of seconds to wait in between checks
            for run completion.

    Example:
        Trigger a Census sync and wait for completion as a subflow:
        ```python
        from prefect import flow
        from prefect_census import CensusSync, run_census_sync

        @flow
        def my_census_flow():
            census_sync = CensusSync.load("BLOCK_NAME")

            # do some setup

            run_census_sync(census_sync)

            # do some cleanup
        ```
    """

    credentials: CensusCredentials = Field(
        default=...,
        description="Credentials for authenticating with Census.",
    )

    sync_id: int = Field(
        default=...,
        description="The ID of the sync to trigger.",
    )

    force_full_sync: bool = Field(
        default=False,
        description="If `True`, a full sync will be triggered.",
    )

    max_wait_seconds: int = Field(
        default=3600,
        description="The maximum number of seconds to wait for the sync to complete.",
    )

    poll_frequency_seconds: int = Field(
        default=5,
        description="Number of seconds to wait between sync status checks.",
    )

    _block_type_name = "Census Sync"
    _description = "Runs a Census sync"
    _documentation_url = "https://prefecthq.github.io/prefect-census/syncs/"
    _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/b2f805555c778f37f861b67c4a861908f66a6e35-700x700.png"  # noqa

    @sync_compatible
    async def trigger(self) -> "CensusSyncRun":
        """
        Trigger a Census sync run.

        Returns:
            A CensusSyncRun instance representing the triggered sync run.
        """
        self.logger.info(f"Triggering Census sync run for sync with ID {self.sync_id}")
        try:
            async with self.credentials.get_client() as client:
                response = await client.trigger_sync_run(
                    sync_id=self.sync_id, force_full_sync=self.force_full_sync
                )
        except HTTPStatusError as e:
            raise CensusSyncTriggerFailed(extract_user_message(e)) from e

        run_data = response.json()["data"]

        if "sync_run_id" in run_data:
            self.logger.info(
                f"Census sync with ID: {self.sync_id} successfully triggered. "
                "You can view the status of this sync run at "
                f"https://app.getcensus.com/sync/{self.sync_id}/sync-history"
            )

        run_id = run_data["sync_run_id"]
        if run_id is None:
            raise RuntimeError("Unable to determine run ID for triggered sync.")

        return CensusSyncRun(census_sync=self, run_id=run_id)

Functions

trigger async

Trigger a Census sync run.

Returns:

Type	Description
CensusSyncRun	A CensusSyncRun instance representing the triggered sync run.

Source code in prefect_census/syncs.py

@sync_compatible
async def trigger(self) -> "CensusSyncRun":
    """
    Trigger a Census sync run.

    Returns:
        A CensusSyncRun instance representing the triggered sync run.
    """
    self.logger.info(f"Triggering Census sync run for sync with ID {self.sync_id}")
    try:
        async with self.credentials.get_client() as client:
            response = await client.trigger_sync_run(
                sync_id=self.sync_id, force_full_sync=self.force_full_sync
            )
    except HTTPStatusError as e:
        raise CensusSyncTriggerFailed(extract_user_message(e)) from e

    run_data = response.json()["data"]

    if "sync_run_id" in run_data:
        self.logger.info(
            f"Census sync with ID: {self.sync_id} successfully triggered. "
            "You can view the status of this sync run at "
            f"https://app.getcensus.com/sync/{self.sync_id}/sync-history"
        )

    run_id = run_data["sync_run_id"]
    if run_id is None:
        raise RuntimeError("Unable to determine run ID for triggered sync.")

    return CensusSyncRun(census_sync=self, run_id=run_id)

CensusSyncResult

Bases: BaseModel

The results for a Census Sync Run.

Source code in prefect_census/syncs.py

class CensusSyncResult(BaseModel):
    """The results for a Census Sync Run."""

    final_status: CensusSyncRunStatus
    run_data: Dict[str, Any]

CensusSyncRun

Bases: JobRun

A Job Run representing a Census sync run.

Attributes:

Name	Type	Description
sync		The CensusSync instance that triggered this run.
run_id		The ID of the Census sync run.
run_data		The data returned by the Census API for this run.
status		The status of the Census sync run.

Source code in prefect_census/syncs.py

class CensusSyncRun(JobRun):
    """
    A Job Run representing a Census sync run.

    Attributes:
        sync: The CensusSync instance that triggered this run.
        run_id: The ID of the Census sync run.
        run_data: The data returned by the Census API for this run.
        status: The status of the Census sync run.
    """

    def __init__(self, census_sync, run_id) -> None:
        self.sync = census_sync
        self.run_id = run_id
        self.run_data = None
        self.status = None

    @sync_compatible
    async def wait_for_completion(self):
        """Wait for the Census sync run to complete."""
        seconds_waited_for_run_completion = 0
        last_status = None
        while seconds_waited_for_run_completion <= self.sync.max_wait_seconds:
            try:
                async with self.sync.credentials.get_client() as client:
                    response = await client.get_run_info(self.run_id)
            except HTTPStatusError as e:
                raise RuntimeError(extract_user_message(e)) from e

            run_data = response.json()["data"]
            self.status = run_data.get("status")

            if CensusSyncRunStatus.is_terminal_status_code(self.status):
                self.run_data = run_data
                if self.status == CensusSyncRunStatus.COMPLETED.value:
                    self.logger.info(
                        "Census sync run with ID %s completed successfully!",
                        self.run_id,
                    )
                    return

                elif self.status == CensusSyncRunStatus.CANCELLED.value:
                    raise CensusSyncRunCancelled(
                        f"Triggered sync run with ID {self.run_id} was cancelled."
                    )
                elif self.status == CensusSyncRunStatus.FAILED.value:
                    raise CensusSyncRunFailed(
                        f"Triggered sync run with ID {self.run_id} has failed."
                    )
                else:
                    raise RuntimeError(
                        f"Sync run with ID: {self.run_id} ended with unexpected "
                        f"status {self.status}"
                    )
            if self.status != last_status:
                self.logger.info(
                    "Census sync run with ID %i has status %s.",
                    self.run_id,
                    CensusSyncRunStatus(self.status).name,
                )
                last_status = self.status

            await asyncio.sleep(self.sync.poll_frequency_seconds)
            seconds_waited_for_run_completion += self.sync.poll_frequency_seconds

        raise CensusSyncRunTimeout(
            f"Timeout of {self.sync.max_wait_seconds} seconds exceeded while "
            f"waiting for sync run with ID {self.run_id} to complete."
        )

    @sync_compatible
    async def fetch_result(self) -> CensusSyncResult:
        """Fetch the result of the Census sync run.

        Returns:
            A CensusSyncResult instance representing the result of the sync run.
        """
        return CensusSyncResult(
            final_status=CensusSyncRunStatus(self.status),
            run_data=self.run_data,
        )

Functions

fetch_result async

Fetch the result of the Census sync run.

Returns:

Type	Description
CensusSyncResult	A CensusSyncResult instance representing the result of the sync run.

Source code in prefect_census/syncs.py

@sync_compatible
async def fetch_result(self) -> CensusSyncResult:
    """Fetch the result of the Census sync run.

    Returns:
        A CensusSyncResult instance representing the result of the sync run.
    """
    return CensusSyncResult(
        final_status=CensusSyncRunStatus(self.status),
        run_data=self.run_data,
    )

wait_for_completion async

Wait for the Census sync run to complete.

Source code in prefect_census/syncs.py

@sync_compatible
async def wait_for_completion(self):
    """Wait for the Census sync run to complete."""
    seconds_waited_for_run_completion = 0
    last_status = None
    while seconds_waited_for_run_completion <= self.sync.max_wait_seconds:
        try:
            async with self.sync.credentials.get_client() as client:
                response = await client.get_run_info(self.run_id)
        except HTTPStatusError as e:
            raise RuntimeError(extract_user_message(e)) from e

        run_data = response.json()["data"]
        self.status = run_data.get("status")

        if CensusSyncRunStatus.is_terminal_status_code(self.status):
            self.run_data = run_data
            if self.status == CensusSyncRunStatus.COMPLETED.value:
                self.logger.info(
                    "Census sync run with ID %s completed successfully!",
                    self.run_id,
                )
                return

            elif self.status == CensusSyncRunStatus.CANCELLED.value:
                raise CensusSyncRunCancelled(
                    f"Triggered sync run with ID {self.run_id} was cancelled."
                )
            elif self.status == CensusSyncRunStatus.FAILED.value:
                raise CensusSyncRunFailed(
                    f"Triggered sync run with ID {self.run_id} has failed."
                )
            else:
                raise RuntimeError(
                    f"Sync run with ID: {self.run_id} ended with unexpected "
                    f"status {self.status}"
                )
        if self.status != last_status:
            self.logger.info(
                "Census sync run with ID %i has status %s.",
                self.run_id,
                CensusSyncRunStatus(self.status).name,
            )
            last_status = self.status

        await asyncio.sleep(self.sync.poll_frequency_seconds)
        seconds_waited_for_run_completion += self.sync.poll_frequency_seconds

    raise CensusSyncRunTimeout(
        f"Timeout of {self.sync.max_wait_seconds} seconds exceeded while "
        f"waiting for sync run with ID {self.run_id} to complete."
    )

CensusSyncTriggerFailed

Bases: RuntimeError

Used to indicate sync triggered.

Source code in prefect_census/syncs.py

class CensusSyncTriggerFailed(RuntimeError):
    """Used to indicate sync triggered."""

Functions

trigger_census_sync async

A task to trigger a Census sync run.

Parameters:

Name	Type	Description	Default
credentials	CensusCredentials	Credentials for authenticating with Census.	required
sync_id	int	The ID of the sync to trigger.	required
force_full_sync	bool	If True, a full sync will be triggered.	False

Returns:

Type	Description
int	The ID of the triggered sync run.

Examples:

Trigger a Census sync run:

from prefect import flow

from prefect_census import CensusCredentials
from prefect_census.syncs import trigger_census_sync

@flow
def trigger_census_sync_flow():
    credentials = CensusCredentials(api_key="my_api_key")
    trigger_census_sync(credentials=credentials, sync_id=42)

trigger_census_sync_flow()

Source code in prefect_census/syncs.py

@task(
    name="Trigger Census sync run",
    description="Triggers a Census sync run for the sync with the given sync_id.",
    retries=3,
    retry_delay_seconds=10,
)
async def trigger_census_sync(
    credentials: CensusCredentials, sync_id: int, force_full_sync: bool = False
) -> int:
    """
    A task to trigger a Census sync run.

    Args:
        credentials: Credentials for authenticating with Census.
        sync_id: The ID of the sync to trigger.
        force_full_sync: If `True`, a full sync will be triggered.

    Returns:
        The ID of the triggered sync run.

    Examples:
        Trigger a Census sync run:
        ```python
        from prefect import flow

        from prefect_census import CensusCredentials
        from prefect_census.syncs import trigger_census_sync

        @flow
        def trigger_census_sync_flow():
            credentials = CensusCredentials(api_key="my_api_key")
            trigger_census_sync(credentials=credentials, sync_id=42)

        trigger_census_sync_flow()
        ```
    """
    logger = get_run_logger()

    logger.info(f"Triggering Census sync run for sync with ID {sync_id}")
    try:
        async with credentials.get_client() as client:
            response = await client.trigger_sync_run(
                sync_id=sync_id, force_full_sync=force_full_sync
            )
    except HTTPStatusError as e:
        raise CensusSyncTriggerFailed(extract_user_message(e)) from e

    run_data = response.json()["data"]

    if "sync_run_id" in run_data:
        logger.info(
            f"Census sync run successfully triggered for sync with ID {sync_id}. "
            "You can view the status of this sync run at "
            f"https://app.getcensus.com/sync/{sync_id}/sync-history"
        )

    return run_data["sync_run_id"]

trigger_census_sync_run_and_wait_for_completion async

Flow that triggers a sync run and waits for the triggered run to complete.

Parameters:

Name	Type	Description	Default
credentials	CensusCredentials	Credentials for authenticating with Census.	required
sync_id	int	The ID of the sync to trigger.	required
force_full_sync	bool	If True, a full sync will be triggered.	False
max_wait_seconds	int	Maximum number of seconds to wait for sync to complete	900
poll_frequency_seconds	int	Number of seconds to wait in between checks for run completion.	10

Raises:

Type	Description
CensusSyncRunCancelled	The triggered Census sync run was cancelled.
CensusSyncRunFailed	The triggered Census sync run failed.
RuntimeError	The triggered Census sync run ended in an unexpected state.

Returns:

Type

Description

Dict[str, Any]

The final run data returned by the Census API as dict with the following shape: { "id": 94, "sync_id": 52, "source_record_count": 1, "records_processed": 1, "records_updated": 1, "records_failed": 0, "records_invalid": 0, "created_at": "2021-10-20T02:51:07.546Z", "updated_at": "2021-10-20T02:52:29.236Z", "completed_at": "2021-10-20T02:52:29.234Z", "scheduled_execution_time": null, "error_code": null, "error_message": null, "error_detail": null, "status": "completed", "canceled": false, "full_sync": true, "sync_trigger_reason": { "ui_tag": "Manual", "ui_detail": "Manually triggered by test@getcensus.com" } }

Examples:

Trigger a Census sync using CensusCredentials instance and wait for completion as a standalone flow:

import asyncio

from prefect_census import CensusCredentials
from prefect_census.syncs import trigger_census_sync_run_and_wait_for_completion

asyncio.run(
    trigger_census_sync_run_and_wait_for_completion(
        credentials=CensusCredentials(
            api_key="my_api_key"
        ),
        sync_id=42
    )
)

Trigger a Census sync and wait for completion as a subflow:

from prefect import flow

from prefect_census import CensusCredentials
from prefect_census.syncs import trigger_census_sync_run_and_wait_for_completion

@flow
def my_flow():
    ...
    creds = CensusCredentials(api_key="my_api_key")
    run_result = trigger_census_sync_run_and_wait_for_completion(
        credentials=creds,
        sync_id=42
    )
    ...

my_flow()

Source code in prefect_census/syncs.py

@flow(
    name="Trigger Census sync run and wait for completion",
    description="Triggers a Census sync run and waits for the"
    "triggered run to complete.",
)
async def trigger_census_sync_run_and_wait_for_completion(
    credentials: CensusCredentials,
    sync_id: int,
    force_full_sync: bool = False,
    max_wait_seconds: int = 900,
    poll_frequency_seconds: int = 10,
) -> Dict[str, Any]:
    """
    Flow that triggers a sync run and waits for the triggered run to complete.

    Args:
        credentials: Credentials for authenticating with Census.
        sync_id: The ID of the sync to trigger.
        force_full_sync: If `True`, a full sync will be triggered.
        max_wait_seconds: Maximum number of seconds to wait for sync to complete
        poll_frequency_seconds: Number of seconds to wait in between checks for run completion.

    Raises:
        CensusSyncRunCancelled: The triggered Census sync run was cancelled.
        CensusSyncRunFailed: The triggered Census sync run failed.
        RuntimeError: The triggered Census sync run ended in an unexpected state.

    Returns:
        The final run data returned by the Census API as dict with the following shape:
            ```
            {
                "id": 94,
                "sync_id": 52,
                "source_record_count": 1,
                "records_processed": 1,
                "records_updated": 1,
                "records_failed": 0,
                "records_invalid": 0,
                "created_at": "2021-10-20T02:51:07.546Z",
                "updated_at": "2021-10-20T02:52:29.236Z",
                "completed_at": "2021-10-20T02:52:29.234Z",
                "scheduled_execution_time": null,
                "error_code": null,
                "error_message": null,
                "error_detail": null,
                "status": "completed",
                "canceled": false,
                "full_sync": true,
                "sync_trigger_reason": {
                    "ui_tag": "Manual",
                    "ui_detail": "Manually triggered by test@getcensus.com"
                }
            }
            ```

    Examples:
        Trigger a Census sync using CensusCredentials instance and wait
        for completion as a standalone flow:
        ```python
        import asyncio

        from prefect_census import CensusCredentials
        from prefect_census.syncs import trigger_census_sync_run_and_wait_for_completion

        asyncio.run(
            trigger_census_sync_run_and_wait_for_completion(
                credentials=CensusCredentials(
                    api_key="my_api_key"
                ),
                sync_id=42
            )
        )
        ```

        Trigger a Census sync and wait for completion as a subflow:
        ```python
        from prefect import flow

        from prefect_census import CensusCredentials
        from prefect_census.syncs import trigger_census_sync_run_and_wait_for_completion

        @flow
        def my_flow():
            ...
            creds = CensusCredentials(api_key="my_api_key")
            run_result = trigger_census_sync_run_and_wait_for_completion(
                credentials=creds,
                sync_id=42
            )
            ...

        my_flow()
        ```
    """  # noqa
    logger = get_run_logger()

    triggered_run_data_future = await trigger_census_sync.submit(
        credentials=credentials, sync_id=sync_id, force_full_sync=force_full_sync
    )

    run_id = await triggered_run_data_future.result()
    if run_id is None:
        raise RuntimeError("Unable to determine run ID for triggered sync.")

    final_run_status, run_data = await wait_census_sync_completion(
        run_id=run_id,
        credentials=credentials,
        max_wait_seconds=max_wait_seconds,
        poll_frequency_seconds=poll_frequency_seconds,
    )

    if final_run_status == CensusSyncRunStatus.COMPLETED:
        logger.info(
            "Census sync run with ID %s completed successfully!",
            run_id,
        )
        return run_data

    elif final_run_status == CensusSyncRunStatus.CANCELLED:
        raise CensusSyncRunCancelled(
            f"Triggered sync run with ID {run_id} was cancelled."
        )
    elif final_run_status == CensusSyncRunStatus.FAILED:
        raise CensusSyncRunFailed(f"Triggered sync run with ID: {run_id} failed.")
    else:
        raise RuntimeError(
            f"Triggered sync run with ID: {run_id} ended with unexpected"
            f"status {final_run_status}"
        )