laburnum 1.17.1

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

//! Partition types and traits for the content-addressed database.
//!
//! This module contains:
//! - Core `Partition` trait and related types
//! - [`partition_store`] - `PartitionStore` and HList traits for partition storage
//! - [`partition_macros`] - Macros for defining partition configurations
//!
//! # The Partition Pattern
//!
//! Laburnum organizes data into **partitions** - logical groupings of related
//! records identified by a partition key ([`Ident`]). Within each partition,
//! records are ordered by a sort key ([`String`]).
//!
//! The partition pattern has three layers:
//!
//! 1. **Storage Layer** - Generic [`Partitions`] that stores all record
//!    types
//! 2. **Partition Layer** - Type-safe enums for each partition's record types
//! 3. **Application Layer** - Domain-specific readers/writers with ergonomic
//!    APIs
//!
//! ## Why This Pattern?
//!
//! - **Type Safety**: Compile-time verification that records match their
//!   partition
//! - **Encapsulation**: Partition keys stay internal to the implementation
//! - **Discoverability**: Enum variants document all valid record types
//! - **Consistency**: Key formatting centralized in one place
//! - **Testability**: Traits can be mocked for unit testing
//!
//! ## Decoupling Partitions from Storage
//!
//! The [`Partition`] trait says nothing about how records are stored. It only
//! defines the partition key, record type, and sort key type. This enables a
//! clean separation of concerns:
//!
//! - **Server crate**: Implements [`Partitions`] with the actual storage
//!   backend (e.g., SQLite, RocksDB, in-memory).
//! - **Feature crates**: Define [`Partition`] types without any knowledge of
//!   the server crate's storage implementation.
//!
//! This means you can have multiple crates defining partitions that all work
//! with any storage backend, as long as that backend implements
//! [`PartitionReader`] and [`PartitionWriter`] for the partition types it
//! wants to support.
//!
//! # Implementing the Pattern
//!
//! ## Step 1: Define Your Partition Record Enum
//!
//! Create an enum representing all record types for your partition. Use `Arc`
//! for heap-allocated types to enable cheap cloning during reads.
//!
//! ```ignore
//! use std::sync::Arc;
//!
//! /// All record types stored in the symbols partition.
//! #[derive(Debug, Clone)]
//! pub enum SymbolRecord {
//!   /// A symbol definition (heap-allocated).
//!   Definition(Arc<SymbolDefinition>),
//!   /// A reference to a symbol (small, inline).
//!   Reference(SymbolReference),
//! }
//!
//! #[derive(Debug)]
//! pub struct SymbolDefinition {
//!   pub name: String,
//!   pub kind: SymbolKind,
//!   pub span: Span,
//! }
//!
//! #[derive(Debug, Clone, Copy)]
//! pub struct SymbolReference {
//!   pub target_id: u64,
//!   pub line: u32,
//!   pub column: u32,
//! }
//! ```
//!
//! **Design tips:**
//! - Use `Arc<T>` for types containing `String`, `Vec`, or other heap
//!   allocations
//! - Use inline storage for small `Copy` types (like small structs or enums)
//! - This allows cheap `Arc::clone()` on reads instead of deep cloning
//!
//! ## Step 2: Define the Partition Type
//!
//! Create a zero-sized marker type and implement [`Partition`] to bind
//! the partition key, record type, and sort key type together:
//!
//! ```ignore
//! use laburnum::{Ident, Partition};
//!
//! /// Zero-sized marker type for the symbols partition.
//! pub struct Symbols;
//!
//! /// Sort keys for the symbols partition.
//! pub enum SymbolKey {
//!   Definition { source_id: u64, symbol_id: u64 },
//!   Reference { source_id: u64, line: u32, column: u32 },
//! }
//!
//! impl std::fmt::Display for SymbolKey {
//!   fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
//!     match self {
//!       Self::Definition { source_id, symbol_id } => {
//!         write!(f, "def|{:010}|{:010}", source_id, symbol_id)
//!       }
//!       Self::Reference { source_id, line, column } => {
//!         write!(f, "ref|{:010}|{:05}|{:05}", source_id, line, column)
//!       }
//!     }
//!   }
//! }
//!
//! impl Partition for Symbols {
//!   const KEY: Ident = Ident::new("symbols");
//!   type Record = SymbolRecord;
//!   type SortKey = SymbolKey;
//! }
//! ```
//!
//! ## Step 3: Implement PartitionWriter
//!
//! Create a writer struct that wraps
//! [`RecordWriter`](crate::database::chunk::RecordWriter) and implements
//! [`PartitionWriter`] for your partition type:
//!
//! ```ignore
//! use laburnum::PartitionWriter;
//! use laburnum::database::chunk::RecordWriter;
//!
//! pub struct SymbolWriter<'a> {
//!   writer: &'a mut RecordWriter<MyStorage>,
//! }
//!
//! impl<'a> SymbolWriter<'a> {
//!   pub fn new(writer: &'a mut RecordWriter<MyStorage>) -> Self {
//!     Self { writer }
//!   }
//! }
//!
//! impl PartitionWriter<Symbols> for SymbolWriter<'_> {
//!   fn write(&mut self, sort_key: SymbolKey, record: SymbolRecord) {
//!     let storage_record = match record {
//!       SymbolRecord::Definition(def) => MyRecord::SymbolDef(def),
//!       SymbolRecord::Reference(reference) => MyRecord::SymbolRef(reference),
//!     };
//!     self.writer.insert(Symbols::KEY, sort_key, storage_record);
//!   }
//! }
//! ```
//!
//! ## Step 4: Implement PartitionReader
//!
//! Implement [`PartitionReader`] on your storage type to convert storage
//! record refs back to partition records:
//!
//! ```ignore
//! use laburnum::PartitionReader;
//! use std::sync::Arc;
//!
//! impl PartitionReader<Symbols> for MyStorage {
//!   fn as_record(record_ref: &Self::RecordRef<'_>) -> Option<SymbolRecord> {
//!     match record_ref {
//!       MyRecordRef::SymbolDef(def) => {
//!         Some(SymbolRecord::Definition(Arc::clone(def)))
//!       },
//!       MyRecordRef::SymbolRef(reference) => {
//!         Some(SymbolRecord::Reference(*reference))
//!       },
//!       _ => None,
//!     }
//!   }
//! }
//! ```
//!
//! ## Step 5: Create a Domain-Specific Reader (Optional)
//!
//! For ergonomic querying, create a reader struct with typed methods:
//!
//! ```ignore
//! use laburnum::database::query::QueryClient;
//!
//! pub struct SymbolReader<'a, S: PartitionReader<Symbols>> {
//!   client: &'a mut QueryClient<P>,
//! }
//!
//! impl<'a, S: PartitionReader<Symbols>> SymbolReader<'a, P> {
//!   pub fn new(client: &'a mut QueryClient<P>) -> Self {
//!     Self { client }
//!   }
//!
//!   pub async fn all_definitions(&mut self) -> Vec<Arc<SymbolDefinition>> {
//!     let results = self.client
//!       .query(Symbols)
//!       .sort_key_begins_with("def|")
//!       .execute()
//!       .await;
//!
//!     results.iter()
//!       .filter_map(|(_, r)| S::as_record(&r))
//!       .filter_map(|r| match r {
//!         SymbolRecord::Definition(def) => Some(def),
//!         _ => None,
//!       })
//!       .collect()
//!   }
//!
//!   pub async fn references_in_source(&mut self, source_id: u64) -> Vec<SymbolReference> {
//!     let prefix = format!("ref|{:010}|", source_id);
//!     let results = self.client
//!       .query(Symbols)
//!       .sort_key_begins_with(&prefix)
//!       .execute()
//!       .await;
//!
//!     results.iter()
//!       .filter_map(|(_, r)| S::as_record(&r))
//!       .filter_map(|r| match r {
//!         SymbolRecord::Reference(reference) => Some(reference),
//!         _ => None,
//!       })
//!       .collect()
//!   }
//! }
//! ```
//!
//! # Usage Example
//!
//! ```ignore
//! // Writing records - partition key derived from Symbols::KEY
//! let mut writer = SymbolWriter::new(record_writer);
//!
//! writer.write(
//!   SymbolKey::Definition { source_id, symbol_id },
//!   SymbolRecord::Definition(Arc::new(def)),
//! );
//!
//! writer.write(
//!   SymbolKey::Reference { source_id, line, column },
//!   SymbolRecord::Reference(reference),
//! );
//!
//! // Reading records
//! let mut reader = SymbolReader::new(query_client);
//! let definitions = reader.all_definitions().await;
//! let refs = reader.references_in_source(source_id).await;
//! ```

pub mod partition_macros;
pub mod partition_store;
pub mod span_index;

pub use partition_store::{
  BluegumStores, CascadingRefCollector, CollectCascadingRefs,
  CollectIndexHashes, HasIndexPartition, HasPartition, HasPartitionAt,
  MergeFrom, MergeFromWithRefcount, MergeResult, NewStores, PartitionStore,
  PartitionStoreInner, RecordRef, RefcountOps,
};
pub use span_index::SpanIndex;

use crate::{
  Ident,
  database::{chunk::RecordWriter, storage::Partitions},
};

pub trait PartitionKey {
  const KEY: Ident;
  /// Human-readable name for debug/display purposes.
  const KEY_NAME: Option<&'static str> = None;
}

pub trait Partition: PartitionKey + Send + Sync {
  /// Content-addressed record type. Deduplicated by ContentHash.
  type Record: crate::record::Record + bluegum::BluegumWithState<dyn crate::SpanResolver>;

  /// Index entry type. Keyed by SortKey, not content-addressed.
  /// Must declare which CAS records it references for GC.
  ///
  /// Use `HandleEntry<Self>` for simple partitions that just need a handle.
  /// Use `()` for index-free partitions (CAS-only).
  type IndexEntry: IndexEntry + bluegum::BluegumWithState<dyn crate::SpanResolver>;

  /// Sort key type for index entries.
  type SortKey: std::fmt::Display;

  /// Create an index entry from just a record handle and content hash.
  ///
  /// This is used by the `define_partitions!` macro for dynamic index insertion
  /// (e.g., during chunk commit for entries added via `RecordWriter::index()`).
  ///
  /// # Default Implementation
  ///
  /// The default implementation panics. This is appropriate for:
  /// - **Index-only partitions** with rich `IndexEntry` types (like `SymbolEntry`)
  ///   that require additional metadata (spans, etc.) not available from just a handle
  /// - These partitions should use `RecordWriter::index_entry()` instead of `index()`
  ///
  /// # When to Override
  ///
  /// Override this method if your partition supports creating index entries from
  /// just a content hash:
  /// - Partitions with `IndexEntry = ()` (CAS-only) - return `()`
  /// - Partitions with `IndexEntry = HandleEntry<Self>` - return `HandleEntry::new(handle)`
  ///
  /// # Example
  ///
  /// ```ignore
  /// impl Partition for MySimplePartition {
  ///     // ...
  ///     fn index_entry_from_handle(handle: RecordHandle<Self>) -> Self::IndexEntry {
  ///         HandleEntry::new(handle)
  ///     }
  /// }
  /// ```
  fn index_entry_from_handle(
    handle: crate::database::RecordHandle<Self>,
  ) -> Self::IndexEntry
  where
    Self: Sized,
  {
    let _ = handle;
    unreachable!(
      "Partition {} does not support creating index entries from just a handle. \
       Use RecordWriter::index_entry() instead of index() for this partition.",
      Self::KEY
    )
  }
}

/// Marker trait for index entry types.
///
/// Index entries must implement `CollectReferences<P>` to declare
/// which CAS records they reference (for GC refcounting).
///
/// Index entries represent "where" and "which" - they say "at this semantic
/// location, this CAS record is currently live." They may carry position-specific
/// metadata like spans that only make sense for a particular source version.
///
/// # GC Role
///
/// Index entries are GC roots. Their referenced CAS records are kept alive
/// via reference counting. When an index entry is overwritten or removed,
/// the old entry's references are decremented.
pub trait IndexEntry: Clone + Send + Sync {
  /// Get the primary content hash referenced by this index entry.
  ///
  /// This is used for GC root collection. For simple entries like `HandleEntry`,
  /// this is just the handle's hash. For richer entries, this should return
  /// the main referenced record's hash.
  ///
  /// Returns `None` for index entries that don't reference any CAS records
  /// (e.g., `()` for index-free partitions).
  fn primary_hash(&self) -> Option<crate::ContentHash> {
    None
  }
}

// Unit type implements IndexEntry for index-free partitions
impl IndexEntry for () {}

/// Marker trait for index entries that participate in a sort_key index.
///
/// Implementing `IndexedEntry` declares that this entry references a CAS
/// record by content hash and can be inserted via `RecordWriter::index_entry`.
/// The unit type `()` (used by CAS-only partitions) deliberately does NOT
/// implement this trait, so calling `index_entry` on a CAS-only partition is
/// a compile error.
pub trait IndexedEntry: IndexEntry {
  /// Content hash of the CAS record this entry references.
  fn content_hash(&self) -> crate::ContentHash;
}

/// A simple index entry that just points to a CAS record.
///
/// This type provides backward compatibility for partitions that don't need
/// rich index entry data - they just want a handle to a record.
///
/// # Example
///
/// ```ignore
/// impl Partition for MyPartition {
///     type Record = MyRecord;
///     type IndexEntry = HandleEntry<Self>;  // Simple handle-only entry
///     type SortKey = String;
/// }
///
/// // Now insert() works the same as before
/// writer.insert::<MyPartition>(sort_key, record);
/// ```
#[derive(Debug)]
pub struct HandleEntry<P: PartitionKey> {
  handle: crate::database::RecordHandle<P>,
}

impl<P: PartitionKey> HandleEntry<P> {
  /// Create a new handle entry.
  pub fn new(handle: crate::database::RecordHandle<P>) -> Self {
    Self { handle }
  }

  /// Get the underlying record handle.
  pub fn handle(&self) -> crate::database::RecordHandle<P> {
    self.handle
  }

  /// Get the content hash of the referenced record.
  pub fn content_hash(&self) -> crate::ContentHash {
    self.handle.content_hash()
  }
}

impl<P: PartitionKey> Clone for HandleEntry<P> {
  fn clone(&self) -> Self {
    *self
  }
}

impl<P: PartitionKey> Copy for HandleEntry<P> {}

impl<P: PartitionKey> PartialEq for HandleEntry<P> {
  fn eq(&self, other: &Self) -> bool {
    self.handle == other.handle
  }
}

impl<P: PartitionKey> Eq for HandleEntry<P> {}

impl<P: PartitionKey> std::hash::Hash for HandleEntry<P> {
  fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
    self.handle.hash(state);
  }
}

impl<P: PartitionKey + Send + Sync> IndexEntry for HandleEntry<P> {
  fn primary_hash(&self) -> Option<crate::ContentHash> {
    Some(self.handle.content_hash())
  }
}

impl<P: PartitionKey + Send + Sync> IndexedEntry for HandleEntry<P> {
  fn content_hash(&self) -> crate::ContentHash {
    self.handle.content_hash()
  }
}

impl<Ps, P> crate::record::CollectReferences<Ps> for HandleEntry<P>
where
  Ps: crate::database::storage::Partitions,
  P: Partition,
  Ps::Stores: HasPartition<P>,
{
  fn collect_references<R: crate::record::References<Ps>>(&self, refs: &mut R) {
    refs.add(self.handle);
  }
}

impl<P: PartitionKey> From<crate::database::RecordHandle<P>>
  for HandleEntry<P>
{
  fn from(handle: crate::database::RecordHandle<P>) -> Self {
    Self { handle }
  }
}

impl<P: PartitionKey> bluegum::Bluegum for HandleEntry<P> {
  fn node(&self, b: &mut bluegum::Builder) {
    b.name("HandleEntry")
      .field("hash", self.handle.content_hash());
  }
}

impl<P: PartitionKey> bluegum::BluegumWithState<dyn crate::SpanResolver> for HandleEntry<P> {}

impl bluegum::BluegumWithState<dyn crate::SpanResolver> for () {}

impl<P> DynPartition for P
where
  P: Partition + PartitionKey,
{
  type DynSortKey = P::SortKey;
  type RecordConstraint = P::Record;
}

pub trait PartitionWriter<P: Partition + PartitionKey> {
  fn write(&mut self, sort_key: P::SortKey, record: P::Record);
}

pub trait PartitionReader<Part: Partition + PartitionKey>: Partitions {
  fn as_record(record_ref: &Self::RecordRef<'_>) -> Option<Part::Record>;
}

pub trait DynPartition: PartitionKey {
  type RecordConstraint: ?Sized;
  type DynSortKey: std::fmt::Display;
}

pub trait DynPartitionRecord<P: DynPartition>: crate::record::Record {}

#[macro_export]
macro_rules! impl_partition_for_dyn {
  ($wrapper:ident, $dyn_partition:ty, $record:ty) => {
    pub struct $wrapper;

    const _: () = {
      fn _assert_record_implements_trait<
        T: $crate::database::partitions::DynPartitionRecord<$dyn_partition>,
      >() {
      }
      fn _check() {
        _assert_record_implements_trait::<$record>();
      }
    };

    impl $crate::database::partitions::PartitionKey for $wrapper {
      const KEY: $crate::Ident =
        <$dyn_partition as $crate::database::partitions::PartitionKey>::KEY;
      const KEY_NAME: Option<&'static str> =
        <$dyn_partition as $crate::database::partitions::PartitionKey>::KEY_NAME;
    }
    impl $crate::database::partitions::Partition for $wrapper {
      type Record = $record;
      type IndexEntry = $crate::database::partitions::HandleEntry<Self>;
      type SortKey =
        <$dyn_partition as $crate::database::partitions::DynPartition>::DynSortKey;

      fn index_entry_from_handle(
        handle: $crate::database::RecordHandle<Self>,
      ) -> Self::IndexEntry {
        $crate::database::partitions::HandleEntry::new(handle)
      }
    }
  };
}

/// Internal owned wrapper for RecordWriter used in async contexts.
///
/// This type is `pub(crate)` and should not be exposed in public APIs.
/// Use `as_ref()` to get a `PartitionWriteContextRef` for actual write operations.
pub(crate) struct InternalPartitionWriteContext<P: Partitions> {
  writer: RecordWriter<P>,
}

impl<P: Partitions> InternalPartitionWriteContext<P> {
  pub(crate) fn new(writer: RecordWriter<P>) -> Self {
    Self { writer }
  }

  pub(crate) fn into_inner(self) -> RecordWriter<P> {
    self.writer
  }

  /// Attach a constructed Source to be committed alongside the chunk's records.
  pub(crate) fn set_source(
    &mut self,
    key: crate::SourceKey,
    source: crate::Source,
  ) {
    self.writer.set_source(key, source);
  }

  /// Get a mutable reference for write operations.
  pub fn as_ref(&mut self) -> PartitionWriteContextRef<'_, P> {
    PartitionWriteContextRef::new(&mut self.writer)
  }
}

pub struct PartitionWriteContextRef<'a, P: Partitions> {
  writer: &'a mut RecordWriter<P>,
  blocked_partition: Option<Ident>,
}

impl<'a, P: Partitions> PartitionWriteContextRef<'a, P> {
  #[cfg(any(test, feature = "test"))]
  pub fn new(writer: &'a mut RecordWriter<P>) -> Self {
    Self {
      writer,
      blocked_partition: None,
    }
  }

  #[cfg(not(any(test, feature = "test")))]
  pub(crate) fn new(writer: &'a mut RecordWriter<P>) -> Self {
    Self {
      writer,
      blocked_partition: None,
    }
  }

  /// Create a write context for a watcher handler that blocks writes to the
  /// partition that triggered the watcher, preventing infinite loops.
  pub(crate) fn new_for_watcher(
    writer: &'a mut RecordWriter<P>,
    trigger_partition: Ident,
  ) -> Self {
    Self {
      writer,
      blocked_partition: Some(trigger_partition),
    }
  }

  /// Returns `true` if the write is allowed, `false` if it targets the
  /// blocked (trigger) partition.
  fn check_write_allowed(&self, target: Ident) -> bool {
    if let Some(blocked) = self.blocked_partition
      && target == blocked
    {
      debug_assert!(
        false,
        "Watcher attempted to write back to its trigger partition '{}'",
        blocked
      );
      otel::error!(
        "watcher.blocked_self_write",
        format!(
          "Watcher attempted to write back to its trigger partition '{}'. Write skipped.",
          blocked
        )
      );
      return false;
    }
    true
  }

  pub fn clear_prefix(
    &mut self,
    partition_key: Ident,
    prefix: impl std::fmt::Display,
  ) {
    if !self.check_write_allowed(partition_key) {
      return;
    }
    self.writer.clear_prefix(partition_key, prefix);
  }

  pub fn write<Part>(&mut self, sort_key: Part::SortKey, record: Part::Record)
  where
    Part: Partition + PartitionKey + 'static,
    P::Stores: HasPartition<Part>,
  {
    if !self.check_write_allowed(Part::KEY) {
      return;
    }
    self.writer.insert::<Part, _>(sort_key, record);
  }

  /// Store a CAS record without creating an index entry.
  ///
  /// Use this when you need to store content-addressed data that will be
  /// referenced by other records or index entries, but doesn't need its
  /// own top-level index entry.
  pub fn store<Part>(
    &mut self,
    record: Part::Record,
  ) -> crate::database::RecordHandle<Part>
  where
    Part: Partition,
    P::Stores: HasPartition<Part>,
  {
    if !self.check_write_allowed(Part::KEY) {
      return crate::database::RecordHandle::new(crate::ContentHash::ZERO);
    }
    self.writer.store::<Part>(record)
  }

  /// Stage a span interval for the span index.
  ///
  /// At commit time, staged intervals are resolved to byte offsets and
  /// built into a `SpanIndex` per URI per partition.
  pub fn index_span<Part>(
    &mut self,
    span: crate::Span,
    content_hash: crate::ContentHash,
  ) where
    Part: Partition + PartitionKey + 'static,
  {
    self.writer.index_span::<Part>(span, content_hash);
  }

  /// Create a rich index entry directly in the partition store.
  ///
  /// Use this when the partition's `IndexEntry` type needs more than just a
  /// `RecordHandle<Part>` - for example, when it includes span information.
  pub fn index_entry<Part>(
    &mut self,
    sort_key: impl std::fmt::Display,
    entry: Part::IndexEntry,
  ) where
    Part: Partition,
    Part::IndexEntry: IndexedEntry,
    P::Stores: HasPartition<Part>,
  {
    if !self.check_write_allowed(Part::KEY) {
      return;
    }
    self.writer.index_entry::<Part, _>(sort_key, entry);
  }
}

#[cfg(test)]
mod tests {
  use super::*;

  #[test]
  fn dyn_partition_record_blanket_impl_works() {
    use crate::{
      ContentHash,
      partitions::{
        DocumentFoldingRange, document_folding_range::FoldingRangeRecord,
      },
      record::Record,
    };

    #[derive(Debug, Hash)]
    struct TestFoldingRange;

    impl Record for TestFoldingRange {
      fn content_hash(&self) -> ContentHash {
        ContentHash::new(&[])
      }
    }

    impl FoldingRangeRecord for TestFoldingRange {
      fn source_key(&self) -> crate::SourceKey {
        crate::SourceKey::new(0, 0)
      }

      fn to_folding_range(&self) -> crate::protocol::lsp::FoldingRange {
        crate::protocol::lsp::FoldingRange {
          start_line: 0,
          start_character: None,
          end_line: 1,
          end_character: None,
          kind: None,
          collapsed_text: None,
        }
      }
    }

    fn assert_dyn_partition_record<
      P: DynPartition,
      R: DynPartitionRecord<P>,
    >() {
    }

    assert_dyn_partition_record::<DocumentFoldingRange, TestFoldingRange>();
  }

  mod watcher_write_guard {
    use crate::{
      Ident,
      database::{
        chunk::RecordWriter,
        partitions::PartitionWriteContextRef,
        tests::storage::{
          Test1Partition, Test2Partition, TestPartitions, TestRecordData,
        },
      },
    };

    fn make_test_record() -> TestRecordData {
      TestRecordData::Module {
        exports: vec![Ident::new("test")],
      }
    }

    #[test]
    #[should_panic(
      expected = "Watcher attempted to write back to its trigger partition"
    )]
    fn watcher_context_blocks_self_write() {
      let pk1 =
        <Test1Partition as crate::database::partitions::PartitionKey>::KEY;
      let mut writer = RecordWriter::<TestPartitions>::new(pk1);
      let mut ctx = PartitionWriteContextRef::new_for_watcher(&mut writer, pk1);

      ctx.write::<Test1Partition>("key".to_string(), make_test_record());
    }

    #[test]
    fn watcher_context_allows_write_to_other_partition() {
      let pk1 =
        <Test1Partition as crate::database::partitions::PartitionKey>::KEY;
      let mut writer = RecordWriter::<TestPartitions>::new(pk1);
      let mut ctx = PartitionWriteContextRef::new_for_watcher(&mut writer, pk1);

      // Writing to Test2Partition (pk2) should succeed when pk1 is blocked
      ctx.write::<Test2Partition>("key".to_string(), make_test_record());
    }

    #[test]
    fn non_watcher_context_allows_all_writes() {
      let pk1 =
        <Test1Partition as crate::database::partitions::PartitionKey>::KEY;
      let mut writer = RecordWriter::<TestPartitions>::new(pk1);
      let mut ctx = PartitionWriteContextRef::new(&mut writer);

      ctx.write::<Test1Partition>("key1".to_string(), make_test_record());
      ctx.write::<Test2Partition>("key2".to_string(), make_test_record());
    }

    #[test]
    #[should_panic(
      expected = "Watcher attempted to write back to its trigger partition"
    )]
    fn watcher_context_blocks_store_to_trigger_partition() {
      let pk1 =
        <Test1Partition as crate::database::partitions::PartitionKey>::KEY;
      let mut writer = RecordWriter::<TestPartitions>::new(pk1);
      let mut ctx = PartitionWriteContextRef::new_for_watcher(&mut writer, pk1);

      ctx.store::<Test1Partition>(make_test_record());
    }

    #[test]
    #[should_panic(
      expected = "Watcher attempted to write back to its trigger partition"
    )]
    fn watcher_context_blocks_clear_prefix_on_trigger_partition() {
      let pk1 =
        <Test1Partition as crate::database::partitions::PartitionKey>::KEY;
      let mut writer = RecordWriter::<TestPartitions>::new(pk1);
      let mut ctx = PartitionWriteContextRef::new_for_watcher(&mut writer, pk1);

      ctx.clear_prefix(pk1, "prefix");
    }
  }
}