apache-airflow-providers-google 10.26.0rc1__py3-none-any.whl → 11.0.0__py3-none-any.whl

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

@@ -20,18 +20,15 @@
 from __future__ import annotations
 
 import inspect
-import ntpath
-import os
 import re
 import time
-import uuid
 import warnings
-from collections.abc import MutableSequence
+from collections.abc import MutableSequence, Sequence
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from enum import Enum
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.exceptions import AlreadyExists, NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -45,7 +42,6 @@ from airflow.providers.google.cloud.hooks.dataproc import (
     DataProcJobBuilder,
     DataprocResourceIsNotReadyError,
 )
-from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.cloud.links.dataproc import (
     DATAPROC_BATCH_LINK,
     DATAPROC_CLUSTER_LINK_DEPRECATED,
@@ -822,7 +818,7 @@ class DataprocCreateClusterOperator(GoogleCloudBaseOperator):
         try:
             # First try to create a new cluster
             operation = self._create_cluster(hook)
-            if not self.deferrable:
+            if not self.deferrable and type(operation) is not str:
                 cluster = hook.wait_for_operation(
                     timeout=self.timeout, result_retry=self.retry, operation=operation
                 )
@@ -1547,632 +1543,6 @@ class DataprocJobBaseOperator(GoogleCloudBaseOperator):
             self.hook.cancel_job(project_id=self.project_id, job_id=self.dataproc_job_id, region=self.region)
 
 
-@deprecated(
-    planned_removal_date="November 01, 2024",
-    use_instead="DataprocSubmitJobOperator",
-    instructions="You can use `generate_job` method to generate dictionary representing your job "
-    "and use it with the new operator.",
-    category=AirflowProviderDeprecationWarning,
-)
-class DataprocSubmitPigJobOperator(DataprocJobBaseOperator):
-    """
-    Start a Pig query Job on a Cloud DataProc cluster.
-
-    .. seealso::
-        This operator is deprecated, please use
-        :class:`~airflow.providers.google.cloud.operators.dataproc.DataprocSubmitJobOperator`:
-
-    The parameters of the operation will be passed to the cluster.
-
-    It's a good practice to define dataproc_* parameters in the default_args of the dag
-    like the cluster name and UDFs.
-
-    .. code-block:: python
-
-        default_args = {
-            "cluster_name": "cluster-1",
-            "dataproc_pig_jars": [
-                "gs://example/udf/jar/datafu/1.2.0/datafu.jar",
-                "gs://example/udf/jar/gpig/1.2/gpig.jar",
-            ],
-        }
-
-    You can pass a pig script as string or file reference. Use variables to pass on
-    variables for the pig script to be resolved on the cluster or use the parameters to
-    be resolved in the script as template parameters.
-
-    .. code-block:: python
-
-        t1 = DataProcPigOperator(
-            task_id="dataproc_pig",
-            query="a_pig_script.pig",
-            variables={"out": "gs://example/output/{{ds}}"},
-        )
-
-    .. seealso::
-        For more detail on about job submission have a look at the reference:
-        https://cloud.google.com/dataproc/reference/rest/v1/projects.regions.jobs
-
-    :param query: The query or reference to the query
-        file (pg or pig extension). (templated)
-    :param query_uri: The HCFS URI of the script that contains the Pig queries.
-    :param variables: Map of named parameters for the query. (templated)
-    """
-
-    template_fields: Sequence[str] = (
-        "query",
-        "variables",
-        "job_name",
-        "cluster_name",
-        "region",
-        "dataproc_jars",
-        "dataproc_properties",
-        "impersonation_chain",
-    )
-    template_ext = (".pg", ".pig")
-    ui_color = "#0273d4"
-    job_type = "pig_job"
-
-    operator_extra_links = (DataprocLink(),)
-
-    def __init__(
-        self,
-        *,
-        query: str | None = None,
-        query_uri: str | None = None,
-        variables: dict | None = None,
-        impersonation_chain: str | Sequence[str] | None = None,
-        region: str,
-        job_name: str = "{{task.task_id}}_{{ds_nodash}}",
-        cluster_name: str = "cluster-1",
-        dataproc_properties: dict | None = None,
-        dataproc_jars: list[str] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            impersonation_chain=impersonation_chain,
-            region=region,
-            job_name=job_name,
-            cluster_name=cluster_name,
-            dataproc_properties=dataproc_properties,
-            dataproc_jars=dataproc_jars,
-            **kwargs,
-        )
-        self.query = query
-        self.query_uri = query_uri
-        self.variables = variables
-
-    def generate_job(self):
-        """
-        Act as a helper method for easier migration to `DataprocSubmitJobOperator`.
-
-        :return: Dict representing Dataproc job
-        """
-        job_template = self.create_job_template()
-
-        if self.query is None:
-            if self.query_uri is None:
-                raise AirflowException("One of query or query_uri should be set here")
-            job_template.add_query_uri(self.query_uri)
-        else:
-            job_template.add_query(self.query)
-        job_template.add_variables(self.variables)
-        return self._generate_job_template()
-
-    def execute(self, context: Context):
-        job_template = self.create_job_template()
-        if self.query is None:
-            if self.query_uri is None:
-                raise AirflowException("One of query or query_uri should be set here")
-            job_template.add_query_uri(self.query_uri)
-        else:
-            job_template.add_query(self.query)
-        job_template.add_variables(self.variables)
-
-        super().execute(context)
-
-
-# TODO: Remove one day
-
-
-# TODO: Remove one day
-@deprecated(
-    planned_removal_date="November 01, 2024",
-    use_instead="DataprocSubmitJobOperator",
-    instructions="You can use `generate_job` method to generate dictionary representing your job "
-    "and use it with the new operator.",
-    category=AirflowProviderDeprecationWarning,
-)
-class DataprocSubmitHiveJobOperator(DataprocJobBaseOperator):
-    """
-    Start a Hive query Job on a Cloud DataProc cluster.
-
-    .. seealso::
-        This operator is deprecated, please use
-        :class:`~airflow.providers.google.cloud.operators.dataproc.DataprocSubmitJobOperator`:
-
-    :param query: The query or reference to the query file (q extension).
-    :param query_uri: The HCFS URI of the script that contains the Hive queries.
-    :param variables: Map of named parameters for the query.
-    """
-
-    template_fields: Sequence[str] = (
-        "query",
-        "variables",
-        "job_name",
-        "cluster_name",
-        "region",
-        "dataproc_jars",
-        "dataproc_properties",
-        "impersonation_chain",
-    )
-    template_ext = (".q", ".hql")
-    ui_color = "#0273d4"
-    job_type = "hive_job"
-
-    def __init__(
-        self,
-        *,
-        query: str | None = None,
-        query_uri: str | None = None,
-        variables: dict | None = None,
-        impersonation_chain: str | Sequence[str] | None = None,
-        region: str,
-        job_name: str = "{{task.task_id}}_{{ds_nodash}}",
-        cluster_name: str = "cluster-1",
-        dataproc_properties: dict | None = None,
-        dataproc_jars: list[str] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            impersonation_chain=impersonation_chain,
-            region=region,
-            job_name=job_name,
-            cluster_name=cluster_name,
-            dataproc_properties=dataproc_properties,
-            dataproc_jars=dataproc_jars,
-            **kwargs,
-        )
-        self.query = query
-        self.query_uri = query_uri
-        self.variables = variables
-        if self.query is not None and self.query_uri is not None:
-            raise AirflowException("Only one of `query` and `query_uri` can be passed.")
-
-    def generate_job(self):
-        """
-        Act as a helper method for easier migration to `DataprocSubmitJobOperator`.
-
-        :return: Dict representing Dataproc job
-        """
-        job_template = self.create_job_template()
-        if self.query is None:
-            if self.query_uri is None:
-                raise AirflowException("One of query or query_uri should be set here")
-            job_template.add_query_uri(self.query_uri)
-        else:
-            job_template.add_query(self.query)
-        job_template.add_variables(self.variables)
-        return self._generate_job_template()
-
-    def execute(self, context: Context):
-        job_template = self.create_job_template()
-        if self.query is None:
-            if self.query_uri is None:
-                raise AirflowException("One of query or query_uri should be set here")
-            job_template.add_query_uri(self.query_uri)
-        else:
-            job_template.add_query(self.query)
-        job_template.add_variables(self.variables)
-        super().execute(context)
-
-
-# TODO: Remove one day
-@deprecated(
-    planned_removal_date="November 01, 2024",
-    use_instead="DataprocSubmitJobOperator",
-    instructions="You can use `generate_job` method to generate dictionary representing your job "
-    "and use it with the new operator.",
-    category=AirflowProviderDeprecationWarning,
-)
-class DataprocSubmitSparkSqlJobOperator(DataprocJobBaseOperator):
-    """
-    Start a Spark SQL query Job on a Cloud DataProc cluster.
-
-    .. seealso::
-        This operator is deprecated, please use
-        :class:`~airflow.providers.google.cloud.operators.dataproc.DataprocSubmitJobOperator`:
-
-    :param query: The query or reference to the query file (q extension). (templated)
-    :param query_uri: The HCFS URI of the script that contains the SQL queries.
-    :param variables: Map of named parameters for the query. (templated)
-    """
-
-    template_fields: Sequence[str] = (
-        "query",
-        "variables",
-        "job_name",
-        "cluster_name",
-        "region",
-        "dataproc_jars",
-        "dataproc_properties",
-        "impersonation_chain",
-    )
-    template_ext = (".q",)
-    template_fields_renderers = {"sql": "sql"}
-    ui_color = "#0273d4"
-    job_type = "spark_sql_job"
-
-    def __init__(
-        self,
-        *,
-        query: str | None = None,
-        query_uri: str | None = None,
-        variables: dict | None = None,
-        impersonation_chain: str | Sequence[str] | None = None,
-        region: str,
-        job_name: str = "{{task.task_id}}_{{ds_nodash}}",
-        cluster_name: str = "cluster-1",
-        dataproc_properties: dict | None = None,
-        dataproc_jars: list[str] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            impersonation_chain=impersonation_chain,
-            region=region,
-            job_name=job_name,
-            cluster_name=cluster_name,
-            dataproc_properties=dataproc_properties,
-            dataproc_jars=dataproc_jars,
-            **kwargs,
-        )
-        self.query = query
-        self.query_uri = query_uri
-        self.variables = variables
-        if self.query is not None and self.query_uri is not None:
-            raise AirflowException("Only one of `query` and `query_uri` can be passed.")
-
-    def generate_job(self):
-        """
-        Act as a helper method for easier migration to `DataprocSubmitJobOperator`.
-
-        :return: Dict representing Dataproc job
-        """
-        job_template = self.create_job_template()
-        if self.query is None:
-            job_template.add_query_uri(self.query_uri)
-        else:
-            job_template.add_query(self.query)
-        job_template.add_variables(self.variables)
-        return self._generate_job_template()
-
-    def execute(self, context: Context):
-        job_template = self.create_job_template()
-        if self.query is None:
-            if self.query_uri is None:
-                raise AirflowException("One of query or query_uri should be set here")
-            job_template.add_query_uri(self.query_uri)
-        else:
-            job_template.add_query(self.query)
-        job_template.add_variables(self.variables)
-        super().execute(context)
-
-
-# TODO: Remove one day
-@deprecated(
-    planned_removal_date="November 01, 2024",
-    use_instead="DataprocSubmitJobOperator",
-    instructions="You can use `generate_job` method to generate dictionary representing your job "
-    "and use it with the new operator.",
-    category=AirflowProviderDeprecationWarning,
-)
-class DataprocSubmitSparkJobOperator(DataprocJobBaseOperator):
-    """
-    Start a Spark Job on a Cloud DataProc cluster.
-
-    .. seealso::
-        This operator is deprecated, please use
-        :class:`~airflow.providers.google.cloud.operators.dataproc.DataprocSubmitJobOperator`:
-
-    :param main_jar: The HCFS URI of the jar file that contains the main class
-        (use this or the main_class, not both together).
-    :param main_class: Name of the job class. (use this or the main_jar, not both
-        together).
-    :param arguments: Arguments for the job. (templated)
-    :param archives: List of archived files that will be unpacked in the work
-        directory. Should be stored in Cloud Storage.
-    :param files: List of files to be copied to the working directory
-    """
-
-    template_fields: Sequence[str] = (
-        "arguments",
-        "job_name",
-        "cluster_name",
-        "region",
-        "dataproc_jars",
-        "dataproc_properties",
-        "impersonation_chain",
-    )
-    ui_color = "#0273d4"
-    job_type = "spark_job"
-
-    def __init__(
-        self,
-        *,
-        main_jar: str | None = None,
-        main_class: str | None = None,
-        arguments: list | None = None,
-        archives: list | None = None,
-        files: list | None = None,
-        impersonation_chain: str | Sequence[str] | None = None,
-        region: str,
-        job_name: str = "{{task.task_id}}_{{ds_nodash}}",
-        cluster_name: str = "cluster-1",
-        dataproc_properties: dict | None = None,
-        dataproc_jars: list[str] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            impersonation_chain=impersonation_chain,
-            region=region,
-            job_name=job_name,
-            cluster_name=cluster_name,
-            dataproc_properties=dataproc_properties,
-            dataproc_jars=dataproc_jars,
-            **kwargs,
-        )
-        self.main_jar = main_jar
-        self.main_class = main_class
-        self.arguments = arguments
-        self.archives = archives
-        self.files = files
-
-    def generate_job(self):
-        """
-        Act as a helper method for easier migration to `DataprocSubmitJobOperator`.
-
-        :return: Dict representing Dataproc job
-        """
-        job_template = self.create_job_template()
-        job_template.set_main(self.main_jar, self.main_class)
-        job_template.add_args(self.arguments)
-        job_template.add_archive_uris(self.archives)
-        job_template.add_file_uris(self.files)
-        return self._generate_job_template()
-
-    def execute(self, context: Context):
-        job_template = self.create_job_template()
-        job_template.set_main(self.main_jar, self.main_class)
-        job_template.add_args(self.arguments)
-        job_template.add_archive_uris(self.archives)
-        job_template.add_file_uris(self.files)
-        super().execute(context)
-
-
-# TODO: Remove one day
-@deprecated(
-    planned_removal_date="November 01, 2024",
-    use_instead="DataprocSubmitJobOperator",
-    instructions="You can use `generate_job` method to generate dictionary representing your job "
-    "and use it with the new operator.",
-    category=AirflowProviderDeprecationWarning,
-)
-class DataprocSubmitHadoopJobOperator(DataprocJobBaseOperator):
-    """
-    Start a Hadoop Job on a Cloud DataProc cluster.
-
-    .. seealso::
-        This operator is deprecated, please use
-        :class:`~airflow.providers.google.cloud.operators.dataproc.DataprocSubmitJobOperator`:
-
-    :param main_jar: The HCFS URI of the jar file containing the main class
-        (use this or the main_class, not both together).
-    :param main_class: Name of the job class. (use this or the main_jar, not both
-        together).
-    :param arguments: Arguments for the job. (templated)
-    :param archives: List of archived files that will be unpacked in the work
-        directory. Should be stored in Cloud Storage.
-    :param files: List of files to be copied to the working directory
-    """
-
-    template_fields: Sequence[str] = (
-        "arguments",
-        "job_name",
-        "cluster_name",
-        "region",
-        "dataproc_jars",
-        "dataproc_properties",
-        "impersonation_chain",
-    )
-    ui_color = "#0273d4"
-    job_type = "hadoop_job"
-
-    def __init__(
-        self,
-        *,
-        main_jar: str | None = None,
-        main_class: str | None = None,
-        arguments: list | None = None,
-        archives: list | None = None,
-        files: list | None = None,
-        impersonation_chain: str | Sequence[str] | None = None,
-        region: str,
-        job_name: str = "{{task.task_id}}_{{ds_nodash}}",
-        cluster_name: str = "cluster-1",
-        dataproc_properties: dict | None = None,
-        dataproc_jars: list[str] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            impersonation_chain=impersonation_chain,
-            region=region,
-            job_name=job_name,
-            cluster_name=cluster_name,
-            dataproc_properties=dataproc_properties,
-            dataproc_jars=dataproc_jars,
-            **kwargs,
-        )
-        self.main_jar = main_jar
-        self.main_class = main_class
-        self.arguments = arguments
-        self.archives = archives
-        self.files = files
-
-    def generate_job(self):
-        """
-        Act as a helper method for easier migration to `DataprocSubmitJobOperator`.
-
-        :return: Dict representing Dataproc job
-        """
-        job_template = self.create_job_template()
-        job_template.set_main(self.main_jar, self.main_class)
-        job_template.add_args(self.arguments)
-        job_template.add_archive_uris(self.archives)
-        job_template.add_file_uris(self.files)
-        return self._generate_job_template()
-
-    def execute(self, context: Context):
-        job_template = self.create_job_template()
-        job_template.set_main(self.main_jar, self.main_class)
-        job_template.add_args(self.arguments)
-        job_template.add_archive_uris(self.archives)
-        job_template.add_file_uris(self.files)
-        super().execute(context)
-
-
-# TODO: Remove one day
-@deprecated(
-    planned_removal_date="November 01, 2024",
-    use_instead="DataprocSubmitJobOperator",
-    instructions="You can use `generate_job` method to generate dictionary representing your job "
-    "and use it with the new operator.",
-    category=AirflowProviderDeprecationWarning,
-)
-class DataprocSubmitPySparkJobOperator(DataprocJobBaseOperator):
-    """
-    Start a PySpark Job on a Cloud DataProc cluster.
-
-    .. seealso::
-        This operator is deprecated, please use
-        :class:`~airflow.providers.google.cloud.operators.dataproc.DataprocSubmitJobOperator`:
-
-    :param main: [Required] The Hadoop Compatible Filesystem (HCFS) URI of the main
-            Python file to use as the driver. Must be a .py file. (templated)
-    :param arguments: Arguments for the job. (templated)
-    :param archives: List of archived files that will be unpacked in the work
-        directory. Should be stored in Cloud Storage.
-    :param files: List of files to be copied to the working directory
-    :param pyfiles: List of Python files to pass to the PySpark framework.
-        Supported file types: .py, .egg, and .zip
-    """
-
-    template_fields: Sequence[str] = (
-        "main",
-        "arguments",
-        "job_name",
-        "cluster_name",
-        "region",
-        "dataproc_jars",
-        "dataproc_properties",
-        "impersonation_chain",
-    )
-    ui_color = "#0273d4"
-    job_type = "pyspark_job"
-
-    @staticmethod
-    def _generate_temp_filename(filename):
-        return f"{time.strftime('%Y%m%d%H%M%S')}_{uuid.uuid4()!s:.8}_{ntpath.basename(filename)}"
-
-    def _upload_file_temp(self, bucket, local_file):
-        """Upload a local file to a Google Cloud Storage bucket."""
-        temp_filename = self._generate_temp_filename(local_file)
-        if not bucket:
-            raise AirflowException(
-                "If you want Airflow to upload the local file to a temporary bucket, set "
-                "the 'temp_bucket' key in the connection string"
-            )
-
-        self.log.info("Uploading %s to %s", local_file, temp_filename)
-
-        GCSHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain).upload(
-            bucket_name=bucket,
-            object_name=temp_filename,
-            mime_type="application/x-python",
-            filename=local_file,
-        )
-        return f"gs://{bucket}/{temp_filename}"
-
-    def __init__(
-        self,
-        *,
-        main: str,
-        arguments: list | None = None,
-        archives: list | None = None,
-        pyfiles: list | None = None,
-        files: list | None = None,
-        impersonation_chain: str | Sequence[str] | None = None,
-        region: str,
-        job_name: str = "{{task.task_id}}_{{ds_nodash}}",
-        cluster_name: str = "cluster-1",
-        dataproc_properties: dict | None = None,
-        dataproc_jars: list[str] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            impersonation_chain=impersonation_chain,
-            region=region,
-            job_name=job_name,
-            cluster_name=cluster_name,
-            dataproc_properties=dataproc_properties,
-            dataproc_jars=dataproc_jars,
-            **kwargs,
-        )
-        self.main = main
-        self.arguments = arguments
-        self.archives = archives
-        self.files = files
-        self.pyfiles = pyfiles
-
-    def generate_job(self):
-        """
-        Act as a helper method for easier migration to :class:`DataprocSubmitJobOperator`.
-
-        :return: Dict representing Dataproc job
-        """
-        job_template = self.create_job_template()
-        #  Check if the file is local, if that is the case, upload it to a bucket
-        if os.path.isfile(self.main):
-            cluster_info = self.hook.get_cluster(
-                project_id=self.project_id, region=self.region, cluster_name=self.cluster_name
-            )
-            bucket = cluster_info["config"]["config_bucket"]
-            self.main = f"gs://{bucket}/{self.main}"
-        job_template.set_python_main(self.main)
-        job_template.add_args(self.arguments)
-        job_template.add_archive_uris(self.archives)
-        job_template.add_file_uris(self.files)
-        job_template.add_python_file_uris(self.pyfiles)
-
-        return self._generate_job_template()
-
-    def execute(self, context: Context):
-        job_template = self.create_job_template()
-        #  Check if the file is local, if that is the case, upload it to a bucket
-        if os.path.isfile(self.main):
-            cluster_info = self.hook.get_cluster(
-                project_id=self.project_id, region=self.region, cluster_name=self.cluster_name
-            )
-            bucket = cluster_info["config"]["config_bucket"]
-            self.main = self._upload_file_temp(bucket, self.main)
-
-        job_template.set_python_main(self.main)
-        job_template.add_args(self.arguments)
-        job_template.add_archive_uris(self.archives)
-        job_template.add_file_uris(self.files)
-        job_template.add_python_file_uris(self.pyfiles)
-        super().execute(context)
-
-
 class DataprocCreateWorkflowTemplateOperator(GoogleCloudBaseOperator):
     """
     Creates new workflow template.

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

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import time
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from google.api_core.exceptions import AlreadyExists
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.datastore import DatastoreHook

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from google.api_core.exceptions import AlreadyExists, InvalidArgument, NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault