Lee-W commented on code in PR #37693:
URL: https://github.com/apache/airflow/pull/37693#discussion_r1550835906
##########
airflow/providers/google/cloud/sensors/dataflow.py:
##########
@@ -115,10 +124,52 @@ def poke(self, context: Context) -> bool:
return False
+ def execute(self, context: Context) -> Any:
+ """Airflow runs this method on the worker and defers using the
trigger."""
+ if not self.deferrable:
+ super().execute(context)
+ else:
+ if self.poke(context=context):
+ return True
+ self.defer(
+ timeout=self.execution_timeout,
+ trigger=DataflowJobStatusTrigger(
+ job_id=self.job_id,
+ expected_statuses=self.expected_statuses,
+ project_id=self.project_id,
+ location=self.location,
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_interval,
+ impersonation_chain=self.impersonation_chain,
+ ),
+ method_name="execute_complete",
+ )
Review Comment:
```suggestion
elif not self.poke(context=context):
self.defer(
timeout=self.execution_timeout,
trigger=DataflowJobStatusTrigger(
job_id=self.job_id,
expected_statuses=self.expected_statuses,
project_id=self.project_id,
location=self.location,
gcp_conn_id=self.gcp_conn_id,
poll_sleep=self.poll_interval,
impersonation_chain=self.impersonation_chain,
),
method_name="execute_complete",
)
return True
```
##########
airflow/providers/google/cloud/hooks/dataflow.py:
##########
@@ -1353,3 +1365,89 @@ async def list_jobs(
)
page_result: ListJobsAsyncPager = await
client.list_jobs(request=request)
return page_result
+
+ async def list_job_messages(
+ self,
+ job_id: str,
+ project_id: str | None = PROVIDE_PROJECT_ID,
+ minimum_importance: int = JobMessageImportance.JOB_MESSAGE_BASIC,
+ page_size: int | None = None,
+ page_token: str | None = None,
+ start_time: Timestamp | None = None,
+ end_time: Timestamp | None = None,
+ location: str | None = DEFAULT_DATAFLOW_LOCATION,
+ ) -> ListJobMessagesAsyncPager:
+ """
+ Return ListJobMessagesAsyncPager object from
MessagesV1Beta3AsyncClient.
+
+ This method wraps around a similar method of
MessagesV1Beta3AsyncClient. ListJobMessagesAsyncPager can be iterated
+ over to extract messages associated with a specific Job ID.
+
+ For more details see the MessagesV1Beta3AsyncClient method description
at:
+
https://cloud.google.com/python/docs/reference/dataflow/latest/google.cloud.dataflow_v1beta3.services.messages_v1_beta3.MessagesV1Beta3AsyncClient
+
+ :param job_id: ID of the Dataflow job to get messages about.
+ :param project_id: Optional. The Google Cloud project ID in which to
start a job.
+ If set to None or missing, the default project_id from the Google
Cloud connection is used.
+ :param page_size: Optional. If specified, determines the maximum
number of messages to return.
+ If unspecified, the service may choose an appropriate default, or
may return an arbitrarily large number of results.
+ :param page_token: Optional. If supplied, this should be the value of
next_page_token returned by an earlier call.
Review Comment:
We miss `minimum_importance` in docstring
##########
airflow/providers/google/cloud/sensors/dataflow.py:
##########
@@ -194,27 +246,72 @@ def poke(self, context: Context) -> bool:
project_id=self.project_id,
location=self.location,
)
+ return result["metrics"] if self.callback is None else
self.callback(result["metrics"])
+
+ def execute(self, context: Context) -> Any:
+ """Airflow runs this method on the worker and defers using the
trigger."""
+ if not self.deferrable:
+ super().execute(context)
+ else:
Review Comment:
```suggestion
elif not self.poke(context=context):
```
##########
airflow/providers/google/cloud/sensors/dataflow.py:
##########
@@ -194,27 +246,72 @@ def poke(self, context: Context) -> bool:
project_id=self.project_id,
location=self.location,
)
+ return result["metrics"] if self.callback is None else
self.callback(result["metrics"])
+
+ def execute(self, context: Context) -> Any:
+ """Airflow runs this method on the worker and defers using the
trigger."""
+ if not self.deferrable:
+ super().execute(context)
+ else:
+ self.defer(
+ timeout=self.execution_timeout,
+ trigger=DataflowJobMetricsTrigger(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_interval,
+ impersonation_chain=self.impersonation_chain,
+ fail_on_terminal_state=self.fail_on_terminal_state,
+ ),
+ method_name="execute_complete",
+ )
- return self.callback(result["metrics"])
+ def execute_complete(self, context: Context, event: dict[str, str | list])
-> Any:
Review Comment:
should we annotate the return type as `str | list`? or the expected type of
`event['result']`
##########
airflow/providers/google/cloud/sensors/dataflow.py:
##########
@@ -151,12 +205,14 @@ def __init__(
self,
*,
job_id: str,
- callback: Callable[[dict], bool],
+ callback: Callable | None = None,
Review Comment:
May I know why are we removing the arg annotation of `Callable` here? Is it
no longer the case?
##########
airflow/providers/google/cloud/triggers/dataflow.py:
##########
@@ -142,3 +152,496 @@ def _get_async_hook(self) -> AsyncDataflowHook:
impersonation_chain=self.impersonation_chain,
cancel_timeout=self.cancel_timeout,
)
+
+
+class DataflowJobStatusTrigger(BaseTrigger):
+ """
+ Trigger that checks for metrics associated with a Dataflow job.
+
+ :param job_id: Required. ID of the job.
+ :param expected_statuses: The expected state(s) of the operation.
+ See:
https://cloud.google.com/dataflow/docs/reference/rest/v1b3/projects.jobs#Job.JobState
+ :param project_id: Required. The Google Cloud project ID in which the job
was started.
+ :param location: Optional. The location where the job is executed. If set
to None then
+ the value of DEFAULT_DATAFLOW_LOCATION will be used.
+ :param gcp_conn_id: The connection ID to use for connecting to Google
Cloud.
+ :param poll_sleep: Time (seconds) to wait between two consecutive calls to
check the job.
+ :param impersonation_chain: Optional. Service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding
identity, with first
+ account from the list granting this role to the originating account
(templated).
+ """
+
+ def __init__(
+ self,
+ job_id: str,
+ expected_statuses: set[str],
+ project_id: str | None,
+ location: str = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ poll_sleep: int = 10,
+ impersonation_chain: str | Sequence[str] | None = None,
+ ):
+ super().__init__()
+ self.job_id = job_id
+ self.expected_statuses = expected_statuses
+ self.project_id = project_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize class arguments and classpath."""
+ return (
+
"airflow.providers.google.cloud.triggers.dataflow.DataflowJobStatusTrigger",
+ {
+ "job_id": self.job_id,
+ "expected_statuses": self.expected_statuses,
+ "project_id": self.project_id,
+ "location": self.location,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.poll_sleep,
+ "impersonation_chain": self.impersonation_chain,
+ },
+ )
+
+ async def run(self):
+ """
+ Loop until the job reaches an expected or terminal state.
+
+ Yields a TriggerEvent with success status, if the client returns an
expected job status.
+
+ Yields a TriggerEvent with error status, if the client returns an
unexpected terminal
+ job status or any exception is raised while looping.
+
+ In any other case the Trigger will wait for a specified amount of time
+ stored in self.poll_sleep variable.
+ """
+ try:
+ while True:
+ job_status = await self.async_hook.get_job_status(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ if job_status.name in self.expected_statuses:
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Job with id '{self.job_id}' has
reached an expected state: {job_status.name}",
+ }
+ )
+ return
+ elif job_status.name in DataflowJobStatus.TERMINAL_STATES:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Job with id '{self.job_id}' is
already in terminal state: {job_status.name}",
+ }
+ )
+ return
+ self.log.info("Sleeping for %s seconds.", self.poll_sleep)
+ await asyncio.sleep(self.poll_sleep)
+ except Exception as e:
+ self.log.error("Exception occurred while checking for job status!")
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": str(e),
+ }
+ )
+
+ @cached_property
+ def async_hook(self) -> AsyncDataflowHook:
+ return AsyncDataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_sleep,
+ impersonation_chain=self.impersonation_chain,
+ )
+
+
+class DataflowJobMetricsTrigger(BaseTrigger):
+ """
+ Trigger that checks for metrics associated with a Dataflow job.
+
+ :param job_id: Required. ID of the job.
+ :param project_id: Required. The Google Cloud project ID in which the job
was started.
+ :param location: Optional. The location where the job is executed. If set
to None then
+ the value of DEFAULT_DATAFLOW_LOCATION will be used.
+ :param gcp_conn_id: The connection ID to use for connecting to Google
Cloud.
+ :param poll_sleep: Time (seconds) to wait between two consecutive calls to
check the job.
+ :param impersonation_chain: Optional. Service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding
identity, with first
+ account from the list granting this role to the originating account
(templated).
+ :param fail_on_terminal_state: If set to True the trigger will yield a
TriggerEvent with
+ error status if the job reaches a terminal state.
+ """
+
+ def __init__(
+ self,
+ job_id: str,
+ project_id: str | None,
+ location: str = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ poll_sleep: int = 10,
+ impersonation_chain: str | Sequence[str] | None = None,
+ fail_on_terminal_state: bool = True,
+ ):
+ super().__init__()
+ self.project_id = project_id
+ self.job_id = job_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+ self.fail_on_terminal_state = fail_on_terminal_state
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize class arguments and classpath."""
+ return (
+
"airflow.providers.google.cloud.triggers.dataflow.DataflowJobMetricsTrigger",
+ {
+ "project_id": self.project_id,
+ "job_id": self.job_id,
+ "location": self.location,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.poll_sleep,
+ "impersonation_chain": self.impersonation_chain,
+ "fail_on_terminal_state": self.fail_on_terminal_state,
+ },
+ )
+
+ async def run(self):
+ """
+ Loop until a terminal job status or any job metrics are returned.
+
+ Yields a TriggerEvent with success status, if the client returns any
job metrics
+ and fail_on_terminal_state attribute is False.
+
+ Yields a TriggerEvent with error status, if the client returns a job
status with
+ a terminal state value and fail_on_terminal_state attribute is True.
+
+ Yields a TriggerEvent with error status, if any exception is raised
while looping.
+
+ In any other case the Trigger will wait for a specified amount of time
+ stored in self.poll_sleep variable.
+ """
+ try:
+ while True:
+ job_status = await self.async_hook.get_job_status(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ job_metrics = await self.get_job_metrics()
+ if self.fail_on_terminal_state:
+ if job_status.name in DataflowJobStatus.TERMINAL_STATES:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Job with id '{self.job_id}' is
already in terminal state: {job_status.name}",
+ "result": None,
+ }
+ )
+ return
+ if job_metrics:
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Detected {len(job_metrics)} metrics
for job '{self.job_id}'",
+ "result": job_metrics,
+ }
+ )
+ return
+ self.log.info("Sleeping for %s seconds.", self.poll_sleep)
+ await asyncio.sleep(self.poll_sleep)
+ except Exception as e:
+ self.log.error("Exception occurred while checking for job's
metrics!")
+ yield TriggerEvent({"status": "error", "message": str(e),
"result": None})
+
+ async def get_job_metrics(self) -> list[dict[str, Any]]:
+ """Wait for the Dataflow client response and then return it in a
serialized list."""
+ job_response: JobMetrics = await self.async_hook.get_job_metrics(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ return self._get_metrics_from_job_response(job_response)
+
+ def _get_metrics_from_job_response(self, job_response: JobMetrics) ->
list[dict[str, Any]]:
+ """Return a list of serialized MetricUpdate objects."""
+ return [MetricUpdate.to_dict(metric) for metric in
job_response.metrics]
+
+ @cached_property
+ def async_hook(self) -> AsyncDataflowHook:
+ return AsyncDataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_sleep,
+ impersonation_chain=self.impersonation_chain,
+ )
+
+
+class DataflowJobAutoScalingEventTrigger(BaseTrigger):
+ """
+ Trigger that checks for autoscaling events associated with a Dataflow job.
+
+ :param job_id: Required. ID of the job.
+ :param project_id: Required. The Google Cloud project ID in which the job
was started.
+ :param location: Optional. The location where the job is executed. If set
to None then
+ the value of DEFAULT_DATAFLOW_LOCATION will be used.
+ :param gcp_conn_id: The connection ID to use for connecting to Google
Cloud.
+ :param poll_sleep: Time (seconds) to wait between two consecutive calls to
check the job.
+ :param impersonation_chain: Optional. Service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding
identity, with first
+ account from the list granting this role to the originating account
(templated).
+ :param fail_on_terminal_state: If set to True the trigger will yield a
TriggerEvent with
+ error status if the job reaches a terminal state.
+ """
+
+ def __init__(
+ self,
+ job_id: str,
+ project_id: str | None,
+ location: str = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ poll_sleep: int = 10,
+ impersonation_chain: str | Sequence[str] | None = None,
+ fail_on_terminal_state: bool = True,
+ ):
+ super().__init__()
+ self.project_id = project_id
+ self.job_id = job_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+ self.fail_on_terminal_state = fail_on_terminal_state
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize class arguments and classpath."""
+ return (
+
"airflow.providers.google.cloud.triggers.dataflow.DataflowJobAutoScalingEventTrigger",
+ {
+ "project_id": self.project_id,
+ "job_id": self.job_id,
+ "location": self.location,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.poll_sleep,
+ "impersonation_chain": self.impersonation_chain,
+ "fail_on_terminal_state": self.fail_on_terminal_state,
+ },
+ )
+
+ async def run(self):
+ """
+ Loop until a terminal job status or any autoscaling events are
returned.
+
+ Yields a TriggerEvent with success status, if the client returns any
autoscaling events
+ and fail_on_terminal_state attribute is False.
+
+ Yields a TriggerEvent with error status, if the client returns a job
status with
+ a terminal state value and fail_on_terminal_state attribute is True.
+
+ Yields a TriggerEvent with error status, if any exception is raised
while looping.
+
+ In any other case the Trigger will wait for a specified amount of time
+ stored in self.poll_sleep variable.
+ """
+ try:
+ while True:
+ job_status = await self.async_hook.get_job_status(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ autoscaling_events = await self.list_job_autoscaling_events()
+ if self.fail_on_terminal_state:
+ if job_status.name in DataflowJobStatus.TERMINAL_STATES:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Job with id '{self.job_id}' is
already in terminal state: {job_status.name}",
+ "result": None,
+ }
+ )
+ return
+ if autoscaling_events:
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Detected {len(autoscaling_events)}
autoscaling events for job '{self.job_id}'",
+ "result": autoscaling_events,
+ }
+ )
+ return
+ self.log.info("Sleeping for %s seconds.", self.poll_sleep)
+ await asyncio.sleep(self.poll_sleep)
+ except Exception as e:
+ self.log.error("Exception occurred while checking for job's
autoscaling events!")
+ yield TriggerEvent({"status": "error", "message": str(e),
"result": None})
+
+ async def list_job_autoscaling_events(self) -> list[dict[str, str | dict]]:
+ """Wait for the Dataflow client response and then return it in a
serialized list."""
+ job_response: ListJobMessagesAsyncPager = await
self.async_hook.list_job_messages(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ return self._get_autoscaling_events_from_job_response(job_response)
+
+ def _get_autoscaling_events_from_job_response(
+ self, job_response: ListJobMessagesAsyncPager
+ ) -> list[dict[str, str | dict]]:
+ """Return a list of serialized AutoscalingEvent objects."""
+ return [AutoscalingEvent.to_dict(event) for event in
job_response.autoscaling_events]
+
+ @cached_property
+ def async_hook(self) -> AsyncDataflowHook:
+ return AsyncDataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_sleep,
+ impersonation_chain=self.impersonation_chain,
+ )
+
+
+class DataflowJobMessagesTrigger(BaseTrigger):
+ """
+ Trigger that checks for job messages associated with a Dataflow job.
+
+ :param job_id: Required. ID of the job.
+ :param project_id: Required. The Google Cloud project ID in which the job
was started.
+ :param location: Optional. The location where the job is executed. If set
to None then
+ the value of DEFAULT_DATAFLOW_LOCATION will be used.
+ :param gcp_conn_id: The connection ID to use for connecting to Google
Cloud.
+ :param poll_sleep: Time (seconds) to wait between two consecutive calls to
check the job.
+ :param impersonation_chain: Optional. Service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding
identity, with first
+ account from the list granting this role to the originating account
(templated).
+ :param fail_on_terminal_state: If set to True the trigger will yield a
TriggerEvent with
+ error status if the job reaches a terminal state.
+ """
+
+ def __init__(
+ self,
+ job_id: str,
+ project_id: str | None,
+ location: str = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ poll_sleep: int = 10,
+ impersonation_chain: str | Sequence[str] | None = None,
+ fail_on_terminal_state: bool = True,
+ ):
+ super().__init__()
+ self.project_id = project_id
+ self.job_id = job_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+ self.fail_on_terminal_state = fail_on_terminal_state
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize class arguments and classpath."""
+ return (
+
"airflow.providers.google.cloud.triggers.dataflow.DataflowJobMessagesTrigger",
+ {
+ "project_id": self.project_id,
+ "job_id": self.job_id,
+ "location": self.location,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.poll_sleep,
+ "impersonation_chain": self.impersonation_chain,
+ "fail_on_terminal_state": self.fail_on_terminal_state,
+ },
+ )
+
+ async def run(self):
+ """
+ Loop until a terminal job status or any job messages are returned.
+
+ Yields a TriggerEvent with success status, if the client returns any
job messages
+ and fail_on_terminal_state attribute is False.
+
+ Yields a TriggerEvent with error status, if the client returns a job
status with
+ a terminal state value and fail_on_terminal_state attribute is True.
+
+ Yields a TriggerEvent with error status, if any exception is raised
while looping.
+
+ In any other case the Trigger will wait for a specified amount of time
+ stored in self.poll_sleep variable.
+ """
+ try:
+ while True:
+ job_status = await self.async_hook.get_job_status(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ job_messages = await self.list_job_messages()
+ if self.fail_on_terminal_state:
+ if job_status.name in DataflowJobStatus.TERMINAL_STATES:
Review Comment:
```suggestion
if self.fail_on_terminal_state and job_status.name in
DataflowJobStatus.TERMINAL_STATES:
```
##########
airflow/providers/google/cloud/triggers/dataflow.py:
##########
@@ -142,3 +152,496 @@ def _get_async_hook(self) -> AsyncDataflowHook:
impersonation_chain=self.impersonation_chain,
cancel_timeout=self.cancel_timeout,
)
+
+
+class DataflowJobStatusTrigger(BaseTrigger):
+ """
+ Trigger that checks for metrics associated with a Dataflow job.
+
+ :param job_id: Required. ID of the job.
+ :param expected_statuses: The expected state(s) of the operation.
+ See:
https://cloud.google.com/dataflow/docs/reference/rest/v1b3/projects.jobs#Job.JobState
+ :param project_id: Required. The Google Cloud project ID in which the job
was started.
+ :param location: Optional. The location where the job is executed. If set
to None then
+ the value of DEFAULT_DATAFLOW_LOCATION will be used.
+ :param gcp_conn_id: The connection ID to use for connecting to Google
Cloud.
+ :param poll_sleep: Time (seconds) to wait between two consecutive calls to
check the job.
+ :param impersonation_chain: Optional. Service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding
identity, with first
+ account from the list granting this role to the originating account
(templated).
+ """
+
+ def __init__(
+ self,
+ job_id: str,
+ expected_statuses: set[str],
+ project_id: str | None,
+ location: str = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ poll_sleep: int = 10,
+ impersonation_chain: str | Sequence[str] | None = None,
+ ):
+ super().__init__()
+ self.job_id = job_id
+ self.expected_statuses = expected_statuses
+ self.project_id = project_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize class arguments and classpath."""
+ return (
+
"airflow.providers.google.cloud.triggers.dataflow.DataflowJobStatusTrigger",
+ {
+ "job_id": self.job_id,
+ "expected_statuses": self.expected_statuses,
+ "project_id": self.project_id,
+ "location": self.location,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.poll_sleep,
+ "impersonation_chain": self.impersonation_chain,
+ },
+ )
+
+ async def run(self):
+ """
+ Loop until the job reaches an expected or terminal state.
+
+ Yields a TriggerEvent with success status, if the client returns an
expected job status.
+
+ Yields a TriggerEvent with error status, if the client returns an
unexpected terminal
+ job status or any exception is raised while looping.
+
+ In any other case the Trigger will wait for a specified amount of time
+ stored in self.poll_sleep variable.
+ """
+ try:
+ while True:
+ job_status = await self.async_hook.get_job_status(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ if job_status.name in self.expected_statuses:
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Job with id '{self.job_id}' has
reached an expected state: {job_status.name}",
+ }
+ )
+ return
+ elif job_status.name in DataflowJobStatus.TERMINAL_STATES:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Job with id '{self.job_id}' is
already in terminal state: {job_status.name}",
+ }
+ )
+ return
+ self.log.info("Sleeping for %s seconds.", self.poll_sleep)
+ await asyncio.sleep(self.poll_sleep)
+ except Exception as e:
+ self.log.error("Exception occurred while checking for job status!")
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": str(e),
+ }
+ )
+
+ @cached_property
+ def async_hook(self) -> AsyncDataflowHook:
+ return AsyncDataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_sleep,
+ impersonation_chain=self.impersonation_chain,
+ )
+
+
+class DataflowJobMetricsTrigger(BaseTrigger):
+ """
+ Trigger that checks for metrics associated with a Dataflow job.
+
+ :param job_id: Required. ID of the job.
+ :param project_id: Required. The Google Cloud project ID in which the job
was started.
+ :param location: Optional. The location where the job is executed. If set
to None then
+ the value of DEFAULT_DATAFLOW_LOCATION will be used.
+ :param gcp_conn_id: The connection ID to use for connecting to Google
Cloud.
+ :param poll_sleep: Time (seconds) to wait between two consecutive calls to
check the job.
+ :param impersonation_chain: Optional. Service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding
identity, with first
+ account from the list granting this role to the originating account
(templated).
+ :param fail_on_terminal_state: If set to True the trigger will yield a
TriggerEvent with
+ error status if the job reaches a terminal state.
+ """
+
+ def __init__(
+ self,
+ job_id: str,
+ project_id: str | None,
+ location: str = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ poll_sleep: int = 10,
+ impersonation_chain: str | Sequence[str] | None = None,
+ fail_on_terminal_state: bool = True,
+ ):
+ super().__init__()
+ self.project_id = project_id
+ self.job_id = job_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+ self.fail_on_terminal_state = fail_on_terminal_state
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize class arguments and classpath."""
+ return (
+
"airflow.providers.google.cloud.triggers.dataflow.DataflowJobMetricsTrigger",
+ {
+ "project_id": self.project_id,
+ "job_id": self.job_id,
+ "location": self.location,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.poll_sleep,
+ "impersonation_chain": self.impersonation_chain,
+ "fail_on_terminal_state": self.fail_on_terminal_state,
+ },
+ )
+
+ async def run(self):
+ """
+ Loop until a terminal job status or any job metrics are returned.
+
+ Yields a TriggerEvent with success status, if the client returns any
job metrics
+ and fail_on_terminal_state attribute is False.
+
+ Yields a TriggerEvent with error status, if the client returns a job
status with
+ a terminal state value and fail_on_terminal_state attribute is True.
+
+ Yields a TriggerEvent with error status, if any exception is raised
while looping.
+
+ In any other case the Trigger will wait for a specified amount of time
+ stored in self.poll_sleep variable.
+ """
+ try:
+ while True:
+ job_status = await self.async_hook.get_job_status(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ job_metrics = await self.get_job_metrics()
+ if self.fail_on_terminal_state:
+ if job_status.name in DataflowJobStatus.TERMINAL_STATES:
Review Comment:
Maybe we could extract this check as a function
##########
airflow/providers/google/cloud/sensors/dataflow.py:
##########
@@ -276,26 +373,72 @@ def poke(self, context: Context) -> bool:
location=self.location,
)
- return self.callback(result)
+ return result if self.callback is None else self.callback(result)
+
+ def execute(self, context: Context) -> Any:
+ """Airflow runs this method on the worker and defers using the
trigger."""
+ if not self.deferrable:
+ super().execute(context)
+ else:
+ self.defer(
+ timeout=self.execution_timeout,
+ trigger=DataflowJobMessagesTrigger(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_interval,
+ impersonation_chain=self.impersonation_chain,
+ fail_on_terminal_state=self.fail_on_terminal_state,
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, str | list])
-> Any:
Review Comment:
```suggestion
def execute_complete(self, context: Context, event: dict[str, str |
list]) -> str | list:
```
##########
airflow/providers/google/cloud/sensors/dataflow.py:
##########
@@ -143,6 +194,9 @@ class DataflowJobMetricsSensor(BaseSensorOperator):
If set as a sequence, the identities from the list must grant
Service Account Token Creator IAM role to the directly preceding
identity, with first
account from the list granting this role to the originating account
(templated).
+ :param deferrable: If True, run the sensor in the deferrable mode.
+ :param poll_interval: Time (seconds) to wait between two consecutive calls
to check the job.
+
Review Comment:
```suggestion
```
##########
airflow/providers/google/cloud/sensors/dataflow.py:
##########
@@ -357,4 +500,45 @@ def poke(self, context: Context) -> bool:
location=self.location,
)
- return self.callback(result)
+ return result if self.callback is None else self.callback(result)
+
+ def execute(self, context: Context) -> Any:
+ """Airflow runs this method on the worker and defers using the
trigger."""
+ if not self.deferrable:
+ super().execute(context)
+ else:
Review Comment:
```suggestion
elif not self.poke(context=context):
```
##########
airflow/providers/google/cloud/sensors/dataflow.py:
##########
@@ -115,10 +124,52 @@ def poke(self, context: Context) -> bool:
return False
+ def execute(self, context: Context) -> Any:
+ """Airflow runs this method on the worker and defers using the
trigger."""
+ if not self.deferrable:
+ super().execute(context)
+ else:
+ if self.poke(context=context):
+ return True
+ self.defer(
+ timeout=self.execution_timeout,
+ trigger=DataflowJobStatusTrigger(
+ job_id=self.job_id,
+ expected_statuses=self.expected_statuses,
+ project_id=self.project_id,
+ location=self.location,
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_interval,
+ impersonation_chain=self.impersonation_chain,
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, str | list])
-> Any:
Review Comment:
Will this change be an mypy error?
```suggestion
def execute_complete(self, context: Context, event: dict[str, str |
list]) -> bool:
```
##########
airflow/providers/google/cloud/sensors/dataflow.py:
##########
@@ -357,4 +500,45 @@ def poke(self, context: Context) -> bool:
location=self.location,
)
- return self.callback(result)
+ return result if self.callback is None else self.callback(result)
+
+ def execute(self, context: Context) -> Any:
+ """Airflow runs this method on the worker and defers using the
trigger."""
+ if not self.deferrable:
+ super().execute(context)
+ else:
+ self.defer(
+ trigger=DataflowJobAutoScalingEventTrigger(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_interval,
+ impersonation_chain=self.impersonation_chain,
+ fail_on_terminal_state=self.fail_on_terminal_state,
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, str | list])
-> Any:
Review Comment:
```suggestion
def execute_complete(self, context: Context, event: dict[str, str |
list]) -> str | list:
```
##########
airflow/providers/google/cloud/sensors/dataflow.py:
##########
@@ -276,26 +373,72 @@ def poke(self, context: Context) -> bool:
location=self.location,
)
- return self.callback(result)
+ return result if self.callback is None else self.callback(result)
+
+ def execute(self, context: Context) -> Any:
+ """Airflow runs this method on the worker and defers using the
trigger."""
+ if not self.deferrable:
+ super().execute(context)
+ else:
Review Comment:
```suggestion
elif not self.poke(context=context):
```
##########
airflow/providers/google/cloud/triggers/dataflow.py:
##########
@@ -142,3 +152,496 @@ def _get_async_hook(self) -> AsyncDataflowHook:
impersonation_chain=self.impersonation_chain,
cancel_timeout=self.cancel_timeout,
)
+
+
+class DataflowJobStatusTrigger(BaseTrigger):
+ """
+ Trigger that checks for metrics associated with a Dataflow job.
+
+ :param job_id: Required. ID of the job.
+ :param expected_statuses: The expected state(s) of the operation.
+ See:
https://cloud.google.com/dataflow/docs/reference/rest/v1b3/projects.jobs#Job.JobState
+ :param project_id: Required. The Google Cloud project ID in which the job
was started.
+ :param location: Optional. The location where the job is executed. If set
to None then
+ the value of DEFAULT_DATAFLOW_LOCATION will be used.
+ :param gcp_conn_id: The connection ID to use for connecting to Google
Cloud.
+ :param poll_sleep: Time (seconds) to wait between two consecutive calls to
check the job.
+ :param impersonation_chain: Optional. Service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding
identity, with first
+ account from the list granting this role to the originating account
(templated).
+ """
+
+ def __init__(
+ self,
+ job_id: str,
+ expected_statuses: set[str],
+ project_id: str | None,
+ location: str = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ poll_sleep: int = 10,
+ impersonation_chain: str | Sequence[str] | None = None,
+ ):
+ super().__init__()
+ self.job_id = job_id
+ self.expected_statuses = expected_statuses
+ self.project_id = project_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize class arguments and classpath."""
+ return (
+
"airflow.providers.google.cloud.triggers.dataflow.DataflowJobStatusTrigger",
+ {
+ "job_id": self.job_id,
+ "expected_statuses": self.expected_statuses,
+ "project_id": self.project_id,
+ "location": self.location,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.poll_sleep,
+ "impersonation_chain": self.impersonation_chain,
+ },
+ )
+
+ async def run(self):
+ """
+ Loop until the job reaches an expected or terminal state.
+
+ Yields a TriggerEvent with success status, if the client returns an
expected job status.
+
+ Yields a TriggerEvent with error status, if the client returns an
unexpected terminal
+ job status or any exception is raised while looping.
+
+ In any other case the Trigger will wait for a specified amount of time
+ stored in self.poll_sleep variable.
+ """
+ try:
+ while True:
+ job_status = await self.async_hook.get_job_status(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ if job_status.name in self.expected_statuses:
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Job with id '{self.job_id}' has
reached an expected state: {job_status.name}",
+ }
+ )
+ return
+ elif job_status.name in DataflowJobStatus.TERMINAL_STATES:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Job with id '{self.job_id}' is
already in terminal state: {job_status.name}",
+ }
+ )
+ return
+ self.log.info("Sleeping for %s seconds.", self.poll_sleep)
+ await asyncio.sleep(self.poll_sleep)
+ except Exception as e:
+ self.log.error("Exception occurred while checking for job status!")
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": str(e),
+ }
+ )
+
+ @cached_property
+ def async_hook(self) -> AsyncDataflowHook:
+ return AsyncDataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_sleep,
+ impersonation_chain=self.impersonation_chain,
+ )
+
+
+class DataflowJobMetricsTrigger(BaseTrigger):
+ """
+ Trigger that checks for metrics associated with a Dataflow job.
+
+ :param job_id: Required. ID of the job.
+ :param project_id: Required. The Google Cloud project ID in which the job
was started.
+ :param location: Optional. The location where the job is executed. If set
to None then
+ the value of DEFAULT_DATAFLOW_LOCATION will be used.
+ :param gcp_conn_id: The connection ID to use for connecting to Google
Cloud.
+ :param poll_sleep: Time (seconds) to wait between two consecutive calls to
check the job.
+ :param impersonation_chain: Optional. Service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding
identity, with first
+ account from the list granting this role to the originating account
(templated).
+ :param fail_on_terminal_state: If set to True the trigger will yield a
TriggerEvent with
+ error status if the job reaches a terminal state.
+ """
+
+ def __init__(
+ self,
+ job_id: str,
+ project_id: str | None,
+ location: str = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ poll_sleep: int = 10,
+ impersonation_chain: str | Sequence[str] | None = None,
+ fail_on_terminal_state: bool = True,
+ ):
+ super().__init__()
+ self.project_id = project_id
+ self.job_id = job_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+ self.fail_on_terminal_state = fail_on_terminal_state
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize class arguments and classpath."""
+ return (
+
"airflow.providers.google.cloud.triggers.dataflow.DataflowJobMetricsTrigger",
+ {
+ "project_id": self.project_id,
+ "job_id": self.job_id,
+ "location": self.location,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.poll_sleep,
+ "impersonation_chain": self.impersonation_chain,
+ "fail_on_terminal_state": self.fail_on_terminal_state,
+ },
+ )
+
+ async def run(self):
+ """
+ Loop until a terminal job status or any job metrics are returned.
+
+ Yields a TriggerEvent with success status, if the client returns any
job metrics
+ and fail_on_terminal_state attribute is False.
+
+ Yields a TriggerEvent with error status, if the client returns a job
status with
+ a terminal state value and fail_on_terminal_state attribute is True.
+
+ Yields a TriggerEvent with error status, if any exception is raised
while looping.
+
+ In any other case the Trigger will wait for a specified amount of time
+ stored in self.poll_sleep variable.
+ """
+ try:
+ while True:
+ job_status = await self.async_hook.get_job_status(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ job_metrics = await self.get_job_metrics()
+ if self.fail_on_terminal_state:
+ if job_status.name in DataflowJobStatus.TERMINAL_STATES:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Job with id '{self.job_id}' is
already in terminal state: {job_status.name}",
+ "result": None,
+ }
+ )
+ return
+ if job_metrics:
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Detected {len(job_metrics)} metrics
for job '{self.job_id}'",
+ "result": job_metrics,
+ }
+ )
+ return
+ self.log.info("Sleeping for %s seconds.", self.poll_sleep)
+ await asyncio.sleep(self.poll_sleep)
+ except Exception as e:
+ self.log.error("Exception occurred while checking for job's
metrics!")
+ yield TriggerEvent({"status": "error", "message": str(e),
"result": None})
+
+ async def get_job_metrics(self) -> list[dict[str, Any]]:
+ """Wait for the Dataflow client response and then return it in a
serialized list."""
+ job_response: JobMetrics = await self.async_hook.get_job_metrics(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ return self._get_metrics_from_job_response(job_response)
+
+ def _get_metrics_from_job_response(self, job_response: JobMetrics) ->
list[dict[str, Any]]:
+ """Return a list of serialized MetricUpdate objects."""
+ return [MetricUpdate.to_dict(metric) for metric in
job_response.metrics]
+
+ @cached_property
+ def async_hook(self) -> AsyncDataflowHook:
+ return AsyncDataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_sleep,
+ impersonation_chain=self.impersonation_chain,
+ )
+
+
+class DataflowJobAutoScalingEventTrigger(BaseTrigger):
+ """
+ Trigger that checks for autoscaling events associated with a Dataflow job.
+
+ :param job_id: Required. ID of the job.
+ :param project_id: Required. The Google Cloud project ID in which the job
was started.
+ :param location: Optional. The location where the job is executed. If set
to None then
+ the value of DEFAULT_DATAFLOW_LOCATION will be used.
+ :param gcp_conn_id: The connection ID to use for connecting to Google
Cloud.
+ :param poll_sleep: Time (seconds) to wait between two consecutive calls to
check the job.
+ :param impersonation_chain: Optional. Service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding
identity, with first
+ account from the list granting this role to the originating account
(templated).
+ :param fail_on_terminal_state: If set to True the trigger will yield a
TriggerEvent with
+ error status if the job reaches a terminal state.
+ """
+
+ def __init__(
+ self,
+ job_id: str,
+ project_id: str | None,
+ location: str = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ poll_sleep: int = 10,
+ impersonation_chain: str | Sequence[str] | None = None,
+ fail_on_terminal_state: bool = True,
+ ):
+ super().__init__()
+ self.project_id = project_id
+ self.job_id = job_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+ self.fail_on_terminal_state = fail_on_terminal_state
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize class arguments and classpath."""
+ return (
+
"airflow.providers.google.cloud.triggers.dataflow.DataflowJobAutoScalingEventTrigger",
+ {
+ "project_id": self.project_id,
+ "job_id": self.job_id,
+ "location": self.location,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.poll_sleep,
+ "impersonation_chain": self.impersonation_chain,
+ "fail_on_terminal_state": self.fail_on_terminal_state,
+ },
+ )
+
+ async def run(self):
+ """
+ Loop until a terminal job status or any autoscaling events are
returned.
+
+ Yields a TriggerEvent with success status, if the client returns any
autoscaling events
+ and fail_on_terminal_state attribute is False.
+
+ Yields a TriggerEvent with error status, if the client returns a job
status with
+ a terminal state value and fail_on_terminal_state attribute is True.
+
+ Yields a TriggerEvent with error status, if any exception is raised
while looping.
+
+ In any other case the Trigger will wait for a specified amount of time
+ stored in self.poll_sleep variable.
+ """
+ try:
+ while True:
+ job_status = await self.async_hook.get_job_status(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ autoscaling_events = await self.list_job_autoscaling_events()
+ if self.fail_on_terminal_state:
+ if job_status.name in DataflowJobStatus.TERMINAL_STATES:
Review Comment:
```suggestion
if self.fail_on_terminal_state and job_status.name in
DataflowJobStatus.TERMINAL_STATES:
```
##########
airflow/providers/google/cloud/triggers/dataflow.py:
##########
@@ -142,3 +152,496 @@ def _get_async_hook(self) -> AsyncDataflowHook:
impersonation_chain=self.impersonation_chain,
cancel_timeout=self.cancel_timeout,
)
+
+
+class DataflowJobStatusTrigger(BaseTrigger):
+ """
+ Trigger that checks for metrics associated with a Dataflow job.
+
+ :param job_id: Required. ID of the job.
+ :param expected_statuses: The expected state(s) of the operation.
+ See:
https://cloud.google.com/dataflow/docs/reference/rest/v1b3/projects.jobs#Job.JobState
+ :param project_id: Required. The Google Cloud project ID in which the job
was started.
+ :param location: Optional. The location where the job is executed. If set
to None then
+ the value of DEFAULT_DATAFLOW_LOCATION will be used.
+ :param gcp_conn_id: The connection ID to use for connecting to Google
Cloud.
+ :param poll_sleep: Time (seconds) to wait between two consecutive calls to
check the job.
+ :param impersonation_chain: Optional. Service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding
identity, with first
+ account from the list granting this role to the originating account
(templated).
+ """
+
+ def __init__(
+ self,
+ job_id: str,
+ expected_statuses: set[str],
+ project_id: str | None,
+ location: str = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ poll_sleep: int = 10,
+ impersonation_chain: str | Sequence[str] | None = None,
+ ):
+ super().__init__()
+ self.job_id = job_id
+ self.expected_statuses = expected_statuses
+ self.project_id = project_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize class arguments and classpath."""
+ return (
+
"airflow.providers.google.cloud.triggers.dataflow.DataflowJobStatusTrigger",
+ {
+ "job_id": self.job_id,
+ "expected_statuses": self.expected_statuses,
+ "project_id": self.project_id,
+ "location": self.location,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.poll_sleep,
+ "impersonation_chain": self.impersonation_chain,
+ },
+ )
+
+ async def run(self):
+ """
+ Loop until the job reaches an expected or terminal state.
+
+ Yields a TriggerEvent with success status, if the client returns an
expected job status.
+
+ Yields a TriggerEvent with error status, if the client returns an
unexpected terminal
+ job status or any exception is raised while looping.
+
+ In any other case the Trigger will wait for a specified amount of time
+ stored in self.poll_sleep variable.
+ """
+ try:
+ while True:
+ job_status = await self.async_hook.get_job_status(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ if job_status.name in self.expected_statuses:
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Job with id '{self.job_id}' has
reached an expected state: {job_status.name}",
+ }
+ )
+ return
+ elif job_status.name in DataflowJobStatus.TERMINAL_STATES:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Job with id '{self.job_id}' is
already in terminal state: {job_status.name}",
+ }
+ )
+ return
+ self.log.info("Sleeping for %s seconds.", self.poll_sleep)
+ await asyncio.sleep(self.poll_sleep)
+ except Exception as e:
+ self.log.error("Exception occurred while checking for job status!")
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": str(e),
+ }
+ )
+
+ @cached_property
+ def async_hook(self) -> AsyncDataflowHook:
+ return AsyncDataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ poll_sleep=self.poll_sleep,
+ impersonation_chain=self.impersonation_chain,
+ )
+
+
+class DataflowJobMetricsTrigger(BaseTrigger):
+ """
+ Trigger that checks for metrics associated with a Dataflow job.
+
+ :param job_id: Required. ID of the job.
+ :param project_id: Required. The Google Cloud project ID in which the job
was started.
+ :param location: Optional. The location where the job is executed. If set
to None then
+ the value of DEFAULT_DATAFLOW_LOCATION will be used.
+ :param gcp_conn_id: The connection ID to use for connecting to Google
Cloud.
+ :param poll_sleep: Time (seconds) to wait between two consecutive calls to
check the job.
+ :param impersonation_chain: Optional. Service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding
identity, with first
+ account from the list granting this role to the originating account
(templated).
+ :param fail_on_terminal_state: If set to True the trigger will yield a
TriggerEvent with
+ error status if the job reaches a terminal state.
+ """
+
+ def __init__(
+ self,
+ job_id: str,
+ project_id: str | None,
+ location: str = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ poll_sleep: int = 10,
+ impersonation_chain: str | Sequence[str] | None = None,
+ fail_on_terminal_state: bool = True,
+ ):
+ super().__init__()
+ self.project_id = project_id
+ self.job_id = job_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+ self.fail_on_terminal_state = fail_on_terminal_state
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize class arguments and classpath."""
+ return (
+
"airflow.providers.google.cloud.triggers.dataflow.DataflowJobMetricsTrigger",
+ {
+ "project_id": self.project_id,
+ "job_id": self.job_id,
+ "location": self.location,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_sleep": self.poll_sleep,
+ "impersonation_chain": self.impersonation_chain,
+ "fail_on_terminal_state": self.fail_on_terminal_state,
+ },
+ )
+
+ async def run(self):
+ """
+ Loop until a terminal job status or any job metrics are returned.
+
+ Yields a TriggerEvent with success status, if the client returns any
job metrics
+ and fail_on_terminal_state attribute is False.
+
+ Yields a TriggerEvent with error status, if the client returns a job
status with
+ a terminal state value and fail_on_terminal_state attribute is True.
+
+ Yields a TriggerEvent with error status, if any exception is raised
while looping.
+
+ In any other case the Trigger will wait for a specified amount of time
+ stored in self.poll_sleep variable.
+ """
+ try:
+ while True:
+ job_status = await self.async_hook.get_job_status(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+ job_metrics = await self.get_job_metrics()
+ if self.fail_on_terminal_state:
+ if job_status.name in DataflowJobStatus.TERMINAL_STATES:
Review Comment:
```suggestion
if self.fail_on_terminal_state and job_status.name in
DataflowJobStatus.TERMINAL_STATES:
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
