Msudhir/add vector update functionality

sdk/python/docs/source/feast.protos.feast.core.rst

@@ -164,6 +164,22 @@ feast.protos.feast.core.FeatureView\_pb2\_grpc module
    :undoc-members:
    :show-inheritance:
 
+feast.protos.feast.core.VectorFeatureView\_pb2 module
+-----------------------------------------------
+
+.. automodule:: feast.protos.feast.core.VectorFeatureView_pb2
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+feast.protos.feast.core.VectorFeatureView\_pb2\_grpc module
+-----------------------------------------------------
+
+.. automodule:: feast.protos.feast.core.VectorFeatureView_pb2_grpc
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 feast.protos.feast.core.Feature\_pb2 module
 -------------------------------------------
 

sdk/python/feast/expediagroup/vectordb/milvus_online_store.py

@@ -1,14 +1,19 @@
+import logging
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
 from pydantic.typing import Literal
+from pymilvus import Collection, connections, utility
 
 from feast import Entity, RepoConfig
 from feast.expediagroup.vectordb.vector_feature_view import VectorFeatureView
 from feast.expediagroup.vectordb.vector_online_store import VectorOnlineStore
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
 from feast.repo_config import FeastConfigBaseModel
+from feast.usage import log_exceptions_and_usage
+
+logger = logging.getLogger(__name__)
 
 
 class MilvusOnlineStoreConfig(FeastConfigBaseModel):
@@ -17,13 +22,47 @@ class MilvusOnlineStoreConfig(FeastConfigBaseModel):
     type: Literal["milvus"] = "milvus"
     """Online store type selector"""
 
+    alias: str = "default"
+    """ alias for milvus connection"""
+
     host: str
     """ the host URL """
 
+    username: str
+    """ username to connect to Milvus """
+
+    password: str
+    """ password to connect to Milvus """
+
     port: int = 19530
     """ the port to connect to a Milvus instance. Should be the one used for GRPC (default: 19530) """
 
 
+class MilvusConnectionManager:
+    def __init__(self, online_config: RepoConfig):
+        self.online_config = online_config
+
+    def __enter__(self):
+        # Connecting to Milvus
+        logger.info(
+            f"Connecting to Milvus with alias {self.online_config.alias} and host {self.online_config.host} and default port {self.online_config.port}."
+        )
+        connections.connect(
+            host=self.online_config.host,
+            username=self.online_config.username,
+            password=self.online_config.password,
+            use_secure=True,
+        )
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # Disconnecting from Milvus
+        logger.info("Closing the connection to Milvus")
+        connections.disconnect(self.online_config.alias)
+        logger.info("Connection Closed")
+        if exc_type is not None:
+            logger.error(f"An exception of type {exc_type} occurred: {exc_value}")
+
+
 class MilvusOnlineStore(VectorOnlineStore):
     def online_write_batch(
         self,
@@ -49,6 +88,7 @@ def online_read(
             "to be implemented in https://jira.expedia.biz/browse/EAPC-7972"
         )
 
+    @log_exceptions_and_usage(online_store="milvus")
     def update(
         self,
         config: RepoConfig,
@@ -58,9 +98,34 @@ def update(
         entities_to_keep: Sequence[Entity],
         partial: bool,
     ):
-        raise NotImplementedError(
-            "to be implemented in https://jira.expedia.biz/browse/EAPC-7970"
-        )
+        with MilvusConnectionManager(config.online_store):
+            for table_to_keep in tables_to_keep:
+                collection_available = utility.has_collection(table_to_keep.name)
+                try:
+                    if collection_available:
+                        logger.error(f"Collection {table_to_keep.name} already exists.")
+                    else:
+                        Collection(name=table_to_keep.name, schema=table_to_keep.schema)
+                        logger.info(
+                            f"Collection {table_to_keep.name} has been created successfully."
+                        )
+                except Exception as e:
+                    logger.error(f"Collection update failed due to {e}")
+
+            for table_to_delete in tables_to_delete:
+                collection_available = utility.has_collection(table_to_delete.name)
+                try:
+                    if collection_available:
+                        utility.drop_collection(table_to_delete.name)
+                        logger.info(
+                            f"Collection {table_to_delete.name} has been deleted successfully."
+                        )
+                    else:
+                        return logger.error(
+                            f"Collection {table_to_delete.name} does not exist or is already deleted."
+                        )
+                except Exception as e:
+                    logger.error(f"Collection deletion failed due to {e}")
 
     def teardown(
         self,

sdk/python/feast/repo_config.py

@@ -203,6 +203,8 @@ def __init__(self, **data: Any):
                 self._offline_config = "redshift"
             elif data["provider"] == "azure":
                 self._offline_config = "mssql"
+            elif data["provider"] == "milvus":
+                self._online_config = "milvus"
 
         self._online_store = None
         if "online_store" in data:
@@ -216,6 +218,8 @@ def __init__(self, **data: Any):
                 self._online_config = "dynamodb"
             elif data["provider"] == "rockset":
                 self._online_config = "rockset"
+            elif data["provider"] == "milvus":
+                self._online_config = "milvus"
 
         self._batch_engine = None
         if "batch_engine" in data: