[ce] repository control plane and data plane

pinecone/pinecone.py

@@ -16,6 +16,7 @@
 if TYPE_CHECKING:
     from pinecone.config import Config, OpenApiConfiguration
     from pinecone.db_data import _Index as Index, _IndexAsyncio as IndexAsyncio
+    from pinecone.repository_data import _Repository as Repository
     from pinecone.db_control.index_host_store import IndexHostStore
     from pinecone.core.openapi.db_control.api.manage_indexes_api import ManageIndexesApi
     from pinecone.db_control.types import CreateIndexForModelEmbedTypedDict, ConfigureIndexEmbed
@@ -42,6 +43,7 @@
         RestoreJobModel,
         RestoreJobList,
     )
+    from pinecone.repository_control.models import RepositoryModel, RepositoryList, DocumentSchema
 
 
 class Pinecone(PluginAware, LegacyPineconeDBControlInterface):
@@ -241,6 +243,9 @@ def __init__(
         self._db_control = None  # Lazy initialization
         """ :meta private: """
 
+        self._repository_control = None  # Lazy initialization
+        """ :meta private: """
+
         super().__init__()  # Initialize PluginAware
 
     @property
@@ -273,6 +278,21 @@ def db(self):
             )
         return self._db_control
 
+    @property
+    def repository_ctrl(self):
+        """
+        RepositoryControl is a namespace where an instance of the `pinecone.repository_control.RepositoryControl` class is lazily created and cached.
+        """
+        if self._repository_control is None:
+            from pinecone.repository_control.repository_control import RepositoryControl
+
+            self._repository_control = RepositoryControl(
+                config=self._config,
+                openapi_config=self._openapi_config,
+                pool_threads=self._pool_threads,
+            )
+        return self._repository_control
+
     @property
     def index_host_store(self) -> "IndexHostStore":
         """:meta private:"""
@@ -460,6 +480,26 @@ def list_restore_jobs(
     def describe_restore_job(self, *, job_id: str) -> "RestoreJobModel":
         return self.db.restore_job.describe(job_id=job_id)
 
+    def create_repository(
+        self,
+        name: str,
+        spec: Union[Dict, "ServerlessSpec"],
+        schema: Union[Dict, "DocumentSchema"],
+        timeout: Optional[int] = None,
+    ) -> "RepositoryModel":
+        return self.repository_ctrl.repository.create(
+            name=name, spec=spec, schema=schema, timeout=timeout
+        )
+
+    def describe_repository(self, name: str) -> "RepositoryModel":
+        return self.repository_ctrl.repository.describe(name=name)
+
+    def list_repositories(self) -> "RepositoryList":
+        return self.repository_ctrl.repository.list()
+
+    def delete_repository(self, name: str, timeout: Optional[int] = None):
+        return self.repository_ctrl.repository.delete(name=name, timeout=timeout)
+
     @staticmethod
     def from_texts(*args, **kwargs):
         """:meta private:"""
@@ -518,6 +558,34 @@ def IndexAsyncio(self, host: str, **kwargs) -> "IndexAsyncio":
             **kwargs,
         )
 
+    def Repository(self, name: str = "", host: str = "", **kwargs) -> "Repository":
+        from pinecone.repository_data import _Repository
+
+        if name == "" and host == "":
+            raise ValueError("Either name or host must be specified")
+
+        pt = kwargs.pop("pool_threads", None) or self._pool_threads
+        api_key = self._config.api_key
+        openapi_config = self._openapi_config
+
+        if host != "":
+            check_realistic_host(host)
+
+            # Use host url if it is provided
+            repository_host = normalize_host(host)
+        else:
+            # Otherwise, get host url from describe_repository using the repo name
+            repository_host = self.repository_ctrl.repository._get_host(name)
+
+        return _Repository(
+            host=repository_host,
+            api_key=api_key,
+            pool_threads=pt,
+            openapi_config=openapi_config,
+            source_tag=self.config.source_tag,
+            **kwargs,
+        )
+
 
 def check_realistic_host(host: str) -> None:
     """:meta private:

pinecone/repository_control/__init__.py

@@ -0,0 +1,18 @@
+from .models import *
+from .repository_control import RepositoryControl
+from pinecone.db_control.enums import *
+
+__all__ = [
+    # from pinecone.db_control.enums
+    "CloudProvider",
+    "AwsRegion",
+    "GcpRegion",
+    "AzureRegion",
+    # from .models
+    "ServerlessSpec",
+    "ServerlessSpecDefinition",
+    "RepositoryList",
+    "RepositoryModel",
+    # direct imports
+    "RepositoryControl",
+]

pinecone/repository_control/models/__init__.py

@@ -0,0 +1,14 @@
+from .document_schema import DocumentSchema
+from .repository_description import ServerlessSpecDefinition
+from .repository_list import RepositoryList
+from .repository_model import RepositoryModel
+from .serverless_spec import ServerlessSpec
+
+
+__all__ = [
+    "DocumentSchema",
+    "ServerlessSpec",
+    "ServerlessSpecDefinition",
+    "RepositoryList",
+    "RepositoryModel",
+]

pinecone/repository_control/models/document_schema.py

@@ -0,0 +1,24 @@
+from pinecone.core.openapi.repository_control.model.document_schema import (
+    DocumentSchema as OpenAPIDocumentSchema,
+)
+import json
+
+
+class DocumentSchema:
+    def __init__(self, schema: OpenAPIDocumentSchema):
+        self.schema = schema
+
+    def __str__(self):
+        return str(self.schema)
+
+    def __getattr__(self, attr):
+        return getattr(self.schema, attr)
+
+    def __getitem__(self, key):
+        return self.__getattr__(key)
+
+    def __repr__(self):
+        return json.dumps(self.to_dict(), indent=4)
+
+    def to_dict(self):
+        return self.schema.to_dict()

pinecone/repository_control/models/repository_description.py

@@ -0,0 +1,10 @@
+from typing import NamedTuple, Dict, Literal
+
+
+class ServerlessSpecDefinition(NamedTuple):
+    cloud: str
+    region: str
+
+
+ServerlessKey = Literal["serverless"]
+ServerlessSpec = Dict[ServerlessKey, ServerlessSpecDefinition]

pinecone/repository_control/models/repository_list.py

@@ -0,0 +1,34 @@
+import json
+from pinecone.core.openapi.repository_control.model.repository_list import (
+    RepositoryList as OpenAPIRepositoryList,
+)
+from .repository_model import RepositoryModel
+from typing import List
+
+
+class RepositoryList:
+    def __init__(self, repository_list: OpenAPIRepositoryList):
+        self.repository_list = repository_list
+        self.repositories = [RepositoryModel(i) for i in self.repository_list.repositories]
+        self.current = 0
+
+    def names(self) -> List[str]:
+        return [i.name for i in self.repositories]
+
+    def __getitem__(self, key):
+        return self.repositories[key]
+
+    def __len__(self):
+        return len(self.repositories)
+
+    def __iter__(self):
+        return iter(self.repositories)
+
+    def __str__(self):
+        return str(self.repositories)
+
+    def __repr__(self):
+        return json.dumps([i.to_dict() for i in self.repositories], indent=4)
+
+    def __getattr__(self, attr):
+        return getattr(self.repository_list, attr)

pinecone/repository_control/models/repository_model.py

@@ -0,0 +1,24 @@
+from pinecone.core.openapi.repository_control.model.repository_model import (
+    RepositoryModel as OpenAPIRepositoryModel,
+)
+import json
+
+
+class RepositoryModel:
+    def __init__(self, repository: OpenAPIRepositoryModel):
+        self.repository = repository
+
+    def __str__(self):
+        return str(self.repository)
+
+    def __getattr__(self, attr):
+        return getattr(self.repository, attr)
+
+    def __getitem__(self, key):
+        return self.__getattr__(key)
+
+    def __repr__(self):
+        return json.dumps(self.to_dict(), indent=4)
+
+    def to_dict(self):
+        return self.repository.to_dict()

pinecone/repository_control/models/serverless_spec.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+from typing import Union
+from enum import Enum
+
+from pinecone.db_control.enums import CloudProvider, AwsRegion, GcpRegion, AzureRegion
+
+
+@dataclass(frozen=True)
+class ServerlessSpec:
+    cloud: str
+    region: str
+
+    def __init__(
+        self,
+        cloud: Union[CloudProvider, str],
+        region: Union[AwsRegion, GcpRegion, AzureRegion, str],
+    ):
+        # Convert Enums to their string values if necessary
+        object.__setattr__(self, "cloud", cloud.value if isinstance(cloud, Enum) else str(cloud))
+        object.__setattr__(
+            self, "region", region.value if isinstance(region, Enum) else str(region)
+        )
+
+    def asdict(self):
+        return {"serverless": {"cloud": self.cloud, "region": self.region}}

pinecone/repository_control/repository_control.py

@@ -0,0 +1,60 @@
+import logging
+from typing import Optional, TYPE_CHECKING
+
+from pinecone.core.openapi.repository_control.api.manage_repositories_api import (
+    ManageRepositoriesApi,
+)
+from pinecone.openapi_support.api_client import ApiClient
+
+from pinecone.utils import setup_openapi_client, PluginAware
+from pinecone.core.openapi.repository_control import API_VERSION
+
+
+logger = logging.getLogger(__name__)
+""" :meta private: """
+
+if TYPE_CHECKING:
+    from .resources.sync.repository import RepositoryResource
+    from pinecone.config import Config, OpenApiConfiguration
+
+
+class RepositoryControl(PluginAware):
+    def __init__(
+        self, config: "Config", openapi_config: "OpenApiConfiguration", pool_threads: int
+    ) -> None:
+        self.config = config
+        """ :meta private: """
+
+        self._openapi_config = openapi_config
+        """ :meta private: """
+
+        self._pool_threads = pool_threads
+        """ :meta private: """
+
+        self._repository_api = setup_openapi_client(
+            api_client_klass=ApiClient,
+            api_klass=ManageRepositoriesApi,
+            config=self.config,
+            openapi_config=self._openapi_config,
+            pool_threads=self._pool_threads,
+            api_version=API_VERSION,
+        )
+        """ :meta private: """
+
+        self._repository_resource: Optional["RepositoryResource"] = None
+        """ :meta private: """
+
+        super().__init__()  # Initialize PluginAware
+
+    @property
+    def repository(self) -> "RepositoryResource":
+        if self._repository_resource is None:
+            from .resources.sync.repository import RepositoryResource
+
+            self._repository_resource = RepositoryResource(
+                repository_api=self._repository_api,
+                config=self.config,
+                openapi_config=self._openapi_config,
+                pool_threads=self._pool_threads,
+            )
+        return self._repository_resource

pinecone/repository_control/repository_host_store.py

@@ -0,0 +1,53 @@
+from typing import Dict
+from pinecone.config import Config
+from pinecone.core.openapi.repository_control.api.manage_repositories_api import (
+    ManageRepositoriesApi as RepositoriesOperationsApi,
+)
+from pinecone.openapi_support.exceptions import PineconeException
+from pinecone.utils import normalize_host
+
+
+class SingletonMeta(type):
+    _instances: Dict[str, str] = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            instance = super().__call__(*args, **kwargs)
+            cls._instances[cls] = instance
+        return cls._instances[cls]
+
+
+class RepositoryHostStore(metaclass=SingletonMeta):
+    _repositoryHosts: Dict[str, str]
+
+    def __init__(self) -> None:
+        self._repositoryHosts = {}
+
+    def _key(self, config: Config, repository_name: str) -> str:
+        return ":".join([config.api_key, repository_name])
+
+    def delete_host(self, config: Config, repository_name: str):
+        key = self._key(config, repository_name)
+        if key in self._repositoryHosts:
+            del self._repositoryHosts[key]
+
+    def key_exists(self, key: str) -> bool:
+        return key in self._repositoryHosts
+
+    def set_host(self, config: Config, repository_name: str, host: str):
+        if host:
+            key = self._key(config, repository_name)
+            self._repositoryHosts[key] = normalize_host(host)
+
+    def get_host(self, api: RepositoriesOperationsApi, config: Config, repository_name: str) -> str:
+        key = self._key(config, repository_name)
+        if self.key_exists(key):
+            return self._repositoryHosts[key]
+        else:
+            description = api.describe_repository(repository_name)
+            self.set_host(config, repository_name, description.host)
+            if not self.key_exists(key):
+                raise PineconeException(
+                    f"Could not get host for repository: {repository_name}. Call describe_repository('{repository_name}') to check the current status."
+                )
+            return self._repositoryHosts[key]