1  use crate::models::commands::{CachedParsedCommand, ParsedCommand, ParsedCommands};
  2  use crate::models::world::scenes::{Scene, Stage, StageOrStub};
  3  use crate::models::{Content, ContentContainer, Entity, Insertable};
  4  use anyhow::Result;
  5  use arangors::document::options::InsertOptions;
  6  use arangors::graph::{EdgeDefinition, Graph};
  7  use arangors::transaction::{TransactionCollections, TransactionSettings};
  8  use arangors::uclient::reqwest::ReqwestClient;
  9  use arangors::{
 10      AqlQuery, ClientError, Collection, Database as ArangoDatabase, Document, GenericConnection,
 11  };
 12  use serde::{Deserialize, Serialize};
 13  use serde_json::value::to_value as to_json_value;
 14  use serde_json::Value as JsonValue;
 15  use std::collections::HashMap;
 16  
 17  mod queries;
 18  
 19  /// Type alias for how we're storing IDs in DB. WOuld prefer to have
 20  /// strong UUIDs in code and strings in DB.
 21  pub type Key = String;
 22  
 23  enum CollectionType {
 24      Document,
 25      Edge,
 26  }
 27  
 28  // Document Collections
 29  const CMD_COLLECTION: &'static str = "command_cache";
 30  const SCENE_COLLECTION: &'static str = "scenes";
 31  const REGION_COLLECTION: &'static str = "regions";
 32  const PEOPLE_COLLECTION: &'static str = "people";
 33  const ITEMS_COLLECTION: &'static str = "items";
 34  const PROPS_COLLECTION: &'static str = "props";
 35  const RACES_COLLECTION: &'static str = "races";
 36  const OCCUPATIONS_COLLECTION: &'static str = "occupations";
 37  
 38  // Edge collections
 39  const GAME_WORLD_EDGES: &'static str = "game_world";
 40  const PERSON_ATTRS: &'static str = "person_attributes";
 41  
 42  // Graphs
 43  const GAME_WORLD_GRAPH: &'static str = "world";
 44  
 45  const DOC_COLLECTIONS: &'static [&str] = &[
 46      CMD_COLLECTION,
 47      SCENE_COLLECTION,
 48      REGION_COLLECTION,
 49      PEOPLE_COLLECTION,
 50      ITEMS_COLLECTION,
 51      PROPS_COLLECTION,
 52      RACES_COLLECTION,
 53      OCCUPATIONS_COLLECTION,
 54  ];
 55  
 56  const EDGE_COLLECTIONS: &'static [&str] = &[GAME_WORLD_EDGES, PERSON_ATTRS];
 57  
 58  // Change if we decide to use a different HTTP client.
 59  type ArangoHttp = ReqwestClient;
 60  type ActiveDatabase = ArangoDatabase<ArangoHttp>;
 61  type ArangoResult<T> = std::result::Result<T, ClientError>;
 62  
 63  /// Generic edge that relates things back and forth, where the
 64  /// relation property determines what kind of relation we actually
 65  /// have.
 66  #[derive(Serialize, Deserialize, Debug)]
 67  struct Edge {
 68      _from: String,
 69      _to: String,
 70      relation: String,
 71  }
 72  
 73  /// Convert an Arango response for a single document, which may be
 74  /// missing, into an Option type. Bubble up any other errors.
 75  fn extract_document<T>(document: ArangoResult<Document<T>>) -> Result<Option<T>> {
 76      match document {
 77          Ok(doc) => Ok(Some(doc.document)),
 78          Err(db_err) => match db_err {
 79              ClientError::Arango(ref arr_err) => {
 80                  if arr_err.error_num() == 1202 {
 81                      Ok(None)
 82                  } else {
 83                      Err(db_err.into())
 84                  }
 85              }
 86              _ => Err(db_err.into()),
 87          },
 88      }
 89  }
 90  
 91  fn take_first<T>(mut vec: Vec<T>) -> Option<T> {
 92      if vec.get(0).is_none() {
 93          None
 94      } else {
 95          Some(vec.swap_remove(0))
 96      }
 97  }
 98  
 99  fn insert_opts() -> InsertOptions {
100      InsertOptions::builder()
101          .silent(false)
102          .return_new(true)
103          .build()
104  }
105  
106  fn is_scene_stub(value: &JsonValue) -> bool {
107      value
108          .as_object()
109          .and_then(|v| v.get("scene"))
110          .and_then(|scene| scene.get("isStub"))
111          .and_then(|is_stub| is_stub.as_bool())
112          .unwrap_or(false)
113  }
114  
115  async fn insert_single<T>(collection: &Collection<ArangoHttp>, value: &mut T) -> Result<()>
116  where
117      T: Insertable + Clone + Serialize,
118  {
119      let doc = to_json_value(&value)?;
120      let resp = collection.create_document(doc, insert_opts()).await?;
121  
122      let header = resp.header().unwrap();
123  
124      value.set_key(header._key.clone());
125      value.set_id(header._id.clone());
126  
127      Ok(())
128  }
129  
130  #[derive(Serialize, Deserialize)]
131  struct UpsertResponse {
132      pub _id: String,
133      pub _key: String,
134  }
135  
136  async fn upsert_scene(db: &ActiveDatabase, scene: &mut Scene) -> Result<()> {
137      let scene_json = serde_json::to_string(&scene)?;
138      let query = queries::UPSERT_SCENE.replace("<SCENE_JSON>", &scene_json);
139  
140      let aql = AqlQuery::builder()
141          .query(&query)
142          .bind_var("@scene_collection", SCENE_COLLECTION)
143          .bind_var("scene_key", to_json_value(&scene._key).unwrap())
144          .build();
145  
146      //db.aql_bind_vars::<JsonValue>(&query, vars).await?;
147      let resp = take_first(db.aql_query::<UpsertResponse>(aql).await?)
148          .expect("did not get upsert response");
149  
150      scene._id = Some(resp._id);
151      scene._key = Some(resp._key);
152  
153      Ok(())
154  }
155  
156  fn content_collection(content: &Content) -> &'static str {
157      match content {
158          Content::Scene(_) => SCENE_COLLECTION,
159          Content::SceneStub(_) => SCENE_COLLECTION,
160          Content::Person(_) => PEOPLE_COLLECTION,
161          Content::Item(_) => ITEMS_COLLECTION,
162      }
163  }
164  
165  pub struct Database {
166      conn: arangors::GenericConnection<ArangoHttp>,
167      world_name: String,
168  }
169  
170  impl Database {
171      pub async fn new(conn: GenericConnection<ArangoHttp>, world_name: &str) -> Result<Database> {
172          let db = Database {
173              conn,
174              world_name: world_name.to_string(),
175          };
176  
177          db.init().await?;
178          Ok(db)
179      }
180  
181      async fn init(&self) -> Result<()> {
182          let dbs = self.conn.accessible_databases().await?;
183  
184          if !dbs.contains_key(&self.world_name) {
185              self.conn.create_database(&self.world_name).await?;
186          }
187  
188          self.create_collections(CollectionType::Document, DOC_COLLECTIONS)
189              .await?;
190          self.create_collections(CollectionType::Edge, EDGE_COLLECTIONS)
191              .await?;
192  
193          self.create_graphs().await?;
194  
195          Ok(())
196      }
197  
198      async fn create_collections(&self, coll_type: CollectionType, names: &[&str]) -> Result<()> {
199          let db = self.db().await?;
200          let in_db = db.accessible_collections().await?;
201  
202          for name in names {
203              if in_db.iter().find(|info| info.name == *name).is_none() {
204                  match coll_type {
205                      CollectionType::Document => db.create_collection(&name).await?,
206                      CollectionType::Edge => db.create_edge_collection(&name).await?,
207                  };
208              }
209          }
210  
211          Ok(())
212      }
213  
214      async fn create_graphs(&self) -> Result<()> {
215          let db = self.db().await?;
216  
217          let in_db = db.graphs().await?.graphs;
218  
219          if in_db
220              .iter()
221              .find(|graph| graph.name == GAME_WORLD_GRAPH)
222              .is_none()
223          {
224              let edge_def = EdgeDefinition {
225                  collection: GAME_WORLD_EDGES.to_string(),
226                  from: vec![SCENE_COLLECTION.to_string()],
227                  to: vec![
228                      ITEMS_COLLECTION.to_string(),
229                      REGION_COLLECTION.to_string(),
230                      OCCUPATIONS_COLLECTION.to_string(),
231                      PEOPLE_COLLECTION.to_string(),
232                      PROPS_COLLECTION.to_string(),
233                      RACES_COLLECTION.to_string(),
234                  ],
235              };
236  
237              let world_graph = Graph::builder()
238                  .edge_definitions(vec![edge_def])
239                  .name(GAME_WORLD_GRAPH.to_string())
240                  .build();
241  
242              db.create_graph(world_graph, false).await?;
243          }
244  
245          Ok(())
246      }
247  
248      async fn db(&self) -> Result<ArangoDatabase<ArangoHttp>> {
249          let db = self.conn.db(&self.world_name).await?;
250          Ok(db)
251      }
252  
253      async fn collection(&self, name: &str) -> Result<Collection<ArangoHttp>> {
254          let coll = self.db().await?.collection(name).await?;
255          Ok(coll)
256      }
257  
258      pub async fn store_content(&self, container: &mut ContentContainer) -> Result<()> {
259          let txn_settings = TransactionSettings::builder()
260              .collections(
261                  TransactionCollections::builder()
262                      .write(vec![
263                          SCENE_COLLECTION.to_string(),
264                          PEOPLE_COLLECTION.to_string(),
265                          ITEMS_COLLECTION.to_string(),
266                          GAME_WORLD_EDGES.to_string(),
267                      ])
268                      .build(),
269              )
270              .build();
271  
272          let txn = self.db().await?.begin_transaction(txn_settings).await?;
273  
274          // First, all contained content must be inserted.
275          for relation in container.contained.as_mut_slice() {
276              let collection = content_collection(&relation.content);
277              self.store_single_content(collection, &mut relation.content)
278                  .await?;
279          }
280  
281          // Now insert the container/owner content + relations
282          let collection = content_collection(&container.owner);
283          self.store_single_content(collection, &mut container.owner)
284              .await?;
285          self.relate_content(&container).await?;
286  
287          txn.commit_transaction().await?;
288  
289          Ok(())
290      }
291  
292      async fn relate_content(&self, container: &ContentContainer) -> Result<()> {
293          let game_world = self.collection(GAME_WORLD_EDGES).await?;
294  
295          let owner_id = container
296              .owner
297              .id()
298              .expect("Did not get an ID from inserted object!");
299  
300          for relation in container.contained.as_slice() {
301              let content_id = relation
302                  .content
303                  .id()
304                  .expect("Did not get ID from inserted contained object!");
305  
306              let outbound = Edge {
307                  _from: owner_id.to_string(),
308                  _to: content_id.to_string(),
309                  relation: relation.outbound.clone(),
310              };
311  
312              let inbound = Edge {
313                  _from: content_id.to_string(),
314                  _to: owner_id.to_string(),
315                  relation: relation.inbound.clone(),
316              };
317  
318              game_world
319                  .create_document(outbound, InsertOptions::default())
320                  .await?;
321              game_world
322                  .create_document(inbound, InsertOptions::default())
323                  .await?;
324          }
325  
326          Ok(())
327      }
328  
329      pub async fn store_single_content(&self, coll_name: &str, content: &mut Content) -> Result<()> {
330          let collection = self.collection(coll_name).await?;
331  
332          match content {
333              //Content::Scene(ref mut scene) => insert_single(&collection, scene).await?,
334              Content::Scene(ref mut scene) => upsert_scene(&self.db().await?, scene).await?,
335              Content::SceneStub(ref mut stub) => insert_single(&collection, stub).await?,
336              Content::Person(ref mut person) => insert_single(&collection, person).await?,
337              Content::Item(ref mut item) => insert_single(&collection, item).await?,
338          };
339  
340          Ok(())
341      }
342  
343      pub async fn load_stage(&self, scene_key: &str) -> Result<Option<StageOrStub>> {
344          let mut vars = HashMap::new();
345          vars.insert("scene_key", to_json_value(&scene_key).unwrap());
346          vars.insert("@scene_collection", SCENE_COLLECTION.into());
347  
348          let db = self.db().await?;
349  
350          let res = db
351              .aql_bind_vars::<JsonValue>(queries::LOAD_STAGE, vars)
352              .await?;
353  
354          let maybe_stage = take_first(res);
355  
356          if let Some(stage) = maybe_stage {
357              let stage_or_stub = if is_scene_stub(&stage) {
358                  // The stub is embedded in the scene field of the result.
359                  StageOrStub::Stub(serde_json::from_value(
360                      stage.get("scene").cloned().unwrap(),
361                  )?)
362              } else {
363                  StageOrStub::Stage(serde_json::from_value(stage)?)
364              };
365  
366              Ok(Some(stage_or_stub))
367          } else {
368              Ok(None)
369          }
370      }
371  
372      pub async fn stage_exists(&self, scene_key: &str) -> Result<bool> {
373          let mut vars = HashMap::new();
374  
375          vars.insert("scene_key", to_json_value(&scene_key).unwrap());
376          vars.insert("@scene_collection", SCENE_COLLECTION.into());
377  
378          let db = self.db().await?;
379          let stage_count = db
380              .aql_bind_vars::<JsonValue>(queries::LOAD_STAGE, vars)
381              .await?
382              .len();
383  
384          Ok(stage_count > 0)
385      }
386  
387      pub async fn entity_exists(&self, entity_key: &str) -> Result<bool> {
388          let mut vars = HashMap::new();
389  
390          vars.insert("entity_key", to_json_value(entity_key).unwrap());
391  
392          let db = self.db().await?;
393          let entity_count = db
394              .aql_bind_vars::<JsonValue>(queries::LOAD_ENTITY, vars)
395              .await?
396              .len();
397  
398          Ok(entity_count > 0)
399      }
400  
401      pub async fn load_entity_in_scene(
402          &self,
403          scene_key: &str,
404          entity_key: &str,
405      ) -> Result<Option<Entity>> {
406          let aql = AqlQuery::builder()
407              .query(queries::LOAD_ENTITY_IN_SCENE)
408              .bind_var("@scene_collection", SCENE_COLLECTION)
409              .bind_var("scene_key", to_json_value(scene_key)?)
410              .bind_var("entity_key", to_json_value(entity_key)?)
411              .build();
412  
413          let results = self.db().await?.aql_query(aql).await?;
414          Ok(take_first(results))
415      }
416  
417      pub async fn load_entity(&self, entity_key: &str) -> Result<Option<Entity>> {
418          let aql = AqlQuery::builder()
419              .query(queries::LOAD_ENTITY)
420              .bind_var("entity_key", to_json_value(entity_key)?)
421              .build();
422  
423          let results = self.db().await?.aql_query(aql).await?;
424          Ok(take_first(results))
425      }
426  
427      pub async fn cache_command(
428          &self,
429          raw_cmd: &str,
430          scene: &Scene,
431          parsed_cmds: &ParsedCommands,
432      ) -> Result<()> {
433          let collection = self.collection(CMD_COLLECTION).await?;
434          let doc = CachedParsedCommand {
435              raw: raw_cmd.to_string(),
436              scene_key: scene._key.as_ref().cloned().expect("scene is missing key"),
437              commands: parsed_cmds.clone(),
438          };
439  
440          collection.create_document(doc, insert_opts()).await?;
441          Ok(())
442      }
443  
444      pub async fn load_cached_command(
445          &self,
446          raw_cmd: &str,
447          scene: &Scene,
448      ) -> Result<Option<CachedParsedCommand>> {
449          let scene_key = scene._key.as_deref();
450          let aql = AqlQuery::builder()
451              .query(queries::LOAD_CACHED_COMMAND)
452              .bind_var("@cache_collection", CMD_COLLECTION)
453              .bind_var("raw_cmd", to_json_value(raw_cmd)?)
454              .bind_var("scene_key", to_json_value(scene_key)?)
455              .build();
456  
457          let results = self.db().await?.aql_query(aql).await?;
458          Ok(take_first(results))
459      }
460  }