jspg/src/database/mod.rs

pub mod r#enum;
pub mod executors;
pub mod formats;
pub mod page;
pub mod punc;
pub mod relation;
pub mod schema;
pub mod r#type;

// External mock exports inside the executor sub-folder

use r#enum::Enum;
use executors::DatabaseExecutor;

#[cfg(not(test))]
use executors::pgrx::SpiExecutor;

#[cfg(test)]
use executors::mock::MockExecutor;

pub mod stem;
use punc::Punc;
use relation::Relation;
use schema::Schema;
use serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use stem::Stem;
use r#type::Type;

pub struct Database {
  pub enums: HashMap<String, Enum>,
  pub types: HashMap<String, Type>,
  pub puncs: HashMap<String, Punc>,
  pub relations: HashMap<String, Relation>,
  pub schemas: HashMap<String, Schema>,
  // Map of Schema ID -> { Entity Type -> Target Subschema Arc }
  pub stems: HashMap<String, HashMap<String, Arc<Stem>>>,
  pub descendants: HashMap<String, Vec<String>>,
  pub depths: HashMap<String, usize>,
  pub executor: Box<dyn DatabaseExecutor + Send + Sync>,
}

impl Database {
  pub fn new(val: &serde_json::Value) -> Result<Self, crate::drop::Drop> {
    let mut db = Self {
      enums: HashMap::new(),
      types: HashMap::new(),
      relations: HashMap::new(),
      puncs: HashMap::new(),
      schemas: HashMap::new(),
      stems: HashMap::new(),
      descendants: HashMap::new(),
      depths: HashMap::new(),
      #[cfg(not(test))]
      executor: Box::new(SpiExecutor::new()),
      #[cfg(test)]
      executor: Box::new(MockExecutor::new()),
    };

    if let Some(arr) = val.get("enums").and_then(|v| v.as_array()) {
      for item in arr {
        if let Ok(def) = serde_json::from_value::<Enum>(item.clone()) {
          db.enums.insert(def.name.clone(), def);
        }
      }
    }

    if let Some(arr) = val.get("types").and_then(|v| v.as_array()) {
      for item in arr {
        if let Ok(def) = serde_json::from_value::<Type>(item.clone()) {
          db.types.insert(def.name.clone(), def);
        }
      }
    }

    if let Some(arr) = val.get("relations").and_then(|v| v.as_array()) {
      for item in arr {
        match serde_json::from_value::<Relation>(item.clone()) {
          Ok(def) => {
            db.relations.insert(def.constraint.clone(), def);
          }
          Err(e) => println!("DATABASE RELATION PARSE FAILED: {:?}", e),
        }
      }
    }

    if let Some(arr) = val.get("puncs").and_then(|v| v.as_array()) {
      for item in arr {
        if let Ok(def) = serde_json::from_value::<Punc>(item.clone()) {
          db.puncs.insert(def.name.clone(), def);
        }
      }
    }

    if let Some(arr) = val.get("schemas").and_then(|v| v.as_array()) {
      for (i, item) in arr.iter().enumerate() {
        if let Ok(mut schema) = serde_json::from_value::<Schema>(item.clone()) {
          let id = schema
            .obj
            .id
            .clone()
            .unwrap_or_else(|| format!("schema_{}", i));
          schema.obj.id = Some(id.clone());
          db.schemas.insert(id, schema);
        }
      }
    }

    db.compile()?;
    Ok(db)
  }

  /// Override the default executor for unit testing
  pub fn with_executor(mut self, executor: Box<dyn DatabaseExecutor + Send + Sync>) -> Self {
    self.executor = executor;
    self
  }

  /// Executes a query expecting a single JSONB array return, representing rows.
  pub fn query(&self, sql: &str, args: Option<&[Value]>) -> Result<Value, String> {
    self.executor.query(sql, args)
  }

  /// Executes an operation (INSERT, UPDATE, DELETE, or pg_notify) that does not return rows.
  pub fn execute(&self, sql: &str, args: Option<&[Value]>) -> Result<(), String> {
    self.executor.execute(sql, args)
  }

  /// Returns the current authenticated user's ID
  pub fn auth_user_id(&self) -> Result<String, String> {
    self.executor.auth_user_id()
  }

  /// Returns the current transaction timestamp
  pub fn timestamp(&self) -> Result<String, String> {
    self.executor.timestamp()
  }

  /// Organizes the graph of the database, compiling regex, format functions, and caching relationships.
  pub fn compile(&mut self) -> Result<(), crate::drop::Drop> {
    self.collect_schemas();
    self.collect_depths();
    self.collect_descendants();
    self.compile_schemas();
    self.collect_stems()?;

    Ok(())
  }

  fn collect_schemas(&mut self) {
    let mut to_insert = Vec::new();

    // Pass 1: Extract all Schemas structurally off top level definitions into the master registry.
    for type_def in self.types.values() {
      for mut schema in type_def.schemas.clone() {
        schema.harvest(&mut to_insert);
      }
    }
    for punc_def in self.puncs.values() {
      for mut schema in punc_def.schemas.clone() {
        schema.harvest(&mut to_insert);
      }
    }
    for enum_def in self.enums.values() {
      for mut schema in enum_def.schemas.clone() {
        schema.harvest(&mut to_insert);
      }
    }

    for (id, schema) in to_insert {
      self.schemas.insert(id, schema);
    }
  }

  fn collect_depths(&mut self) {
    let mut depths: HashMap<String, usize> = HashMap::new();
    let schema_ids: Vec<String> = self.schemas.keys().cloned().collect();

    for id in schema_ids {
      let mut current_id = id.clone();
      let mut depth = 0;
      let mut visited = HashSet::new();

      while let Some(schema) = self.schemas.get(&current_id) {
        if !visited.insert(current_id.clone()) {
          break; // Cycle detected
        }
        if let Some(ref_str) = &schema.obj.r#ref {
          current_id = ref_str.clone();
          depth += 1;
        } else {
          break;
        }
      }
      depths.insert(id, depth);
    }
    self.depths = depths;
  }

  fn collect_descendants(&mut self) {
    let mut direct_refs: HashMap<String, Vec<String>> = HashMap::new();
    for (id, schema) in &self.schemas {
      if let Some(ref_str) = &schema.obj.r#ref {
        direct_refs
          .entry(ref_str.clone())
          .or_default()
          .push(id.clone());
      }
    }

    // Cache generic descendants for $family runtime lookups
    let mut descendants = HashMap::new();
    for (id, schema) in &self.schemas {
      if let Some(family_target) = &schema.obj.family {
        let mut desc_set = HashSet::new();
        Self::collect_descendants_recursively(family_target, &direct_refs, &mut desc_set);
        let mut desc_vec: Vec<String> = desc_set.into_iter().collect();
        desc_vec.sort();

        // By placing all descendants directly onto the ID mapped location of the Family declaration,
        // we can lookup descendants natively in ValidationContext without AST replacement overrides.
        descendants.insert(id.clone(), desc_vec);
      }
    }
    self.descendants = descendants;
  }

  fn collect_descendants_recursively(
    target: &str,
    direct_refs: &HashMap<String, Vec<String>>,
    descendants: &mut HashSet<String>,
  ) {
    if let Some(children) = direct_refs.get(target) {
      for child in children {
        if descendants.insert(child.clone()) {
          Self::collect_descendants_recursively(child, direct_refs, descendants);
        }
      }
    }
  }

  fn compile_schemas(&mut self) {
    // Pass 3: compile_internals across pure structure
    let schema_ids: Vec<String> = self.schemas.keys().cloned().collect();
    for id in schema_ids {
      if let Some(schema) = self.schemas.get_mut(&id) {
        schema.compile_internals();
      }
    }
  }

  fn collect_stems(&mut self) -> Result<(), crate::drop::Drop> {
    let mut db_stems: HashMap<String, HashMap<String, Arc<Stem>>> = HashMap::new();
    let mut errors: Vec<crate::drop::Error> = Vec::new();

    let schema_ids: Vec<String> = self.schemas.keys().cloned().collect();
    for schema_id in schema_ids {
      if let Some(schema) = self.schemas.get(&schema_id) {
        let mut inner_map = HashMap::new();
        Self::discover_stems(
          self,
          &schema_id,
          schema,
          String::from(""),
          None,
          None,
          false,
          &mut inner_map,
          Vec::new(),
          &mut errors,
        );
        if !inner_map.is_empty() {
          db_stems.insert(schema_id, inner_map);
        }
      }
    }

    self.stems = db_stems;

    if !errors.is_empty() {
      return Err(crate::drop::Drop::with_errors(errors));
    }
    Ok(())
  }

  fn discover_stems(
    db: &Database,
    root_schema_id: &str,
    schema: &Schema,
    current_path: String,
    parent_type: Option<String>,
    property_name: Option<String>,
    is_polymorphic: bool,
    inner_map: &mut HashMap<String, Arc<Stem>>,
    seen_entities: Vec<String>,
    errors: &mut Vec<crate::drop::Error>,
  ) {
    let mut is_entity = false;
    let mut entity_type = String::new();

    // First check if the Schema's $id is a native Database Type
    if let Some(ref id) = schema.obj.id {
      let parts: Vec<&str> = id.split('.').collect();
      if let Some(last_seg) = parts.last() {
        if db.types.contains_key(*last_seg) {
          is_entity = true;
          entity_type = last_seg.to_string();
        }
      }
    }

    // If not found via $id, check the $ref pointer
    // This allows ad-hoc schemas (like `save_person.response`) to successfully adopt the Type of what they $ref
    if !is_entity {
      if let Some(ref r) = schema.obj.r#ref {
        let parts: Vec<&str> = r.split('.').collect();
        if let Some(last_seg) = parts.last() {
          if db.types.contains_key(*last_seg) {
            is_entity = true;
            entity_type = last_seg.to_string();
          }
        }
      }
    }

    if is_entity {
      if seen_entities.contains(&entity_type) {
        return; // Break cyclical schemas!
      }
    }

    let mut relation_col = None;
    if is_entity {
      if let (Some(pt), Some(prop)) = (&parent_type, &property_name) {
        let expected_col = format!("{}_id", prop);
        let mut found = false;
        for rel in db.relations.values() {
          if (rel.source_type == *pt && rel.destination_type == entity_type)
            || (rel.source_type == entity_type && rel.destination_type == *pt)
          {
            if rel.source_columns.contains(&expected_col) {
              relation_col = Some(expected_col.clone());
              found = true;
              break;
            }
          }
        }
        if !found {
          relation_col = Some(expected_col);
        }
      }

      let mut final_path = current_path.clone();
      if is_polymorphic && !final_path.is_empty() && !final_path.ends_with(&entity_type) {
        if final_path.ends_with(".#") {
          final_path = format!("{}(type==\"{}\")", final_path, entity_type);
        } else {
          final_path = format!("{}#(type==\"{}\")", final_path, entity_type);
        }
      }

      let stem = Stem {
        r#type: entity_type.clone(),
        relation: relation_col,
        schema: Arc::new(schema.clone()),
      };
      inner_map.insert(final_path, Arc::new(stem));
    }

    let next_parent = if is_entity {
      Some(entity_type.clone())
    } else {
      parent_type.clone()
    };

    let pass_seen = if is_entity {
      let mut ns = seen_entities.clone();
      ns.push(entity_type.clone());
      ns
    } else {
      seen_entities.clone()
    };

    // Properties branch
    if let Some(props) = &schema.obj.properties {
      for (k, v) in props {
        // Standard Property Pathing
        let next_path = if current_path.is_empty() {
          k.clone()
        } else {
          format!("{}.{}", current_path, k)
        };

        Self::discover_stems(
          db,
          root_schema_id,
          v,
          next_path,
          next_parent.clone(),
          Some(k.clone()),
          false,
          inner_map,
          pass_seen.clone(),
          errors,
        );
      }
    }

    // Array Item branch
    if let Some(items) = &schema.obj.items {
      let next_path = if current_path.is_empty() {
        String::from("#")
      } else {
        format!("{}.#", current_path)
      };

      Self::discover_stems(
        db,
        root_schema_id,
        items,
        next_path,
        next_parent.clone(),
        property_name.clone(),
        false,
        inner_map,
        pass_seen.clone(),
        errors,
      );
    }

    // Follow external reference if we didn't just crawl local properties
    if schema.obj.properties.is_none() && schema.obj.items.is_none() && schema.obj.one_of.is_none()
    {
      if let Some(ref r) = schema.obj.r#ref {
        if let Some(target_schema) = db.schemas.get(r) {
          Self::discover_stems(
            db,
            root_schema_id,
            target_schema,
            current_path.clone(),
            next_parent.clone(),
            property_name.clone(),
            is_polymorphic,
            inner_map,
            seen_entities.clone(),
            errors,
          );
        }
      }
    }

    // Polymorphism branch
    if let Some(arr) = &schema.obj.one_of {
      for v in arr {
        Self::discover_stems(
          db,
          root_schema_id,
          v.as_ref(),
          current_path.clone(),
          next_parent.clone(),
          property_name.clone(),
          true,
          inner_map,
          pass_seen.clone(),
          errors,
        );
      }
    }
    if let Some(arr) = &schema.obj.all_of {
      for v in arr {
        Self::discover_stems(
          db,
          root_schema_id,
          v.as_ref(),
          current_path.clone(),
          next_parent.clone(),
          property_name.clone(),
          is_polymorphic,
          inner_map,
          pass_seen.clone(),
          errors,
        );
      }
    }
  }
}