queries.rs 12 KB

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