laburnum 1.17.1

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

use {
  crate::{
    ContentHash, Ident,
    database::{
      Database, DynPartition, GenerationEpoch, HasPartition, Partition, Partitions, RecordRef,
      query::{PartitionQueryBuilder, QueryBuilder, TypedPartitionQueryBuilder},
      query_results::{QueryResults, RecordMetadata},
    },
    prelude::PartitionKey,
  },
  std::sync::Arc,
};

/// Snapshot-isolated read interface for the compilation database.
///
/// `QueryClient` is the single public API for reading data from the database.
/// It provides three categories of read access:
///
/// # Typed index entry queries
///
/// For partitions with rich index entries (e.g. `SymbolEntry`, `ResolutionEntry`),
/// use the typed index methods to get the full `Part::IndexEntry`:
///
/// ```ignore
/// // Exact lookup
/// let entry: Option<Part::IndexEntry> = client.index_get::<MyPartition>("sort_key");
///
/// // Prefix scan
/// let entries: Vec<(String, Part::IndexEntry)> = client.index_range::<MyPartition>("prefix");
///
/// // Range queries
/// let entries = client.index_less_than::<MyPartition>("value", true);
/// let entries = client.index_greater_than::<MyPartition>("value", false);
/// let entries = client.index_between::<MyPartition>("from", "to");
/// ```
///
/// # Typed CAS record lookup
///
/// To retrieve a content-addressed record by its hash (e.g. following a
/// `RecordHandle` from an index entry):
///
/// ```ignore
/// let record: Option<RecordRef<'_, MyPartition>> = client.get::<MyPartition>(content_hash);
/// ```
///
/// # Query builder API
///
/// For partitions with CAS records, the query builder returns `QueryResults`
/// which resolves index entries to their underlying records:
///
/// ```ignore
/// let results = client.query(MyPartition)
///     .sort_key_begins_with(prefix)
///     .execute()
///     .await;
///
/// for record_ref in results.iter() {
///     // work with resolved P::RecordRef
/// }
/// ```
///
/// Use `query_partition()` for standard partitions (`DynPartition` types)
/// and `query()` for custom partitions (`Partition` types).
#[derive(Clone)]
pub struct QueryClient<P: Partitions> {
  pub(crate) db: Database<P>,
  snapshot_epoch: GenerationEpoch,
}

impl<P: Partitions> QueryClient<P> {
  pub fn new(db: Database<P>) -> Self {
    let snapshot_epoch = db.get_current_epoch();
    Self {
      db,
      snapshot_epoch,
    }
  }

  pub(crate) fn query_any<'a>(
    &'a mut self,
    partition_key: Ident,
  ) -> QueryBuilder<'a, P> {
    QueryBuilder::new(self, partition_key)
  }

  pub fn query_partition<'a, Part: DynPartition + PartitionKey>(
    &'a mut self,
    _partition: Part,
  ) -> PartitionQueryBuilder<'a, P, Part> {
    PartitionQueryBuilder::new(self)
  }

  pub fn query<'a, Part: Partition + PartitionKey>(
    &'a mut self,
    _partition: Part,
  ) -> TypedPartitionQueryBuilder<'a, P, Part> {
    TypedPartitionQueryBuilder::new(self)
  }

  pub async fn get_record(
    &mut self,
    partition_key: Ident,
    sort_key: String,
  ) -> QueryResults<P> {
    self
      .get_record_internal(partition_key, Some(sort_key))
      .await
  }

  /// Query the span index for the record at a given byte offset in a URI.
  ///
  /// Returns the innermost record whose indexed span contains the given
  /// byte offset, or `None` if no record covers that position.
  pub fn span_index_get<D: PartitionKey>(
    &self,
    uri: &crate::Uri,
    byte_offset: u64,
  ) -> Option<P::RecordRef<'_>> {
    let hash = self.db.span_index_query(D::KEY, uri, byte_offset)?;
    self.db.cas.get_any(D::KEY, hash)
  }

  /// Get a CAS record by content hash from a specific partition.
  ///
  /// Use this to resolve a `ContentHash` obtained from one partition's record
  /// into a record in another partition. The type parameter `D` provides the
  /// target partition key at compile time.
  pub fn get_by_hash<D: PartitionKey>(
    &self,
    hash: ContentHash,
  ) -> Option<P::RecordRef<'_>> {
    self.db.cas.get_any(D::KEY, hash)
  }

  pub async fn batch_get_items<T>(
    &mut self,
    items: Vec<(Ident, T)>,
  ) -> QueryResults<P>
  where
    T: Into<String> + Clone,
  {
    let mut all_records = Vec::new();

    for (partition_key, sort_key) in items {
      let results = self
        .get_record_internal(partition_key, Some(sort_key))
        .await;
      all_records.extend(results.records);
    }

    all_records.sort_by(|a, b| a.sort_key.cmp(&b.sort_key));

    QueryResults::new(all_records, Arc::clone(&self.db.cas))
  }

  pub(crate) async fn get_record_internal<T>(
    &mut self,
    partition_key: Ident,
    sort_key: Option<T>,
  ) -> QueryResults<P>
  where
    T: Into<String> + Clone,
  {
    match sort_key {
      | Some(sk) => {
        let sk_string = sk.into();
        match self.db.index_get(partition_key, &sk_string) {
          | Some(content_hash) => {
            let records = vec![RecordMetadata {
              partition_key,
              sort_key: sk_string,
              content_hash,
            }];
            QueryResults::new(records, Arc::clone(&self.db.cas))
          },
          | None => {
            QueryResults::new(Vec::new(), Arc::clone(&self.db.cas))
          },
        }
      },
      | None => {
        let entries = self.db.index_range(partition_key, "");
        if entries.is_empty() {
          return QueryResults::new(Vec::new(), Arc::clone(&self.db.cas));
        }

        let records: Vec<RecordMetadata> = entries
          .into_iter()
          .map(|(_, sort_key, content_hash)| RecordMetadata {
            partition_key,
            sort_key,
            content_hash,
          })
          .collect();

        QueryResults::new(records, Arc::clone(&self.db.cas))
      },
    }
  }

  pub(crate) async fn less_than_internal(
    &mut self,
    partition_key: Ident,
    value: String,
    inclusive: bool,
  ) -> QueryResults<P> {
    let entries = self.db.index_less_than(partition_key, &value, inclusive);

    if entries.is_empty() {
      return QueryResults::new(Vec::new(), Arc::clone(&self.db.cas));
    }

    let records: Vec<RecordMetadata> = entries
      .into_iter()
      .map(|(_, sort_key, content_hash)| RecordMetadata {
        partition_key,
        sort_key,
        content_hash,
      })
      .collect();

    QueryResults::new(records, Arc::clone(&self.db.cas))
  }

  pub(crate) async fn greater_than_internal(
    &mut self,
    partition_key: Ident,
    value: String,
    inclusive: bool,
  ) -> QueryResults<P> {
    let entries = self.db.index_greater_than(partition_key, &value, inclusive);

    if entries.is_empty() {
      return QueryResults::new(Vec::new(), Arc::clone(&self.db.cas));
    }

    let records: Vec<RecordMetadata> = entries
      .into_iter()
      .map(|(_, sort_key, content_hash)| RecordMetadata {
        partition_key,
        sort_key,
        content_hash,
      })
      .collect();

    QueryResults::new(records, Arc::clone(&self.db.cas))
  }

  pub(crate) async fn between_internal(
    &mut self,
    partition_key: Ident,
    from: String,
    to: String,
  ) -> QueryResults<P> {
    let entries = self.db.index_between(partition_key, &from, &to);

    if entries.is_empty() {
      return QueryResults::new(Vec::new(), Arc::clone(&self.db.cas));
    }

    let records: Vec<RecordMetadata> = entries
      .into_iter()
      .map(|(_, sort_key, content_hash)| RecordMetadata {
        partition_key,
        sort_key,
        content_hash,
      })
      .collect();

    QueryResults::new(records, Arc::clone(&self.db.cas))
  }

  pub(crate) async fn prefix_internal(
    &mut self,
    partition_key: Ident,
    prefix: String,
  ) -> QueryResults<P> {
    let entries = self.db.index_range(partition_key, &prefix);

    if entries.is_empty() {
      return QueryResults::new(Vec::new(), Arc::clone(&self.db.cas));
    }

    let records: Vec<RecordMetadata> = entries
      .into_iter()
      .map(|(_, sort_key, content_hash)| RecordMetadata {
        partition_key,
        sort_key,
        content_hash,
      })
      .collect();

    QueryResults::new(records, Arc::clone(&self.db.cas))
  }

  // -- Typed index entry queries ----------------------------------------

  /// Get an index entry by exact sort key, returning the full `Part::IndexEntry`.
  pub fn index_get<Part>(&self, sort_key: &str) -> Option<Part::IndexEntry>
  where
    Part: Partition + PartitionKey,
    P::Stores: HasPartition<Part>,
  {
    <P::Stores as HasPartition<Part>>::store(self.db.cas.stores())
      .index_get(sort_key)
  }

  /// Get all index entries whose sort key starts with `prefix`.
  pub fn index_range<Part>(&self, prefix: &str) -> Vec<(String, Part::IndexEntry)>
  where
    Part: Partition + PartitionKey,
    P::Stores: HasPartition<Part>,
  {
    <P::Stores as HasPartition<Part>>::store(self.db.cas.stores())
      .index_range(prefix)
  }

  /// Get all index entries with sort key less than (or equal to) `value`.
  pub fn index_less_than<Part>(
    &self,
    value: &str,
    inclusive: bool,
  ) -> Vec<(String, Part::IndexEntry)>
  where
    Part: Partition + PartitionKey,
    P::Stores: HasPartition<Part>,
  {
    <P::Stores as HasPartition<Part>>::store(self.db.cas.stores())
      .index_less_than(value, inclusive)
  }

  /// Get all index entries with sort key greater than (or equal to) `value`.
  pub fn index_greater_than<Part>(
    &self,
    value: &str,
    inclusive: bool,
  ) -> Vec<(String, Part::IndexEntry)>
  where
    Part: Partition + PartitionKey,
    P::Stores: HasPartition<Part>,
  {
    <P::Stores as HasPartition<Part>>::store(self.db.cas.stores())
      .index_greater_than(value, inclusive)
  }

  /// Get all index entries with sort key between `from` and `to` (inclusive).
  pub fn index_between<Part>(
    &self,
    from: &str,
    to: &str,
  ) -> Vec<(String, Part::IndexEntry)>
  where
    Part: Partition + PartitionKey,
    P::Stores: HasPartition<Part>,
  {
    <P::Stores as HasPartition<Part>>::store(self.db.cas.stores())
      .index_between(from, to)
  }

  // -- Typed CAS record lookup ------------------------------------------

  /// Get a CAS record by content hash from a specific partition.
  ///
  /// Use this to follow a `RecordHandle` from an index entry to the
  /// underlying record.
  pub fn get<Part>(&self, hash: ContentHash) -> Option<RecordRef<'_, Part>>
  where
    Part: Partition + 'static,
    P::Stores: HasPartition<Part>,
  {
    self.db.cas.get::<Part>(hash)
  }

  pub(crate) fn refresh_snapshot(&mut self) {
    self.snapshot_epoch = self.db.get_current_epoch();
  }
}