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

airflow/providers/google/cloud/operators/vertex_ai/feature_store.py

@@ -0,0 +1,163 @@
+#
+# 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 Vertex AI Feature Store operators."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.google.cloud.hooks.vertex_ai.feature_store import FeatureStoreHook
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class SyncFeatureViewOperator(GoogleCloudBaseOperator):
+    """
+    Initiate a synchronization operation for a Feature View in Vertex AI Feature Store.
+
+    This operator triggers a sync operation that updates the online serving data for a feature view
+    based on the latest data in the underlying batch source. The sync operation ensures that
+    the online feature values are up-to-date for real-time serving.
+
+    :param project_id: Required. The ID of the Google Cloud project that contains the feature store.
+        This is used to identify which project's resources to interact with.
+    :param location: Required. The location of the feature store (e.g., 'us-central1', 'us-east1').
+        This specifies the Google Cloud region where the feature store resources are located.
+    :param feature_online_store_id: Required. The ID of the online feature store that contains
+        the feature view to be synchronized. This store serves as the online serving layer.
+    :param feature_view_id: Required. The ID of the feature view to synchronize. This identifies
+        the specific view that needs to have its online values updated from the batch source.
+    :param gcp_conn_id: The connection ID to use for connecting to Google Cloud Platform.
+        Defaults to 'google_cloud_default'.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials. Can be either a single account or a chain 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.
+    """
+
+    template_fields: Sequence[str] = (
+        "project_id",
+        "location",
+        "feature_online_store_id",
+        "feature_view_id",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        feature_online_store_id: str,
+        feature_view_id: str,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.feature_online_store_id = feature_online_store_id
+        self.feature_view_id = feature_view_id
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> str:
+        """Execute the feature view sync operation."""
+        self.hook = FeatureStoreHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        self.log.info("Submitting Feature View sync job now...")
+        response = self.hook.sync_feature_view(
+            project_id=self.project_id,
+            location=self.location,
+            feature_online_store_id=self.feature_online_store_id,
+            feature_view_id=self.feature_view_id,
+        )
+        self.log.info("Retrieved Feature View sync: %s", response)
+
+        return response
+
+
+class GetFeatureViewSyncOperator(GoogleCloudBaseOperator):
+    """
+    Retrieve the status and details of a Feature View synchronization operation.
+
+    This operator fetches information about a specific feature view sync operation,
+    including its current status, timing information, and synchronization metrics.
+    It's typically used to monitor the progress of a sync operation initiated by
+    the SyncFeatureViewOperator.
+
+    :param location: Required. The location of the feature store (e.g., 'us-central1', 'us-east1').
+        This specifies the Google Cloud region where the feature store resources are located.
+    :param feature_view_sync_name: Required. The full resource name of the feature view
+        sync operation to retrieve. This is typically the return value from a
+        SyncFeatureViewOperator execution.
+    :param gcp_conn_id: The connection ID to use for connecting to Google Cloud Platform.
+        Defaults to 'google_cloud_default'.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials. Can be either a single account or a chain 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.
+    """
+
+    template_fields: Sequence[str] = (
+        "location",
+        "feature_view_sync_name",
+    )
+
+    def __init__(
+        self,
+        *,
+        location: str,
+        feature_view_sync_name: str,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.location = location
+        self.feature_view_sync_name = feature_view_sync_name
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """Execute the get feature view sync operation."""
+        self.hook = FeatureStoreHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        self.log.info("Retrieving Feature View sync job now...")
+        response = self.hook.get_feature_view_sync(
+            location=self.location, feature_view_sync_name=self.feature_view_sync_name
+        )
+        self.log.info("Retrieved Feature View sync: %s", self.feature_view_sync_name)
+        self.log.info(response)
+
+        return response

airflow/providers/google/cloud/operators/vertex_ai/generative_model.py

@@ -353,6 +353,11 @@ class PromptMultimodalModelWithMediaOperator(GoogleCloudBaseOperator):
         return response
 
 
+@deprecated(
+    planned_removal_date="April 09, 2025",
+    use_instead="GenerativeModelGenerateContentOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class TextGenerationModelPredictOperator(GoogleCloudBaseOperator):
     """
     Uses the Vertex AI PaLM API to generate natural language text.

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

@@ -107,10 +107,10 @@ class DataprocJobSensor(BaseSensorOperator):
         }:
             message = f"Job was cancelled:\n{job}"
             raise AirflowException(message)
-        elif JobStatus.State.DONE == state:
+        elif state == JobStatus.State.DONE:
             self.log.debug("Job %s completed successfully.", self.dataproc_job_id)
             return True
-        elif JobStatus.State.ATTEMPT_FAILURE == state:
+        elif state == JobStatus.State.ATTEMPT_FAILURE:
             self.log.debug("Job %s attempt has failed.", self.dataproc_job_id)
 
         self.log.info("Waiting for job %s to complete.", self.dataproc_job_id)

airflow/providers/google/cloud/sensors/vertex_ai/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py

@@ -0,0 +1,112 @@
+# 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 a Vertex AI Feature Store sensor."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.cloud.hooks.vertex_ai.feature_store import FeatureStoreHook
+from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class FeatureViewSyncSensor(BaseSensorOperator):
+    """
+    Sensor to monitor the state of a Vertex AI Feature View sync operation.
+
+    :param feature_view_sync_name: The name of the feature view sync operation to monitor. (templated)
+    :param location: Required. The Cloud region in which to handle the request. (templated)
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud Platform.
+    :param wait_timeout: How many seconds to wait for sync to complete.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials.
+    """
+
+    template_fields: Sequence[str] = ("location", "feature_view_sync_name")
+    ui_color = "#f0eee4"
+
+    def __init__(
+        self,
+        *,
+        feature_view_sync_name: str,
+        location: str,
+        gcp_conn_id: str = "google_cloud_default",
+        wait_timeout: int | None = None,
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.feature_view_sync_name = feature_view_sync_name
+        self.location = location
+        self.gcp_conn_id = gcp_conn_id
+        self.wait_timeout = wait_timeout
+        self.impersonation_chain = impersonation_chain
+        self.start_sensor_time: float | None = None
+
+    def execute(self, context: Context) -> None:
+        self.start_sensor_time = time.monotonic()
+        super().execute(context)
+
+    def _duration(self):
+        return time.monotonic() - self.start_sensor_time
+
+    def poke(self, context: Context) -> bool:
+        hook = FeatureStoreHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        try:
+            response = hook.get_feature_view_sync(
+                location=self.location,
+                feature_view_sync_name=self.feature_view_sync_name,
+            )
+
+            # Check if the sync has completed by verifying end_time exists
+            if response.get("end_time", 0) > 0:
+                self.log.info(
+                    "Feature View sync %s completed. Rows synced: %d, Total slots: %d",
+                    self.feature_view_sync_name,
+                    int(response.get("sync_summary", "").get("row_synced", "")),
+                    int(response.get("sync_summary", "").get("total_slot", "")),
+                )
+                return True
+
+            if self.wait_timeout and self._duration() > self.wait_timeout:
+                raise AirflowException(
+                    f"Timeout: Feature View sync {self.feature_view_sync_name} "
+                    f"not completed after {self.wait_timeout}s"
+                )
+
+            self.log.info("Waiting for Feature View sync %s to complete.", self.feature_view_sync_name)
+            return False
+
+        except Exception as e:
+            if self.wait_timeout and self._duration() > self.wait_timeout:
+                raise AirflowException(
+                    f"Timeout: Feature View sync {self.feature_view_sync_name} "
+                    f"not completed after {self.wait_timeout}s"
+                )
+            self.log.info("Error checking sync status, will retry: %s", str(e))
+            return False

airflow/providers/google/cloud/transfers/bigquery_to_gcs.py

@@ -292,8 +292,6 @@ class BigQueryToGCSOperator(BaseOperator):
 
     def get_openlineage_facets_on_complete(self, task_instance):
         """Implement on_complete as we will include final BQ job id."""
-        from pathlib import Path
-
         from airflow.providers.common.compat.openlineage.facet import (
             BaseFacet,
             Dataset,
@@ -303,6 +301,8 @@ class BigQueryToGCSOperator(BaseOperator):
         )
         from airflow.providers.google.cloud.hooks.gcs import _parse_gcs_url
         from airflow.providers.google.cloud.openlineage.utils import (
+            WILDCARD,
+            extract_ds_name_from_gcs_path,
             get_facets_from_bq_table,
             get_identity_column_lineage_facet,
         )
@@ -333,24 +333,19 @@ class BigQueryToGCSOperator(BaseOperator):
         output_datasets = []
         for uri in sorted(self.destination_cloud_storage_uris):
             bucket, blob = _parse_gcs_url(uri)
-            additional_facets = {}
 
-            if "*" in blob:
-                # If wildcard ("*") is used in gcs path, we want the name of dataset to be directory name,
-                # but we create a symlink to the full object path with wildcard.
+            additional_facets = {}
+            if WILDCARD in blob:
+                # For path with wildcard we attach a symlink with unmodified path.
                 additional_facets = {
                     "symlink": SymlinksDatasetFacet(
                         identifiers=[Identifier(namespace=f"gs://{bucket}", name=blob, type="file")]
                     ),
                 }
-                blob = Path(blob).parent.as_posix()
-                if blob == ".":
-                    # blob path does not have leading slash, but we need root dataset name to be "/"
-                    blob = "/"
 
             dataset = Dataset(
                 namespace=f"gs://{bucket}",
-                name=blob,
+                name=extract_ds_name_from_gcs_path(blob),
                 facets=merge_dicts(output_dataset_facets, additional_facets),
             )
             output_datasets.append(dataset)

airflow/providers/google/cloud/transfers/bigquery_to_mssql.py

@@ -45,6 +45,9 @@ class BigQueryToMsSqlOperator(BigQueryToSqlBaseOperator):
     :param mssql_table: target MsSQL table. It is deprecated: use target_table_name instead. (templated)
     :param target_table_name: target MsSQL table. It takes precedence over mssql_table. (templated)
     :param mssql_conn_id: reference to a specific mssql hook
+
+    .. warning::
+        The `mssql_table` parameter has been deprecated. Use `target_table_name` instead.
     """
 
     template_fields: Sequence[str] = (

airflow/providers/google/cloud/transfers/bigquery_to_mysql.py

@@ -39,6 +39,9 @@ class BigQueryToMySqlOperator(BigQueryToSqlBaseOperator):
         specific database. It is deprecated: use target_table_name instead. (templated)
     :param target_table_name: target MySQL table. It takes precedence over mysql_table. (templated)
     :param mysql_conn_id: Reference to :ref:`mysql connection id <howto/connection:mysql>`.
+
+    .. warning::
+        The `mysql_table` parameter has been deprecated. Use `target_table_name` instead.
     """
 
     template_fields: Sequence[str] = (*BigQueryToSqlBaseOperator.template_fields, "dataset_id", "table_id")

airflow/providers/google/cloud/transfers/gcs_to_bigquery.py

@@ -756,8 +756,6 @@ class GCSToBigQueryOperator(BaseOperator):
 
     def get_openlineage_facets_on_complete(self, task_instance):
         """Implement on_complete as we will include final BQ job id."""
-        from pathlib import Path
-
         from airflow.providers.common.compat.openlineage.facet import (
             Dataset,
             ExternalQueryRunFacet,
@@ -765,6 +763,8 @@ class GCSToBigQueryOperator(BaseOperator):
             SymlinksDatasetFacet,
         )
         from airflow.providers.google.cloud.openlineage.utils import (
+            WILDCARD,
+            extract_ds_name_from_gcs_path,
             get_facets_from_bq_table,
             get_identity_column_lineage_facet,
         )
@@ -793,22 +793,17 @@ class GCSToBigQueryOperator(BaseOperator):
         for blob in sorted(source_objects):
             additional_facets = {}
 
-            if "*" in blob:
-                # If wildcard ("*") is used in gcs path, we want the name of dataset to be directory name,
-                # but we create a symlink to the full object path with wildcard.
+            if WILDCARD in blob:
+                # For path with wildcard we attach a symlink with unmodified path.
                 additional_facets = {
                     "symlink": SymlinksDatasetFacet(
                         identifiers=[Identifier(namespace=f"gs://{self.bucket}", name=blob, type="file")]
                     ),
                 }
-                blob = Path(blob).parent.as_posix()
-                if blob == ".":
-                    # blob path does not have leading slash, but we need root dataset name to be "/"
-                    blob = "/"
 
             dataset = Dataset(
                 namespace=f"gs://{self.bucket}",
-                name=blob,
+                name=extract_ds_name_from_gcs_path(blob),
                 facets=merge_dicts(input_dataset_facets, additional_facets),
             )
             input_datasets.append(dataset)

airflow/providers/google/cloud/transfers/gcs_to_gcs.py

@@ -551,28 +551,16 @@ class GCSToGCSOperator(BaseOperator):
         This means we won't have to normalize self.source_object and self.source_objects,
         destination bucket and so on.
         """
-        from pathlib import Path
-
         from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.openlineage.utils import extract_ds_name_from_gcs_path
         from airflow.providers.openlineage.extractors import OperatorLineage
 
-        def _process_prefix(pref):
-            if WILDCARD in pref:
-                pref = pref.split(WILDCARD)[0]
-            # Use parent if not a file (dot not in name) and not a dir (ends with slash)
-            if "." not in pref.split("/")[-1] and not pref.endswith("/"):
-                pref = Path(pref).parent.as_posix()
-            return ["/" if pref in ("", "/", ".") else pref.rstrip("/")]  # Adjust root path
-
-        inputs = []
-        for prefix in self.source_objects:
-            result = _process_prefix(prefix)
-            inputs.extend(result)
+        inputs = [extract_ds_name_from_gcs_path(path) for path in self.source_objects]
 
         if self.destination_object is None:
             outputs = inputs.copy()
         else:
-            outputs = _process_prefix(self.destination_object)
+            outputs = [extract_ds_name_from_gcs_path(self.destination_object)]
 
         return OperatorLineage(
             inputs=[

airflow/providers/google/cloud/transfers/gcs_to_local.py

@@ -113,3 +113,12 @@ class GCSToLocalFilesystemOperator(BaseOperator):
                 raise AirflowException("The size of the downloaded file is too large to push to XCom!")
         else:
             hook.download(bucket_name=self.bucket, object_name=self.object_name, filename=self.filename)
+
+    def get_openlineage_facets_on_start(self):
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        return OperatorLineage(
+            inputs=[Dataset(namespace=f"gs://{self.bucket}", name=self.object_name)],
+            outputs=[Dataset(namespace="file", name=self.filename)] if self.filename else [],
+        )

airflow/providers/google/cloud/transfers/local_to_gcs.py

@@ -69,12 +69,12 @@ class LocalFilesystemToGCSOperator(BaseOperator):
     def __init__(
         self,
         *,
-        src,
-        dst,
-        bucket,
-        gcp_conn_id="google_cloud_default",
-        mime_type="application/octet-stream",
-        gzip=False,
+        src: str | list[str],
+        dst: str,
+        bucket: str,
+        gcp_conn_id: str = "google_cloud_default",
+        mime_type: str = "application/octet-stream",
+        gzip: bool = False,
         chunk_size: int | None = None,
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
@@ -120,3 +120,38 @@ class LocalFilesystemToGCSOperator(BaseOperator):
                 gzip=self.gzip,
                 chunk_size=self.chunk_size,
             )
+
+    def get_openlineage_facets_on_start(self):
+        from airflow.providers.common.compat.openlineage.facet import (
+            Dataset,
+            Identifier,
+            SymlinksDatasetFacet,
+        )
+        from airflow.providers.google.cloud.openlineage.utils import WILDCARD, extract_ds_name_from_gcs_path
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        source_facets = {}
+        if isinstance(self.src, str):  # Single path provided, possibly relative or with wildcard
+            original_src = f"{self.src}"
+            absolute_src = os.path.abspath(self.src)
+            resolved_src = extract_ds_name_from_gcs_path(absolute_src)
+            if original_src.startswith("/") and not resolved_src.startswith("/"):
+                resolved_src = "/" + resolved_src
+            source_objects = [resolved_src]
+
+            if WILDCARD in original_src or absolute_src != resolved_src:
+                # We attach a symlink with unmodified path.
+                source_facets = {
+                    "symlink": SymlinksDatasetFacet(
+                        identifiers=[Identifier(namespace="file", name=original_src, type="file")]
+                    ),
+                }
+        else:
+            source_objects = self.src
+
+        dest_object = self.dst if os.path.basename(self.dst) else extract_ds_name_from_gcs_path(self.dst)
+
+        return OperatorLineage(
+            inputs=[Dataset(namespace="file", name=src, facets=source_facets) for src in source_objects],
+            outputs=[Dataset(namespace=f"gs://{self.bucket}", name=dest_object)],
+        )

airflow/providers/google/cloud/transfers/s3_to_gcs.py

@@ -345,3 +345,18 @@ class S3ToGCSOperator(S3ListOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.google_impersonation_chain,
         )
+
+    def get_openlineage_facets_on_start(self):
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.hooks.gcs import _parse_gcs_url
+        from airflow.providers.google.cloud.openlineage.utils import extract_ds_name_from_gcs_path
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        gcs_bucket, gcs_blob = _parse_gcs_url(self.dest_gcs)
+        if not self.apply_gcs_prefix:
+            gcs_blob += self.prefix
+
+        return OperatorLineage(
+            inputs=[Dataset(namespace=f"s3://{self.bucket}", name=self.prefix.strip("/") or "/")],
+            outputs=[Dataset(namespace=f"gs://{gcs_bucket}", name=extract_ds_name_from_gcs_path(gcs_blob))],
+        )