queries.rs 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279
  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 User {
  46. pub name: String,
  47. pub age: u32,
  48. pub email: String,
  49. pub created_at: i32,
  50. pub updated_at: i32,
  51. }
  52. pub struct Post {
  53. pub content: String,
  54. pub created_at: i32,
  55. pub updated_at: i32,
  56. }
  57. pub struct Follows {
  58. pub from: User,
  59. pub to: User,
  60. pub since: i32,
  61. }
  62. pub struct Created {
  63. pub from: User,
  64. pub to: Post,
  65. pub created_at: i32,
  66. }
  67. #[derive(Serialize, Deserialize)]
  68. pub struct create_postInput {
  69. pub user_id: ID,
  70. pub content: String,
  71. pub now: i32
  72. }
  73. #[handler]
  74. pub fn create_post (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  75. let data: create_postInput = 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 user = G::new(Arc::clone(&db), &txn)
  83. .n_from_id(&data.user_id).collect_to::<Vec<_>>();
  84. let post = G::new_mut(Arc::clone(&db), &mut txn)
  85. .add_n("Post", Some(props! { "created_at" => data.now.clone(), "updated_at" => data.now.clone(), "content" => data.content.clone() }), None).collect_to::<Vec<_>>();
  86. G::new_mut(Arc::clone(&db), &mut txn)
  87. .add_e("Created", Some(props! { "created_at" => data.now.clone() }), user.id(), post.id(), true, EdgeType::Node).collect_to::<Vec<_>>();
  88. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  89. return_vals.insert("post".to_string(), ReturnValue::from_traversal_value_array_with_mixin(post.clone(), remapping_vals.borrow_mut()));
  90. txn.commit().unwrap();
  91. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  92. Ok(())
  93. }
  94. #[derive(Serialize, Deserialize)]
  95. pub struct create_userInput {
  96. pub name: String,
  97. pub age: u32,
  98. pub email: String,
  99. pub now: i32
  100. }
  101. #[handler]
  102. pub fn create_user (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  103. let data: create_userInput = match sonic_rs::from_slice(&input.request.body) {
  104. Ok(data) => data,
  105. Err(err) => return Err(GraphError::from(err)),
  106. };
  107. let mut remapping_vals: RefCell<HashMap<u128, ResponseRemapping>> = RefCell::new(HashMap::new());
  108. let db = Arc::clone(&input.graph.storage);
  109. let mut txn = db.graph_env.write_txn().unwrap();
  110. let user = G::new_mut(Arc::clone(&db), &mut txn)
  111. .add_n("User", Some(props! { "name" => data.name.clone(), "email" => data.email.clone(), "age" => data.age.clone(), "created_at" => data.now.clone(), "updated_at" => data.now.clone() }), None).collect_to::<Vec<_>>();
  112. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  113. return_vals.insert("user".to_string(), ReturnValue::from_traversal_value_array_with_mixin(user.clone(), remapping_vals.borrow_mut()));
  114. txn.commit().unwrap();
  115. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  116. Ok(())
  117. }
  118. #[derive(Serialize, Deserialize)]
  119. pub struct get_followed_usersInput {
  120. pub user_id: ID
  121. }
  122. #[handler]
  123. pub fn get_followed_users (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  124. let data: get_followed_usersInput = match sonic_rs::from_slice(&input.request.body) {
  125. Ok(data) => data,
  126. Err(err) => return Err(GraphError::from(err)),
  127. };
  128. let mut remapping_vals: RefCell<HashMap<u128, ResponseRemapping>> = RefCell::new(HashMap::new());
  129. let db = Arc::clone(&input.graph.storage);
  130. let txn = db.graph_env.read_txn().unwrap();
  131. let followed = G::new(Arc::clone(&db), &txn)
  132. .n_from_id(&data.user_id)
  133. .out("Follows",&EdgeType::Node).collect_to::<Vec<_>>();
  134. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  135. return_vals.insert("followed".to_string(), ReturnValue::from_traversal_value_array_with_mixin(followed.clone(), remapping_vals.borrow_mut()));
  136. txn.commit().unwrap();
  137. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  138. Ok(())
  139. }
  140. #[derive(Serialize, Deserialize)]
  141. pub struct create_followInput {
  142. pub follower_id: ID,
  143. pub followed_id: ID,
  144. pub now: i32
  145. }
  146. #[handler]
  147. pub fn create_follow (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  148. let data: create_followInput = match sonic_rs::from_slice(&input.request.body) {
  149. Ok(data) => data,
  150. Err(err) => return Err(GraphError::from(err)),
  151. };
  152. let mut remapping_vals: RefCell<HashMap<u128, ResponseRemapping>> = RefCell::new(HashMap::new());
  153. let db = Arc::clone(&input.graph.storage);
  154. let mut txn = db.graph_env.write_txn().unwrap();
  155. let follower = G::new(Arc::clone(&db), &txn)
  156. .n_from_id(&data.follower_id).collect_to::<Vec<_>>();
  157. let followed = G::new(Arc::clone(&db), &txn)
  158. .n_from_id(&data.followed_id).collect_to::<Vec<_>>();
  159. G::new_mut(Arc::clone(&db), &mut txn)
  160. .add_e("Follows", Some(props! { "since" => data.now.clone() }), follower.id(), followed.id(), true, EdgeType::Node).collect_to::<Vec<_>>();
  161. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  162. return_vals.insert("success".to_string(), ReturnValue::from(Value::from("success")));
  163. txn.commit().unwrap();
  164. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  165. Ok(())
  166. }
  167. #[handler]
  168. pub fn find_users_access (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  169. let mut remapping_vals: RefCell<HashMap<u128, ResponseRemapping>> = RefCell::new(HashMap::new());
  170. let db = Arc::clone(&input.graph.storage);
  171. let txn = db.graph_env.read_txn().unwrap();
  172. let users = G::new(Arc::clone(&db), &txn)
  173. .n_from_type("User").collect_to::<Vec<_>>();
  174. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  175. return_vals.insert("users".to_string(), ReturnValue::from_traversal_value_array_with_mixin(G::new_from(Arc::clone(&db), &txn, users.clone())
  176. .map_traversal(|item, txn| { traversal_remapping!(remapping_vals, item.clone(), "name" => G::new_from(Arc::clone(&db), &txn, vec![item.clone()])
  177. .check_property("name").collect_to::<Vec<_>>())?;
  178. traversal_remapping!(remapping_vals, item.clone(), "age" => G::new_from(Arc::clone(&db), &txn, vec![item.clone()])
  179. .check_property("age").collect_to::<Vec<_>>())?;
  180. Ok(item) }).collect_to::<Vec<_>>().clone(), remapping_vals.borrow_mut()));
  181. txn.commit().unwrap();
  182. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  183. Ok(())
  184. }
  185. #[handler]
  186. pub fn get_name_remapping_simple (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  187. let mut remapping_vals: RefCell<HashMap<u128, ResponseRemapping>> = RefCell::new(HashMap::new());
  188. let db = Arc::clone(&input.graph.storage);
  189. let txn = db.graph_env.read_txn().unwrap();
  190. let users = G::new(Arc::clone(&db), &txn)
  191. .n_from_type("User").collect_to::<Vec<_>>();
  192. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  193. return_vals.insert("users".to_string(), ReturnValue::from_traversal_value_array_with_mixin(G::new_from(Arc::clone(&db), &txn, users.clone())
  194. .map_traversal(|item, txn| { traversal_remapping!(remapping_vals, item.clone(), "givenName" => G::new_from(Arc::clone(&db), &txn, vec![item.clone()])
  195. .check_property("name").collect_to::<Vec<_>>())?;
  196. Ok(item) }).collect_to::<Vec<_>>().clone(), remapping_vals.borrow_mut()));
  197. txn.commit().unwrap();
  198. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  199. Ok(())
  200. }
  201. #[derive(Serialize, Deserialize)]
  202. pub struct find_user_posts_with_creator_detailsInput {
  203. pub userID: ID
  204. }
  205. #[handler]
  206. pub fn find_user_posts_with_creator_details (input: &HandlerInput, response: &mut Response) -> Result<(), GraphError> {
  207. let data: find_user_posts_with_creator_detailsInput = match sonic_rs::from_slice(&input.request.body) {
  208. Ok(data) => data,
  209. Err(err) => return Err(GraphError::from(err)),
  210. };
  211. let mut remapping_vals: RefCell<HashMap<u128, ResponseRemapping>> = RefCell::new(HashMap::new());
  212. let db = Arc::clone(&input.graph.storage);
  213. let txn = db.graph_env.read_txn().unwrap();
  214. let user = G::new(Arc::clone(&db), &txn)
  215. .n_from_id(&data.userID).collect_to::<Vec<_>>();
  216. let posts = G::new_from(Arc::clone(&db), &txn, user.clone())
  217. .out("Created",&EdgeType::Node).collect_to::<Vec<_>>();
  218. let mut return_vals: HashMap<String, ReturnValue> = HashMap::new();
  219. return_vals.insert("user".to_string(), ReturnValue::from_traversal_value_array_with_mixin(G::new_from(Arc::clone(&db), &txn, user.clone())
  220. .map_traversal(|creator, txn| { traversal_remapping!(remapping_vals, creator.clone(), "creatorName" => G::new_from(Arc::clone(&db), &txn, vec![creator.clone()])
  221. .check_property("name").collect_to::<Vec<_>>())?;
  222. traversal_remapping!(remapping_vals, creator.clone(), "createdPosts" => G::new_from(Arc::clone(&db), &txn, posts.clone())
  223. .map_traversal(|item, txn| { traversal_remapping!(remapping_vals, item.clone(), "postContent" => G::new_from(Arc::clone(&db), &txn, vec![item.clone()])
  224. .check_property("content").collect_to::<Vec<_>>())?;
  225. traversal_remapping!(remapping_vals, item.clone(), "createdAt" => G::new_from(Arc::clone(&db), &txn, vec![item.clone()])
  226. .check_property("created_at").collect_to::<Vec<_>>())?;
  227. traversal_remapping!(remapping_vals, item.clone(), "updatedAt" => G::new_from(Arc::clone(&db), &txn, vec![item.clone()])
  228. .check_property("updated_at").collect_to::<Vec<_>>())?;
  229. Ok(item) }).collect_to::<Vec<_>>())?;
  230. Ok(creator) }).collect_to::<Vec<_>>().clone(), remapping_vals.borrow_mut()));
  231. txn.commit().unwrap();
  232. response.body = sonic_rs::to_vec(&return_vals).unwrap();
  233. Ok(())
  234. }