From 5eaffe33f5e1b30a37b6a0c8eb3467e6af42b29e Mon Sep 17 00:00:00 2001 From: Austin DeNoble Date: Mon, 30 Oct 2023 16:06:49 -0400 Subject: [PATCH] review feedback: rename singleton, new indexOperationsBuilder, add PINECONE_CONTROLLER_HOST environment value and use over the default if provided, update unit tests --- .env.example | 2 +- src/__tests__/pinecone.test.ts | 41 +++++++----- src/control/describeIndex.ts | 12 +--- src/control/index.ts | 1 + src/control/indexOperationsBuilder.ts | 27 ++++++++ ...ton.test.ts => indexHostSingleton.test.ts} | 49 ++++++++++++-- .../vectorOperationsProvider.test.ts | 18 ++--- ...tUrlSingleton.ts => indexHostSingleton.ts} | 38 +++-------- src/data/types.ts | 14 +++- src/data/vectorOperationsProvider.ts | 12 ++-- src/pinecone.ts | 66 ++++++++----------- utils/cleanupResources.ts | 2 +- 12 files changed, 160 insertions(+), 122 deletions(-) create mode 100644 src/control/indexOperationsBuilder.ts rename src/data/__tests__/{hostUrlSingleton.test.ts => indexHostSingleton.test.ts} (65%) rename src/data/{hostUrlSingleton.ts => indexHostSingleton.ts} (65%) diff --git a/.env.example b/.env.example index 86c06146..df67b19c 100644 --- a/.env.example +++ b/.env.example @@ -1,3 +1,3 @@ # For testing purposes only PINECONE_API_KEY="" -PINECONE_ENVIRONMENT="" \ No newline at end of file +PINECONE_CONTROLLER_HOST="" \ No newline at end of file diff --git a/src/__tests__/pinecone.test.ts b/src/__tests__/pinecone.test.ts index 07a07149..472e36cb 100644 --- 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 index 4afc1d33..d2b61220 100644 --- 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 index 60dec0ee..95b9e697 100644 --- 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 new file mode 100644 index 00000000..8e8f2dc5 --- /dev/null +++ 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 b/src/data/__tests__/indexHostSingleton.test.ts similarity index 65% rename from src/data/__tests__/hostUrlSingleton.test.ts rename to src/data/__tests__/indexHostSingleton.test.ts index 3911c356..deb0f2b5 100644 --- a/src/data/__tests__/hostUrlSingleton.test.ts +++ b/src/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,7 +78,7 @@ describe('HostUrlSingleton', () => { status: { ready: true, state: 'Ready', host: testHost2 }, }); - const hostUrl = await HostUrlSingleton.getHostUrl( + const hostUrl = await IndexHostSingleton.getHostUrl( pineconeConfig, testIndex ); @@ -86,7 +86,7 @@ describe('HostUrlSingleton', () => { 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 index 2e7d8e8c..99e533e9 100644 --- a/src/data/__tests__/vectorOperationsProvider.test.ts +++ b/src/data/__tests__/vectorOperationsProvider.test.ts @@ -1,20 +1,20 @@ 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 () => { @@ -22,7 +22,7 @@ describe('VectorOperationsProvider', () => { 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 () => { @@ -30,13 +30,13 @@ describe('VectorOperationsProvider', () => { 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 b/src/data/indexHostSingleton.ts similarity index 65% rename from src/data/hostUrlSingleton.ts rename to src/data/indexHostSingleton.ts index 49230dff..e36997eb 100644 --- a/src/data/hostUrlSingleton.ts +++ b/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 | 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 => { 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 index ec2ddc4c..56aa6d98 100644 --- 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,9 +26,9 @@ 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. @@ -36,6 +36,14 @@ export type PineconeConfiguration = { 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( diff --git a/src/data/vectorOperationsProvider.ts b/src/data/vectorOperationsProvider.ts index 1b911f42..75b3c515 100644 --- a/src/data/vectorOperationsProvider.ts +++ b/src/data/vectorOperationsProvider.ts @@ -1,19 +1,19 @@ -import type { PineconeConfiguration } from './types'; +import type { IndexConfiguration } from './types'; import { Configuration, ConfigurationParameters, VectorOperationsApi, } from '../pinecone-generated-ts-fetch'; import { queryParamsStringify, buildUserAgent, getFetch } from '../utils'; -import { HostUrlSingleton } from './hostUrlSingleton'; +import { IndexHostSingleton } from './indexHostSingleton'; import { middleware } from '../utils/middleware'; export class VectorOperationsProvider { - private config: PineconeConfiguration; + private config: IndexConfiguration; private indexName: string; private vectorOperations?: VectorOperationsApi; - constructor(config: PineconeConfiguration, indexName: string) { + constructor(config: IndexConfiguration, indexName: string) { this.config = config; this.indexName = indexName; } @@ -23,10 +23,10 @@ export class VectorOperationsProvider { return this.vectorOperations; } - if (this.config.hostUrl) { + if (this.config.indexHostUrl) { this.vectorOperations = this.buildVectorOperationsConfig(this.config); } else { - this.config.hostUrl = await HostUrlSingleton.getHostUrl( + this.config.indexHostUrl = await IndexHostSingleton.getHostUrl( this.config, this.indexName ); diff --git a/src/pinecone.ts b/src/pinecone.ts index 4c7b5daa..7abac419 100644 --- a/src/pinecone.ts +++ b/src/pinecone.ts @@ -1,8 +1,3 @@ -import type { ConfigurationParameters as IndexOperationsApiConfigurationParameters } from './pinecone-generated-ts-fetch'; -import { - IndexOperationsApi, - Configuration as ApiConfiguration, -} from './pinecone-generated-ts-fetch'; import { describeIndex, listIndexes, @@ -17,18 +12,16 @@ import { CreateCollectionOptions, CreateIndexOptions, IndexName, + indexOperationsBuilder, CollectionName, } from './control'; -import type { IndexDescription } from './control'; -import { HostUrlSingleton } from './data/hostUrlSingleton'; +import { IndexHostSingleton } from './data/indexHostSingleton'; import { PineconeConfigurationError, PineconeEnvironmentVarsNotSupportedError, } from './errors'; -import { middleware } from './utils/middleware'; import { Index, PineconeConfigurationSchema } from './data'; import { buildValidator } from './validator'; -import { queryParamsStringify, buildUserAgent, getFetch } from './utils'; import type { PineconeConfiguration, RecordMetadata } from './data'; /** @@ -115,50 +108,26 @@ export class Pinecone { this.config = options; - const { apiKey } = options; - const controllerPath = `https://api.pinecone.io`; - const apiConfig: IndexOperationsApiConfigurationParameters = { - basePath: controllerPath, - apiKey, - queryParamsStringify, - headers: { - 'User-Agent': buildUserAgent(false), - }, - fetchApi: getFetch(options), - middleware, - }; - const api = new IndexOperationsApi(new ApiConfiguration(apiConfig)); + const api = indexOperationsBuilder(this.config); this._configureIndex = configureIndex(api); this._createCollection = createCollection(api); this._createIndex = createIndex(api); this._describeCollection = describeCollection(api); this._deleteCollection = deleteCollection(api); - this._describeIndex = describeIndex(api, this._handleDescribeIndex); + this._describeIndex = describeIndex(api); this._deleteIndex = deleteIndex(api); this._listCollections = listCollections(api); this._listIndexes = listIndexes(api); } - /** - * @internal - * Used as a callback by {@link Pinecone.describeIndex} to cache any index host URLs that are returned. - * This bypasses the need to fetch the host on subsequent dataplane calls. - */ - _handleDescribeIndex = (response: IndexDescription, indexName: IndexName) => { - if (response.status?.host) { - HostUrlSingleton._set(this.config, indexName, response.status.host); - } - }; - /** * @internal * This method is used by {@link Pinecone.constructor} to read configuration from environment variables. * * It looks for the following environment variables: - * - `PINECONE_ENVIRONMENT` * - `PINECONE_API_KEY` - * - `PINECONE_PROJECT_ID` + * - `PINECONE_CONTROLLER_HOST` * * @returns A {@link PineconeConfiguration} object populated with values found in environment variables. */ @@ -189,7 +158,9 @@ export class Pinecone { ); } - const optionalEnvVarMap = { projectId: 'PINECONE_PROJECT_ID' }; + const optionalEnvVarMap = { + controllerHostUrl: 'PINECONE_CONTROLLER_HOST', + }; for (const [key, envVar] of Object.entries(optionalEnvVarMap)) { const value = process.env[envVar]; if (value !== undefined) { @@ -229,7 +200,17 @@ export class Pinecone { * @returns A promise that resolves to {@link IndexMeta} */ describeIndex(indexName: IndexName) { - return this._describeIndex(indexName); + const describeIndexPromise = this._describeIndex(indexName); + + // For any describeIndex calls we want to update the IndexHostSingleton cache. + // This prevents unneeded calls to describeIndex for resolving the host for vector operations. + describeIndexPromise.then((indexMeta) => { + if (indexMeta.status?.host) { + IndexHostSingleton._set(this.config, indexName, indexMeta.status.host); + } + }); + + return describeIndexPromise; } /** @@ -330,7 +311,14 @@ export class Pinecone { * @throws {@link Errors.PineconeArgumentError} when invalid arguments are provided */ deleteIndex(indexName: IndexName) { - return this._deleteIndex(indexName); + const deleteIndexPromise = this._deleteIndex(indexName); + + // When an index is deleted, we need to evict the host from the IndexHostSingleton cache. + deleteIndexPromise.then(() => { + IndexHostSingleton._delete(this.config, indexName); + }); + + return deleteIndexPromise; } /** diff --git a/utils/cleanupResources.ts b/utils/cleanupResources.ts index 8210aaf3..4e3b5322 100644 --- a/utils/cleanupResources.ts +++ b/utils/cleanupResources.ts @@ -1,7 +1,7 @@ var dotenv = require('dotenv'); dotenv.config(); -for (const envVar of ['PINECONE_ENVIRONMENT', 'PINECONE_API_KEY']) { +for (const envVar of ['PINECONE_API_KEY']) { if (!process.env[envVar]) { console.warn(`WARNING Missing environment variable ${envVar} in .env file`); } else {