feature to allow to link two extraction graphs

crates/indexify_internal_api/src/lib.rs

@@ -85,6 +85,35 @@ impl ExtractionGraphBuilder {
     }
 }
 
+#[derive(Debug, Clone, Serialize, PartialEq, Eq, Deserialize, Hash)]
+pub struct ExtractionGraphNode {
+    pub namespace: String,
+    pub graph_name: String,
+    pub source: ContentSource,
+}
+
+/// Links a node in extraction graph to another graph.
+/// All top level policies in the linked graph will be applied to the
+/// specified node.
+#[derive(Debug, Clone, Serialize, PartialEq, Eq, Deserialize)]
+pub struct ExtractionGraphLink {
+    pub node: ExtractionGraphNode,
+    pub graph_name: String,
+}
+
+impl From<indexify_coordinator::LinkExtractionGraphsRequest> for ExtractionGraphLink {
+    fn from(value: indexify_coordinator::LinkExtractionGraphsRequest) -> Self {
+        Self {
+            node: ExtractionGraphNode {
+                namespace: value.namespace,
+                graph_name: value.source_graph_name,
+                source: value.content_source.into(),
+            },
+            graph_name: value.linked_graph_name,
+        }
+    }
+}
+
 pub type IndexName = String;
 pub type IndexId = String;
 
@@ -794,7 +823,7 @@ impl Default for ContentMetadataId {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub enum ContentSource {
     Ingestion,
     ExtractionPolicyName(ExtractionPolicyName),

crates/indexify_proto/src/indexify_coordinator.rs

@@ -834,6 +834,21 @@ pub struct ExecutorsHeartbeatRequest {
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ExecutorsHeartbeatResponse {}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct LinkExtractionGraphsRequest {
+    #[prost(string, tag = "1")]
+    pub namespace: ::prost::alloc::string::String,
+    #[prost(string, tag = "2")]
+    pub source_graph_name: ::prost::alloc::string::String,
+    #[prost(string, tag = "3")]
+    pub content_source: ::prost::alloc::string::String,
+    #[prost(string, tag = "4")]
+    pub linked_graph_name: ::prost::alloc::string::String,
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct LinkExtractionGraphsResponse {}
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
 #[repr(i32)]
 pub enum TaskOutcome {
@@ -2075,6 +2090,36 @@ pub mod coordinator_service_client {
                 );
             self.inner.unary(req, path, codec).await
         }
+        pub async fn link_extraction_graphs(
+            &mut self,
+            request: impl tonic::IntoRequest<super::LinkExtractionGraphsRequest>,
+        ) -> std::result::Result<
+            tonic::Response<super::LinkExtractionGraphsResponse>,
+            tonic::Status,
+        > {
+            self.inner
+                .ready()
+                .await
+                .map_err(|e| {
+                    tonic::Status::new(
+                        tonic::Code::Unknown,
+                        format!("Service was not ready: {}", e.into()),
+                    )
+                })?;
+            let codec = tonic::codec::ProstCodec::default();
+            let path = http::uri::PathAndQuery::from_static(
+                "/indexify_coordinator.CoordinatorService/LinkExtractionGraphs",
+            );
+            let mut req = request.into_request();
+            req.extensions_mut()
+                .insert(
+                    GrpcMethod::new(
+                        "indexify_coordinator.CoordinatorService",
+                        "LinkExtractionGraphs",
+                    ),
+                );
+            self.inner.unary(req, path, codec).await
+        }
     }
 }
 /// Generated server implementations.
@@ -2332,6 +2377,13 @@ pub mod coordinator_service_server {
             tonic::Response<super::ExecutorsHeartbeatResponse>,
             tonic::Status,
         >;
+        async fn link_extraction_graphs(
+            &self,
+            request: tonic::Request<super::LinkExtractionGraphsRequest>,
+        ) -> std::result::Result<
+            tonic::Response<super::LinkExtractionGraphsResponse>,
+            tonic::Status,
+        >;
     }
     #[derive(Debug)]
     pub struct CoordinatorServiceServer<T: CoordinatorService> {
@@ -4113,6 +4165,56 @@ pub mod coordinator_service_server {
                     };
                     Box::pin(fut)
                 }
+                "/indexify_coordinator.CoordinatorService/LinkExtractionGraphs" => {
+                    #[allow(non_camel_case_types)]
+                    struct LinkExtractionGraphsSvc<T: CoordinatorService>(pub Arc<T>);
+                    impl<
+                        T: CoordinatorService,
+                    > tonic::server::UnaryService<super::LinkExtractionGraphsRequest>
+                    for LinkExtractionGraphsSvc<T> {
+                        type Response = super::LinkExtractionGraphsResponse;
+                        type Future = BoxFuture<
+                            tonic::Response<Self::Response>,
+                            tonic::Status,
+                        >;
+                        fn call(
+                            &mut self,
+                            request: tonic::Request<super::LinkExtractionGraphsRequest>,
+                        ) -> Self::Future {
+                            let inner = Arc::clone(&self.0);
+                            let fut = async move {
+                                <T as CoordinatorService>::link_extraction_graphs(
+                                        &inner,
+                                        request,
+                                    )
+                                    .await
+                            };
+                            Box::pin(fut)
+                        }
+                    }
+                    let accept_compression_encodings = self.accept_compression_encodings;
+                    let send_compression_encodings = self.send_compression_encodings;
+                    let max_decoding_message_size = self.max_decoding_message_size;
+                    let max_encoding_message_size = self.max_encoding_message_size;
+                    let inner = self.inner.clone();
+                    let fut = async move {
+                        let inner = inner.0;
+                        let method = LinkExtractionGraphsSvc(inner);
+                        let codec = tonic::codec::ProstCodec::default();
+                        let mut grpc = tonic::server::Grpc::new(codec)
+                            .apply_compression_config(
+                                accept_compression_encodings,
+                                send_compression_encodings,
+                            )
+                            .apply_max_message_size_config(
+                                max_decoding_message_size,
+                                max_encoding_message_size,
+                            );
+                        let res = grpc.unary(method, req).await;
+                        Ok(res)
+                    };
+                    Box::pin(fut)
+                }
                 _ => {
                     Box::pin(async move {
                         Ok(

protos/coordinator_service.proto

@@ -73,6 +73,8 @@ service CoordinatorService {
     rpc UpdateLabels(UpdateLabelsRequest) returns (UpdateLabelsResponse) {}
 
     rpc ExecutorsHeartbeat(ExecutorsHeartbeatRequest) returns (ExecutorsHeartbeatResponse) {}
+
+    rpc LinkExtractionGraphs(LinkExtractionGraphsRequest) returns (LinkExtractionGraphsResponse) {}
 }
 
 message GetContentMetadataRequest {
@@ -575,3 +577,12 @@ message ExecutorsHeartbeatRequest {
 }
 
 message ExecutorsHeartbeatResponse {}
+
+message LinkExtractionGraphsRequest {
+    string namespace = 1;
+    string source_graph_name = 2;
+    string content_source = 3;
+    string linked_graph_name = 4;
+}
+
+message LinkExtractionGraphsResponse {}

src/api.rs

@@ -15,6 +15,13 @@ use utoipa::{IntoParams, ToSchema};
 
 use crate::{api_utils, metadata_storage, vectordbs};
 
+#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
+pub struct ExtractionGraphLink {
+    pub source_graph_name: String,
+    pub content_source: String,
+    pub linked_graph_name: String,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
 pub struct ExtractionGraph {
     pub id: String,