test: Add Pinecone source integration tests (#46)

* chore(deps): Use Anush008/chroma-go until amikos-tech/chroma-go # 221 is merged

Signed-off-by: Anush008 <[email protected]>

* chore: Revert back to amikos-tech/chroma-go

Signed-off-by: Anush008 <[email protected]>

* chore: tests

Signed-off-by: Anush008 <[email protected]>

* Chroma test

Signed-off-by: Anush008 <[email protected]>

* chore: Addressed lint issues

Signed-off-by: Anush008 <[email protected]>

* chore: No close client

Signed-off-by: Anush008 <[email protected]>

* test: Pinecone

Signed-off-by: Anush008 <[email protected]>

* chore: go mod tidy

Signed-off-by: Anush008 <[email protected]>

* test: Sparse vectors

Signed-off-by: Anush008 <[email protected]>

* chore: Run as from shell

Signed-off-by: Anush008 <[email protected]>

* chore: Test as binary

Signed-off-by: Anush008 <[email protected]>

* test: Updated implementation

Signed-off-by: Anush008 <[email protected]>

---------

Signed-off-by: Anush008 <[email protected]>

Author: Anush008

README.md

@@ -102,7 +102,8 @@ Migrate data from a **Pinecone** database to **Qdrant**:
 
 ```bash
 migration pinecone \
-    --pinecone.host 'https://example-index-12345.svc.region.pinecone.io' \
+    --pinecone.index-host 'https://example-index.svc.region.pinecone.io' \
+    --pinecone.index-name 'example-index' \
     --pinecone.api-key 'optional-pinecone-api-key' \
     --qdrant.url 'https://example.cloud-region.cloud-provider.cloud.qdrant.io:6334' \
     --qdrant.api-key 'optional-qdrant-api-key' \
@@ -114,17 +115,19 @@ With Docker:
 
 ```bash
 docker run --net=host --rm -it registry.cloud.qdrant.io/library/qdrant-migration pinecone \
-    --pinecone.host 'https://example-index-12345.svc.region.pinecone.io' \
+    --pinecone.index-host 'https://example-index.svc.region.pinecone.io' \
     ...
 ```
 
 #### Pinecone Options
 
 | Flag                            | Description                                                     |
 | ------------------------------- | --------------------------------------------------------------- |
-| `--pinecone.api-key`            | Pinecone API key for authentication                             |
-| `--pinecone.host`               | Pinecone index host URL (e.g., `https://your-pinecone-url`)     |
-| `--pinecone.namespace`          | Namespace of the partition to migrate                           |
+| `--pinecone.index-name`              | Pinecone index name.                                       |
+| `--pinecone.index-host`         | Pinecone index host URL (e.g., `https://your-pinecone-url`)     |
+| `--pinecone.api-key`            | Pinecone API key for authentication.                            |
+| `--pinecone.namespace`          | Namespace of the partition to migrate. Optional.                |
+| `--pinecone.service-host`       | Pinecone service host URL. Optional.                            |
 
 #### Qdrant Options
 

cmd/migrate_from_pinecone.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"os"
 	"os/signal"
-	"strings"
 	"syscall"
 
 	"github.com/pinecone-io/go-pinecone/v3/pinecone"
@@ -29,22 +28,13 @@ type MigrateFromPineconeCmd struct {
 	targetTLS  bool
 }
 
-func stripScheme(input string) string {
-	if idx := strings.Index(input, "://"); idx != -1 {
-		return input[idx+3:]
-	}
-	return input
-}
-
 func (r *MigrateFromPineconeCmd) Parse() error {
 	var err error
 	r.targetHost, r.targetPort, r.targetTLS, err = parseQdrantUrl(r.Qdrant.Url)
 	if err != nil {
 		return fmt.Errorf("failed to parse target URL: %w", err)
 	}
 
-	r.Pinecone.Host = stripScheme(r.Pinecone.Host)
-
 	return nil
 }
 
@@ -88,7 +78,7 @@ func (r *MigrateFromPineconeCmd) Run(globals *Globals) error {
 		return fmt.Errorf("error preparing target collection: %w", err)
 	}
 
-	displayMigrationStart("pinecone", r.Pinecone.Host, r.Qdrant.Collection)
+	displayMigrationStart("pinecone", r.Pinecone.IndexHost, r.Qdrant.Collection)
 
 	err = r.migrateData(ctx, sourceIndexConn, targetClient, sourcePointCount)
 	if err != nil {
@@ -110,14 +100,15 @@ func (r *MigrateFromPineconeCmd) Run(globals *Globals) error {
 
 func (r *MigrateFromPineconeCmd) connectToPinecone() (*pinecone.Client, *pinecone.IndexConnection, error) {
 	client, err := pinecone.NewClient(pinecone.NewClientParams{
+		Host:   r.Pinecone.ServiceHost,
 		ApiKey: r.Pinecone.APIKey,
 	})
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to create Pinecone client: %w", err)
 	}
 
 	indexConn, err := client.Index(pinecone.NewIndexConnParams{
-		Host:      r.Pinecone.Host,
+		Host:      r.Pinecone.IndexHost,
 		Namespace: r.Pinecone.Namespace,
 	})
 	if err != nil {
@@ -158,14 +149,14 @@ func (r *MigrateFromPineconeCmd) prepareTargetCollection(ctx context.Context, so
 
 	var foundIndex *pinecone.Index
 	for i := range indexes {
-		if indexes[i].Host == r.Pinecone.Host {
+		if indexes[i].Name == r.Pinecone.IndexName {
 			foundIndex = indexes[i]
 			break
 		}
 	}
 
 	if foundIndex == nil {
-		return fmt.Errorf("index %q not found in Pinecone", r.Pinecone.Host)
+		return fmt.Errorf("index %q not found in Pinecone", r.Pinecone.IndexName)
 	}
 
 	distanceMapping := map[pinecone.IndexMetric]qdrant.Distance{
@@ -213,7 +204,7 @@ func (r *MigrateFromPineconeCmd) migrateData(ctx context.Context, sourceIndexCon
 	offsetCount := uint64(0)
 
 	if !r.Migration.Restart {
-		id, offsetStored, err := commons.GetStartOffset(ctx, r.Migration.OffsetsCollection, targetClient, r.Pinecone.Host)
+		id, offsetStored, err := commons.GetStartOffset(ctx, r.Migration.OffsetsCollection, targetClient, r.Pinecone.IndexHost)
 		if err != nil {
 			return fmt.Errorf("failed to get start offset: %w", err)
 		}
@@ -239,7 +230,7 @@ func (r *MigrateFromPineconeCmd) migrateData(ctx context.Context, sourceIndexCon
 		}
 
 		if len(listRes.VectorIds) < 1 {
-			return fmt.Errorf("pinecone.ListVectors returned no IDs")
+			break
 		}
 
 		ids := make([]string, 0, len(listRes.VectorIds))
@@ -298,7 +289,7 @@ func (r *MigrateFromPineconeCmd) migrateData(ctx context.Context, sourceIndexCon
 		if listRes.NextPaginationToken != nil {
 			offsetCount += uint64(len(targetPoints))
 			offsetId = qdrant.NewID(*listRes.NextPaginationToken)
-			err = commons.StoreStartOffset(ctx, r.Migration.OffsetsCollection, targetClient, r.Pinecone.Host, offsetId, offsetCount)
+			err = commons.StoreStartOffset(ctx, r.Migration.OffsetsCollection, targetClient, r.Pinecone.IndexHost, offsetId, offsetCount)
 			if err != nil {
 				return fmt.Errorf("failed to store offset: %w", err)
 			}

cmd/root.go

@@ -42,3 +42,14 @@ func Execute(projectVersion, projectBuild string) {
 		ctx.Exit(1)
 	}
 }
+
+func NewParser(args []string) (*kong.Context, error) {
+	cli := &CLI{}
+
+	parser, err := kong.New(cli, kong.Bind(&cli.Globals))
+	if err != nil {
+		return nil, err
+	}
+
+	return parser.Parse(args)
+}

go.mod

@@ -18,6 +18,7 @@ require (
 	github.com/weaviate/weaviate v1.27.0
 	github.com/weaviate/weaviate-go-client/v4 v4.16.1
 	google.golang.org/grpc v1.72.2
+	google.golang.org/protobuf v1.36.6
 )
 
 require (
@@ -173,7 +174,6 @@ require (
 	google.golang.org/genproto v0.0.0-20250512202823-5a2f75b736a9 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20250512202823-5a2f75b736a9 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20250512202823-5a2f75b736a9 // indirect
-	google.golang.org/protobuf v1.36.6 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect

integration_tests/commons_test.go

@@ -15,6 +15,7 @@ const (
 	qdrantAPIKey       = "00000000"
 	totalEntries       = 100
 	dimension          = 384
+	idField            = "__id__"
 )
 
 func randFloat32Values(n int) []float32 {
@@ -25,6 +26,14 @@ func randFloat32Values(n int) []float32 {
 	return values
 }
 
+func randIndices(n int) []uint32 {
+	indices := make([]uint32, n)
+	for i := range indices {
+		indices[i] = rand.Uint32()
+	}
+	return indices
+}
+
 func runMigrationBinary(t *testing.T, args []string) {
 	binaryPath := filepath.Join(t.TempDir(), "migration")
 	cmd := exec.Command("go", "build", "-o", binaryPath, "main.go")

integration_tests/image_test.go

@@ -69,3 +69,24 @@ func weaviateContainer(ctx context.Context, t *testing.T) testcontainers.Contain
 
 	return container
 }
+
+func pineconeContainer(ctx context.Context, t *testing.T) testcontainers.Container {
+
+	req := testcontainers.ContainerRequest{
+		Image:        "ghcr.io/pinecone-io/pinecone-local:latest",
+		ExposedPorts: []string{"5081/tcp", "5082/tcp"},
+		Env: map[string]string{
+			"PORT": "5081",
+		},
+		WaitingFor: wait.ForAll(wait.ForListeningPort("5081/tcp").WithStartupTimeout(30*time.Second),
+			wait.ForListeningPort("5082/tcp").WithStartupTimeout(30*time.Second)),
+	}
+
+	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
+		ContainerRequest: req,
+		Started:          true,
+	})
+	require.NoError(t, err)
+
+	return container
+}

integration_tests/migrate_from_chroma_test.go

@@ -18,7 +18,6 @@ const (
 	sourceField      = "source"
 	denseVectorField = "dense_vector"
 	distance         = "euclid"
-	idField          = "__id__"
 )
 
 func TestMigrateFromChroma(t *testing.T) {