feat(attachment): support target attachment in the core setup engine

Author: georgeh0

python/cocoindex/flow.py

@@ -405,7 +405,7 @@ def export(
         /,
         *,
         primary_key_fields: Sequence[str],
-        attachments: Sequence[op.TargetAttachmentSpec],
+        attachments: Sequence[op.TargetAttachmentSpec] = (),
         vector_indexes: Sequence[index.VectorIndexDef] = (),
         vector_index: Sequence[tuple[str, index.VectorSimilarityMetric]] = (),
         setup_by_user: bool = False,

src/builder/analyzer.rs

@@ -1,5 +1,7 @@
 use crate::builder::exec_ctx::AnalyzedSetupState;
-use crate::ops::{get_function_factory, get_source_factory, get_target_factory};
+use crate::ops::{
+    get_attachment_factory, get_function_factory, get_source_factory, get_target_factory,
+};
 use crate::prelude::*;
 
 use super::plan::*;
@@ -913,6 +915,27 @@ impl AnalyzerContext {
                 let op_name = export_op.name.clone();
                 let export_target_factory = export_op_group.target_factory.clone();
 
+                let attachments = export_op
+                    .spec
+                    .attachments
+                    .iter()
+                    .map(|attachment| {
+                        let attachment_factory = get_attachment_factory(&attachment.kind)?;
+                        let attachment_state = attachment_factory.get_state(
+                            &op_name,
+                            &export_op.spec.target.spec,
+                            serde_json::Value::Object(attachment.spec.clone()),
+                        )?;
+                        Ok((
+                            interface::AttachmentSetupKey(
+                                attachment.kind.clone(),
+                                attachment_state.setup_key,
+                            ),
+                            attachment_state.setup_state,
+                        ))
+                    })
+                    .collect::<Result<IndexMap<_, _>>>()?;
+
                 let export_op_ss = exec_ctx::AnalyzedTargetSetupState {
                     target_kind: target_kind.to_string(),
                     setup_key: data_coll_output.setup_key,
@@ -925,6 +948,7 @@ impl AnalyzerContext {
                             .map(|field| field.value_type.typ.clone())
                             .collect::<Box<[_]>>(),
                     ),
+                    attachments,
                 };
                 targets_analyzed_ss[*idx] = Some(export_op_ss);
 
@@ -956,6 +980,7 @@ impl AnalyzerContext {
                 desired_setup_state,
                 setup_by_user: false,
                 key_type: None,
+                attachments: IndexMap::new(),
             };
             declarations_analyzed_ss.push(decl_ss);
         }

src/builder/exec_ctx.rs

@@ -26,6 +26,8 @@ pub struct AnalyzedTargetSetupState {
     pub setup_by_user: bool,
     /// None for declarations.
     pub key_type: Option<Box<[schema::ValueType]>>,
+
+    pub attachments: IndexMap<interface::AttachmentSetupKey, serde_json::Value>,
 }
 
 pub struct AnalyzedSetupState {
@@ -176,6 +178,7 @@ fn build_export_op_exec_ctx(
     } else {
         max_schema_version_id + 1
     };
+
     match target_states.entry(resource_id) {
         indexmap::map::Entry::Occupied(entry) => {
             api_bail!(
@@ -194,6 +197,7 @@ fn build_export_op_exec_ctx(
                     key_type: analyzed_target_ss.key_type.clone(),
                 },
                 state: analyzed_target_ss.desired_setup_state.clone(),
+                attachments: analyzed_target_ss.attachments.clone(),
             });
         }
     }

src/ops/interface.rs

@@ -322,6 +322,12 @@ pub struct TargetAttachmentState {
 }
 
 #[async_trait]
+pub trait AttachmentSetupChangeAction {
+    fn describe_change(&self) -> String;
+
+    async fn apply_change(&self) -> Result<()>;
+}
+
 pub trait TargetAttachmentFactory: Send + Sync {
     /// Normalize the key. e.g. the JSON format may change (after code change, e.g. new optional field or field ordering), even if the underlying value is not changed.
     /// This should always return the canonical serialized form.
@@ -330,22 +336,17 @@ pub trait TargetAttachmentFactory: Send + Sync {
     fn get_state(
         &self,
         target_name: &str,
-        target_spec: &spec::OpSpec,
+        target_spec: &serde_json::Map<String, serde_json::Value>,
         attachment_spec: serde_json::Value,
     ) -> Result<TargetAttachmentState>;
 
-    fn describe(&self, key: &serde_json::Value, state: &serde_json::Value) -> Result<String>;
-
-    fn has_update(&self, old: &serde_json::Value, new: &serde_json::Value) -> bool {
-        old != new
-    }
-
-    async fn apply_setup_change(
+    /// Should return Some if and only if any changes are needed.
+    fn diff_setup_states(
         &self,
         key: &serde_json::Value,
-        old: Option<serde_json::Value>,
-        new: Option<serde_json::Value>,
-    ) -> Result<()>;
+        new_state: Option<serde_json::Value>,
+        existing_states: setup::CombinedState<serde_json::Value>,
+    ) -> Result<Option<Box<dyn AttachmentSetupChangeAction + Send + Sync>>>;
 }
 
 #[derive(Clone)]
@@ -355,3 +356,12 @@ pub enum ExecutorFactory {
     ExportTarget(Arc<dyn TargetFactory + Send + Sync>),
     TargetAttachment(Arc<dyn TargetAttachmentFactory + Send + Sync>),
 }
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
+pub struct AttachmentSetupKey(pub String, pub serde_json::Value);
+
+impl std::fmt::Display for AttachmentSetupKey {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}:{}", self.0, self.1)
+    }
+}

src/ops/registration.rs

@@ -56,6 +56,13 @@ pub fn get_optional_target_factory(
     registry.get_target(kind).cloned()
 }
 
+pub fn get_optional_attachment_factory(
+    kind: &str,
+) -> Option<std::sync::Arc<dyn super::interface::TargetAttachmentFactory + Send + Sync>> {
+    let registry = EXECUTOR_FACTORY_REGISTRY.read().unwrap();
+    registry.get_target_attachment(kind).cloned()
+}
+
 pub fn get_source_factory(
     kind: &str,
 ) -> Result<std::sync::Arc<dyn super::interface::SourceFactory + Send + Sync>> {
@@ -77,6 +84,13 @@ pub fn get_target_factory(
         .ok_or_else(|| anyhow::anyhow!("Target factory not found for op kind: {}", kind))
 }
 
+pub fn get_attachment_factory(
+    kind: &str,
+) -> Result<std::sync::Arc<dyn super::interface::TargetAttachmentFactory + Send + Sync>> {
+    get_optional_attachment_factory(kind)
+        .ok_or_else(|| anyhow::anyhow!("Attachment factory not found for op kind: {}", kind))
+}
+
 pub fn register_factory(name: String, factory: ExecutorFactory) -> Result<()> {
     let mut registry = EXECUTOR_FACTORY_REGISTRY.write().unwrap();
     registry.register(name, factory)

src/ops/registry.rs

@@ -99,4 +99,11 @@ impl ExecutorFactoryRegistry {
     ) -> Option<&Arc<dyn super::interface::TargetFactory + Send + Sync>> {
         self.target_factories.get(name)
     }
+
+    pub fn get_target_attachment(
+        &self,
+        name: &str,
+    ) -> Option<&Arc<dyn super::interface::TargetAttachmentFactory + Send + Sync>> {
+        self.target_attachment_factories.get(name)
+    }
 }