remove duplicate type (#690)

ContentSource and ExtractionPolicyContentSource were defined to have identical content, remove one of them. Compare values directly instead of converting to string.
tensorlakeai · Jun 26, 2024 · 5849644 · 5849644
1 parent 90a4a22
commit 5849644
Show file tree

Hide file tree

Showing 6 changed files with 11 additions and 69 deletions.
diff --git a/crates/indexify_internal_api/src/lib.rs b/crates/indexify_internal_api/src/lib.rs
@@ -599,64 +599,6 @@ impl From<GarbageCollectionTask> for indexify_coordinator::GcTask {
 pub type ExtractionPolicyId = String;
 pub type ExtractionPolicyName = String;
 
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
-pub enum ExtractionPolicyContentSource {
-    Ingestion,
-    ExtractionPolicyName(ExtractionPolicyName),
-}
-
-impl Display for ExtractionPolicyContentSource {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match self {
-            ExtractionPolicyContentSource::Ingestion => write!(f, ""),
-            ExtractionPolicyContentSource::ExtractionPolicyName(name) => {
-                write!(f, "{}", name)
-            }
-        }
-    }
-}
-
-impl From<&ExtractionPolicyContentSource> for ContentSource {
-    fn from(value: &ExtractionPolicyContentSource) -> Self {
-        match value {
-            ExtractionPolicyContentSource::Ingestion => ContentSource::Ingestion,
-            ExtractionPolicyContentSource::ExtractionPolicyName(name) => {
-                ContentSource::ExtractionPolicyName(name.to_string())
-            }
-        }
-    }
-}
-
-impl Default for ExtractionPolicyContentSource {
-    fn default() -> Self {
-        ExtractionPolicyContentSource::ExtractionPolicyName(ExtractionPolicyId::default())
-    }
-}
-
-impl From<ExtractionPolicyContentSource> for String {
-    fn from(source: ExtractionPolicyContentSource) -> Self {
-        String::from(&source)
-    }
-}
-
-impl From<&ExtractionPolicyContentSource> for String {
-    fn from(value: &ExtractionPolicyContentSource) -> Self {
-        match value {
-            ExtractionPolicyContentSource::Ingestion => "".to_string(),
-            ExtractionPolicyContentSource::ExtractionPolicyName(name) => name.clone(),
-        }
-    }
-}
-
-impl From<&str> for ExtractionPolicyContentSource {
-    fn from(value: &str) -> Self {
-        if value.is_empty() {
-            return ExtractionPolicyContentSource::Ingestion;
-        }
-        ExtractionPolicyContentSource::ExtractionPolicyName(value.to_string())
-    }
-}
-
 #[derive(Debug, Clone, Serialize, PartialEq, Eq, Deserialize, Default, Builder)]
 #[builder(build_fn(skip))]
 pub struct ExtractionPolicy {
@@ -671,7 +613,7 @@ pub struct ExtractionPolicy {
     pub output_table_mapping: HashMap<String, String>,
     // The source of the content this policy will match against. Will either be the graph id or a
     // parent policy id
-    pub content_source: ExtractionPolicyContentSource,
+    pub content_source: ContentSource,
 }
 
 impl TryFrom<ExtractionPolicy> for indexify_coordinator::ExtractionPolicy {

diff --git a/crates/indexify_internal_api/src/v1.rs b/crates/indexify_internal_api/src/v1.rs
@@ -37,7 +37,7 @@ pub struct ExtractionPolicy {
     pub filters: HashMap<String, String>,
     pub input_params: serde_json::Value,
     pub output_table_mapping: HashMap<String, String>,
-    pub content_source: super::ExtractionPolicyContentSource,
+    pub content_source: super::ContentSource,
 }
 
 impl From<ExtractionPolicy> for super::ExtractionPolicy {

diff --git a/src/coordinator_service.rs b/src/coordinator_service.rs
@@ -169,9 +169,9 @@ impl CoordinatorServiceServer {
                 .map_err(|e| anyhow!(format!("unable to parse input_params: {}", e)))?;
             let extractor = self.coordinator.get_extractor(&policy_request.extractor)?;
             let content_source = if policy_request.content_source.eq("") {
-                internal_api::ExtractionPolicyContentSource::Ingestion
+                internal_api::ContentSource::Ingestion
             } else {
-                internal_api::ExtractionPolicyContentSource::ExtractionPolicyName(
+                internal_api::ContentSource::ExtractionPolicyName(
                     policy_request.content_source.clone(),
                 )
             };

diff --git a/src/ingest_extracted_content.rs b/src/ingest_extracted_content.rs
@@ -389,10 +389,10 @@ mod tests {
     use anyhow::Result;
     use indexify_internal_api::{
         ContentMetadata,
+        ContentSource,
         ExtractedEmbeddings,
         ExtractionGraph,
         ExtractionPolicy,
-        ExtractionPolicyContentSource,
         ExtractorDescription,
         StructuredDataSchema,
         Task,
@@ -808,7 +808,7 @@ mod tests {
             extractor: "extractor_name".to_string(),
             graph_name: "extraction_graph_id".to_string(),
             filters: HashMap::new(),
-            content_source: ExtractionPolicyContentSource::Ingestion,
+            content_source: ContentSource::Ingestion,
             output_table_mapping: vec![("test_output".to_string(), "test_table".to_string())]
                 .into_iter()
                 .collect(),
@@ -938,7 +938,7 @@ mod tests {
             extractor: "extractor_name".to_string(),
             graph_name: "extraction_graph_id".to_string(),
             filters: HashMap::new(),
-            content_source: ExtractionPolicyContentSource::Ingestion,
+            content_source: ContentSource::Ingestion,
             output_table_mapping: vec![("test_output".to_string(), "test_table".to_string())]
                 .into_iter()
                 .collect(),

diff --git a/src/state/mod.rs b/src/state/mod.rs
@@ -433,7 +433,7 @@ impl App {
         }
         let mut matched_policies = Vec::new();
         for extraction_policy in all_extraction_policies {
-            if content_metadata.source.to_string() != extraction_policy.content_source.to_string() {
+            if content_metadata.source != extraction_policy.content_source {
                 continue;
             }
             if !extraction_policy.filters.iter().all(|(name, value)| {

diff --git a/src/test_util.rs b/src/test_util.rs
@@ -72,7 +72,7 @@ pub mod db_utils {
                     "test_output".to_string(),
                     "test_table".to_string(),
                 )]),
-                content_source: internal_api::ExtractionPolicyContentSource::Ingestion,
+                content_source: internal_api::ContentSource::Ingestion,
             };
             extraction_policies.push(ep);
         }
@@ -114,9 +114,9 @@ pub mod db_utils {
                     "test_table".to_string(),
                 )]),
                 content_source: match parent {
-                    Parent::Root => internal_api::ExtractionPolicyContentSource::Ingestion,
+                    Parent::Root => internal_api::ContentSource::Ingestion,
                     Parent::Child(parent_index) => {
-                        internal_api::ExtractionPolicyContentSource::ExtractionPolicyName(
+                        internal_api::ContentSource::ExtractionPolicyName(
                             extraction_policy_names[*parent_index].to_string(),
                         )
                     }