diff --git a/ast/src/lang/asg.rs b/ast/src/lang/asg.rs index e3c4124ad..15267aa6f 100644 --- a/ast/src/lang/asg.rs +++ b/ast/src/lang/asg.rs @@ -297,6 +297,34 @@ impl FromStr for NodeType { } } } +impl NodeType { + pub fn all_labels() -> &'static [&'static str] { + &[ + "Repository", + "Package", + "Language", + "Directory", + "File", + "Import", + "Library", + "Class", + "Trait", + "Instance", + "Function", + "Endpoint", + "Request", + "Datamodel", + "Feature", + "Page", + "Var", + "UnitTest", + "IntegrationTest", + "E2etest", + "Mock", + ] + } +} + impl std::fmt::Display for NodeType { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let s = match self { diff --git a/ast/src/lang/graphs/neo4j/executor.rs b/ast/src/lang/graphs/neo4j/executor.rs index a46f5fed8..960ad5a77 100644 --- a/ast/src/lang/graphs/neo4j/executor.rs +++ b/ast/src/lang/graphs/neo4j/executor.rs @@ -260,28 +260,7 @@ impl<'a> TransactionManager<'a> { async move { let mut txn = conn.start_txn().await?; for (query_str, bolt_map) in queries { - let mut query_obj = query(&query_str); - if query_str.contains("$properties") { - if let Some(BoltType::String(node_key)) = bolt_map.value.get("node_key") { - query_obj = query_obj.param("node_key", node_key.value.as_str()); - } - let properties = boltmap_to_bolttype_map(bolt_map); - query_obj = query_obj.param("properties", properties); - if query_str.contains("$now") { - use std::time::{SystemTime, UNIX_EPOCH}; - if let Ok(dur) = SystemTime::now().duration_since(UNIX_EPOCH) { - let ts = dur.as_secs_f64(); - query_obj = query_obj.param( - "now", - neo4rs::BoltType::String(format!("{:.7}", ts).into()), - ); - } - } - } else { - for (key, value) in bolt_map.value.iter() { - query_obj = query_obj.param(key.value.as_str(), value.clone()); - } - } + let query_obj = bind_parameters(&query_str, bolt_map); txn.run(query_obj).await?; } txn.commit().await?; diff --git a/ast/src/lang/graphs/neo4j/queries/edges.rs b/ast/src/lang/graphs/neo4j/queries/edges.rs index 95004670b..5ab122304 100644 --- a/ast/src/lang/graphs/neo4j/queries/edges.rs +++ b/ast/src/lang/graphs/neo4j/queries/edges.rs @@ -292,30 +292,7 @@ pub fn find_dynamic_edges_for_file_query(file: &str) -> (String, BoltMap) { let mut params = BoltMap::new(); boltmap_insert_str(&mut params, "file", file); - let static_types = vec![ - "Repository", - "Package", - "Language", - "Directory", - "File", - "Import", - "Library", - "Class", - "Trait", - "Instance", - "Function", - "Endpoint", - "Request", - "Datamodel", - "Feature", - "Page", - "Var", - "UnitTest", - "IntegrationTest", - "E2etest", - ]; - - let static_labels = static_types + let static_labels = NodeType::all_labels() .iter() .map(|t| format!("source:{}", t)) .collect::>() @@ -336,30 +313,7 @@ pub fn find_dynamic_edges_for_file_query(file: &str) -> (String, BoltMap) { pub fn find_all_dynamic_edges_query() -> (String, BoltMap) { let params = BoltMap::new(); - let static_types = vec![ - "Repository", - "Package", - "Language", - "Directory", - "File", - "Import", - "Library", - "Class", - "Trait", - "Instance", - "Function", - "Endpoint", - "Request", - "Datamodel", - "Feature", - "Page", - "Var", - "UnitTest", - "IntegrationTest", - "E2etest", - ]; - - let static_labels = static_types + let static_labels = NodeType::all_labels() .iter() .map(|t| format!("source:{}", t)) .collect::>() diff --git a/ast/src/lang/graphs/neo4j/queries/nodes.rs b/ast/src/lang/graphs/neo4j/queries/nodes.rs index 157cb5b40..3ce7df8f0 100644 --- a/ast/src/lang/graphs/neo4j/queries/nodes.rs +++ b/ast/src/lang/graphs/neo4j/queries/nodes.rs @@ -500,12 +500,21 @@ pub fn get_muted_nodes_for_files_query(files: &[String]) -> (String, BoltMap) { .collect::>(); boltmap_insert_list(&mut params, "files", files_list); - let query = "MATCH (n:Data_Bank) + let labels_list = NodeType::all_labels() + .iter() + .map(|t| format!("'{}'", t)) + .collect::>() + .join(", "); + + let query = format!( + "MATCH (n:Data_Bank) WHERE n.file IN $files AND (n.is_muted = true OR n.is_muted = 'true') - WITH n, [label IN labels(n) WHERE label IN ['Function', 'Class', 'DataModel', 'Endpoint', 'Request', 'File', 'Directory', 'Repository', 'Language', 'Library', 'Import', 'Instance', 'Page', 'Var', 'UnitTest', 'IntegrationTest', 'E2eTest', 'Trait']][0] as node_type + WITH n, [label IN labels(n) WHERE label IN [{}]][0] as node_type WHERE node_type IS NOT NULL - RETURN node_type, n.name as name, n.file as file".to_string(); + RETURN node_type, n.name as name, n.file as file", + labels_list + ); (query, params) }