Add Pinecone Provider

.github/ISSUE_TEMPLATE/airflow_providers_bug_report.yml

@@ -87,6 +87,7 @@ body:
         - oracle
         - pagerduty
         - papermill
+        - pinecone
         - plexus
         - postgres
         - presto

CONTRIBUTING.rst

@@ -677,9 +677,9 @@ doc, doc_gen, docker, druid, elasticsearch, exasol, facebook, ftp, gcp, gcp_api,
 github_enterprise, google, google_auth, grpc, hashicorp, hdfs, hive, http, imap, influxdb, jdbc,
 jenkins, kerberos, kubernetes, ldap, leveldb, microsoft.azure, microsoft.mssql, microsoft.psrp,
 microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas, openlineage, opensearch, opsgenie,
-oracle, otel, pagerduty, pandas, papermill, password, pinot, plexus, postgres, presto, rabbitmq,
-redis, s3, s3fs, salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack, smtp,
-snowflake, spark, sqlite, ssh, statsd, tableau, tabular, telegram, trino, vertica, virtualenv,
+oracle, otel, pagerduty, pandas, papermill, password, pinecone, pinot, plexus, postgres, presto,
+rabbitmq, redis, s3, s3fs, salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack,
+smtp, snowflake, spark, sqlite, ssh, statsd, tableau, tabular, telegram, trino, vertica, virtualenv,
 weaviate, webhdfs, winrm, yandex, zendesk
   .. END EXTRAS HERE
 

INSTALL

@@ -104,9 +104,9 @@ doc, doc_gen, docker, druid, elasticsearch, exasol, facebook, ftp, gcp, gcp_api,
 github_enterprise, google, google_auth, grpc, hashicorp, hdfs, hive, http, imap, influxdb, jdbc,
 jenkins, kerberos, kubernetes, ldap, leveldb, microsoft.azure, microsoft.mssql, microsoft.psrp,
 microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas, openlineage, opensearch, opsgenie,
-oracle, otel, pagerduty, pandas, papermill, password, pinot, plexus, postgres, presto, rabbitmq,
-redis, s3, s3fs, salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack, smtp,
-snowflake, spark, sqlite, ssh, statsd, tableau, tabular, telegram, trino, vertica, virtualenv,
+oracle, otel, pagerduty, pandas, papermill, password, pinecone, pinot, plexus, postgres, presto,
+rabbitmq, redis, s3, s3fs, salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack,
+smtp, snowflake, spark, sqlite, ssh, statsd, tableau, tabular, telegram, trino, vertica, virtualenv,
 weaviate, webhdfs, winrm, yandex, zendesk
 # END EXTRAS HERE
 

airflow/providers/pinecone/CHANGELOG.rst

@@ -0,0 +1,26 @@
+ .. 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.
+
+``apache-airflow-providers-pinecone``
+
+Changelog
+---------
+
+1.0.0
+.....
+
+Initial version of the provider.

airflow/providers/pinecone/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/pinecone/hooks/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/pinecone/hooks/pinecone.py

@@ -0,0 +1,128 @@
+# 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.
+
+"""Hook for Pinecone."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import pinecone
+
+from airflow.hooks.base import BaseHook
+
+if TYPE_CHECKING:
+    from pinecone.core.client.models import UpsertResponse
+
+
+class PineconeHook(BaseHook):
+    """
+    Interact with Pinecone. This hook uses the Pinecone conn_id.
+
+    :param conn_id: Optional, default connection id is `pinecone_default`. The connection id to use when
+        connecting to Pinecone.
+    """
+
+    conn_name_attr = "conn_id"
+    default_conn_name = "pinecone_default"
+    conn_type = "pinecone"
+    hook_name = "Pinecone"
+
+    @staticmethod
+    def get_connection_form_widgets() -> dict[str, Any]:
+        """Returns connection widgets to add to connection form."""
+        from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
+        from flask_babel import lazy_gettext
+        from wtforms import StringField
+
+        return {
+            "log_level": StringField(lazy_gettext("Log Level"), widget=BS3TextFieldWidget(), default=None),
+            "project_name": StringField(
+                lazy_gettext("Project Name"),
+                widget=BS3TextFieldWidget(),
+            ),
+        }
+
+    @classmethod
+    def get_ui_field_behaviour(cls) -> dict[str, Any]:
+        """Returns custom field behaviour."""
+        return {
+            "hidden_fields": ["port", "schema"],
+            "relabeling": {"login": "Pinecone Environment", "password": "Pinecone API key"},
+        }
+
+    def __init__(self, conn_id: str = default_conn_name) -> None:
+        self.conn_id = conn_id
+        self.get_conn()
+
+    def get_conn(self) -> None:
+        pinecone_connection = self.get_connection(self.conn_id)
+        api_key = pinecone_connection.password
+        pinecone_environment = pinecone_connection.login
+        pinecone_host = pinecone_connection.host
+        extras = pinecone_connection.extra_dejson
+        pinecone_project_name = extras.get("project_name")
+        log_level = extras.get("log_level", None)
+        pinecone.init(
+            api_key=api_key,
+            environment=pinecone_environment,
+            host=pinecone_host,
+            project_name=pinecone_project_name,
+            log_level=log_level,
+        )
+
+    def test_connection(self) -> tuple[bool, str]:
+        try:
+            pinecone.list_indexes()
+            return True, "Connection established"
+        except Exception as e:
+            return False, str(e)
+
+    @staticmethod
+    def upsert(
+        index_name: str,
+        vectors: list[Any],
+        namespace: str = "",
+        batch_size: int | None = None,
+        show_progress: bool = True,
+        **kwargs: Any,
+    ) -> UpsertResponse:
+        """
+        The upsert operation writes vectors into a namespace.
+
+        If a new value is upserted for an existing vector id, it will overwrite the previous value.
+
+         .. seealso:: https://docs.pinecone.io/reference/upsert
+
+        To upsert in parallel follow
+
+        .. seealso:: https://docs.pinecone.io/docs/insert-data#sending-upserts-in-parallel
+
+        :param index_name: The name of the index to describe.
+        :param vectors: A list of vectors to upsert.
+        :param namespace: The namespace to write to. If not specified, the default namespace - "" is used.
+        :param batch_size: The number of vectors to upsert in each batch.
+        :param show_progress: Whether to show a progress bar using tqdm. Applied only
+            if batch_size is provided.
+        """
+        index = pinecone.Index(index_name)
+        return index.upsert(
+            vectors=vectors,
+            namespace=namespace,
+            batch_size=batch_size,
+            show_progress=show_progress,
+            **kwargs,
+        )

airflow/providers/pinecone/operators/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/pinecone/operators/pinecone.py

@@ -0,0 +1,83 @@
+# 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.
+
+from __future__ import annotations
+
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.pinecone.hooks.pinecone import PineconeHook
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class PineconeIngestOperator(BaseOperator):
+    """
+    Ingest vector embeddings into Pinecone.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:PineconeIngestOperator`
+
+    :param conn_id: The connection id to use when connecting to Pinecone.
+    :param index_name: Name of the Pinecone index.
+    :param input_vectors: Data to be ingested, in the form of a list of tuples where each tuple
+        contains (id, vector_embedding, metadata).
+    :param namespace: The namespace to write to. If not specified, the default namespace is used.
+    :param batch_size: The number of vectors to upsert in each batch.
+    :param upsert_kwargs: .. seealso:: https://docs.pinecone.io/reference/upsert
+    """
+
+    template_fields: Sequence[str] = ("index_name", "input_vectors", "namespace")
+
+    def __init__(
+        self,
+        *,
+        conn_id: str = PineconeHook.default_conn_name,
+        index_name: str,
+        input_vectors: list[tuple],
+        namespace: str = "",
+        batch_size: int | None = None,
+        upsert_kwargs: dict | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self.upsert_kwargs = upsert_kwargs or {}
+        super().__init__(**kwargs)
+        self.conn_id = conn_id
+        self.index_name = index_name
+        self.namespace = namespace
+        self.batch_size = batch_size
+        self.input_vectors = input_vectors
+
+    @cached_property
+    def hook(self) -> PineconeHook:
+        """Return an instance of the PineconeHook."""
+        return PineconeHook(conn_id=self.conn_id)
+
+    def execute(self, context: Context) -> None:
+        """Ingest data into Pinecone using the PineconeHook."""
+        self.hook.upsert(
+            index_name=self.index_name,
+            vectors=self.input_vectors,
+            namespace=self.namespace,
+            batch_size=self.batch_size,
+            **self.upsert_kwargs,
+        )
+
+        self.log.info("Successfully ingested data into Pinecone index %s.", self.index_name)

airflow/providers/pinecone/provider.yaml

@@ -0,0 +1,54 @@
+# 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-name: apache-airflow-providers-pinecone
+
+name: Pinecone
+
+description: |
+    `Pinecone <https://docs.pinecone.io/docs/overview>`__
+
+suspended: false
+
+versions:
+  - 1.0.0
+
+integrations:
+  - integration-name: Pinecone
+    external-doc-url: https://docs.pinecone.io/docs/overview
+    how-to-guide:
+      - /docs/apache-airflow-providers-pinecone/operators/pinecone.rst
+    tags: [software]
+
+dependencies:
+  - apache-airflow>=2.5.0
+  - pinecone-client>=2.2.4
+
+hooks:
+  - integration-name: Pinecone
+    python-modules:
+      - airflow.providers.pinecone.hooks.pinecone
+
+connection-types:
+  - hook-class-name: airflow.providers.pinecone.hooks.pinecone.PineconeHook
+    connection-type: pinecone
+
+operators:
+  - integration-name: Pinecone
+    python-modules:
+      - airflow.providers.pinecone.operators.pinecone

docs/apache-airflow-providers-pinecone/changelog.rst

@@ -0,0 +1,19 @@
+
+ .. 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.
+
+.. include:: ../../airflow/providers/pinecone/CHANGELOG.rst