Add API version and custom headers

codegen/build-oas.sh

@@ -1,13 +1,20 @@
 #!/bin/bash
 
+version=$1
+
+if [ -z "$version" ]; then
+	echo "Version is required"
+	exit 1
+fi
+
 pushd codegen/apis
 	just build
 popd
 
 outdir="openapi"
 
 docker run --rm -v $(pwd):/workspace openapitools/openapi-generator-cli:v7.6.0 generate \
-	--input-spec /workspace/codegen/apis/_build/2024-07/control_2024-07.oas.yaml \
+	--input-spec /workspace/codegen/apis/_build/$version/control_$version.oas.yaml \
 	--generator-name rust \
 	--output /workspace/$outdir \
 	--additional-properties "packageVersion=0.0.1"

codegen/build-proto.sh

@@ -1,14 +1,21 @@
 #!/bin/bash
 
+version=$1
+
 SCRIPT_DIR=$(dirname "$(realpath "$0")")
 outdir="protos"
 
+if [ -z "$version" ]; then
+	echo "Version is required"
+	exit 1
+fi
+
 OUT_DIR=$SCRIPT_DIR/../$outdir
 
 pushd $SCRIPT_DIR/apis
 	just build
 popd
 
 pushd $SCRIPT_DIR/proto_build
-	cargo run -- $OUT_DIR
+	cargo run -- $OUT_DIR $version
 popd

codegen/proto_build/src/main.rs

@@ -5,14 +5,18 @@ fn main() -> Result<(), Box<dyn Error>> {
     let args: Vec<String> = std::env::args().collect();
 
     let out_dir: &str;
-    if args.len() > 1 {
+    let version: &str;
+    if args.len() == 3 {
         out_dir = &args[1];
+        version = &args[2];
         println!("OUT_DIR: {:?}", out_dir);
+        println!("version: {:?}", version);
     } else {
-        return Err("missing out_dir argument".into());
+        return Err("Required 2 arguments: out_dir version".into());
     }
 
-    let proto_path: &Path = "../apis/_build/2024-07/data_2024-07.proto".as_ref();
+    let proto_path = format!("../apis/_build/{version}/data_{version}.proto");
+    let proto_path: &Path = proto_path.as_ref();
 
     // directory the main .proto file resides in
     let proto_dir = proto_path

justfile

@@ -1,6 +1,17 @@
+api_version := "2024-07"
+
+# Build the OpenAPI and Protobuf definitions in `codegen/apis`
+build-apis:
+  cd codegen/apis && just build
+
+# Generate the control plane OpenAPI code based on the yaml files in `codegen/apis/_build`
 build-openapi:
-  ./codegen/build-oas.sh
-  cargo build -p openapi
+  ./codegen/build-oas.sh {{api_version}}
 
+# Generate the data plane protobuf code based on the yaml files in `codegen/apis/_build`
 build-proto:
-  ./codegen/build-proto.sh
+  ./codegen/build-proto.sh {{api_version}}
+
+# Generate all OpenAPI and protobuf code
+build-client: build-apis build-openapi build-proto
+  echo "/// Pinecone API version\npub const API_VERSION: &str = \"{{api_version}}\";" > pinecone_sdk/src/version.rs

pinecone_sdk/src/lib.rs

@@ -14,3 +14,6 @@ pub mod pinecone;
 
 /// Utility modules.
 pub mod utils;
+
+/// Version information.
+pub mod version;

pinecone_sdk/src/pinecone/mod.rs

@@ -6,6 +6,9 @@
 use serde_json;
 use std::collections::HashMap;
 
+/// The `PINECONE_API_VERSION_KEY` is the key for the Pinecone API version header.
+pub const PINECONE_API_VERSION_KEY: &str = "x-pinecone-api-version";
+
 /// Control plane module.
 pub mod control;
 
@@ -18,7 +21,7 @@
     api_key: String,
     controller_url: String,
     additional_headers: HashMap<String, String>,
     source_tag: Option<String>,
     user_agent: Option<String>,
 }
 
@@ -74,9 +77,8 @@
         let controller_host = control_plane_host.unwrap_or(env_controller);
 
         // get additional headers
-        let additional_headers = match additional_headers {
-            Some(headers) => headers,
-            None => match std::env::var("PINECONE_ADDITIONAL_HEADERS") {
+        let additional_headers =
+            additional_headers.unwrap_or(match std::env::var("PINECONE_ADDITIONAL_HEADERS") {
                 Ok(headers) => match serde_json::from_str(&headers) {
                     Ok(headers) => headers,
                     Err(_) => {
@@ -87,8 +89,7 @@
                     }
                 },
                 Err(_) => HashMap::new(),
-            },
-        };
+            });
 
         // create config
         let config = Config {
@@ -113,13 +114,37 @@
 
     /// Returns the OpenAPI configuration object.
     pub fn openapi_config(&self) -> Configuration {
+        // create reqwest headers
+        // init an empty one if self.additional_headers is empty, otherwise make it
+        let mut headers = reqwest::header::HeaderMap::new();
+        for (k, v) in self.additional_headers.iter() {
+            let key = reqwest::header::HeaderName::from_bytes(k.as_bytes()).unwrap();
+            let value = reqwest::header::HeaderValue::from_str(v).unwrap();
+            headers.insert(key, value);
+        }
+
+        // add X-Pinecone-Api-Version header if not present
+        if !headers.contains_key(PINECONE_API_VERSION_KEY) {
+            headers.insert(
+                reqwest::header::HeaderName::from_static(PINECONE_API_VERSION_KEY),
+                reqwest::header::HeaderValue::from_static(crate::version::API_VERSION),
+            );
+        }
+
+        // create reqwest client
+        let client = reqwest::Client::builder()
+            .default_headers(headers)
+            .build()
+            .unwrap_or(reqwest::Client::new());
+
         Configuration {
             base_path: self.controller_url.clone(),
             user_agent: self.user_agent.clone(),
             api_key: Some(ApiKey {
                 prefix: None,
                 key: self.api_key.clone(),
             }),
+            client,
             ..Default::default()
         }
     }
@@ -161,13 +186,9 @@
         let mock_controller_host = "mock-arg-controller-host";
 
         temp_env::with_var("PINECONE_API_KEY", Some(mock_api_key), || {
-            let pinecone = PineconeClient::new(
-                None,
-                Some(mock_controller_host),
-                Some(HashMap::new()),
-                None,
-            )
-            .expect("Expected to successfully create Pinecone instance");
+            let pinecone =
+                PineconeClient::new(None, Some(mock_controller_host), Some(HashMap::new()), None)
+                    .expect("Expected to successfully create Pinecone instance");
 
             assert_eq!(pinecone.api_key, mock_api_key);
             assert_eq!(pinecone.controller_url, mock_controller_host);
@@ -187,13 +208,11 @@
         let mock_controller_host = "mock-arg-controller-host";
 
         temp_env::with_var_unset("PINECONE_API_KEY", || {
-            let pinecone = PineconeClient::new(
-                None,
-                Some(mock_controller_host),
-                Some(HashMap::new()),
-                None,
-            )
-            .expect_err("Expected to fail creating Pinecone instance due to missing API key");
+            let pinecone =
+                PineconeClient::new(None, Some(mock_controller_host), Some(HashMap::new()), None)
+                    .expect_err(
+                        "Expected to fail creating Pinecone instance due to missing API key",
+                    );
 
             assert!(matches!(pinecone, PineconeError::APIKeyMissingError { .. }));
         });
@@ -227,13 +246,9 @@
             "PINECONE_CONTROLLER_HOST",
             Some(mock_controller_host),
             || {
-                let pinecone = PineconeClient::new(
-                    Some(mock_api_key),
-                    None,
-                    Some(HashMap::new()),
-                    None,
-                )
-                .expect("Expected to successfully create Pinecone instance with env host");
+                let pinecone =
+                    PineconeClient::new(Some(mock_api_key), None, Some(HashMap::new()), None)
+                        .expect("Expected to successfully create Pinecone instance with env host");
 
                 assert_eq!(pinecone.controller_url, mock_controller_host);
             },
@@ -301,13 +316,11 @@
             "PINECONE_ADDITIONAL_HEADERS",
             Some(serde_json::to_string(&mock_headers).unwrap().as_str()),
             || {
-                let pinecone = PineconeClient::new(
-                    Some(mock_api_key),
-                    Some(mock_controller_host),
-                    None,
-                    None,
-                )
-                .expect("Expected to successfully create Pinecone instance with env headers");
+                let pinecone =
+                    PineconeClient::new(Some(mock_api_key), Some(mock_controller_host), None, None)
+                        .expect(
+                            "Expected to successfully create Pinecone instance with env headers",
+                        );
 
                 assert_eq!(pinecone.additional_headers, mock_headers);
             },
@@ -322,13 +335,11 @@
         let mock_controller_host = "mock-arg-controller-host";
 
         temp_env::with_var("PINECONE_ADDITIONAL_HEADERS", Some("invalid-json"), || {
-            let pinecone = PineconeClient::new(
-                Some(mock_api_key),
-                Some(mock_controller_host),
-                None,
-                None,
-            )
-            .expect_err("Expected to fail creating Pinecone instance due to invalid headers");
+            let pinecone =
+                PineconeClient::new(Some(mock_api_key), Some(mock_controller_host), None, None)
+                    .expect_err(
+                        "Expected to fail creating Pinecone instance due to invalid headers",
+                    );
 
             assert!(matches!(
                 pinecone,
@@ -377,10 +388,7 @@
         temp_env::with_vars(
             [
                 ("PINECONE_API_KEY", Some(mock_env_api_key)),
-                (
-                    "PINECONE_CONTROLLER_HOST",
-                    Some(mock_env_controller_host),
-                ),
+                ("PINECONE_CONTROLLER_HOST", Some(mock_env_controller_host)),
                 (
                     "PINECONE_ADDITIONAL_HEADERS",
                     Some(serde_json::to_string(&mock_env_headers).unwrap().as_str()),

pinecone_sdk/src/version.rs

@@ -0,0 +1,2 @@
+/// Pinecone API version
+pub const API_VERSION: &str = "2024-07";

pinecone_sdk/tests/integration_test.rs

@@ -5,7 +5,7 @@ use pinecone_sdk::pinecone::data::{Kind, Metadata, Namespace, SparseValues, Valu
 use pinecone_sdk::pinecone::PineconeClient;
 use pinecone_sdk::utils::errors::PineconeError;
 use rand::Rng;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashMap};
 use std::time::Duration;
 use std::vec;
 
@@ -435,7 +435,7 @@ async fn test_list_collections() -> Result<(), PineconeError> {
 }
 
 #[tokio::test]
-async fn test_delete_collection_err() -> Result<(), PineconeError> {
+async fn test_delete_collection_invalid_collection() -> Result<(), PineconeError> {
     let pinecone =
         PineconeClient::new(None, None, None, None).expect("Failed to create Pinecone instance");
 
@@ -447,6 +447,27 @@ async fn test_delete_collection_err() -> Result<(), PineconeError> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_list_collections_invalid_api_version() -> Result<(), PineconeError> {
+    let headers: HashMap<String, String> = [(
+        pinecone_sdk::pinecone::PINECONE_API_VERSION_KEY.to_string(),
+        "invalid".to_string(),
+    )]
+    .iter()
+    .cloned()
+    .collect();
+
+    let pinecone = PineconeClient::new(None, None, Some(headers), None)
+        .expect("Failed to create Pinecone instance");
+
+    let _ = pinecone
+        .list_collections()
+        .await
+        .expect_err("Expected to fail listing collections due to invalid api version");
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_index() -> Result<(), PineconeError> {
     let pinecone = PineconeClient::new(None, None, None, None).unwrap();
@@ -982,10 +1003,7 @@ async fn test_fetch_vectors() -> Result<(), PineconeError> {
     std::thread::sleep(std::time::Duration::from_secs(5));
 
     let fetch_response = index
-        .fetch(
-            &["1".to_string(), "2".to_string()],
-            namespace,
-        )
+        .fetch(&["1".to_string(), "2".to_string()], namespace)
         .await
         .expect("Failed to fetch vectors");
 
@@ -1034,7 +1052,10 @@ async fn test_fetch_no_match() -> Result<(), PineconeError> {
         .expect("Failed to target index");
 
     let fetch_response = index
-        .fetch(&["invalid-id1".to_string(), "invalid-id2".to_string()], &Default::default())
+        .fetch(
+            &["invalid-id1".to_string(), "invalid-id2".to_string()],
+            &Default::default(),
+        )
         .await
         .expect("Failed to fetch vectors");