Add collections operations to PineconeControlPlaneClient with integration tests

.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 }}
-        
+
 

build.gradle

@@ -1,3 +1,6 @@
+import org.gradle.api.tasks.testing.logging.TestExceptionFormat;
+import org.gradle.api.tasks.testing.logging.TestLogEvent;
+
 plugins {
     id 'com.github.johnrengelman.shadow' version '6.1.0'
     id 'java-library'
@@ -93,13 +96,73 @@ tasks.named('build') {
 
 test {
     useJUnitPlatform()
+
+    testLogging {
+        events = [TestLogEvent.FAILED, TestLogEvent.PASSED, TestLogEvent.SKIPPED, TestLogEvent.STANDARD_OUT]
+        exceptionFormat = TestExceptionFormat.FULL
+        showExceptions = true
+        showCauses = true
+        showStackTraces = true
+
+        debugOptions {
+            events = [
+                    TestLogEvent.STARTED,
+                    TestLogEvent.FAILED,
+                    TestLogEvent.PASSED,
+                    TestLogEvent.SKIPPED,
+                    TestLogEvent.STANDARD_ERROR,
+                    TestLogEvent.STANDARD_OUT
+            ]
+            exceptionFormat = TestExceptionFormat.FULL
+        }
+
+        afterSuite { desc, result ->
+            if (!desc.parent) {
+                def output = "Results: ${result.resultType} (${result.testCount}) tests, ${result.successfulTestCount} successful, ${result.failedTestCount} failed, ${result.skippedTestCount} skipped."
+                def startItem = "| ", endItem = " |"
+                def repeatLength = startItem.length() + output.length() + endItem.length()
+                println('\n' + ('-' * repeatLength) + '\n' + startItem + output + endItem + '\n' + ('-' * repeatLength))
+            }
+        }
+    }
 }
 
 task integrationTest(type: Test) {
     useJUnitPlatform()
     testClassesDirs = sourceSets.integrationTest.output.classesDirs
     classpath = sourceSets.integrationTest.runtimeClasspath
     outputs.upToDateWhen { false }
+
+
+    testLogging {
+        events = [TestLogEvent.FAILED, TestLogEvent.PASSED, TestLogEvent.SKIPPED, TestLogEvent.STANDARD_OUT]
+        exceptionFormat = TestExceptionFormat.FULL
+        showExceptions = true
+        showCauses = true
+        showStackTraces = true
+
+        debugOptions {
+            events = [
+                    TestLogEvent.STARTED,
+                    TestLogEvent.FAILED,
+                    TestLogEvent.PASSED,
+                    TestLogEvent.SKIPPED,
+                    TestLogEvent.STANDARD_ERROR,
+                    TestLogEvent.STANDARD_OUT
+            ]
+            exceptionFormat = TestExceptionFormat.FULL
+
+        }
+
+        afterSuite { desc, result ->
+            if (!desc.parent) {
+                def output = "Results: ${result.resultType} (${result.testCount}) tests, ${result.successfulTestCount} successful, ${result.failedTestCount} failed, ${result.skippedTestCount} skipped."
+                def startItem = "| ", endItem = " |"
+                def repeatLength = startItem.length() + output.length() + endItem.length()
+                println('\n' + ('-' * repeatLength) + '\n' + startItem + output + endItem + '\n' + ('-' * repeatLength))
+            }
+        }
+    }
 }
 
 // Configure Auto Relocation

src/integration/java/io/pinecone/helpers/AssertRetry.java

@@ -1,18 +1,21 @@
 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 InterruptedException, PineconeException {
         int retryCount = 0;
+        int delayCount = delay;
         boolean success = false;
 
         while (retryCount < maxRetry && !success) {
@@ -21,14 +24,14 @@ public static void assertWithRetry(AssertionRunnable assertionRunnable, int back
                 success = true;
             } catch (AssertionError | ExecutionException | IOException e) {
                 retryCount++;
-                delay*=backOff;
-                Thread.sleep(delay);
+                Thread.sleep(delayCount);
+                delayCount*=backOff;
             }
         }
     }
 
     @FunctionalInterface
     public interface AssertionRunnable {
-        void run() throws AssertionError, ExecutionException, InterruptedException, IOException;
+        void run() throws AssertionError, ExecutionException, InterruptedException, IOException, PineconeException;
     }
 }

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<String> upsertIds, S
                 .build();
     }
 
+    public static UpsertRequest buildRequiredUpsertRequestByDimension(List<String> upsertIds, int dimension, String namespace) {
+        if (upsertIds.isEmpty()) upsertIds = Arrays.asList("v1", "v2", "v3");
+
+        List<Vector> 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<String, List<String>> createAndGetMetadataMap() {
 
         return metadataMap;
     }
+
+    public static ArrayList<Float> generateVectorValuesByDimension(int dimension) {
+        ArrayList<Float> values = new ArrayList<>();
+        Random random = new Random();
+
+        for (int i = 0; i < dimension; i++) {
+            values.add(random.nextFloat());
+        }
+
+        return values;
+    }
 }

src/integration/java/io/pinecone/helpers/IndexManager.java

@@ -1,12 +1,15 @@
 package io.pinecone.helpers;
 
 import io.pinecone.*;
+import io.pinecone.exceptions.PineconeException;
 import org.openapitools.client.model.*;
 
 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;
@@ -45,9 +48,11 @@ private static String findIndexWithDimensionAndType(IndexList indexList, int dim
         List<IndexModel> 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,6 +84,76 @@ 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) {
+                System.out.println("Index " + indexName + " not ready after " + waitedTimeMs + "ms");
+                break;
+            }
+            if (index.getStatus().getReady()) {
+                System.out.println("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) {
+                System.out.println("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];