Resolve merge conflict in src/base/json_schema.rs

Author: aryasoni98

python/cocoindex/flow.py

@@ -405,6 +405,7 @@ def export(
         /,
         *,
         primary_key_fields: Sequence[str],
+        attachments: Sequence[op.TargetAttachmentSpec] = (),
         vector_indexes: Sequence[index.VectorIndexDef] = (),
         vector_index: Sequence[tuple[str, index.VectorSimilarityMetric]] = (),
         setup_by_user: bool = False,
@@ -436,6 +437,10 @@ def export(
             target_name,
             _spec_kind(target_spec),
             dump_engine_object(target_spec),
+            [
+                {"kind": _spec_kind(att), **dump_engine_object(att)}
+                for att in attachments
+            ],
             dump_engine_object(index_options),
             self._engine_data_collector,
             setup_by_user,

python/cocoindex/op.py

@@ -46,6 +46,7 @@ class OpCategory(Enum):
     SOURCE = "source"
     TARGET = "target"
     DECLARATION = "declaration"
+    TARGET_ATTACHMENT = "target_attachment"
 
 
 @dataclass_transform()
@@ -81,6 +82,10 @@ class TargetSpec(metaclass=SpecMeta, category=OpCategory.TARGET):  # pylint: dis
     """A target spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
 
 
+class TargetAttachmentSpec(metaclass=SpecMeta, category=OpCategory.TARGET_ATTACHMENT):  # pylint: disable=too-few-public-methods
+    """A target attachment spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
+
+
 class DeclarationSpec(metaclass=SpecMeta, category=OpCategory.DECLARATION):  # pylint: disable=too-few-public-methods
     """A declaration spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
 

python/cocoindex/targets/_engine_builtin_specs.py

@@ -16,6 +16,14 @@ class Postgres(op.TargetSpec):
     table_name: str | None = None
 
 
+class PostgresSqlAttachment(op.TargetAttachmentSpec):
+    """Attachment to execute specified SQL statements for Postgres targets."""
+
+    name: str
+    setup_sql: str
+    teardown_sql: str | None = None
+
+
 @dataclass
 class QdrantConnection:
     """Connection spec for Qdrant."""

src/base/json_schema.rs

@@ -395,7 +395,6 @@ pub fn build_json_schema(
 #[cfg(test)]
 mod tests {
     use super::*;
-<<<<<<< HEAD
     use crate::base::schema::*;
     use expect_test::expect;
     use serde_json::json;
@@ -1368,11 +1367,7 @@ mod tests {
               "type": "string"
             }"#]]
         .assert_eq(&serde_json::to_string_pretty(&json_schema).unwrap());
-=======
-    use crate::base::schema::{
-        BasicValueType, EnrichedValueType, FieldSchema, StructSchema, ValueType,
-    };
-    use std::sync::Arc;
+    }
 
     #[test]
     fn test_description_concatenation() {
@@ -1421,6 +1416,5 @@ mod tests {
         } else {
             panic!("No description found in the schema");
         }
->>>>>>> ef87b68dac10345e67361c4d5d6a3edf44a4af0f
     }
 }

src/base/spec.rs

@@ -493,6 +493,10 @@ impl fmt::Display for IndexOptions {
 pub struct ExportOpSpec {
     pub collector_name: FieldName,
     pub target: OpSpec,
+
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub attachments: Vec<OpSpec>,
+
     pub index_options: IndexOptions,
     pub setup_by_user: bool,
 }

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/builder/flow_builder.rs

@@ -549,12 +549,13 @@ impl FlowBuilder {
         Ok(())
     }
 
-    #[pyo3(signature = (name, kind, op_spec, index_options, input, setup_by_user=false))]
+    #[pyo3(signature = (name, kind, op_spec, attachments, index_options, input, setup_by_user=false))]
     pub fn export(
         &mut self,
         name: String,
         kind: String,
         op_spec: py::Pythonized<serde_json::Map<String, serde_json::Value>>,
+        attachments: py::Pythonized<Vec<spec::OpSpec>>,
         index_options: py::Pythonized<spec::IndexOptions>,
         input: &DataCollector,
         setup_by_user: bool,
@@ -574,6 +575,7 @@ impl FlowBuilder {
             spec: spec::ExportOpSpec {
                 collector_name: input.name.clone(),
                 target: spec,
+                attachments: attachments.into_inner(),
                 index_options: index_options.into_inner(),
                 setup_by_user,
             },

src/execution/db_tracking_setup.rs

@@ -103,6 +103,8 @@ pub struct TrackingTableSetupChange {
     pub legacy_source_state_table_names: BTreeSet<String>,
 
     pub source_names_need_state_cleanup: BTreeMap<i32, BTreeSet<String>>,
+
+    has_state_change: bool,
 }
 
 impl TrackingTableSetupChange {
@@ -136,6 +138,7 @@ impl TrackingTableSetupChange {
                 legacy_source_state_table_names,
                 min_existing_version_id,
                 source_names_need_state_cleanup,
+                has_state_change: existing.has_state_diff(desired, |v| v),
             })
         } else {
             None
@@ -148,6 +151,7 @@ impl TrackingTableSetupChange {
         ResourceSetupInfo {
             key: (),
             state: self.desired_state.clone(),
+            has_tracked_state_change: self.has_state_change,
             description: "Internal Storage for Tracking".to_string(),
             setup_change: Some(self),
             legacy_key: None,

src/ops/factory_bases.rs

@@ -374,7 +374,7 @@ pub struct TypedResourceSetupChangeItem<'a, F: TargetFactoryBase + ?Sized> {
 }
 
 #[async_trait]
-pub trait TargetFactoryBase: TargetFactory + Send + Sync + 'static {
+pub trait TargetFactoryBase: Send + Sync + 'static {
     type Spec: DeserializeOwned + Send + Sync;
     type DeclarationSpec: DeserializeOwned + Send + Sync;
 
@@ -635,3 +635,101 @@ fn from_json_combined_state<T: Debug + Clone + Serialize + DeserializeOwned>(
         legacy_state_key: existing_states.legacy_state_key,
     })
 }
+
+pub struct TypedTargetAttachmentState<F: TargetSpecificAttachmentFactoryBase + ?Sized> {
+    pub setup_key: F::SetupKey,
+    pub setup_state: F::SetupState,
+}
+
+/// A factory for target-specific attachments.
+#[async_trait]
+pub trait TargetSpecificAttachmentFactoryBase: Send + Sync + 'static {
+    type TargetKey: Debug + Clone + Serialize + DeserializeOwned + Eq + Hash + Send + Sync;
+    type TargetSpec: DeserializeOwned + Send + Sync;
+    type Spec: DeserializeOwned + Send + Sync;
+    type SetupKey: Debug + Clone + Serialize + DeserializeOwned + Eq + Hash + Send + Sync;
+    type SetupState: Debug + Clone + Serialize + DeserializeOwned + Send + Sync;
+    type SetupChange: interface::AttachmentSetupChange + Send + Sync;
+
+    fn name(&self) -> &str;
+
+    fn get_state(
+        &self,
+        target_name: &str,
+        target_spec: &Self::TargetSpec,
+        attachment_spec: Self::Spec,
+    ) -> Result<TypedTargetAttachmentState<Self>>;
+
+    async fn diff_setup_states(
+        &self,
+        target_key: &Self::TargetKey,
+        attachment_key: &Self::SetupKey,
+        new_state: Option<Self::SetupState>,
+        existing_states: setup::CombinedState<Self::SetupState>,
+        context: &interface::FlowInstanceContext,
+    ) -> Result<Option<Self::SetupChange>>;
+
+    /// Deserialize the setup key from a JSON value.
+    /// You can override this method to provide a custom deserialization logic, e.g. to perform backward compatible deserialization.
+    fn deserialize_setup_key(key: serde_json::Value) -> Result<Self::SetupKey> {
+        Ok(utils::deser::from_json_value(key)?)
+    }
+
+    fn register(self, registry: &mut ExecutorFactoryRegistry) -> Result<()>
+    where
+        Self: Sized,
+    {
+        registry.register(
+            self.name().to_string(),
+            ExecutorFactory::TargetAttachment(Arc::new(self)),
+        )
+    }
+}
+
+#[async_trait]
+impl<T: TargetSpecificAttachmentFactoryBase> TargetAttachmentFactory for T {
+    fn normalize_setup_key(&self, key: &serde_json::Value) -> Result<serde_json::Value> {
+        let key: T::SetupKey = Self::deserialize_setup_key(key.clone())?;
+        Ok(serde_json::to_value(key)?)
+    }
+
+    fn get_state(
+        &self,
+        target_name: &str,
+        target_spec: &serde_json::Map<String, serde_json::Value>,
+        attachment_spec: serde_json::Value,
+    ) -> Result<interface::TargetAttachmentState> {
+        let state = TargetSpecificAttachmentFactoryBase::get_state(
+            self,
+            target_name,
+            &utils::deser::from_json_value(serde_json::Value::Object(target_spec.clone()))?,
+            utils::deser::from_json_value(attachment_spec)?,
+        )?;
+        Ok(interface::TargetAttachmentState {
+            setup_key: serde_json::to_value(state.setup_key)?,
+            setup_state: serde_json::to_value(state.setup_state)?,
+        })
+    }
+
+    async fn diff_setup_states(
+        &self,
+        target_key: &serde_json::Value,
+        attachment_key: &serde_json::Value,
+        new_state: Option<serde_json::Value>,
+        existing_states: setup::CombinedState<serde_json::Value>,
+        context: &interface::FlowInstanceContext,
+    ) -> Result<Option<Box<dyn AttachmentSetupChange + Send + Sync>>> {
+        let setup_change = self
+            .diff_setup_states(
+                &utils::deser::from_json_value(target_key.clone())?,
+                &utils::deser::from_json_value(attachment_key.clone())?,
+                new_state
+                    .map(|v| utils::deser::from_json_value(v))
+                    .transpose()?,
+                from_json_combined_state(existing_states)?,
+                context,
+            )
+            .await?;
+        Ok(setup_change.map(|s| Box::new(s) as Box<dyn AttachmentSetupChange + Send + Sync>))
+    }
+}