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

airflow/providers/google/cloud/links/translate.py

@@ -50,6 +50,12 @@ TRANSLATION_NATIVE_DATASET_LINK = (
 )
 TRANSLATION_NATIVE_LIST_LINK = TRANSLATION_BASE_LINK + "/datasets?project={project_id}"
 
+TRANSLATION_NATIVE_MODEL_LINK = (
+    TRANSLATION_BASE_LINK
+    + "/locations/{location}/datasets/{dataset_id}/evaluate;modelId={model_id}?project={project_id}"
+)
+TRANSLATION_MODELS_LIST_LINK = TRANSLATION_BASE_LINK + "/models/list?project={project_id}"
+
 
 class TranslationLegacyDatasetLink(BaseGoogleLink):
     """
@@ -270,3 +276,95 @@ class TranslationDatasetsListLink(BaseGoogleLink):
                 "project_id": project_id,
             },
         )
+
+
+class TranslationModelLink(BaseGoogleLink):
+    """
+    Helper class for constructing Translation Model link.
+
+    Link for legacy and native models.
+    """
+
+    name = "Translation Model"
+    key = "translation_model"
+    format_str = TRANSLATION_NATIVE_MODEL_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+        dataset_id: str,
+        model_id: str,
+        project_id: str,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=TranslationLegacyModelLink.key,
+            value={
+                "location": task_instance.location,
+                "dataset_id": dataset_id,
+                "model_id": model_id,
+                "project_id": project_id,
+            },
+        )
+
+
+class TranslationModelsListLink(BaseGoogleLink):
+    """
+    Helper class for constructing Translation Models List link.
+
+    Both legacy and native models are available under this link.
+    """
+
+    name = "Translation Models List"
+    key = "translation_models_list"
+    format_str = TRANSLATION_MODELS_LIST_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+        project_id: str,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=TranslationModelsListLink.key,
+            value={
+                "project_id": project_id,
+            },
+        )
+
+
+class TranslateResultByOutputConfigLink(BaseGoogleLink):
+    """
+    Helper class for constructing Translation results Link.
+
+    Provides link to gcs destination output translation results, by provided output_config
+    with gcs destination specified.
+    """
+
+    name = "Translate Results By Output Config"
+    key = "translate_results_by_output_config"
+    format_str = TRANSLATION_TRANSLATE_TEXT_BATCH
+
+    @staticmethod
+    def extract_output_uri_prefix(output_config):
+        return output_config["gcs_destination"]["output_uri_prefix"].rpartition("gs://")[-1]
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+        project_id: str,
+        output_config: dict,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=TranslateResultByOutputConfigLink.key,
+            value={
+                "project_id": project_id,
+                "output_uri_prefix": TranslateResultByOutputConfigLink.extract_output_uri_prefix(
+                    output_config
+                ),
+            },
+        )

airflow/providers/google/cloud/log/stackdriver_task_handler.py

@@ -30,18 +30,14 @@ from google.cloud.logging import Resource
 from google.cloud.logging.handlers.transports import BackgroundThreadTransport, Transport
 from google.cloud.logging_v2.services.logging_service_v2 import LoggingServiceV2Client
 from google.cloud.logging_v2.types import ListLogEntriesRequest, ListLogEntriesResponse
-from packaging.version import Version
 
-from airflow import __version__ as airflow_version
 from airflow.exceptions import RemovedInAirflow3Warning
 from airflow.providers.google.cloud.utils.credentials_provider import get_credentials_and_project_id
 from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.utils.log.trigger_handler import ctx_indiv_trigger
 from airflow.utils.types import NOTSET, ArgNotSet
 
-AIRFLOW_VERSION = Version(airflow_version)
-AIRFLOW_V_3_0_PLUS = Version(AIRFLOW_VERSION.base_version) >= Version("3.0.0")
-
 if TYPE_CHECKING:
     from google.auth.credentials import Credentials
 

airflow/providers/google/cloud/openlineage/mixins.py

@@ -88,7 +88,7 @@ class _BigQueryOpenLineageMixin:
 
         job_facets = {"sql": SQLJobFacet(query=SQLParser.normalize_sql(self.sql))}
 
-        self.client = self.hook.get_client(project_id=self.hook.project_id)
+        self.client = self.hook.get_client(project_id=self.hook.project_id, location=self.location)
         job_ids = self.job_id
         if isinstance(self.job_id, str):
             job_ids = [self.job_id]
@@ -108,10 +108,7 @@ class _BigQueryOpenLineageMixin:
 
     def get_facets(self, job_id: str):
         from airflow.providers.common.compat.openlineage.facet import ErrorMessageRunFacet
-        from airflow.providers.google.cloud.openlineage.utils import (
-            BigQueryErrorRunFacet,
-            get_from_nullable_chain,
-        )
+        from airflow.providers.google.cloud.openlineage.utils import get_from_nullable_chain
 
         inputs = []
         outputs = []
@@ -125,8 +122,7 @@ class _BigQueryOpenLineageMixin:
             if get_from_nullable_chain(props, ["status", "state"]) != "DONE":
                 raise ValueError(f"Trying to extract data from running bigquery job: `{job_id}`")
 
-            # TODO: remove bigQuery_job in next release
-            run_facets["bigQuery_job"] = run_facets["bigQueryJob"] = self._get_bigquery_job_run_facet(props)
+            run_facets["bigQueryJob"] = self._get_bigquery_job_run_facet(props)
 
             if get_from_nullable_chain(props, ["statistics", "numChildJobs"]):
                 if hasattr(self, "log"):
@@ -145,16 +141,12 @@ class _BigQueryOpenLineageMixin:
             if hasattr(self, "log"):
                 self.log.warning("Cannot retrieve job details from BigQuery.Client. %s", e, exc_info=True)
             exception_msg = traceback.format_exc()
-            # TODO: remove BigQueryErrorRunFacet in next release
             run_facets.update(
                 {
                     "errorMessage": ErrorMessageRunFacet(
                         message=f"{e}: {exception_msg}",
                         programmingLanguage="python",
-                    ),
-                    "bigQuery_error": BigQueryErrorRunFacet(
-                        clientError=f"{e}: {exception_msg}",
-                    ),
+                    )
                 }
             )
         deduplicated_outputs = self._deduplicate_outputs(outputs)

airflow/providers/google/cloud/openlineage/utils.py

@@ -17,6 +17,9 @@
 # under the License.
 from __future__ import annotations
 
+import logging
+import os
+import pathlib
 from typing import TYPE_CHECKING, Any
 
 from attr import define, field
@@ -25,21 +28,84 @@ if TYPE_CHECKING:
     from google.cloud.bigquery.table import Table
 
     from airflow.providers.common.compat.openlineage.facet import Dataset
+    from airflow.utils.context import Context
 
 from airflow.providers.common.compat.openlineage.facet import (
     BaseFacet,
     ColumnLineageDatasetFacet,
     DocumentationDatasetFacet,
     Fields,
+    Identifier,
     InputField,
     RunFacet,
     SchemaDatasetFacet,
     SchemaDatasetFacetFields,
+    SymlinksDatasetFacet,
+)
+from airflow.providers.common.compat.openlineage.utils.spark import (
+    inject_parent_job_information_into_spark_properties,
 )
 from airflow.providers.google import __version__ as provider_version
+from airflow.providers.google.cloud.hooks.gcs import _parse_gcs_url
+
+log = logging.getLogger(__name__)
 
 BIGQUERY_NAMESPACE = "bigquery"
 BIGQUERY_URI = "bigquery"
+WILDCARD = "*"
+
+
+def extract_ds_name_from_gcs_path(path: str) -> str:
+    """
+    Extract and process the dataset name from a given path.
+
+    Args:
+        path: The path to process e.g. of a gcs file.
+
+    Returns:
+        The processed dataset name.
+
+    Examples:
+        >>> extract_ds_name_from_gcs_path("/dir/file.*")
+        'dir'
+        >>> extract_ds_name_from_gcs_path("/dir/pre_")
+        'dir'
+        >>> extract_ds_name_from_gcs_path("/dir/file.txt")
+        'dir/file.txt'
+        >>> extract_ds_name_from_gcs_path("/dir/file.")
+        'dir'
+        >>> extract_ds_name_from_gcs_path("/dir/")
+        'dir'
+        >>> extract_ds_name_from_gcs_path("")
+        '/'
+        >>> extract_ds_name_from_gcs_path("/")
+        '/'
+        >>> extract_ds_name_from_gcs_path(".")
+        '/'
+    """
+    if WILDCARD in path:
+        path = path.split(WILDCARD, maxsplit=1)[0]
+
+    # We want to end up with parent directory if the path:
+    # - does not refer to a file (no dot in the last segment)
+    #     and does not explicitly end with a slash, it is treated as a prefix and removed.
+    #     Example: "/dir/pre_" -> "/dir/"
+    # - contains a dot at the end, then it is treated as a prefix (created after removing the wildcard).
+    #     Example: "/dir/file." (was "/dir/file.*" with wildcard) -> "/dir/"
+    last_path_segment = os.path.basename(path).rstrip(".")
+    if "." not in last_path_segment and not path.endswith("/"):
+        path = pathlib.Path(path).parent.as_posix()
+
+    # Normalize the path:
+    # - Remove trailing slashes.
+    # - Remove leading slashes.
+    # - Handle edge cases for empty paths or single-dot paths.
+    path = path.rstrip("/")
+    path = path.lstrip("/")
+    if path in ("", "."):
+        path = "/"
+
+    return path
 
 
 def get_facets_from_bq_table(table: Table) -> dict[str, BaseFacet]:
@@ -57,6 +123,20 @@ def get_facets_from_bq_table(table: Table) -> dict[str, BaseFacet]:
     if table.description:
         facets["documentation"] = DocumentationDatasetFacet(description=table.description)
 
+    if table.external_data_configuration:
+        symlinks = set()
+        for uri in table.external_data_configuration.source_uris:
+            if uri.startswith("gs://"):
+                bucket, blob = _parse_gcs_url(uri)
+                blob = extract_ds_name_from_gcs_path(blob)
+                symlinks.add((f"gs://{bucket}", blob))
+
+        facets["symlink"] = SymlinksDatasetFacet(
+            identifiers=[
+                Identifier(namespace=namespace, name=name, type="file")
+                for namespace, name in sorted(symlinks)
+            ]
+        )
     return facets
 
 
@@ -145,28 +225,6 @@ class BigQueryJobRunFacet(RunFacet):
         )
 
 
-# TODO: remove BigQueryErrorRunFacet in next release
-@define
-class BigQueryErrorRunFacet(RunFacet):
-    """
-    Represents errors that can happen during execution of BigqueryExtractor.
-
-    :param clientError: represents errors originating in bigquery client
-    :param parserError: represents errors that happened during parsing SQL provided to bigquery
-    """
-
-    clientError: str | None = field(default=None)
-    parserError: str | None = field(default=None)
-
-    @staticmethod
-    def _get_schema() -> str:
-        return (
-            "https://raw.githubusercontent.com/apache/airflow/"
-            f"providers-google/{provider_version}/airflow/providers/google/"
-            "openlineage/BigQueryErrorRunFacet.json"
-        )
-
-
 def get_from_nullable_chain(source: Any, chain: list[str]) -> Any | None:
     """
     Get object from nested structure of objects, where it's not guaranteed that all keys in the nested structure exist.
@@ -208,3 +266,123 @@ def get_from_nullable_chain(source: Any, chain: list[str]) -> Any | None:
         return source
     except AttributeError:
         return None
+
+
+def _is_openlineage_provider_accessible() -> bool:
+    """
+    Check if the OpenLineage provider is accessible.
+
+    This function attempts to import the necessary OpenLineage modules and checks if the provider
+    is enabled and the listener is available.
+
+    Returns:
+        bool: True if the OpenLineage provider is accessible, False otherwise.
+    """
+    try:
+        from airflow.providers.openlineage.conf import is_disabled
+        from airflow.providers.openlineage.plugins.listener import get_openlineage_listener
+    except ImportError:
+        log.debug("OpenLineage provider could not be imported.")
+        return False
+
+    if is_disabled():
+        log.debug("OpenLineage provider is disabled.")
+        return False
+
+    if not get_openlineage_listener():
+        log.debug("OpenLineage listener could not be found.")
+        return False
+
+    return True
+
+
+def _extract_supported_job_type_from_dataproc_job(job: dict) -> str | None:
+    """
+    Extract job type from a Dataproc job definition.
+
+    Args:
+        job: The Dataproc job definition.
+
+    Returns:
+        The job type for which the automatic OL injection is supported, if found, otherwise None.
+    """
+    supported_job_types = ("sparkJob", "pysparkJob", "spark_job", "pyspark_job")
+    return next((job_type for job_type in supported_job_types if job_type in job), None)
+
+
+def _replace_dataproc_job_properties(job: dict, job_type: str, new_properties: dict) -> dict:
+    """
+    Replace the properties of a specific job type in a Dataproc job definition.
+
+    Args:
+        job: The original Dataproc job definition.
+        job_type: The key representing the job type (e.g., "sparkJob").
+        new_properties: The new properties to replace the existing ones.
+
+    Returns:
+        A modified copy of the job with updated properties.
+
+    Raises:
+        KeyError: If the job_type does not exist in the job or lacks a "properties" field.
+    """
+    if job_type not in job:
+        raise KeyError(f"Job type '{job_type}' is missing in the job definition.")
+
+    updated_job = job.copy()
+    updated_job[job_type] = job[job_type].copy()
+    updated_job[job_type]["properties"] = new_properties
+
+    return updated_job
+
+
+def inject_openlineage_properties_into_dataproc_job(
+    job: dict, context: Context, inject_parent_job_info: bool
+) -> dict:
+    """
+    Inject OpenLineage properties into Spark job definition.
+
+    Function is not removing any configuration or modifying the job in any other way,
+    apart from adding desired OpenLineage properties to Dataproc job definition if not already present.
+
+    Note:
+        Any modification to job will be skipped if:
+            - OpenLineage provider is not accessible.
+            - The job type is not supported.
+            - Automatic parent job information injection is disabled.
+            - Any OpenLineage properties with parent job information are already present
+              in the Spark job definition.
+
+    Args:
+        job: The original Dataproc job definition.
+        context: The Airflow context in which the job is running.
+        inject_parent_job_info: Flag indicating whether to inject parent job information.
+
+    Returns:
+        The modified job definition with OpenLineage properties injected, if applicable.
+    """
+    if not inject_parent_job_info:
+        log.debug("Automatic injection of OpenLineage information is disabled.")
+        return job
+
+    if not _is_openlineage_provider_accessible():
+        log.warning(
+            "Could not access OpenLineage provider for automatic OpenLineage "
+            "properties injection. No action will be performed."
+        )
+        return job
+
+    if (job_type := _extract_supported_job_type_from_dataproc_job(job)) is None:
+        log.warning(
+            "Could not find a supported Dataproc job type for automatic OpenLineage "
+            "properties injection. No action will be performed.",
+        )
+        return job
+
+    properties = job[job_type].get("properties", {})
+
+    properties = inject_parent_job_information_into_spark_properties(properties=properties, context=context)
+
+    job_with_ol_config = _replace_dataproc_job_properties(
+        job=job, job_type=job_type, new_properties=properties
+    )
+    return job_with_ol_config