apache-airflow-providers-google 10.13.0rc1__py3-none-any.whl → 10.13.0rc2__py3-none-any.whl

airflow/providers/google/cloud/operators/dataplex.py

@@ -22,7 +22,10 @@ import time
 from typing import TYPE_CHECKING, Any, Sequence
 
 from airflow.exceptions import AirflowException
-from airflow.providers.google.cloud.triggers.dataplex import DataplexDataQualityJobTrigger
+from airflow.providers.google.cloud.triggers.dataplex import (
+    DataplexDataProfileJobTrigger,
+    DataplexDataQualityJobTrigger,
+)
 
 if TYPE_CHECKING:
     from google.protobuf.field_mask_pb2 import FieldMask
@@ -1204,6 +1207,532 @@ class DataplexGetDataQualityScanResultOperator(GoogleCloudBaseOperator):
         return job
 
 
+class DataplexCreateOrUpdateDataProfileScanOperator(GoogleCloudBaseOperator):
+    """
+    Creates a DataScan Data Profile resource.
+
+    :param project_id: Required. The ID of the Google Cloud project that the lake belongs to.
+    :param region: Required. The ID of the Google Cloud region that the lake belongs to.
+    :param body:  Required. The Request body contains an instance of DataScan.
+    :param data_scan_id: Required. Data Profile scan identifier.
+    :param update_mask: Mask of fields to update.
+    :param api_version: The version of the api that will be requested for example 'v1'.
+    :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.
+    :param gcp_conn_id: The connection ID to use when fetching connection info.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+
+    :return: Dataplex data profile id
+    """
+
+    template_fields = ("project_id", "data_scan_id", "body", "impersonation_chain")
+    template_fields_renderers = {"body": "json"}
+
+    def __init__(
+        self,
+        project_id: str,
+        region: str,
+        data_scan_id: str,
+        body: dict[str, Any] | DataScan,
+        api_version: str = "v1",
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        update_mask: dict | FieldMask | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.project_id = project_id
+        self.region = region
+        self.data_scan_id = data_scan_id
+        self.body = body
+        self.update_mask = update_mask
+        self.api_version = api_version
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        hook = DataplexHook(
+            gcp_conn_id=self.gcp_conn_id,
+            api_version=self.api_version,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        self.log.info("Creating Dataplex Data Profile scan %s", self.data_scan_id)
+        try:
+            operation = hook.create_data_scan(
+                project_id=self.project_id,
+                region=self.region,
+                data_scan_id=self.data_scan_id,
+                body=self.body,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+            hook.wait_for_operation(timeout=self.timeout, operation=operation)
+            self.log.info("Dataplex Data Profile scan %s created successfully!", self.data_scan_id)
+        except AlreadyExists:
+            self.log.info("Dataplex Data Profile scan already exists: %s", {self.data_scan_id})
+
+            operation = hook.update_data_scan(
+                project_id=self.project_id,
+                region=self.region,
+                data_scan_id=self.data_scan_id,
+                body=self.body,
+                update_mask=self.update_mask,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+            hook.wait_for_operation(timeout=self.timeout, operation=operation)
+            self.log.info("Dataplex Data Profile scan %s updated successfully!", self.data_scan_id)
+        except GoogleAPICallError as e:
+            raise AirflowException(f"Error creating Data Profile scan {self.data_scan_id}", e)
+
+        return self.data_scan_id
+
+
+class DataplexGetDataProfileScanOperator(GoogleCloudBaseOperator):
+    """
+    Gets a DataScan DataProfile resource.
+
+    :param project_id: Required. The ID of the Google Cloud project that the lake belongs to.
+    :param region: Required. The ID of the Google Cloud region that the lake belongs to.
+    :param data_scan_id: Required. Data Profile scan identifier.
+    :param api_version: The version of the api that will be requested for example 'v1'.
+    :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.
+    :param gcp_conn_id: The connection ID to use when fetching connection info.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+
+    :return: Dataplex data profile
+    """
+
+    template_fields = ("project_id", "data_scan_id", "impersonation_chain")
+
+    def __init__(
+        self,
+        project_id: str,
+        region: str,
+        data_scan_id: str,
+        api_version: str = "v1",
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.project_id = project_id
+        self.region = region
+        self.data_scan_id = data_scan_id
+        self.api_version = api_version
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        hook = DataplexHook(
+            gcp_conn_id=self.gcp_conn_id,
+            api_version=self.api_version,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        self.log.info("Retrieving the details of Dataplex Data Profile scan %s", self.data_scan_id)
+        data_profile_scan = hook.get_data_scan(
+            project_id=self.project_id,
+            region=self.region,
+            data_scan_id=self.data_scan_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+
+        return DataScan.to_dict(data_profile_scan)
+
+
+class DataplexDeleteDataProfileScanOperator(GoogleCloudBaseOperator):
+    """
+    Deletes a DataScan DataProfile resource.
+
+    :param project_id: Required. The ID of the Google Cloud project that the lake belongs to.
+    :param region: Required. The ID of the Google Cloud region that the lake belongs to.
+    :param data_scan_id: Required. Data Profile scan identifier.
+    :param api_version: The version of the api that will be requested for example 'v1'.
+    :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.
+    :param gcp_conn_id: The connection ID to use when fetching connection info.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    :return: None
+    """
+
+    template_fields = ("project_id", "data_scan_id", "impersonation_chain")
+
+    def __init__(
+        self,
+        project_id: str,
+        region: str,
+        data_scan_id: str,
+        api_version: str = "v1",
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.project_id = project_id
+        self.region = region
+        self.data_scan_id = data_scan_id
+        self.api_version = api_version
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> None:
+        hook = DataplexHook(
+            gcp_conn_id=self.gcp_conn_id,
+            api_version=self.api_version,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        self.log.info("Deleting Dataplex Data Profile Scan: %s", self.data_scan_id)
+
+        operation = hook.delete_data_scan(
+            project_id=self.project_id,
+            region=self.region,
+            data_scan_id=self.data_scan_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        hook.wait_for_operation(timeout=self.timeout, operation=operation)
+        self.log.info("Dataplex Data Profile scan %s deleted successfully!", self.data_scan_id)
+
+
+class DataplexRunDataProfileScanOperator(GoogleCloudBaseOperator):
+    """
+    Runs an on-demand execution of a DataScan Data Profile Scan.
+
+    :param project_id: Required. The ID of the Google Cloud project that the lake belongs to.
+    :param region: Required. The ID of the Google Cloud region that the lake belongs to.
+    :param data_scan_id: Required. Data Profile scan identifier.
+    :param api_version: The version of the api that will be requested for example 'v1'.
+    :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.
+    :param gcp_conn_id: The connection ID to use when fetching connection info.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    :param asynchronous: Flag informing that the Dataplex job should be run asynchronously.
+        This is useful for submitting long-running jobs and
+        waiting on them asynchronously using the DataplexDataProfileJobStatusSensor
+    :param result_timeout: Value in seconds for which operator will wait for the Data Profile scan result
+        when the flag `asynchronous = False`.
+        Throws exception if there is no result found after specified amount of seconds.
+    :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 deferrable: Run operator in the deferrable mode.
+
+    :return: Dataplex Data Profile scan job id.
+    """
+
+    template_fields = ("project_id", "data_scan_id", "impersonation_chain")
+
+    def __init__(
+        self,
+        project_id: str,
+        region: str,
+        data_scan_id: str,
+        api_version: str = "v1",
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        asynchronous: bool = False,
+        result_timeout: float = 60.0 * 10,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        polling_interval_seconds: int = 10,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.project_id = project_id
+        self.region = region
+        self.data_scan_id = data_scan_id
+        self.api_version = api_version
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+        self.asynchronous = asynchronous
+        self.result_timeout = result_timeout
+        self.deferrable = deferrable
+        self.polling_interval_seconds = polling_interval_seconds
+
+    def execute(self, context: Context) -> dict:
+        hook = DataplexHook(
+            gcp_conn_id=self.gcp_conn_id,
+            api_version=self.api_version,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        result = hook.run_data_scan(
+            project_id=self.project_id,
+            region=self.region,
+            data_scan_id=self.data_scan_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        job_id = result.job.name.split("/")[-1]
+
+        if self.deferrable:
+            if self.asynchronous:
+                raise AirflowException(
+                    "Both asynchronous and deferrable parameters were passed. Please, provide only one."
+                )
+            self.defer(
+                trigger=DataplexDataProfileJobTrigger(
+                    job_id=job_id,
+                    data_scan_id=self.data_scan_id,
+                    project_id=self.project_id,
+                    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",
+            )
+        if not self.asynchronous:
+            job = hook.wait_for_data_scan_job(
+                job_id=job_id,
+                data_scan_id=self.data_scan_id,
+                project_id=self.project_id,
+                region=self.region,
+                result_timeout=self.result_timeout,
+            )
+
+            if job.state == DataScanJob.State.FAILED:
+                raise AirflowException(f"Data Profile job failed: {job_id}")
+            if job.state == DataScanJob.State.SUCCEEDED:
+                self.log.info("Data Profile job executed successfully.")
+            else:
+                self.log.info("Data Profile job execution returned status: %s", job.status)
+
+        return job_id
+
+    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_state = event["job_state"]
+        job_id = event["job_id"]
+        if job_state == DataScanJob.State.FAILED:
+            raise AirflowException(f"Job failed:\n{job_id}")
+        if job_state == DataScanJob.State.CANCELLED:
+            raise AirflowException(f"Job was cancelled:\n{job_id}")
+        if job_state == DataScanJob.State.SUCCEEDED:
+            self.log.info("Data Profile job executed successfully.")
+        return job_id
+
+
+class DataplexGetDataProfileScanResultOperator(GoogleCloudBaseOperator):
+    """
+    Gets a DataScan Data Profile Job resource.
+
+    :param project_id: Required. The ID of the Google Cloud project that the lake belongs to.
+    :param region: Required. The ID of the Google Cloud region that the lake belongs to.
+    :param data_scan_id: Required. Data Profile scan identifier.
+    :param job_id: Optional. Data Profile scan job identifier.
+    :param api_version: The version of the api that will be requested for example 'v1'.
+    :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.
+    :param gcp_conn_id: The connection ID to use when fetching connection info.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    :param wait_for_results: Flag indicating whether to wait for the result of a job execution
+        or to return the job in its current state.
+    :param result_timeout: Value in seconds for which operator will wait for the Data Profile scan result
+        when the flag `wait_for_results = True`.
+        Throws exception if there is no result found after specified amount of seconds.
+
+    :return: Dict representing DataScanJob.
+        When the job completes with a successful status, information about the Data Profile result
+        is available.
+    """
+
+    template_fields = ("project_id", "data_scan_id", "impersonation_chain")
+
+    def __init__(
+        self,
+        project_id: str,
+        region: str,
+        data_scan_id: str,
+        job_id: str | None = None,
+        api_version: str = "v1",
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        wait_for_results: bool = True,
+        result_timeout: float = 60.0 * 10,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.project_id = project_id
+        self.region = region
+        self.data_scan_id = data_scan_id
+        self.job_id = job_id
+        self.api_version = api_version
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+        self.wait_for_results = wait_for_results
+        self.result_timeout = result_timeout
+
+    def execute(self, context: Context) -> dict:
+        hook = DataplexHook(
+            gcp_conn_id=self.gcp_conn_id,
+            api_version=self.api_version,
+            impersonation_chain=self.impersonation_chain,
+        )
+        # fetch the last job
+        if not self.job_id:
+            jobs = hook.list_data_scan_jobs(
+                project_id=self.project_id,
+                region=self.region,
+                data_scan_id=self.data_scan_id,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+            job_ids = [DataScanJob.to_dict(job) for job in jobs]
+            if not job_ids:
+                raise AirflowException("There are no jobs, you should create one before.")
+            job_id = job_ids[0]["name"]
+            self.job_id = job_id.split("/")[-1]
+
+        if self.wait_for_results:
+            job = hook.wait_for_data_scan_job(
+                job_id=self.job_id,
+                data_scan_id=self.data_scan_id,
+                project_id=self.project_id,
+                region=self.region,
+                result_timeout=self.result_timeout,
+            )
+        else:
+            job = hook.get_data_scan_job(
+                project_id=self.project_id,
+                region=self.region,
+                job_id=self.job_id,
+                data_scan_id=self.data_scan_id,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+        if job.state == DataScanJob.State.SUCCEEDED:
+            self.log.info("Data Profile job executed successfully")
+        else:
+            self.log.info("Data Profile job execution returned status: %s", job.state)
+
+        result = DataScanJob.to_dict(job)
+        result["state"] = DataScanJob.State(result["state"]).name
+
+        return result
+
+    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_state = event["job_state"]
+        job_id = event["job_id"]
+        job = event["job"]
+        if job_state == DataScanJob.State.FAILED:
+            raise AirflowException(f"Job failed:\n{job_id}")
+        if job_state == DataScanJob.State.CANCELLED:
+            raise AirflowException(f"Job was cancelled:\n{job_id}")
+        if job_state == DataScanJob.State.SUCCEEDED:
+            self.log.info("Data Profile job executed successfully")
+        else:
+            self.log.info("Data Profile job execution returned status: %s", job_state)
+
+        return job
+
+
 class DataplexCreateZoneOperator(GoogleCloudBaseOperator):
     """
     Creates a Zone resource within a Lake.

airflow/providers/google/cloud/operators/dataproc.py

@@ -63,6 +63,7 @@ from airflow.utils import timezone
 
 if TYPE_CHECKING:
     from google.api_core import operation
+    from google.api_core.retry_async import AsyncRetry
     from google.protobuf.duration_pb2 import Duration
     from google.protobuf.field_mask_pb2 import FieldMask
 
@@ -592,7 +593,7 @@ class DataprocCreateClusterOperator(GoogleCloudBaseOperator):
         request_id: str | None = None,
         delete_on_error: bool = True,
         use_if_exists: bool = True,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float = 1 * 60 * 60,
         metadata: Sequence[tuple[str, str]] = (),
         gcp_conn_id: str = "google_cloud_default",
@@ -985,7 +986,7 @@ class DataprocDeleteClusterOperator(GoogleCloudBaseOperator):
         project_id: str | None = None,
         cluster_uuid: str | None = None,
         request_id: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float = 1 * 60 * 60,
         metadata: Sequence[tuple[str, str]] = (),
         gcp_conn_id: str = "google_cloud_default",
@@ -1891,7 +1892,7 @@ class DataprocInstantiateWorkflowTemplateOperator(GoogleCloudBaseOperator):
         version: int | None = None,
         request_id: str | None = None,
         parameters: dict[str, str] | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
         gcp_conn_id: str = "google_cloud_default",
@@ -2341,7 +2342,7 @@ class DataprocUpdateClusterOperator(GoogleCloudBaseOperator):
         region: str,
         request_id: str | None = None,
         project_id: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
         gcp_conn_id: str = "google_cloud_default",
@@ -2481,7 +2482,7 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
         metadata: Sequence[tuple[str, str]] = (),
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
-        result_retry: Retry | _MethodDefault = DEFAULT,
+        result_retry: AsyncRetry | _MethodDefault = DEFAULT,
         asynchronous: bool = False,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         polling_interval_seconds: int = 5,