Add Google Vertex AI Feature Store - Feature View Sync Operators, Sen…

…sor (#44891)

docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst

@@ -644,6 +644,38 @@ The operator returns the cached content response in :ref:`XCom <concepts:xcom>`
     :start-after: [START how_to_cloud_vertex_ai_generate_from_cached_content_operator]
     :end-before: [END how_to_cloud_vertex_ai_generate_from_cached_content_operator]
 
+Interacting with Vertex AI Feature Store
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To get a feature view sync job you can use
+:class:`~airflow.providers.google.cloud.operators.vertex_ai.feature_store.GetFeatureViewSyncOperator`.
+The operator returns sync job results in :ref:`XCom <concepts:xcom>` under ``return_value`` key.
+
+.. exampleinclude:: /../../providers/tests/system/google/cloud/vertex_ai/example_vertex_ai_feature_store.py
+    :language: python
+    :dedent: 4
+    :start-after: [START how_to_cloud_vertex_ai_feature_store_get_feature_view_sync_operator]
+    :end-before: [END how_to_cloud_vertex_ai_feature_store_get_feature_view_sync_operator]
+
+To sync a feature view you can use
+:class:`~airflow.providers.google.cloud.operators.vertex_ai.feature_store.SyncFeatureViewOperator`.
+The operator returns the sync job name in :ref:`XCom <concepts:xcom>` under ``return_value`` key.
+
+.. exampleinclude:: /../../providers/tests/system/google/cloud/vertex_ai/example_vertex_ai_feature_store.py
+    :language: python
+    :dedent: 4
+    :start-after: [START how_to_cloud_vertex_ai_feature_store_sync_feature_view_operator]
+    :end-before: [END how_to_cloud_vertex_ai_feature_store_sync_feature_view_operator]
+
+To check if Feature View Sync succeeded you can use
+:class:`~airflow.providers.google.cloud.sensors.vertex_ai.FeatureViewSyncSensor`.
+
+.. exampleinclude:: /../../providers/tests/system/google/cloud/vertex_ai/example_vertex_ai_feature_store.py
+    :language: python
+    :dedent: 4
+    :start-after: [START how_to_cloud_vertex_ai_feature_store_feature_view_sync_sensor]
+    :end-before: [END how_to_cloud_vertex_ai_feature_store_feature_view_sync_sensor]
+
 Reference
 ^^^^^^^^^
 

generated/provider_dependencies.json

@@ -646,7 +646,7 @@
       "google-api-python-client>=2.0.2",
       "google-auth-httplib2>=0.0.1",
       "google-auth>=2.29.0",
-      "google-cloud-aiplatform>=1.70.0",
+      "google-cloud-aiplatform>=1.73.0",
       "google-cloud-automl>=2.12.0",
       "google-cloud-batch>=0.13.0",
       "google-cloud-bigquery-datatransfer>=3.13.0",

providers/src/airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py

@@ -0,0 +1,147 @@
+# 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 Google Cloud Vertex AI Feature Store hook."""
+
+from __future__ import annotations
+
+from google.api_core.client_options import ClientOptions
+from google.cloud.aiplatform_v1beta1 import (
+    FeatureOnlineStoreAdminServiceClient,
+)
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
+
+class FeatureStoreHook(GoogleBaseHook):
+    """
+    Hook for interacting with Google Cloud Vertex AI Feature Store.
+
+    This hook provides an interface to manage Feature Store resources in Vertex AI,
+    including feature views and their synchronization operations. It handles authentication
+    and provides methods for common Feature Store operations.
+
+    :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.
+    """
+
+    def get_feature_online_store_admin_service_client(
+        self,
+        location: str | None = None,
+    ) -> FeatureOnlineStoreAdminServiceClient:
+        """
+        Create and returns a FeatureOnlineStoreAdminServiceClient object.
+
+        This method initializes a client for interacting with the Feature Store API,
+        handling proper endpoint configuration based on the specified location.
+
+        :param location: Optional. The Google Cloud region where the service is located.
+            If provided and not 'global', the client will be configured to use the
+            region-specific API endpoint.
+        """
+        if location and location != "global":
+            client_options = ClientOptions(api_endpoint=f"{location}-aiplatform.googleapis.com:443")
+        else:
+            client_options = ClientOptions()
+        return FeatureOnlineStoreAdminServiceClient(
+            credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
+        )
+
+    def get_feature_view_sync(
+        self,
+        location: str,
+        feature_view_sync_name: str,
+    ) -> dict:
+        """
+        Retrieve the status and details of a Feature View synchronization operation.
+
+        This method fetches information about a specific feature view sync operation,
+        including its current status, timing information, and synchronization metrics.
+
+        :param location: The Google Cloud region where the feature store is located
+            (e.g., 'us-central1', 'us-east1').
+        :param feature_view_sync_name: The full resource name of the feature view
+            sync operation to retrieve.
+        """
+        client = self.get_feature_online_store_admin_service_client(location)
+
+        try:
+            response = client.get_feature_view_sync(name=feature_view_sync_name)
+
+            report = {
+                "name": feature_view_sync_name,
+                "start_time": int(response.run_time.start_time.seconds),
+            }
+
+            if hasattr(response.run_time, "end_time") and response.run_time.end_time.seconds:
+                report["end_time"] = int(response.run_time.end_time.seconds)
+                report["sync_summary"] = {
+                    "row_synced": int(response.sync_summary.row_synced),
+                    "total_slot": int(response.sync_summary.total_slot),
+                }
+
+            return report
+
+        except Exception as e:
+            self.log.error("Failed to get feature view sync: %s", str(e))
+            raise AirflowException(str(e))
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def sync_feature_view(
+        self,
+        location: str,
+        feature_online_store_id: str,
+        feature_view_id: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> str:
+        """
+        Initiate a synchronization operation for a Feature View.
+
+        This method 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 location: The Google Cloud region where the feature store is located
+            (e.g., 'us-central1', 'us-east1').
+        :param feature_online_store_id: The ID of the online feature store that
+            contains the feature view to be synchronized.
+        :param feature_view_id: The ID of the feature view to synchronize.
+        :param project_id: The ID of the Google Cloud project that contains the
+            feature store. If not provided, will attempt to determine from the
+            environment.
+        """
+        client = self.get_feature_online_store_admin_service_client(location)
+        feature_view = f"projects/{project_id}/locations/{location}/featureOnlineStores/{feature_online_store_id}/featureViews/{feature_view_id}"
+
+        try:
+            response = client.sync_feature_view(feature_view=feature_view)
+
+            return str(response.feature_view_sync)
+
+        except Exception as e:
+            self.log.error("Failed to sync feature view: %s", str(e))
+            raise AirflowException(str(e))

providers/src/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

providers/src/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.