review feedback: rename singleton, new indexOperationsBuilder, add PI…

…NECONE_CONTROLLER_HOST environment value and use over the default if provided, update unit tests
pinecone-io · Oct 30, 2023 · 2b21bb4 · 2b21bb4
1 parent aebf975
commit 2b21bb4
Show file tree

Hide file tree

Showing 12 changed files with 160 additions and 122 deletions.
diff --git a/.env.example b/.env.example
@@ -1,3 +1,3 @@
 # For testing purposes only
 PINECONE_API_KEY=""
-PINECONE_ENVIRONMENT=""
+PINECONE_CONTROLLER_HOST=""
diff --git a/src/__tests__/pinecone.test.ts b/src/__tests__/pinecone.test.ts
@@ -1,5 +1,5 @@
 import { Pinecone } from '../pinecone';
-import { HostUrlSingleton } from '../data/hostUrlSingleton';
+import { IndexHostSingleton } from '../data/indexHostSingleton';
 import type { PineconeConfiguration } from '../data';
 import * as utils from '../utils';
 
@@ -16,21 +16,18 @@ jest.mock('../utils', () => {
 });
 jest.mock('../data/fetch');
 jest.mock('../data/upsert');
-jest.mock('../data/hostUrlSingleton');
-
+jest.mock('../data/indexHostSingleton');
 jest.mock('../control', () => {
   const realControl = jest.requireActual('../control');
   return {
     ...realControl,
-    describeIndex: (_, callback) =>
-      jest.fn((indexName) =>
-        callback(
-          {
-            status: { ready: true, state: 'Ready', host: fakeHost },
-          },
-          indexName
-        )
-      ),
+    describeIndex: () =>
+      jest
+        .fn()
+        .mockImplementation(() =>
+          Promise.resolve({ status: { host: fakeHost } })
+        ),
+    deleteIndex: () => jest.fn().mockImplementation(() => Promise.resolve()),
   };
 });
 
@@ -71,12 +68,12 @@ describe('Pinecone', () => {
     });
 
     describe('optional properties', () => {
-      test('should not throw when optional properties provided: fetchAPI, hostUrl', () => {
+      test('should not throw when optional properties provided: fetchAPI, controllerHostUrl', () => {
         expect(() => {
           new Pinecone({
             apiKey: 'test-key',
             fetchApi: utils.getFetch({} as PineconeConfiguration),
-            hostUrl: 'https://foo-bar.io',
+            controllerHostUrl: 'https://foo-bar.io',
           } as PineconeConfiguration);
         }).not.toThrow();
       });
@@ -154,15 +151,25 @@ describe('Pinecone', () => {
   });
 
   describe('control plane operations', () => {
-    test('describeIndex callback triggers calling HostUrlSingleton._set', async () => {
+    test('describeIndex triggers calling IndexHostSingleton._set', async () => {
       const p = new Pinecone({ apiKey: 'foo' });
-      p.describeIndex('test-index');
+      await p.describeIndex('test-index');
 
-      expect(HostUrlSingleton._set).toHaveBeenCalledWith(
+      expect(IndexHostSingleton._set).toHaveBeenCalledWith(
         { apiKey: 'foo' },
         'test-index',
         fakeHost
       );
     });
+
+    test('deleteIndex trigger calling IndexHostSingleton._delete', async () => {
+      const p = new Pinecone({ apiKey: 'foo' });
+      await p.deleteIndex('test-index');
+
+      expect(IndexHostSingleton._delete).toHaveBeenCalledWith(
+        { apiKey: 'foo' },
+        'test-index'
+      );
+    });
   });
 });
diff --git a/src/control/describeIndex.ts b/src/control/describeIndex.ts
@@ -9,13 +9,7 @@ export type DescribeIndexOptions = IndexName;
 /** The description of your index returned from { @link Pinecone.describeIndex } */
 export type IndexDescription = IndexMeta;
 
-export const describeIndex = (
-  api: IndexOperationsApi,
-  callback?: (
-    descriptionResponse: IndexDescription,
-    indexName: IndexName
-  ) => void
-) => {
+export const describeIndex = (api: IndexOperationsApi) => {
   const validator = buildConfigValidator(IndexNameSchema, 'describeIndex');
 
   const removeDeprecatedFields = (result: any) => {
@@ -34,10 +28,6 @@ export const describeIndex = (
     const result = await api.describeIndex({ indexName: name });
     removeDeprecatedFields(result);
 
-    if (callback) {
-      callback(result, name);
-    }
-
     return result;
   };
 };
diff --git a/src/control/index.ts b/src/control/index.ts
@@ -1,4 +1,5 @@
 // Index Operations
+export { indexOperationsBuilder } from './indexOperationsBuilder';
 export type { IndexName, PodType } from './types';
 export { configureIndex } from './configureIndex';
 export type { ConfigureIndexOptions } from './configureIndex';

diff --git a/src/control/indexOperationsBuilder.ts b/src/control/indexOperationsBuilder.ts
@@ -0,0 +1,27 @@
+import {
+  IndexOperationsApi,
+  Configuration,
+} from '../pinecone-generated-ts-fetch';
+import { queryParamsStringify, buildUserAgent, getFetch } from '../utils';
+import { middleware } from '../utils/middleware';
+import type { PineconeConfiguration } from '../data/types';
+import type { ConfigurationParameters as IndexOperationsApiConfigurationParameters } from '../pinecone-generated-ts-fetch';
+
+export const indexOperationsBuilder = (
+  config: PineconeConfiguration
+): IndexOperationsApi => {
+  const { apiKey } = config;
+  const controllerPath = config.controllerHostUrl || 'https://api.pinecone.io';
+  const apiConfig: IndexOperationsApiConfigurationParameters = {
+    basePath: controllerPath,
+    apiKey,
+    queryParamsStringify,
+    headers: {
+      'User-Agent': buildUserAgent(false),
+    },
+    fetchApi: getFetch(config),
+    middleware,
+  };
+
+  return new IndexOperationsApi(new Configuration(apiConfig));
+};
diff --git a/src/data/__tests__/hostUrlSingleton.test.ts → ...data/__tests__/indexHostSingleton.test.ts b/src/data/__tests__/hostUrlSingleton.test.ts → ...data/__tests__/indexHostSingleton.test.ts
@@ -1,4 +1,4 @@
-import { HostUrlSingleton } from '../hostUrlSingleton';
+import { IndexHostSingleton } from '../indexHostSingleton';
 
 const mockDescribeIndex = jest.fn();
 
@@ -11,9 +11,9 @@ jest.mock('../../control', () => {
   };
 });
 
-describe('HostUrlSingleton', () => {
+describe('IndexHostSingleton', () => {
   afterEach(() => {
-    HostUrlSingleton._reset();
+    IndexHostSingleton._reset();
     mockDescribeIndex.mockReset();
   });
 
@@ -36,7 +36,7 @@ describe('HostUrlSingleton', () => {
       status: { ready: true, state: 'Ready', host: testHost },
     });
 
-    const hostUrl = await HostUrlSingleton.getHostUrl(
+    const hostUrl = await IndexHostSingleton.getHostUrl(
       pineconeConfig,
       testIndex
     );
@@ -78,15 +78,15 @@ describe('HostUrlSingleton', () => {
         status: { ready: true, state: 'Ready', host: testHost2 },
       });
 
-    const hostUrl = await HostUrlSingleton.getHostUrl(
+    const hostUrl = await IndexHostSingleton.getHostUrl(
       pineconeConfig,
       testIndex
     );
     expect(mockDescribeIndex).toHaveBeenCalledTimes(1);
     expect(hostUrl).toEqual(`https://${testHost}`);
 
     // call again for same indexName, no additional calls
-    const hostUrl2 = await HostUrlSingleton.getHostUrl(
+    const hostUrl2 = await IndexHostSingleton.getHostUrl(
       pineconeConfig,
       testIndex
     );
@@ -95,11 +95,46 @@ describe('HostUrlSingleton', () => {
 
     // new indexName means we call describeIndex again
     // to resolve the new host
-    const hostUrl3 = await HostUrlSingleton.getHostUrl(
+    const hostUrl3 = await IndexHostSingleton.getHostUrl(
       pineconeConfig,
       testIndex2
     );
     expect(mockDescribeIndex).toHaveBeenCalledTimes(2);
     expect(hostUrl3).toEqual(`https://${testHost2}`);
   });
+
+  test('_set, _delete, and _reset work as expected', async () => {
+    const pineconeConfig = { apiKey: 'test-key' };
+
+    mockDescribeIndex.mockResolvedValue({
+      database: {
+        name: 'index-1',
+        dimensions: 10,
+        metric: 'cosine',
+        pods: 1,
+        replicas: 1,
+        shards: 1,
+        podType: 'p1.x1',
+      },
+      status: { ready: true, state: 'Ready', host: 'test-host' },
+    });
+
+    // _set test
+    IndexHostSingleton._set(pineconeConfig, 'index-1', 'test-host');
+    const host1 = await IndexHostSingleton.getHostUrl(
+      pineconeConfig,
+      'index-1'
+    );
+    expect(mockDescribeIndex).toHaveBeenCalledTimes(0);
+    expect(host1).toEqual('https://test-host');
+
+    // _delete test
+    IndexHostSingleton._delete(pineconeConfig, 'index-1');
+    const host2 = await IndexHostSingleton.getHostUrl(
+      pineconeConfig,
+      'index-1'
+    );
+    expect(mockDescribeIndex).toHaveBeenCalledTimes(1);
+    expect(host2).toBe('https://test-host');
+  });
 });
diff --git a/src/data/__tests__/vectorOperationsProvider.test.ts b/src/data/__tests__/vectorOperationsProvider.test.ts
@@ -1,42 +1,42 @@
 import { VectorOperationsProvider } from '../vectorOperationsProvider';
-import { HostUrlSingleton } from '../hostUrlSingleton';
+import { IndexHostSingleton } from '../indexHostSingleton';
 
 describe('VectorOperationsProvider', () => {
   let real;
 
   beforeAll(() => {
-    real = HostUrlSingleton.getHostUrl;
+    real = IndexHostSingleton.getHostUrl;
   });
   afterAll(() => {
-    HostUrlSingleton.getHostUrl = real;
+    IndexHostSingleton.getHostUrl = real;
   });
   beforeEach(() => {
-    HostUrlSingleton.getHostUrl = jest.fn();
+    IndexHostSingleton.getHostUrl = jest.fn();
   });
   afterEach(() => {
-    HostUrlSingleton._reset();
+    IndexHostSingleton._reset();
   });
 
   test('makes no API calls on instantiation', async () => {
     const config = {
       apiKey: 'test-api-key',
     };
     new VectorOperationsProvider(config, 'index-name');
-    expect(HostUrlSingleton.getHostUrl).not.toHaveBeenCalled();
+    expect(IndexHostSingleton.getHostUrl).not.toHaveBeenCalled();
   });
 
   test('api calls occur only the first time the provide method is called', async () => {
     const config = {
       apiKey: 'test-api-key',
     };
     const provider = new VectorOperationsProvider(config, 'index-name');
-    expect(HostUrlSingleton.getHostUrl).not.toHaveBeenCalled();
+    expect(IndexHostSingleton.getHostUrl).not.toHaveBeenCalled();
 
     const api = await provider.provide();
-    expect(HostUrlSingleton.getHostUrl).toHaveBeenCalled();
+    expect(IndexHostSingleton.getHostUrl).toHaveBeenCalled();
 
     const api2 = await provider.provide();
-    expect(HostUrlSingleton.getHostUrl).toHaveBeenCalledTimes(1);
+    expect(IndexHostSingleton.getHostUrl).toHaveBeenCalledTimes(1);
     expect(api).toEqual(api2);
   });
 });
diff --git a/src/data/hostUrlSingleton.ts → src/data/indexHostSingleton.ts b/src/data/hostUrlSingleton.ts → src/data/indexHostSingleton.ts
@@ -1,44 +1,21 @@
-import type { ConfigurationParameters as IndexOperationsApiConfigurationParameters } from '../pinecone-generated-ts-fetch';
-import {
-  IndexOperationsApi,
-  Configuration as ApiConfiguration,
-} from '../pinecone-generated-ts-fetch';
+import { IndexOperationsApi } from '../pinecone-generated-ts-fetch';
 import type { PineconeConfiguration } from './types';
-import { middleware } from '../utils/middleware';
-import { queryParamsStringify, buildUserAgent, getFetch } from '../utils';
-import { describeIndex } from '../control';
+import { describeIndex, indexOperationsBuilder } from '../control';
 import { PineconeUnableToResolveHostError } from '../errors';
 
 // We use describeIndex to retrieve the data plane url (host) for a given API key
 // and index. We only ever want to call describeIndex a maximum of once per API key
-// and index so we cache them in a singleton for reuse.
-export const HostUrlSingleton = (function () {
+// and index, so we cache them in a singleton for reuse.
+export const IndexHostSingleton = (function () {
   const hostUrls = {}; // map of apiKey-indexName to hostUrl
   let indexOperationsApi: InstanceType<typeof IndexOperationsApi> | null = null;
 
-  const _buildIndexOperationsApi = (config: PineconeConfiguration) => {
-    const { apiKey } = config;
-    const controllerPath = `https://api.pinecone.io`;
-    const apiConfig: IndexOperationsApiConfigurationParameters = {
-      basePath: controllerPath,
-      apiKey,
-      queryParamsStringify,
-      headers: {
-        'User-Agent': buildUserAgent(false),
-      },
-      fetchApi: getFetch(config),
-      middleware,
-    };
-
-    return new IndexOperationsApi(new ApiConfiguration(apiConfig));
-  };
-
   const _describeIndex = async (
     config: PineconeConfiguration,
     indexName: string
   ): Promise<string> => {
     if (!indexOperationsApi) {
-      indexOperationsApi = _buildIndexOperationsApi(config);
+      indexOperationsApi = indexOperationsBuilder(config);
     }
 
     const describeResponse = await describeIndex(indexOperationsApi)(indexName);
@@ -85,5 +62,10 @@ export const HostUrlSingleton = (function () {
       const cacheKey = key(config, indexName);
       hostUrls[cacheKey] = hostUrl;
     },
+
+    _delete: (config, indexName) => {
+      const cacheKey = key(config, indexName);
+      delete hostUrls[cacheKey];
+    },
   };
 })();
diff --git a/src/data/types.ts b/src/data/types.ts
@@ -4,7 +4,7 @@ import type { FetchAPI } from '../pinecone-generated-ts-fetch';
 export const PineconeConfigurationSchema = Type.Object(
   {
     apiKey: Type.String({ minLength: 1 }),
-    hostUrl: Type.Optional(Type.String({ minLength: 1 })),
+    controllerHostUrl: Type.Optional(Type.String({ minLength: 1 })),
 
     // fetchApi is a complex type that I don't really want to recreate in the
     // form of a json schema (seems difficult and error prone). So we will
@@ -26,16 +26,24 @@ export type PineconeConfiguration = {
   apiKey: string;
 
   /**
-   * The host URL for the Index.
+   * Optional configuration field for specifying the controller host. If not specified, the client will use the default controller host: https://api.pinecone.io.
    */
-  hostUrl?: string;
+  controllerHostUrl?: string;
 
   /**
    * Optional configuration field for specifying the fetch implementation. If not specified, the client will look for fetch in the global scope and if none is found it will fall back to a [cross-fetch](https://www.npmjs.com/package/cross-fetch) polyfill.
    */
   fetchApi?: FetchAPI;
 };
 
+/** Configuration for a single Pinecone Index */
+export type IndexConfiguration = PineconeConfiguration & {
+  /**
+   * The host URL for the Index.
+   */
+  indexHostUrl?: string;
+};
+
 export const RecordIdSchema = Type.String({ minLength: 1 });
 export const RecordValuesSchema = Type.Array(Type.Number());
 export const RecordSparseValuesSchema = Type.Object(