apache-airflow-providers-google 15.1.0rc1__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

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

@@ -21,7 +21,7 @@ from typing import TYPE_CHECKING
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 VERTEX_AI_BASE_LINK = "/vertex-ai"
 VERTEX_AI_MODEL_LINK = (
@@ -54,6 +54,10 @@ VERTEX_AI_PIPELINE_JOB_LINK = (
     VERTEX_AI_BASE_LINK + "/locations/{region}/pipelines/runs/{pipeline_id}?project={project_id}"
 )
 VERTEX_AI_PIPELINE_JOB_LIST_LINK = VERTEX_AI_BASE_LINK + "/pipelines/runs?project={project_id}"
+VERTEX_AI_RAY_CLUSTER_LINK = (
+    VERTEX_AI_BASE_LINK + "/locations/{location}/ray-clusters/{cluster_id}?project={project_id}"
+)
+VERTEX_AI_RAY_CLUSTER_LIST_LINK = VERTEX_AI_BASE_LINK + "/ray?project={project_id}"
 
 
 class VertexAIModelLink(BaseGoogleLink):
@@ -63,22 +67,6 @@ class VertexAIModelLink(BaseGoogleLink):
     key = "model_conf"
     format_str = VERTEX_AI_MODEL_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        model_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIModelLink.key,
-            value={
-                "model_id": model_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIModelListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Models Link."""
@@ -87,19 +75,6 @@ class VertexAIModelListLink(BaseGoogleLink):
     key = "models_conf"
     format_str = VERTEX_AI_MODEL_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIModelListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIModelExportLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Model Export Link."""
@@ -113,19 +88,15 @@ class VertexAIModelExportLink(BaseGoogleLink):
         """Return bucket name from output configuration."""
         return config["artifact_destination"]["output_uri_prefix"].rpartition("gs://")[-1]
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
+    @classmethod
+    def persist(cls, context: Context, **value):
+        output_config = value.get("output_config")
+        bucket_name = cls.extract_bucket_name(output_config)
+        super().persist(
             context=context,
-            key=VertexAIModelExportLink.key,
-            value={
-                "project_id": task_instance.project_id,
-                "model_id": task_instance.model_id,
-                "bucket_name": VertexAIModelExportLink.extract_bucket_name(task_instance.output_config),
-            },
+            project_id=value.get("project_id"),
+            model_id=value.get("model_id"),
+            bucket_name=bucket_name,
         )
 
 
@@ -136,22 +107,6 @@ class VertexAITrainingLink(BaseGoogleLink):
     key = "training_conf"
     format_str = VERTEX_AI_TRAINING_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        training_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAITrainingLink.key,
-            value={
-                "training_id": training_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAITrainingPipelinesLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Training Pipelines link."""
@@ -160,19 +115,6 @@ class VertexAITrainingPipelinesLink(BaseGoogleLink):
     key = "pipelines_conf"
     format_str = VERTEX_AI_TRAINING_PIPELINES_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAITrainingPipelinesLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIDatasetLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Dataset link."""
@@ -181,18 +123,6 @@ class VertexAIDatasetLink(BaseGoogleLink):
     key = "dataset_conf"
     format_str = VERTEX_AI_DATASET_LINK
 
-    @staticmethod
-    def persist(context: Context, task_instance, dataset_id: str):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIDatasetLink.key,
-            value={
-                "dataset_id": dataset_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIDatasetListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Datasets Link."""
@@ -201,19 +131,6 @@ class VertexAIDatasetListLink(BaseGoogleLink):
     key = "datasets_conf"
     format_str = VERTEX_AI_DATASET_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIDatasetListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIHyperparameterTuningJobListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI HyperparameterTuningJobs Link."""
@@ -222,19 +139,6 @@ class VertexAIHyperparameterTuningJobListLink(BaseGoogleLink):
     key = "hyperparameter_tuning_jobs_conf"
     format_str = VERTEX_AI_HYPERPARAMETER_TUNING_JOB_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIHyperparameterTuningJobListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIBatchPredictionJobLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI BatchPredictionJob link."""
@@ -243,22 +147,6 @@ class VertexAIBatchPredictionJobLink(BaseGoogleLink):
     key = "batch_prediction_job_conf"
     format_str = VERTEX_AI_BATCH_PREDICTION_JOB_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        batch_prediction_job_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIBatchPredictionJobLink.key,
-            value={
-                "batch_prediction_job_id": batch_prediction_job_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIBatchPredictionJobListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI BatchPredictionJobList link."""
@@ -267,19 +155,6 @@ class VertexAIBatchPredictionJobListLink(BaseGoogleLink):
     key = "batch_prediction_jobs_conf"
     format_str = VERTEX_AI_BATCH_PREDICTION_JOB_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIBatchPredictionJobListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIEndpointLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Endpoint link."""
@@ -288,22 +163,6 @@ class VertexAIEndpointLink(BaseGoogleLink):
     key = "endpoint_conf"
     format_str = VERTEX_AI_ENDPOINT_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        endpoint_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIEndpointLink.key,
-            value={
-                "endpoint_id": endpoint_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIEndpointListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI EndpointList link."""
@@ -312,19 +171,6 @@ class VertexAIEndpointListLink(BaseGoogleLink):
     key = "endpoints_conf"
     format_str = VERTEX_AI_ENDPOINT_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIEndpointListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIPipelineJobLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI PipelineJob link."""
@@ -333,22 +179,6 @@ class VertexAIPipelineJobLink(BaseGoogleLink):
     key = "pipeline_job_conf"
     format_str = VERTEX_AI_PIPELINE_JOB_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        pipeline_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIPipelineJobLink.key,
-            value={
-                "pipeline_id": pipeline_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIPipelineJobListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI PipelineJobList link."""
@@ -357,15 +187,18 @@ class VertexAIPipelineJobListLink(BaseGoogleLink):
     key = "pipeline_job_list_conf"
     format_str = VERTEX_AI_PIPELINE_JOB_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIPipelineJobListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
+
+class VertexAIRayClusterLink(BaseGoogleLink):
+    """Helper class for constructing Vertex AI Ray Cluster link."""
+
+    name = "Ray Cluster"
+    key = "ray_cluster_conf"
+    format_str = VERTEX_AI_RAY_CLUSTER_LINK
+
+
+class VertexAIRayClusterListLink(BaseGoogleLink):
+    """Helper class for constructing Vertex AI Ray Cluster List link."""
+
+    name = "Ray Cluster List"
+    key = "ray_cluster_list_conf"
+    format_str = VERTEX_AI_RAY_CLUSTER_LIST_LINK

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

@@ -19,14 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.utils.context import Context
-
 WORKFLOWS_BASE_LINK = "/workflows"
 WORKFLOW_LINK = WORKFLOWS_BASE_LINK + "/workflow/{location_id}/{workflow_id}/executions?project={project_id}"
 WORKFLOWS_LINK = WORKFLOWS_BASE_LINK + "?project={project_id}"
@@ -43,20 +37,6 @@ class WorkflowsWorkflowDetailsLink(BaseGoogleLink):
     key = "workflow_details"
     format_str = WORKFLOW_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        location_id: str,
-        workflow_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=WorkflowsWorkflowDetailsLink.key,
-            value={"location_id": location_id, "workflow_id": workflow_id, "project_id": project_id},
-        )
-
 
 class WorkflowsListOfWorkflowsLink(BaseGoogleLink):
     """Helper class for constructing list of Workflows Link."""
@@ -65,18 +45,6 @@ class WorkflowsListOfWorkflowsLink(BaseGoogleLink):
     key = "list_of_workflows"
     format_str = WORKFLOWS_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=WorkflowsListOfWorkflowsLink.key,
-            value={"project_id": project_id},
-        )
-
 
 class WorkflowsExecutionLink(BaseGoogleLink):
     """Helper class for constructing Workflows Execution Link."""
@@ -84,23 +52,3 @@ class WorkflowsExecutionLink(BaseGoogleLink):
     name = "Workflow Execution"
     key = "workflow_execution"
     format_str = EXECUTION_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        location_id: str,
-        workflow_id: str,
-        execution_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=WorkflowsExecutionLink.key,
-            value={
-                "location_id": location_id,
-                "workflow_id": workflow_id,
-                "execution_id": execution_id,
-                "project_id": project_id,
-            },
-        )

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

@@ -27,8 +27,8 @@ from typing import TYPE_CHECKING
 
 import attrs
 
-# not sure why but mypy complains on missing `storage` but it is clearly there and is importable
-from google.cloud import storage  # type: ignore[attr-defined]
+# Make mypy happy by importing as aliases
+import google.cloud.storage as storage
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowNotFoundException
@@ -61,13 +61,15 @@ class GCSRemoteLogIO(LoggingMixin):  # noqa: D101
     remote_base: str
     base_log_folder: Path = attrs.field(converter=Path)
     delete_local_copy: bool
+    project_id: str | None = None
 
-    gcp_key_path: str | None
-    gcp_keyfile_dict: dict | None
-    scopes: Collection[str] | None
-    project_id: str
+    gcp_key_path: str | None = None
+    gcp_keyfile_dict: dict | None = None
+    scopes: Collection[str] | None = _DEFAULT_SCOPESS
 
-    def upload(self, path: os.PathLike, ti: RuntimeTI):
+    processors = ()
+
+    def upload(self, path: os.PathLike | str, ti: RuntimeTI):
         """Upload the given log path to the remote storage."""
         path = Path(path)
         if path.is_absolute():
@@ -211,9 +213,15 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
         gcp_keyfile_dict: dict | None = None,
         gcp_scopes: Collection[str] | None = _DEFAULT_SCOPESS,
         project_id: str = PROVIDE_PROJECT_ID,
+        max_bytes: int = 0,
+        backup_count: int = 0,
+        delay: bool = False,
         **kwargs,
-    ):
-        super().__init__(base_log_folder)
+    ) -> None:
+        # support log file size handling of FileTaskHandler
+        super().__init__(
+            base_log_folder=base_log_folder, max_bytes=max_bytes, backup_count=backup_count, delay=delay
+        )
         self.handler: logging.FileHandler | None = None
         self.log_relative_path = ""
         self.closed = False

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

@@ -35,17 +35,20 @@ from airflow.exceptions import AirflowProviderDeprecationWarning
 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.types import NOTSET, ArgNotSet
+
+try:
+    from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
+except ImportError:
+    from airflow.utils.types import NOTSET, ArgNotSet  # type: ignore[attr-defined,no-redef]
+
+if not AIRFLOW_V_3_0_PLUS:
+    from airflow.utils.log.trigger_handler import ctx_indiv_trigger
 
 if TYPE_CHECKING:
     from google.auth.credentials import Credentials
 
     from airflow.models import TaskInstance
 
-
-if not AIRFLOW_V_3_0_PLUS:
-    from airflow.utils.log.trigger_handler import ctx_indiv_trigger
-
 DEFAULT_LOGGER_NAME = "airflow"
 _GLOBAL_RESOURCE = Resource(type="global", labels={})
 
@@ -159,7 +162,7 @@ class StackdriverTaskHandler(logging.Handler):
         """Object responsible for sending data to Stackdriver."""
         # The Transport object is badly defined (no init) but in the docs client/name as constructor
         # arguments are a requirement for any class that derives from Transport class, hence ignore:
-        return self.transport_type(self._client, self.gcp_log_name)  # type: ignore[call-arg]
+        return self.transport_type(self._client, self.gcp_log_name)
 
     def _get_labels(self, task_instance=None):
         if task_instance:

airflow/providers/google/cloud/openlineage/CloudStorageTransferJobFacet.json

@@ -0,0 +1,68 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$defs": {
+    "CloudStorageTransferJobFacet": {
+      "allOf": [
+        {
+          "$ref": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/JobFacet"
+        },
+        {
+          "type": "object",
+          "properties": {
+            "jobName": {
+              "type": "string",
+              "description": "Transfer job name assigned by GCP Storage Transfer Service."
+            },
+            "projectId": {
+              "type": "string",
+              "description": "GCP project ID."
+            },
+            "description": {
+              "type": "string",
+              "description": "Optional description of the transfer job."
+            },
+            "status": {
+              "type": "string",
+              "description": "Status of the transfer job (ENABLED, DISABLED)."
+            },
+            "sourceBucket": {
+              "type": "string",
+              "description": "Source AWS S3 bucket."
+            },
+            "sourcePath": {
+              "type": "string",
+              "description": "Prefix path inside the source bucket."
+            },
+            "targetBucket": {
+              "type": "string",
+              "description": "Target GCS bucket."
+            },
+            "targetPath": {
+              "type": "string",
+              "description": "Prefix path inside the target bucket."
+            },
+            "objectConditions": {
+              "type": "object",
+              "description": "Filtering conditions for objects transferred."
+            },
+            "transferOptions": {
+              "type": "object",
+              "description": "Transfer options such as overwrite or delete."
+            },
+            "schedule": {
+              "type": "object",
+              "description": "Transfer schedule details."
+            }
+          }
+        }
+      ],
+      "type": "object"
+    }
+  },
+  "type": "object",
+  "properties": {
+    "cloudStorageTransferJob": {
+      "$ref": "#/$defs/CloudStorageTransferJobFacet"
+    }
+  }
+}

airflow/providers/google/cloud/openlineage/CloudStorageTransferRunFacet.json

@@ -0,0 +1,60 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$defs": {
+    "CloudStorageTransferRunFacet": {
+      "allOf": [
+        {
+          "$ref": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/RunFacet"
+        },
+        {
+          "type": "object",
+          "properties": {
+            "jobName": {
+              "type": "string",
+              "description": "Transfer job name associated with this run."
+            },
+            "operationName": {
+              "type": "string",
+              "description": "Transfer operation name if available."
+            },
+            "status": {
+              "type": "string",
+              "description": "Run status if available."
+            },
+            "startTime": {
+              "type": "string",
+              "description": "Start time of the transfer operation."
+            },
+            "endTime": {
+              "type": "string",
+              "description": "End time of the transfer operation."
+            },
+            "wait": {
+              "type": "boolean",
+              "description": "Whether the operator waited for completion."
+            },
+            "timeout": {
+              "type": ["number", "null"],
+              "description": "Timeout in seconds."
+            },
+            "deferrable": {
+              "type": "boolean",
+              "description": "Whether the operator used deferrable mode."
+            },
+            "deleteJobAfterCompletion": {
+              "type": "boolean",
+              "description": "Whether the transfer job was deleted after completion."
+            }
+          }
+        }
+      ],
+      "type": "object"
+    }
+  },
+  "type": "object",
+  "properties": {
+    "cloudStorageTransferRun": {
+      "$ref": "#/$defs/CloudStorageTransferRunFacet"
+    }
+  }
+}

airflow/providers/google/cloud/openlineage/DataFusionRunFacet.json

@@ -0,0 +1,32 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$defs": {
+    "DataFusionRunFacet": {
+      "allOf": [
+        {
+          "$ref": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/RunFacet"
+        },
+        {
+          "type": "object",
+          "properties": {
+            "runId": {
+              "type": "string",
+              "description": "Pipeline run ID assigned by Cloud Data Fusion."
+            },
+            "runtimeArgs": {
+              "type": "object",
+              "description": "Runtime arguments provided when starting the pipeline."
+            }
+          }
+        }
+      ],
+      "type": "object"
+    }
+  },
+  "type": "object",
+  "properties": {
+    "dataFusionRun": {
+      "$ref": "#/$defs/DataFusionRunFacet"
+    }
+  }
+}