From 2c363baad9886204da4c64b61b9d2775e60b6ea4 Mon Sep 17 00:00:00 2001
From: Jen Hamon <jhamon@pinecone.io>
Date: Tue, 22 Oct 2024 14:46:15 -0400
Subject: [PATCH] Install plugins on Index and IndexGRPC

---
 pinecone/data/index.py | 35 ++++++++++++++++++++++++++++++-----
 pinecone/grpc/base.py  | 18 ++++++++++++++++++
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/pinecone/data/index.py b/pinecone/data/index.py
index cc1bae6c..e542b962 100644
--- a/pinecone/data/index.py
+++ b/pinecone/data/index.py
@@ -1,11 +1,11 @@
 from tqdm.autonotebook import tqdm
 
+import logging
 from typing import Union, List, Optional, Dict, Any
 
 from pinecone.config import ConfigBuilder
 
 from pinecone.core.openapi.shared import API_VERSION
-from pinecone.core.openapi.data.models import SparseValues
 from pinecone.core.openapi.data import ApiClient
 from pinecone.core.openapi.data.models import (
     FetchResponse,
@@ -22,12 +22,22 @@
     UpdateRequest,
     DescribeIndexStatsRequest,
     ListResponse,
+    SparseValues,
 )
-from .features.bulk_import import ImportFeatureMixin
 from pinecone.core.openapi.data.api.data_plane_api import DataPlaneApi
-from ..utils import setup_openapi_client, parse_non_empty_args
+from ..utils import (
+    setup_openapi_client,
+    parse_non_empty_args,
+    build_plugin_setup_client,
+    validate_and_convert_errors,
+)
+from .features.bulk_import import ImportFeatureMixin
 from .vector_factory import VectorFactory
 
+from pinecone_plugin_interface import load_and_install as install_plugins
+
+logger = logging.getLogger(__name__)
+
 __all__ = [
     "Index",
     "FetchResponse",
@@ -47,8 +57,6 @@
     "SparseValues",
 ]
 
-from ..utils.error_handling import validate_and_convert_errors
-
 _OPENAPI_ENDPOINT_PARAMS = (
     "_return_http_data_only",
     "_preload_content",
@@ -103,6 +111,23 @@ def __init__(
             api_version=API_VERSION,
         )
 
+        self._load_plugins()
+
+    def _load_plugins(self):
+        """@private"""
+        try:
+            # I don't expect this to ever throw, but wrapping this in a
+            # try block just in case to make sure a bad plugin doesn't
+            # halt client initialization.
+            openapi_client_builder = build_plugin_setup_client(
+                config=self.config,
+                openapi_config=self.openapi_config,
+                pool_threads=self.pool_threads,
+            )
+            install_plugins(self, openapi_client_builder)
+        except Exception as e:
+            logger.error(f"Error loading plugins in Index: {e}")
+
     def __enter__(self):
         return self
 
diff --git a/pinecone/grpc/base.py b/pinecone/grpc/base.py
index 17580d7e..187afd58 100644
--- a/pinecone/grpc/base.py
+++ b/pinecone/grpc/base.py
@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Optional
 
+import logging
 import grpc
 from grpc._channel import Channel
 
@@ -10,6 +11,10 @@
 from .config import GRPCClientConfig
 from .grpc_runner import GrpcRunner
 
+from pinecone_plugin_interface import load_and_install as install_plugins
+
+_logger = logging.getLogger(__name__)
+
 
 class GRPCIndexBase(ABC):
     """
@@ -40,6 +45,19 @@ def __init__(
         self._channel = channel or self._gen_channel()
         self.stub = self.stub_class(self._channel)
 
+        self._load_plugins()
+
+    def _load_plugins(self):
+        """@private"""
+        try:
+
+            def stub_openapi_client_builder(**kwargs):
+                pass
+
+            install_plugins(self, stub_openapi_client_builder)
+        except Exception as e:
+            _logger.error(f"Error loading plugins in GRPCIndex: {e}")
+
     @property
     @abstractmethod
     def stub_class(self):