From 20a38acb8baa676ef461b81c0cf1a847cedbc245 Mon Sep 17 00:00:00 2001 From: Austin DeNoble Date: Tue, 20 Feb 2024 20:35:41 -0500 Subject: [PATCH] Add collections operations to `PineconeControlPlaneClient` with integration tests (#65) ## Problem The Java SDK is currently missing collections functionality. Now that we've added generated code for the OpenAPI spec, we can hook up collections operations in `PineconeControlPlaneClient`. We'd also like to add integration tests to cover collections. ## Solution - Update `PineconeControlPlaneClient` to include collections operations for create, list, delete, and describe. - Add new integration test files: `CollectionTest` and `CollectionErrorTest`. - Add new helpers to `IndexManager` for creating and connecting to an index by `indexName`, creating a collection and waiting for it to be ready, and polling until an index is ready. - Add new helper to `BuildUpsertRequest` for generating vectors by dimension. Bonuses: - I added some logging configs in `build.gradle` for `test` and the `integrationTest` task. This was primarily to help me debug things in CI and log a bit more info to the console. We can tweak as needed, but I think having something like this will be really helpful as opposed to the generic output. - Set `max-parallel: 1` in the `integration-test` job in the `pr.yml` workflow. Because of the way `findIndexWithDimensionAndType` works we cannot run integration tests in parallel without flakiness. - Fixed an issue in `AssertRetry`. We were incrementing `delay` at the class level which is a static variable, meaning if you re-used the function in one place it would continue increasing the delay for all subsequent calls based on the `backOff`. ## Type of Change - [X] New feature (non-breaking change which adds functionality) ## Test Plan Run integration tests to make sure `CollectionTest` and `CollectionErrorTest` pass as expected. Integration tests should also be passing in CI barring any flakiness. --- .github/workflows/pr.yml | 3 +- .../java/io/pinecone/helpers/AssertRetry.java | 21 +- .../pinecone/helpers/BuildUpsertRequest.java | 33 ++- .../io/pinecone/helpers/IndexManager.java | 84 +++++++- .../controlPlane/pod/CollectionErrorTest.java | 163 ++++++++++++++ .../controlPlane/pod/CollectionTest.java | 198 ++++++++++++++++++ .../{index => }/pod/ConfigureIndexTest.java | 16 +- .../CreateDescribeListAndDeleteIndexTest.java | 2 +- .../CreateDescribeListAndDeleteIndexTest.java | 2 +- .../PineconeClientLiveIntegTest.java | 10 +- .../dataplane/UpdateAndQueryTest.java | 10 +- .../dataplane/UpsertAndDeleteTest.java | 9 +- .../UpsertAndDescribeIndexStatsTest.java | 8 +- .../pinecone/PineconeControlPlaneClient.java | 58 ++++- .../pinecone/exceptions/HttpErrorMapper.java | 2 +- 15 files changed, 577 insertions(+), 42 deletions(-) create mode 100644 src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionErrorTest.java create mode 100644 src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionTest.java rename src/integration/java/io/pinecone/integration/controlPlane/{index => }/pod/ConfigureIndexTest.java (95%) rename src/integration/java/io/pinecone/integration/controlPlane/{index => }/pod/CreateDescribeListAndDeleteIndexTest.java (97%) rename src/integration/java/io/pinecone/integration/controlPlane/{index => }/serverless/CreateDescribeListAndDeleteIndexTest.java (96%) diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 0533b18d..2e4fa8e9 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -58,6 +58,7 @@ jobs: { java: 8, gradle: 6.8 }, { java: 17, gradle: 7.3.1 } ] + max-parallel: 1 steps: - uses: actions/checkout@v4 @@ -83,5 +84,5 @@ jobs: env: PINECONE_API_KEY: ${{ secrets.PINECONE_API_KEY }} PINECONE_ENVIRONMENT: ${{ secrets.PINECONE_ENVIRONMENT }} - + diff --git a/src/integration/java/io/pinecone/helpers/AssertRetry.java b/src/integration/java/io/pinecone/helpers/AssertRetry.java index 1212893f..c33daa1e 100644 --- a/src/integration/java/io/pinecone/helpers/AssertRetry.java +++ b/src/integration/java/io/pinecone/helpers/AssertRetry.java @@ -1,34 +1,43 @@ package io.pinecone.helpers; +import io.pinecone.exceptions.PineconeException; + import java.io.IOException; import java.util.concurrent.ExecutionException; public class AssertRetry { private static final int maxRetry = 4; - private static int delay = 1500; + private static final int delay = 1500; - public static void assertWithRetry(AssertionRunnable assertionRunnable) throws InterruptedException { + public static void assertWithRetry(AssertionRunnable assertionRunnable) throws InterruptedException, PineconeException { assertWithRetry(assertionRunnable, 2); } - public static void assertWithRetry(AssertionRunnable assertionRunnable, int backOff) throws InterruptedException { + public static void assertWithRetry(AssertionRunnable assertionRunnable, int backOff) throws AssertionError, InterruptedException { int retryCount = 0; + int delayCount = delay; boolean success = false; + String errorMessage = null; while (retryCount < maxRetry && !success) { try { assertionRunnable.run(); success = true; } catch (AssertionError | ExecutionException | IOException e) { + errorMessage = e.getLocalizedMessage(); retryCount++; - delay*=backOff; - Thread.sleep(delay); + delayCount*=backOff; + Thread.sleep(delayCount); } } + + if (!success) { + throw new AssertionError(errorMessage); + } } @FunctionalInterface public interface AssertionRunnable { - void run() throws AssertionError, ExecutionException, InterruptedException, IOException; + void run() throws AssertionError, ExecutionException, InterruptedException, IOException, PineconeException; } } diff --git a/src/integration/java/io/pinecone/helpers/BuildUpsertRequest.java b/src/integration/java/io/pinecone/helpers/BuildUpsertRequest.java index 3fa5acd0..6e1cca36 100644 --- a/src/integration/java/io/pinecone/helpers/BuildUpsertRequest.java +++ b/src/integration/java/io/pinecone/helpers/BuildUpsertRequest.java @@ -7,10 +7,7 @@ import io.pinecone.proto.UpsertRequest; import io.pinecone.proto.Vector; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.HashMap; -import java.util.List; +import java.util.*; public class BuildUpsertRequest { private static final float[][] upsertData = {{1.0F, 2.0F, 3.0F}, {4.0F, 5.0F, 6.0F}, {7.0F, 8.0F, 9.0F}}; @@ -44,6 +41,23 @@ public static UpsertRequest buildRequiredUpsertRequest(List upsertIds, S .build(); } + public static UpsertRequest buildRequiredUpsertRequestByDimension(List upsertIds, int dimension, String namespace) { + if (upsertIds.isEmpty()) upsertIds = Arrays.asList("v1", "v2", "v3"); + + List upsertVectors = new ArrayList<>(); + for (String upsertId : upsertIds) { + upsertVectors.add(Vector.newBuilder() + .addAllValues(generateVectorValuesByDimension(dimension)) + .setId(upsertId) + .build()); + } + + return UpsertRequest.newBuilder() + .addAllVectors(upsertVectors) + .setNamespace(namespace) + .build(); + } + public static UpsertRequest buildOptionalUpsertRequest() { return buildOptionalUpsertRequest(new ArrayList<>(), ""); } @@ -108,4 +122,15 @@ public static HashMap> createAndGetMetadataMap() { return metadataMap; } + + public static ArrayList generateVectorValuesByDimension(int dimension) { + ArrayList values = new ArrayList<>(); + Random random = new Random(); + + for (int i = 0; i < dimension; i++) { + values.add(random.nextFloat()); + } + + return values; + } } diff --git a/src/integration/java/io/pinecone/helpers/IndexManager.java b/src/integration/java/io/pinecone/helpers/IndexManager.java index c99818a9..6d5d3900 100644 --- a/src/integration/java/io/pinecone/helpers/IndexManager.java +++ b/src/integration/java/io/pinecone/helpers/IndexManager.java @@ -1,15 +1,21 @@ package io.pinecone.helpers; import io.pinecone.*; +import io.pinecone.exceptions.PineconeException; import org.openapitools.client.model.*; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.io.IOException; import java.util.List; import static io.pinecone.helpers.AssertRetry.assertWithRetry; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.fail; public class IndexManager { private static PineconeClientConfig config; + private static final Logger logger = LoggerFactory.getLogger(IndexManager.class); public static PineconeConnection createIndexIfNotExistsDataPlane(int dimension, String indexType) throws IOException, InterruptedException { String apiKey = System.getenv("PINECONE_API_KEY"); @@ -45,9 +51,11 @@ private static String findIndexWithDimensionAndType(IndexList indexList, int dim List indexModels = indexList.getIndexes(); while (i < indexModels.size()) { IndexModel indexModel = isIndexReady(indexModels.get(i).getName(), controlPlaneClient); - // ToDo: add pod type support if (indexModel.getDimension() == dimension - && ((indexType.equalsIgnoreCase(IndexModelSpec.SERIALIZED_NAME_POD) && indexModel.getSpec().getPod() != null && indexModel.getSpec().getPod().getReplicas() == 1 && indexModel.getSpec().getPod().getPodType().equalsIgnoreCase("p1.x1")) + && ((indexType.equalsIgnoreCase(IndexModelSpec.SERIALIZED_NAME_POD) + && indexModel.getSpec().getPod() != null + && indexModel.getSpec().getPod().getReplicas() == 1 + && indexModel.getSpec().getPod().getPodType().equalsIgnoreCase("p1.x1")) || (indexType.equalsIgnoreCase(IndexModelSpec.SERIALIZED_NAME_SERVERLESS)))) { return indexModel.getName(); } @@ -79,13 +87,83 @@ private static String createNewIndex(PineconeControlPlaneClient controlPlaneClie return indexName; } + public static IndexModel waitUntilIndexIsReady(PineconeControlPlaneClient controlPlaneClient, String indexName, Integer totalMsToWait) throws InterruptedException { + IndexModel index = controlPlaneClient.describeIndex(indexName); + int waitedTimeMs = 0; + int intervalMs = 1500; + + while (!index.getStatus().getReady()) { + index = controlPlaneClient.describeIndex(indexName); + if (waitedTimeMs >= totalMsToWait) { + logger.info("Index " + indexName + " not ready after " + waitedTimeMs + "ms"); + break; + } + if (index.getStatus().getReady()) { + logger.info("Index " + indexName + " is ready after " + waitedTimeMs + "ms"); + break; + } + Thread.sleep(intervalMs); + waitedTimeMs += intervalMs; + } + return index; + } + + public static IndexModel waitUntilIndexIsReady(PineconeControlPlaneClient controlPlaneClient, String indexName) throws InterruptedException { + return waitUntilIndexIsReady(controlPlaneClient, indexName, 120000); + } + + public static PineconeConnection createNewIndexAndConnect(PineconeControlPlaneClient controlPlaneClient, String indexName, int dimension, IndexMetric metric, CreateIndexRequestSpec spec) throws InterruptedException, PineconeException { + CreateIndexRequest createIndexRequest = new CreateIndexRequest().name(indexName).dimension(dimension).metric(metric).spec(spec); + controlPlaneClient.createIndex(createIndexRequest); + + // Wait until index is ready + waitUntilIndexIsReady(controlPlaneClient, indexName, 200000); + // wait a bit more before we connect... + Thread.sleep(15000); + + String host = controlPlaneClient.describeIndex(indexName).getHost(); + + PineconeClientConfig specificConfig = new PineconeClientConfig().withApiKey(System.getenv("PINECONE_API_KEY")); + PineconeClient dataPlaneClient = new PineconeClient(specificConfig); + + return dataPlaneClient.connect( + new PineconeConnectionConfig() + .withConnectionUrl("https://" + host)); + } + + public static CollectionModel createCollection(PineconeControlPlaneClient controlPlaneClient, String collectionName, String indexName, boolean waitUntilReady) throws InterruptedException { + CreateCollectionRequest createCollectionRequest = new CreateCollectionRequest().name(collectionName).source(indexName); + CollectionModel collection = controlPlaneClient.createCollection(createCollectionRequest); + + assertEquals(collection.getStatus(), CollectionModel.StatusEnum.INITIALIZING); + + // Wait until collection is ready + if (waitUntilReady) { + int timeWaited = 0; + CollectionModel.StatusEnum collectionReady = collection.getStatus(); + while (collectionReady != CollectionModel.StatusEnum.READY && timeWaited < 120000) { + logger.info("Waiting for collection" + collectionName + " to be ready. Waited " + timeWaited + " milliseconds..."); + Thread.sleep(5000); + timeWaited += 5000; + collection = controlPlaneClient.describeCollection(collectionName); + collectionReady = collection.getStatus(); + } + + if (timeWaited > 120000) { + fail("Collection: " + collectionName + " is not ready after 120 seconds"); + } + } + + return collection; + } + public static IndexModel isIndexReady(String indexName, PineconeControlPlaneClient controlPlaneClient) throws InterruptedException { final IndexModel[] indexModels = new IndexModel[1]; assertWithRetry(() -> { indexModels[0] = controlPlaneClient.describeIndex(indexName); assert (indexModels[0].getStatus().getReady()); - }, 1); + }, 4); return indexModels[0]; } diff --git a/src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionErrorTest.java b/src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionErrorTest.java new file mode 100644 index 00000000..c7a08b5f --- /dev/null +++ b/src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionErrorTest.java @@ -0,0 +1,163 @@ +package io.pinecone.integration.controlPlane.pod; + +import io.pinecone.PineconeConnection; +import io.pinecone.PineconeControlPlaneClient; +import io.pinecone.exceptions.PineconeException; +import io.pinecone.helpers.RandomStringBuilder; +import io.pinecone.proto.VectorServiceGrpc; +import org.openapitools.client.model.*; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.ArrayList; +import java.util.LinkedList; +import java.util.List; +import java.util.Random; + +import static io.pinecone.helpers.IndexManager.createNewIndexAndConnect; +import static io.pinecone.helpers.IndexManager.createCollection; +import static io.pinecone.helpers.IndexManager.waitUntilIndexIsReady; +import static io.pinecone.helpers.BuildUpsertRequest.*; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class CollectionErrorTest { + private static final String apiKey = System.getenv("PINECONE_API_KEY"); + private static final String environment = System.getenv("PINECONE_ENVIRONMENT"); + private static final String indexName = RandomStringBuilder.build("collection-error-test", 8); + private static final String collectionName = RandomStringBuilder.build("reusable-coll", 8); + private static final ArrayList indexes = new ArrayList<>(); + private static final ArrayList collections = new ArrayList<>(); + private static final List upsertIds = Arrays.asList("v1", "v2", "v3"); + private static final int dimension = 4; + private static PineconeControlPlaneClient controlPlaneClient; + private static final Logger logger = LoggerFactory.getLogger(CollectionErrorTest.class); + + @BeforeAll + public static void setUpIndexAndCollection() throws InterruptedException { + controlPlaneClient = new PineconeControlPlaneClient(apiKey); + CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().pods(1).podType("p1.x1").replicas(1).environment(environment); + CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec); + PineconeConnection dataPlaneConnection = createNewIndexAndConnect(controlPlaneClient, indexName, dimension, IndexMetric.COSINE, spec); + VectorServiceGrpc.VectorServiceBlockingStub blockingStub = dataPlaneConnection.getBlockingStub(); + indexes.add(indexName); + + // Upsert vectors to index and sleep for freshness + blockingStub.upsert(buildRequiredUpsertRequestByDimension(upsertIds, dimension, "")); + dataPlaneConnection.close(); + Thread.sleep(3500); + + // Create collection from index + createCollection(controlPlaneClient, collectionName, indexName, true); + collections.add(collectionName); + } + + @AfterAll + public static void cleanUp() throws InterruptedException { + // wait for things to settle before cleanup... + Thread.sleep(2500); + for (String index : indexes) { + controlPlaneClient.deleteIndex(index); + } + for (String collection : collections) { + controlPlaneClient.deleteCollection(collection); + } + Thread.sleep(2500); + } + + @Test + public void testCreateCollectionFromInvalidIndex() { + try { + CreateCollectionRequest createCollectionRequest = new CreateCollectionRequest().name(RandomStringBuilder.build("coll1", 8)).source("invalid-index"); + controlPlaneClient.createCollection(createCollectionRequest); + } catch (PineconeException exception) { + logger.info("Exception: " + exception.getMessage()); + assertTrue(exception.getMessage().contains("Resource invalid-index not found")); + } + } + @Test + public void testIndexFromNonExistentCollection() { + try { + CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().environment(environment).sourceCollection("non-existent-collection"); + CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec); + CreateIndexRequest newCreateIndexRequest = new CreateIndexRequest().name(RandomStringBuilder.build("from-nonexistent-coll", 8)).dimension(dimension).metric(IndexMetric.COSINE).spec(spec); + controlPlaneClient.createIndex(newCreateIndexRequest); + } catch (PineconeException exception) { + logger.info("Exception: " + exception.getMessage()); + assertTrue(exception.getMessage().contains("Resource non-existent-collection not found")); + } + } + + @Test + public void testCreateIndexInMismatchedEnvironment() { + try { + List environments = new LinkedList<>(Arrays.asList( + "eastus-azure", + "eu-west4-gcp", + "northamerica-northeast1-gcp", + "us-central1-gcp", + "us-west4-gcp", + "asia-southeast1-gcp", + "us-east-1-aws", + "asia-northeast1-gcp", + "eu-west1-gcp", + "us-east1-gcp", + "us-east4-gcp", + "us-west1-gcp" + )); + CollectionModel collection = controlPlaneClient.describeCollection(collectionName); + environments.remove(collection.getEnvironment()); + String mismatchedEnv = environments.get(new Random().nextInt(environments.size())); + + CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().sourceCollection(collection.getName()).environment(mismatchedEnv); + CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec); + CreateIndexRequest createIndexRequest = new CreateIndexRequest().name(RandomStringBuilder.build("from-coll", 8)).dimension(dimension).metric(IndexMetric.COSINE).spec(spec); + controlPlaneClient.createIndex(createIndexRequest); + } catch (PineconeException exception) { + logger.info("Exception: " + exception.getMessage()); + assertTrue(exception.getMessage().contains("Source collection must be in the same environment as the index")); + } + } + + @Test + @Disabled("Bug reported in #global-cps") + public void testCreateIndexWithMismatchedDimension() { + try { + CollectionModel collection = controlPlaneClient.describeCollection(collectionName); + CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().sourceCollection(collection.getName()).environment(collection.getEnvironment()); + CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec); + CreateIndexRequest createIndexRequest = new CreateIndexRequest().name(RandomStringBuilder.build("from-coll", 8)).dimension(dimension + 1).metric(IndexMetric.COSINE).spec(spec); + controlPlaneClient.createIndex(createIndexRequest); + } catch (PineconeException exception) { + logger.info("Exception: " + exception.getMessage()); + assertTrue(exception.getMessage().contains("Index and collection must have the same dimension")); + } + } + + @Test + public void testCreateCollectionFromNotReadyIndex() throws InterruptedException { + String notReadyIndexName = RandomStringBuilder.build("from-coll4", 8); + String newCollectionName = RandomStringBuilder.build("coll4-", 8); + try { + CreateIndexRequestSpecPod specPod = new CreateIndexRequestSpecPod().pods(1).podType("p1.x1").replicas(1).environment(environment); + CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(specPod); + CreateIndexRequest createIndexRequest = new CreateIndexRequest().name(notReadyIndexName).dimension(dimension).metric(IndexMetric.COSINE).spec(spec); + controlPlaneClient.createIndex(createIndexRequest); + indexes.add(notReadyIndexName); + + createCollection(controlPlaneClient, newCollectionName, notReadyIndexName, true); + collections.add(newCollectionName); + } catch (PineconeException exception) { + logger.info("Exception: " + exception.getMessage()); + assert (exception.getMessage().contains("Source index is not ready")); + } finally { + // Wait for index to initialize so it can be deleted in @AfterAll + waitUntilIndexIsReady(controlPlaneClient, notReadyIndexName); + } + } +} \ No newline at end of file diff --git a/src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionTest.java b/src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionTest.java new file mode 100644 index 00000000..f18bb920 --- /dev/null +++ b/src/integration/java/io/pinecone/integration/controlPlane/pod/CollectionTest.java @@ -0,0 +1,198 @@ +package io.pinecone.integration.controlPlane.pod; + +import io.pinecone.*; +import io.pinecone.helpers.RandomStringBuilder; +import io.pinecone.proto.*; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.openapitools.client.model.*; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +import static io.pinecone.helpers.IndexManager.createNewIndexAndConnect; +import static io.pinecone.helpers.IndexManager.waitUntilIndexIsReady; +import static io.pinecone.helpers.IndexManager.createCollection; +import static io.pinecone.helpers.BuildUpsertRequest.*; +import static org.junit.jupiter.api.Assertions.*; + +public class CollectionTest { + + private static PineconeControlPlaneClient controlPlaneClient; + private static final String indexName = RandomStringBuilder.build("collection-test", 8); + private static final ArrayList indexes = new ArrayList<>(); + private static final ArrayList collections = new ArrayList<>(); + private static final IndexMetric indexMetric = IndexMetric.COSINE; + private static final List upsertIds = Arrays.asList("v1", "v2", "v3"); + private static final String namespace = RandomStringBuilder.build("ns", 8); + private static final String apiKey = System.getenv("PINECONE_API_KEY"); + private static final String environment = System.getenv("PINECONE_ENVIRONMENT"); + private static final int dimension = 4; + + private static final Logger logger = LoggerFactory.getLogger(CollectionTest.class); + + @BeforeAll + public static void setUp() throws InterruptedException { + controlPlaneClient = new PineconeControlPlaneClient(apiKey); + CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().pods(1).podType("p1.x1").replicas(1).environment(environment); + CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec); + PineconeConnection dataPlaneConnection = createNewIndexAndConnect(controlPlaneClient, indexName, dimension, indexMetric, spec); + VectorServiceGrpc.VectorServiceBlockingStub blockingStub = dataPlaneConnection.getBlockingStub(); + indexes.add(indexName); + + blockingStub.upsert(buildRequiredUpsertRequestByDimension(upsertIds, dimension, namespace)); + dataPlaneConnection.close(); + + } + + @AfterAll + public static void cleanUp() throws InterruptedException { + // wait for things to settle before cleanup... + Thread.sleep(2500); + // Clean up indexes + for (String index : indexes) { + controlPlaneClient.deleteIndex(index); + } + // Clean up collections + for (String collection : collections) { + controlPlaneClient.deleteCollection(collection); + } + Thread.sleep(2500); + } + + @Test + public void testIndexToCollectionHappyPath() throws InterruptedException { + String collectionName = RandomStringBuilder.build("collection-test", 8); + + // Create collection from index + CollectionModel collection = createCollection(controlPlaneClient, collectionName, indexName, true); + collections.add(collectionName); + + assertEquals(collection.getName(), collectionName); + assertEquals(collection.getEnvironment(), environment); + assertEquals(collection.getStatus(), CollectionModel.StatusEnum.READY); + + // Verify collection is listed + List collectionList = controlPlaneClient.listCollections().getCollections(); + boolean collectionFound = false; + if (collectionList != null && !collectionList.isEmpty()) { + for (CollectionModel col : collectionList) { + if (col.getName().equals(collectionName)) { + collectionFound = true; + break; + } + } + } + + if (!collectionFound) { + fail("Collection " + collectionName + " was not found when listing collections"); + } + + // Verify collection can be described + collection = controlPlaneClient.describeCollection(collectionName); + + assertEquals(collection.getStatus(), CollectionModel.StatusEnum.READY); + assertEquals(collection.getDimension(), dimension); + assertEquals(collection.getVectorCount(), 3); + assertNotEquals(collection.getVectorCount(), null); + assertTrue(collection.getSize() > 0); + + // Create index from collection + String newIndexName = RandomStringBuilder.build("index-from-col", 5); + logger.info("Creating index " + newIndexName + " from collection " + collectionName); + + CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().environment(environment).sourceCollection(collectionName); + CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec); + CreateIndexRequest newCreateIndexRequest = new CreateIndexRequest().name(newIndexName).dimension(dimension).metric(indexMetric).spec(spec); + controlPlaneClient.createIndex(newCreateIndexRequest); + indexes.add(newIndexName); + logger.info("Index " + newIndexName + " created from collection " + collectionName + ". Waiting until index is ready..."); + waitUntilIndexIsReady(controlPlaneClient, newIndexName, 250000); + // wait a bit more to make sure index is ready... + Thread.sleep(30000); + + IndexModel indexDescription = controlPlaneClient.describeIndex(newIndexName); + assertEquals(indexDescription.getName(), newIndexName); + assertEquals(indexDescription.getSpec().getPod().getSourceCollection(), collectionName); + assertEquals(indexDescription.getStatus().getReady(), true); + + // Set up new index data plane connection + PineconeClient newIndexClient = new PineconeClient(new PineconeClientConfig().withApiKey(apiKey).withEnvironment(environment)); + PineconeConnection newIndexDataPlaneClient = newIndexClient.connect(new PineconeConnectionConfig().withConnectionUrl("https://" + indexDescription.getHost())); + VectorServiceGrpc.VectorServiceBlockingStub newIndexBlockingStub = newIndexDataPlaneClient.getBlockingStub(); + DescribeIndexStatsResponse describeResponse = newIndexBlockingStub.describeIndexStats(DescribeIndexStatsRequest.newBuilder().build()); + + // Verify stats reflect the vectors in the collection + assertEquals(describeResponse.getTotalVectorCount(), 3); + + // Verify the vectors from the collection -> new index can be fetched + FetchResponse fetchedVectors = newIndexBlockingStub.fetch(FetchRequest.newBuilder().addAllIds(upsertIds).setNamespace(namespace).build()); + + for (String key : upsertIds) { + assert (fetchedVectors.containsVectors(key)); + } + + // Verify we can delete the collection + controlPlaneClient.deleteCollection(collectionName); + collections.remove(collectionName); + Thread.sleep(2500); + + collectionList = controlPlaneClient.listCollections().getCollections(); + + + if (collectionList != null) { + boolean isCollectionDeleted = true; + for (CollectionModel col : collectionList) { + if (col.getName().equals(collectionName)) { + isCollectionDeleted = false; + break; + } + } + + if (!isCollectionDeleted) { + fail("Collection " + collectionName + " was not successfully deleted"); + } + } + + newIndexDataPlaneClient.close(); + } + + @Test + public void testIndexFromDifferentMetricCollection() throws InterruptedException { + String collectionName = RandomStringBuilder.build("collection-test", 8); + + // Create collection from index + CollectionModel collection = createCollection(controlPlaneClient, collectionName, indexName, true); + collections.add(collectionName); + + assertEquals(collection.getName(), collectionName); + assertEquals(collection.getEnvironment(), environment); + assertEquals(collection.getStatus(), CollectionModel.StatusEnum.READY); + + // Use a different metric than the source index + IndexMetric[] metrics = { IndexMetric.COSINE, IndexMetric.EUCLIDEAN, IndexMetric.DOTPRODUCT }; + IndexMetric targetMetric = IndexMetric.COSINE; + for (IndexMetric metric : metrics) { + if (!metric.equals(indexMetric)) { + targetMetric = metric; + } + } + + String newIndexName = RandomStringBuilder.build("from-coll", 8); + CreateIndexRequestSpecPod podSpec = new CreateIndexRequestSpecPod().environment(environment).sourceCollection(collectionName); + CreateIndexRequestSpec spec = new CreateIndexRequestSpec().pod(podSpec); + PineconeConnection dataPlaneConnection = createNewIndexAndConnect(controlPlaneClient, newIndexName, dimension, targetMetric, spec); + indexes.add(newIndexName); + + IndexModel newIndex = controlPlaneClient.describeIndex(newIndexName); + assertEquals(newIndex.getName(), newIndexName); + assertEquals(newIndex.getMetric(), targetMetric); + + dataPlaneConnection.close(); + } + +} \ No newline at end of file diff --git a/src/integration/java/io/pinecone/integration/controlPlane/index/pod/ConfigureIndexTest.java b/src/integration/java/io/pinecone/integration/controlPlane/pod/ConfigureIndexTest.java similarity index 95% rename from src/integration/java/io/pinecone/integration/controlPlane/index/pod/ConfigureIndexTest.java rename to src/integration/java/io/pinecone/integration/controlPlane/pod/ConfigureIndexTest.java index 4e49cdc5..148e1a16 100644 --- a/src/integration/java/io/pinecone/integration/controlPlane/index/pod/ConfigureIndexTest.java +++ b/src/integration/java/io/pinecone/integration/controlPlane/pod/ConfigureIndexTest.java @@ -1,4 +1,4 @@ -package io.pinecone.integration.controlPlane.index.pod; +package io.pinecone.integration.controlPlane.pod; import io.pinecone.PineconeControlPlaneClient; import io.pinecone.exceptions.PineconeException; @@ -20,16 +20,12 @@ public class ConfigureIndexTest { private static PineconeControlPlaneClient controlPlaneClient; - private String indexName; - private static final Logger logger = LoggerFactory.getLogger(PineconeClientLiveIntegTest.class); + private static String indexName; + private static final Logger logger = LoggerFactory.getLogger(ConfigureIndexTest.class); @BeforeAll - public static void defineControlPlaneClient() { + public static void setUp() throws InterruptedException, IOException { controlPlaneClient = new PineconeControlPlaneClient(System.getenv("PINECONE_API_KEY")); - } - - @BeforeEach - public void setUp() throws IOException, InterruptedException { indexName = createIndexIfNotExistsControlPlane(controlPlaneClient, 5, IndexModelSpec.SERIALIZED_NAME_POD); } @@ -83,7 +79,7 @@ public void scaleUpAndDown() { }); // Scaling down - pod = new ConfigureIndexRequestSpecPod().replicas(3); + pod = new ConfigureIndexRequestSpecPod().replicas(1); spec = new ConfigureIndexRequestSpec().pod(pod); configureIndexRequest = new ConfigureIndexRequest().spec(spec); controlPlaneClient.configureIndex(indexName, configureIndexRequest); @@ -95,7 +91,7 @@ public void scaleUpAndDown() { assertEquals(podSpec.getReplicas(), 1); }); } catch (Exception exception) { - throw new PineconeException("Test failed: " + exception.getStackTrace()); + throw new PineconeException("Test failed: " + exception.getLocalizedMessage()); } } diff --git a/src/integration/java/io/pinecone/integration/controlPlane/index/pod/CreateDescribeListAndDeleteIndexTest.java b/src/integration/java/io/pinecone/integration/controlPlane/pod/CreateDescribeListAndDeleteIndexTest.java similarity index 97% rename from src/integration/java/io/pinecone/integration/controlPlane/index/pod/CreateDescribeListAndDeleteIndexTest.java rename to src/integration/java/io/pinecone/integration/controlPlane/pod/CreateDescribeListAndDeleteIndexTest.java index b031e7a6..18333ee6 100644 --- a/src/integration/java/io/pinecone/integration/controlPlane/index/pod/CreateDescribeListAndDeleteIndexTest.java +++ b/src/integration/java/io/pinecone/integration/controlPlane/pod/CreateDescribeListAndDeleteIndexTest.java @@ -1,4 +1,4 @@ -package io.pinecone.integration.controlPlane.index.pod; +package io.pinecone.integration.controlPlane.pod; import io.pinecone.PineconeControlPlaneClient; import io.pinecone.helpers.RandomStringBuilder; diff --git a/src/integration/java/io/pinecone/integration/controlPlane/index/serverless/CreateDescribeListAndDeleteIndexTest.java b/src/integration/java/io/pinecone/integration/controlPlane/serverless/CreateDescribeListAndDeleteIndexTest.java similarity index 96% rename from src/integration/java/io/pinecone/integration/controlPlane/index/serverless/CreateDescribeListAndDeleteIndexTest.java rename to src/integration/java/io/pinecone/integration/controlPlane/serverless/CreateDescribeListAndDeleteIndexTest.java index df9a4dba..d39683cf 100644 --- a/src/integration/java/io/pinecone/integration/controlPlane/index/serverless/CreateDescribeListAndDeleteIndexTest.java +++ b/src/integration/java/io/pinecone/integration/controlPlane/serverless/CreateDescribeListAndDeleteIndexTest.java @@ -1,4 +1,4 @@ -package io.pinecone.integration.controlPlane.index.serverless; +package io.pinecone.integration.controlPlane.serverless; import io.pinecone.PineconeControlPlaneClient; import io.pinecone.helpers.RandomStringBuilder; diff --git a/src/integration/java/io/pinecone/integration/dataplane/PineconeClientLiveIntegTest.java b/src/integration/java/io/pinecone/integration/dataplane/PineconeClientLiveIntegTest.java index e1d2b5be..67d8ec57 100644 --- a/src/integration/java/io/pinecone/integration/dataplane/PineconeClientLiveIntegTest.java +++ b/src/integration/java/io/pinecone/integration/dataplane/PineconeClientLiveIntegTest.java @@ -6,6 +6,7 @@ import io.pinecone.PineconeConnection; import io.pinecone.helpers.RandomStringBuilder; import io.pinecone.proto.*; +import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -24,14 +25,21 @@ public class PineconeClientLiveIntegTest { private static final Logger logger = LoggerFactory.getLogger(PineconeClientLiveIntegTest.class); + + private static PineconeConnection connection; private static VectorServiceGrpc.VectorServiceBlockingStub blockingStub; @BeforeAll public static void defineConfig() throws IOException, InterruptedException { - PineconeConnection connection = createIndexIfNotExistsDataPlane(3, IndexModelSpec.SERIALIZED_NAME_POD); + connection = createIndexIfNotExistsDataPlane(3, IndexModelSpec.SERIALIZED_NAME_POD); blockingStub = connection.getBlockingStub(); } + @AfterAll + public static void cleanUp() { + connection.close(); + } + @Test public void sanity() throws InterruptedException { String namespace = RandomStringBuilder.build("ns", 8); diff --git a/src/integration/java/io/pinecone/integration/dataplane/UpdateAndQueryTest.java b/src/integration/java/io/pinecone/integration/dataplane/UpdateAndQueryTest.java index fe182f30..c98e33f7 100644 --- a/src/integration/java/io/pinecone/integration/dataplane/UpdateAndQueryTest.java +++ b/src/integration/java/io/pinecone/integration/dataplane/UpdateAndQueryTest.java @@ -6,6 +6,8 @@ import io.pinecone.PineconeConnection; import io.pinecone.helpers.RandomStringBuilder; import io.pinecone.proto.*; +import org.junit.After; +import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.openapitools.client.model.IndexModelSpec; @@ -21,17 +23,23 @@ import static org.junit.jupiter.api.Assertions.assertEquals; public class UpdateAndQueryTest { + private static PineconeConnection connection; private static VectorServiceGrpc.VectorServiceBlockingStub blockingStub; private static VectorServiceGrpc.VectorServiceFutureStub futureStub; private static final int dimension = 3; @BeforeAll public static void setUp() throws IOException, InterruptedException { - PineconeConnection connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD); + connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD); blockingStub = connection.getBlockingStub(); futureStub = connection.getFutureStub(); } + @AfterAll + public static void cleanUp() { + connection.close(); + } + @Test public void UpdateRequiredParamsFetchAndQuerySync() throws InterruptedException { // Upsert vectors with required parameters diff --git a/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDeleteTest.java b/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDeleteTest.java index 5fa2629b..c7276325 100644 --- a/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDeleteTest.java +++ b/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDeleteTest.java @@ -5,6 +5,7 @@ import io.pinecone.PineconeConnection; import io.pinecone.helpers.RandomStringBuilder; import io.pinecone.proto.*; +import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.openapitools.client.model.IndexModelSpec; @@ -20,17 +21,23 @@ import static org.junit.jupiter.api.Assertions.assertThrows; public class UpsertAndDeleteTest { + private static PineconeConnection connection; private static VectorServiceGrpc.VectorServiceBlockingStub blockingStub; private static VectorServiceGrpc.VectorServiceFutureStub futureStub; private static final int dimension = 3; @BeforeAll public static void setUp() throws IOException, InterruptedException { - PineconeConnection connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD); + connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD); blockingStub = connection.getBlockingStub(); futureStub = connection.getFutureStub(); } + @AfterAll + public static void cleanUp() { + connection.close(); + } + @Test public void UpsertVectorsAndDeleteByIdSyncTest() throws InterruptedException { // Upsert vectors with required parameters diff --git a/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDescribeIndexStatsTest.java b/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDescribeIndexStatsTest.java index 0b065d52..5f639ded 100644 --- a/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDescribeIndexStatsTest.java +++ b/src/integration/java/io/pinecone/integration/dataplane/UpsertAndDescribeIndexStatsTest.java @@ -14,17 +14,23 @@ import java.util.concurrent.ExecutionException; public class UpsertAndDescribeIndexStatsTest { + private static PineconeConnection connection; private static VectorServiceGrpc.VectorServiceBlockingStub blockingStub; private static VectorServiceGrpc.VectorServiceFutureStub futureStub; private static final int dimension = 3; @BeforeAll public static void setUp() throws IOException, InterruptedException { - PineconeConnection connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD); + connection = createIndexIfNotExistsDataPlane(dimension, IndexModelSpec.SERIALIZED_NAME_POD); blockingStub = connection.getBlockingStub(); futureStub = connection.getFutureStub(); } + @AfterAll + public static void cleanUp() { + connection.close(); + } + @Test public void UpsertRequiredVectorsAndDescribeIndexStatsSyncTest() throws InterruptedException { // Get vector and namespace counts before upserting vectors with required parameters diff --git a/src/main/java/io/pinecone/PineconeControlPlaneClient.java b/src/main/java/io/pinecone/PineconeControlPlaneClient.java index f8c3e09d..9052e561 100644 --- a/src/main/java/io/pinecone/PineconeControlPlaneClient.java +++ b/src/main/java/io/pinecone/PineconeControlPlaneClient.java @@ -2,15 +2,13 @@ import io.pinecone.exceptions.FailedRequestInfo; import io.pinecone.exceptions.HttpErrorMapper; +import io.pinecone.exceptions.PineconeException; import io.pinecone.exceptions.PineconeValidationException; import okhttp3.*; import org.openapitools.client.ApiClient; import org.openapitools.client.ApiException; import org.openapitools.client.api.ManageIndexesApi; -import org.openapitools.client.model.ConfigureIndexRequest; -import org.openapitools.client.model.CreateIndexRequest; -import org.openapitools.client.model.IndexList; -import org.openapitools.client.model.IndexModel; +import org.openapitools.client.model.*; public class PineconeControlPlaneClient { private ManageIndexesApi manageIndexesApi; @@ -29,17 +27,17 @@ public PineconeControlPlaneClient(String apiKey, OkHttpClient okHttpClient) { manageIndexesApi.setApiClient(apiClient); } - public IndexModel createIndex(CreateIndexRequest createIndexRequest) { + public IndexModel createIndex(CreateIndexRequest createIndexRequest) throws PineconeException { IndexModel indexModel = new IndexModel(); try { - manageIndexesApi.createIndex(createIndexRequest); + indexModel = manageIndexesApi.createIndex(createIndexRequest); } catch (ApiException apiException) { handleApiException(apiException); } return indexModel; } - public IndexModel describeIndex(String indexName) { + public IndexModel describeIndex(String indexName) throws PineconeException { IndexModel indexModel = new IndexModel(); try { indexModel = manageIndexesApi.describeIndex(indexName); @@ -49,7 +47,7 @@ public IndexModel describeIndex(String indexName) { return indexModel; } - public void configureIndex(String indexName, ConfigureIndexRequest configureIndexRequest) { + public void configureIndex(String indexName, ConfigureIndexRequest configureIndexRequest) throws PineconeException { try { manageIndexesApi.configureIndex(indexName, configureIndexRequest); } catch (ApiException apiException) { @@ -57,7 +55,7 @@ public void configureIndex(String indexName, ConfigureIndexRequest configureInde } } - public IndexList listIndexes() { + public IndexList listIndexes() throws PineconeException { IndexList indexList = new IndexList(); try { indexList = manageIndexesApi.listIndexes(); @@ -67,7 +65,7 @@ public IndexList listIndexes() { return indexList; } - public void deleteIndex(String indexName) { + public void deleteIndex(String indexName) throws PineconeException { try { manageIndexesApi.deleteIndex(indexName); } catch (ApiException apiException) { @@ -75,7 +73,45 @@ public void deleteIndex(String indexName) { } } - private void handleApiException(ApiException apiException) { + public CollectionModel createCollection(CreateCollectionRequest createCollectionRequest) throws PineconeException { + CollectionModel collection = null; + try { + collection = manageIndexesApi.createCollection(createCollectionRequest); + } catch (ApiException apiException) { + handleApiException(apiException); + } + return collection; + } + + public CollectionModel describeCollection(String collectionName) throws PineconeException { + CollectionModel collection = null; + try { + collection = manageIndexesApi.describeCollection(collectionName); + } catch (ApiException apiException) { + handleApiException(apiException); + } + return collection; + } + + public CollectionList listCollections() throws PineconeException { + CollectionList collections = null; + try { + collections = manageIndexesApi.listCollections(); + } catch (ApiException apiException) { + handleApiException(apiException); + } + return collections; + } + + public void deleteCollection(String collectionName) throws PineconeException { + try { + manageIndexesApi.deleteCollection(collectionName); + } catch (ApiException apiException) { + handleApiException(apiException); + } + } + + private void handleApiException(ApiException apiException) throws PineconeException { int statusCode = apiException.getCode(); String responseBody = apiException.getResponseBody(); FailedRequestInfo failedRequestInfo = new FailedRequestInfo(statusCode, responseBody); diff --git a/src/main/java/io/pinecone/exceptions/HttpErrorMapper.java b/src/main/java/io/pinecone/exceptions/HttpErrorMapper.java index 4a4ebb73..58c7e2c8 100644 --- a/src/main/java/io/pinecone/exceptions/HttpErrorMapper.java +++ b/src/main/java/io/pinecone/exceptions/HttpErrorMapper.java @@ -2,7 +2,7 @@ public class HttpErrorMapper { - public static void mapHttpStatusError(FailedRequestInfo failedRequestInfo) { + public static void mapHttpStatusError(FailedRequestInfo failedRequestInfo) throws PineconeException { int statusCode = failedRequestInfo.getStatus(); switch (statusCode) { case 400: