[ENH] Add rust protobufs and conversion. Add build.rs, protobufs, and conversions

.github/workflows/chroma-worker-test.yml

@@ -19,6 +19,8 @@ jobs:
     steps:
     - name: Checkout
       uses: actions/checkout@v3
+    - name: Install Protoc
+      uses: arduino/setup-protoc@v2
     - name: Build
       run: cargo build --verbose
     - name: Test

rust/worker/Cargo.toml

@@ -5,6 +5,8 @@ edition = "2021"
 
 [dependencies]
 tonic = "0.10"
+prost = "0.12"
+prost-types = "0.12"
 tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] }
 tokio-util = "0.7.10"
 rand = "0.8.5"
@@ -16,6 +18,7 @@ serde = { version = "1.0.193", features = ["derive"] }
 num_cpus = "1.16.0"
 murmur3 = "0.5.2"
 thiserror = "1.0.50"
+num-bigint = "0.4.4"
 
 [build-dependencies]
 tonic-build = "0.10"

rust/worker/Dockerfile

@@ -1,8 +1,14 @@
 FROM rust:1.74.1 as builder
 
-WORKDIR /chroma
+WORKDIR /chroma/
 COPY . .
 
+ENV PROTOC_ZIP=protoc-25.1-linux-x86_64.zip
+RUN curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v25.1/$PROTOC_ZIP \
+    && unzip -o $PROTOC_ZIP -d /usr/local bin/protoc \
+    && unzip -o $PROTOC_ZIP -d /usr/local 'include/*' \
+    && rm -f $PROTOC_ZIP
+
 RUN cargo build
 
 # For now this runs cargo test since we have no main binary

rust/worker/build.rs

@@ -0,0 +1,10 @@
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    tonic_build::configure().compile(
+        &[
+            "../../idl/chromadb/proto/chroma.proto",
+            "../../idl/chromadb/proto/coordinator.proto",
+        ],
+        &["../../idl/"],
+    )?;
+    Ok(())
+}

rust/worker/src/config.rs

@@ -84,6 +84,9 @@ impl RootConfig {
 /// ## Description of parameters
 /// - my_ip: The IP address of the worker service. Used for memberlist assignment. Must be provided
 /// - num_indexing_threads: The number of indexing threads to use. If not provided, defaults to the number of cores on the machine.
+/// - pulsar_tenant: The pulsar tenant to use. Must be provided.
+/// - pulsar_namespace: The pulsar namespace to use. Must be provided.
+/// - assignment_policy: The assignment policy to use. Must be provided.
 /// # Notes
 /// In order to set the enviroment variables, you must prefix them with CHROMA_WORKER__<FIELD_NAME>.
 /// For example, to set my_ip, you would set CHROMA_WORKER__MY_IP.

rust/worker/src/lib.rs

@@ -1,3 +1,8 @@
 mod assignment;
 mod config;
 mod errors;
+mod types;
+
+mod chroma_proto {
+    tonic::include_proto!("chroma");
+}

rust/worker/src/types/collection.rs

@@ -0,0 +1,88 @@
+use super::{Metadata, MetadataValueConversionError};
+use crate::{
+    chroma_proto,
+    errors::{ChromaError, ErrorCodes},
+};
+use thiserror::Error;
+use uuid::Uuid;
+
+#[derive(Debug, PartialEq)]
+pub(crate) struct Collection {
+    pub(crate) id: Uuid,
+    pub(crate) name: String,
+    pub(crate) topic: String,
+    pub(crate) metadata: Option<Metadata>,
+    pub(crate) dimension: Option<i32>,
+    pub(crate) tenant: String,
+    pub(crate) database: String,
+}
+
+#[derive(Error, Debug)]
+pub(crate) enum CollectionConversionError {
+    #[error("Invalid UUID")]
+    InvalidUuid,
+    #[error(transparent)]
+    MetadataValueConversionError(#[from] MetadataValueConversionError),
+}
+
+impl ChromaError for CollectionConversionError {
+    fn code(&self) -> crate::errors::ErrorCodes {
+        match self {
+            CollectionConversionError::InvalidUuid => ErrorCodes::InvalidArgument,
+            CollectionConversionError::MetadataValueConversionError(e) => e.code(),
+        }
+    }
+}
+
+impl TryFrom<chroma_proto::Collection> for Collection {
+    type Error = CollectionConversionError;
+
+    fn try_from(proto_collection: chroma_proto::Collection) -> Result<Self, Self::Error> {
+        let collection_uuid = match Uuid::try_parse(&proto_collection.id) {
+            Ok(uuid) => uuid,
+            Err(_) => return Err(CollectionConversionError::InvalidUuid),
+        };
+        let collection_metadata: Option<Metadata> = match proto_collection.metadata {
+            Some(proto_metadata) => match proto_metadata.try_into() {
+                Ok(metadata) => Some(metadata),
+                Err(e) => return Err(CollectionConversionError::MetadataValueConversionError(e)),
+            },
+            None => None,
+        };
+        Ok(Collection {
+            id: collection_uuid,
+            name: proto_collection.name,
+            topic: proto_collection.topic,
+            metadata: collection_metadata,
+            dimension: proto_collection.dimension,
+            tenant: proto_collection.tenant,
+            database: proto_collection.database,
+        })
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_collection_try_from() {
+        let proto_collection = chroma_proto::Collection {
+            id: "00000000-0000-0000-0000-000000000000".to_string(),
+            name: "foo".to_string(),
+            topic: "bar".to_string(),
+            metadata: None,
+            dimension: None,
+            tenant: "baz".to_string(),
+            database: "qux".to_string(),
+        };
+        let converted_collection: Collection = proto_collection.try_into().unwrap();
+        assert_eq!(converted_collection.id, Uuid::nil());
+        assert_eq!(converted_collection.name, "foo".to_string());
+        assert_eq!(converted_collection.topic, "bar".to_string());
+        assert_eq!(converted_collection.metadata, None);
+        assert_eq!(converted_collection.dimension, None);
+        assert_eq!(converted_collection.tenant, "baz".to_string());
+        assert_eq!(converted_collection.database, "qux".to_string());
+    }
+}