queries.rs 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. use heed3::RoTxn;
  2. use get_routes::handler;
  3. use helix_db::{field_remapping, identifier_remapping, traversal_remapping, exclude_field};
  4. use helix_db::helix_engine::vector_core::vector::HVector;
  5. use helix_db::{
  6. helix_engine::graph_core::ops::{
  7. g::G,
  8. in_::{in_::InAdapter, in_e::InEdgesAdapter, to_n::ToNAdapter},
  9. out::{from_n::FromNAdapter, out::OutAdapter, out_e::OutEdgesAdapter},
  10. source::{
  11. add_e::{AddEAdapter, EdgeType},
  12. add_n::AddNAdapter,
  13. e_from_id::EFromIdAdapter,
  14. e_from_type::EFromTypeAdapter,
  15. n_from_id::NFromIdAdapter,
  16. n_from_type::NFromTypeAdapter,
  17. n_from_index::NFromIndexAdapter,
  18. },
  19. tr_val::{Traversable, TraversalVal},
  20. util::{
  21. dedup::DedupAdapter, filter_mut::FilterMut,
  22. filter_ref::FilterRefAdapter, range::RangeAdapter, update::UpdateAdapter,
  23. map::MapAdapter, paths::ShortestPathAdapter, props::PropsAdapter, drop::Drop,
  24. },
  25. vectors::{insert::InsertVAdapter, search::SearchVAdapter, brute_force_search::BruteForceSearchVAdapter},
  26. bm25::search_bm25::SearchBM25Adapter,
  27. },
  28. helix_engine::types::GraphError,
  29. helix_gateway::router::router::HandlerInput,
  30. node_matches, props,
  31. protocol::count::Count,
  32. protocol::remapping::ResponseRemapping,
  33. protocol::response::Response,
  34. protocol::traversal_value::TraversalValue,
  35. protocol::{
  36. filterable::Filterable, remapping::Remapping, return_values::ReturnValue, value::Value, id::ID,
  37. },
  38. };
  39. use sonic_rs::{Deserialize, Serialize};
  40. use std::collections::{HashMap, HashSet};
  41. use std::sync::Arc;
  42. use std::time::Instant;
  43. use std::cell::RefCell;
  44. use chrono::{DateTime, Utc};
  45. pub struct Company {
  46. pub company_number: String,
  47. }
  48. #[derive(Serialize, Deserialize)]
  49. pub struct HasCompanyInput {
  50. pub company_number: String
  51. }
  52. #[handler]
  53. pub fn HasCompany (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  54. let data: HasCompanyInput = match sonic_rs::from_slice(&input.request.body) {
  55. Ok(data) => data,
  56. Err(err) => return Err(GraphError::from(err)),
  57. };
  58. let mut remapping_vals: RefCell<HashMap<u128, ResponseRemapping>> = RefCell::new(HashMap::new());
  59. let db = Arc::clone(&input.graph.storage);
  60. let txn = db.graph_env.read_txn().unwrap();
  61. let c = G::new(Arc::clone(&db), &txn)
  62. .n_from_index("company_number", &data.company_number).collect_to::<Vec<_>>();
  63. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  64. return_vals.insert("c".to_string(), ReturnValue::from_traversal_value_array_with_mixin(c.clone(), remapping_vals.borrow_mut()));
  65. txn.commit().unwrap();
  66. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  67. Ok(())
  68. }
  69. #[derive(Serialize, Deserialize)]
  70. pub struct AddCompanyInput {
  71. pub company_number: String
  72. }
  73. #[handler]
  74. pub fn AddCompany (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  75. let data: AddCompanyInput = match sonic_rs::from_slice(&input.request.body) {
  76. Ok(data) => data,
  77. Err(err) => return Err(GraphError::from(err)),
  78. };
  79. let mut remapping_vals: RefCell<HashMap<u128, ResponseRemapping>> = RefCell::new(HashMap::new());
  80. let db = Arc::clone(&input.graph.storage);
  81. let mut txn = db.graph_env.write_txn().unwrap();
  82. let c = G::new_mut(Arc::clone(&db), &mut txn)
  83. .add_n("Company", Some(props! { "company_number" => data.company_number.clone() }), Some(&["company_number"])).collect_to::<Vec<_>>();
  84. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  85. return_vals.insert("c".to_string(), ReturnValue::from_traversal_value_array_with_mixin(c.clone(), remapping_vals.borrow_mut()));
  86. txn.commit().unwrap();
  87. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  88. Ok(())
  89. }