From cf8fab404bdebb381c7433e2af8430ffd7683a25 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 12 Jul 2022 10:25:56 +0200 Subject: [PATCH 01/29] Initial changes --- .../providers/google/cloud/hooks/dataproc.py | 15 ++- .../google/cloud/operators/dataproc.py | 95 ++++++++++++++++++- 2 files changed, 106 insertions(+), 4 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py index d870d80726669..22f3844a37c9a 100644 --- a/airflow/providers/google/cloud/hooks/dataproc.py +++ b/airflow/providers/google/cloud/hooks/dataproc.py @@ -200,6 +200,16 @@ class DataprocHook(GoogleBaseHook): keyword arguments rather than positional. """ + def __init__( + self, + gcp_conn_id: str = 'google_cloud_default', + delegate_to: Optional[str] = None, + impersonation_chain: Optional[Union[str, Sequence[str]]] = None, + async_client: bool = False + ) -> None: + super().__init__(gcp_conn_id, delegate_to, impersonation_chain) + self.async_client = async_client + def get_cluster_client(self, region: Optional[str] = None) -> ClusterControllerClient: """Returns ClusterControllerClient.""" client_options = None @@ -227,7 +237,10 @@ def get_job_client(self, region: Optional[str] = None) -> JobControllerClient: client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443') return JobControllerClient( - credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options + credentials=self._get_credentials(), + transport="grpc_asyncio" if self.async_client else "grpc", + client_info=CLIENT_INFO, + client_options=client_options ) def get_batch_client(self, region: Optional[str] = None) -> BatchControllerClient: diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index 69ba3943e83ab..c4cc720ab85b3 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -18,6 +18,7 @@ # """This module contains Google Dataproc operators.""" +import asyncio import inspect import ntpath import os @@ -32,9 +33,10 @@ from google.api_core.exceptions import AlreadyExists, NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault from google.api_core.retry import Retry, exponential_sleep_generator -from google.cloud.dataproc_v1 import Batch, Cluster +from google.cloud.dataproc_v1 import Batch, Cluster, JobStatus, Job from google.protobuf.duration_pb2 import Duration from google.protobuf.field_mask_pb2 import FieldMask +from libcst import And from airflow.exceptions import AirflowException from airflow.models import BaseOperator @@ -50,6 +52,7 @@ DataprocLink, DataprocListLink, ) +from airflow.triggers.base import BaseTrigger, TriggerEvent from airflow.utils import timezone if TYPE_CHECKING: @@ -894,6 +897,7 @@ def __init__( job_error_states: Optional[Set[str]] = None, impersonation_chain: Optional[Union[str, Sequence[str]]] = None, asynchronous: bool = False, + deferred: bool = False, **kwargs, ) -> None: super().__init__(**kwargs) @@ -914,6 +918,7 @@ def __init__( self.job: Optional[dict] = None self.dataproc_job_id = None self.asynchronous = asynchronous + self.deferred = deferred def create_job_template(self) -> DataProcJobBuilder: """Initialize `self.job_template` with default values""" @@ -958,6 +963,14 @@ def execute(self, context: 'Context'): context=context, task_instance=self, url=DATAPROC_JOB_LOG_LINK, resource=job_id ) + if self.deferred: + self.defer( + trigger=DataprocBaseTrigger( + job_id=self.job_id, + project_id=self.project_id, + region=self.region + ), + method_name="execute_complete") if not self.asynchronous: self.log.info('Waiting for job %s to complete', job_id) self.hook.wait_for_job(job_id=job_id, region=self.region, project_id=self.project_id) @@ -966,13 +979,64 @@ def execute(self, context: 'Context'): else: raise AirflowException("Create a job template before") + def execute_complete(self, context: "Context", event: Any = None) -> None: + """ + Callback for when the trigger fires - returns immediately. + Relies on trigger to throw an exception, otherwise it assumes execution was + successful. + """ + job: Job = event["job"] + if job.status.state == JobStatus.State.ERROR: + raise AirflowException(f'Job failed:\n{job}') + if job.status.state == JobStatus.State.CANCELLED: + raise AirflowException(f'Job was cancelled:\n{job}') + self.log.info("%s completed successfully.", self.task_id) + def on_kill(self) -> None: """ Callback called when the operator is killed. Cancel any running job. """ if self.dataproc_job_id: - self.hook.cancel_job(project_id=self.project_id, job_id=self.dataproc_job_id, region=self.region) + self.hook.cancel_job(project_id=self.project_id, + job_id=self.dataproc_job_id, region=self.region) + + +class DataprocBaseTrigger(BaseTrigger): + """ + TODO: description + """ + + def __init__(self, job_id, project_id, region, gcp_conn_id, impersonation_chain): + super().__init__() + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + self.job_id = job_id + self.project_id = project_id + self.region = region + self.hook = DataprocHook( + gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, async_client=True) + + def serialize(self): + return ("airflow.providers.google.cloud.operators.dataproc.DataprocBaseTrigger", + {"job_id": self.job_id, "project_id": self.project_id, "region": self.region}) + + def get_job(self) -> Job: + return self.hook.get_job( + job_id=self.job_id, + project_id=self.project_id, + region=self.region + ) + + def job_finished(self) -> bool: + job = self.get_job() + state = job.status.state + return state in (JobStatus.State.ERROR, JobStatus.State.DONE, JobStatus.State.CANCELLED) + + async def run(self): + while self.job_finished(): + await asyncio.sleep(10) + yield TriggerEvent({"job": self.get_job()}) class DataprocSubmitPigJobOperator(DataprocJobBaseOperator): @@ -1771,6 +1835,8 @@ class DataprocSubmitJobOperator(BaseOperator): :param asynchronous: Flag to return after submitting the job to the Dataproc API. This is useful for submitting long running jobs and waiting on them asynchronously using the DataprocJobSensor + :param deferred: Flag to deffer the execution after submitting the job to the Dataproc API. + Use this flag to run the operation in asynchronuous mode based on Defferable operators API. :param cancel_on_kill: Flag which indicates whether cancel the hook's job or not, when on_kill is called :param wait_timeout: How many seconds wait for job to be ready. Used only if ``asynchronous`` is False """ @@ -1793,6 +1859,7 @@ def __init__( gcp_conn_id: str = "google_cloud_default", impersonation_chain: Optional[Union[str, Sequence[str]]] = None, asynchronous: bool = False, + deferred: bool = False, cancel_on_kill: bool = True, wait_timeout: Optional[int] = None, **kwargs, @@ -1808,6 +1875,7 @@ def __init__( self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain self.asynchronous = asynchronous + self.deferred = deferred self.cancel_on_kill = cancel_on_kill self.hook: Optional[DataprocHook] = None self.job_id: Optional[str] = None @@ -1833,7 +1901,15 @@ def execute(self, context: 'Context'): ) self.job_id = new_job_id - if not self.asynchronous: + if self.deferred: + self.defer( + trigger=DataprocBaseTrigger( + job_id=self.job_id, + project_id=self.project_id, + region=self.region + ), + method_name="execute_complete") + elif not self.asynchronous: self.log.info('Waiting for job %s to complete', new_job_id) self.hook.wait_for_job( job_id=new_job_id, region=self.region, project_id=self.project_id, timeout=self.wait_timeout @@ -1842,6 +1918,19 @@ def execute(self, context: 'Context'): return self.job_id + def execute_complete(self, context: "Context", event: Any = None) -> None: + """ + Callback for when the trigger fires - returns immediately. + Relies on trigger to throw an exception, otherwise it assumes execution was + successful. + """ + job: Job = event["job"] + if job.status.state == JobStatus.State.ERROR: + raise AirflowException(f'Job failed:\n{job}') + if job.status.state == JobStatus.State.CANCELLED: + raise AirflowException(f'Job was cancelled:\n{job}') + self.log.info("%s completed successfully.", self.task_id) + def on_kill(self): if self.job_id and self.cancel_on_kill: self.hook.cancel_job(job_id=self.job_id, project_id=self.project_id, region=self.region) From 1ab3ac8cf0e29f3e2371ddd1b56d53b1959a50f0 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Wed, 13 Jul 2022 07:32:48 +0200 Subject: [PATCH 02/29] Move trigger to a dedicated file --- .../google/cloud/operators/dataproc.py | 41 +----------------- .../google/cloud/triggers/dataproc.py | 42 +++++++++++++++++++ 2 files changed, 43 insertions(+), 40 deletions(-) create mode 100644 airflow/providers/google/cloud/triggers/dataproc.py diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index c4cc720ab85b3..da380550ae8be 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -18,7 +18,6 @@ # """This module contains Google Dataproc operators.""" -import asyncio import inspect import ntpath import os @@ -52,7 +51,7 @@ DataprocLink, DataprocListLink, ) -from airflow.triggers.base import BaseTrigger, TriggerEvent +from airflow.providers.google.cloud.triggers.dataproc import DataprocBaseTrigger from airflow.utils import timezone if TYPE_CHECKING: @@ -1001,44 +1000,6 @@ def on_kill(self) -> None: self.hook.cancel_job(project_id=self.project_id, job_id=self.dataproc_job_id, region=self.region) - -class DataprocBaseTrigger(BaseTrigger): - """ - TODO: description - """ - - def __init__(self, job_id, project_id, region, gcp_conn_id, impersonation_chain): - super().__init__() - self.gcp_conn_id = gcp_conn_id - self.impersonation_chain = impersonation_chain - self.job_id = job_id - self.project_id = project_id - self.region = region - self.hook = DataprocHook( - gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, async_client=True) - - def serialize(self): - return ("airflow.providers.google.cloud.operators.dataproc.DataprocBaseTrigger", - {"job_id": self.job_id, "project_id": self.project_id, "region": self.region}) - - def get_job(self) -> Job: - return self.hook.get_job( - job_id=self.job_id, - project_id=self.project_id, - region=self.region - ) - - def job_finished(self) -> bool: - job = self.get_job() - state = job.status.state - return state in (JobStatus.State.ERROR, JobStatus.State.DONE, JobStatus.State.CANCELLED) - - async def run(self): - while self.job_finished(): - await asyncio.sleep(10) - yield TriggerEvent({"job": self.get_job()}) - - class DataprocSubmitPigJobOperator(DataprocJobBaseOperator): """ Start a Pig query Job on a Cloud DataProc cluster. The parameters of the operation diff --git a/airflow/providers/google/cloud/triggers/dataproc.py b/airflow/providers/google/cloud/triggers/dataproc.py new file mode 100644 index 0000000000000..a818bcda43975 --- /dev/null +++ b/airflow/providers/google/cloud/triggers/dataproc.py @@ -0,0 +1,42 @@ + +import asyncio +from airflow.providers.google.cloud.hooks.dataproc import DataprocHook +from airflow.triggers.base import BaseTrigger, TriggerEvent +from google.cloud.dataproc_v1 import JobStatus, Job + + +class DataprocBaseTrigger(BaseTrigger): + """ + TODO: description + """ + + def __init__(self, job_id, project_id, region, gcp_conn_id, impersonation_chain): + super().__init__() + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + self.job_id = job_id + self.project_id = project_id + self.region = region + self.hook = DataprocHook( + gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, async_client=True) + + def serialize(self): + return ("airflow.providers.google.cloud.operators.dataproc.DataprocBaseTrigger", + {"job_id": self.job_id, "project_id": self.project_id, "region": self.region}) + + def get_job(self) -> Job: + return self.hook.get_job( + job_id=self.job_id, + project_id=self.project_id, + region=self.region + ) + + def job_finished(self) -> bool: + job = self.get_job() + state = job.status.state + return state in (JobStatus.State.ERROR, JobStatus.State.DONE, JobStatus.State.CANCELLED) + + async def run(self): + while self.job_finished(): + await asyncio.sleep(10) + yield TriggerEvent({"job": self.get_job()}) From 2f390cd2e286ce50249c100fce29b82b4a7cf88a Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Wed, 13 Jul 2022 07:59:05 +0200 Subject: [PATCH 03/29] Trigger fixes --- .../google/cloud/operators/dataproc.py | 10 +++- .../google/cloud/triggers/dataproc.py | 54 ++++++++++++++----- 2 files changed, 50 insertions(+), 14 deletions(-) diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index da380550ae8be..641c8b24c35b6 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -967,7 +967,10 @@ def execute(self, context: 'Context'): trigger=DataprocBaseTrigger( job_id=self.job_id, project_id=self.project_id, - region=self.region + region=self.region, + delegate_to=self.delegate_to, + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain ), method_name="execute_complete") if not self.asynchronous: @@ -1867,7 +1870,10 @@ def execute(self, context: 'Context'): trigger=DataprocBaseTrigger( job_id=self.job_id, project_id=self.project_id, - region=self.region + region=self.region, + delegate_to=self.delegate_to, + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain ), method_name="execute_complete") elif not self.asynchronous: diff --git a/airflow/providers/google/cloud/triggers/dataproc.py b/airflow/providers/google/cloud/triggers/dataproc.py index a818bcda43975..c3f5c59b56c38 100644 --- a/airflow/providers/google/cloud/triggers/dataproc.py +++ b/airflow/providers/google/cloud/triggers/dataproc.py @@ -1,8 +1,11 @@ import asyncio + +from airflow import AirflowException from airflow.providers.google.cloud.hooks.dataproc import DataprocHook from airflow.triggers.base import BaseTrigger, TriggerEvent from google.cloud.dataproc_v1 import JobStatus, Job +from typing import Any, Dict, Optional, Sequence, Tuple, Union class DataprocBaseTrigger(BaseTrigger): @@ -10,19 +13,43 @@ class DataprocBaseTrigger(BaseTrigger): TODO: description """ - def __init__(self, job_id, project_id, region, gcp_conn_id, impersonation_chain): + def __init__( + self, + job_id: str, + project_id: str, + region: str, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: Optional[Union[str, Sequence[str]]] = None, + delegate_to: Optional[str] = None, + pooling_period_seconds: int = 30 + ): super().__init__() self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain self.job_id = job_id self.project_id = project_id self.region = region + self.pooling_period_seconds = pooling_period_seconds + self.delegate_to = delegate_to self.hook = DataprocHook( - gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, async_client=True) + delegate_to=self.delegate_to, + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + async_client=True + ) def serialize(self): - return ("airflow.providers.google.cloud.operators.dataproc.DataprocBaseTrigger", - {"job_id": self.job_id, "project_id": self.project_id, "region": self.region}) + return ( + "airflow.providers.google.cloud.operators.dataproc.DataprocBaseTrigger", + { + "job_id": self.job_id, + "project_id": self.project_id, + "region": self.region, + "gcp_conn_id": self.gcp_conn_id, + "delegate_to": self.delegate_to, + "impersonation_chain": self.impersonation_chain, + "pooling_period_seconds": self.pooling_period_seconds + }) def get_job(self) -> Job: return self.hook.get_job( @@ -31,12 +58,15 @@ def get_job(self) -> Job: region=self.region ) - def job_finished(self) -> bool: - job = self.get_job() - state = job.status.state - return state in (JobStatus.State.ERROR, JobStatus.State.DONE, JobStatus.State.CANCELLED) - async def run(self): - while self.job_finished(): - await asyncio.sleep(10) - yield TriggerEvent({"job": self.get_job()}) + while True: + job = self.get_job() + if job.done.state: + if job.status.state in (JobStatus.State.DONE, JobStatus.State.CANCELLED): + break + elif job.status.state == JobStatus.State.ERROR: + raise AirflowException(f"Dataproc job execution failed {self.job_id}") + else: + raise AirflowException(f"Dataproc job execution finished in uknknown stat {job.staus.state} {self.job_id}") + await asyncio.sleep(self.pooling_period_seconds) + yield TriggerEvent({"job": job}) From e431d3400d2ff71900105f404d503f31dbf1e835 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 12 Jul 2022 10:25:56 +0200 Subject: [PATCH 04/29] Deferrable operators for Dataproc --- .../providers/google/cloud/hooks/dataproc.py | 5 +- .../google/cloud/operators/dataproc.py | 57 +++++++------- .../google/cloud/triggers/dataproc.py | 75 ++++++++++--------- 3 files changed, 72 insertions(+), 65 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py index 22f3844a37c9a..a3f0089fc5849 100644 --- a/airflow/providers/google/cloud/hooks/dataproc.py +++ b/airflow/providers/google/cloud/hooks/dataproc.py @@ -205,7 +205,7 @@ def __init__( gcp_conn_id: str = 'google_cloud_default', delegate_to: Optional[str] = None, impersonation_chain: Optional[Union[str, Sequence[str]]] = None, - async_client: bool = False + async_client: bool = False, ) -> None: super().__init__(gcp_conn_id, delegate_to, impersonation_chain) self.async_client = async_client @@ -240,7 +240,8 @@ def get_job_client(self, region: Optional[str] = None) -> JobControllerClient: credentials=self._get_credentials(), transport="grpc_asyncio" if self.async_client else "grpc", client_info=CLIENT_INFO, - client_options=client_options + client_options=client_options, + transport="grpc_asyncio" if self.async_client else "grpc", ) def get_batch_client(self, region: Optional[str] = None) -> BatchControllerClient: diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index 641c8b24c35b6..c91ce767d516d 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -869,6 +869,7 @@ class DataprocJobBaseOperator(BaseOperator): :param asynchronous: Flag to return after submitting the job to the Dataproc API. This is useful for submitting long running jobs and waiting on them asynchronously using the DataprocJobSensor + :param deferrable: Run operator in the deferrable mode :var dataproc_job_id: The actual "jobId" as submitted to the Dataproc API. This is useful for identifying or linking to the job in the Google Cloud Console @@ -896,7 +897,7 @@ def __init__( job_error_states: Optional[Set[str]] = None, impersonation_chain: Optional[Union[str, Sequence[str]]] = None, asynchronous: bool = False, - deferred: bool = False, + deferrable: bool = False, **kwargs, ) -> None: super().__init__(**kwargs) @@ -917,7 +918,7 @@ def __init__( self.job: Optional[dict] = None self.dataproc_job_id = None self.asynchronous = asynchronous - self.deferred = deferred + self.deferrable = deferrable def create_job_template(self) -> DataProcJobBuilder: """Initialize `self.job_template` with default values""" @@ -962,7 +963,7 @@ def execute(self, context: 'Context'): context=context, task_instance=self, url=DATAPROC_JOB_LOG_LINK, resource=job_id ) - if self.deferred: + if self.deferrable: self.defer( trigger=DataprocBaseTrigger( job_id=self.job_id, @@ -970,9 +971,10 @@ def execute(self, context: 'Context'): region=self.region, delegate_to=self.delegate_to, gcp_conn_id=self.gcp_conn_id, - impersonation_chain=self.impersonation_chain + impersonation_chain=self.impersonation_chain, ), - method_name="execute_complete") + method_name="execute_complete", + ) if not self.asynchronous: self.log.info('Waiting for job %s to complete', job_id) self.hook.wait_for_job(job_id=job_id, region=self.region, project_id=self.project_id) @@ -981,17 +983,18 @@ def execute(self, context: 'Context'): else: raise AirflowException("Create a job template before") - def execute_complete(self, context: "Context", event: Any = None) -> None: + def execute_complete(self, context, event=None) -> None: """ Callback for when the trigger fires - returns immediately. Relies on trigger to throw an exception, otherwise it assumes execution was successful. """ - job: Job = event["job"] - if job.status.state == JobStatus.State.ERROR: - raise AirflowException(f'Job failed:\n{job}') - if job.status.state == JobStatus.State.CANCELLED: - raise AirflowException(f'Job was cancelled:\n{job}') + job_state = event["job_state"] + job_id = event["job_id"] + if job_state == JobStatus.State.ERROR: + raise AirflowException(f'Job failed:\n{job_id}') + if job_state == JobStatus.State.CANCELLED: + raise AirflowException(f'Job was cancelled:\n{job_id}') self.log.info("%s completed successfully.", self.task_id) def on_kill(self) -> None: @@ -1000,8 +1003,8 @@ def on_kill(self) -> None: Cancel any running job. """ if self.dataproc_job_id: - self.hook.cancel_job(project_id=self.project_id, - job_id=self.dataproc_job_id, region=self.region) + self.hook.cancel_job(project_id=self.project_id, job_id=self.dataproc_job_id, region=self.region) + class DataprocSubmitPigJobOperator(DataprocJobBaseOperator): """ @@ -1799,8 +1802,7 @@ class DataprocSubmitJobOperator(BaseOperator): :param asynchronous: Flag to return after submitting the job to the Dataproc API. This is useful for submitting long running jobs and waiting on them asynchronously using the DataprocJobSensor - :param deferred: Flag to deffer the execution after submitting the job to the Dataproc API. - Use this flag to run the operation in asynchronuous mode based on Defferable operators API. + :param deferrable: Run operator in the deferrable mode :param cancel_on_kill: Flag which indicates whether cancel the hook's job or not, when on_kill is called :param wait_timeout: How many seconds wait for job to be ready. Used only if ``asynchronous`` is False """ @@ -1823,7 +1825,7 @@ def __init__( gcp_conn_id: str = "google_cloud_default", impersonation_chain: Optional[Union[str, Sequence[str]]] = None, asynchronous: bool = False, - deferred: bool = False, + deferrable: bool = False, cancel_on_kill: bool = True, wait_timeout: Optional[int] = None, **kwargs, @@ -1839,7 +1841,7 @@ def __init__( self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain self.asynchronous = asynchronous - self.deferred = deferred + self.deferrable = deferrable self.cancel_on_kill = cancel_on_kill self.hook: Optional[DataprocHook] = None self.job_id: Optional[str] = None @@ -1865,17 +1867,17 @@ def execute(self, context: 'Context'): ) self.job_id = new_job_id - if self.deferred: + if self.deferrable: self.defer( trigger=DataprocBaseTrigger( job_id=self.job_id, project_id=self.project_id, region=self.region, - delegate_to=self.delegate_to, gcp_conn_id=self.gcp_conn_id, - impersonation_chain=self.impersonation_chain + impersonation_chain=self.impersonation_chain, ), - method_name="execute_complete") + method_name="execute_complete", + ) elif not self.asynchronous: self.log.info('Waiting for job %s to complete', new_job_id) self.hook.wait_for_job( @@ -1885,17 +1887,18 @@ def execute(self, context: 'Context'): return self.job_id - def execute_complete(self, context: "Context", event: Any = None) -> None: + def execute_complete(self, context, event=None) -> None: """ Callback for when the trigger fires - returns immediately. Relies on trigger to throw an exception, otherwise it assumes execution was successful. """ - job: Job = event["job"] - if job.status.state == JobStatus.State.ERROR: - raise AirflowException(f'Job failed:\n{job}') - if job.status.state == JobStatus.State.CANCELLED: - raise AirflowException(f'Job was cancelled:\n{job}') + job_state = event["job_state"] + job_id = event["job_id"] + if job_state == JobStatus.State.ERROR: + raise AirflowException(f'Job failed:\n{job_id}') + if job_state == JobStatus.State.CANCELLED: + raise AirflowException(f'Job was cancelled:\n{job_id}') self.log.info("%s completed successfully.", self.task_id) def on_kill(self): diff --git a/airflow/providers/google/cloud/triggers/dataproc.py b/airflow/providers/google/cloud/triggers/dataproc.py index c3f5c59b56c38..1a2a1a952aa1d 100644 --- a/airflow/providers/google/cloud/triggers/dataproc.py +++ b/airflow/providers/google/cloud/triggers/dataproc.py @@ -1,4 +1,3 @@ - import asyncio from airflow import AirflowException @@ -10,18 +9,19 @@ class DataprocBaseTrigger(BaseTrigger): """ - TODO: description + Trigger that periodically pollls information from Dataproc API to verify job status. + Implementation leverages asynchronous transport. """ def __init__( - self, - job_id: str, - project_id: str, - region: str, - gcp_conn_id: str = "google_cloud_default", - impersonation_chain: Optional[Union[str, Sequence[str]]] = None, - delegate_to: Optional[str] = None, - pooling_period_seconds: int = 30 + self, + job_id: str, + project_id: str, + region: str, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: Optional[Union[str, Sequence[str]]] = None, + delegate_to: Optional[str] = None, + pooling_period_seconds: int = 30, ): super().__init__() self.gcp_conn_id = gcp_conn_id @@ -35,38 +35,41 @@ def __init__( delegate_to=self.delegate_to, gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, - async_client=True + async_client=True, ) def serialize(self): return ( - "airflow.providers.google.cloud.operators.dataproc.DataprocBaseTrigger", - { - "job_id": self.job_id, - "project_id": self.project_id, - "region": self.region, - "gcp_conn_id": self.gcp_conn_id, - "delegate_to": self.delegate_to, - "impersonation_chain": self.impersonation_chain, - "pooling_period_seconds": self.pooling_period_seconds - }) - - def get_job(self) -> Job: - return self.hook.get_job( - job_id=self.job_id, - project_id=self.project_id, - region=self.region + "airflow.providers.google.cloud.operators.dataproc.DataprocBaseTrigger", + { + "job_id": self.job_id, + "project_id": self.project_id, + "region": self.region, + "gcp_conn_id": self.gcp_conn_id, + "delegate_to": self.delegate_to, + "impersonation_chain": self.impersonation_chain, + "pooling_period_seconds": self.pooling_period_seconds, + }, ) + async def get_job(self) -> Job: + return await self.hook.get_job(job_id=self.job_id, project_id=self.project_id, region=self.region) + async def run(self): while True: - job = self.get_job() - if job.done.state: - if job.status.state in (JobStatus.State.DONE, JobStatus.State.CANCELLED): - break - elif job.status.state == JobStatus.State.ERROR: - raise AirflowException(f"Dataproc job execution failed {self.job_id}") - else: - raise AirflowException(f"Dataproc job execution finished in uknknown stat {job.staus.state} {self.job_id}") + job: Job = await self.hook.get_job( + project_id=self.project_id, region=self.region, job_id=self.job_id + ) + state = job.status.state + self.log.info("Dataproc job: %s is in state: %s", self.job_id, state) + if state in (JobStatus.State.ERROR, JobStatus.State.DONE, JobStatus.State.CANCELLED): + if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED): + break + elif state == JobStatus.State.ERROR: + raise AirflowException(f"Dataproc job execution failed {self.job_id}") + else: + raise AirflowException( + f"Dataproc job execution finished in uknknown state {state} {self.job_id}" + ) await asyncio.sleep(self.pooling_period_seconds) - yield TriggerEvent({"job": job}) + yield TriggerEvent({"job_id": self.job_id, "job_state": state}) From 452d4032fb451103b78ce9e31ac0e870418e876e Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Mon, 25 Jul 2022 15:55:48 +0200 Subject: [PATCH 05/29] Tests --- .../google/cloud/operators/dataproc.py | 1 + .../google/cloud/operators/test_dataproc.py | 42 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index c91ce767d516d..be22e88d1b2e4 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -972,6 +972,7 @@ def execute(self, context: 'Context'): delegate_to=self.delegate_to, gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, + pooling_period_seconds=10, ), method_name="execute_complete", ) diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py index a7101ca55875c..827392c4b09b0 100644 --- a/tests/providers/google/cloud/operators/test_dataproc.py +++ b/tests/providers/google/cloud/operators/test_dataproc.py @@ -62,6 +62,7 @@ AIRFLOW_VERSION = "v" + airflow_version.replace(".", "-").replace("+", "-") DATAPROC_PATH = "airflow.providers.google.cloud.operators.dataproc.{}" +COMPOSER_TRIGGERS_STRING = "airflow.providers.google.cloud.triggers.dataproc.{}" TASK_ID = "task-id" GCP_PROJECT = "test-project" @@ -904,6 +905,47 @@ def test_execute_async(self, mock_hook): key="conf", value=DATAPROC_JOB_CONF_EXPECTED, execution_date=None ) + @mock.patch(DATAPROC_PATH.format("DataprocHook")) + @mock.patch(COMPOSER_TRIGGERS_STRING.format("DataprocHook")) + def test_execute_deferrable(self, mock_trigger_hook, mock_hook): + job = {} + mock_hook.return_value.submit_job.return_value.reference.job_id = TEST_JOB_ID + + op = DataprocSubmitJobOperator( + task_id=TASK_ID, + region=GCP_LOCATION, + project_id=GCP_PROJECT, + job=job, + gcp_conn_id=GCP_CONN_ID, + retry=RETRY, + asynchronous=True, + timeout=TIMEOUT, + metadata=METADATA, + request_id=REQUEST_ID, + impersonation_chain=IMPERSONATION_CHAIN, + deferrable=True, + ) + op.execute(context=self.mock_context) + + mock_hook.assert_called_once_with( + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + ) + mock_hook.return_value.submit_job.assert_called_once_with( + project_id=GCP_PROJECT, + region=GCP_LOCATION, + job=job, + request_id=REQUEST_ID, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + mock_hook.return_value.wait_for_job.assert_not_called() + + self.mock_ti.xcom_push.assert_called_once_with( + key="conf", value=DATAPROC_JOB_CONF_EXPECTED, execution_date=None + ) + @mock.patch(DATAPROC_PATH.format("DataprocHook")) def test_on_kill(self, mock_hook): job = {} From e8cec425e68a76a9be75fe7cb05261da36bb93c4 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 26 Jul 2022 07:41:44 +0200 Subject: [PATCH 06/29] Remove unused import --- airflow/providers/google/cloud/operators/dataproc.py | 1 - 1 file changed, 1 deletion(-) diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index be22e88d1b2e4..ae1d7e7969e1c 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -35,7 +35,6 @@ from google.cloud.dataproc_v1 import Batch, Cluster, JobStatus, Job from google.protobuf.duration_pb2 import Duration from google.protobuf.field_mask_pb2 import FieldMask -from libcst import And from airflow.exceptions import AirflowException from airflow.models import BaseOperator From a0c515c79dd6cd5e3f6f71cdfecaccaff853babd Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 26 Jul 2022 07:43:47 +0200 Subject: [PATCH 07/29] Fixes --- airflow/providers/google/cloud/hooks/dataproc.py | 1 - 1 file changed, 1 deletion(-) diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py index a3f0089fc5849..f4380ccb1dd37 100644 --- a/airflow/providers/google/cloud/hooks/dataproc.py +++ b/airflow/providers/google/cloud/hooks/dataproc.py @@ -238,7 +238,6 @@ def get_job_client(self, region: Optional[str] = None) -> JobControllerClient: return JobControllerClient( credentials=self._get_credentials(), - transport="grpc_asyncio" if self.async_client else "grpc", client_info=CLIENT_INFO, client_options=client_options, transport="grpc_asyncio" if self.async_client else "grpc", From c1125c88eddde444e530645cf2b695fd66b7fc19 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 26 Jul 2022 07:53:26 +0200 Subject: [PATCH 08/29] Tests --- tests/providers/google/cloud/operators/test_dataproc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py index 827392c4b09b0..7f39727021381 100644 --- a/tests/providers/google/cloud/operators/test_dataproc.py +++ b/tests/providers/google/cloud/operators/test_dataproc.py @@ -925,7 +925,7 @@ def test_execute_deferrable(self, mock_trigger_hook, mock_hook): impersonation_chain=IMPERSONATION_CHAIN, deferrable=True, ) - op.execute(context=self.mock_context) + op.execute(mock.MagicMock()) mock_hook.assert_called_once_with( gcp_conn_id=GCP_CONN_ID, From 125b54af311ce8d04bc3988117ffed82f21838dd Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 26 Jul 2022 07:56:55 +0200 Subject: [PATCH 09/29] Tests --- .../providers/google/cloud/operators/test_dataproc.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py index 7f39727021381..5ecd61dfaf541 100644 --- a/tests/providers/google/cloud/operators/test_dataproc.py +++ b/tests/providers/google/cloud/operators/test_dataproc.py @@ -26,7 +26,7 @@ from google.cloud.dataproc_v1 import Batch from airflow import AirflowException -from airflow.exceptions import AirflowTaskTimeout +from airflow.exceptions import AirflowTaskTimeout, TaskDeferred from airflow.models import DAG, DagBag from airflow.providers.google.cloud.operators.dataproc import ( DATAPROC_CLUSTER_LINK, @@ -52,6 +52,8 @@ DataprocSubmitSparkSqlJobOperator, DataprocUpdateClusterOperator, ) +from airflow.providers.google.cloud.triggers.dataproc import DataprocBaseTrigger +from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME from airflow.serialization.serialized_objects import SerializedDAG from airflow.utils.timezone import datetime from airflow.version import version as airflow_version @@ -925,7 +927,8 @@ def test_execute_deferrable(self, mock_trigger_hook, mock_hook): impersonation_chain=IMPERSONATION_CHAIN, deferrable=True, ) - op.execute(mock.MagicMock()) + with pytest.raises(TaskDeferred) as exc: + op.execute(mock.MagicMock()) mock_hook.assert_called_once_with( gcp_conn_id=GCP_CONN_ID, @@ -946,6 +949,9 @@ def test_execute_deferrable(self, mock_trigger_hook, mock_hook): key="conf", value=DATAPROC_JOB_CONF_EXPECTED, execution_date=None ) + assert isinstance(exc.value.trigger, DataprocBaseTrigger) + assert exc.value.method_name == GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME + @mock.patch(DATAPROC_PATH.format("DataprocHook")) def test_on_kill(self, mock_hook): job = {} From 98ddda347fe60647e592abbdd8663ffd839de98a Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 26 Jul 2022 07:59:20 +0200 Subject: [PATCH 10/29] Tests --- tests/providers/google/cloud/operators/test_dataproc.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py index 5ecd61dfaf541..8cb93f8d19acb 100644 --- a/tests/providers/google/cloud/operators/test_dataproc.py +++ b/tests/providers/google/cloud/operators/test_dataproc.py @@ -945,9 +945,7 @@ def test_execute_deferrable(self, mock_trigger_hook, mock_hook): ) mock_hook.return_value.wait_for_job.assert_not_called() - self.mock_ti.xcom_push.assert_called_once_with( - key="conf", value=DATAPROC_JOB_CONF_EXPECTED, execution_date=None - ) + self.mock_ti.xcom_push.assert_not_called() assert isinstance(exc.value.trigger, DataprocBaseTrigger) assert exc.value.method_name == GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME From 6830cfc120f57deff8f0bc11e4d2e87b16fd1430 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 26 Jul 2022 08:59:18 +0200 Subject: [PATCH 11/29] DataProcAsyncHook implementation --- .../providers/google/cloud/hooks/dataproc.py | 740 +++++++++++++++++- .../google/cloud/triggers/dataproc.py | 12 +- 2 files changed, 740 insertions(+), 12 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py index f4380ccb1dd37..fde23877ee8f4 100644 --- a/airflow/providers/google/cloud/hooks/dataproc.py +++ b/airflow/providers/google/cloud/hooks/dataproc.py @@ -30,13 +30,17 @@ from google.cloud.dataproc_v1 import ( Batch, BatchControllerClient, + BatchControllerAsyncClient, Cluster, ClusterControllerClient, + ClusterControllerAsyncClient, Job, JobControllerClient, + JobControllerAsyncClient, JobStatus, WorkflowTemplate, WorkflowTemplateServiceClient, + WorkflowTemplateServiceAsyncClient, ) from google.protobuf.duration_pb2 import Duration from google.protobuf.field_mask_pb2 import FieldMask @@ -205,10 +209,8 @@ def __init__( gcp_conn_id: str = 'google_cloud_default', delegate_to: Optional[str] = None, impersonation_chain: Optional[Union[str, Sequence[str]]] = None, - async_client: bool = False, ) -> None: super().__init__(gcp_conn_id, delegate_to, impersonation_chain) - self.async_client = async_client def get_cluster_client(self, region: Optional[str] = None) -> ClusterControllerClient: """Returns ClusterControllerClient.""" @@ -240,7 +242,6 @@ def get_job_client(self, region: Optional[str] = None) -> JobControllerClient: credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options, - transport="grpc_asyncio" if self.async_client else "grpc", ) def get_batch_client(self, region: Optional[str] = None) -> BatchControllerClient: @@ -971,3 +972,736 @@ def list_batches( metadata=metadata, ) return result + + +class DataprocAsyncHook(GoogleBaseHook): + """ + Asynchronuous Hook for Google Cloud Dataproc APIs. + + All the methods in the hook where project_id is used must be called with + keyword arguments rather than positional. + """ + + def __init__( + self, + gcp_conn_id: str = 'google_cloud_default', + delegate_to: Optional[str] = None, + impersonation_chain: Optional[Union[str, Sequence[str]]] = None, + ) -> None: + super().__init__(gcp_conn_id, delegate_to, impersonation_chain) + + def get_cluster_client(self, region: Optional[str] = None) -> ClusterControllerAsyncClient: + """Returns ClusterControllerAsyncClient.""" + client_options = None + if region and region != 'global': + client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443') + + return ClusterControllerAsyncClient( + credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options + ) + + def get_template_client(self, region: Optional[str] = None) -> WorkflowTemplateServiceAsyncClient: + """Returns WorkflowTemplateServiceAsyncClient.""" + client_options = None + if region and region != 'global': + client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443') + + return WorkflowTemplateServiceAsyncClient( + credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options + ) + + def get_job_client(self, region: Optional[str] = None) -> JobControllerAsyncClient: + """Returns JobControllerAsyncClient.""" + client_options = None + if region and region != 'global': + client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443') + + return JobControllerAsyncClient( + credentials=self._get_credentials(), + client_info=CLIENT_INFO, + client_options=client_options, + ) + + def get_batch_client(self, region: Optional[str] = None) -> BatchControllerAsyncClient: + """Returns BatchControllerAsyncClient""" + client_options = None + if region and region != 'global': + client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443') + + return BatchControllerAsyncClient( + credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options + ) + + @GoogleBaseHook.fallback_to_default_project_id + async def create_cluster( + self, + region: str, + project_id: str, + cluster_name: str, + cluster_config: Union[Dict, Cluster, None] = None, + virtual_cluster_config: Optional[Dict] = None, + labels: Optional[Dict[str, str]] = None, + request_id: Optional[str] = None, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ): + """ + Creates a cluster in a project. + + :param project_id: Required. The ID of the Google Cloud project that the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param cluster_name: Name of the cluster to create + :param labels: Labels that will be assigned to created cluster + :param cluster_config: Required. The cluster config to create. + If a dict is provided, it must be of the same form as the protobuf message + :class:`~google.cloud.dataproc_v1.types.ClusterConfig` + :param virtual_cluster_config: Optional. The virtual cluster config, used when creating a Dataproc + cluster that does not directly control the underlying compute resources, for example, when + creating a `Dataproc-on-GKE cluster` + :class:`~google.cloud.dataproc_v1.types.VirtualClusterConfig` + :param request_id: Optional. A unique id used to identify the request. If the server receives two + ``CreateClusterRequest`` requests with the same id, then the second request will be ignored and + the first ``google.longrunning.Operation`` created and stored in the backend is returned. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + # Dataproc labels must conform to the following regex: + # [a-z]([-a-z0-9]*[a-z0-9])? (current airflow version string follows + # semantic versioning spec: x.y.z). + labels = labels or {} + labels.update({'airflow-version': 'v' + airflow_version.replace('.', '-').replace('+', '-')}) + + cluster = { + "project_id": project_id, + "cluster_name": cluster_name, + } + if virtual_cluster_config is not None: + cluster['virtual_cluster_config'] = virtual_cluster_config # type: ignore + if cluster_config is not None: + cluster['config'] = cluster_config # type: ignore + cluster['labels'] = labels # type: ignore + + client = self.get_cluster_client(region=region) + result = await client.create_cluster( + request={ + 'project_id': project_id, + 'region': region, + 'cluster': cluster, + 'request_id': request_id, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + async def delete_cluster( + self, + region: str, + cluster_name: str, + project_id: str, + cluster_uuid: Optional[str] = None, + request_id: Optional[str] = None, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ): + """ + Deletes a cluster in a project. + + :param project_id: Required. The ID of the Google Cloud project that the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param cluster_name: Required. The cluster name. + :param cluster_uuid: Optional. Specifying the ``cluster_uuid`` means the RPC should fail + if cluster with specified UUID does not exist. + :param request_id: Optional. A unique id used to identify the request. If the server receives two + ``DeleteClusterRequest`` requests with the same id, then the second request will be ignored and + the first ``google.longrunning.Operation`` created and stored in the backend is returned. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + client = self.get_cluster_client(region=region) + result = client.delete_cluster( + request={ + 'project_id': project_id, + 'region': region, + 'cluster_name': cluster_name, + 'cluster_uuid': cluster_uuid, + 'request_id': request_id, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + async def diagnose_cluster( + self, + region: str, + cluster_name: str, + project_id: str, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ): + """ + Gets cluster diagnostic information. After the operation completes GCS uri to + diagnose is returned + + :param project_id: Required. The ID of the Google Cloud project that the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param cluster_name: Required. The cluster name. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + client = self.get_cluster_client(region=region) + operation = await client.diagnose_cluster( + request={'project_id': project_id, 'region': region, 'cluster_name': cluster_name}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + operation.result() + gcs_uri = str(operation.operation.response.value) + return gcs_uri + + @GoogleBaseHook.fallback_to_default_project_id + async def get_cluster( + self, + region: str, + cluster_name: str, + project_id: str, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ): + """ + Gets the resource representation for a cluster in a project. + + :param project_id: Required. The ID of the Google Cloud project that the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param cluster_name: Required. The cluster name. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + client = self.get_cluster_client(region=region) + result = await client.get_cluster( + request={'project_id': project_id, 'region': region, 'cluster_name': cluster_name}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + async def list_clusters( + self, + region: str, + filter_: str, + project_id: str, + page_size: Optional[int] = None, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ): + """ + Lists all regions/{region}/clusters in a project. + + :param project_id: Required. The ID of the Google Cloud project that the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param filter_: Optional. A filter constraining the clusters to list. Filters are case-sensitive. + :param page_size: The maximum number of resources contained in the underlying API response. If page + streaming is performed per- resource, this parameter does not affect the return value. If page + streaming is performed per-page, this determines the maximum number of resources in a page. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + client = self.get_cluster_client(region=region) + result = await client.list_clusters( + request={'project_id': project_id, 'region': region, 'filter': filter_, 'page_size': page_size}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + async def update_cluster( + self, + cluster_name: str, + cluster: Union[Dict, Cluster], + update_mask: Union[Dict, FieldMask], + project_id: str, + region: str, + graceful_decommission_timeout: Optional[Union[Dict, Duration]] = None, + request_id: Optional[str] = None, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ): + """ + Updates a cluster in a project. + + :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param cluster_name: Required. The cluster name. + :param cluster: Required. The changes to the cluster. + + If a dict is provided, it must be of the same form as the protobuf message + :class:`~google.cloud.dataproc_v1.types.Cluster` + :param update_mask: Required. Specifies the path, relative to ``Cluster``, of the field to update. For + example, to change the number of workers in a cluster to 5, the ``update_mask`` parameter would be + specified as ``config.worker_config.num_instances``, and the ``PATCH`` request body would specify + the new value, as follows: + + :: + + { "config":{ "workerConfig":{ "numInstances":"5" } } } + + Similarly, to change the number of preemptible workers in a cluster to 5, the ``update_mask`` + parameter would be ``config.secondary_worker_config.num_instances``, and the ``PATCH`` request + body would be set as follows: + + :: + + { "config":{ "secondaryWorkerConfig":{ "numInstances":"5" } } } + + If a dict is provided, it must be of the same form as the protobuf message + :class:`~google.cloud.dataproc_v1.types.FieldMask` + :param graceful_decommission_timeout: Optional. Timeout for graceful YARN decommissioning. Graceful + decommissioning allows removing nodes from the cluster without interrupting jobs in progress. + Timeout specifies how long to wait for jobs in progress to finish before forcefully removing nodes + (and potentially interrupting jobs). Default timeout is 0 (for forceful decommission), and the + maximum allowed timeout is 1 day. + + Only supported on Dataproc image versions 1.2 and higher. + + If a dict is provided, it must be of the same form as the protobuf message + :class:`~google.cloud.dataproc_v1.types.Duration` + :param request_id: Optional. A unique id used to identify the request. If the server receives two + ``UpdateClusterRequest`` requests with the same id, then the second request will be ignored and + the first ``google.longrunning.Operation`` created and stored in the backend is returned. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + if region is None: + raise TypeError("missing 1 required keyword argument: 'region'") + client = self.get_cluster_client(region=region) + operation = await client.update_cluster( + request={ + 'project_id': project_id, + 'region': region, + 'cluster_name': cluster_name, + 'cluster': cluster, + 'update_mask': update_mask, + 'graceful_decommission_timeout': graceful_decommission_timeout, + 'request_id': request_id, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return operation + + @GoogleBaseHook.fallback_to_default_project_id + async def create_workflow_template( + self, + template: Union[Dict, WorkflowTemplate], + project_id: str, + region: str, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> WorkflowTemplate: + """ + Creates new workflow template. + + :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param template: The Dataproc workflow template to create. If a dict is provided, + it must be of the same form as the protobuf message WorkflowTemplate. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + if region is None: + raise TypeError("missing 1 required keyword argument: 'region'") + metadata = metadata or () + client = self.get_template_client(region) + parent = f'projects/{project_id}/regions/{region}' + return await client.create_workflow_template( + request={'parent': parent, 'template': template}, retry=retry, timeout=timeout, metadata=metadata + ) + + @GoogleBaseHook.fallback_to_default_project_id + async def instantiate_workflow_template( + self, + template_name: str, + project_id: str, + region: str, + version: Optional[int] = None, + request_id: Optional[str] = None, + parameters: Optional[Dict[str, str]] = None, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ): + """ + Instantiates a template and begins execution. + + :param template_name: Name of template to instantiate. + :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param version: Optional. The version of workflow template to instantiate. If specified, + the workflow will be instantiated only if the current version of + the workflow template has the supplied version. + This option cannot be used to instantiate a previous version of + workflow template. + :param request_id: Optional. A tag that prevents multiple concurrent workflow instances + with the same tag from running. This mitigates risk of concurrent + instances started due to retries. + :param parameters: Optional. Map from parameter names to values that should be used for those + parameters. Values may not exceed 100 characters. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + if region is None: + raise TypeError("missing 1 required keyword argument: 'region'") + metadata = metadata or () + client = self.get_template_client(region) + name = f'projects/{project_id}/regions/{region}/workflowTemplates/{template_name}' + operation = await client.instantiate_workflow_template( + request={'name': name, 'version': version, 'request_id': request_id, 'parameters': parameters}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return operation + + @GoogleBaseHook.fallback_to_default_project_id + async def instantiate_inline_workflow_template( + self, + template: Union[Dict, WorkflowTemplate], + project_id: str, + region: str, + request_id: Optional[str] = None, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ): + """ + Instantiates a template and begins execution. + + :param template: The workflow template to instantiate. If a dict is provided, + it must be of the same form as the protobuf message WorkflowTemplate + :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param request_id: Optional. A tag that prevents multiple concurrent workflow instances + with the same tag from running. This mitigates risk of concurrent + instances started due to retries. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + if region is None: + raise TypeError("missing 1 required keyword argument: 'region'") + metadata = metadata or () + client = self.get_template_client(region) + parent = f'projects/{project_id}/regions/{region}' + operation = await client.instantiate_inline_workflow_template( + request={'parent': parent, 'template': template, 'request_id': request_id}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return operation + + @GoogleBaseHook.fallback_to_default_project_id + async def get_job( + self, + job_id: str, + project_id: str, + region: str, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> Job: + """ + Gets the resource representation for a job in a project. + + :param job_id: Id of the Dataproc job + :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + if region is None: + raise TypeError("missing 1 required keyword argument: 'region'") + client = self.get_job_client(region=region) + job = await client.get_job( + request={'project_id': project_id, 'region': region, 'job_id': job_id}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return job + + @GoogleBaseHook.fallback_to_default_project_id + async def submit_job( + self, + job: Union[dict, Job], + project_id: str, + region: str, + request_id: Optional[str] = None, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> Job: + """ + Submits a job to a cluster. + + :param job: The job resource. If a dict is provided, + it must be of the same form as the protobuf message Job + :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param request_id: Optional. A tag that prevents multiple concurrent workflow instances + with the same tag from running. This mitigates risk of concurrent + instances started due to retries. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + if region is None: + raise TypeError("missing 1 required keyword argument: 'region'") + client = self.get_job_client(region=region) + return await client.submit_job( + request={'project_id': project_id, 'region': region, 'job': job, 'request_id': request_id}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + @GoogleBaseHook.fallback_to_default_project_id + async def cancel_job( + self, + job_id: str, + project_id: str, + region: Optional[str] = None, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> Job: + """ + Starts a job cancellation request. + + :param project_id: Required. The ID of the Google Cloud project that the job belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param job_id: Required. The job ID. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + client = self.get_job_client(region=region) + + job = await client.cancel_job( + request={'project_id': project_id, 'region': region, 'job_id': job_id}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return job + + @GoogleBaseHook.fallback_to_default_project_id + async def create_batch( + self, + region: str, + project_id: str, + batch: Union[Dict, Batch], + batch_id: Optional[str] = None, + request_id: Optional[str] = None, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> Operation: + """ + Creates a batch workload. + + :param project_id: Required. The ID of the Google Cloud project that the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param batch: Required. The batch to create. + :param batch_id: Optional. The ID to use for the batch, which will become the final component + of the batch's resource name. + This value must be 4-63 characters. Valid characters are /[a-z][0-9]-/. + :param request_id: Optional. A unique id used to identify the request. If the server receives two + ``CreateBatchRequest`` requests with the same id, then the second request will be ignored and + the first ``google.longrunning.Operation`` created and stored in the backend is returned. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + client = self.get_batch_client(region) + parent = f'projects/{project_id}/regions/{region}' + + result = await client.create_batch( + request={ + 'parent': parent, + 'batch': batch, + 'batch_id': batch_id, + 'request_id': request_id, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + async def delete_batch( + self, + batch_id: str, + region: str, + project_id: str, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> None: + """ + Deletes the batch workload resource. + + :param batch_id: Required. The ID to use for the batch, which will become the final component + of the batch's resource name. + This value must be 4-63 characters. Valid characters are /[a-z][0-9]-/. + :param project_id: Required. The ID of the Google Cloud project that the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + client = self.get_batch_client(region) + name = f"projects/{project_id}/regions/{region}/batches/{batch_id}" + + await client.delete_batch( + request={ + 'name': name, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + @GoogleBaseHook.fallback_to_default_project_id + async def get_batch( + self, + batch_id: str, + region: str, + project_id: str, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> Batch: + """ + Gets the batch workload resource representation. + + :param batch_id: Required. The ID to use for the batch, which will become the final component + of the batch's resource name. + This value must be 4-63 characters. Valid characters are /[a-z][0-9]-/. + :param project_id: Required. The ID of the Google Cloud project that the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + client = self.get_batch_client(region) + name = f"projects/{project_id}/regions/{region}/batches/{batch_id}" + + result = await client.get_batch( + request={ + 'name': name, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + async def list_batches( + self, + region: str, + project_id: str, + page_size: Optional[int] = None, + page_token: Optional[str] = None, + retry: Union[Retry, _MethodDefault] = DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ): + """ + Lists batch workloads. + + :param project_id: Required. The ID of the Google Cloud project that the cluster belongs to. + :param region: Required. The Cloud Dataproc region in which to handle the request. + :param page_size: Optional. The maximum number of batches to return in each response. The service may + return fewer than this value. The default page size is 20; the maximum page size is 1000. + :param page_token: Optional. A page token received from a previous ``ListBatches`` call. + Provide this token to retrieve the subsequent page. + :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be + retried. + :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if + ``retry`` is specified, the timeout applies to each individual attempt. + :param metadata: Additional metadata that is provided to the method. + """ + client = self.get_batch_client(region) + parent = f'projects/{project_id}/regions/{region}' + + result = await client.list_batches( + request={ + 'parent': parent, + 'page_size': page_size, + 'page_token': page_token, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result diff --git a/airflow/providers/google/cloud/triggers/dataproc.py b/airflow/providers/google/cloud/triggers/dataproc.py index 1a2a1a952aa1d..5467ff91b65b4 100644 --- a/airflow/providers/google/cloud/triggers/dataproc.py +++ b/airflow/providers/google/cloud/triggers/dataproc.py @@ -1,7 +1,7 @@ import asyncio from airflow import AirflowException -from airflow.providers.google.cloud.hooks.dataproc import DataprocHook +from airflow.providers.google.cloud.hooks.dataproc import DataprocAsyncHook from airflow.triggers.base import BaseTrigger, TriggerEvent from google.cloud.dataproc_v1 import JobStatus, Job from typing import Any, Dict, Optional, Sequence, Tuple, Union @@ -31,11 +31,10 @@ def __init__( self.region = region self.pooling_period_seconds = pooling_period_seconds self.delegate_to = delegate_to - self.hook = DataprocHook( + self.hook = DataprocAsyncHook( delegate_to=self.delegate_to, gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, - async_client=True, ) def serialize(self): @@ -52,14 +51,9 @@ def serialize(self): }, ) - async def get_job(self) -> Job: - return await self.hook.get_job(job_id=self.job_id, project_id=self.project_id, region=self.region) - async def run(self): while True: - job: Job = await self.hook.get_job( - project_id=self.project_id, region=self.region, job_id=self.job_id - ) + job = await self.hook.get_job(project_id=self.project_id, region=self.region, job_id=self.job_id) state = job.status.state self.log.info("Dataproc job: %s is in state: %s", self.job_id, state) if state in (JobStatus.State.ERROR, JobStatus.State.DONE, JobStatus.State.CANCELLED): From 0ad986839a6a2fb18a02b8dc5f88324db0617fdb Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 26 Jul 2022 09:13:40 +0200 Subject: [PATCH 12/29] DataprocAsyncHook tests --- .../google/cloud/hooks/test_dataproc.py | 405 +++++++++++++++++- 1 file changed, 404 insertions(+), 1 deletion(-) diff --git a/tests/providers/google/cloud/hooks/test_dataproc.py b/tests/providers/google/cloud/hooks/test_dataproc.py index a8d91daf409d7..85bc28242c640 100644 --- a/tests/providers/google/cloud/hooks/test_dataproc.py +++ b/tests/providers/google/cloud/hooks/test_dataproc.py @@ -26,7 +26,7 @@ from parameterized import parameterized from airflow.exceptions import AirflowException -from airflow.providers.google.cloud.hooks.dataproc import DataprocHook, DataProcJobBuilder +from airflow.providers.google.cloud.hooks.dataproc import DataprocAsyncHook, DataprocHook, DataProcJobBuilder from airflow.providers.google.common.consts import CLIENT_INFO from airflow.version import version @@ -462,6 +462,409 @@ def test_list_batches(self, mock_client): ) +class TestDataprocAsyncHook(unittest.TestCase): + def setUp(self): + with mock.patch(BASE_STRING.format("GoogleBaseHook.__init__"), new=mock_init): + self.hook = DataprocAsyncHook(gcp_conn_id="test") + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("ClusterControllerAsyncClient")) + def test_get_cluster_client(self, mock_client, mock_get_credentials): + self.hook.get_cluster_client(region=GCP_LOCATION) + mock_client.assert_called_once_with( + credentials=mock_get_credentials.return_value, + client_info=CLIENT_INFO, + client_options=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("ClusterControllerAsyncClient")) + def test_get_cluster_client_region(self, mock_client, mock_get_credentials): + self.hook.get_cluster_client(region='region1') + mock_client.assert_called_once_with( + credentials=mock_get_credentials.return_value, + client_info=CLIENT_INFO, + client_options=ANY, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("WorkflowTemplateServiceAsyncClient")) + def test_get_template_client_global(self, mock_client, mock_get_credentials): + _ = self.hook.get_template_client() + mock_client.assert_called_once_with( + credentials=mock_get_credentials.return_value, + client_info=CLIENT_INFO, + client_options=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("WorkflowTemplateServiceAsyncClient")) + def test_get_template_client_region(self, mock_client, mock_get_credentials): + _ = self.hook.get_template_client(region='region1') + mock_client.assert_called_once_with( + credentials=mock_get_credentials.return_value, + client_info=CLIENT_INFO, + client_options=ANY, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("JobControllerAsyncClient")) + def test_get_job_client(self, mock_client, mock_get_credentials): + self.hook.get_job_client(region=GCP_LOCATION) + mock_client.assert_called_once_with( + credentials=mock_get_credentials.return_value, + client_info=CLIENT_INFO, + client_options=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("JobControllerAsyncClient")) + def test_get_job_client_region(self, mock_client, mock_get_credentials): + self.hook.get_job_client(region='region1') + mock_client.assert_called_once_with( + credentials=mock_get_credentials.return_value, + client_info=CLIENT_INFO, + client_options=ANY, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("BatchControllerAsyncClient")) + def test_get_batch_client(self, mock_client, mock_get_credentials): + self.hook.get_batch_client(region=GCP_LOCATION) + mock_client.assert_called_once_with( + credentials=mock_get_credentials.return_value, + client_info=CLIENT_INFO, + client_options=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("BatchControllerAsyncClient")) + def test_get_batch_client_region(self, mock_client, mock_get_credentials): + self.hook.get_batch_client(region='region1') + mock_client.assert_called_once_with( + credentials=mock_get_credentials.return_value, client_info=CLIENT_INFO, client_options=ANY + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) + def test_create_cluster(self, mock_client): + self.hook.create_cluster( + project_id=GCP_PROJECT, + region=GCP_LOCATION, + cluster_name=CLUSTER_NAME, + cluster_config=CLUSTER_CONFIG, + labels=LABELS, + ) + mock_client.assert_called_once_with(region=GCP_LOCATION) + mock_client.return_value.create_cluster.assert_called_once_with( + request=dict( + project_id=GCP_PROJECT, + region=GCP_LOCATION, + cluster=CLUSTER, + request_id=None, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) + def test_delete_cluster(self, mock_client): + self.hook.delete_cluster(project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME) + mock_client.assert_called_once_with(region=GCP_LOCATION) + mock_client.return_value.delete_cluster.assert_called_once_with( + request=dict( + project_id=GCP_PROJECT, + region=GCP_LOCATION, + cluster_name=CLUSTER_NAME, + cluster_uuid=None, + request_id=None, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) + def test_diagnose_cluster(self, mock_client): + self.hook.diagnose_cluster(project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME) + mock_client.assert_called_once_with(region=GCP_LOCATION) + mock_client.return_value.diagnose_cluster.assert_called_once_with( + request=dict( + project_id=GCP_PROJECT, + region=GCP_LOCATION, + cluster_name=CLUSTER_NAME, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + mock_client.return_value.diagnose_cluster.return_value.result.assert_called_once_with() + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) + def test_get_cluster(self, mock_client): + self.hook.get_cluster(project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME) + mock_client.assert_called_once_with(region=GCP_LOCATION) + mock_client.return_value.get_cluster.assert_called_once_with( + request=dict( + project_id=GCP_PROJECT, + region=GCP_LOCATION, + cluster_name=CLUSTER_NAME, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) + def test_list_clusters(self, mock_client): + filter_ = "filter" + + self.hook.list_clusters(project_id=GCP_PROJECT, region=GCP_LOCATION, filter_=filter_) + mock_client.assert_called_once_with(region=GCP_LOCATION) + mock_client.return_value.list_clusters.assert_called_once_with( + request=dict( + project_id=GCP_PROJECT, + region=GCP_LOCATION, + filter=filter_, + page_size=None, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) + def test_update_cluster(self, mock_client): + update_mask = "update-mask" + self.hook.update_cluster( + project_id=GCP_PROJECT, + region=GCP_LOCATION, + cluster=CLUSTER, + cluster_name=CLUSTER_NAME, + update_mask=update_mask, + ) + mock_client.assert_called_once_with(region=GCP_LOCATION) + mock_client.return_value.update_cluster.assert_called_once_with( + request=dict( + project_id=GCP_PROJECT, + region=GCP_LOCATION, + cluster=CLUSTER, + cluster_name=CLUSTER_NAME, + update_mask=update_mask, + graceful_decommission_timeout=None, + request_id=None, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) + def test_update_cluster_missing_region(self, mock_client): + with pytest.raises(TypeError): + self.hook.update_cluster( + project_id=GCP_PROJECT, + cluster=CLUSTER, + cluster_name=CLUSTER_NAME, + update_mask="update-mask", + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_template_client")) + def test_create_workflow_template(self, mock_client): + template = {"test": "test"} + parent = f'projects/{GCP_PROJECT}/regions/{GCP_LOCATION}' + self.hook.create_workflow_template(region=GCP_LOCATION, template=template, project_id=GCP_PROJECT) + mock_client.return_value.create_workflow_template.assert_called_once_with( + request=dict(parent=parent, template=template), retry=DEFAULT, timeout=None, metadata=() + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_template_client")) + def test_instantiate_workflow_template(self, mock_client): + template_name = "template_name" + name = f'projects/{GCP_PROJECT}/regions/{GCP_LOCATION}/workflowTemplates/{template_name}' + self.hook.instantiate_workflow_template( + region=GCP_LOCATION, template_name=template_name, project_id=GCP_PROJECT + ) + mock_client.return_value.instantiate_workflow_template.assert_called_once_with( + request=dict(name=name, version=None, parameters=None, request_id=None), + retry=DEFAULT, + timeout=None, + metadata=(), + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_template_client")) + def test_instantiate_workflow_template_missing_region(self, mock_client): + with pytest.raises(TypeError): + self.hook.instantiate_workflow_template(template_name="template_name", project_id=GCP_PROJECT) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_template_client")) + def test_instantiate_inline_workflow_template(self, mock_client): + template = {"test": "test"} + parent = f'projects/{GCP_PROJECT}/regions/{GCP_LOCATION}' + self.hook.instantiate_inline_workflow_template( + region=GCP_LOCATION, template=template, project_id=GCP_PROJECT + ) + mock_client.return_value.instantiate_inline_workflow_template.assert_called_once_with( + request=dict(parent=parent, template=template, request_id=None), + retry=DEFAULT, + timeout=None, + metadata=(), + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_template_client")) + def test_instantiate_inline_workflow_template_missing_region(self, mock_client): + with pytest.raises(TypeError): + self.hook.instantiate_inline_workflow_template(template={"test": "test"}, project_id=GCP_PROJECT) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job")) + def test_wait_for_job(self, mock_get_job): + mock_get_job.side_effect = [ + mock.MagicMock(status=mock.MagicMock(state=JobStatus.State.RUNNING)), + mock.MagicMock(status=mock.MagicMock(state=JobStatus.State.ERROR)), + ] + with pytest.raises(AirflowException): + self.hook.wait_for_job(job_id=JOB_ID, region=GCP_LOCATION, project_id=GCP_PROJECT, wait_time=0) + calls = [ + mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), + mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), + ] + mock_get_job.assert_has_calls(calls) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job")) + def test_wait_for_job_missing_region(self, mock_get_job): + with pytest.raises(TypeError): + self.hook.wait_for_job(job_id=JOB_ID, project_id=GCP_PROJECT, wait_time=0) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job_client")) + def test_get_job(self, mock_client): + self.hook.get_job(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT) + mock_client.assert_called_once_with(region=GCP_LOCATION) + mock_client.return_value.get_job.assert_called_once_with( + request=dict( + region=GCP_LOCATION, + job_id=JOB_ID, + project_id=GCP_PROJECT, + ), + retry=DEFAULT, + timeout=None, + metadata=(), + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job_client")) + def test_get_job_missing_region(self, mock_client): + with pytest.raises(TypeError): + self.hook.get_job(job_id=JOB_ID, project_id=GCP_PROJECT) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job_client")) + def test_submit_job(self, mock_client): + self.hook.submit_job(region=GCP_LOCATION, job=JOB, project_id=GCP_PROJECT) + mock_client.assert_called_once_with(region=GCP_LOCATION) + mock_client.return_value.submit_job.assert_called_once_with( + request=dict( + region=GCP_LOCATION, + job=JOB, + project_id=GCP_PROJECT, + request_id=None, + ), + retry=DEFAULT, + timeout=None, + metadata=(), + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job_client")) + def test_submit_job_missing_region(self, mock_client): + with pytest.raises(TypeError): + self.hook.submit_job(job=JOB, project_id=GCP_PROJECT) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job_client")) + def test_cancel_job(self, mock_client): + self.hook.cancel_job(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT) + mock_client.assert_called_once_with(region=GCP_LOCATION) + mock_client.return_value.cancel_job.assert_called_once_with( + request=dict( + region=GCP_LOCATION, + job_id=JOB_ID, + project_id=GCP_PROJECT, + ), + retry=DEFAULT, + timeout=None, + metadata=(), + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_batch_client")) + def test_create_batch(self, mock_client): + self.hook.create_batch( + project_id=GCP_PROJECT, + region=GCP_LOCATION, + batch=BATCH, + batch_id=BATCH_ID, + ) + mock_client.assert_called_once_with(GCP_LOCATION) + mock_client.return_value.create_batch.assert_called_once_with( + request=dict( + parent=PARENT.format(GCP_PROJECT, GCP_LOCATION), + batch=BATCH, + batch_id=BATCH_ID, + request_id=None, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_batch_client")) + def test_delete_batch(self, mock_client): + self.hook.delete_batch( + batch_id=BATCH_ID, + region=GCP_LOCATION, + project_id=GCP_PROJECT, + ) + mock_client.assert_called_once_with(GCP_LOCATION) + mock_client.return_value.delete_batch.assert_called_once_with( + request=dict( + name=BATCH_NAME.format(GCP_PROJECT, GCP_LOCATION, BATCH_ID), + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_batch_client")) + def test_get_batch(self, mock_client): + self.hook.get_batch( + batch_id=BATCH_ID, + region=GCP_LOCATION, + project_id=GCP_PROJECT, + ) + mock_client.assert_called_once_with(GCP_LOCATION) + mock_client.return_value.get_batch.assert_called_once_with( + request=dict( + name=BATCH_NAME.format(GCP_PROJECT, GCP_LOCATION, BATCH_ID), + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_batch_client")) + def test_list_batches(self, mock_client): + self.hook.list_batches( + project_id=GCP_PROJECT, + region=GCP_LOCATION, + ) + mock_client.assert_called_once_with(GCP_LOCATION) + mock_client.return_value.list_batches.assert_called_once_with( + request=dict( + parent=PARENT.format(GCP_PROJECT, GCP_LOCATION), + page_size=None, + page_token=None, + ), + metadata=(), + retry=DEFAULT, + timeout=None, + ) + + class TestDataProcJobBuilder(unittest.TestCase): def setUp(self) -> None: self.job_type = "test" From 8ec99be25fca3f61be10d349682fd7dfbf974f75 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 26 Jul 2022 09:25:23 +0200 Subject: [PATCH 13/29] Async tests --- tests/providers/google/cloud/hooks/test_dataproc.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/providers/google/cloud/hooks/test_dataproc.py b/tests/providers/google/cloud/hooks/test_dataproc.py index 85bc28242c640..c8807d2c0e897 100644 --- a/tests/providers/google/cloud/hooks/test_dataproc.py +++ b/tests/providers/google/cloud/hooks/test_dataproc.py @@ -776,9 +776,10 @@ def test_submit_job_missing_region(self, mock_client): with pytest.raises(TypeError): self.hook.submit_job(job=JOB, project_id=GCP_PROJECT) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job_client")) - def test_cancel_job(self, mock_client): - self.hook.cancel_job(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT) + async def test_cancel_job(self, mock_client): + await self.hook.cancel_job(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT) mock_client.assert_called_once_with(region=GCP_LOCATION) mock_client.return_value.cancel_job.assert_called_once_with( request=dict( From d0122cd0c6f117d8ededfc1757e9a8221e9f0dc1 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 26 Jul 2022 09:31:04 +0200 Subject: [PATCH 14/29] Async tests --- .../google/cloud/hooks/test_dataproc.py | 94 +++++++++---------- 1 file changed, 46 insertions(+), 48 deletions(-) diff --git a/tests/providers/google/cloud/hooks/test_dataproc.py b/tests/providers/google/cloud/hooks/test_dataproc.py index c8807d2c0e897..828e979404040 100644 --- a/tests/providers/google/cloud/hooks/test_dataproc.py +++ b/tests/providers/google/cloud/hooks/test_dataproc.py @@ -545,9 +545,10 @@ def test_get_batch_client_region(self, mock_client, mock_get_credentials): credentials=mock_get_credentials.return_value, client_info=CLIENT_INFO, client_options=ANY ) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) - def test_create_cluster(self, mock_client): - self.hook.create_cluster( + async def test_create_cluster(self, mock_client): + await self.hook.create_cluster( project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME, @@ -567,9 +568,10 @@ def test_create_cluster(self, mock_client): timeout=None, ) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) - def test_delete_cluster(self, mock_client): - self.hook.delete_cluster(project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME) + async def test_delete_cluster(self, mock_client): + await self.hook.delete_cluster(project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME) mock_client.assert_called_once_with(region=GCP_LOCATION) mock_client.return_value.delete_cluster.assert_called_once_with( request=dict( @@ -584,9 +586,10 @@ def test_delete_cluster(self, mock_client): timeout=None, ) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) def test_diagnose_cluster(self, mock_client): - self.hook.diagnose_cluster(project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME) + async self.hook.diagnose_cluster(project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME) mock_client.assert_called_once_with(region=GCP_LOCATION) mock_client.return_value.diagnose_cluster.assert_called_once_with( request=dict( @@ -600,9 +603,10 @@ def test_diagnose_cluster(self, mock_client): ) mock_client.return_value.diagnose_cluster.return_value.result.assert_called_once_with() + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) - def test_get_cluster(self, mock_client): - self.hook.get_cluster(project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME) + async def test_get_cluster(self, mock_client): + await self.hook.get_cluster(project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME) mock_client.assert_called_once_with(region=GCP_LOCATION) mock_client.return_value.get_cluster.assert_called_once_with( request=dict( @@ -615,11 +619,12 @@ def test_get_cluster(self, mock_client): timeout=None, ) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) - def test_list_clusters(self, mock_client): + async def test_list_clusters(self, mock_client): filter_ = "filter" - self.hook.list_clusters(project_id=GCP_PROJECT, region=GCP_LOCATION, filter_=filter_) + await self.hook.list_clusters(project_id=GCP_PROJECT, region=GCP_LOCATION, filter_=filter_) mock_client.assert_called_once_with(region=GCP_LOCATION) mock_client.return_value.list_clusters.assert_called_once_with( request=dict( @@ -633,10 +638,11 @@ def test_list_clusters(self, mock_client): timeout=None, ) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) - def test_update_cluster(self, mock_client): + async def test_update_cluster(self, mock_client): update_mask = "update-mask" - self.hook.update_cluster( + await self.hook.update_cluster( project_id=GCP_PROJECT, region=GCP_LOCATION, cluster=CLUSTER, @@ -669,20 +675,24 @@ def test_update_cluster_missing_region(self, mock_client): update_mask="update-mask", ) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_template_client")) - def test_create_workflow_template(self, mock_client): + async def test_create_workflow_template(self, mock_client): template = {"test": "test"} parent = f'projects/{GCP_PROJECT}/regions/{GCP_LOCATION}' - self.hook.create_workflow_template(region=GCP_LOCATION, template=template, project_id=GCP_PROJECT) + await self.hook.create_workflow_template( + region=GCP_LOCATION, template=template, project_id=GCP_PROJECT + ) mock_client.return_value.create_workflow_template.assert_called_once_with( request=dict(parent=parent, template=template), retry=DEFAULT, timeout=None, metadata=() ) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_template_client")) - def test_instantiate_workflow_template(self, mock_client): + async def test_instantiate_workflow_template(self, mock_client): template_name = "template_name" name = f'projects/{GCP_PROJECT}/regions/{GCP_LOCATION}/workflowTemplates/{template_name}' - self.hook.instantiate_workflow_template( + await self.hook.instantiate_workflow_template( region=GCP_LOCATION, template_name=template_name, project_id=GCP_PROJECT ) mock_client.return_value.instantiate_workflow_template.assert_called_once_with( @@ -697,11 +707,12 @@ def test_instantiate_workflow_template_missing_region(self, mock_client): with pytest.raises(TypeError): self.hook.instantiate_workflow_template(template_name="template_name", project_id=GCP_PROJECT) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_template_client")) - def test_instantiate_inline_workflow_template(self, mock_client): + async def test_instantiate_inline_workflow_template(self, mock_client): template = {"test": "test"} parent = f'projects/{GCP_PROJECT}/regions/{GCP_LOCATION}' - self.hook.instantiate_inline_workflow_template( + await self.hook.instantiate_inline_workflow_template( region=GCP_LOCATION, template=template, project_id=GCP_PROJECT ) mock_client.return_value.instantiate_inline_workflow_template.assert_called_once_with( @@ -716,28 +727,10 @@ def test_instantiate_inline_workflow_template_missing_region(self, mock_client): with pytest.raises(TypeError): self.hook.instantiate_inline_workflow_template(template={"test": "test"}, project_id=GCP_PROJECT) - @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job")) - def test_wait_for_job(self, mock_get_job): - mock_get_job.side_effect = [ - mock.MagicMock(status=mock.MagicMock(state=JobStatus.State.RUNNING)), - mock.MagicMock(status=mock.MagicMock(state=JobStatus.State.ERROR)), - ] - with pytest.raises(AirflowException): - self.hook.wait_for_job(job_id=JOB_ID, region=GCP_LOCATION, project_id=GCP_PROJECT, wait_time=0) - calls = [ - mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), - mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), - ] - mock_get_job.assert_has_calls(calls) - - @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job")) - def test_wait_for_job_missing_region(self, mock_get_job): - with pytest.raises(TypeError): - self.hook.wait_for_job(job_id=JOB_ID, project_id=GCP_PROJECT, wait_time=0) - + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job_client")) - def test_get_job(self, mock_client): - self.hook.get_job(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT) + async def test_get_job(self, mock_client): + await self.hook.get_job(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT) mock_client.assert_called_once_with(region=GCP_LOCATION) mock_client.return_value.get_job.assert_called_once_with( request=dict( @@ -755,9 +748,10 @@ def test_get_job_missing_region(self, mock_client): with pytest.raises(TypeError): self.hook.get_job(job_id=JOB_ID, project_id=GCP_PROJECT) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_job_client")) - def test_submit_job(self, mock_client): - self.hook.submit_job(region=GCP_LOCATION, job=JOB, project_id=GCP_PROJECT) + async def test_submit_job(self, mock_client): + await self.hook.submit_job(region=GCP_LOCATION, job=JOB, project_id=GCP_PROJECT) mock_client.assert_called_once_with(region=GCP_LOCATION) mock_client.return_value.submit_job.assert_called_once_with( request=dict( @@ -792,9 +786,10 @@ async def test_cancel_job(self, mock_client): metadata=(), ) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_batch_client")) - def test_create_batch(self, mock_client): - self.hook.create_batch( + async def test_create_batch(self, mock_client): + await self.hook.create_batch( project_id=GCP_PROJECT, region=GCP_LOCATION, batch=BATCH, @@ -813,9 +808,10 @@ def test_create_batch(self, mock_client): timeout=None, ) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_batch_client")) - def test_delete_batch(self, mock_client): - self.hook.delete_batch( + async def test_delete_batch(self, mock_client): + await self.hook.delete_batch( batch_id=BATCH_ID, region=GCP_LOCATION, project_id=GCP_PROJECT, @@ -830,9 +826,10 @@ def test_delete_batch(self, mock_client): timeout=None, ) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_batch_client")) - def test_get_batch(self, mock_client): - self.hook.get_batch( + async def test_get_batch(self, mock_client): + await self.hook.get_batch( batch_id=BATCH_ID, region=GCP_LOCATION, project_id=GCP_PROJECT, @@ -847,9 +844,10 @@ def test_get_batch(self, mock_client): timeout=None, ) + @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_batch_client")) - def test_list_batches(self, mock_client): - self.hook.list_batches( + async def test_list_batches(self, mock_client): + await self.hook.list_batches( project_id=GCP_PROJECT, region=GCP_LOCATION, ) From 3531692cbd2b54b95c5d00a6e8c003194d1ee061 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 26 Jul 2022 09:32:47 +0200 Subject: [PATCH 15/29] Async tests --- tests/providers/google/cloud/hooks/test_dataproc.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/providers/google/cloud/hooks/test_dataproc.py b/tests/providers/google/cloud/hooks/test_dataproc.py index 828e979404040..293cc0aa0763c 100644 --- a/tests/providers/google/cloud/hooks/test_dataproc.py +++ b/tests/providers/google/cloud/hooks/test_dataproc.py @@ -588,8 +588,10 @@ async def test_delete_cluster(self, mock_client): @pytest.mark.asyncio @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_cluster_client")) - def test_diagnose_cluster(self, mock_client): - async self.hook.diagnose_cluster(project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME) + async def test_diagnose_cluster(self, mock_client): + await self.hook.diagnose_cluster( + project_id=GCP_PROJECT, region=GCP_LOCATION, cluster_name=CLUSTER_NAME + ) mock_client.assert_called_once_with(region=GCP_LOCATION) mock_client.return_value.diagnose_cluster.assert_called_once_with( request=dict( From 3d2f15961a59e27795555356d35471694f9e3d7a Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 26 Jul 2022 11:49:25 +0200 Subject: [PATCH 16/29] Fixes --- tests/providers/google/cloud/operators/test_dataproc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py index 8cb93f8d19acb..827b36f33a2d5 100644 --- a/tests/providers/google/cloud/operators/test_dataproc.py +++ b/tests/providers/google/cloud/operators/test_dataproc.py @@ -64,7 +64,7 @@ AIRFLOW_VERSION = "v" + airflow_version.replace(".", "-").replace("+", "-") DATAPROC_PATH = "airflow.providers.google.cloud.operators.dataproc.{}" -COMPOSER_TRIGGERS_STRING = "airflow.providers.google.cloud.triggers.dataproc.{}" +DATAPROC_TRIGGERS_STRING = "airflow.providers.google.cloud.triggers.dataproc.{}" TASK_ID = "task-id" GCP_PROJECT = "test-project" @@ -908,7 +908,7 @@ def test_execute_async(self, mock_hook): ) @mock.patch(DATAPROC_PATH.format("DataprocHook")) - @mock.patch(COMPOSER_TRIGGERS_STRING.format("DataprocHook")) + @mock.patch(DATAPROC_TRIGGERS_STRING.format("DataprocHook")) def test_execute_deferrable(self, mock_trigger_hook, mock_hook): job = {} mock_hook.return_value.submit_job.return_value.reference.job_id = TEST_JOB_ID From f640fe8e52ed7a11f118ce988c058596b1ca5776 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Fri, 29 Jul 2022 10:39:23 +0200 Subject: [PATCH 17/29] Add examples and documentation --- .../operators/cloud/dataproc.rst | 8 ++ .../example_dataproc_spark_deferrable.py | 111 ++++++++++++++++++ 2 files changed, 119 insertions(+) create mode 100644 tests/system/providers/google/cloud/dataproc/example_dataproc_spark_deferrable.py diff --git a/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst b/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst index 4bcb3dc089eec..eb3da2657878a 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst @@ -174,6 +174,14 @@ Example of the configuration for a Spark Job: :start-after: [START how_to_cloud_dataproc_spark_config] :end-before: [END how_to_cloud_dataproc_spark_config] +Example of the configuration for a Spark Job running in `deferrable mode `__: + +.. exampleinclude:: /../../tests/system/providers/google/dataproc/example_dataproc_spark_deferrable.py + :language: python + :dedent: 0 + :start-after: [START how_to_cloud_dataproc_spark_deferrable_config] + :end-before: [END how_to_cloud_dataproc_spark_deferrable_config] + Example of the configuration for a Hive Job: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_hive.py diff --git a/tests/system/providers/google/cloud/dataproc/example_dataproc_spark_deferrable.py b/tests/system/providers/google/cloud/dataproc/example_dataproc_spark_deferrable.py new file mode 100644 index 0000000000000..65c8621b056f7 --- /dev/null +++ b/tests/system/providers/google/cloud/dataproc/example_dataproc_spark_deferrable.py @@ -0,0 +1,111 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Example Airflow DAG for DataprocSubmitJobOperator with spark job +in deferrable mode. +""" + +import os +from datetime import datetime + +from airflow import models +from airflow.providers.google.cloud.operators.dataproc import ( + DataprocCreateClusterOperator, + DataprocDeleteClusterOperator, + DataprocSubmitJobOperator, +) +from airflow.utils.trigger_rule import TriggerRule + +ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") +DAG_ID = "dataproc_spark" +PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "") + +CLUSTER_NAME = f"cluster-dataproc-spark-{ENV_ID}" +REGION = "europe-west1" +ZONE = "europe-west1-b" + + +# Cluster definition +CLUSTER_CONFIG = { + "master_config": { + "num_instances": 1, + "machine_type_uri": "n1-standard-4", + "disk_config": {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 1024}, + }, + "worker_config": { + "num_instances": 2, + "machine_type_uri": "n1-standard-4", + "disk_config": {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 1024}, + }, +} + +TIMEOUT = {"seconds": 1 * 24 * 60 * 60} + +# Jobs definitions +# [START how_to_cloud_dataproc_spark_deferrable_config] +SPARK_JOB = { + "reference": {"project_id": PROJECT_ID}, + "placement": {"cluster_name": CLUSTER_NAME}, + "spark_job": { + "jar_file_uris": ["file:///usr/lib/spark/examples/jars/spark-examples.jar"], + "main_class": "org.apache.spark.examples.SparkPi", + }, +} +# [END how_to_cloud_dataproc_spark_deferrable_config] + + +with models.DAG( + DAG_ID, + schedule_interval='@once', + start_date=datetime(2021, 1, 1), + catchup=False, + tags=["example", "dataproc"], +) as dag: + create_cluster = DataprocCreateClusterOperator( + task_id="create_cluster", + project_id=PROJECT_ID, + cluster_config=CLUSTER_CONFIG, + region=REGION, + cluster_name=CLUSTER_NAME, + ) + + spark_task = DataprocSubmitJobOperator( + task_id="spark_task", job=SPARK_JOB, region=REGION, project_id=PROJECT_ID, deferrable=True + ) + + delete_cluster = DataprocDeleteClusterOperator( + task_id="delete_cluster", + project_id=PROJECT_ID, + cluster_name=CLUSTER_NAME, + region=REGION, + trigger_rule=TriggerRule.ALL_DONE, + ) + + create_cluster >> spark_task >> delete_cluster + + from tests.system.utils.watcher import watcher + + # This test needs watcher in order to properly mark success/failure + # when "teardown" task with trigger rule is part of the DAG + list(dag.tasks) >> watcher() + + +from tests.system.utils import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) From 1b2ce1184898cd78904fa84e586d9c49e318ef87 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Mon, 1 Aug 2022 15:27:01 +0200 Subject: [PATCH 18/29] Code review fixes --- .../google/cloud/operators/dataproc.py | 15 +++++++++++- .../google/cloud/triggers/dataproc.py | 16 +++++-------- .../operators/cloud/dataproc.rst | 24 +++++++++---------- 3 files changed, 32 insertions(+), 23 deletions(-) diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index ae1d7e7969e1c..8195ddfa45b4c 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -869,6 +869,8 @@ class DataprocJobBaseOperator(BaseOperator): This is useful for submitting long running jobs and waiting on them asynchronously using the DataprocJobSensor :param deferrable: Run operator in the deferrable mode + :param polling_interval_seconds time in seconds between polling for job completion. + The value is considered only when running in deferrable mode. Must be greater than 0. :var dataproc_job_id: The actual "jobId" as submitted to the Dataproc API. This is useful for identifying or linking to the job in the Google Cloud Console @@ -897,9 +899,12 @@ def __init__( impersonation_chain: Optional[Union[str, Sequence[str]]] = None, asynchronous: bool = False, deferrable: bool = False, + polling_interval_seconds: int = 10, **kwargs, ) -> None: super().__init__(**kwargs) + if deferrable and polling_interval_seconds <= 0: + raise ValueError("Invalid value for polling_interval_seconds. Expected value greater than 0") self.gcp_conn_id = gcp_conn_id self.delegate_to = delegate_to self.labels = labels @@ -918,6 +923,7 @@ def __init__( self.dataproc_job_id = None self.asynchronous = asynchronous self.deferrable = deferrable + self.polling_interval_seconds = polling_interval_seconds def create_job_template(self) -> DataProcJobBuilder: """Initialize `self.job_template` with default values""" @@ -971,7 +977,7 @@ def execute(self, context: 'Context'): delegate_to=self.delegate_to, gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, - pooling_period_seconds=10, + polling_interval_seconds=self.polling_interval_seconds, ), method_name="execute_complete", ) @@ -1803,6 +1809,8 @@ class DataprocSubmitJobOperator(BaseOperator): This is useful for submitting long running jobs and waiting on them asynchronously using the DataprocJobSensor :param deferrable: Run operator in the deferrable mode + :param polling_interval_seconds time in seconds between polling for job completion. + The value is considered only when running in deferrable mode. Must be greater than 0. :param cancel_on_kill: Flag which indicates whether cancel the hook's job or not, when on_kill is called :param wait_timeout: How many seconds wait for job to be ready. Used only if ``asynchronous`` is False """ @@ -1826,11 +1834,14 @@ def __init__( impersonation_chain: Optional[Union[str, Sequence[str]]] = None, asynchronous: bool = False, deferrable: bool = False, + polling_interval_seconds: int = 10, cancel_on_kill: bool = True, wait_timeout: Optional[int] = None, **kwargs, ) -> None: super().__init__(**kwargs) + if deferrable and polling_interval_seconds <= 0: + raise ValueError("Invalid value for polling_interval_seconds. Expected value greater than 0") self.project_id = project_id self.region = region self.job = job @@ -1842,6 +1853,7 @@ def __init__( self.impersonation_chain = impersonation_chain self.asynchronous = asynchronous self.deferrable = deferrable + self.polling_interval_seconds = polling_interval_seconds self.cancel_on_kill = cancel_on_kill self.hook: Optional[DataprocHook] = None self.job_id: Optional[str] = None @@ -1875,6 +1887,7 @@ def execute(self, context: 'Context'): region=self.region, gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, + polling_interval_seconds=self.polling_interval_seconds, ), method_name="execute_complete", ) diff --git a/airflow/providers/google/cloud/triggers/dataproc.py b/airflow/providers/google/cloud/triggers/dataproc.py index 5467ff91b65b4..3ccaf19af6127 100644 --- a/airflow/providers/google/cloud/triggers/dataproc.py +++ b/airflow/providers/google/cloud/triggers/dataproc.py @@ -9,7 +9,7 @@ class DataprocBaseTrigger(BaseTrigger): """ - Trigger that periodically pollls information from Dataproc API to verify job status. + Trigger that periodically polls information from Dataproc API to verify job status. Implementation leverages asynchronous transport. """ @@ -21,7 +21,7 @@ def __init__( gcp_conn_id: str = "google_cloud_default", impersonation_chain: Optional[Union[str, Sequence[str]]] = None, delegate_to: Optional[str] = None, - pooling_period_seconds: int = 30, + polling_interval_seconds: int = 30, ): super().__init__() self.gcp_conn_id = gcp_conn_id @@ -29,7 +29,7 @@ def __init__( self.job_id = job_id self.project_id = project_id self.region = region - self.pooling_period_seconds = pooling_period_seconds + self.polling_interval_seconds = polling_interval_seconds self.delegate_to = delegate_to self.hook = DataprocAsyncHook( delegate_to=self.delegate_to, @@ -39,7 +39,7 @@ def __init__( def serialize(self): return ( - "airflow.providers.google.cloud.operators.dataproc.DataprocBaseTrigger", + "airflow.providers.google.cloud.triggers.dataproc.DataprocBaseTrigger", { "job_id": self.job_id, "project_id": self.project_id, @@ -47,7 +47,7 @@ def serialize(self): "gcp_conn_id": self.gcp_conn_id, "delegate_to": self.delegate_to, "impersonation_chain": self.impersonation_chain, - "pooling_period_seconds": self.pooling_period_seconds, + "polling_interval_seconds": self.polling_interval_seconds, }, ) @@ -61,9 +61,5 @@ async def run(self): break elif state == JobStatus.State.ERROR: raise AirflowException(f"Dataproc job execution failed {self.job_id}") - else: - raise AirflowException( - f"Dataproc job execution finished in uknknown state {state} {self.job_id}" - ) - await asyncio.sleep(self.pooling_period_seconds) + await asyncio.sleep(self.polling_interval_seconds) yield TriggerEvent({"job_id": self.job_id, "job_state": state}) diff --git a/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst b/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst index eb3da2657878a..44c19396e0714 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst @@ -45,7 +45,7 @@ A cluster configuration can look as followed: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_hive.py :language: python - :dedent: 0 + :dedent: 4 :start-after: [START how_to_cloud_dataproc_create_cluster] :end-before: [END how_to_cloud_dataproc_create_cluster] @@ -62,7 +62,7 @@ For create Dataproc cluster in Google Kubernetes Engine you should use this clus .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_gke.py :language: python - :dedent: 0 + :dedent: 4 :start-after: [START how_to_cloud_dataproc_create_cluster_in_gke_config] :end-before: [END how_to_cloud_dataproc_create_cluster_in_gke_config] @@ -84,7 +84,7 @@ You can generate and use config as followed: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_cluster_generator.py :language: python - :dedent: 0 + :dedent: 4 :start-after: [START how_to_cloud_dataproc_create_cluster_generate_cluster_config] :end-before: [END how_to_cloud_dataproc_create_cluster_generate_cluster_config] @@ -98,7 +98,7 @@ An example of a new cluster config and the updateMask: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_update.py :language: python - :dedent: 0 + :dedent: 4 :start-after: [START how_to_cloud_dataproc_updatemask_cluster_operator] :end-before: [END how_to_cloud_dataproc_updatemask_cluster_operator] @@ -154,7 +154,7 @@ Example of the configuration for a PySpark Job: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_pyspark.py :language: python - :dedent: 0 + :dedent: 4 :start-after: [START how_to_cloud_dataproc_pyspark_config] :end-before: [END how_to_cloud_dataproc_pyspark_config] @@ -162,7 +162,7 @@ Example of the configuration for a SparkSQl Job: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_spark_sql.py :language: python - :dedent: 0 + :dedent: 4 :start-after: [START how_to_cloud_dataproc_sparksql_config] :end-before: [END how_to_cloud_dataproc_sparksql_config] @@ -170,7 +170,7 @@ Example of the configuration for a Spark Job: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_spark.py :language: python - :dedent: 0 + :dedent: 4 :start-after: [START how_to_cloud_dataproc_spark_config] :end-before: [END how_to_cloud_dataproc_spark_config] @@ -178,7 +178,7 @@ Example of the configuration for a Spark Job running in `deferrable mode Date: Tue, 2 Aug 2022 21:28:25 +0200 Subject: [PATCH 19/29] Fixes --- .../providers/google/cloud/hooks/dataproc.py | 11 +++++----- .../google/cloud/operators/dataproc.py | 2 +- .../google/cloud/triggers/dataproc.py | 22 ++++++++++++++++++- 3 files changed, 28 insertions(+), 7 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py index fde23877ee8f4..4f2c783dfe16d 100644 --- a/airflow/providers/google/cloud/hooks/dataproc.py +++ b/airflow/providers/google/cloud/hooks/dataproc.py @@ -26,21 +26,22 @@ from google.api_core.exceptions import ServerError from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault from google.api_core.operation import Operation +from google.api_core.operation_async import AsyncOperation from google.api_core.retry import Retry from google.cloud.dataproc_v1 import ( Batch, - BatchControllerClient, BatchControllerAsyncClient, + BatchControllerClient, Cluster, - ClusterControllerClient, ClusterControllerAsyncClient, + ClusterControllerClient, Job, - JobControllerClient, JobControllerAsyncClient, + JobControllerClient, JobStatus, WorkflowTemplate, - WorkflowTemplateServiceClient, WorkflowTemplateServiceAsyncClient, + WorkflowTemplateServiceClient, ) from google.protobuf.duration_pb2 import Duration from google.protobuf.field_mask_pb2 import FieldMask @@ -1557,7 +1558,7 @@ async def create_batch( retry: Union[Retry, _MethodDefault] = DEFAULT, timeout: Optional[float] = None, metadata: Sequence[Tuple[str, str]] = (), - ) -> Operation: + ) -> AsyncOperation: """ Creates a batch workload. diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index 8195ddfa45b4c..95727ab4271d7 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -971,7 +971,7 @@ def execute(self, context: 'Context'): if self.deferrable: self.defer( trigger=DataprocBaseTrigger( - job_id=self.job_id, + job_id=job_id, project_id=self.project_id, region=self.region, delegate_to=self.delegate_to, diff --git a/airflow/providers/google/cloud/triggers/dataproc.py b/airflow/providers/google/cloud/triggers/dataproc.py index 3ccaf19af6127..f658d61581814 100644 --- a/airflow/providers/google/cloud/triggers/dataproc.py +++ b/airflow/providers/google/cloud/triggers/dataproc.py @@ -1,3 +1,23 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +"""This module contains Google Dataproc triggers.""" + import asyncio from airflow import AirflowException @@ -16,8 +36,8 @@ class DataprocBaseTrigger(BaseTrigger): def __init__( self, job_id: str, - project_id: str, region: str, + project_id: Optional[str] = None, gcp_conn_id: str = "google_cloud_default", impersonation_chain: Optional[Union[str, Sequence[str]]] = None, delegate_to: Optional[str] = None, From e754dcbd07730381f1c1c777142b38f0bbd2e833 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Tue, 2 Aug 2022 21:35:50 +0200 Subject: [PATCH 20/29] Fixes --- tests/providers/google/cloud/operators/test_dataproc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py index 827b36f33a2d5..88c50cc8dd9b7 100644 --- a/tests/providers/google/cloud/operators/test_dataproc.py +++ b/tests/providers/google/cloud/operators/test_dataproc.py @@ -907,8 +907,8 @@ def test_execute_async(self, mock_hook): key="conf", value=DATAPROC_JOB_CONF_EXPECTED, execution_date=None ) - @mock.patch(DATAPROC_PATH.format("DataprocHook")) - @mock.patch(DATAPROC_TRIGGERS_STRING.format("DataprocHook")) + @mock.patch(DATAPROC_PATH.format("DataprocAsyncHook")) + @mock.patch(DATAPROC_TRIGGERS_STRING.format("DataprocBaseTrigger")) def test_execute_deferrable(self, mock_trigger_hook, mock_hook): job = {} mock_hook.return_value.submit_job.return_value.reference.job_id = TEST_JOB_ID From 619f625f64ad0328058ead6ca39b93810a905434 Mon Sep 17 00:00:00 2001 From: Bartosz Jankiewicz Date: Thu, 4 Aug 2022 19:30:22 +0200 Subject: [PATCH 21/29] Update airflow/providers/google/cloud/operators/dataproc.py Co-authored-by: Pankaj Singh <98807258+pankajastro@users.noreply.github.com> --- airflow/providers/google/cloud/operators/dataproc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index 95727ab4271d7..b18993adee901 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -869,7 +869,7 @@ class DataprocJobBaseOperator(BaseOperator): This is useful for submitting long running jobs and waiting on them asynchronously using the DataprocJobSensor :param deferrable: Run operator in the deferrable mode - :param polling_interval_seconds time in seconds between polling for job completion. + :param polling_interval_seconds: time in seconds between polling for job completion. The value is considered only when running in deferrable mode. Must be greater than 0. :var dataproc_job_id: The actual "jobId" as submitted to the Dataproc API. From 572f7c60c71ff7646bd17347a8ee64a77e10e339 Mon Sep 17 00:00:00 2001 From: Bartosz Jankiewicz Date: Fri, 5 Aug 2022 10:50:02 +0200 Subject: [PATCH 22/29] Update airflow/providers/google/cloud/operators/dataproc.py Co-authored-by: Pankaj Singh <98807258+pankajastro@users.noreply.github.com> --- airflow/providers/google/cloud/operators/dataproc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index b18993adee901..ad37b2ad9a841 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -1809,7 +1809,7 @@ class DataprocSubmitJobOperator(BaseOperator): This is useful for submitting long running jobs and waiting on them asynchronously using the DataprocJobSensor :param deferrable: Run operator in the deferrable mode - :param polling_interval_seconds time in seconds between polling for job completion. + :param polling_interval_seconds: time in seconds between polling for job completion. The value is considered only when running in deferrable mode. Must be greater than 0. :param cancel_on_kill: Flag which indicates whether cancel the hook's job or not, when on_kill is called :param wait_timeout: How many seconds wait for job to be ready. Used only if ``asynchronous`` is False From ff772cfa64b3b2161ed83b76b2fe20e44181a465 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Sun, 7 Aug 2022 22:35:34 +0200 Subject: [PATCH 23/29] Fixing tests --- airflow/providers/google/cloud/operators/dataproc.py | 2 +- airflow/providers/google/cloud/triggers/dataproc.py | 5 +++-- tests/providers/google/cloud/operators/test_dataproc.py | 2 +- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index ad37b2ad9a841..07ce34dc4f60a 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -32,7 +32,7 @@ from google.api_core.exceptions import AlreadyExists, NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault from google.api_core.retry import Retry, exponential_sleep_generator -from google.cloud.dataproc_v1 import Batch, Cluster, JobStatus, Job +from google.cloud.dataproc_v1 import Batch, Cluster, JobStatus from google.protobuf.duration_pb2 import Duration from google.protobuf.field_mask_pb2 import FieldMask diff --git a/airflow/providers/google/cloud/triggers/dataproc.py b/airflow/providers/google/cloud/triggers/dataproc.py index f658d61581814..cdeb1feb8f49a 100644 --- a/airflow/providers/google/cloud/triggers/dataproc.py +++ b/airflow/providers/google/cloud/triggers/dataproc.py @@ -19,12 +19,13 @@ """This module contains Google Dataproc triggers.""" import asyncio +from typing import Optional, Sequence, Union + +from google.cloud.dataproc_v1 import JobStatus from airflow import AirflowException from airflow.providers.google.cloud.hooks.dataproc import DataprocAsyncHook from airflow.triggers.base import BaseTrigger, TriggerEvent -from google.cloud.dataproc_v1 import JobStatus, Job -from typing import Any, Dict, Optional, Sequence, Tuple, Union class DataprocBaseTrigger(BaseTrigger): diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py index 88c50cc8dd9b7..215cad4364376 100644 --- a/tests/providers/google/cloud/operators/test_dataproc.py +++ b/tests/providers/google/cloud/operators/test_dataproc.py @@ -907,7 +907,7 @@ def test_execute_async(self, mock_hook): key="conf", value=DATAPROC_JOB_CONF_EXPECTED, execution_date=None ) - @mock.patch(DATAPROC_PATH.format("DataprocAsyncHook")) + @mock.patch(DATAPROC_PATH.format("DataprocHook")) @mock.patch(DATAPROC_TRIGGERS_STRING.format("DataprocBaseTrigger")) def test_execute_deferrable(self, mock_trigger_hook, mock_hook): job = {} From 848f3218b844ec96cca39de483067ed3f6427de8 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Sun, 7 Aug 2022 22:42:13 +0200 Subject: [PATCH 24/29] Test fixes --- tests/providers/google/cloud/operators/test_dataproc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py index 215cad4364376..5c182d4e40b76 100644 --- a/tests/providers/google/cloud/operators/test_dataproc.py +++ b/tests/providers/google/cloud/operators/test_dataproc.py @@ -908,7 +908,7 @@ def test_execute_async(self, mock_hook): ) @mock.patch(DATAPROC_PATH.format("DataprocHook")) - @mock.patch(DATAPROC_TRIGGERS_STRING.format("DataprocBaseTrigger")) + @mock.patch(DATAPROC_PATH.format("DataprocBaseTrigger")) def test_execute_deferrable(self, mock_trigger_hook, mock_hook): job = {} mock_hook.return_value.submit_job.return_value.reference.job_id = TEST_JOB_ID From edbe50131598c8666494ba998fba5190731cd4d7 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Sun, 7 Aug 2022 22:48:28 +0200 Subject: [PATCH 25/29] Text fixes --- tests/providers/google/cloud/operators/test_dataproc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py index 5c182d4e40b76..ca3173b9e2a81 100644 --- a/tests/providers/google/cloud/operators/test_dataproc.py +++ b/tests/providers/google/cloud/operators/test_dataproc.py @@ -64,7 +64,7 @@ AIRFLOW_VERSION = "v" + airflow_version.replace(".", "-").replace("+", "-") DATAPROC_PATH = "airflow.providers.google.cloud.operators.dataproc.{}" -DATAPROC_TRIGGERS_STRING = "airflow.providers.google.cloud.triggers.dataproc.{}" +DATAPROC_TRIGGERS_PATH = "airflow.providers.google.cloud.triggers.dataproc.{}" TASK_ID = "task-id" GCP_PROJECT = "test-project" @@ -908,7 +908,7 @@ def test_execute_async(self, mock_hook): ) @mock.patch(DATAPROC_PATH.format("DataprocHook")) - @mock.patch(DATAPROC_PATH.format("DataprocBaseTrigger")) + @mock.patch(DATAPROC_TRIGGERS_PATH.format("DataprocAsyncHook")) def test_execute_deferrable(self, mock_trigger_hook, mock_hook): job = {} mock_hook.return_value.submit_job.return_value.reference.job_id = TEST_JOB_ID From f385f57f25d3d048ec30d7d297116549ca749a94 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Mon, 8 Aug 2022 10:28:34 +0200 Subject: [PATCH 26/29] Documentation fix --- .../operators/cloud/dataproc.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst b/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst index 44c19396e0714..65be156574557 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst @@ -176,7 +176,7 @@ Example of the configuration for a Spark Job: Example of the configuration for a Spark Job running in `deferrable mode `__: -.. exampleinclude:: /../../tests/system/providers/google/dataproc/example_dataproc_spark_deferrable.py +.. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_spark_deferrable.py :language: python :dedent: 4 :start-after: [START how_to_cloud_dataproc_spark_deferrable_config] From 9679ca1be1ad0d20561aef585519ab0419e6bd32 Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Fri, 12 Aug 2022 12:59:43 +0200 Subject: [PATCH 27/29] Documentation fixes --- .../providers/google/cloud/hooks/dataproc.py | 2 +- .../operators/cloud/dataproc.rst | 24 +++++++++---------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py index 4f2c783dfe16d..52cb23119ff75 100644 --- a/airflow/providers/google/cloud/hooks/dataproc.py +++ b/airflow/providers/google/cloud/hooks/dataproc.py @@ -977,7 +977,7 @@ def list_batches( class DataprocAsyncHook(GoogleBaseHook): """ - Asynchronuous Hook for Google Cloud Dataproc APIs. + Asynchronous Hook for Google Cloud Dataproc APIs. All the methods in the hook where project_id is used must be called with keyword arguments rather than positional. diff --git a/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst b/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst index 65be156574557..924ee7d24d265 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst @@ -45,7 +45,7 @@ A cluster configuration can look as followed: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_hive.py :language: python - :dedent: 4 + :dedent: 0 :start-after: [START how_to_cloud_dataproc_create_cluster] :end-before: [END how_to_cloud_dataproc_create_cluster] @@ -62,7 +62,7 @@ For create Dataproc cluster in Google Kubernetes Engine you should use this clus .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_gke.py :language: python - :dedent: 4 + :dedent: 0 :start-after: [START how_to_cloud_dataproc_create_cluster_in_gke_config] :end-before: [END how_to_cloud_dataproc_create_cluster_in_gke_config] @@ -84,7 +84,7 @@ You can generate and use config as followed: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_cluster_generator.py :language: python - :dedent: 4 + :dedent: 0 :start-after: [START how_to_cloud_dataproc_create_cluster_generate_cluster_config] :end-before: [END how_to_cloud_dataproc_create_cluster_generate_cluster_config] @@ -98,7 +98,7 @@ An example of a new cluster config and the updateMask: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_update.py :language: python - :dedent: 4 + :dedent: 0 :start-after: [START how_to_cloud_dataproc_updatemask_cluster_operator] :end-before: [END how_to_cloud_dataproc_updatemask_cluster_operator] @@ -154,7 +154,7 @@ Example of the configuration for a PySpark Job: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_pyspark.py :language: python - :dedent: 4 + :dedent: 0 :start-after: [START how_to_cloud_dataproc_pyspark_config] :end-before: [END how_to_cloud_dataproc_pyspark_config] @@ -162,7 +162,7 @@ Example of the configuration for a SparkSQl Job: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_spark_sql.py :language: python - :dedent: 4 + :dedent: 0 :start-after: [START how_to_cloud_dataproc_sparksql_config] :end-before: [END how_to_cloud_dataproc_sparksql_config] @@ -170,7 +170,7 @@ Example of the configuration for a Spark Job: .. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc/example_dataproc_spark.py :language: python - :dedent: 4 + :dedent: 0 :start-after: [START how_to_cloud_dataproc_spark_config] :end-before: [END how_to_cloud_dataproc_spark_config] @@ -178,7 +178,7 @@ Example of the configuration for a Spark Job running in `deferrable mode Date: Mon, 22 Aug 2022 16:51:11 +0200 Subject: [PATCH 28/29] Sync with master fix --- airflow/providers/google/cloud/hooks/dataproc.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py index 2cecb7ef716d7..8e916ff0e5110 100644 --- a/airflow/providers/google/cloud/hooks/dataproc.py +++ b/airflow/providers/google/cloud/hooks/dataproc.py @@ -996,7 +996,7 @@ def get_cluster_client(self, region: Optional[str] = None) -> ClusterControllerA client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443') return ClusterControllerAsyncClient( - credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options + credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options ) def get_template_client(self, region: Optional[str] = None) -> WorkflowTemplateServiceAsyncClient: @@ -1006,7 +1006,7 @@ def get_template_client(self, region: Optional[str] = None) -> WorkflowTemplateS client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443') return WorkflowTemplateServiceAsyncClient( - credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options + credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options ) def get_job_client(self, region: Optional[str] = None) -> JobControllerAsyncClient: @@ -1016,7 +1016,7 @@ def get_job_client(self, region: Optional[str] = None) -> JobControllerAsyncClie client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443') return JobControllerAsyncClient( - credentials=self._get_credentials(), + credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options, ) @@ -1028,7 +1028,7 @@ def get_batch_client(self, region: Optional[str] = None) -> BatchControllerAsync client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443') return BatchControllerAsyncClient( - credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options + credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options ) @GoogleBaseHook.fallback_to_default_project_id From 00afe22b85a7d48693e2c1eb8ce8dafc15b07c2c Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Mon, 22 Aug 2022 16:55:10 +0200 Subject: [PATCH 29/29] Fix test after merge --- .../google/cloud/hooks/test_dataproc.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/providers/google/cloud/hooks/test_dataproc.py b/tests/providers/google/cloud/hooks/test_dataproc.py index bc0a8007e9780..536bb017ba87c 100644 --- a/tests/providers/google/cloud/hooks/test_dataproc.py +++ b/tests/providers/google/cloud/hooks/test_dataproc.py @@ -467,7 +467,7 @@ def setUp(self): with mock.patch(BASE_STRING.format("GoogleBaseHook.__init__"), new=mock_init): self.hook = DataprocAsyncHook(gcp_conn_id="test") - @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_credentials")) @mock.patch(DATAPROC_STRING.format("ClusterControllerAsyncClient")) def test_get_cluster_client(self, mock_client, mock_get_credentials): self.hook.get_cluster_client(region=GCP_LOCATION) @@ -477,7 +477,7 @@ def test_get_cluster_client(self, mock_client, mock_get_credentials): client_options=None, ) - @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_credentials")) @mock.patch(DATAPROC_STRING.format("ClusterControllerAsyncClient")) def test_get_cluster_client_region(self, mock_client, mock_get_credentials): self.hook.get_cluster_client(region='region1') @@ -487,7 +487,7 @@ def test_get_cluster_client_region(self, mock_client, mock_get_credentials): client_options=ANY, ) - @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_credentials")) @mock.patch(DATAPROC_STRING.format("WorkflowTemplateServiceAsyncClient")) def test_get_template_client_global(self, mock_client, mock_get_credentials): _ = self.hook.get_template_client() @@ -497,7 +497,7 @@ def test_get_template_client_global(self, mock_client, mock_get_credentials): client_options=None, ) - @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_credentials")) @mock.patch(DATAPROC_STRING.format("WorkflowTemplateServiceAsyncClient")) def test_get_template_client_region(self, mock_client, mock_get_credentials): _ = self.hook.get_template_client(region='region1') @@ -507,7 +507,7 @@ def test_get_template_client_region(self, mock_client, mock_get_credentials): client_options=ANY, ) - @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_credentials")) @mock.patch(DATAPROC_STRING.format("JobControllerAsyncClient")) def test_get_job_client(self, mock_client, mock_get_credentials): self.hook.get_job_client(region=GCP_LOCATION) @@ -517,7 +517,7 @@ def test_get_job_client(self, mock_client, mock_get_credentials): client_options=None, ) - @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_credentials")) @mock.patch(DATAPROC_STRING.format("JobControllerAsyncClient")) def test_get_job_client_region(self, mock_client, mock_get_credentials): self.hook.get_job_client(region='region1') @@ -527,7 +527,7 @@ def test_get_job_client_region(self, mock_client, mock_get_credentials): client_options=ANY, ) - @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_credentials")) @mock.patch(DATAPROC_STRING.format("BatchControllerAsyncClient")) def test_get_batch_client(self, mock_client, mock_get_credentials): self.hook.get_batch_client(region=GCP_LOCATION) @@ -537,7 +537,7 @@ def test_get_batch_client(self, mock_client, mock_get_credentials): client_options=None, ) - @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook._get_credentials")) + @mock.patch(DATAPROC_STRING.format("DataprocAsyncHook.get_credentials")) @mock.patch(DATAPROC_STRING.format("BatchControllerAsyncClient")) def test_get_batch_client_region(self, mock_client, mock_get_credentials): self.hook.get_batch_client(region='region1')