Simplify ResourceSetupStatusCheck. (#286)

Author: badmonster0

src/execution/db_tracking_setup.rs

@@ -1,6 +1,6 @@
 use crate::prelude::*;
 
-use crate::setup::{CombinedState, ResourceSetupStatusCheck, SetupChangeType};
+use crate::setup::{CombinedState, ResourceSetupInfo, ResourceSetupStatusCheck, SetupChangeType};
 use serde::{Deserialize, Serialize};
 use sqlx::PgPool;
 
@@ -79,22 +79,21 @@ impl TrackingTableSetupStatusCheck {
             source_ids_to_delete,
         }
     }
-}
 
-#[async_trait]
-impl ResourceSetupStatusCheck<(), TrackingTableSetupState> for TrackingTableSetupStatusCheck {
-    fn describe_resource(&self) -> String {
-        "Tracking Table".to_string()
-    }
-
-    fn key(&self) -> &() {
-        &()
-    }
-
-    fn desired_state(&self) -> Option<&TrackingTableSetupState> {
-        self.desired_state.as_ref()
+    pub fn into_setup_info(
+        self,
+    ) -> ResourceSetupInfo<(), TrackingTableSetupState, TrackingTableSetupStatusCheck> {
+        ResourceSetupInfo {
+            key: (),
+            state: self.desired_state.clone(),
+            description: "Tracking Table".to_string(),
+            status_check: Some(self),
+        }
     }
+}
 
+#[async_trait]
+impl ResourceSetupStatusCheck for TrackingTableSetupStatusCheck {
     fn describe_changes(&self) -> Vec<String> {
         let mut changes: Vec<String> = vec![];
         if self.desired_state.is_some() && !self.legacy_table_names.is_empty() {

src/ops/factory_bases.rs

@@ -296,14 +296,16 @@ pub trait StorageFactoryBase: ExportTargetFactory + Send + Sync + 'static {
         desired_state: Option<Self::SetupState>,
         existing_states: setup::CombinedState<Self::SetupState>,
         auth_registry: &Arc<AuthRegistry>,
-    ) -> Result<impl setup::ResourceSetupStatusCheck<Self::Key, Self::SetupState> + 'static>;
+    ) -> Result<impl setup::ResourceSetupStatusCheck + 'static>;
 
     fn check_state_compatibility(
         &self,
         desired_state: &Self::SetupState,
         existing_state: &Self::SetupState,
     ) -> Result<SetupStateCompatibility>;
 
+    fn describe_resource(&self, key: &Self::Key) -> Result<String>;
+
     fn register(self, registry: &mut ExecutorFactoryRegistry) -> Result<()>
     where
         Self: Sized,
@@ -315,62 +317,6 @@ pub trait StorageFactoryBase: ExportTargetFactory + Send + Sync + 'static {
     }
 }
 
-struct ResourceSetupStatusCheckWrapper<T: StorageFactoryBase> {
-    inner: Box<dyn setup::ResourceSetupStatusCheck<T::Key, T::SetupState> + Send + Sync>,
-    key_json: serde_json::Value,
-    state_json: Option<serde_json::Value>,
-}
-
-impl<T: StorageFactoryBase> ResourceSetupStatusCheckWrapper<T> {
-    fn new(
-        inner: Box<dyn setup::ResourceSetupStatusCheck<T::Key, T::SetupState> + Send + Sync>,
-    ) -> Result<Self> {
-        Ok(Self {
-            key_json: serde_json::to_value(inner.key())?,
-            state_json: inner
-                .desired_state()
-                .map(serde_json::to_value)
-                .transpose()?,
-            inner,
-        })
-    }
-}
-
-impl<T: StorageFactoryBase> Debug for ResourceSetupStatusCheckWrapper<T> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        std::fmt::Debug::fmt(&self.inner, f)
-    }
-}
-
-#[async_trait]
-impl<T: StorageFactoryBase> setup::ResourceSetupStatusCheck<serde_json::Value, serde_json::Value>
-    for ResourceSetupStatusCheckWrapper<T>
-{
-    fn describe_resource(&self) -> String {
-        self.inner.describe_resource()
-    }
-
-    fn key(&self) -> &serde_json::Value {
-        &self.key_json
-    }
-
-    fn desired_state(&self) -> Option<&serde_json::Value> {
-        self.state_json.as_ref()
-    }
-
-    fn describe_changes(&self) -> Vec<String> {
-        self.inner.describe_changes()
-    }
-
-    fn change_type(&self) -> setup::SetupChangeType {
-        self.inner.change_type()
-    }
-
-    async fn apply_change(&self) -> Result<()> {
-        self.inner.apply_change().await
-    }
-}
-
 impl<T: StorageFactoryBase> ExportTargetFactory for T {
     fn build(
         self: Arc<Self>,
@@ -404,11 +350,7 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
         desired_state: Option<serde_json::Value>,
         existing_states: setup::CombinedState<serde_json::Value>,
         auth_registry: &Arc<AuthRegistry>,
-    ) -> Result<
-        Box<
-            dyn setup::ResourceSetupStatusCheck<serde_json::Value, serde_json::Value> + Send + Sync,
-        >,
-    > {
+    ) -> Result<Box<dyn setup::ResourceSetupStatusCheck>> {
         let key: T::Key = serde_json::from_value(key.clone())?;
         let desired_state: Option<T::SetupState> = desired_state
             .map(|v| serde_json::from_value(v.clone()))
@@ -421,9 +363,12 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
             existing_states,
             auth_registry,
         )?;
-        Ok(Box::new(ResourceSetupStatusCheckWrapper::<T>::new(
-            Box::new(status_check),
-        )?))
+        Ok(Box::new(status_check))
+    }
+
+    fn describe_resource(&self, key: &serde_json::Value) -> Result<String> {
+        let key: T::Key = serde_json::from_value(key.clone())?;
+        StorageFactoryBase::describe_resource(self, &key)
     }
 
     fn check_state_compatibility(

src/ops/interface.rs

@@ -182,17 +182,15 @@ pub trait ExportTargetFactory {
         desired_state: Option<serde_json::Value>,
         existing_states: setup::CombinedState<serde_json::Value>,
         auth_registry: &Arc<AuthRegistry>,
-    ) -> Result<
-        Box<
-            dyn setup::ResourceSetupStatusCheck<serde_json::Value, serde_json::Value> + Send + Sync,
-        >,
-    >;
+    ) -> Result<Box<dyn setup::ResourceSetupStatusCheck>>;
 
     fn check_state_compatibility(
         &self,
         desired_state: &serde_json::Value,
         existing_state: &serde_json::Value,
     ) -> Result<SetupStateCompatibility>;
+
+    fn describe_resource(&self, key: &serde_json::Value) -> Result<String>;
 }
 
 #[derive(Clone)]

src/ops/storages/neo4j.rs

@@ -427,9 +427,6 @@ struct SetupStatusCheck {
     graph_pool: Arc<GraphPool>,
     conn_spec: Neo4jConnectionSpec,
 
-    key: GraphRelationship,
-    desired_state: Option<RelationshipSetupState>,
-
     data_clear: Option<DataClearAction>,
     rel_constraint_to_delete: IndexSet<String>,
     rel_constraint_to_create: IndexMap<String, KeyConstraint>,
@@ -473,7 +470,7 @@ impl SetupStatusCheck {
 
         let mut rel_constraint_to_create = IndexMap::new();
         let mut node_constraint_to_create = IndexMap::new();
-        if let Some(desired_state) = &desired_state {
+        if let Some(desired_state) = desired_state {
             let rel_constraint = KeyConstraint {
                 label: key.relationship.clone(),
                 field_name: desired_state.key_field_name.clone(),
@@ -485,8 +482,7 @@ impl SetupStatusCheck {
                 .map(|c| rel_constraint.field_name == c.key_field_name)
                 .unwrap_or(false)
             {
-                rel_constraint_to_create
-                    .insert(desired_state.key_constraint_name.clone(), rel_constraint);
+                rel_constraint_to_create.insert(desired_state.key_constraint_name, rel_constraint);
             }
 
             old_node_constraints.swap_remove(&desired_state.src_node.key_constraint_name);
@@ -543,8 +539,6 @@ impl SetupStatusCheck {
         Self {
             graph_pool,
             conn_spec,
-            key,
-            desired_state,
             data_clear,
             rel_constraint_to_delete,
             rel_constraint_to_create,
@@ -556,19 +550,7 @@ impl SetupStatusCheck {
 }
 
 #[async_trait]
-impl ResourceSetupStatusCheck<GraphRelationship, RelationshipSetupState> for SetupStatusCheck {
-    fn describe_resource(&self) -> String {
-        format!("Neo4j relationship {}", self.key.relationship)
-    }
-
-    fn key(&self) -> &GraphRelationship {
-        &self.key
-    }
-
-    fn desired_state(&self) -> Option<&RelationshipSetupState> {
-        self.desired_state.as_ref()
-    }
-
+impl ResourceSetupStatusCheck for SetupStatusCheck {
     fn describe_changes(&self) -> Vec<String> {
         let mut result = vec![];
         if let Some(data_clear) = &self.data_clear {
@@ -753,8 +735,7 @@ impl StorageFactoryBase for RelationshipFactory {
         desired: Option<RelationshipSetupState>,
         existing: CombinedState<RelationshipSetupState>,
         auth_registry: &Arc<AuthRegistry>,
-    ) -> Result<impl ResourceSetupStatusCheck<GraphRelationship, RelationshipSetupState> + 'static>
-    {
+    ) -> Result<impl ResourceSetupStatusCheck + 'static> {
         let conn_spec = auth_registry.get::<Neo4jConnectionSpec>(&key.connection)?;
         Ok(SetupStatusCheck::new(
             key,
@@ -777,4 +758,8 @@ impl StorageFactoryBase for RelationshipFactory {
         };
         Ok(compatibility)
     }
+
+    fn describe_resource(&self, key: &GraphRelationship) -> Result<String> {
+        Ok(format!("Neo4j relationship {}", key.relationship))
+    }
 }

src/ops/storages/postgres.rs

@@ -736,19 +736,7 @@ fn describe_index_spec(index_name: &str, index_spec: &VectorIndexDef) -> String
 }
 
 #[async_trait]
-impl setup::ResourceSetupStatusCheck<TableId, SetupState> for SetupStatusCheck {
-    fn describe_resource(&self) -> String {
-        format!("Postgres table {}", self.table_id)
-    }
-
-    fn key(&self) -> &TableId {
-        &self.table_id
-    }
-
-    fn desired_state(&self) -> Option<&SetupState> {
-        self.desired_state.as_ref()
-    }
-
+impl setup::ResourceSetupStatusCheck for SetupStatusCheck {
     fn describe_changes(&self) -> Vec<String> {
         let mut descriptions = vec![];
         if self.drop_existing {
@@ -958,7 +946,7 @@ impl StorageFactoryBase for Arc<Factory> {
         desired: Option<SetupState>,
         existing: setup::CombinedState<SetupState>,
         _auth_registry: &Arc<AuthRegistry>,
-    ) -> Result<impl setup::ResourceSetupStatusCheck<TableId, SetupState> + 'static> {
+    ) -> Result<impl setup::ResourceSetupStatusCheck + 'static> {
         Ok(SetupStatusCheck::new(self.clone(), key, desired, existing))
     }
 
@@ -987,6 +975,10 @@ impl StorageFactoryBase for Arc<Factory> {
         };
         Ok(compatibility)
     }
+
+    fn describe_resource(&self, key: &TableId) -> Result<String> {
+        Ok(format!("Postgres table {}", key.table_name))
+    }
 }
 
 impl Factory {

src/setup/db_metadata.rs

@@ -1,6 +1,6 @@
 use crate::prelude::*;
 
-use super::{ResourceSetupStatusCheck, SetupChangeType, StateChange};
+use super::{ResourceSetupInfo, ResourceSetupStatusCheck, SetupChangeType, StateChange};
 use crate::utils::db::WriteAction;
 use axum::http::StatusCode;
 use sqlx::PgPool;
@@ -294,20 +294,19 @@ pub struct MetadataTableSetup {
     pub metadata_table_missing: bool,
 }
 
-#[async_trait]
-impl ResourceSetupStatusCheck<(), ()> for MetadataTableSetup {
-    fn key(&self) -> &() {
-        &()
-    }
-
-    fn desired_state(&self) -> Option<&()> {
-        Some(&())
-    }
-
-    fn describe_resource(&self) -> String {
-        "CocoIndex Metadata Table".to_string()
+impl MetadataTableSetup {
+    pub fn into_setup_info(self) -> ResourceSetupInfo<(), (), MetadataTableSetup> {
+        ResourceSetupInfo {
+            key: (),
+            state: None,
+            description: "CocoIndex Metadata Table".to_string(),
+            status_check: Some(self),
+        }
     }
+}
 
+#[async_trait]
+impl ResourceSetupStatusCheck for MetadataTableSetup {
     fn describe_changes(&self) -> Vec<String> {
         if self.metadata_table_missing {
             vec![format!(