Skip to content

Commit

Permalink
add pinecone source connector to read data from pinecone
Browse files Browse the repository at this point in the history
  • Loading branch information
nianliuu committed Oct 25, 2024
1 parent 10c37ac commit e70d1d3
Show file tree
Hide file tree
Showing 17 changed files with 1,199 additions and 0 deletions.
1 change: 1 addition & 0 deletions config/plugin_config
Original file line number Diff line number Diff line change
Expand Up @@ -92,3 +92,4 @@ connector-sls
connector-qdrant
connector-typesense
connector-cdc-opengauss
connector-pinecone
41 changes: 41 additions & 0 deletions docs/en/connector-v2/source/Pinecone.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
# Pinecone

> Pinecone source connector
## Description

This Pinecone source connector reads data from Pinecone
## Key Features

- [x] [batch](../../concept/connector-v2-features.md)
- [x] [exactly-once](../../concept/connector-v2-features.md)
- [ ] [column projection](../../concept/connector-v2-features.md)

## Data Type Mapping

| Pinecone Data Type | SeaTunnel Data Type |
|---------------------|---------------------|
| FLOAT_VECTOR | FLOAT_VECTOR |
| SPARSE_FLOAT_VECTOR | SPARSE_FLOAT_VECTOR |

## Source Options

| Name | Type | Required | Default | Description |
|------------|---------|----------|---------|------------------------------------|
| api_key | String | Yes | - | Pinecone token for authentication. |
| index | String | Yes | - | Pinecone index name |
| batch_size | Integer | No | 100 | Batch size to read |

## Task Example

```bash
source {
Milvus {
token = "*****"
index = "movie_reviews"
}
}
```

## Changelog

1 change: 1 addition & 0 deletions plugin-mapping.properties
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,7 @@ seatunnel.sink.Sls = connector-sls
seatunnel.source.Typesense = connector-typesense
seatunnel.sink.Typesense = connector-typesense
seatunnel.source.Opengauss-CDC = connector-cdc-opengauss
seatunnel.source.Pinecone = connector-pinecone

seatunnel.transform.Sql = seatunnel-transforms-v2
seatunnel.transform.FieldMapper = seatunnel-transforms-v2
Expand Down
85 changes: 85 additions & 0 deletions seatunnel-connectors-v2/connector-pinecone/pom.xml
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements. See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You under the Apache License, Version 2.0
(the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>org.apache.seatunnel</groupId>
<artifactId>seatunnel-connectors-v2</artifactId>
<version>${revision}</version>
</parent>
<artifactId>connector-pinecone</artifactId>
<name>SeaTunnel : Connectors V2 : Pinecone</name>

<properties>
<!-- Specify the gRPC version -->
<grpc.version>1.59.1</grpc.version>
<protobuf.version>3.25.1</protobuf.version>
</properties>

<dependencyManagement>
<dependencies>
<dependency>
<groupId>com.google.protobuf</groupId>
<artifactId>protobuf-java</artifactId>
<version>${protobuf.version}</version>
</dependency>
<dependency>
<groupId>io.grpc</groupId>
<artifactId>grpc-netty</artifactId>
<version>${grpc.version}</version>
</dependency>
<dependency>
<groupId>io.grpc</groupId>
<artifactId>grpc-netty-shaded</artifactId>
<version>${grpc.version}</version>
</dependency>
<dependency>
<groupId>io.grpc</groupId>
<artifactId>grpc-protobuf</artifactId>
<version>${grpc.version}</version>
</dependency>
<dependency>
<groupId>io.grpc</groupId>
<artifactId>grpc-stub</artifactId>
<version>${grpc.version}</version>
</dependency>
</dependencies>
</dependencyManagement>

<dependencies>
<dependency>
<groupId>io.pinecone</groupId>
<artifactId>pinecone-client</artifactId>
<version>2.1.0</version>
</dependency>
<dependency>
<groupId>com.google.code.gson</groupId>
<artifactId>gson</artifactId>
<version>2.10.1</version>
</dependency>
<dependency>
<groupId>com.google.protobuf</groupId>
<artifactId>protobuf-java</artifactId>
<version>3.25.1</version>
</dependency>
</dependencies>

</project>
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.apache.seatunnel.connectors.pinecone.config;

import org.apache.seatunnel.api.configuration.Option;
import org.apache.seatunnel.api.configuration.Options;

public class PineconeSourceConfig {
public static final String CONNECTOR_IDENTITY = "Pinecone";

public static final Option<String> API_KEY =
Options.key("api_key")
.stringType()
.noDefaultValue()
.withDescription("Pinecone token for authentication");

public static final Option<String> INDEX =
Options.key("index")
.stringType()
.noDefaultValue()
.withDescription("Pinecone index name");

public static final Option<Integer> BATCH_SIZE =
Options.key("batch_size")
.intType()
.defaultValue(100)
.withDescription("writer batch size");
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.apache.seatunnel.connectors.pinecone.exception;

import org.apache.seatunnel.common.exception.SeaTunnelErrorCode;

import lombok.Getter;

@Getter
public enum PineconeConnectionErrorCode implements SeaTunnelErrorCode {
SOURCE_TABLE_SCHEMA_IS_NULL("PINECONE-01", "Source table schema is null"),
READ_DATA_FAIL("PINECONE-02", "Read data fail");

private final String code;
private final String description;

PineconeConnectionErrorCode(String code, String description) {
this.code = code;
this.description = description;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.apache.seatunnel.connectors.pinecone.exception;

import org.apache.seatunnel.common.exception.SeaTunnelErrorCode;
import org.apache.seatunnel.common.exception.SeaTunnelRuntimeException;

public class PineconeConnectorException extends SeaTunnelRuntimeException {
public PineconeConnectorException(SeaTunnelErrorCode seaTunnelErrorCode) {
super(seaTunnelErrorCode, seaTunnelErrorCode.getErrorMessage());
}

public PineconeConnectorException(SeaTunnelErrorCode seaTunnelErrorCode, Throwable cause) {
super(seaTunnelErrorCode, seaTunnelErrorCode.getErrorMessage(), cause);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.apache.seatunnel.connectors.pinecone.source;

import org.apache.seatunnel.api.configuration.ReadonlyConfig;
import org.apache.seatunnel.api.source.Boundedness;
import org.apache.seatunnel.api.source.SeaTunnelSource;
import org.apache.seatunnel.api.source.SourceReader;
import org.apache.seatunnel.api.source.SourceSplitEnumerator;
import org.apache.seatunnel.api.source.SupportColumnProjection;
import org.apache.seatunnel.api.source.SupportParallelism;
import org.apache.seatunnel.api.table.catalog.CatalogTable;
import org.apache.seatunnel.api.table.catalog.TablePath;
import org.apache.seatunnel.api.table.type.SeaTunnelRow;
import org.apache.seatunnel.connectors.pinecone.config.PineconeSourceConfig;
import org.apache.seatunnel.connectors.pinecone.utils.PineconeUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class PineconeSource
implements SeaTunnelSource<SeaTunnelRow, PineconeSourceSplit, PineconeSourceState>,
SupportParallelism,
SupportColumnProjection {
private final ReadonlyConfig config;
private final Map<TablePath, CatalogTable> sourceTables;

public PineconeSource(ReadonlyConfig config) {
this.config = config;
PineconeUtils pineconeUtils = new PineconeUtils(config);
this.sourceTables = pineconeUtils.getSourceTables();
}

/**
* Get the boundedness of this source.
*
* @return the boundedness of this source.
*/
@Override
public Boundedness getBoundedness() {
return Boundedness.BOUNDED;
}

/**
* Create source reader, used to produce data.
*
* @param readerContext reader context.
* @return source reader.
* @throws Exception when create reader failed.
*/
@Override
public SourceReader<SeaTunnelRow, PineconeSourceSplit> createReader(
SourceReader.Context readerContext) throws Exception {
return new PineconeSourceReader(readerContext, config, sourceTables);
}

@Override
public List<CatalogTable> getProducedCatalogTables() {
return new ArrayList<>(sourceTables.values());
}

/**
* Create source split enumerator, used to generate splits. This method will be called only once
* when start a source.
*
* @param enumeratorContext enumerator context.
* @return source split enumerator.
* @throws Exception when create enumerator failed.
*/
@Override
public SourceSplitEnumerator<PineconeSourceSplit, PineconeSourceState> createEnumerator(
SourceSplitEnumerator.Context<PineconeSourceSplit> enumeratorContext) throws Exception {
return new PineconeSourceSplitEnumertor(enumeratorContext, config, sourceTables, null);
}

/**
* Create source split enumerator, used to generate splits. This method will be called when
* restore from checkpoint.
*
* @param enumeratorContext enumerator context.
* @param checkpointState checkpoint state.
* @return source split enumerator.
* @throws Exception when create enumerator failed.
*/
@Override
public SourceSplitEnumerator<PineconeSourceSplit, PineconeSourceState> restoreEnumerator(
SourceSplitEnumerator.Context<PineconeSourceSplit> enumeratorContext,
PineconeSourceState checkpointState)
throws Exception {
return new PineconeSourceSplitEnumertor(
enumeratorContext, config, sourceTables, checkpointState);
}

/**
* Returns a unique identifier among same factory interfaces.
*
* <p>For consistency, an identifier should be declared as one lower case word (e.g. {@code
* kafka}). If multiple factories exist for different versions, a version should be appended
* using "-" (e.g. {@code elasticsearch-7}).
*/
@Override
public String getPluginName() {
return PineconeSourceConfig.CONNECTOR_IDENTITY;
}
}
Loading

0 comments on commit e70d1d3

Please sign in to comment.